Merge branch 'main' into remove_legacy_visualization

docs/source/user_manual/search_params.rst

@@ -38,7 +38,8 @@ This document serves to provide a quick overview of the existing parameters and
 | ``coadds``             | []                          | A list of additional coadds to create. |
 |                        |                             | These are not used in filtering, but   |
 |                        |                             | saved to columns for analysis. Can     |
-|                        |                             | include: "sum", "mean", and "median".  |
+|                        |                             | include: "sum", "mean", "median", and  |
+|                        |                             | "weighted".
 |                        |                             | The filtering coadd is controlled by   |
 |                        |                             | the ``stamp_type`` parameter.          |
 +------------------------+-----------------------------+----------------------------------------+
@@ -135,8 +136,10 @@ This document serves to provide a quick overview of the existing parameters and
 |                        |                             | filtering (if ``do_stamp_filter=True``)|
 |                        |                             | if:                                    |
 |                        |                             | * ``sum`` - (default) Per pixel sum    |
-|                        |                             | * ``median`` - A per pixel median      |
-|                        |                             | * ``mean`` - A per pixel mean          |
+|                        |                             | * ``median`` - Per pixel median        |
+|                        |                             | * ``mean`` - Per pixel mean            |
+|                        |                             | * ``weighted`` - Per pixel mean        |
+|                        |                             |   weighted by 1.0 / variance.          |
 +------------------------+-----------------------------+----------------------------------------+
 | ``track_filtered``     | False                       | A Boolean indicating whether to track  |
 |                        |                             | the filtered trajectories. Warning     |

src/kbmod/filters/stamp_filters.py

@@ -57,6 +57,8 @@ def extract_search_parameters_from_config(config):
         params.stamp_type = StampType.STAMP_MEAN
     elif stamp_type == "cpp_sum" or stamp_type == "sum":
         params.stamp_type = StampType.STAMP_SUM
+    elif stamp_type == "weighted":
+        params.stamp_type = StampType.STAMP_VAR_WEIGHTED
     else:
         raise ValueError(f"Unrecognized stamp type: {stamp_type}")
 
@@ -204,6 +206,8 @@ def append_coadds(result_data, im_stack, coadd_types, radius, chunk_size=100_000
             params.stamp_type = StampType.STAMP_MEAN
         elif coadd_type == "sum":
             params.stamp_type = StampType.STAMP_SUM
+        elif coadd_type == "weighted":
+            params.stamp_type = StampType.STAMP_VAR_WEIGHTED
         else:
             raise ValueError(f"Unrecognized stamp type: {coadd_type}")
 

src/kbmod/search/bindings.cpp

@@ -28,6 +28,7 @@ PYBIND11_MODULE(search, m) {
             .value("STAMP_SUM", search::StampType::STAMP_SUM)
             .value("STAMP_MEAN", search::StampType::STAMP_MEAN)
             .value("STAMP_MEDIAN", search::StampType::STAMP_MEDIAN)
+            .value("STAMP_VAR_WEIGHTED", search::StampType::STAMP_VAR_WEIGHTED)
             .export_values();
     logging::logging_bindings(m);
     indexing::index_bindings(m);

src/kbmod/search/common.h

@@ -24,7 +24,7 @@ constexpr unsigned short THREAD_DIM_Y = 2;
 // The NO_DATA flag indicates masked values in the image.
 constexpr float NO_DATA = NAN;
 
-enum StampType { STAMP_SUM = 0, STAMP_MEAN, STAMP_MEDIAN };
+enum StampType { STAMP_SUM = 0, STAMP_MEAN, STAMP_MEDIAN, STAMP_VAR_WEIGHTED };
 
 // A helper function to check that a pixel value is valid. This should include
 // both masked pixel values (NO_DATA above) and other invalid values (e.g. inf).

src/kbmod/search/pydocs/stamp_creator_docs.h

@@ -165,6 +165,49 @@ static const auto DOC_StampCreator_filter_stamp = R"doc(
       Whether or not to filter the stamp.
   )doc";
 
+static const auto DOC_StampCreator_create_variance_stamps = R"doc(
+  Create a vector of stamps from the variance layer centered on the
+  predicted position of an Trajectory at different times.
+
+  Parameters
+  ----------
+  stack : `ImageStack`
+      The stack of images to use.
+  trj : `Trajectory`
+      The trajectory to project to each time.
+  radius : `int`
+      The stamp radius. Width = 2*radius+1.
+  use_index : `list` of `bool`
+      A list (vector) of Booleans indicating whether or not to use each time step.
+      An empty (size=0) vector will use all time steps.
+
+  Returns
+  -------
+  `list` of `RawImage`
+      The stamps.
+  )doc";
+
+static const auto DOC_StampCreator_get_variance_weighted_stamp = R"doc(
+  Create a weighted-mean stamp where the weight for each pixel is 1.0 / variance.
+
+  Parameters
+  ----------
+  stack : `ImageStack`
+      The stack of images to use.
+  trj : `Trajectory`
+      The trajectory to project to each time.
+  radius : `int`
+      The stamp radius. Width = 2*radius+1.
+  use_index : `list` of `bool`
+      A list (vector) of Booleans indicating whether or not to use each time step.
+      An empty (size=0) vector will use all time steps.
+
+  Returns
+  -------
+  `RawImage`
+      The co-added stamp.
+  )doc";
+
 }  // namespace pydocs
 
 #endif /* STAMP_CREATOR_DOCS */

