diff --git a/notebooks/TrajectoryExplorer.ipynb b/notebooks/TrajectoryExplorer.ipynb index aa0cb1604..9ef168643 100644 --- a/notebooks/TrajectoryExplorer.ipynb +++ b/notebooks/TrajectoryExplorer.ipynb @@ -244,6 +244,30 @@ "for i in range(len(fake_times)):\n", " print(f\"Time {i} is valid = {result['obs_valid'][0][i]}.\")" ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Neighborhood Searches\n", + "\n", + "The `TrajectoryExplorer` can also be used to perform a hyper-localized search. This search effectively uses a small neighborhood around a given trajectory (both in terms of starting pixel and velocities) and returns all results from this neighborhood. This localized set can be used to:\n", + "1) refine trajectories by searching a finer parameter space around the best results found by the initial search, or\n", + "2) collect a distribution of trajectories and their likelihoods around a single result.\n", + "For this search, the `TrajectoryExplorer` does not perform any filtering, so it will return all trajectories and their likelihoods (even one <= -1.0)\n", + "\n", + "Only basic statistics, such as likelihood and flux, are returned from the search. Stamps and lightcurves are not computed." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "samples = explorer.evaluate_around_linear_trajectory(50, 60, 5.0, -2.0, pixel_radius=5)\n", + "print(samples[0:10])" + ] } ], "metadata": { diff --git a/src/kbmod/configuration.py b/src/kbmod/configuration.py index f406d8dce..64ab3da43 100644 --- a/src/kbmod/configuration.py +++ b/src/kbmod/configuration.py @@ -1,3 +1,4 @@ +import copy import math from astropy.io import fits @@ -84,6 +85,10 @@ def __str__(self): result += f"{key}: {value}\n" return result + def copy(self): + """Create a new deep copy of the configuration.""" + return copy.deepcopy(self) + def set(self, param, value, warn_on_unknown=False): """Sets the value of a specific parameter. diff --git a/src/kbmod/run_search.py b/src/kbmod/run_search.py index 8333b7e53..d03b15f6d 100644 --- a/src/kbmod/run_search.py +++ b/src/kbmod/run_search.py @@ -20,6 +20,54 @@ logger = kb.Logging.getLogger(__name__) +def configure_kb_search_stack(search, config): + """Configure the kbmod SearchStack object from a search configuration. + + Parameters + ---------- + search : `kb.StackSearch` + The SearchStack object. + config : `SearchConfiguration` + The configuration parameters + """ + width = search.get_image_width() + height = search.get_image_height() + + # Set the search bounds. + if config["x_pixel_bounds"] and len(config["x_pixel_bounds"]) == 2: + search.set_start_bounds_x(config["x_pixel_bounds"][0], config["x_pixel_bounds"][1]) + elif config["x_pixel_buffer"] and config["x_pixel_buffer"] > 0: + search.set_start_bounds_x(-config["x_pixel_buffer"], width + config["x_pixel_buffer"]) + + if config["y_pixel_bounds"] and len(config["y_pixel_bounds"]) == 2: + search.set_start_bounds_y(config["y_pixel_bounds"][0], config["y_pixel_bounds"][1]) + elif config["y_pixel_buffer"] and config["y_pixel_buffer"] > 0: + search.set_start_bounds_y(-config["y_pixel_buffer"], height + config["y_pixel_buffer"]) + + # Set the results per pixel. + search.set_results_per_pixel(config["results_per_pixel"]) + + # If we are using gpu_filtering, enable it and set the parameters. + if config["gpu_filter"]: + logger.debug("Using in-line GPU sigmaG filtering methods") + coeff = SigmaGClipping.find_sigma_g_coeff( + config["sigmaG_lims"][0], + config["sigmaG_lims"][1], + ) + search.enable_gpu_sigmag_filter( + np.array(config["sigmaG_lims"]) / 100.0, + coeff, + config["lh_level"], + ) + else: + search.disable_gpu_sigmag_filter() + + # If we are using an encoded image representation on GPU, enable it and + # set the parameters. + if config["encode_num_bytes"] > 0: + search.enable_gpu_encoding(config["encode_num_bytes"]) + + class SearchRunner: """A class to run the KBMOD grid search.""" @@ -140,45 +188,11 @@ def do_gpu_search(self, config, stack, trj_generator): """ # Create the search object which will hold intermediate data and results. search = kb.StackSearch(stack) - - width = search.get_image_width() - height = search.get_image_height() - - # Set the search bounds. - if config["x_pixel_bounds"] and len(config["x_pixel_bounds"]) == 2: - search.set_start_bounds_x(config["x_pixel_bounds"][0], config["x_pixel_bounds"][1]) - elif config["x_pixel_buffer"] and config["x_pixel_buffer"] > 0: - search.set_start_bounds_x(-config["x_pixel_buffer"], width + config["x_pixel_buffer"]) - - if config["y_pixel_bounds"] and len(config["y_pixel_bounds"]) == 2: - search.set_start_bounds_y(config["y_pixel_bounds"][0], config["y_pixel_bounds"][1]) - elif config["y_pixel_buffer"] and config["y_pixel_buffer"] > 0: - search.set_start_bounds_y(-config["y_pixel_buffer"], height + config["y_pixel_buffer"]) - - # Set the results per pixel. - search.set_results_per_pixel(config["results_per_pixel"]) + configure_kb_search_stack(search, config) search_timer = kb.DebugTimer("grid search", logger) logger.debug(f"{trj_generator}") - # If we are using gpu_filtering, enable it and set the parameters. - if config["gpu_filter"]: - logger.debug("Using in-line GPU sigmaG filtering methods") - coeff = SigmaGClipping.find_sigma_g_coeff( - config["sigmaG_lims"][0], - config["sigmaG_lims"][1], - ) - search.enable_gpu_sigmag_filter( - np.array(config["sigmaG_lims"]) / 100.0, - coeff, - config["lh_level"], - ) - - # If we are using an encoded image representation on GPU, enable it and - # set the parameters. - if config["encode_num_bytes"] > 0: - search.enable_gpu_encoding(config["encode_num_bytes"]) - # Do the actual search. candidates = [trj for trj in trj_generator] search.search_all(candidates, int(config["num_obs"])) diff --git a/src/kbmod/search/kernels/kernels.cu b/src/kbmod/search/kernels/kernels.cu index ed559d421..b4b2e3a6d 100644 --- a/src/kbmod/search/kernels/kernels.cu +++ b/src/kbmod/search/kernels/kernels.cu @@ -251,7 +251,7 @@ __global__ void searchFilterImages(PsiPhiArrayMeta psi_phi_meta, void *psi_phi_v for (int r = 0; r < params.results_per_pixel; ++r) { results[base_index + r].x = x; results[base_index + r].y = y; - results[base_index + r].lh = -1.0; + results[base_index + r].lh = -FLT_MAX; } // For each trajectory we'd like to search @@ -274,10 +274,9 @@ __global__ void searchFilterImages(PsiPhiArrayMeta psi_phi_meta, void *psi_phi_v continue; // Insert the new trajectory into the sorted list of final results. - // Only sort the values with valid likelihoods. Trajectory temp; for (unsigned int r = 0; r < params.results_per_pixel; ++r) { - if (curr_trj.lh > results[base_index + r].lh && curr_trj.lh > -1.0) { + if (curr_trj.lh > results[base_index + r].lh) { temp = results[base_index + r]; results[base_index + r] = curr_trj; curr_trj = temp; diff --git a/src/kbmod/search/pydocs/stack_search_docs.h b/src/kbmod/search/pydocs/stack_search_docs.h index e64c030d5..4d394c86b 100644 --- a/src/kbmod/search/pydocs/stack_search_docs.h +++ b/src/kbmod/search/pydocs/stack_search_docs.h @@ -30,6 +30,10 @@ static const auto DOC_StackSearch_set_min_lh = R"doc( The minimum likelihood value for a trajectory to be returned. )doc"; +static const auto DOC_StackSearch_disable_gpu_sigmag_filter = R"doc( + Turns off the on-GPU sigma-G filtering. + )doc"; + static const auto DOC_StackSearch_enable_gpu_sigmag_filter = R"doc( Enable on-GPU sigma-G filtering. diff --git a/src/kbmod/search/stack_search.cpp b/src/kbmod/search/stack_search.cpp index 0123b9dc9..6e991f999 100644 --- a/src/kbmod/search/stack_search.cpp +++ b/src/kbmod/search/stack_search.cpp @@ -66,6 +66,10 @@ void StackSearch::enable_gpu_sigmag_filter(std::vector percentiles, float params.sigmag_coeff = sigmag_coeff; params.min_lh = min_lh; } + +void StackSearch::disable_gpu_sigmag_filter() { + params.do_sigmag_filter = false; +} void StackSearch::enable_gpu_encoding(int encode_num_bytes) { // If changing a setting that would impact the search data encoding, clear the cached values. @@ -310,6 +314,8 @@ static void stack_search_bindings(py::module& m) { .def("set_min_lh", &ks::set_min_lh, pydocs::DOC_StackSearch_set_min_lh) .def("set_results_per_pixel", &ks::set_results_per_pixel, pydocs::DOC_StackSearch_set_results_per_pixel) + .def("disable_gpu_sigmag_filter", &ks::disable_gpu_sigmag_filter, + pydocs::DOC_StackSearch_disable_gpu_sigmag_filter) .def("enable_gpu_sigmag_filter", &ks::enable_gpu_sigmag_filter, pydocs::DOC_StackSearch_enable_gpu_sigmag_filter) .def("enable_gpu_encoding", &ks::enable_gpu_encoding, pydocs::DOC_StackSearch_enable_gpu_encoding) diff --git a/src/kbmod/search/stack_search.h b/src/kbmod/search/stack_search.h index 7a629bea2..3c8cb2ac6 100644 --- a/src/kbmod/search/stack_search.h +++ b/src/kbmod/search/stack_search.h @@ -40,6 +40,7 @@ class StackSearch { // Parameter setters used to control the searches. void set_min_obs(int new_value); void set_min_lh(float new_value); + void disable_gpu_sigmag_filter(); void enable_gpu_sigmag_filter(std::vector percentiles, float sigmag_coeff, float min_lh); void enable_gpu_encoding(int num_bytes); void set_start_bounds_x(int x_min, int x_max); diff --git a/src/kbmod/trajectory_explorer.py b/src/kbmod/trajectory_explorer.py index 244d40af7..87a976b88 100644 --- a/src/kbmod/trajectory_explorer.py +++ b/src/kbmod/trajectory_explorer.py @@ -3,8 +3,10 @@ from kbmod.configuration import SearchConfiguration from kbmod.filters.sigma_g_filter import apply_clipped_sigma_g, SigmaGClipping from kbmod.results import Results -from kbmod.search import StackSearch, Logging +from kbmod.run_search import configure_kb_search_stack +from kbmod.search import DebugTimer, StackSearch, Logging from kbmod.filters.stamp_filters import append_all_stamps, append_coadds +from kbmod.trajectory_generator import PencilSearch from kbmod.trajectory_utils import make_trajectory_from_ra_dec @@ -42,9 +44,24 @@ def __init__(self, img_stack, config=None): # Allocate and configure the StackSearch object. self.search = None - def initialize_data(self): - """Perform any needed initialization and preprocessing on the images.""" + def initialize_data(self, config=None): + """Initialize the data, including applying the configuration parameters. + + Parameters + ---------- + config : `SearchConfiguration`, optional + Any custom configuration parameters to use for this run. + If ``None`` uses the default configuration parameters. + """ + if config is None: + config = self.config + if self._data_initalized: + # Always reapply the configuration parameters if in case we used custom + # ones on a previous search. + configure_kb_search_stack(self.search, config) + + # Nothing else to do return # If we are using an encoded image representation on GPU, enable it and @@ -55,6 +72,7 @@ def initialize_data(self): # Allocate the search structure. self.search = StackSearch(self.im_stack) + configure_kb_search_stack(self.search, config) self._data_initalized = True @@ -121,6 +139,82 @@ def evaluate_angle_trajectory(self, ra, dec, v_ra, v_dec, wcs): trj = make_trajectory_from_ra_dec(ra, dec, v_ra, v_dec, wcs) return self.evaluate_linear_trajectory(trj.x, trj.y, trj.vx, trj.vy) + def evaluate_around_linear_trajectory( + self, + x, + y, + vx, + vy, + pixel_radius=5, + max_ang_offset=0.2618, + ang_step=0.035, + max_vel_offset=10.0, + vel_step=0.5, + ): + """Evaluate all the trajectories within a local neighborhood of the given trajectory. + No filtering is done at all. + + Parameters + ---------- + x : `int` + The starting x pixel of the trajectory. + y : `int` + The starting y pixel of the trajectory. + vx : `float` + The x velocity of the trajectory in pixels per day. + vy : `float` + The y velocity of the trajectory in pixels per day. + pixel_radius : `int` + The number of pixels to evaluate to each side of the Trajectory's starting pixel. + max_ang_offset : `float` + The maximum offset of a candidate trajectory from the original (in radians) + ang_step : `float` + The step size to explore for each angle (in radians) + max_vel_offset : `float` + The maximum offset of the velocity's magnitude from the original (in pixels per day) + vel_step : `float` + The step size to explore for each velocity magnitude (in pixels per day) + + Returns + ------- + result : `Results` + The results table with a single row and all the columns filled out. + """ + if pixel_radius < 0: + raise ValueError(f"Pixel radius must be >= 0. Got {pixel_radius}") + num_pixels = (2 * pixel_radius + 1) * (2 * pixel_radius + 1) + logger.debug(f"Testing {num_pixels} starting pixels.") + + # Create a pencil search around the given trajectory. + trj_generator = PencilSearch(vx, vy, max_ang_offset, ang_step, max_vel_offset, vel_step) + num_trj = len(trj_generator) + logger.debug(f"Exploring {num_trj} trajectories per starting pixel.") + + # Set the search bounds to right around the trajectory's starting position and + # turn off all filtering. + reduced_config = self.config.copy() + reduced_config.set("x_pixel_bounds", [x - pixel_radius, x + pixel_radius + 1]) + reduced_config.set("y_pixel_bounds", [y - pixel_radius, y + pixel_radius + 1]) + reduced_config.set("results_per_pixel", min(num_trj, 10_000)) + reduced_config.set("gpu_filter", False) + reduced_config.set("num_obs", 1) + reduced_config.set("max_lh", 1e25) + reduced_config.set("lh_level", -1e25) + self.initialize_data(config=reduced_config) + + # Do the actual search. + search_timer = DebugTimer("grid search", logger) + candidates = [trj for trj in trj_generator] + self.search.search_all(candidates, int(reduced_config["num_obs"])) + search_timer.stop() + + # Load all of the results without any filtering. + logger.debug(f"Loading {num_pixels * num_trj} results.") + trjs = self.search.get_results(0, num_pixels * num_trj) + results = Results.from_trajectories(trjs) + + return results + def apply_sigma_g(self, result): """Apply sigma G clipping to a single ResultRow. Modifies the row in-place. diff --git a/src/kbmod/trajectory_generator.py b/src/kbmod/trajectory_generator.py index b57632daf..828667928 100644 --- a/src/kbmod/trajectory_generator.py +++ b/src/kbmod/trajectory_generator.py @@ -157,6 +157,9 @@ def __repr__(self): def __str__(self): return f"SingleVelocitySearch: vx={self.vx}, vy={self.vy}" + def __len__(self): + return 1 + def generate(self, *args, **kwargs): """Produces a single candidate trajectory to test. @@ -220,6 +223,9 @@ def __str__(self): f" Vel Y: [{self.min_vy}, {self.max_vy}] in {self.vy_steps} steps." ) + def __len__(self): + return self.vy_steps * self.vx_steps + def generate(self, *args, **kwargs): """Produces a single candidate trajectory to test. @@ -280,11 +286,13 @@ def __init__( self.min_ang = self.center_ang - max_ang_offset self.max_ang = self.center_ang + max_ang_offset self.ang_step = ang_step + self.ang_array = np.arange(self.min_ang, self.max_ang + 1e-8, self.ang_step) self.center_vel = np.sqrt(vx * vx + vy * vy) self.min_vel = np.max([self.center_vel - max_vel_offset, 0.0]) self.max_vel = self.center_vel + max_vel_offset self.vel_step = vel_step + self.vel_array = np.arange(self.min_vel, self.max_vel + 1e-8, self.vel_step) def __repr__(self): return ( @@ -300,6 +308,9 @@ def __str__(self): f" Ang: [{self.min_ang}, {self.max_ang}) in {self.ang_step} sized steps." ) + def __len__(self): + return len(self.ang_array) * len(self.vel_array) + def generate(self, *args, **kwargs): """Produces a single candidate trajectory to test. @@ -308,8 +319,8 @@ def generate(self, *args, **kwargs): candidate : `Trajectory` A ``Trajectory`` to test at each pixel. """ - for ang in np.arange(self.min_ang, self.max_ang + 1e-8, self.ang_step): - for vel in np.arange(self.min_vel, self.max_vel + 1e-8, self.vel_step): + for ang in self.ang_array: + for vel in self.vel_array: vx = np.cos(ang) * vel vy = np.sin(ang) * vel @@ -368,6 +379,9 @@ def __str__(self): f" Ang: [{self.min_ang}, {self.max_ang}) in {self.ang_steps} steps." ) + def __len__(self): + return self.ang_steps * self.vel_steps + def generate(self, *args, **kwargs): """Produces a single candidate trajectory to test. @@ -532,6 +546,9 @@ def __str__(self): Offsets = {self.angles[0]} to {self.angles[1]} [{self.min_ang}, {self.max_ang}] in {self.angles[2]} steps.""" + def __len__(self): + return self.angles[2] * self.velocities[2] + def generate(self, *args, **kwargs): """Produces a single candidate trajectory to test. @@ -593,6 +610,9 @@ def __repr__(self): def __str__(self): return self.__repr__() + def __len__(self): + return self.samples_left + def reset_sample_count(self, max_samples): """Reset the counter of samples left. diff --git a/tests/test_configuration.py b/tests/test_configuration.py index db4237dbe..a38e06d92 100644 --- a/tests/test_configuration.py +++ b/tests/test_configuration.py @@ -46,6 +46,21 @@ def test_from_dict(self): self.assertEqual(config["im_filepath"], "Here2") self.assertEqual(config["num_obs"], 5) + def test_copy(self): + d = {"im_filepath": "Here2", "encode_num_bytes": -1} + config = SearchConfiguration.from_dict(d) + + # Create a copy and change values. + config2 = config.copy() + config2.set("im_filepath", "who knows?") + config2.set("encode_num_bytes", 2000) + self.assertEqual(config2["im_filepath"], "who knows?") + self.assertEqual(config2["encode_num_bytes"], 2000) + + # Confirm the original configuration is unchanged. + self.assertEqual(config["im_filepath"], "Here2") + self.assertEqual(config["encode_num_bytes"], -1) + def test_from_hdu(self): t = Table( [ diff --git a/tests/test_trajectory_explorer.py b/tests/test_trajectory_explorer.py index d3f329c74..caf493cec 100644 --- a/tests/test_trajectory_explorer.py +++ b/tests/test_trajectory_explorer.py @@ -33,7 +33,7 @@ def setUp(self): ) fake_ds.insert_object(self.trj) - # Remove at least observation from the trajectory. + # Remove at least one observation from the trajectory. pred_x = self.trj.get_x_index(fake_times[10]) pred_y = self.trj.get_y_index(fake_times[10]) sci_t10 = fake_ds.stack.get_single_image(10).get_science() @@ -80,6 +80,40 @@ def test_evaluate_trajectory(self): self.explorer.apply_sigma_g(result) self.assertFalse(result["obs_valid"][0][10]) + @unittest.skipIf(not HAS_GPU, "Skipping test (no GPU detected)") + def test_evaluate_around_linear_trajectory(self): + radius = 3 + edge_length = 2 * radius + 1 + num_pixels = edge_length * edge_length + + results = self.explorer.evaluate_around_linear_trajectory( + self.x0, + self.y0, + self.vx, + self.vy, + pixel_radius=radius, + max_ang_offset=0.2618, + ang_step=0.035, + max_vel_offset=10.0, + vel_step=0.5, + ) + + # Using the above settings should provide 615 trajectories per starting pixel. + self.assertEqual(len(results), num_pixels * 615) + + # Count the number of results we have per starting pixel. + counts = np.zeros((edge_length, edge_length)) + for row in range(len(results)): + self.assertGreaterEqual(results["x"][row], self.x0 - 3) + self.assertLessEqual(results["x"][row], self.x0 + 3) + self.assertGreaterEqual(results["y"][row], self.y0 - 3) + self.assertLessEqual(results["y"][row], self.y0 + 3) + + x = results["x"][row] - self.x0 + 3 + y = results["y"][row] - self.y0 + 3 + counts[y, x] += 1 + self.assertTrue(np.all(counts == 615)) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_trajectory_generator.py b/tests/test_trajectory_generator.py index 7a26a1897..9853e5d59 100644 --- a/tests/test_trajectory_generator.py +++ b/tests/test_trajectory_generator.py @@ -25,11 +25,13 @@ def test_SingleVelocitySearch(self): self.assertEqual(len(trjs), 1) self.assertEqual(trjs[0].vx, 10.0) self.assertEqual(trjs[0].vy, 5.0) + self.assertEqual(len(gen), 1) def test_VelocityGridSearch(self): gen = VelocityGridSearch(3, 0.0, 2.0, 3, -0.25, 0.25) expected_x = [0.0, 1.0, 2.0, 0.0, 1.0, 2.0, 0.0, 1.0, 2.0] expected_y = [-0.25, -0.25, -0.25, 0.0, 0.0, 0.0, 0.25, 0.25, 0.25] + self.assertEqual(len(gen), 9) trjs = [trj for trj in gen] self.assertEqual(len(trjs), 9) @@ -59,6 +61,7 @@ def test_PencilSearch(self): max_vel_offset=5.0, vel_step=2.5, ) + self.assertEqual(len(gen), 25) trjs = [trj for trj in gen] self.assertEqual(len(trjs), 25) @@ -78,6 +81,7 @@ def test_KBMODV1Search(self): gen = KBMODV1Search(3, 0.0, 3.0, 2, -0.25, 0.25) expected_x = [0.0, 0.9689, 1.9378, 0.0, 1.0, 2.0] expected_y = [0.0, -0.247, -0.4948, 0.0, 0.0, 0.0] + self.assertEqual(len(gen), 6) trjs = [trj for trj in gen] self.assertEqual(len(trjs), 6) @@ -102,6 +106,7 @@ def test_EclipticCenteredSearch(self): gen = EclipticCenteredSearch( [0.0, 2.0, 3], [-45.0, 45.0, 3], angle_units="degree", given_ecliptic=0.0 ) + self.assertEqual(len(gen), 9) expected_x = [0.0, 0.707107, 1.41421, 0.0, 1.0, 2.0, 0.0, 0.707107, 1.41421] expected_y = [0.0, -0.707107, -1.41421, 0.0, 0.0, 0.0, 0.0, 0.707107, 1.41421]