Add predicted (x, y) values to Results

docs/source/user_manual/output_files.rst

@@ -38,6 +38,10 @@ The mapped RA, dec information consists of up to four columns. The columns `glob
 
 The columns `img_ra` and `img_dec` indicate the positions in the original images. These could be the same or different from the global (RA, dec) even for reprojected images. If the reprojection consists of aligning the images, such as correcting for rotation, the coordinates will be the same. In that case, the RA and dec are not actually changing, just the mappping from RA, dec to pixels. However if the reprojection includes a shift of the viewing location, such as with the barycentric reprojection, we would expect the RA and dec to also change.
 
+**Predicted x, y Information**
+
+KBMOD will also listed the predicted (x, y) pixel coordinates of the object for each time step. The columns `pred_x` and `pred_y` list the predicted x and y positions in the common WCS frame that KBMOD used for the search.  The columns `img_x` and `img_y` list the predicted x and y positions in each image's original WCS frame. The `img_` columns may be identical to the `pred_` columns if the images were not reprojected.
+
 **Metadata**
 
 The table also includes some basic metadata about the set of images, including the number of images (`num_img`), the image dimensions (`dims`), and the midpoint times of the observations (`mid_mjd`).

src/kbmod/run_search.py

@@ -270,7 +270,7 @@ def run_search(
         # including a global WCS and per-time (RA, dec) predictions for each image.
         if workunit is not None:
             keep.table.wcs = workunit.wcs
-            append_ra_dec_to_results(workunit, keep)
+            append_positions_to_results(workunit, keep)
 
         # Create and save any additional meta data that should be saved with the results.
         num_img = stack.img_count()
@@ -322,8 +322,9 @@ def run_search_from_work_unit(self, work):
         )
 
 
-def append_ra_dec_to_results(workunit, results):
-    """Append predicted (RA, dec) positions to the results.
+def append_positions_to_results(workunit, results):
+    """Append predicted (x, y) and (RA, dec) positions in the original images. If
+    the images were reprojected, also appends the (RA, dec) in the common frame.
 
     Parameters
     ----------
@@ -340,51 +341,69 @@ def append_ra_dec_to_results(workunit, results):
     num_times = workunit.im_stack.img_count()
     times = workunit.im_stack.build_zeroed_times()
 
-    # Predict where each candidate trajectory will be at each time step.
+    # Predict where each candidate trajectory will be at each time step in the
+    # common WCS frame. These are the pixel locations used to assess the trajectory.
     xp = predict_pixel_locations(times, results["x"], results["vx"], as_int=False)
     yp = predict_pixel_locations(times, results["y"], results["vy"], as_int=False)
-
-    # Compute the predicted (RA, dec) positions for each trajectory in global space.
+    results.table["pred_x"] = xp
+    results.table["pred_y"] = yp
+
+    # Compute the predicted (RA, dec) positions for each trajectory the common WCS
+    # frame and original image WCS frames.
+    all_inds = np.arange(num_times)
+    all_ra = np.zeros((len(results), num_times))
+    all_dec = np.zeros((len(results), num_times))
     if workunit.wcs is not None:
         logger.info("Found common WCS. Adding global_ra and global_dec columns.")
 
+        # Compute the (RA, dec) for each result x time in the common WCS frame.
         skypos = workunit.wcs.pixel_to_world(xp, yp)
         results.table["global_ra"] = skypos.ra.degree
         results.table["global_dec"] = skypos.dec.degree
 
-        # Loop over the trajectories to build the original positions.
-        all_ra = []
-        all_dec = []
+        # Loop over the trajectories to build the (RA, dec) positions in each image's WCS frame.
         for idx in range(num_results):
-            pos_tuples = [(xp[idx, j], yp[idx, j]) for j in range(num_times)]
-            skypos = workunit.image_positions_to_original_icrs(
-                image_indices=np.arange(num_times),  # Compute for all times.
-                positions=pos_tuples,
-                input_format="xy",
+            # Build a list of this trajectory's RA, dec position at each time.
+            pos_list = [skypos[idx, j] for j in range(num_times)]
+            img_skypos = workunit.image_positions_to_original_icrs(
+                image_indices=all_inds,  # Compute for all times.
+                positions=pos_list,
+                input_format="radec",
                 output_format="radec",
                 filter_in_frame=False,
             )
 
             # We get back a list of SkyCoord, because we gave a list.
             # So we flatten it and extract the coordinate values.
-            all_ra.append([skypos[j].ra.degree for j in range(num_times)])
-            all_dec.append([skypos[j].dec.degree for j in range(num_times)])
+            for time_idx in range(num_times):
+                all_ra[idx, time_idx] = img_skypos[time_idx].ra.degree
+                all_dec[idx, time_idx] = img_skypos[time_idx].dec.degree
 
-        results.table["img_ra"] = all_ra
-        results.table["img_dec"] = all_dec
     else:
         logger.info("No common WCS found. Skipping global_ra and global_dec columns.")
 
         # If there are no global WCS, we just predict per image.
-        all_ra = np.zeros((len(results), num_times))
-        all_dec = np.zeros((len(results), num_times))
-
         for time_idx in range(num_times):
             wcs = workunit.get_wcs(time_idx)
             if wcs is not None:
                 skypos = wcs.pixel_to_world(xp[:, time_idx], yp[:, time_idx])
                 all_ra[:, time_idx] = skypos.ra.degree
                 all_dec[:, time_idx] = skypos.dec.degree
 
-        results.table["img_ra"] = all_ra
-        results.table["img_dec"] = all_dec
+    # Add the per-image coordinates to the results table.
+    results.table["img_ra"] = all_ra
+    results.table["img_dec"] = all_dec
+
+    # If we have have per-image WCSes, compute the pixel location in the original image.
+    if "per_image_wcs" in workunit.org_img_meta.colnames:
+        img_x = np.zeros((len(results), num_times))
+        img_y = np.zeros((len(results), num_times))
+        for time_idx in range(num_times):
+            wcs = workunit.org_img_meta["per_image_wcs"][time_idx]
+            if wcs is not None:
+                xy_pos = wcs.world_to_pixel_values(all_ra[:, time_idx], all_dec[:, time_idx])
+                img_x[:, time_idx] = xy_pos[0]
+                img_y[:, time_idx] = xy_pos[1]
+
+        results.table["img_x"] = img_x
+        results.table["img_y"] = img_y

src/kbmod/work_unit.py

@@ -89,7 +89,7 @@ class WorkUnit:
     per_image_wcs : `list`, optional
         A list with one WCS for each image in the WorkUnit. Used for when
         the images have *not* been standardized to the same pixel space. If provided
-        this will the WCS values in org_image_meta
+        this will overwrite the WCS values in org_image_meta
     reprojected : `bool`, optional
         Whether or not the WorkUnit image data has been reprojected.
     reprojection_frame : `str`, optional

tests/test_run_search.py

@@ -1,45 +1,74 @@
 """Test some of the functions needed for running the search."""
 
+from astropy.coordinates import EarthLocation, SkyCoord
+from astropy.table import Table
+from astropy.time import Time
+
 import unittest
 
 import numpy as np
 
 from kbmod.configuration import SearchConfiguration
 from kbmod.fake_data.fake_data_creator import create_fake_times, FakeDataSet
+from kbmod.reprojection_utils import fit_barycentric_wcs
 from kbmod.results import Results
-from kbmod.run_search import append_ra_dec_to_results
+from kbmod.run_search import append_positions_to_results
 from kbmod.search import *
 from kbmod.wcs_utils import make_fake_wcs
 from kbmod.work_unit import WorkUnit
 
 
 class test_run_search(unittest.TestCase):
-    def test_append_ra_dec_global(self):
+    def test_append_positions_to_results_global(self):
         # Create a fake WorkUnit with 20 times, a completely random ImageStack,
         # and no trajectories.
         num_times = 20
-        fake_times = create_fake_times(num_times, t0=60676.0)
-        fake_ds = FakeDataSet(800, 600, fake_times)
+        width = 800
+        height = 600
+        t0 = 60676.0
+        helio_dist = 500.0
 
-        # Append a global fake WCS and one for each time.
+        fake_times = create_fake_times(num_times, t0=t0)
+        fake_ds = FakeDataSet(width, height, fake_times)
+
+        # Create a global fake WCS,  one for each time (slightly shifted), and the EBD information.
         global_wcs = make_fake_wcs(20.0, 0.0, 800, 600, deg_per_pixel=0.5 / 3600.0)
-        all_wcs = []
+
+        per_image_wcs = []
         for idx in range(num_times):
+            # Each WCS is slight shifted from the global one.
             curr = make_fake_wcs(
-                20.01 + idx / 100.0, 0.01 + idx / 100.0, 800, 600, deg_per_pixel=0.5 / 3600.0
+                20.001 + idx / 1000.0, 0.001 + idx / 1000.0, 800, 600, deg_per_pixel=0.5 / 3600.0
             )
-            all_wcs.append(curr)
+            per_image_wcs.append(curr)
+
+        ebd_wcs, geo_dist = fit_barycentric_wcs(
+            global_wcs,
+            width,
+            height,
+            helio_dist,
+            Time(t0, format="mjd"),
+            EarthLocation.of_site("ctio"),
+        )
 
+        # Create the fake WorkUnit with this information.
+        org_image_meta = Table(
+            {
+                "ebd_wcs": np.array([ebd_wcs] * num_times),
+                "geocentric_distance": np.array([geo_dist] * num_times),
+                "per_image_wcs": np.array(per_image_wcs),
+            }
+        )
         fake_wu = WorkUnit(
             im_stack=fake_ds.stack,
             config=SearchConfiguration(),
-            wcs=global_wcs,
-            per_image_wcs=all_wcs,
+            wcs=ebd_wcs,
             reprojected=True,
             reprojection_frame="ebd",
             per_image_indices=[i for i in range(num_times)],
-            heliocentric_distance=np.full(num_times, 100.0),
+            heliocentric_distance=helio_dist,
             obstimes=fake_times,
+            org_image_meta=org_image_meta,
         )
 
         # Create three fake trajectories in the bounds of the images. We don't
@@ -52,7 +81,7 @@ def test_append_ra_dec_global(self):
         results = Results.from_trajectories(trjs)
         self.assertEqual(len(results), 3)
 
-        append_ra_dec_to_results(fake_wu, results)
+        append_positions_to_results(fake_wu, results)
 
         # The global RA should exist and be close to 20.0 for all observations.
         self.assertEqual(len(results["global_ra"]), 3)
@@ -68,24 +97,45 @@ def test_append_ra_dec_global(self):
             self.assertTrue(np.all(results["global_dec"][i] > -1.0))
             self.assertTrue(np.all(results["global_dec"][i] < 1.0))
 
-        # The per-image RA should exist, be close to 20.0 for all observations,
-        # and be different from the global RA
+        # The per-image RA should exist, be close to (but not the same as)
+        # the global RA for all observations.
         self.assertEqual(len(results["img_ra"]), 3)
         for i in range(3):
             self.assertEqual(len(results["img_ra"][i]), num_times)
-            self.assertTrue(np.all(results["img_ra"][i] > 19.0))
-            self.assertTrue(np.all(results["img_ra"][i] < 21.0))
-            self.assertFalse(np.any(results["img_ra"][i] == results["global_ra"][i]))
+            ra_diffs = np.abs(results["img_ra"][i] - results["global_ra"][i])
+            self.assertTrue(np.all(ra_diffs > 0.0))
+            self.assertTrue(np.all(ra_diffs < 1.0))
 
-        # The global Dec should exist and be close to 0.0 for all observations.
+        # The per-image dec should exist, be close to (but not the same as)
+        # the global dec for all observations.
         self.assertEqual(len(results["img_dec"]), 3)
         for i in range(3):
             self.assertEqual(len(results["img_dec"][i]), num_times)
-            self.assertTrue(np.all(results["img_dec"][i] > -1.0))
-            self.assertTrue(np.all(results["img_dec"][i] < 1.0))
-            self.assertFalse(np.any(results["img_dec"][i] == results["global_dec"][i]))
+            dec_diffs = np.abs(results["img_dec"][i] - results["global_dec"][i])
+            self.assertTrue(np.all(dec_diffs > 0.0))
+            self.assertTrue(np.all(dec_diffs < 1.0))
+
+        # The per-image x should exist and be within some delta of the global predicted x.
+        self.assertEqual(len(results["pred_x"]), 3)
+        self.assertEqual(len(results["img_x"]), 3)
+        for i in range(3):
+            self.assertEqual(len(results["img_x"][i]), num_times)
+            self.assertEqual(len(results["pred_x"][i]), num_times)
+            x_diffs = np.abs(results["img_x"][i] - results["pred_x"][i])
+            self.assertTrue(np.all(x_diffs > 0.0))
+            self.assertTrue(np.all(x_diffs < 1000.0))
+
+        # The per-image y should exist and be within some delta of the global predicted y.
+        self.assertEqual(len(results["pred_y"]), 3)
+        self.assertEqual(len(results["img_y"]), 3)
+        for i in range(3):
+            self.assertEqual(len(results["img_y"][i]), num_times)
+            self.assertEqual(len(results["pred_y"][i]), num_times)
+            y_diffs = np.abs(results["img_y"][i] - results["pred_y"][i])
+            self.assertTrue(np.all(y_diffs > 0.0))
+            self.assertTrue(np.all(y_diffs < 1000.0))
 
-    def test_append_ra_dec_no_global(self):
+    def test_append_positions_to_results_no_global(self):
         # Create a fake WorkUnit with 20 times, a completely random ImageStack,
         # and no trajectories.
         num_times = 20
@@ -120,7 +170,7 @@ def test_append_ra_dec_no_global(self):
         results = Results.from_trajectories(trjs)
         self.assertEqual(len(results), 3)
 
-        append_ra_dec_to_results(fake_wu, results)
+        append_positions_to_results(fake_wu, results)
 
         # The global RA and global dec should not exist.
         self.assertFalse("global_ra" in results.colnames)
@@ -140,6 +190,22 @@ def test_append_ra_dec_no_global(self):
             self.assertTrue(np.all(results["img_dec"][i] > -1.0))
             self.assertTrue(np.all(results["img_dec"][i] < 1.0))
 
+        # The per-image x should exist and the same as global predicted x.
+        self.assertEqual(len(results["pred_x"]), 3)
+        self.assertEqual(len(results["img_x"]), 3)
+        for i in range(3):
+            self.assertEqual(len(results["img_x"][i]), num_times)
+            self.assertEqual(len(results["pred_x"][i]), num_times)
+            self.assertTrue(np.allclose(results["img_x"][i], results["pred_x"][i]))
+
+        # The per-image y should exist and the same as global predicted y.
+        self.assertEqual(len(results["pred_y"]), 3)
+        self.assertEqual(len(results["img_y"]), 3)
+        for i in range(3):
+            self.assertEqual(len(results["img_y"][i]), num_times)
+            self.assertEqual(len(results["pred_y"][i]), num_times)
+            self.assertTrue(np.allclose(results["img_y"][i], results["pred_y"][i]))
+
 
 if __name__ == "__main__":
     unittest.main()