src/kbmod/search/stamp_creator.cpp

@@ -68,7 +68,8 @@ std::vector<RawImage> StampCreator::get_coadded_stamps(ImageStack& stack, std::v
     rs_logger->info("Performing stamp filtering on " + std::to_string(t_array.size()) + " trajectories.");
     DebugTimer timer = DebugTimer("stamp filtering", rs_logger);
 
-    if (use_gpu) {
+    // We use the GPU if we have it for everything except STAMP_VAR_WEIGHTED which is CPU only.
+    if (use_gpu && (params.stamp_type != STAMP_VAR_WEIGHTED)) {
 #ifdef HAVE_CUDA
         rs_logger->info("Performing co-adds on the GPU.");
         return get_coadded_stamps_gpu(stack, t_array, use_index_vect, params);
@@ -87,19 +88,19 @@ std::vector<RawImage> StampCreator::get_coadded_stamps_cpu(ImageStack& stack,
     std::vector<RawImage> results(num_trajectories);
 
     for (uint64_t i = 0; i < num_trajectories; ++i) {
-        std::vector<RawImage> stamps =
-                StampCreator::create_stamps(stack, t_array[i], params.radius, true, use_index_vect[i]);
-
         RawImage coadd(1, 1);
         switch (params.stamp_type) {
             case STAMP_MEDIAN:
-                coadd = create_median_image(stamps);
+                coadd = get_median_stamp(stack, t_array[i], params.radius, use_index_vect[i]);
                 break;
             case STAMP_MEAN:
-                coadd = create_mean_image(stamps);
+                coadd = get_mean_stamp(stack, t_array[i], params.radius, use_index_vect[i]);
                 break;
             case STAMP_SUM:
-                coadd = create_summed_image(stamps);
+                coadd = get_summed_stamp(stack, t_array[i], params.radius, use_index_vect[i]);
+                break;
+            case STAMP_VAR_WEIGHTED:
+                coadd = get_variance_weighted_stamp(stack, t_array[i], params.radius, use_index_vect[i]);
                 break;
             default:
                 throw std::runtime_error("Invalid stamp coadd type.");
@@ -272,6 +273,64 @@ std::vector<RawImage> StampCreator::get_coadded_stamps_gpu(ImageStack& stack,
     return results;
 }
 
+std::vector<RawImage> StampCreator::create_variance_stamps(ImageStack& stack, const Trajectory& trj,
+                                                           int radius, const std::vector<bool>& 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;
+
+    std::vector<RawImage> stamps;
+    unsigned int num_times = stack.img_count();
+    for (unsigned int i = 0; i < num_times; ++i) {
+        if (use_all_stamps || use_index[i]) {
+            // Calculate the trajectory position.
+            double time = stack.get_zeroed_time(i);
+            Point pos{trj.get_x_pos(time), trj.get_y_pos(time)};
+            RawImage& img = stack.get_single_image(i).get_variance();
+
+            RawImage stamp = img.create_stamp(pos, radius, true /* keep_no_data */);
+            stamps.push_back(std::move(stamp));
+        }
+    }
+    return stamps;
+}
+
+RawImage StampCreator::get_variance_weighted_stamp(ImageStack& stack, const Trajectory& trj, int radius,
+                                                   const std::vector<bool>& use_index) {
+    unsigned int num_images = stack.img_count();
+    if (num_images == 0) throw std::runtime_error("Unable to create mean image given 0 images.");
+    unsigned int stamp_width = 2 * radius + 1;
+
+    // Make the stamps for each time step.
+    std::vector<bool> empty_vect;
+    std::vector<RawImage> sci_stamps = create_stamps(stack, trj, radius, true /*=keep_no_data*/, use_index);
+    std::vector<RawImage> var_stamps = create_variance_stamps(stack, trj, radius, use_index);
+
+    // Do the weighted mean.
+    Image result = Image::Zero(stamp_width, stamp_width);
+    for (int y = 0; y < stamp_width; ++y) {
+        for (int x = 0; x < stamp_width; ++x) {
+            float sum = 0.0;
+            float scale = 0.0;
+            for (int i = 0; i < num_images; ++i) {
+                float sci_val = sci_stamps[i].get_pixel({y, x});
+                float var_val = var_stamps[i].get_pixel({y, x});
+                if (pixel_value_valid(sci_val) && pixel_value_valid(var_val) && (var_val != 0.0)) {
+                    sum += sci_val / var_val;
+                    scale += 1.0 / var_val;
+                }
+            }
+
+            if (scale > 0.0) {
+                result(y, x) = sum / scale;
+            } else {
+                result(y, x) = 0.0;
+            }
+        }  // for x
+    }      // for y
+    return RawImage(result);
+}
+
 #ifdef Py_PYTHON_H
 static void stamp_creator_bindings(py::module& m) {
     using sc = search::StampCreator;
@@ -284,6 +343,11 @@ static void stamp_creator_bindings(py::module& m) {
             .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);
 }
 #endif /* Py_PYTHON_H */

src/kbmod/search/stamp_creator.h

@@ -53,6 +53,14 @@ class StampCreator {
     // Function to do the actual stamp filtering.
     static bool filter_stamp(const RawImage& img, const StampParameters& params);
 
+    // Function for generating variance stamps. All times are returned and NO_DATA values are preserved.
+    static std::vector<RawImage> create_variance_stamps(ImageStack& stack, const Trajectory& trj, int radius,
+                                                        const std::vector<bool>& use_index);
+
+    // 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<bool>& use_index);
+
     virtual ~StampCreator(){};
 };
 

tests/test_stamp_creator.py

@@ -0,0 +1,122 @@
+import numpy as np
+import unittest
+
+from kbmod.fake_data.fake_data_creator import create_fake_times, FakeDataSet
+from kbmod.search import ImageStack, LayeredImage, PSF, StampCreator, Trajectory
+
+
+class test_stamp_creator(unittest.TestCase):
+    def setUp(self):
+        # Create a fake data set to use in the tests.
+        self.image_count = 10
+        self.fake_times = create_fake_times(self.image_count, 57130.2, 1, 0.01, 1)
+        self.ds = FakeDataSet(
+            25,  # width
+            35,  # height
+            self.fake_times,  # time stamps
+            1.0,  # noise level
+            0.5,  # psf value
+            True,  # Use a fixed seed for testing
+        )
+
+        # Insert a single fake object with known parameters.
+        self.trj = 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.trj, 1, True, [])
+        self.assertEqual(len(stamps), self.image_count)
+        for i in range(self.image_count):
+            self.assertEqual(stamps[i].image.shape, (3, 3))
+
+            pix_val = self.ds.stack.get_single_image(i).get_science().get_pixel(7 + i, 8 + 2 * i)
+            if np.isnan(pix_val):
+                self.assertTrue(np.isnan(stamps[i].get_pixel(1, 1)))
+            else:
+                self.assertAlmostEqual(pix_val, stamps[i].get_pixel(1, 1))
+
+        # 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.trj, 1, True, use_times)
+        self.assertEqual(len(stamps), np.count_nonzero(use_times))
+
+        stamp_count = 0
+        for i in range(self.image_count):
+            if use_times[i]:
+                self.assertEqual(stamps[stamp_count].image.shape, (3, 3))
+
+                pix_val = self.ds.stack.get_single_image(i).get_science().get_pixel(7 + i, 8 + 2 * i)
+                if np.isnan(pix_val):
+                    self.assertTrue(np.isnan(stamps[stamp_count].get_pixel(1, 1)))
+                else:
+                    self.assertAlmostEqual(pix_val, stamps[stamp_count].get_pixel(1, 1))
+
+                stamp_count += 1
+
+    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.trj, 1, [])
+        self.assertEqual(len(stamps), self.image_count)
+        for i in range(self.image_count):
+            self.assertEqual(stamps[i].image.shape, (3, 3))
+
+            pix_val = self.ds.stack.get_single_image(i).get_variance().get_pixel(7 + i, 8 + 2 * i)
+            if np.isnan(pix_val):
+                self.assertTrue(np.isnan(stamps[i].get_pixel(1, 1)))
+            else:
+                self.assertAlmostEqual(pix_val, stamps[i].get_pixel(1, 1))
+
+        # 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.trj, 1, use_times)
+        self.assertEqual(len(stamps), np.count_nonzero(use_times))
+
+        stamp_count = 0
+        for i in range(self.image_count):
+            if use_times[i]:
+                self.assertEqual(stamps[stamp_count].image.shape, (3, 3))
+
+                pix_val = self.ds.stack.get_single_image(i).get_variance().get_pixel(7 + i, 8 + 2 * i)
+                if np.isnan(pix_val):
+                    self.assertTrue(np.isnan(stamps[stamp_count].get_pixel(1, 1)))
+                else:
+                    self.assertAlmostEqual(pix_val, stamps[stamp_count].get_pixel(1, 1))
+
+                stamp_count += 1
+
+    def test_get_variance_weighted_stamp(self):
+        sci1 = np.array([[1.0, 1.0, 1.0], [1.0, 1.0, 1.0], [1.0, 1.0, 1.0]], dtype=np.single)
+        var1 = np.array([[1.0, 1.0, 1.0], [1.0, 1.0, 1.0], [1.0, 1.0, 1.0]], dtype=np.single)
+        msk1 = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 1.0]], dtype=np.single)
+        layer1 = LayeredImage(sci1, var1, msk1, PSF(1e-12), 0.0)
+        layer1.apply_mask(0xFFFFFF)
+
+        sci2 = np.array([[2.0, 2.0, 2.0], [2.0, 2.0, 2.0], [2.0, 2.0, 2.0]], dtype=np.single)
+        var2 = np.array([[0.5, 0.5, 0.5], [0.5, 0.5, 0.5], [0.5, 0.5, 0.5]], dtype=np.single)
+        msk2 = np.array([[1.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], dtype=np.single)
+        layer2 = LayeredImage(sci2, var2, msk2, PSF(1e-12), 0.0)
+        layer2.apply_mask(0xFFFFFF)
+
+        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, [])
+        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, [])
+        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, [])
+        self.assertEqual(stamp.image.shape, (1, 1))
+        self.assertAlmostEqual(stamp.get_pixel(0, 0), 2.0)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_stamp_filters.py

