From 01592b5fc577b4339bac1984f0bf0611be0aabbe Mon Sep 17 00:00:00 2001 From: Jeremy Kubica <104161096+jeremykubica@users.noreply.github.com> Date: Fri, 10 Jan 2025 13:09:27 -0500 Subject: [PATCH 1/2] Have KBMOD append RA, dec columns if possible --- src/kbmod/run_search.py | 87 +++++++++++++++++++++-- tests/test_run_search.py | 144 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 225 insertions(+), 6 deletions(-) create mode 100644 tests/test_run_search.py diff --git a/src/kbmod/run_search.py b/src/kbmod/run_search.py index e9309d18..8baf9d35 100644 --- a/src/kbmod/run_search.py +++ b/src/kbmod/run_search.py @@ -12,6 +12,7 @@ from .results import Results from .trajectory_generator import create_trajectory_generator +from .trajectory_utils import predict_pixel_locations from .wcs_utils import wcs_to_dict from .work_unit import WorkUnit @@ -187,7 +188,14 @@ def do_gpu_search(self, config, stack, trj_generator): keep = self.load_and_filter_results(search, config) return keep - def run_search(self, config, stack, trj_generator=None, wcs=None, extra_meta=None): + def run_search( + self, + config, + stack, + trj_generator=None, + workunit=None, + extra_meta=None, + ): """This function serves as the highest-level python interface for starting a KBMOD search given an ImageStack and SearchConfiguration. @@ -200,8 +208,8 @@ def run_search(self, config, stack, trj_generator=None, wcs=None, extra_meta=Non trj_generator : `TrajectoryGenerator`, optional The object to generate the candidate trajectories for each pixel. If None uses the default EclipticCenteredSearch - wcs : `astropy.wcs.WCS`, optional - A global WCS for all images in the search. + workunit : `WorkUnit`, optional + An optional WorkUnit with additional meta-data, including the per-image WCS. extra_meta : `dict`, optional Any additional metadata to save as part of the results file. @@ -253,8 +261,11 @@ def run_search(self, config, stack, trj_generator=None, wcs=None, extra_meta=Non if config["save_all_stamps"]: append_all_stamps(keep, stack, config["stamp_radius"]) - # Append the WCS information if it is provided. This will be saved with the results. - keep.table.wcs = wcs + # Append additional information derived from the WorkUnit if one is provided, + # 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) # Create and save any additional meta data that should be saved with the results. num_img = stack.img_count() @@ -306,6 +317,70 @@ def run_search_from_work_unit(self, work): work.config, work.im_stack, trj_generator=trj_generator, - wcs=work.wcs, + workunit=work, extra_meta=extra_meta, ) + + +def append_ra_dec_to_results(workunit, results): + """Append predicted (RA, dec) positions to the results. + + Parameters + ---------- + workunit : `WorkUnit` + The WorkUnit with all the WCS information. + results : `Results` + The current table of results including the per-pixel trajectories. + This is modified in-place. + """ + num_results = len(results) + if num_results == 0: + return # Nothing to do + + 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. + 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. + if workunit.wcs is not None: + 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 = [] + 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", + 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)]) + + results.table["img_ra"] = all_ra + results.table["img_dec"] = all_dec + else: + # 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 diff --git a/tests/test_run_search.py b/tests/test_run_search.py new file mode 100644 index 00000000..1ad9cf4f --- /dev/null +++ b/tests/test_run_search.py @@ -0,0 +1,144 @@ +"""Test some of the functions needed for running the search.""" + +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.results import Results +from kbmod.run_search import append_ra_dec_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): + # 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) + + # Append a global fake WCS and one for each time. + global_wcs = make_fake_wcs(20.0, 0.0, 800, 600, deg_per_pixel=0.5 / 3600.0) + all_wcs = [] + for idx in range(num_times): + curr = make_fake_wcs( + 20.01 + idx / 100.0, 0.01 + idx / 100.0, 800, 600, deg_per_pixel=0.5 / 3600.0 + ) + all_wcs.append(curr) + + fake_wu = WorkUnit( + im_stack=fake_ds.stack, + config=SearchConfiguration(), + wcs=global_wcs, + per_image_wcs=all_wcs, + reprojected=True, + per_image_indices=[i for i in range(num_times)], + heliocentric_distance=np.full(num_times, 100.0), + obstimes=fake_times, + ) + + # Create three fake trajectories in the bounds of the images. We don't + # bother actually inserting them into the pixels. + trjs = [ + Trajectory(x=5, y=10, vx=1, vy=1, flux=1000.0, lh=1000.0, obs_count=num_times), + Trajectory(x=400, y=300, vx=-5, vy=-2, flux=1000.0, lh=1000.0, obs_count=num_times), + Trajectory(x=100, y=500, vx=10, vy=-10, flux=1000.0, lh=1000.0, obs_count=num_times), + ] + results = Results.from_trajectories(trjs) + self.assertEqual(len(results), 3) + + append_ra_dec_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) + for i in range(3): + self.assertEqual(len(results["global_ra"][i]), num_times) + self.assertTrue(np.all(results["global_ra"][i] > 19.0)) + self.assertTrue(np.all(results["global_ra"][i] < 21.0)) + + # The global Dec should exist and be close to 0.0 for all observations. + self.assertEqual(len(results["global_dec"]), 3) + for i in range(3): + self.assertEqual(len(results["global_dec"][i]), num_times) + 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 + 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])) + + # The global Dec should exist and be close to 0.0 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])) + + def test_append_ra_dec_no_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) + + # Append a global fake WCS and one for each time. + all_wcs = [] + for idx in range(num_times): + curr = make_fake_wcs( + 20.01 + idx / 100.0, 0.01 + idx / 100.0, 800, 600, deg_per_pixel=0.5 / 3600.0 + ) + all_wcs.append(curr) + + fake_wu = WorkUnit( + im_stack=fake_ds.stack, + config=SearchConfiguration(), + wcs=None, + per_image_wcs=all_wcs, + reprojected=False, + per_image_indices=[i for i in range(num_times)], + obstimes=fake_times, + ) + + # Create three fake trajectories in the bounds of the images. We don't + # bother actually inserting them into the pixels. + trjs = [ + Trajectory(x=5, y=10, vx=1, vy=1, flux=1000.0, lh=1000.0, obs_count=num_times), + Trajectory(x=400, y=300, vx=-5, vy=-2, flux=1000.0, lh=1000.0, obs_count=num_times), + Trajectory(x=100, y=500, vx=10, vy=-10, flux=1000.0, lh=1000.0, obs_count=num_times), + ] + results = Results.from_trajectories(trjs) + self.assertEqual(len(results), 3) + + append_ra_dec_to_results(fake_wu, results) + + # The global RA and global dec should not exist. + self.assertFalse("global_ra" in results.colnames) + self.assertFalse("global_dec" in results.colnames) + + # The per-image RA should exist, be close to 20.0 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)) + + # The global Dec should exist and be close to 0.0 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)) + + +if __name__ == "__main__": + unittest.main() From c599f41a7c1b1589d4c1406e722e59b29527264d Mon Sep 17 00:00:00 2001 From: Jeremy Kubica <104161096+jeremykubica@users.noreply.github.com> Date: Sat, 11 Jan 2025 10:36:56 -0500 Subject: [PATCH 2/2] Add basic logging --- src/kbmod/run_search.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/kbmod/run_search.py b/src/kbmod/run_search.py index 8baf9d35..8333b7e5 100644 --- a/src/kbmod/run_search.py +++ b/src/kbmod/run_search.py @@ -346,6 +346,8 @@ def append_ra_dec_to_results(workunit, results): # Compute the predicted (RA, dec) positions for each trajectory in global space. if workunit.wcs is not None: + logger.info("Found common WCS. Adding global_ra and global_dec columns.") + skypos = workunit.wcs.pixel_to_world(xp, yp) results.table["global_ra"] = skypos.ra.degree results.table["global_dec"] = skypos.dec.degree @@ -371,6 +373,8 @@ def append_ra_dec_to_results(workunit, results): 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))