From 4625ec152bb8e7152b9f584ef922be370980073d Mon Sep 17 00:00:00 2001 From: Jeremy Kubica <104161096+jeremykubica@users.noreply.github.com> Date: Tue, 17 Sep 2024 13:27:28 -0400 Subject: [PATCH 01/12] Re-add filter on likelihood --- src/kbmod/search/pydocs/trajectory_list_docs.h | 14 ++++++++++++++ src/kbmod/search/trajectory_list.cpp | 15 +++++++++++++++ src/kbmod/search/trajectory_list.h | 4 +++- tests/test_trajectory_list.py | 16 ++++++++++++++++ 4 files changed, 48 insertions(+), 1 deletion(-) diff --git a/src/kbmod/search/pydocs/trajectory_list_docs.h b/src/kbmod/search/pydocs/trajectory_list_docs.h index 23ae5f43f..52ddb747c 100644 --- a/src/kbmod/search/pydocs/trajectory_list_docs.h +++ b/src/kbmod/search/pydocs/trajectory_list_docs.h @@ -143,6 +143,20 @@ static const auto DOC_TrajectoryList_sort_by_likelihood = R"doc( Raises a ``RuntimeError`` the data is on GPU. )doc"; +static const auto DOC_TrajectoryList_filter_by_likelihood = R"doc( + Sort the data in order of decreasing likelihood and drop everything less than + a given threshold. The data must reside on the CPU. + + Parameters + ---------- + min_likelihood : `float` + The threshold on minimum likelihood. + + Raises + ------ + Raises a ``RuntimeError`` the data is on GPU. + )doc"; + } // namespace pydocs diff --git a/src/kbmod/search/trajectory_list.cpp b/src/kbmod/search/trajectory_list.cpp index 33f09b3db..357d7a40f 100644 --- a/src/kbmod/search/trajectory_list.cpp +++ b/src/kbmod/search/trajectory_list.cpp @@ -73,6 +73,19 @@ void TrajectoryList::sort_by_likelihood() { [](const Trajectory& a, const Trajectory& b) { return b.lh < a.lh; }); } +void TrajectoryList::filter_by_likelihood(float min_likelihood) { + sort_by_likelihood(); + + // Find the first index that does not meet the threshold. + uint64_t index = 0; + while ((index < max_size) && (cpu_list[index].lh >= min_likelihood)) { + ++index; + } + + // Drop the values below the threshold. + resize(index); +} + void TrajectoryList::move_to_gpu() { if (data_on_gpu) return; // Nothing to do. @@ -119,6 +132,8 @@ static void trajectory_list_binding(py::module& m) { .def("get_batch", &trjl::get_batch, pydocs::DOC_TrajectoryList_get_batch) .def("sort_by_likelihood", &trjl::sort_by_likelihood, pydocs::DOC_TrajectoryList_sort_by_likelihood) + .def("filter_by_likelihood", &trjl::filter_by_likelihood, + pydocs::DOC_TrajectoryList_filter_by_likelihood) .def("move_to_cpu", &trjl::move_to_cpu, pydocs::DOC_TrajectoryList_move_to_cpu) .def("move_to_gpu", &trjl::move_to_gpu, pydocs::DOC_TrajectoryList_move_to_gpu); } diff --git a/src/kbmod/search/trajectory_list.h b/src/kbmod/search/trajectory_list.h index 6ee36d17d..f710eafc4 100644 --- a/src/kbmod/search/trajectory_list.h +++ b/src/kbmod/search/trajectory_list.h @@ -58,8 +58,10 @@ class TrajectoryList { // Get a batch of results. std::vector<Trajectory> get_batch(uint64_t start, uint64_t count); - // Processing functions for sorting. + // Processing functions for sorting/filtering. void sort_by_likelihood(); + void filter_by_likelihood(float min_likelihood); + // Data allocation functions. inline bool on_gpu() const { return data_on_gpu; } diff --git a/tests/test_trajectory_list.py b/tests/test_trajectory_list.py index a8b008445..d70d80694 100644 --- a/tests/test_trajectory_list.py +++ b/tests/test_trajectory_list.py @@ -91,6 +91,22 @@ def test_sort(self): for i in range(5): self.assertEqual(trjs.get_trajectory(i).x, lh_order[i]) + def test_filter_on_lh(self): + lh = [100.0, 110.0, 90.0, 120.0, 125.0, 121.0, 10.0] + trjs = TrajectoryList(len(lh)) + for i in range(len(lh)): + trjs.set_trajectory(i, Trajectory(x=i, lh=lh[i])) + + trjs.filter_by_likelihood(110.0) + expected = set([4, 5, 3, 1]) + + # Test that each remaining result appears once in the expected set. + self.assertEqual(len(trjs), len(expected)) + for i in range(len(trjs)): + idx = trjs.get_trajectory(i).x + self.assertTrue(idx in expected) + expected.remove(idx) + @unittest.skipIf(not HAS_GPU, "Skipping test (no GPU detected)") def test_move_to_from_gpu(self): for i in range(self.max_size): From 342029facbc0b8d591ab3d4a7040a8911b247de7 Mon Sep 17 00:00:00 2001 From: Jeremy Kubica <104161096+jeremykubica@users.noreply.github.com> Date: Tue, 17 Sep 2024 13:52:51 -0400 Subject: [PATCH 02/12] Add functions for bulk psi/phi generation --- src/kbmod/search/image_stack.cpp | 8 ++++---- src/kbmod/search/image_stack.h | 4 ++-- src/kbmod/search/stack_search.cpp | 26 ++++++++++++++++++++++++-- src/kbmod/search/stack_search.h | 2 ++ src/kbmod/search/trajectory_list.cpp | 2 +- src/kbmod/search/trajectory_list.h | 1 - 6 files changed, 33 insertions(+), 10 deletions(-) diff --git a/src/kbmod/search/image_stack.cpp b/src/kbmod/search/image_stack.cpp index 6b89db3f8..ea64ce36e 100644 --- a/src/kbmod/search/image_stack.cpp +++ b/src/kbmod/search/image_stack.cpp @@ -39,9 +39,9 @@ void ImageStack::set_single_image(int index, LayeredImage& img, bool force_move) assert_sizes_equal(img.get_height(), height, "ImageStack image height"); if (force_move) { - images[index] = img; + images[index] = img; } else { - images[index] = std::move(img); + images[index] = std::move(img); } } @@ -186,8 +186,8 @@ static void image_stack_bindings(py::module& m) { .def("get_single_image", &is::get_single_image, py::return_value_policy::reference_internal, pydocs::DOC_ImageStack_get_single_image) .def("set_single_image", &is::set_single_image, py::arg("index"), py::arg("img"), - py::arg("force_move")=false, pydocs::DOC_ImageStack_set_single_image) - .def("append_image", &is::append_image, py::arg("img"), py::arg("force_move")=false, + py::arg("force_move") = false, pydocs::DOC_ImageStack_set_single_image) + .def("append_image", &is::append_image, py::arg("img"), py::arg("force_move") = false, pydocs::DOC_ImageStack_append_image) .def("get_obstime", &is::get_obstime, pydocs::DOC_ImageStack_get_obstime) .def("get_zeroed_time", &is::get_zeroed_time, pydocs::DOC_ImageStack_get_zeroed_time) diff --git a/src/kbmod/search/image_stack.h b/src/kbmod/search/image_stack.h index f27e3b77c..20c59abc4 100644 --- a/src/kbmod/search/image_stack.h +++ b/src/kbmod/search/image_stack.h @@ -36,8 +36,8 @@ class ImageStack { // Functions for setting or appending a single LayeredImage. If force_move is true, // then the code uses move semantics and destroys the input object. - void set_single_image(int index, LayeredImage& img, bool force_move=false); - void append_image(LayeredImage& img, bool force_move=false); + void set_single_image(int index, LayeredImage& img, bool force_move = false); + void append_image(LayeredImage& img, bool force_move = false); // Functions for getting or using times. double get_obstime(int index) const; diff --git a/src/kbmod/search/stack_search.cpp b/src/kbmod/search/stack_search.cpp index 951ec8ca6..f2b2239c0 100644 --- a/src/kbmod/search/stack_search.cpp +++ b/src/kbmod/search/stack_search.cpp @@ -256,17 +256,33 @@ std::vector<float> StackSearch::extract_psi_or_phi_curve(Trajectory& trj, bool e return result; } +std::vector<std::vector<float> > StackSearch::get_psi_curves(std::vector<Trajectory>& trajectories) { + std::vector<std::vector<float> > all_results; + for (auto& trj : trajectories) { + all_results.push_back(extract_psi_or_phi_curve(trj, true)); + } + return all_results; +} + std::vector<float> StackSearch::get_psi_curves(Trajectory& trj) { return extract_psi_or_phi_curve(trj, true); } +std::vector<std::vector<float> > StackSearch::get_phi_curves(std::vector<Trajectory>& trajectories) { + std::vector<std::vector<float> > all_results; + for (auto& trj : trajectories) { + all_results.push_back(extract_psi_or_phi_curve(trj, false)); + } + return all_results; +} + std::vector<float> StackSearch::get_phi_curves(Trajectory& trj) { return extract_psi_or_phi_curve(trj, false); } std::vector<Trajectory> StackSearch::get_results(uint64_t start, uint64_t count) { - rs_logger->debug("Reading results [" + std::to_string(start) + ", " + - std::to_string(start + count) + ")"); + rs_logger->debug("Reading results [" + std::to_string(start) + ", " + std::to_string(start + count) + + ")"); return results.get_batch(start, count); } @@ -310,6 +326,12 @@ static void stack_search_bindings(py::module& m) { pydocs::DOC_StackSearch_get_psi_curves) .def("get_phi_curves", (std::vector<float>(ks::*)(tj&)) & ks::get_phi_curves, pydocs::DOC_StackSearch_get_phi_curves) + .def("get_psi_curves", + (std::vector<std::vector<float> >(ks::*)(std::vector<tj>&)) & ks::get_psi_curves, + pydocs::DOC_StackSearch_get_psi_curves) + .def("get_phi_curves", + (std::vector<std::vector<float> >(ks::*)(std::vector<tj>&)) & ks::get_phi_curves, + pydocs::DOC_StackSearch_get_phi_curves) .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_number_total_results", &ks::get_number_total_results, diff --git a/src/kbmod/search/stack_search.h b/src/kbmod/search/stack_search.h index 43b0bac03..6e00b8c22 100644 --- a/src/kbmod/search/stack_search.h +++ b/src/kbmod/search/stack_search.h @@ -62,6 +62,8 @@ class StackSearch { // Getters for the Psi and Phi data. std::vector<float> get_psi_curves(Trajectory& t); std::vector<float> get_phi_curves(Trajectory& t); + std::vector<std::vector<float> > get_psi_curves(std::vector<Trajectory>& trajectories); + std::vector<std::vector<float> > get_phi_curves(std::vector<Trajectory>& trajectories); // Helper functions for computing Psi and Phi void prepare_psi_phi(); diff --git a/src/kbmod/search/trajectory_list.cpp b/src/kbmod/search/trajectory_list.cpp index 357d7a40f..10608fba8 100644 --- a/src/kbmod/search/trajectory_list.cpp +++ b/src/kbmod/search/trajectory_list.cpp @@ -85,7 +85,7 @@ void TrajectoryList::filter_by_likelihood(float min_likelihood) { // Drop the values below the threshold. resize(index); } - + void TrajectoryList::move_to_gpu() { if (data_on_gpu) return; // Nothing to do. diff --git a/src/kbmod/search/trajectory_list.h b/src/kbmod/search/trajectory_list.h index f710eafc4..fc6d6c4b4 100644 --- a/src/kbmod/search/trajectory_list.h +++ b/src/kbmod/search/trajectory_list.h @@ -61,7 +61,6 @@ class TrajectoryList { // Processing functions for sorting/filtering. void sort_by_likelihood(); void filter_by_likelihood(float min_likelihood); - // Data allocation functions. inline bool on_gpu() const { return data_on_gpu; } From 6f82c714d5cc44ff23475267bb92a2193ae74b6a Mon Sep 17 00:00:00 2001 From: Jeremy Kubica <104161096+jeremykubica@users.noreply.github.com> Date: Tue, 17 Sep 2024 15:50:15 -0400 Subject: [PATCH 03/12] Update stack_search_docs.h --- src/kbmod/search/pydocs/stack_search_docs.h | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/src/kbmod/search/pydocs/stack_search_docs.h b/src/kbmod/search/pydocs/stack_search_docs.h index adb9946a3..576fadee3 100644 --- a/src/kbmod/search/pydocs/stack_search_docs.h +++ b/src/kbmod/search/pydocs/stack_search_docs.h @@ -130,12 +130,12 @@ static const auto DOC_StackSearch_get_psi_curves = R"doc( Parameters ---------- - trj : `kb.Trajectory` - The input trajectory. + trj : `kb.Trajectory` or `list` of `kb.Trajectory` + The input trajectory or trajectories. Returns ------- - result : `list` of `float` + result : `list` of `float` or `list` of `list` of `float` The psi values at each time step with NO_DATA replaced by 0.0. )doc"; @@ -144,12 +144,12 @@ static const auto DOC_StackSearch_get_phi_curves = R"doc( Parameters ---------- - trj : `kb.Trajectory` - The input trajectory. + trj : `kb.Trajectory` or `list` of `kb.Trajectory` + The input trajectory or trajectories. Returns ------- - result : `list` of `float` + result : `list` of `float` or `list` of `list` of `float` The phi values at each time step with NO_DATA replaced by 0.0. )doc"; From 64be778f302514fe00a4bc88c172caee96699db0 Mon Sep 17 00:00:00 2001 From: Jeremy Kubica <104161096+jeremykubica@users.noreply.github.com> Date: Wed, 18 Sep 2024 10:07:33 -0400 Subject: [PATCH 04/12] Add a test --- tests/test_search.py | 100 -------------------- tests/test_stack_search_results.py | 147 +++++++++++++++++++++++++++++ 2 files changed, 147 insertions(+), 100 deletions(-) create mode 100644 tests/test_stack_search_results.py diff --git a/tests/test_search.py b/tests/test_search.py index 40837640f..4072e8eb3 100644 --- a/tests/test_search.py +++ b/tests/test_search.py @@ -103,106 +103,6 @@ def setUp(self): self.max_angle, ) - def test_set_get_results(self): - results = self.search.get_results(0, 10) - self.assertEqual(len(results), 0) - - trjs = [Trajectory(i, i, 0.0, 0.0) for i in range(10)] - self.search.set_results(trjs) - - # Check that we extract them all. - results = self.search.get_results(0, 10) - self.assertEqual(len(results), 10) - for i in range(10): - self.assertEqual(results[i].x, i) - - # Check that we can run past the end of the results. - results = self.search.get_results(0, 100) - self.assertEqual(len(results), 10) - - # Check that we can pull a subset. - results = self.search.get_results(2, 2) - self.assertEqual(len(results), 2) - self.assertEqual(results[0].x, 2) - self.assertEqual(results[1].x, 3) - - # Check that we can pull a subset aligned with the end. - results = self.search.get_results(8, 2) - self.assertEqual(len(results), 2) - self.assertEqual(results[0].x, 8) - self.assertEqual(results[1].x, 9) - - # Check invalid settings - self.assertRaises(RuntimeError, self.search.get_results, 0, 0) - - def test_load_and_filter_results_lh(self): - time_list = [i / self.img_count for i in range(self.img_count)] - fake_ds = FakeDataSet( - self.dim_x, - self.dim_y, - time_list, - noise_level=1.0, - psf_val=0.5, - use_seed=True, - ) - - # Create fake result trajectories with given initial likelihoods. The 1st is - # filtered by max likelihood. The two final ones are filtered by min likelihood. - trjs = [ - Trajectory(10, 10, 0, 0, 500.0, 9000.0, self.img_count), - Trajectory(20, 20, 0, 0, 110.0, 110.0, self.img_count), - Trajectory(30, 30, 0, 0, 100.0, 100.0, self.img_count), - Trajectory(40, 40, 0, 0, 50.0, 50.0, self.img_count), - Trajectory(41, 41, 0, 0, 50.0, 50.0, self.img_count), - Trajectory(42, 42, 0, 0, 50.0, 50.0, self.img_count), - Trajectory(43, 43, 0, 0, 50.0, 50.0, self.img_count), - Trajectory(50, 50, 0, 0, 1.0, 2.0, self.img_count), - Trajectory(60, 60, 0, 0, 1.0, 1.0, self.img_count), - ] - for trj in trjs: - fake_ds.insert_object(trj) - - # Create the stack search and insert the fake results. - search = StackSearch(fake_ds.stack) - search.set_results(trjs) - - # Do the loading and filtering. - config = SearchConfiguration() - overrides = { - "clip_negative": False, - "chunk_size": 500000, - "lh_level": 10.0, - "max_lh": 1000.0, - "num_cores": 1, - "num_obs": 5, - "sigmaG_lims": [25, 75], - } - config.set_multiple(overrides) - - runner = SearchRunner() - results = runner.load_and_filter_results(search, config) - - # Only two of the middle results should pass the filtering. - self.assertEqual(len(results), 6) - self.assertEqual(results["y"][0], 20) - self.assertEqual(results["y"][1], 30) - self.assertEqual(results["y"][2], 40) - self.assertEqual(results["y"][3], 41) - self.assertEqual(results["y"][4], 42) - self.assertEqual(results["y"][5], 43) - - # Rerun the search with a small chunk_size to make sure we still - # find everything. - overrides["chunk_size"] = 2 - results = runner.load_and_filter_results(search, config) - self.assertEqual(len(results), 6) - self.assertEqual(results["y"][0], 20) - self.assertEqual(results["y"][1], 30) - self.assertEqual(results["y"][2], 40) - self.assertEqual(results["y"][3], 41) - self.assertEqual(results["y"][4], 42) - self.assertEqual(results["y"][5], 43) - @unittest.skipIf(not HAS_GPU, "Skipping test (no GPU detected)") def test_evaluate_single_trajectory(self): test_trj = Trajectory( diff --git a/tests/test_stack_search_results.py b/tests/test_stack_search_results.py new file mode 100644 index 000000000..1a8de0edc --- /dev/null +++ b/tests/test_stack_search_results.py @@ -0,0 +1,147 @@ +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.run_search import SearchRunner +from kbmod.search import StackSearch, Trajectory + + +class test_search(unittest.TestCase): + def setUp(self): + self.num_times = 10 + self.width = 256 + self.height = 256 + self.num_objs = 5 + + self.times = create_fake_times(self.num_times, obs_per_day=3) + self.fake_ds = FakeDataSet(self.width, self.height, self.times) + for _ in range(self.num_objs): + self.fake_ds.insert_random_object(500) + + self.search = StackSearch(self.fake_ds.stack) + self.fake_trjs = self.fake_ds.trajectories + + def test_set_get_results(self): + results = self.search.get_results(0, 10) + self.assertEqual(len(results), 0) + + trjs = [Trajectory(i, i, 0.0, 0.0) for i in range(10)] + self.search.set_results(trjs) + + # Check that we extract them all. + results = self.search.get_results(0, 10) + self.assertEqual(len(results), 10) + for i in range(10): + self.assertEqual(results[i].x, i) + + # Check that we can run past the end of the results. + results = self.search.get_results(0, 100) + self.assertEqual(len(results), 10) + + # Check that we can pull a subset. + results = self.search.get_results(2, 2) + self.assertEqual(len(results), 2) + self.assertEqual(results[0].x, 2) + self.assertEqual(results[1].x, 3) + + # Check that we can pull a subset aligned with the end. + results = self.search.get_results(8, 2) + self.assertEqual(len(results), 2) + self.assertEqual(results[0].x, 8) + self.assertEqual(results[1].x, 9) + + # Check invalid settings + self.assertRaises(RuntimeError, self.search.get_results, 0, 0) + + def test_psi_phi_curves(self): + psi_curves = np.array(self.search.get_psi_curves(self.fake_trjs)) + self.assertEqual(psi_curves.shape[0], self.num_objs) + self.assertEqual(psi_curves.shape[1], self.num_times) + self.assertTrue(np.all(psi_curves > 0.0)) + + phi_curves = np.array(self.search.get_phi_curves(self.fake_trjs)) + self.assertEqual(phi_curves.shape[0], self.num_objs) + self.assertEqual(phi_curves.shape[1], self.num_times) + self.assertTrue(np.all(phi_curves > 0.0)) + + # Check that the batch getters give the same results as the iterative ones. + for i in range(self.num_objs): + current_psi = self.search.get_psi_curves(self.fake_trjs[i]) + self.assertTrue(np.allclose(psi_curves[i], current_psi)) + + current_phi = self.search.get_phi_curves(self.fake_trjs[i]) + self.assertTrue(np.allclose(phi_curves[i], current_phi)) + + def test_load_and_filter_results_lh(self): + time_list = [i / self.num_times for i in range(self.num_times)] + fake_ds = FakeDataSet( + self.width, + self.height, + time_list, + noise_level=1.0, + psf_val=0.5, + use_seed=True, + ) + + # Create fake result trajectories with given initial likelihoods. The 1st is + # filtered by max likelihood. The two final ones are filtered by min likelihood. + trjs = [ + Trajectory(10, 10, 0, 0, 500.0, 9000.0, self.num_times), + Trajectory(20, 20, 0, 0, 110.0, 110.0, self.num_times), + Trajectory(30, 30, 0, 0, 100.0, 100.0, self.num_times), + Trajectory(40, 40, 0, 0, 50.0, 50.0, self.num_times), + Trajectory(41, 41, 0, 0, 50.0, 50.0, self.num_times), + Trajectory(42, 42, 0, 0, 50.0, 50.0, self.num_times), + Trajectory(43, 43, 0, 0, 50.0, 50.0, self.num_times), + Trajectory(50, 50, 0, 0, 1.0, 2.0, self.num_times), + Trajectory(60, 60, 0, 0, 1.0, 1.0, self.num_times), + ] + for trj in trjs: + fake_ds.insert_object(trj) + + # Create the stack search and insert the fake results. + search = StackSearch(fake_ds.stack) + search.set_results(trjs) + + # Do the loading and filtering. + config = SearchConfiguration() + overrides = { + "clip_negative": False, + "chunk_size": 500000, + "lh_level": 10.0, + "max_lh": 1000.0, + "num_cores": 1, + "num_obs": 5, + "sigmaG_lims": [25, 75], + } + config.set_multiple(overrides) + + runner = SearchRunner() + results = runner.load_and_filter_results(search, config) + + # Only two of the middle results should pass the filtering. + self.assertEqual(len(results), 6) + self.assertEqual(results["y"][0], 20) + self.assertEqual(results["y"][1], 30) + self.assertEqual(results["y"][2], 40) + self.assertEqual(results["y"][3], 41) + self.assertEqual(results["y"][4], 42) + self.assertEqual(results["y"][5], 43) + + # Rerun the search with a small chunk_size to make sure we still + # find everything. + overrides["chunk_size"] = 2 + results = runner.load_and_filter_results(search, config) + self.assertEqual(len(results), 6) + self.assertEqual(results["y"][0], 20) + self.assertEqual(results["y"][1], 30) + self.assertEqual(results["y"][2], 40) + self.assertEqual(results["y"][3], 41) + self.assertEqual(results["y"][4], 42) + self.assertEqual(results["y"][5], 43) + + +if __name__ == "__main__": + unittest.main() From c0f1becd8980c0a98ffc1e14ae028cc3d79f25a5 Mon Sep 17 00:00:00 2001 From: Jeremy Kubica <104161096+jeremykubica@users.noreply.github.com> Date: Wed, 18 Sep 2024 10:23:33 -0400 Subject: [PATCH 05/12] Revert C++ filtering changes --- src/kbmod/search/pydocs/trajectory_list_docs.h | 14 -------------- src/kbmod/search/trajectory_list.cpp | 15 --------------- src/kbmod/search/trajectory_list.h | 3 +-- tests/test_trajectory_list.py | 16 ---------------- 4 files changed, 1 insertion(+), 47 deletions(-) diff --git a/src/kbmod/search/pydocs/trajectory_list_docs.h b/src/kbmod/search/pydocs/trajectory_list_docs.h index 52ddb747c..23ae5f43f 100644 --- a/src/kbmod/search/pydocs/trajectory_list_docs.h +++ b/src/kbmod/search/pydocs/trajectory_list_docs.h @@ -143,20 +143,6 @@ static const auto DOC_TrajectoryList_sort_by_likelihood = R"doc( Raises a ``RuntimeError`` the data is on GPU. )doc"; -static const auto DOC_TrajectoryList_filter_by_likelihood = R"doc( - Sort the data in order of decreasing likelihood and drop everything less than - a given threshold. The data must reside on the CPU. - - Parameters - ---------- - min_likelihood : `float` - The threshold on minimum likelihood. - - Raises - ------ - Raises a ``RuntimeError`` the data is on GPU. - )doc"; - } // namespace pydocs diff --git a/src/kbmod/search/trajectory_list.cpp b/src/kbmod/search/trajectory_list.cpp index 10608fba8..33f09b3db 100644 --- a/src/kbmod/search/trajectory_list.cpp +++ b/src/kbmod/search/trajectory_list.cpp @@ -73,19 +73,6 @@ void TrajectoryList::sort_by_likelihood() { [](const Trajectory& a, const Trajectory& b) { return b.lh < a.lh; }); } -void TrajectoryList::filter_by_likelihood(float min_likelihood) { - sort_by_likelihood(); - - // Find the first index that does not meet the threshold. - uint64_t index = 0; - while ((index < max_size) && (cpu_list[index].lh >= min_likelihood)) { - ++index; - } - - // Drop the values below the threshold. - resize(index); -} - void TrajectoryList::move_to_gpu() { if (data_on_gpu) return; // Nothing to do. @@ -132,8 +119,6 @@ static void trajectory_list_binding(py::module& m) { .def("get_batch", &trjl::get_batch, pydocs::DOC_TrajectoryList_get_batch) .def("sort_by_likelihood", &trjl::sort_by_likelihood, pydocs::DOC_TrajectoryList_sort_by_likelihood) - .def("filter_by_likelihood", &trjl::filter_by_likelihood, - pydocs::DOC_TrajectoryList_filter_by_likelihood) .def("move_to_cpu", &trjl::move_to_cpu, pydocs::DOC_TrajectoryList_move_to_cpu) .def("move_to_gpu", &trjl::move_to_gpu, pydocs::DOC_TrajectoryList_move_to_gpu); } diff --git a/src/kbmod/search/trajectory_list.h b/src/kbmod/search/trajectory_list.h index fc6d6c4b4..6ee36d17d 100644 --- a/src/kbmod/search/trajectory_list.h +++ b/src/kbmod/search/trajectory_list.h @@ -58,9 +58,8 @@ class TrajectoryList { // Get a batch of results. std::vector<Trajectory> get_batch(uint64_t start, uint64_t count); - // Processing functions for sorting/filtering. + // Processing functions for sorting. void sort_by_likelihood(); - void filter_by_likelihood(float min_likelihood); // Data allocation functions. inline bool on_gpu() const { return data_on_gpu; } diff --git a/tests/test_trajectory_list.py b/tests/test_trajectory_list.py index d70d80694..a8b008445 100644 --- a/tests/test_trajectory_list.py +++ b/tests/test_trajectory_list.py @@ -91,22 +91,6 @@ def test_sort(self): for i in range(5): self.assertEqual(trjs.get_trajectory(i).x, lh_order[i]) - def test_filter_on_lh(self): - lh = [100.0, 110.0, 90.0, 120.0, 125.0, 121.0, 10.0] - trjs = TrajectoryList(len(lh)) - for i in range(len(lh)): - trjs.set_trajectory(i, Trajectory(x=i, lh=lh[i])) - - trjs.filter_by_likelihood(110.0) - expected = set([4, 5, 3, 1]) - - # Test that each remaining result appears once in the expected set. - self.assertEqual(len(trjs), len(expected)) - for i in range(len(trjs)): - idx = trjs.get_trajectory(i).x - self.assertTrue(idx in expected) - expected.remove(idx) - @unittest.skipIf(not HAS_GPU, "Skipping test (no GPU detected)") def test_move_to_from_gpu(self): for i in range(self.max_size): From 37d42393a048a241a70c247ef557ba606f35b4af Mon Sep 17 00:00:00 2001 From: Jeremy Kubica <104161096+jeremykubica@users.noreply.github.com> Date: Wed, 18 Sep 2024 10:25:36 -0400 Subject: [PATCH 06/12] Do batched curve lookup --- src/kbmod/run_search.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/src/kbmod/run_search.py b/src/kbmod/run_search.py index 642c452a9..77fcecff7 100644 --- a/src/kbmod/run_search.py +++ b/src/kbmod/run_search.py @@ -82,8 +82,6 @@ def load_and_filter_results(self, search, config): logger.info(f"Chunk Min. Likelihood = {results[-1].lh}") trj_batch = [] - psi_batch = [] - phi_batch = [] for i, trj in enumerate(results): # Stop as soon as we hit a result below our limit, because anything after # that is not guarrenteed to be valid due to potential on-GPU filtering. @@ -93,14 +91,15 @@ def load_and_filter_results(self, search, config): if trj.lh < max_lh: trj_batch.append(trj) - psi_batch.append(search.get_psi_curves(trj)) - phi_batch.append(search.get_phi_curves(trj)) total_count += 1 batch_size = len(trj_batch) logger.info(f"Extracted batch of {batch_size} results for total of {total_count}") if batch_size > 0: + psi_batch = search.get_psi_curves(trj_batch) + phi_batch = search.get_phi_curves(trj_batch) + result_batch = Results.from_trajectories(trj_batch, track_filtered=do_tracking) result_batch.add_psi_phi_data(psi_batch, phi_batch) From 3f9e86b0dcdf0eadc455b9407dfe675ac9d3282a Mon Sep 17 00:00:00 2001 From: Jeremy Kubica <104161096+jeremykubica@users.noreply.github.com> Date: Wed, 18 Sep 2024 10:41:57 -0400 Subject: [PATCH 07/12] Fix repeated doc string --- src/kbmod/search/stack_search.cpp | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/src/kbmod/search/stack_search.cpp b/src/kbmod/search/stack_search.cpp index f2b2239c0..711637261 100644 --- a/src/kbmod/search/stack_search.cpp +++ b/src/kbmod/search/stack_search.cpp @@ -327,11 +327,9 @@ static void stack_search_bindings(py::module& m) { .def("get_phi_curves", (std::vector<float>(ks::*)(tj&)) & ks::get_phi_curves, pydocs::DOC_StackSearch_get_phi_curves) .def("get_psi_curves", - (std::vector<std::vector<float> >(ks::*)(std::vector<tj>&)) & ks::get_psi_curves, - pydocs::DOC_StackSearch_get_psi_curves) + (std::vector<std::vector<float> >(ks::*)(std::vector<tj>&)) & ks::get_psi_curves) .def("get_phi_curves", - (std::vector<std::vector<float> >(ks::*)(std::vector<tj>&)) & ks::get_phi_curves, - pydocs::DOC_StackSearch_get_phi_curves) + (std::vector<std::vector<float> >(ks::*)(std::vector<tj>&)) & ks::get_phi_curves) .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_number_total_results", &ks::get_number_total_results, From e838ca1c8f868fb8746760e36815766653ff8f87 Mon Sep 17 00:00:00 2001 From: Jeremy Kubica <104161096+jeremykubica@users.noreply.github.com> Date: Wed, 18 Sep 2024 12:34:22 -0400 Subject: [PATCH 08/12] Make the vectors const --- src/kbmod/search/stack_search.cpp | 22 +++++++++++----------- src/kbmod/search/stack_search.h | 10 +++++----- 2 files changed, 16 insertions(+), 16 deletions(-) diff --git a/src/kbmod/search/stack_search.cpp b/src/kbmod/search/stack_search.cpp index 711637261..0123b9dc9 100644 --- a/src/kbmod/search/stack_search.cpp +++ b/src/kbmod/search/stack_search.cpp @@ -237,7 +237,7 @@ uint64_t StackSearch::compute_max_results() { return num_search_pixels * params.results_per_pixel; } -std::vector<float> StackSearch::extract_psi_or_phi_curve(Trajectory& trj, bool extract_psi) { +std::vector<float> StackSearch::extract_psi_or_phi_curve(const Trajectory& trj, bool extract_psi) { prepare_psi_phi(); const unsigned int num_times = stack.img_count(); @@ -256,27 +256,27 @@ std::vector<float> StackSearch::extract_psi_or_phi_curve(Trajectory& trj, bool e return result; } -std::vector<std::vector<float> > StackSearch::get_psi_curves(std::vector<Trajectory>& trajectories) { +std::vector<std::vector<float> > StackSearch::get_psi_curves(const std::vector<Trajectory>& trajectories) { std::vector<std::vector<float> > all_results; - for (auto& trj : trajectories) { + for (const auto& trj : trajectories) { all_results.push_back(extract_psi_or_phi_curve(trj, true)); } return all_results; } -std::vector<float> StackSearch::get_psi_curves(Trajectory& trj) { +std::vector<float> StackSearch::get_psi_curves(const Trajectory& trj) { return extract_psi_or_phi_curve(trj, true); } -std::vector<std::vector<float> > StackSearch::get_phi_curves(std::vector<Trajectory>& trajectories) { +std::vector<std::vector<float> > StackSearch::get_phi_curves(const std::vector<Trajectory>& trajectories) { std::vector<std::vector<float> > all_results; - for (auto& trj : trajectories) { + for (const auto& trj : trajectories) { all_results.push_back(extract_psi_or_phi_curve(trj, false)); } return all_results; } -std::vector<float> StackSearch::get_phi_curves(Trajectory& trj) { +std::vector<float> StackSearch::get_phi_curves(const Trajectory& trj) { return extract_psi_or_phi_curve(trj, false); } @@ -322,14 +322,14 @@ static void stack_search_bindings(py::module& m) { .def("get_imagestack", &ks::get_imagestack, py::return_value_policy::reference_internal, pydocs::DOC_StackSearch_get_imagestack) // For testings - .def("get_psi_curves", (std::vector<float>(ks::*)(tj&)) & ks::get_psi_curves, + .def("get_psi_curves", (std::vector<float>(ks::*)(const tj&)) & ks::get_psi_curves, pydocs::DOC_StackSearch_get_psi_curves) - .def("get_phi_curves", (std::vector<float>(ks::*)(tj&)) & ks::get_phi_curves, + .def("get_phi_curves", (std::vector<float>(ks::*)(const tj&)) & ks::get_phi_curves, pydocs::DOC_StackSearch_get_phi_curves) .def("get_psi_curves", - (std::vector<std::vector<float> >(ks::*)(std::vector<tj>&)) & ks::get_psi_curves) + (std::vector<std::vector<float> >(ks::*)(const std::vector<tj>&)) & ks::get_psi_curves) .def("get_phi_curves", - (std::vector<std::vector<float> >(ks::*)(std::vector<tj>&)) & ks::get_phi_curves) + (std::vector<std::vector<float> >(ks::*)(const std::vector<tj>&)) & ks::get_phi_curves) .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_number_total_results", &ks::get_number_total_results, diff --git a/src/kbmod/search/stack_search.h b/src/kbmod/search/stack_search.h index 6e00b8c22..7a629bea2 100644 --- a/src/kbmod/search/stack_search.h +++ b/src/kbmod/search/stack_search.h @@ -60,10 +60,10 @@ class StackSearch { std::vector<Trajectory> get_results(uint64_t start, uint64_t count); // Getters for the Psi and Phi data. - std::vector<float> get_psi_curves(Trajectory& t); - std::vector<float> get_phi_curves(Trajectory& t); - std::vector<std::vector<float> > get_psi_curves(std::vector<Trajectory>& trajectories); - std::vector<std::vector<float> > get_phi_curves(std::vector<Trajectory>& trajectories); + std::vector<float> get_psi_curves(const Trajectory& t); + std::vector<float> get_phi_curves(const Trajectory& t); + std::vector<std::vector<float> > get_psi_curves(const std::vector<Trajectory>& trajectories); + std::vector<std::vector<float> > get_phi_curves(const std::vector<Trajectory>& trajectories); // Helper functions for computing Psi and Phi void prepare_psi_phi(); @@ -75,7 +75,7 @@ class StackSearch { virtual ~StackSearch(){}; protected: - std::vector<float> extract_psi_or_phi_curve(Trajectory& trj, bool extract_psi); + std::vector<float> extract_psi_or_phi_curve(const Trajectory& trj, bool extract_psi); // Core data and search parameters. Note the StackSearch does not own // the ImageStack and it must exist for the duration of the object's life. From 6f4288980f3a63a9c129f56d577299dfa2db792a Mon Sep 17 00:00:00 2001 From: Jeremy Kubica <104161096+jeremykubica@users.noreply.github.com> Date: Fri, 20 Sep 2024 10:15:49 -0400 Subject: [PATCH 09/12] Fix a few warnings that come up in doc generation --- src/kbmod/reprojection.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/src/kbmod/reprojection.py b/src/kbmod/reprojection.py index 9657df41c..11cd1d7ae 100644 --- a/src/kbmod/reprojection.py +++ b/src/kbmod/reprojection.py @@ -34,7 +34,7 @@ def reproject_image(image, original_wcs, common_wcs): The WCS to reproject all the images into. Returns - ---------- + ------- new_image : `numpy.ndarray` The image data reprojected with a common `astropy.wcs.WCS`. footprint : `numpy.ndarray` @@ -115,7 +115,7 @@ def reproject_work_unit( displayed or hidden. Returns - ---------- + ------- A `kbmod.WorkUnit` reprojected with a common `astropy.wcs.WCS`, or `None` in the case where `write_output` is set to True. """ @@ -186,8 +186,9 @@ def _reproject_work_unit( The base filename where output will be written if `write_output` is set to True. disable_show_progress : `bool` Whether or not to disable the `tqdm` show_progress bar. + Returns - ---------- + ------- A `kbmod.WorkUnit` reprojected with a common `astropy.wcs.WCS`, or `None` in the case where `write_output` is set to True. """ @@ -344,7 +345,7 @@ def _reproject_work_unit_in_parallel( Whether or not to enable the `tqdm` show_progress bar. Returns - ---------- + ------- A `kbmod.WorkUnit` reprojected with a common `astropy.wcs.WCS`, or `None` in the case where `write_output` is set to True. """ From e0c30eddc5b37c994a1bb21e88f87aa3964adb68 Mon Sep 17 00:00:00 2001 From: Jeremy Kubica <104161096+jeremykubica@users.noreply.github.com> Date: Fri, 20 Sep 2024 10:21:48 -0400 Subject: [PATCH 10/12] Fix more sphinx warnings --- src/kbmod/search/pydocs/stack_search_docs.h | 8 ++++---- src/kbmod/wcs_utils.py | 6 +++--- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/src/kbmod/search/pydocs/stack_search_docs.h b/src/kbmod/search/pydocs/stack_search_docs.h index adb9946a3..2289f3c13 100644 --- a/src/kbmod/search/pydocs/stack_search_docs.h +++ b/src/kbmod/search/pydocs/stack_search_docs.h @@ -249,8 +249,8 @@ static const auto DOC_StackSearch_evaluate_single_trajectory = R"doc( Performs the evaluation of a single Trajectory object. Modifies the object in-place. - Note - ---- + Notes + ----- Runs on the CPU, but requires CUDA compiler. Parameters @@ -262,8 +262,8 @@ static const auto DOC_StackSearch_evaluate_single_trajectory = R"doc( static const auto DOC_StackSearch_search_linear_trajectory = R"doc( Performs the evaluation of a linear trajectory in pixel space. - Note - ---- + Notes + ----- Runs on the CPU, but requires CUDA compiler. Parameters diff --git a/src/kbmod/wcs_utils.py b/src/kbmod/wcs_utils.py index 2df5b4491..60656e16f 100644 --- a/src/kbmod/wcs_utils.py +++ b/src/kbmod/wcs_utils.py @@ -301,8 +301,8 @@ def calc_ecliptic_angle(wcs, center_pixel=(1000, 2000), step=12): with the image axes. Used to transform the specified search angles, with respect to the ecliptic, to search angles within the image. - Note - ---- + Notes + ----- It is not neccessary to calculate this angle for each image in an image set if they have all been warped to a common WCS. """ @@ -336,7 +336,7 @@ def extract_wcs_from_hdu_header(header): The header from which to read the data. Returns - -------- + ------- curr_wcs : `astropy.wcs.WCS` The WCS or None if it does not exist. """ From 1b6bab799cdc6e778a667820575142d0d5f31541 Mon Sep 17 00:00:00 2001 From: Jeremy Kubica <104161096+jeremykubica@users.noreply.github.com> Date: Fri, 20 Sep 2024 12:43:37 -0400 Subject: [PATCH 11/12] Remove tqdm_utils.py --- src/kbmod/reprojection.py | 8 ++++---- src/kbmod/tqdm_utils.py | 4 ---- src/kbmod/work_unit.py | 8 +++++--- 3 files changed, 9 insertions(+), 11 deletions(-) delete mode 100644 src/kbmod/tqdm_utils.py diff --git a/src/kbmod/reprojection.py b/src/kbmod/reprojection.py index 9657df41c..0f602a290 100644 --- a/src/kbmod/reprojection.py +++ b/src/kbmod/reprojection.py @@ -8,7 +8,6 @@ from kbmod import is_interactive from kbmod.search import KB_NO_DATA, PSF, ImageStack, LayeredImage, RawImage from kbmod.work_unit import WorkUnit -from kbmod.tqdm_utils import TQDMUtils from kbmod.wcs_utils import append_wcs_to_hdu_header from astropy.io import fits import os @@ -17,6 +16,7 @@ # The number of executors to use in the parallel reprojecting function. MAX_PROCESSES = 8 +_DEFAULT_TQDM_BAR = ("{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}]",) def reproject_image(image, original_wcs, common_wcs): @@ -197,7 +197,7 @@ def _reproject_work_unit( stack = ImageStack() for obstime_index, o_i in tqdm( enumerate(zip(unique_obstimes, unique_obstime_indices)), - bar_format=TQDMUtils.DEFAULT_TQDM_BAR_FORMAT, + bar_format=_DEFAULT_TQDM_BAR, desc="Reprojecting", disable=not show_progress, ): @@ -406,7 +406,7 @@ def _reproject_work_unit_in_parallel( tqdm( concurrent.futures.as_completed(future_reprojections), total=len(future_reprojections), - bar_format=TQDMUtils.DEFAULT_TQDM_BAR_FORMAT, + bar_format=_DEFAULT_TQDM_BAR, desc="Reprojecting", disable=not show_progress, ) @@ -538,7 +538,7 @@ def reproject_lazy_work_unit( tqdm( concurrent.futures.as_completed(future_reprojections), total=len(future_reprojections), - bar_format=TQDMUtils.DEFAULT_TQDM_BAR_FORMAT, + bar_format=_DEFAULT_TQDM_BAR, desc="Reprojecting", disable=not show_progress, ) diff --git a/src/kbmod/tqdm_utils.py b/src/kbmod/tqdm_utils.py deleted file mode 100644 index f3d34d8c3..000000000 --- a/src/kbmod/tqdm_utils.py +++ /dev/null @@ -1,4 +0,0 @@ -class TQDMUtils: - """Holds configuration for TQDM progress bars.""" - - DEFAULT_TQDM_BAR_FORMAT = "{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}]" diff --git a/src/kbmod/work_unit.py b/src/kbmod/work_unit.py index bf9e932bb..60eb5e58c 100644 --- a/src/kbmod/work_unit.py +++ b/src/kbmod/work_unit.py @@ -23,7 +23,9 @@ wcs_to_dict, ) from kbmod.reprojection_utils import invert_correct_parallax -from kbmod.tqdm_utils import TQDMUtils + + +_DEFAULT_WORKUNIT_TQDM_BAR = ("{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}]",) logger = Logging.getLogger(__name__) @@ -311,7 +313,7 @@ def from_fits(cls, filename, show_progress=None): # Read in all the image files. for i in tqdm( range(num_images), - bar_format=TQDMUtils.DEFAULT_TQDM_BAR_FORMAT, + bar_format=_DEFAULT_WORKUNIT_TQDM_BAR, desc="Loading images", disable=not show_progress, ): @@ -340,7 +342,7 @@ def from_fits(cls, filename, show_progress=None): constituent_images = [] for i in tqdm( range(n_constituents), - bar_format=TQDMUtils.DEFAULT_TQDM_BAR_FORMAT, + bar_format=_DEFAULT_WORKUNIT_TQDM_BAR, desc="Loading WCS", disable=not show_progress, ): From 136408d347fd5e2247511f828c1e45514d0d0839 Mon Sep 17 00:00:00 2001 From: Jeremy Kubica <104161096+jeremykubica@users.noreply.github.com> Date: Fri, 20 Sep 2024 13:24:42 -0400 Subject: [PATCH 12/12] Fix string --- src/kbmod/reprojection.py | 2 +- src/kbmod/work_unit.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/kbmod/reprojection.py b/src/kbmod/reprojection.py index 0f602a290..565f4d1b2 100644 --- a/src/kbmod/reprojection.py +++ b/src/kbmod/reprojection.py @@ -16,7 +16,7 @@ # The number of executors to use in the parallel reprojecting function. MAX_PROCESSES = 8 -_DEFAULT_TQDM_BAR = ("{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}]",) +_DEFAULT_TQDM_BAR = "{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}]" def reproject_image(image, original_wcs, common_wcs): diff --git a/src/kbmod/work_unit.py b/src/kbmod/work_unit.py index 60eb5e58c..b28902488 100644 --- a/src/kbmod/work_unit.py +++ b/src/kbmod/work_unit.py @@ -25,7 +25,7 @@ from kbmod.reprojection_utils import invert_correct_parallax -_DEFAULT_WORKUNIT_TQDM_BAR = ("{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}]",) +_DEFAULT_WORKUNIT_TQDM_BAR = "{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}]" logger = Logging.getLogger(__name__)