diff --git a/src/kbmod/analysis_utils.py b/src/kbmod/analysis_utils.py index 79582482..86c1cc34 100644 --- a/src/kbmod/analysis_utils.py +++ b/src/kbmod/analysis_utils.py @@ -470,8 +470,11 @@ def apply_stamp_filter( all_true = [True] * num_times bool_slice = [all_true for _ in inds_to_use] - # Create and filter the results. - stamps_slice = search.gpu_coadded_stamps(trj_slice, bool_slice, params) + # Create and filter the results, using the GPU if there is one and enough + # trajectories to make it worthwhile. + stamps_slice = search.coadded_stamps( + trj_slice, bool_slice, params, kb.HAS_GPU and len(trj_slice) > 100 + ) for ind, stamp in enumerate(stamps_slice): if stamp.get_width() > 1: result_list.results[ind + start_idx].stamp = np.array(stamp) diff --git a/src/kbmod/search/KBMOSearch.cpp b/src/kbmod/search/KBMOSearch.cpp index 749b694e..8def1b79 100644 --- a/src/kbmod/search/KBMOSearch.cpp +++ b/src/kbmod/search/KBMOSearch.cpp @@ -373,6 +373,56 @@ bool KBMOSearch::filterStamp(const RawImage& img, const stampParameters& params) return false; } +std::vector KBMOSearch::coaddedScienceStamps(std::vector& t_array, + std::vector >& use_index_vect, + const stampParameters& params, + bool use_gpu) { + if (use_gpu) { + #ifdef HAVE_CUDA + return coaddedScienceStampsGPU(t_array, use_index_vect, params); + #else + print("WARNING: GPU is not enabled. Performing co-adds on the CPU."); + #endif + } + return coaddedScienceStampsCPU(t_array, use_index_vect, params); +} + +std::vector KBMOSearch::coaddedScienceStampsCPU(std::vector& t_array, + std::vector >& use_index_vect, + const stampParameters& params) { + const int num_trajectories = t_array.size(); + std::vector results(num_trajectories); + std::vector empty_pixels(1, NO_DATA); + + for (int i = 0; i < num_trajectories; ++i) { + std::vector stamps = scienceStamps(t_array[i], params.radius, false, true, use_index_vect[i]); + + RawImage coadd(1, 1); + switch (params.stamp_type) { + case STAMP_MEDIAN: + coadd = createMedianImage(stamps); + break; + case STAMP_MEAN: + coadd = createMeanImage(stamps); + break; + case STAMP_SUM: + coadd = createSummedImage(stamps); + break; + default: + throw std::runtime_error("Invalid stamp coadd type."); + } + + // Do the filtering if needed. + if (params.do_filtering && filterStamp(coadd, params)) { + results[i] = RawImage(1, 1, empty_pixels); + } else { + results[i] = coadd; + } + } + + return results; +} + std::vector KBMOSearch::coaddedScienceStampsGPU(std::vector& t_array, std::vector >& use_index_vect, const stampParameters& params) { diff --git a/src/kbmod/search/KBMOSearch.h b/src/kbmod/search/KBMOSearch.h index 30118c27..6668b732 100644 --- a/src/kbmod/search/KBMOSearch.h +++ b/src/kbmod/search/KBMOSearch.h @@ -69,13 +69,14 @@ class KBMOSearch { RawImage meanScienceStamp(const trajectory& trj, int radius, const std::vector& use_index); RawImage summedScienceStamp(const trajectory& trj, int radius, const std::vector& use_index); - // Compute a mean or summed stamp for each trajectory on the GPU. This is slower than the - // above for small numbers of trajectories (< 500), but performs relatively better as the - // number of trajectories increases. If filtering is applied then will use a 1x1 image with - // NO_DATA to represent filtered images. - std::vector coaddedScienceStampsGPU(std::vector& t_array, - std::vector >& use_index_vect, - const stampParameters& params); + // 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. + std::vector coaddedScienceStamps(std::vector& t_array, + std::vector >& use_index_vect, + const stampParameters& params, + bool use_cpu); // Function to do the actual stamp filtering. bool filterStamp(const RawImage& img, const stampParameters& params); @@ -120,6 +121,13 @@ class KBMOSearch { void createSearchList(int angleSteps, int veloctiySteps, float minAngle, float maxAngle, float minVelocity, float maxVelocity); + std::vector coaddedScienceStampsGPU(std::vector& t_array, + std::vector >& use_index_vect, + const stampParameters& params); + + std::vector coaddedScienceStampsCPU(std::vector& t_array, + std::vector >& use_index_vect, + const stampParameters& params); // Helper functions for timing operations of the search. void startTimer(const std::string& message); void endTimer(); diff --git a/src/kbmod/search/bindings.cpp b/src/kbmod/search/bindings.cpp index bf620a4a..b6508ce3 100644 --- a/src/kbmod/search/bindings.cpp +++ b/src/kbmod/search/bindings.cpp @@ -199,10 +199,10 @@ PYBIND11_MODULE(search, m) { .def("median_sci_stamp", &ks::medianScienceStamp) .def("mean_sci_stamp", &ks::meanScienceStamp) .def("summed_sci_stamp", &ks::summedScienceStamp) - .def("gpu_coadded_stamps", + .def("coadded_stamps", (std::vector(ks::*)(std::vector &, std::vector> &, - const search::stampParameters &)) & - ks::coaddedScienceStampsGPU) + const search::stampParameters &, bool)) & + ks::coaddedScienceStamps) // For testing .def("filter_stamp", &ks::filterStamp) .def("get_traj_pos", &ks::getTrajPos) diff --git a/tests/test_search.py b/tests/test_search.py index b088c1a5..2ad095db 100644 --- a/tests/test_search.py +++ b/tests/test_search.py @@ -456,6 +456,68 @@ def test_filter_stamp(self): stamp.set_pixel(4, 5, 20.0) self.assertFalse(self.search.filter_stamp(stamp, self.params)) + def test_coadd_cpu_simple(self): + # Create an image set with three images. + imlist = [] + for i in range(3): + time = i + im = layered_image(str(i), 3, 3, 0.1, 0.01, i, self.p, i) + + # Overwrite the middle row to be i + 1. + sci = im.get_science() + for x in range(3): + sci.set_pixel(x, 1, i + 1) + im.set_science(sci) + + # Mask out the row's first pixel twice and second pixel once. + mask = im.get_mask() + if i == 0: + mask.set_pixel(0, 1, 1) + mask.set_pixel(1, 1, 1) + if i == 1: + mask.set_pixel(0, 1, 1) + im.set_mask(mask) + im.apply_mask_flags(1, []) + + imlist.append(im) + stack = image_stack(imlist) + search = stack_search(stack) + all_valid = [True, True, True] # convenience array + + # One trajectory right in the image's middle. + trj = trajectory() + trj.x = 1 + trj.y = 1 + trj.x_v = 0 + trj.y_v = 0 + + # Basic Stamp parameters. + params = stamp_parameters() + params.radius = 1 + params.do_filtering = False + + # Test summed. + params.stamp_type = StampType.STAMP_SUM + stamps = search.coadded_stamps([trj], [all_valid], params, False) + 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 = search.coadded_stamps([trj], [all_valid], params, False) + 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 = search.coadded_stamps([trj], [all_valid], params, False) + 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) + + @unittest.skipIf(not HAS_GPU, "Skipping test (no GPU detected)") def test_coadd_gpu_simple(self): # Create an image set with three images. imlist = [] @@ -498,25 +560,76 @@ def test_coadd_gpu_simple(self): # Test summed. params.stamp_type = StampType.STAMP_SUM - stamps = search.gpu_coadded_stamps([trj], [all_valid], params) + stamps = search.coadded_stamps([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 = search.gpu_coadded_stamps([trj], [all_valid], params) + stamps = search.coadded_stamps([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 = search.gpu_coadded_stamps([trj], [all_valid], params) + stamps = search.coadded_stamps([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) + def test_coadd_cpu(self): + params = stamp_parameters() + params.radius = 3 + params.do_filtering = False + + # Compute the stacked science (summed and mean) from a single trajectory. + params.stamp_type = StampType.STAMP_SUM + summedStamps = self.search.coadded_stamps([self.trj], [self.all_valid], params, False) + self.assertEqual(summedStamps[0].get_width(), 2 * params.radius + 1) + self.assertEqual(summedStamps[0].get_height(), 2 * params.radius + 1) + + params.stamp_type = StampType.STAMP_MEAN + meanStamps = self.search.coadded_stamps([self.trj], [self.all_valid], params, False) + self.assertEqual(meanStamps[0].get_width(), 2 * params.radius + 1) + self.assertEqual(meanStamps[0].get_height(), 2 * params.radius + 1) + + params.stamp_type = StampType.STAMP_MEDIAN + medianStamps = self.search.coadded_stamps([self.trj], [self.all_valid], params, False) + self.assertEqual(medianStamps[0].get_width(), 2 * params.radius + 1) + self.assertEqual(medianStamps[0].get_height(), 2 * params.radius + 1) + + # Compute the true summed and mean pixels for all of the pixels in the stamp. + times = self.stack.get_times() + for stamp_x in range(2 * params.radius + 1): + for stamp_y in range(2 * params.radius + 1): + x_offset = stamp_x - params.radius + y_offset = stamp_y - params.radius + + pix_sum = 0.0 + pix_count = 0.0 + pix_vals = [] + for i in range(self.imCount): + t = times[i] + x = int(self.trj.x + self.trj.x_v * t) + x_offset + y = int(self.trj.y + self.trj.y_v * t) + y_offset + pixVal = self.imlist[i].get_science().get_pixel(x, y) + if pixVal != KB_NO_DATA: + pix_sum += pixVal + pix_count += 1.0 + pix_vals.append(pixVal) + + # Check that we get the correct answers. + self.assertAlmostEqual(pix_sum, summedStamps[0].get_pixel(stamp_x, stamp_y), delta=1e-3) + self.assertAlmostEqual( + pix_sum / pix_count, meanStamps[0].get_pixel(stamp_x, stamp_y), delta=1e-3 + ) + self.assertAlmostEqual( + np.median(pix_vals), medianStamps[0].get_pixel(stamp_x, stamp_y), delta=1e-3 + ) + + @unittest.skipIf(not HAS_GPU, "Skipping test (no GPU detected)") def test_coadd_gpu(self): params = stamp_parameters() params.radius = 3 @@ -524,17 +637,17 @@ def test_coadd_gpu(self): # Compute the stacked science (summed and mean) from a single trajectory. params.stamp_type = StampType.STAMP_SUM - summedStamps = self.search.gpu_coadded_stamps([self.trj], [self.all_valid], params) + summedStamps = self.search.coadded_stamps([self.trj], [self.all_valid], params, True) self.assertEqual(summedStamps[0].get_width(), 2 * params.radius + 1) self.assertEqual(summedStamps[0].get_height(), 2 * params.radius + 1) params.stamp_type = StampType.STAMP_MEAN - meanStamps = self.search.gpu_coadded_stamps([self.trj], [self.all_valid], params) + meanStamps = self.search.coadded_stamps([self.trj], [self.all_valid], params, True) self.assertEqual(meanStamps[0].get_width(), 2 * params.radius + 1) self.assertEqual(meanStamps[0].get_height(), 2 * params.radius + 1) params.stamp_type = StampType.STAMP_MEDIAN - medianStamps = self.search.gpu_coadded_stamps([self.trj], [self.all_valid], params) + medianStamps = self.search.coadded_stamps([self.trj], [self.all_valid], params, True) self.assertEqual(medianStamps[0].get_width(), 2 * params.radius + 1) self.assertEqual(medianStamps[0].get_height(), 2 * params.radius + 1) @@ -567,6 +680,55 @@ def test_coadd_gpu(self): np.median(pix_vals), medianStamps[0].get_pixel(stamp_x, stamp_y), delta=1e-3 ) + def test_coadd_cpu_use_inds(self): + params = stamp_parameters() + params.radius = 1 + params.do_filtering = False + params.stamp_type = StampType.STAMP_MEAN + + # Mark a few of the observations as "do not use" + inds = [[True] * self.imCount, [True] * self.imCount] + inds[0][5] = False + inds[1][3] = False + inds[1][6] = False + inds[1][7] = False + inds[1][11] = False + + # Compute the stacked science (summed and mean) from a single trajectory. + meanStamps = self.search.coadded_stamps([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.get_times() + for stamp_x in range(2 * params.radius + 1): + for stamp_y in range(2 * params.radius + 1): + x_offset = stamp_x - params.radius + y_offset = stamp_y - params.radius + + sum_0 = 0.0 + sum_1 = 0.0 + count_0 = 0.0 + count_1 = 0.0 + for i in range(self.imCount): + t = times[i] + x = int(self.trj.x + self.trj.x_v * t) + x_offset + y = int(self.trj.y + self.trj.y_v * t) + y_offset + pixVal = self.imlist[i].get_science().get_pixel(x, y) + + if pixVal != KB_NO_DATA and inds[0][i] > 0: + sum_0 += pixVal + count_0 += 1.0 + + if pixVal != KB_NO_DATA and inds[1][i] > 0: + sum_1 += pixVal + count_1 += 1.0 + + # Check that we get the correct answers. + self.assertAlmostEqual(count_0, 19.0) + self.assertAlmostEqual(count_1, 16.0) + self.assertAlmostEqual(sum_0 / count_0, meanStamps[0].get_pixel(stamp_x, stamp_y), delta=1e-3) + self.assertAlmostEqual(sum_1 / count_1, meanStamps[1].get_pixel(stamp_x, stamp_y), delta=1e-3) + + @unittest.skipIf(not HAS_GPU, "Skipping test (no GPU detected)") def test_coadd_gpu_use_inds(self): params = stamp_parameters() params.radius = 1 @@ -582,7 +744,7 @@ def test_coadd_gpu_use_inds(self): inds[1][11] = False # Compute the stacked science (summed and mean) from a single trajectory. - meanStamps = self.search.gpu_coadded_stamps([self.trj, self.trj], inds, params) + meanStamps = self.search.coadded_stamps([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.get_times() @@ -615,6 +777,47 @@ def test_coadd_gpu_use_inds(self): self.assertAlmostEqual(sum_0 / count_0, meanStamps[0].get_pixel(stamp_x, stamp_y), delta=1e-3) self.assertAlmostEqual(sum_1 / count_1, meanStamps[1].get_pixel(stamp_x, stamp_y), delta=1e-3) + def test_coadd_filter_cpu(self): + # Create a second trajectory that isn't any good. + trj2 = trajectory() + trj2.x = 1 + trj2.y = 1 + trj2.x_v = 0 + trj2.y_v = 0 + + # Create a third trajectory that is close to good, but offset. + trj3 = trajectory() + trj3.x = self.trj.x + 2 + trj3.y = self.trj.y + 2 + trj3.x_v = self.trj.x_v + trj3.y_v = self.trj.y_v + + # Create a fourth trajectory that is close enough + trj4 = trajectory() + trj4.x = self.trj.x + 1 + trj4.y = self.trj.y + 1 + trj4.x_v = self.trj.x_v + trj4.y_v = self.trj.y_v + + # Compute the stacked science from a single trajectory. + all_valid_vect = [(self.all_valid) for i in range(4)] + meanStamps = self.search.coadded_stamps( + [self.trj, trj2, trj3, trj4], all_valid_vect, self.params, False + ) + + # The first and last are unfiltered + self.assertEqual(meanStamps[0].get_width(), 2 * self.params.radius + 1) + self.assertEqual(meanStamps[0].get_height(), 2 * self.params.radius + 1) + self.assertEqual(meanStamps[3].get_width(), 2 * self.params.radius + 1) + self.assertEqual(meanStamps[3].get_height(), 2 * self.params.radius + 1) + + # The second and third are filtered. + self.assertEqual(meanStamps[1].get_width(), 1) + self.assertEqual(meanStamps[1].get_height(), 1) + self.assertEqual(meanStamps[2].get_width(), 1) + self.assertEqual(meanStamps[2].get_height(), 1) + + @unittest.skipIf(not HAS_GPU, "Skipping test (no GPU detected)") def test_coadd_filter_gpu(self): # Create a second trajectory that isn't any good. trj2 = trajectory() @@ -639,7 +842,9 @@ 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 = self.search.gpu_coadded_stamps([self.trj, trj2, trj3, trj4], all_valid_vect, self.params) + meanStamps = self.search.coadded_stamps( + [self.trj, trj2, trj3, trj4], all_valid_vect, self.params, True + ) # The first and last are unfiltered self.assertEqual(meanStamps[0].get_width(), 2 * self.params.radius + 1) diff --git a/tests/test_stamp_parity.py b/tests/test_stamp_parity.py index 225bd7dc..c65915d6 100644 --- a/tests/test_stamp_parity.py +++ b/tests/test_stamp_parity.py @@ -1,3 +1,9 @@ +""" +KBMOD provides a series of different wrapper functions for extracting +coadded stamps from trajectories. These tests confirm that the behavior +of the different approaches is consistent. +""" + import unittest import numpy as np @@ -70,6 +76,7 @@ def setUp(self): self.stack = image_stack(self.imlist) self.search = stack_search(self.stack) + @unittest.skipIf(not HAS_GPU, "Skipping test (no GPU detected)") def test_coadd_gpu_parity(self): radius = 2 width = 2 * radius + 1 @@ -90,13 +97,11 @@ def test_coadd_gpu_parity(self): self.search.summed_sci_stamp(self.trj, radius, all_valid), self.search.summed_sci_stamp(self.trj, radius, all_valid), ] - stamps_new = self.search.gpu_coadded_stamps(results, [all_valid, all_valid], params) + stamps_gpu = self.search.coadded_stamps(results, [all_valid, all_valid], params, True) + stamps_cpu = self.search.coadded_stamps(results, [all_valid, all_valid], params, False) for r in range(2): - for x in range(width): - for y in range(width): - self.assertAlmostEqual( - stamps_old[r].get_pixel(x, y), stamps_new[r].get_pixel(x, y), delta=1e-5 - ) + self.assertTrue(stamps_old[r].approx_equal(stamps_gpu[r], 1e-5)) + self.assertTrue(stamps_old[r].approx_equal(stamps_cpu[r], 1e-5)) # Check the mean stamps. params.stamp_type = StampType.STAMP_MEAN @@ -104,13 +109,11 @@ def test_coadd_gpu_parity(self): self.search.mean_sci_stamp(self.trj, radius, goodIdx[0]), self.search.mean_sci_stamp(self.trj, radius, goodIdx[1]), ] - stamps_new = self.search.gpu_coadded_stamps(results, goodIdx, params) + stamps_gpu = self.search.coadded_stamps(results, goodIdx, params, True) + stamps_cpu = self.search.coadded_stamps(results, goodIdx, params, False) for r in range(2): - for x in range(width): - for y in range(width): - self.assertAlmostEqual( - stamps_old[r].get_pixel(x, y), stamps_new[r].get_pixel(x, y), delta=1e-5 - ) + self.assertTrue(stamps_old[r].approx_equal(stamps_gpu[r], 1e-5)) + self.assertTrue(stamps_old[r].approx_equal(stamps_cpu[r], 1e-5)) # Check the median stamps. params.stamp_type = StampType.STAMP_MEDIAN @@ -118,13 +121,11 @@ def test_coadd_gpu_parity(self): self.search.median_sci_stamp(self.trj, radius, goodIdx[0]), self.search.median_sci_stamp(self.trj, radius, goodIdx[1]), ] - stamps_new = self.search.gpu_coadded_stamps(results, goodIdx, params) + stamps_gpu = self.search.coadded_stamps(results, goodIdx, params, True) + stamps_cpu = self.search.coadded_stamps(results, goodIdx, params, False) for r in range(2): - for x in range(width): - for y in range(width): - self.assertAlmostEqual( - stamps_old[r].get_pixel(x, y), stamps_new[r].get_pixel(x, y), delta=1e-5 - ) + self.assertTrue(stamps_old[r].approx_equal(stamps_gpu[r], 1e-5)) + self.assertTrue(stamps_old[r].approx_equal(stamps_cpu[r], 1e-5)) if __name__ == "__main__":