@@ -158,20 +158,31 @@ def test_append_coadds(self):
         self.assertFalse("coadd_sum" in keep.colnames)
         self.assertFalse("coadd_mean" in keep.colnames)
         self.assertFalse("coadd_median" in keep.colnames)
+        self.assertFalse("coadd_weighted" in keep.colnames)
         self.assertFalse("stamp" in keep.colnames)
 
         # Adding nothing does nothing.
         append_coadds(keep, self.ds.stack, [], 3)
         self.assertFalse("coadd_sum" in keep.colnames)
         self.assertFalse("coadd_mean" in keep.colnames)
         self.assertFalse("coadd_median" in keep.colnames)
+        self.assertFalse("coadd_weighted" in keep.colnames)
         self.assertFalse("stamp" in keep.colnames)
 
         # Adding "mean" and "median" does only those.
         append_coadds(keep, self.ds.stack, ["median", "mean"], 3)
         self.assertFalse("coadd_sum" in keep.colnames)
         self.assertTrue("coadd_mean" in keep.colnames)
         self.assertTrue("coadd_median" in keep.colnames)
+        self.assertFalse("coadd_weighted" in keep.colnames)
+        self.assertFalse("stamp" in keep.colnames)
+
+        # We can add "weighted" later.
+        append_coadds(keep, self.ds.stack, ["weighted"], 3)
+        self.assertFalse("coadd_sum" in keep.colnames)
+        self.assertTrue("coadd_mean" in keep.colnames)
+        self.assertTrue("coadd_median" in keep.colnames)
+        self.assertTrue("coadd_weighted" in keep.colnames)
         self.assertFalse("stamp" in keep.colnames)
 
         # Check that all coadds are generated without filtering.