diff --git a/README.md b/README.md index 1c74ca704..560c2a22a 100644 --- a/README.md +++ b/README.md @@ -140,7 +140,7 @@ candidates = [trj for trj in gen] # Do the actual search. search = kb.StackSearch(stack) -search.search( +search.search_all( strategy, 7, # The minimum number of observations ) diff --git a/notebooks/Kbmod_Reference.ipynb b/notebooks/Kbmod_Reference.ipynb index be6a67c9b..6e2db6a78 100644 --- a/notebooks/Kbmod_Reference.ipynb +++ b/notebooks/Kbmod_Reference.ipynb @@ -406,7 +406,7 @@ "candidates = [trj for trj in gen]\n", "print(f\"Created {len(candidates)} candidate trajectories per pixel.\")\n", "\n", - "search.search(candidates, 2)" + "search.search_all(candidates, 2)" ] }, { diff --git a/src/kbmod/batch_search.py b/src/kbmod/batch_search.py new file mode 100644 index 000000000..f34c33c74 --- /dev/null +++ b/src/kbmod/batch_search.py @@ -0,0 +1,31 @@ +class BatchSearchManager: + def __init__(self, stack_search, search_list, min_observations): + """Manages a batch search over a list of Trajectory instances using a StackSearch instance. + + Parameters + ---------- + stack_search: `StackSearch` + StackSearch instance to use for the batch search. + search_list: `list[Trajectory]` + List of Trajectory instances to search along the stack. + min_observations: `int` + Minimum number of observations required to consider a candidate. + """ + self.stack_search = stack_search + self.search_list = search_list + self.min_observations = min_observations + + def __enter__(self): + # Prepare memory for the batch search. + self.stack_search.prepare_search(self.search_list, self.min_observations) + return self.stack_search + + def __exit__(self, *_): + """ + This method is called when exiting the context. + It cleans up resources by calling `finish_search` on the StackSearch instance. + We return False to indicate that we do not want to suppress any exceptions that may have been raised. + """ + # Clean up + self.stack_search.finish_search() + return False diff --git a/src/kbmod/run_search.py b/src/kbmod/run_search.py index 4812f258c..412ec267b 100644 --- a/src/kbmod/run_search.py +++ b/src/kbmod/run_search.py @@ -197,7 +197,7 @@ def do_gpu_search(self, config, stack, trj_generator): # Do the actual search. candidates = [trj for trj in trj_generator] - search.search(candidates, int(config["num_obs"])) + search.search_all(candidates, int(config["num_obs"])) search_timer.stop() # Load the results. diff --git a/src/kbmod/search/pydocs/stack_search_docs.h b/src/kbmod/search/pydocs/stack_search_docs.h index 9ddb520b3..99b29a470 100644 --- a/src/kbmod/search/pydocs/stack_search_docs.h +++ b/src/kbmod/search/pydocs/stack_search_docs.h @@ -177,6 +177,46 @@ static const auto DOC_StackSearch_get_results = R"doc( ``RunTimeError`` if start < 0 or count <= 0. )doc"; +static const auto DOC_StackSearch_prepare_batch_search = R"doc( + Prepare the search for a batch of trajectories. + + Parameters + ---------- + search_list : `List` + A list of ``Trajectory`` objects to search. + min_observations : `int` + The minimum number of observations for a trajectory to be considered. + )doc"; + +static const auto DOC_StackSearch_compute_max_results = R"doc( + Compute the maximum number of results according to the x, y bounds and the RESULTS_PER_PIXEL constant + + Returns + ------- + max_results : `int` + The maximum number of results that a search will return according to the current bounds and the RESULTS_PER_PIXEL constant. + )doc"; + +static const auto DOC_StackSearch_search_single_batch = R"doc( + Perform a search on the given trajectories for the current batch. + Batch is defined by the parameters set `set_start_bounds_x` & `set_start_bounds_y`. + + Returns + ------- + results : `List` + A list of ``Trajectory`` search results + )doc"; + +static const auto DOC_StackSearch_finish_search = R"doc( + Clears memory used for the batch search. + + This method should be called after a batch search is completed to ensure that any resources allocated during the search are properly freed. + + Returns + ------- + None + )doc"; + static const auto DOC_StackSearch_set_results = R"doc( Set the cached results. Used for testing. diff --git a/src/kbmod/search/stack_search.cpp b/src/kbmod/search/stack_search.cpp index 086015d42..6d3aeaa2a 100644 --- a/src/kbmod/search/stack_search.cpp +++ b/src/kbmod/search/stack_search.cpp @@ -15,7 +15,7 @@ extern "C" void evaluateTrajectory(PsiPhiArrayMeta psi_phi_meta, void* psi_phi_v // I'd imaging... auto rs_logger = logging::getLogger("kbmod.search.run_search"); -StackSearch::StackSearch(ImageStack& imstack) : stack(imstack), results(0) { +StackSearch::StackSearch(ImageStack& imstack) : stack(imstack), results(0), gpu_search_list(0) { debug_info = false; psi_phi_generated = false; @@ -155,19 +155,40 @@ Trajectory StackSearch::search_linear_trajectory(short x, short y, float vx, flo return result; } -void StackSearch::search(std::vector& search_list, int min_observations) { - DebugTimer core_timer = DebugTimer("core search", rs_logger); +void StackSearch::finish_search(){ + psi_phi_array.clear_from_gpu(); + gpu_search_list.move_to_cpu(); +} - DebugTimer psi_phi_timer = DebugTimer("creating psi/phi buffers", rs_logger); +void StackSearch::prepare_search(std::vector& search_list, int min_observations){ + DebugTimer psi_phi_timer = DebugTimer("Creating psi/phi buffers", rs_logger); prepare_psi_phi(); psi_phi_array.move_to_gpu(); psi_phi_timer.stop(); - // Allocate a vector for the results and move it onto the GPU. - int search_width = params.x_start_max - params.x_start_min; - int search_height = params.y_start_max - params.y_start_min; - int num_search_pixels = search_width * search_height; - int max_results = num_search_pixels * RESULTS_PER_PIXEL; + + int num_to_search = search_list.size(); + if (debug_info) std::cout << "Preparing to search " << num_to_search << " trajectories... \n" << std::flush; + gpu_search_list.set_trajectories(search_list); + gpu_search_list.move_to_gpu(); + + params.min_observations = min_observations; +} + +void StackSearch::search_all(std::vector& search_list, int min_observations) { + prepare_search(search_list, min_observations); + search_batch(); + finish_search(); +} + +void StackSearch::search_batch(){ + if(!psi_phi_array.gpu_array_allocated()){ + throw std::runtime_error("PsiPhiArray array not allocated on GPU. Did you forget to call prepare_search?"); + } + + DebugTimer core_timer = DebugTimer("Running batch search", rs_logger); + int max_results = compute_max_results(); + // staple C++ std::stringstream logmsg; logmsg << "Searching X=[" << params.x_start_min << ", " << params.x_start_max << "] " @@ -175,24 +196,12 @@ void StackSearch::search(std::vector& search_list, int min_observati << "Allocating space for " << max_results << " results."; rs_logger->info(logmsg.str()); + results.resize(max_results); results.move_to_gpu(); - // Allocate space for the search list and move that to the GPU. - int num_to_search = search_list.size(); - - logmsg.str(""); - logmsg << search_list.size() << " trajectories..."; - rs_logger->info(logmsg.str()); - - TrajectoryList gpu_search_list(search_list); - gpu_search_list.move_to_gpu(); - - // Set the minimum number of observations. - params.min_observations = min_observations; - // Do the actual search on the GPU. - DebugTimer search_timer = DebugTimer("search execution", rs_logger); + DebugTimer search_timer = DebugTimer("Running search", rs_logger); #ifdef HAVE_CUDA deviceSearchFilter(psi_phi_array, params, gpu_search_list, results); #else @@ -200,18 +209,27 @@ void StackSearch::search(std::vector& search_list, int min_observati #endif search_timer.stop(); - // Move data back to CPU to unallocate GPU space (this will happen automatically - // for gpu_search_list when the object goes out of scope, but we do it explicitly here). - psi_phi_array.clear_from_gpu(); results.move_to_cpu(); - gpu_search_list.move_to_cpu(); - DebugTimer sort_timer = DebugTimer("Sorting results", rs_logger); results.sort_by_likelihood(); sort_timer.stop(); core_timer.stop(); } +std::vector StackSearch::search_single_batch(){ + int max_results = compute_max_results(); + search_batch(); + return results.get_batch(0, max_results); +} + + +int StackSearch::compute_max_results(){ + int search_width = params.x_start_max - params.x_start_min; + int search_height = params.y_start_max - params.y_start_min; + int num_search_pixels = search_width * search_height; + return num_search_pixels * RESULTS_PER_PIXEL; +} + std::vector StackSearch::extract_psi_or_phi_curve(Trajectory& trj, bool extract_psi) { prepare_psi_phi(); @@ -261,7 +279,7 @@ static void stack_search_bindings(py::module& m) { py::class_(m, "StackSearch", pydocs::DOC_StackSearch) .def(py::init()) - .def("search", &ks::search, pydocs::DOC_StackSearch_search) + .def("search_all", &ks::search_all, pydocs::DOC_StackSearch_search) .def("evaluate_single_trajectory", &ks::evaluate_single_trajectory, pydocs::DOC_StackSearch_evaluate_single_trajectory) .def("search_linear_trajectory", &ks::search_linear_trajectory, @@ -288,7 +306,11 @@ static void stack_search_bindings(py::module& m) { .def("prepare_psi_phi", &ks::prepare_psi_phi, pydocs::DOC_StackSearch_prepare_psi_phi) .def("clear_psi_phi", &ks::clear_psi_phi, pydocs::DOC_StackSearch_clear_psi_phi) .def("get_results", &ks::get_results, pydocs::DOC_StackSearch_get_results) - .def("set_results", &ks::set_results, pydocs::DOC_StackSearch_set_results); + .def("set_results", &ks::set_results, pydocs::DOC_StackSearch_set_results) + .def("compute_max_results", &ks::compute_max_results, pydocs::DOC_StackSearch_compute_max_results) + .def("search_single_batch", &ks::search_single_batch, pydocs::DOC_StackSearch_search_single_batch) + .def("prepare_search", &ks::prepare_search, pydocs::DOC_StackSearch_prepare_batch_search) + .def("finish_search", &ks::finish_search, pydocs::DOC_StackSearch_finish_search); } #endif /* Py_PYTHON_H */ diff --git a/src/kbmod/search/stack_search.h b/src/kbmod/search/stack_search.h index 83192a24a..5b7906f7d 100644 --- a/src/kbmod/search/stack_search.h +++ b/src/kbmod/search/stack_search.h @@ -31,7 +31,7 @@ using Image = search::Image; class StackSearch { public: StackSearch(ImageStack& imstack); - + int compute_max_results(); int num_images() const { return stack.img_count(); } int get_image_width() const { return stack.get_width(); } int get_image_height() const { return stack.get_height(); } @@ -50,7 +50,11 @@ class StackSearch { // The primary search functions void evaluate_single_trajectory(Trajectory& trj); Trajectory search_linear_trajectory(short x, short y, float vx, float vy); - void search(std::vector& search_list, int min_observations); + void prepare_search(std::vector& search_list, int min_observations); + std::vector search_single_batch(); + void search_batch(); + void search_all(std::vector& search_list, int min_observations); + void finish_search(); // Gets the vector of result trajectories from the grid search. std::vector get_results(int start, int end); @@ -82,6 +86,9 @@ class StackSearch { // Results from the grid search. TrajectoryList results; + + // Trajectories that are being searched. + TrajectoryList gpu_search_list; }; } /* namespace search */ diff --git a/tests/test_readme_example.py b/tests/test_readme_example.py index 0e2a704ad..62f19d24c 100644 --- a/tests/test_readme_example.py +++ b/tests/test_readme_example.py @@ -53,7 +53,7 @@ def test_make_and_copy(self): # Do the actual search. search = kb.StackSearch(stack) - search.search( + search.search_all( candidates, 7, # The minimum number of observations ) diff --git a/tests/test_search.py b/tests/test_search.py index 36cbeb94b..3eefe8837 100644 --- a/tests/test_search.py +++ b/tests/test_search.py @@ -2,6 +2,7 @@ import numpy as np +from kbmod.batch_search import BatchSearchManager from kbmod.configuration import SearchConfiguration from kbmod.fake_data.fake_data_creator import add_fake_object, make_fake_layered_image, FakeDataSet from kbmod.run_search import SearchRunner @@ -210,7 +211,7 @@ def test_search_linear_trajectory(self): @unittest.skipIf(not HAS_GPU, "Skipping test (no GPU detected)") def test_results(self): candidates = [trj for trj in self.trj_gen] - self.search.search(candidates, int(self.img_count / 2)) + self.search.search_all(candidates, int(self.img_count / 2)) results = self.search.get_results(0, 10) best = results[0] @@ -226,7 +227,7 @@ def test_results_extended_bounds(self): self.search.set_start_bounds_y(-10, self.dim_y + 10) candidates = [trj for trj in self.trj_gen] - self.search.search(candidates, int(self.img_count / 2)) + self.search.search_all(candidates, int(self.img_count / 2)) results = self.search.get_results(0, 10) best = results[0] @@ -242,7 +243,7 @@ def test_results_reduced_bounds(self): self.search.set_start_bounds_y(5, self.dim_y - 5) candidates = [trj for trj in self.trj_gen] - self.search.search(candidates, int(self.img_count / 2)) + self.search.search_all(candidates, int(self.img_count / 2)) results = self.search.get_results(0, 10) best = results[0] @@ -290,7 +291,7 @@ def test_results_off_chip(self): search.set_start_bounds_x(-10, self.dim_x + 10) search.set_start_bounds_y(-10, self.dim_y + 10) candidates = [trj for trj in self.trj_gen] - search.search(candidates, int(self.img_count / 2)) + search.search_all(candidates, int(self.img_count / 2)) # Check the results. results = search.get_results(0, 10) @@ -932,6 +933,70 @@ def test_coadd_filter_gpu(self): self.assertEqual(meanStamps[2].width, 1) self.assertEqual(meanStamps[2].height, 1) + @staticmethod + def result_hash(res): + return hash((res.x, res.y, res.vx, res.vy, res.lh, res.obs_count)) + + def test_search_batch(self): + width = 50 + height = 50 + results_per_pixel = 8 + min_observations = 2 + + # Simple average PSF + psf_data = np.zeros((5, 5), dtype=np.single) + psf_data[1:4, 1:4] = 0.1111111 + p = PSF(psf_data) + + # Create a stack with 10 20x20 images with random noise and times ranging from 0 to 1 + count = 10 + imlist = [make_fake_layered_image(width, height, 5.0, 25.0, n / count, p) for n in range(count)] + stack = ImageStack(imlist) + + for i in range(count): + im = stack.get_single_image(i) + time = stack.get_zeroed_time(i) + add_fake_object(im, 5.0 + (time * 8.0), 35.0 + (time * 0.0), 25000.0) + + search = StackSearch(stack) + + # Sample generator + gen = KBMODV1Search( + 10, 5, 15, 10, -0.1, 0.1 + ) # velocity_steps, min_vel, max_vel, angle_steps, min_ang, max_ang, + candidates = [trj for trj in gen] + + # Peform complete in-memory search + search.search_all(candidates, min_observations) + total_results = width * height * results_per_pixel + # Need to filter as the fields are undefined otherwise + results = [ + result + for result in search.get_results(0, total_results) + if result.lh > -1 and result.obs_count >= min_observations + ] + + with BatchSearchManager(StackSearch(stack), candidates, min_observations) as batch_search: + + batch_results = [] + for i in range(0, width, 5): + batch_search.set_start_bounds_x(i, i + 5) + for j in range(0, height, 5): + batch_search.set_start_bounds_y(j, j + 5) + batch_results.extend(batch_search.search_single_batch()) + + # Need to filter as the fields are undefined otherwise + batch_results = [ + result for result in batch_results if result.lh > -1 and result.obs_count >= min_observations + ] + + # Check that the results are the same. + results_hash_set = {test_search.result_hash(result) for result in results} + batch_results_hash_set = {test_search.result_hash(result) for result in batch_results} + + for res_hash in results_hash_set: + self.assertTrue(res_hash in batch_results_hash_set) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_search_encode.py b/tests/test_search_encode.py index fc1dbd0e8..5de0aaaac 100644 --- a/tests/test_search_encode.py +++ b/tests/test_search_encode.py @@ -72,7 +72,7 @@ def test_different_encodings(self): search = StackSearch(self.stack) search.enable_gpu_encoding(encoding_bytes) candidates = [trj for trj in self.trj_gen] - search.search(candidates, int(self.img_count / 2)) + search.search_all(candidates, int(self.img_count / 2)) results = search.get_results(0, 10) best = results[0] diff --git a/tests/test_search_filter.py b/tests/test_search_filter.py index dbbbe99b6..ccbbb321c 100644 --- a/tests/test_search_filter.py +++ b/tests/test_search_filter.py @@ -66,7 +66,7 @@ def setUp(self): self.max_angle, ) candidates = [trj for trj in trj_gen] - self.search.search(candidates, int(self.img_count / 2)) + self.search.search_all(candidates, int(self.img_count / 2)) @unittest.skipIf(not HAS_GPU, "Skipping test (no GPU detected)") def test_results(self):