From 9f092b0747a937f75ed17994e47b7115aacb56fb Mon Sep 17 00:00:00 2001 From: Max West Date: Tue, 17 Dec 2024 14:57:37 -0800 Subject: [PATCH 1/3] functionalize and remove the StampCreator class --- benchmarks/bench_filter_stamps.py | 2 - .../kbmod_search_results_for_fakes.ipynb | 2 +- notebooks/kbmod_visualize.ipynb | 20 +++-- src/kbmod/analysis/visualizer.py | 4 +- src/kbmod/filters/stamp_filters.py | 7 +- src/kbmod/search/pydocs/stamp_creator_docs.h | 5 -- src/kbmod/search/stamp_creator.cpp | 52 +++++------ src/kbmod/search/stamp_creator.h | 87 +++++++++--------- src/kbmod/trajectory_explorer.py | 2 +- tests/test_stamp_creator.py | 89 ++++++++++--------- tests/test_stamp_parity.py | 24 ++--- 11 files changed, 144 insertions(+), 150 deletions(-) diff --git a/benchmarks/bench_filter_stamps.py b/benchmarks/bench_filter_stamps.py index 3e9b45176..43dae1b03 100644 --- a/benchmarks/bench_filter_stamps.py +++ b/benchmarks/bench_filter_stamps.py @@ -11,7 +11,6 @@ StampParameters, StampType, Trajectory, - StampCreator, ) @@ -55,7 +54,6 @@ def run_search_benchmark(params): # Create an empty search stack. # im_stack = ImageStack([]) - sc = StampCreator() # Do the timing runs. tmr = timeit.Timer(stmt="sc.filter_stamp(stamp, params)", globals=locals()) diff --git a/notebooks/kbmod_search_results_for_fakes.ipynb b/notebooks/kbmod_search_results_for_fakes.ipynb index a7b2c52db..52fe7a576 100644 --- a/notebooks/kbmod_search_results_for_fakes.ipynb +++ b/notebooks/kbmod_search_results_for_fakes.ipynb @@ -32,7 +32,7 @@ "import os\n", "\n", "from kbmod.analysis.plotting import *\n", - "from kbmod.search import ImageStack, PSF, StampCreator, Trajectory\n", + "from kbmod.search import ImageStack, PSF, Trajectory\n", "from kbmod.results import Results\n", "from kbmod.work_unit import WorkUnit\n", "\n", diff --git a/notebooks/kbmod_visualize.ipynb b/notebooks/kbmod_visualize.ipynb index ad7138040..2b24ad9e0 100644 --- a/notebooks/kbmod_visualize.ipynb +++ b/notebooks/kbmod_visualize.ipynb @@ -31,7 +31,15 @@ "\n", "from kbmod.analysis.plotting import *\n", "from kbmod.util_functions import load_deccam_layered_image\n", - "from kbmod.search import ImageStack, PSF, StampCreator, Trajectory\n", + "from kbmod.search import (\n", + " ImageStack,\n", + " PSF,\n", + " Trajectory,\n", + " get_stamps,\n", + " get_median_stamp,\n", + " get_mean_stamp,\n", + " get_summed_stamp,\n", + ")\n", "from kbmod.results import Results\n", "\n", "# Data paths\n", @@ -113,7 +121,7 @@ "source": [ "# Create and Visualize Stamps\n", "\n", - "Stamps are a critical tool for analyzing and debugging proposed detections. They can be created automatically using the `StampCreator` class. It requires a few pieces of data:\n", + "Stamps are a critical tool for analyzing and debugging proposed detections. They can be created automatically using the stamp creation utilities. It requires a few pieces of data:\n", "* search_stack - provides the machinery for making predictions on the image (needed to handle the various corrections).\n", "* trajectory - Contains the information about where to place the stamps (the underlying trajectory).\n", "* stamp_radius - The radius in pixels." @@ -133,7 +141,7 @@ "trj.vy = 3.3\n", "\n", "# Create the stamps around this trajectory.\n", - "stamps = StampCreator.get_stamps(stack, trj, 20)" + "stamps = get_stamps(stack, trj, 20)" ] }, { @@ -174,9 +182,9 @@ "\n", "plot_multiple_images(\n", " [\n", - " StampCreator.get_summed_stamp(stack, trj, 10, inds),\n", - " StampCreator.get_mean_stamp(stack, trj, 10, inds),\n", - " StampCreator.get_median_stamp(stack, trj, 10, inds),\n", + " get_summed_stamp(stack, trj, 10, inds),\n", + " get_mean_stamp(stack, trj, 10, inds),\n", + " get_median_stamp(stack, trj, 10, inds),\n", " ],\n", " labels=[\"Summed\", \"Mean\", \"Median\"],\n", " norm=True,\n", diff --git a/src/kbmod/analysis/visualizer.py b/src/kbmod/analysis/visualizer.py index 875e70ff4..a9030f8fc 100644 --- a/src/kbmod/analysis/visualizer.py +++ b/src/kbmod/analysis/visualizer.py @@ -1,5 +1,5 @@ from kbmod.analysis.plotting import plot_multiple_images -from kbmod.search import StampCreator +from kbmod.search import get_stamps from kbmod.util_functions import mjd_to_day import numpy as np @@ -33,7 +33,7 @@ def generate_all_stamps(self, radius=10): radius of the stamp. """ self.results.table["all_stamps"] = [ - StampCreator.get_stamps(self.im_stack, trj, radius) for trj in self.trajectories + get_stamps(self.im_stack, trj, radius) for trj in self.trajectories ] def count_num_days(self): diff --git a/src/kbmod/filters/stamp_filters.py b/src/kbmod/filters/stamp_filters.py index c4f537e8f..0fd7d08df 100644 --- a/src/kbmod/filters/stamp_filters.py +++ b/src/kbmod/filters/stamp_filters.py @@ -14,10 +14,11 @@ DebugTimer, ImageStack, RawImage, - StampCreator, StampParameters, StampType, Logging, + get_stamps, + get_coadded_stamps, ) @@ -141,7 +142,7 @@ def get_coadds_and_filter_results(result_data, im_stack, stamp_params, chunk_siz # Create and filter the results, using the GPU if there is one and enough # trajectories to make it worthwhile. - stamps_slice = StampCreator.get_coadded_stamps( + stamps_slice = get_coadded_stamps( im_stack, trj_slice, bool_slice, @@ -240,7 +241,7 @@ def append_all_stamps(result_data, im_stack, stamp_radius): all_stamps = [] for trj in result_data.make_trajectory_list(): - stamps = StampCreator.get_stamps(im_stack, trj, stamp_radius) + stamps = get_stamps(im_stack, trj, stamp_radius) all_stamps.append(np.array([stamp.image for stamp in stamps])) # We add the column even if it is empty so we can have consistent diff --git a/src/kbmod/search/pydocs/stamp_creator_docs.h b/src/kbmod/search/pydocs/stamp_creator_docs.h index 1b648397b..de0a0bea7 100644 --- a/src/kbmod/search/pydocs/stamp_creator_docs.h +++ b/src/kbmod/search/pydocs/stamp_creator_docs.h @@ -2,11 +2,6 @@ #define STAMP_CREATOR_DOCS namespace pydocs { -static const auto DOC_StampCreator = R"doc( - A class for creating a set of stamps or a co-added stamp - from an ImageStack and Trajectory. - )doc"; - static const auto DOC_StampCreator_create_stamps = R"doc( Create a vector of stamps centered on the predicted position of an Trajectory at different times. diff --git a/src/kbmod/search/stamp_creator.cpp b/src/kbmod/search/stamp_creator.cpp index c5d1341af..123264784 100644 --- a/src/kbmod/search/stamp_creator.cpp +++ b/src/kbmod/search/stamp_creator.cpp @@ -8,9 +8,7 @@ void deviceGetCoadds(const uint64_t num_images, const uint64_t width, const uint GPUArray& use_index_vect, GPUArray& results); #endif -StampCreator::StampCreator() {} - -std::vector StampCreator::create_stamps(ImageStack& stack, const Trajectory& trj, int radius, +std::vector create_stamps(ImageStack& stack, const Trajectory& trj, int radius, bool keep_no_data, const std::vector& use_index) { if (use_index.size() > 0) assert_sizes_equal(use_index.size(), stack.img_count(), "create_stamps() use_index"); @@ -35,33 +33,33 @@ std::vector StampCreator::create_stamps(ImageStack& stack, const Traje // For stamps used for visualization we replace invalid pixels with zeros // and return all the stamps (regardless of whether individual timesteps // have been filtered). -std::vector StampCreator::get_stamps(ImageStack& stack, const Trajectory& t, int radius) { +std::vector get_stamps(ImageStack& stack, const Trajectory& t, int radius) { std::vector empty_vect; return create_stamps(stack, t, radius, false /*=keep_no_data*/, empty_vect); } // For creating coadded stamps, we do not interpolate the pixel values and keep // invalid pixels tagged (so we can filter it out of mean/median). -RawImage StampCreator::get_median_stamp(ImageStack& stack, const Trajectory& trj, int radius, +RawImage get_median_stamp(ImageStack& stack, const Trajectory& trj, int radius, const std::vector& use_index) { return create_median_image(create_stamps(stack, trj, radius, true /*=keep_no_data*/, use_index)); } // For creating coadded stamps, we do not interpolate the pixel values and keep // invalid pixels tagged (so we can filter it out of mean/median). -RawImage StampCreator::get_mean_stamp(ImageStack& stack, const Trajectory& trj, int radius, +RawImage get_mean_stamp(ImageStack& stack, const Trajectory& trj, int radius, const std::vector& use_index) { return create_mean_image(create_stamps(stack, trj, radius, true /*=keep_no_data*/, use_index)); } // For creating summed stamps, we do not interpolate the pixel values and replace // invalid pixels with zero (which is the same as filtering it out for the sum). -RawImage StampCreator::get_summed_stamp(ImageStack& stack, const Trajectory& trj, int radius, +RawImage get_summed_stamp(ImageStack& stack, const Trajectory& trj, int radius, const std::vector& use_index) { return create_summed_image(create_stamps(stack, trj, radius, false /*=keep_no_data*/, use_index)); } -std::vector StampCreator::get_coadded_stamps(ImageStack& stack, std::vector& t_array, +std::vector get_coadded_stamps(ImageStack& stack, std::vector& t_array, std::vector>& use_index_vect, const StampParameters& params, bool use_gpu) { logging::Logger* rs_logger = logging::getLogger("kbmod.search.stamp_creator"); @@ -80,7 +78,7 @@ std::vector StampCreator::get_coadded_stamps(ImageStack& stack, std::v return get_coadded_stamps_cpu(stack, t_array, use_index_vect, params); } -std::vector StampCreator::get_coadded_stamps_cpu(ImageStack& stack, +std::vector get_coadded_stamps_cpu(ImageStack& stack, std::vector& t_array, std::vector>& use_index_vect, const StampParameters& params) { @@ -117,7 +115,7 @@ std::vector StampCreator::get_coadded_stamps_cpu(ImageStack& stack, return results; } -bool StampCreator::filter_stamp(const RawImage& img, const StampParameters& params) { +bool filter_stamp(const RawImage& img, const StampParameters& params) { if (params.radius <= 0) throw std::runtime_error("Invalid stamp radius=" + std::to_string(params.radius)); // Allocate space for the coadd information and initialize to zero. @@ -159,7 +157,7 @@ bool StampCreator::filter_stamp(const RawImage& img, const StampParameters& para return false; } -std::vector StampCreator::get_coadded_stamps_gpu(ImageStack& stack, +std::vector get_coadded_stamps_gpu(ImageStack& stack, std::vector& t_array, std::vector>& use_index_vect, const StampParameters& params) { @@ -273,7 +271,7 @@ std::vector StampCreator::get_coadded_stamps_gpu(ImageStack& stack, return results; } -std::vector StampCreator::create_variance_stamps(ImageStack& stack, const Trajectory& trj, +std::vector create_variance_stamps(ImageStack& stack, const Trajectory& trj, int radius, const std::vector& use_index) { if (use_index.size() > 0) assert_sizes_equal(use_index.size(), stack.img_count(), "create_stamps() use_index"); @@ -295,7 +293,7 @@ std::vector StampCreator::create_variance_stamps(ImageStack& stack, co return stamps; } -RawImage StampCreator::get_variance_weighted_stamp(ImageStack& stack, const Trajectory& trj, int radius, +RawImage get_variance_weighted_stamp(ImageStack& stack, const Trajectory& trj, int radius, const std::vector& use_index) { if (radius < 0) throw std::runtime_error("Invalid stamp radius. Must be >= 0."); unsigned int num_images = stack.img_count(); @@ -337,22 +335,18 @@ RawImage StampCreator::get_variance_weighted_stamp(ImageStack& stack, const Traj #ifdef Py_PYTHON_H static void stamp_creator_bindings(py::module& m) { - using sc = search::StampCreator; - - py::class_(m, "StampCreator", pydocs::DOC_StampCreator) - .def(py::init<>()) - .def_static("get_stamps", &sc::get_stamps, pydocs::DOC_StampCreator_get_stamps) - .def_static("get_median_stamp", &sc::get_median_stamp, pydocs::DOC_StampCreator_get_median_stamp) - .def_static("get_mean_stamp", &sc::get_mean_stamp, pydocs::DOC_StampCreator_get_mean_stamp) - .def_static("get_summed_stamp", &sc::get_summed_stamp, pydocs::DOC_StampCreator_get_summed_stamp) - .def_static("get_coadded_stamps", &sc::get_coadded_stamps, - pydocs::DOC_StampCreator_get_coadded_stamps) - .def_static("get_variance_weighted_stamp", &sc::get_variance_weighted_stamp, - pydocs::DOC_StampCreator_get_variance_weighted_stamp) - .def_static("create_stamps", &sc::create_stamps, pydocs::DOC_StampCreator_create_stamps) - .def_static("create_variance_stamps", &sc::create_variance_stamps, - pydocs::DOC_StampCreator_create_variance_stamps) - .def_static("filter_stamp", &sc::filter_stamp, pydocs::DOC_StampCreator_filter_stamp); + m.def("get_stamps", &search::get_stamps, pydocs::DOC_StampCreator_get_stamps); + m.def("get_median_stamp", &search::get_median_stamp, pydocs::DOC_StampCreator_get_median_stamp); + m.def("get_mean_stamp", &search::get_mean_stamp, pydocs::DOC_StampCreator_get_mean_stamp); + m.def("get_summed_stamp", &search::get_summed_stamp, pydocs::DOC_StampCreator_get_summed_stamp); + m.def("get_coadded_stamps", &search::get_coadded_stamps, + pydocs::DOC_StampCreator_get_coadded_stamps); + m.def("get_variance_weighted_stamp", &search::get_variance_weighted_stamp, + pydocs::DOC_StampCreator_get_variance_weighted_stamp); + m.def("create_stamps", &search::create_stamps, pydocs::DOC_StampCreator_create_stamps); + m.def("create_variance_stamps", &search::create_variance_stamps, + pydocs::DOC_StampCreator_create_variance_stamps); + m.def("filter_stamp", &search::filter_stamp, pydocs::DOC_StampCreator_filter_stamp); } #endif /* Py_PYTHON_H */ diff --git a/src/kbmod/search/stamp_creator.h b/src/kbmod/search/stamp_creator.h index 0cc924f8e..5174d404c 100644 --- a/src/kbmod/search/stamp_creator.h +++ b/src/kbmod/search/stamp_creator.h @@ -11,58 +11,51 @@ namespace search { * Utility class for functions used for creating science stamps for * filtering, visualization, etc. */ -class StampCreator { -public: - StampCreator(); - - // Functions science stamps for filtering, visualization, etc. User can specify - // the radius of the stamp, whether to keep no data values (e.g. NaN) or replace - // them with zero, and what indices to use. - // The indices to use are indicated by use_index: a vector indicating whether to use - // each time step. An empty (size=0) vector will use all time steps. - static std::vector create_stamps(ImageStack& stack, const Trajectory& trj, int radius, - bool keep_no_data, const std::vector& use_index); - - static std::vector get_stamps(ImageStack& stack, const Trajectory& t, int radius); - - static RawImage get_median_stamp(ImageStack& stack, const Trajectory& trj, int radius, - const std::vector& use_index); - - static RawImage get_mean_stamp(ImageStack& stack, const Trajectory& trj, int radius, - const std::vector& use_index); - - static RawImage get_summed_stamp(ImageStack& stack, const Trajectory& trj, int radius, - const std::vector& use_index); - - // Compute a mean or summed stamp for each trajectory on the GPU or CPU. - // The GPU implementation is slower for small numbers of trajectories (< 500), but performs - // relatively better as the number of trajectories increases. If filtering is applied then - // the code will return a 1x1 image with NO_DATA to represent each filtered image. - static std::vector get_coadded_stamps(ImageStack& stack, std::vector& t_array, +// Functions science stamps for filtering, visualization, etc. User can specify +// the radius of the stamp, whether to keep no data values (e.g. NaN) or replace +// them with zero, and what indices to use. +// The indices to use are indicated by use_index: a vector indicating whether to use +// each time step. An empty (size=0) vector will use all time steps. +static std::vector create_stamps(ImageStack& stack, const Trajectory& trj, int radius, + bool keep_no_data, const std::vector& use_index); + +static std::vector get_stamps(ImageStack& stack, const Trajectory& t, int radius); + +static RawImage get_median_stamp(ImageStack& stack, const Trajectory& trj, int radius, + const std::vector& use_index); + +static RawImage get_mean_stamp(ImageStack& stack, const Trajectory& trj, int radius, + const std::vector& use_index); + +static RawImage get_summed_stamp(ImageStack& stack, const Trajectory& trj, int radius, + const std::vector& use_index); + +// Compute a mean or summed stamp for each trajectory on the GPU or CPU. +// The GPU implementation is slower for small numbers of trajectories (< 500), but performs +// relatively better as the number of trajectories increases. If filtering is applied then +// the code will return a 1x1 image with NO_DATA to represent each filtered image. +static std::vector get_coadded_stamps(ImageStack& stack, std::vector& t_array, + std::vector >& use_index_vect, + const StampParameters& params, bool use_gpu); + +static std::vector get_coadded_stamps_gpu(ImageStack& stack, std::vector& t_array, std::vector >& use_index_vect, - const StampParameters& params, bool use_gpu); - - static std::vector get_coadded_stamps_gpu(ImageStack& stack, std::vector& t_array, - std::vector >& use_index_vect, - const StampParameters& params); + const StampParameters& params); - static std::vector get_coadded_stamps_cpu(ImageStack& stack, std::vector& t_array, - std::vector >& use_index_vect, - const StampParameters& params); - - // Function to do the actual stamp filtering. - static bool filter_stamp(const RawImage& img, const StampParameters& params); +static std::vector get_coadded_stamps_cpu(ImageStack& stack, std::vector& t_array, + std::vector >& use_index_vect, + const StampParameters& params); - // Function for generating variance stamps. All times are returned and NO_DATA values are preserved. - static std::vector create_variance_stamps(ImageStack& stack, const Trajectory& trj, int radius, - const std::vector& use_index); +// Function to do the actual stamp filtering. +static bool filter_stamp(const RawImage& img, const StampParameters& params); - // Function for generating variance weighted stamps. All times are used and NO_DATA values are skipped. - static RawImage get_variance_weighted_stamp(ImageStack& stack, const Trajectory& trj, int radius, - const std::vector& use_index); +// Function for generating variance stamps. All times are returned and NO_DATA values are preserved. +static std::vector create_variance_stamps(ImageStack& stack, const Trajectory& trj, int radius, + const std::vector& use_index); - virtual ~StampCreator(){}; -}; +// Function for generating variance weighted stamps. All times are used and NO_DATA values are skipped. +static RawImage get_variance_weighted_stamp(ImageStack& stack, const Trajectory& trj, int radius, + const std::vector& use_index); } /* namespace search */ diff --git a/src/kbmod/trajectory_explorer.py b/src/kbmod/trajectory_explorer.py index 06a3f8fca..56f7f398a 100644 --- a/src/kbmod/trajectory_explorer.py +++ b/src/kbmod/trajectory_explorer.py @@ -3,7 +3,7 @@ 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, StampCreator, Logging +from kbmod.search import StackSearch, Logging from kbmod.filters.stamp_filters import append_all_stamps, append_coadds from kbmod.trajectory_utils import make_trajectory_from_ra_dec diff --git a/tests/test_stamp_creator.py b/tests/test_stamp_creator.py index 58a9567ab..a61a1171d 100644 --- a/tests/test_stamp_creator.py +++ b/tests/test_stamp_creator.py @@ -14,10 +14,18 @@ pixel_value_valid, PSF, RawImage, - StampCreator, StampParameters, StampType, Trajectory, + get_stamps, + get_median_stamp, + get_mean_stamp, + get_summed_stamp, + get_coadded_stamps, + get_variance_weighted_stamp, + create_stamps, + create_variance_stamps, + filter_stamp, ) @@ -115,11 +123,8 @@ def setUp(self): self.trj2 = Trajectory(8, 7, 2.0, 1.0, flux=250.0) self.ds.insert_object(self.trj) - # Make a StampCreator. - self.stamp_creator = StampCreator() - def test_create_stamps(self): - stamps = self.stamp_creator.create_stamps(self.ds.stack, self.trj2, 1, True, []) + stamps = create_stamps(self.ds.stack, self.trj2, 1, True, []) self.assertEqual(len(stamps), self.img_count2) for i in range(self.img_count2): self.assertEqual(stamps[i].image.shape, (3, 3)) @@ -132,7 +137,7 @@ def test_create_stamps(self): # Check that we can set use_indices to produce only some stamps. use_times = [False, True, False, True, True, False, False, False, True, False] - stamps = self.stamp_creator.create_stamps(self.ds.stack, self.trj2, 1, True, use_times) + stamps = create_stamps(self.ds.stack, self.trj2, 1, True, use_times) self.assertEqual(len(stamps), np.count_nonzero(use_times)) stamp_count = 0 @@ -150,7 +155,7 @@ def test_create_stamps(self): def test_create_variance_stamps(self): test_trj = Trajectory(8, 7, 1.0, 2.0) - stamps = self.stamp_creator.create_variance_stamps(self.ds.stack, self.trj2, 1, []) + stamps = create_variance_stamps(self.ds.stack, self.trj2, 1, []) self.assertEqual(len(stamps), self.img_count2) for i in range(self.img_count2): self.assertEqual(stamps[i].image.shape, (3, 3)) @@ -163,7 +168,7 @@ def test_create_variance_stamps(self): # Check that we can set use_indices to produce only some stamps. use_times = [False, True, False, True, True, False, False, False, True, False] - stamps = self.stamp_creator.create_variance_stamps(self.ds.stack, self.trj2, 1, use_times) + stamps = create_variance_stamps(self.ds.stack, self.trj2, 1, use_times) self.assertEqual(len(stamps), np.count_nonzero(use_times)) stamp_count = 0 @@ -195,22 +200,22 @@ def test_get_variance_weighted_stamp(self): stack = ImageStack([layer1, layer2]) # Unmoving point in the center. Result should be (1.0 / 1.0 + 2.0 / 0.5) / (1.0 / 1.0 + 1.0 / 0.5) - stamp = self.stamp_creator.get_variance_weighted_stamp(stack, Trajectory(1, 1, 0.0, 0.0), 0, []) + stamp = get_variance_weighted_stamp(stack, Trajectory(1, 1, 0.0, 0.0), 0, []) self.assertEqual(stamp.image.shape, (1, 1)) self.assertAlmostEqual(stamp.get_pixel(0, 0), 5.0 / 3.0) # Unmoving point in the top corner. Should ignore the point in the second image. - stamp = self.stamp_creator.get_variance_weighted_stamp(stack, Trajectory(0, 0, 0.0, 0.0), 0, []) + stamp = get_variance_weighted_stamp(stack, Trajectory(0, 0, 0.0, 0.0), 0, []) self.assertEqual(stamp.image.shape, (1, 1)) self.assertAlmostEqual(stamp.get_pixel(0, 0), 1.0) # Unmoving point in the bottom corner. Should ignore the point in the first image. - stamp = self.stamp_creator.get_variance_weighted_stamp(stack, Trajectory(2, 2, 0.0, 0.0), 0, []) + stamp = get_variance_weighted_stamp(stack, Trajectory(2, 2, 0.0, 0.0), 0, []) self.assertEqual(stamp.image.shape, (1, 1)) self.assertAlmostEqual(stamp.get_pixel(0, 0), 2.0) def test_sci_viz_stamps(self): - sci_stamps = StampCreator.get_stamps(self.stack, self.trj, 2) + sci_stamps = get_stamps(self.stack, self.trj, 2) self.assertEqual(len(sci_stamps), self.img_count) times = self.stack.build_zeroed_times() @@ -231,7 +236,7 @@ def test_sci_viz_stamps(self): def test_stacked_sci(self): # Compute the stacked science from a single Trajectory. - sci = StampCreator.get_summed_stamp(self.stack, self.trj, 2, []) + sci = get_summed_stamp(self.stack, self.trj, 2, []) self.assertEqual(sci.width, 5) self.assertEqual(sci.height, 5) @@ -258,11 +263,11 @@ def test_median_stamps_trj(self): goodIdx[1][5] = 0 goodIdx[1][9] = 0 - medianStamps0 = StampCreator.get_median_stamp(self.stack, self.trj, 2, goodIdx[0]) + medianStamps0 = get_median_stamp(self.stack, self.trj, 2, goodIdx[0]) self.assertEqual(medianStamps0.width, 5) self.assertEqual(medianStamps0.height, 5) - medianStamps1 = StampCreator.get_median_stamp(self.stack, self.trj, 2, goodIdx[1]) + medianStamps1 = get_median_stamp(self.stack, self.trj, 2, goodIdx[1]) self.assertEqual(medianStamps1.width, 5) self.assertEqual(medianStamps1.height, 5) @@ -291,7 +296,7 @@ def test_median_stamps_no_data(self): trj = Trajectory(x=self.masked_x, y=self.masked_y, vx=0.0, vy=0.0) # Compute the stacked science from a single Trajectory. - medianStamp = StampCreator.get_median_stamp(self.stack, trj, 2, self.all_valid) + medianStamp = get_median_stamp(self.stack, trj, 2, self.all_valid) self.assertEqual(medianStamp.width, 5) self.assertEqual(medianStamp.height, 5) @@ -313,11 +318,11 @@ def test_mean_stamps_trj(self): goodIdx[1][5] = 0 goodIdx[1][9] = 0 - meanStamp0 = StampCreator.get_mean_stamp(self.stack, self.trj, 2, goodIdx[0]) + meanStamp0 = get_mean_stamp(self.stack, self.trj, 2, goodIdx[0]) self.assertEqual(meanStamp0.width, 5) self.assertEqual(meanStamp0.height, 5) - meanStamp1 = StampCreator.get_mean_stamp(self.stack, self.trj, 2, goodIdx[1]) + meanStamp1 = get_mean_stamp(self.stack, self.trj, 2, goodIdx[1]) self.assertEqual(meanStamp1.width, 5) self.assertEqual(meanStamp1.height, 5) @@ -350,7 +355,7 @@ def test_mean_stamps_no_data(self): trj = Trajectory(x=self.masked_x, y=self.masked_y, vx=0.0, vy=0.0) # Compute the stacked science from a single Trajectory - meanStamp = StampCreator.get_mean_stamp(self.stack, trj, 2, self.all_valid) + meanStamp = get_mean_stamp(self.stack, trj, 2, self.all_valid) self.assertEqual(meanStamp.width, 5) self.assertEqual(meanStamp.height, 5) @@ -373,11 +378,11 @@ def test_filter_stamp(self): # Test a stamp with nothing in it. stamp = RawImage(stamp_width, stamp_width) stamp.set_all(1.0) - self.assertTrue(StampCreator.filter_stamp(stamp, self.params)) + self.assertTrue(filter_stamp(stamp, self.params)) # Test a stamp with a bright spot in the center. stamp.set_pixel(5, 5, 100.0) - self.assertFalse(StampCreator.filter_stamp(stamp, self.params)) + self.assertFalse(filter_stamp(stamp, self.params)) # A little noise around the pixel does not hurt as long as the shape is # roughly Gaussian. @@ -385,11 +390,11 @@ def test_filter_stamp(self): stamp.set_pixel(5, 4, 10.0) stamp.set_pixel(6, 5, 10.0) stamp.set_pixel(5, 6, 20.0) - self.assertFalse(StampCreator.filter_stamp(stamp, self.params)) + self.assertFalse(filter_stamp(stamp, self.params)) # A bright peak far from the center is bad. stamp.set_pixel(1, 1, 500.0) - self.assertTrue(StampCreator.filter_stamp(stamp, self.params)) + self.assertTrue(filter_stamp(stamp, self.params)) stamp.set_pixel(1, 1, 1.0) # A non-Gaussian bright spot is also bad. Blur to the -x direction. @@ -402,17 +407,17 @@ def test_filter_stamp(self): stamp.set_pixel(4, 6, 55.0) stamp.set_pixel(3, 6, 55.0) stamp.set_pixel(2, 6, 65.0) - self.assertTrue(StampCreator.filter_stamp(stamp, self.params)) + self.assertTrue(filter_stamp(stamp, self.params)) # A very dim peak at the center is invalid. stamp.set_all(1.0) stamp.set_pixel(5, 5, 1.0001) - self.assertTrue(StampCreator.filter_stamp(stamp, self.params)) + self.assertTrue(filter_stamp(stamp, self.params)) # A slightly offset peak of sufficient brightness is okay. stamp.set_pixel(5, 5, 15.0) stamp.set_pixel(4, 5, 20.0) - self.assertFalse(StampCreator.filter_stamp(stamp, self.params)) + self.assertFalse(filter_stamp(stamp, self.params)) def test_coadd_cpu_simple(self): # Create an image set with three images. @@ -455,7 +460,7 @@ def test_coadd_cpu_simple(self): # Test summed. params.stamp_type = StampType.STAMP_SUM - stamps = StampCreator.get_coadded_stamps(stack, [trj], [all_valid], params, False) + stamps = get_coadded_stamps(stack, [trj], [all_valid], params, False) self.assertAlmostEqual(stamps[0].get_pixel(0, 0), 3.0) self.assertAlmostEqual(stamps[0].get_pixel(1, 0), 3.0) self.assertAlmostEqual(stamps[0].get_pixel(2, 0), 3.0) @@ -468,7 +473,7 @@ def test_coadd_cpu_simple(self): # Test mean. params.stamp_type = StampType.STAMP_MEAN - stamps = StampCreator.get_coadded_stamps(stack, [trj], [all_valid], params, False) + stamps = get_coadded_stamps(stack, [trj], [all_valid], params, False) self.assertAlmostEqual(stamps[0].get_pixel(0, 0), 1.0) self.assertAlmostEqual(stamps[0].get_pixel(1, 0), 1.0) self.assertAlmostEqual(stamps[0].get_pixel(2, 0), 1.0) @@ -481,7 +486,7 @@ def test_coadd_cpu_simple(self): # Test median. params.stamp_type = StampType.STAMP_MEDIAN - stamps = StampCreator.get_coadded_stamps(stack, [trj], [all_valid], params, False) + stamps = get_coadded_stamps(stack, [trj], [all_valid], params, False) self.assertAlmostEqual(stamps[0].get_pixel(0, 0), 1.0) self.assertAlmostEqual(stamps[0].get_pixel(1, 0), 1.0) self.assertAlmostEqual(stamps[0].get_pixel(2, 0), 1.0) @@ -528,21 +533,21 @@ def test_coadd_gpu_simple(self): # Test summed. params.stamp_type = StampType.STAMP_SUM - stamps = StampCreator.get_coadded_stamps(stack, [trj], [all_valid], params, True) + stamps = get_coadded_stamps(stack, [trj], [all_valid], params, True) self.assertAlmostEqual(stamps[0].get_pixel(0, 1), 3.0) self.assertAlmostEqual(stamps[0].get_pixel(1, 1), 5.0) self.assertAlmostEqual(stamps[0].get_pixel(2, 1), 6.0) # Test mean. params.stamp_type = StampType.STAMP_MEAN - stamps = StampCreator.get_coadded_stamps(stack, [trj], [all_valid], params, True) + stamps = get_coadded_stamps(stack, [trj], [all_valid], params, True) self.assertAlmostEqual(stamps[0].get_pixel(0, 1), 3.0) self.assertAlmostEqual(stamps[0].get_pixel(1, 1), 2.5) self.assertAlmostEqual(stamps[0].get_pixel(2, 1), 2.0) # Test median. params.stamp_type = StampType.STAMP_MEDIAN - stamps = StampCreator.get_coadded_stamps(stack, [trj], [all_valid], params, True) + stamps = get_coadded_stamps(stack, [trj], [all_valid], params, True) self.assertAlmostEqual(stamps[0].get_pixel(0, 1), 3.0) self.assertAlmostEqual(stamps[0].get_pixel(1, 1), 2.5) self.assertAlmostEqual(stamps[0].get_pixel(2, 1), 2.0) @@ -554,19 +559,19 @@ def test_coadd_cpu(self): # Compute the stacked science (summed and mean) from a single Trajectory. params.stamp_type = StampType.STAMP_SUM - summedStamps = StampCreator.get_coadded_stamps( + summedStamps = get_coadded_stamps( self.stack, [self.trj], [self.all_valid], params, False ) self.assertEqual(summedStamps[0].width, 2 * params.radius + 1) self.assertEqual(summedStamps[0].height, 2 * params.radius + 1) params.stamp_type = StampType.STAMP_MEAN - meanStamps = StampCreator.get_coadded_stamps(self.stack, [self.trj], [self.all_valid], params, False) + meanStamps = get_coadded_stamps(self.stack, [self.trj], [self.all_valid], params, False) self.assertEqual(meanStamps[0].width, 2 * params.radius + 1) self.assertEqual(meanStamps[0].height, 2 * params.radius + 1) params.stamp_type = StampType.STAMP_MEDIAN - medianStamps = StampCreator.get_coadded_stamps( + medianStamps = get_coadded_stamps( self.stack, [self.trj], [self.all_valid], params, False ) self.assertEqual(medianStamps[0].width, 2 * params.radius + 1) @@ -609,17 +614,17 @@ def test_coadd_gpu(self): # Compute the stacked science (summed and mean) from a single Trajectory. params.stamp_type = StampType.STAMP_SUM - summedStamps = StampCreator.get_coadded_stamps(self.stack, [self.trj], [self.all_valid], params, True) + summedStamps = get_coadded_stamps(self.stack, [self.trj], [self.all_valid], params, True) self.assertEqual(summedStamps[0].width, 2 * params.radius + 1) self.assertEqual(summedStamps[0].height, 2 * params.radius + 1) params.stamp_type = StampType.STAMP_MEAN - meanStamps = StampCreator.get_coadded_stamps(self.stack, [self.trj], [self.all_valid], params, True) + meanStamps = get_coadded_stamps(self.stack, [self.trj], [self.all_valid], params, True) self.assertEqual(meanStamps[0].width, 2 * params.radius + 1) self.assertEqual(meanStamps[0].height, 2 * params.radius + 1) params.stamp_type = StampType.STAMP_MEDIAN - medianStamps = StampCreator.get_coadded_stamps(self.stack, [self.trj], [self.all_valid], params, True) + medianStamps = get_coadded_stamps(self.stack, [self.trj], [self.all_valid], params, True) self.assertEqual(medianStamps[0].width, 2 * params.radius + 1) self.assertEqual(medianStamps[0].height, 2 * params.radius + 1) @@ -667,7 +672,7 @@ def test_coadd_cpu_use_inds(self): inds[1][11] = False # Compute the stacked science (summed and mean) from a single Trajectory. - meanStamps = StampCreator.get_coadded_stamps(self.stack, [self.trj, self.trj], inds, params, False) + meanStamps = get_coadded_stamps(self.stack, [self.trj, self.trj], inds, params, False) # Compute the true summed and mean pixels for all of the pixels in the stamp. times = self.stack.build_zeroed_times() @@ -716,7 +721,7 @@ def test_coadd_gpu_use_inds(self): inds[1][11] = False # Compute the stacked science (summed and mean) from a single Trajectory. - meanStamps = StampCreator.get_coadded_stamps(self.stack, [self.trj, self.trj], inds, params, True) + meanStamps = get_coadded_stamps(self.stack, [self.trj, self.trj], inds, params, True) # Compute the true summed and mean pixels for all of the pixels in the stamp. times = self.stack.build_zeroed_times() @@ -771,7 +776,7 @@ def test_coadd_filter_cpu(self): # Compute the stacked science from a single Trajectory. all_valid_vect = [(self.all_valid) for i in range(4)] - meanStamps = StampCreator.get_coadded_stamps( + meanStamps = get_coadded_stamps( self.stack, [self.trj, trj2, trj3, trj4], all_valid_vect, self.params, True ) @@ -810,7 +815,7 @@ def test_coadd_filter_gpu(self): # Compute the stacked science from a single Trajectory. all_valid_vect = [(self.all_valid) for i in range(4)] - meanStamps = StampCreator.get_coadded_stamps( + meanStamps = get_coadded_stamps( self.stack, [self.trj, trj2, trj3, trj4], all_valid_vect, self.params, True ) diff --git a/tests/test_stamp_parity.py b/tests/test_stamp_parity.py index b4b580643..110738b3a 100644 --- a/tests/test_stamp_parity.py +++ b/tests/test_stamp_parity.py @@ -64,13 +64,13 @@ def test_coadd_gpu_parity(self): # Check the summed stamps. Note summed stamp does not use goodIdx. params.stamp_type = StampType.STAMP_SUM stamps_old = [ - StampCreator.get_summed_stamp(self.search.get_imagestack(), self.trj, radius, all_valid), - StampCreator.get_summed_stamp(self.search.get_imagestack(), self.trj, radius, all_valid), + get_summed_stamp(self.search.get_imagestack(), self.trj, radius, all_valid), + get_summed_stamp(self.search.get_imagestack(), self.trj, radius, all_valid), ] - stamps_gpu = StampCreator.get_coadded_stamps( + stamps_gpu = get_coadded_stamps( self.search.get_imagestack(), results, [all_valid, all_valid], params, True ) - stamps_cpu = StampCreator.get_coadded_stamps( + stamps_cpu = get_coadded_stamps( self.search.get_imagestack(), results, [all_valid, all_valid], params, False ) for r in range(2): @@ -80,13 +80,13 @@ def test_coadd_gpu_parity(self): # Check the mean stamps. params.stamp_type = StampType.STAMP_MEAN stamps_old = [ - StampCreator.get_mean_stamp(self.search.get_imagestack(), self.trj, radius, goodIdx[0]), - StampCreator.get_mean_stamp(self.search.get_imagestack(), self.trj, radius, goodIdx[1]), + get_mean_stamp(self.search.get_imagestack(), self.trj, radius, goodIdx[0]), + get_mean_stamp(self.search.get_imagestack(), self.trj, radius, goodIdx[1]), ] - stamps_gpu = StampCreator.get_coadded_stamps( + stamps_gpu = get_coadded_stamps( self.search.get_imagestack(), results, goodIdx, params, True ) - stamps_cpu = StampCreator.get_coadded_stamps( + stamps_cpu = get_coadded_stamps( self.search.get_imagestack(), results, goodIdx, params, False ) for r in range(2): @@ -96,13 +96,13 @@ def test_coadd_gpu_parity(self): # Check the median stamps. params.stamp_type = StampType.STAMP_MEDIAN stamps_old = [ - StampCreator.get_median_stamp(self.search.get_imagestack(), self.trj, radius, goodIdx[0]), - StampCreator.get_median_stamp(self.search.get_imagestack(), self.trj, radius, goodIdx[1]), + get_median_stamp(self.search.get_imagestack(), self.trj, radius, goodIdx[0]), + get_median_stamp(self.search.get_imagestack(), self.trj, radius, goodIdx[1]), ] - stamps_gpu = StampCreator.get_coadded_stamps( + stamps_gpu = get_coadded_stamps( self.search.get_imagestack(), results, goodIdx, params, True ) - stamps_cpu = StampCreator.get_coadded_stamps( + stamps_cpu = get_coadded_stamps( self.search.get_imagestack(), results, goodIdx, params, False ) for r in range(2): From c3da8cb9a4f66acc60c6bf97c84d0c7bd394b3b4 Mon Sep 17 00:00:00 2001 From: Max West Date: Tue, 17 Dec 2024 15:02:38 -0800 Subject: [PATCH 2/3] black formatting --- tests/test_stamp_creator.py | 8 ++------ tests/test_stamp_parity.py | 16 ++++------------ 2 files changed, 6 insertions(+), 18 deletions(-) diff --git a/tests/test_stamp_creator.py b/tests/test_stamp_creator.py index a61a1171d..e1d0dc08c 100644 --- a/tests/test_stamp_creator.py +++ b/tests/test_stamp_creator.py @@ -559,9 +559,7 @@ def test_coadd_cpu(self): # Compute the stacked science (summed and mean) from a single Trajectory. params.stamp_type = StampType.STAMP_SUM - summedStamps = get_coadded_stamps( - self.stack, [self.trj], [self.all_valid], params, False - ) + summedStamps = get_coadded_stamps(self.stack, [self.trj], [self.all_valid], params, False) self.assertEqual(summedStamps[0].width, 2 * params.radius + 1) self.assertEqual(summedStamps[0].height, 2 * params.radius + 1) @@ -571,9 +569,7 @@ def test_coadd_cpu(self): self.assertEqual(meanStamps[0].height, 2 * params.radius + 1) params.stamp_type = StampType.STAMP_MEDIAN - medianStamps = get_coadded_stamps( - self.stack, [self.trj], [self.all_valid], params, False - ) + medianStamps = get_coadded_stamps(self.stack, [self.trj], [self.all_valid], params, False) self.assertEqual(medianStamps[0].width, 2 * params.radius + 1) self.assertEqual(medianStamps[0].height, 2 * params.radius + 1) diff --git a/tests/test_stamp_parity.py b/tests/test_stamp_parity.py index 110738b3a..dc584ed82 100644 --- a/tests/test_stamp_parity.py +++ b/tests/test_stamp_parity.py @@ -83,12 +83,8 @@ def test_coadd_gpu_parity(self): get_mean_stamp(self.search.get_imagestack(), self.trj, radius, goodIdx[0]), get_mean_stamp(self.search.get_imagestack(), self.trj, radius, goodIdx[1]), ] - stamps_gpu = get_coadded_stamps( - self.search.get_imagestack(), results, goodIdx, params, True - ) - stamps_cpu = get_coadded_stamps( - self.search.get_imagestack(), results, goodIdx, params, False - ) + stamps_gpu = get_coadded_stamps(self.search.get_imagestack(), results, goodIdx, params, True) + stamps_cpu = get_coadded_stamps(self.search.get_imagestack(), results, goodIdx, params, False) for r in range(2): self.assertTrue(np.allclose(stamps_old[r].image, stamps_gpu[r].image, atol=1e-5)) self.assertTrue(np.allclose(stamps_old[r].image, stamps_cpu[r].image, atol=1e-5)) @@ -99,12 +95,8 @@ def test_coadd_gpu_parity(self): get_median_stamp(self.search.get_imagestack(), self.trj, radius, goodIdx[0]), get_median_stamp(self.search.get_imagestack(), self.trj, radius, goodIdx[1]), ] - stamps_gpu = get_coadded_stamps( - self.search.get_imagestack(), results, goodIdx, params, True - ) - stamps_cpu = get_coadded_stamps( - self.search.get_imagestack(), results, goodIdx, params, False - ) + stamps_gpu = get_coadded_stamps(self.search.get_imagestack(), results, goodIdx, params, True) + stamps_cpu = get_coadded_stamps(self.search.get_imagestack(), results, goodIdx, params, False) for r in range(2): self.assertTrue(np.allclose(stamps_old[r].image, stamps_gpu[r].image, 1e-5)) self.assertTrue(np.allclose(stamps_old[r].image, stamps_cpu[r].image, 1e-5)) From 824a4423eaa1b5bd27b9d098d0712c74046cf9c9 Mon Sep 17 00:00:00 2001 From: Max West Date: Wed, 18 Dec 2024 12:27:50 -0800 Subject: [PATCH 3/3] address comments from #761 --- benchmarks/bench_filter_stamps.py | 2 +- src/kbmod/search/stamp_creator.cpp | 28 ++++++++++++++-------------- 2 files changed, 15 insertions(+), 15 deletions(-) diff --git a/benchmarks/bench_filter_stamps.py b/benchmarks/bench_filter_stamps.py index 43dae1b03..ca66729d3 100644 --- a/benchmarks/bench_filter_stamps.py +++ b/benchmarks/bench_filter_stamps.py @@ -56,7 +56,7 @@ def run_search_benchmark(params): # im_stack = ImageStack([]) # Do the timing runs. - tmr = timeit.Timer(stmt="sc.filter_stamp(stamp, params)", globals=locals()) + tmr = timeit.Timer(stmt="filter_stamp(stamp, params)", globals=locals()) res_time = np.mean(tmr.repeat(repeat=10, number=20)) return res_time diff --git a/src/kbmod/search/stamp_creator.cpp b/src/kbmod/search/stamp_creator.cpp index 123264784..affa8854c 100644 --- a/src/kbmod/search/stamp_creator.cpp +++ b/src/kbmod/search/stamp_creator.cpp @@ -9,7 +9,7 @@ void deviceGetCoadds(const uint64_t num_images, const uint64_t width, const uint #endif std::vector create_stamps(ImageStack& stack, const Trajectory& trj, int radius, - bool keep_no_data, const std::vector& use_index) { + bool keep_no_data, const std::vector& use_index) { if (use_index.size() > 0) assert_sizes_equal(use_index.size(), stack.img_count(), "create_stamps() use_index"); bool use_all_stamps = (use_index.size() == 0); @@ -41,27 +41,27 @@ std::vector get_stamps(ImageStack& stack, const Trajectory& t, int rad // For creating coadded stamps, we do not interpolate the pixel values and keep // invalid pixels tagged (so we can filter it out of mean/median). RawImage get_median_stamp(ImageStack& stack, const Trajectory& trj, int radius, - const std::vector& use_index) { + const std::vector& use_index) { return create_median_image(create_stamps(stack, trj, radius, true /*=keep_no_data*/, use_index)); } // For creating coadded stamps, we do not interpolate the pixel values and keep // invalid pixels tagged (so we can filter it out of mean/median). RawImage get_mean_stamp(ImageStack& stack, const Trajectory& trj, int radius, - const std::vector& use_index) { + const std::vector& use_index) { return create_mean_image(create_stamps(stack, trj, radius, true /*=keep_no_data*/, use_index)); } // For creating summed stamps, we do not interpolate the pixel values and replace // invalid pixels with zero (which is the same as filtering it out for the sum). RawImage get_summed_stamp(ImageStack& stack, const Trajectory& trj, int radius, - const std::vector& use_index) { + const std::vector& use_index) { return create_summed_image(create_stamps(stack, trj, radius, false /*=keep_no_data*/, use_index)); } std::vector get_coadded_stamps(ImageStack& stack, std::vector& t_array, - std::vector>& use_index_vect, - const StampParameters& params, bool use_gpu) { + std::vector>& use_index_vect, + const StampParameters& params, bool use_gpu) { logging::Logger* rs_logger = logging::getLogger("kbmod.search.stamp_creator"); rs_logger->info("Generating co_added stamps on " + std::to_string(t_array.size()) + " trajectories."); DebugTimer timer = DebugTimer("coadd generating", rs_logger); @@ -79,9 +79,9 @@ std::vector get_coadded_stamps(ImageStack& stack, std::vector get_coadded_stamps_cpu(ImageStack& stack, - std::vector& t_array, - std::vector>& use_index_vect, - const StampParameters& params) { + std::vector& t_array, + std::vector>& use_index_vect, + const StampParameters& params) { const uint64_t num_trajectories = t_array.size(); std::vector results(num_trajectories); @@ -158,9 +158,9 @@ bool filter_stamp(const RawImage& img, const StampParameters& params) { } std::vector get_coadded_stamps_gpu(ImageStack& stack, - std::vector& t_array, - std::vector>& use_index_vect, - const StampParameters& params) { + std::vector& t_array, + std::vector>& use_index_vect, + const StampParameters& params) { logging::Logger* rs_logger = logging::getLogger("kbmod.search.stamp_creator"); // Right now only limited stamp sizes are allowed. @@ -272,7 +272,7 @@ std::vector get_coadded_stamps_gpu(ImageStack& stack, } std::vector create_variance_stamps(ImageStack& stack, const Trajectory& trj, - int radius, const std::vector& use_index) { + int radius, const std::vector& use_index) { if (use_index.size() > 0) assert_sizes_equal(use_index.size(), stack.img_count(), "create_stamps() use_index"); bool use_all_stamps = (use_index.size() == 0); @@ -294,7 +294,7 @@ std::vector create_variance_stamps(ImageStack& stack, const Trajectory } RawImage get_variance_weighted_stamp(ImageStack& stack, const Trajectory& trj, int radius, - const std::vector& use_index) { + const std::vector& use_index) { if (radius < 0) throw std::runtime_error("Invalid stamp radius. Must be >= 0."); unsigned int num_images = stack.img_count(); if (num_images == 0) throw std::runtime_error("Unable to create mean image given 0 images.");