Merge pull request #431 from dirac-institute/cleanups

A series of small cleanups for KBMOD

src/kbmod/search/bindings.cpp

@@ -18,7 +18,6 @@ namespace py = pybind11;
 #include "kernel_testing_helpers.cpp"
 #include "psi_phi_array.cpp"
 
-
 PYBIND11_MODULE(search, m) {
     m.attr("KB_NO_DATA") = pybind11::float_(search::NO_DATA);
     m.attr("HAS_GPU") = pybind11::bool_(search::HAVE_GPU);

src/kbmod/search/layered_image.cpp

@@ -229,7 +229,7 @@ RawImage LayeredImage::generate_psi_image() {
     const int num_pixels = get_npixels();
     for (int p = 0; p < num_pixels; ++p) {
         float var_pix = var_array[p];
-        if (var_pix != NO_DATA) {
+        if (var_pix != NO_DATA && var_pix != 0.0 && sci_array[p] != NO_DATA) {
             result_arr[p] = sci_array[p] / var_pix;
         } else {
             result_arr[p] = NO_DATA;
@@ -251,7 +251,7 @@ RawImage LayeredImage::generate_phi_image() {
     const int num_pixels = get_npixels();
     for (int p = 0; p < num_pixels; ++p) {
         float var_pix = var_array[p];
-        if (var_pix != NO_DATA) {
+        if (var_pix != NO_DATA && var_pix != 0.0) {
             result_arr[p] = 1.0 / var_pix;
         } else {
             result_arr[p] = NO_DATA;
@@ -280,28 +280,26 @@ static void layered_image_bindings(py::module& m) {
             .def("contains", &li::contains, pydocs::DOC_LayeredImage_cointains)
             .def("get_science_pixel", &li::get_science_pixel, pydocs::DOC_LayeredImage_get_science_pixel)
             .def("get_variance_pixel", &li::get_variance_pixel, pydocs::DOC_LayeredImage_get_variance_pixel)
-            .def(
-                    "contains",
-                    [](li& cls, int i, int j) {
-                        return cls.contains({i, j});
-                    })
-            .def(
-                    "get_science_pixel",
-                    [](li& cls, int i, int j) {
-                        return cls.get_science_pixel({i, j});
-                    })
-            .def(
-                    "get_variance_pixel",
-                    [](li& cls, int i, int j) {
-                        return cls.get_variance_pixel({i, j});
-                    })
+            .def("contains",
+                 [](li& cls, int i, int j) {
+                     return cls.contains({i, j});
+                 })
+            .def("get_science_pixel",
+                 [](li& cls, int i, int j) {
+                     return cls.get_science_pixel({i, j});
+                 })
+            .def("get_variance_pixel",
+                 [](li& cls, int i, int j) {
+                     return cls.get_variance_pixel({i, j});
+                 })
             .def("set_psf", &li::set_psf, pydocs::DOC_LayeredImage_set_psf)
             .def("get_psf", &li::get_psf, py::return_value_policy::reference_internal,
                  pydocs::DOC_LayeredImage_get_psf)
             .def("binarize_mask", &li::binarize_mask, pydocs::DOC_LayeredImage_binarize_mask)
             .def("apply_mask", &li::apply_mask, pydocs::DOC_LayeredImage_apply_mask)
             .def("union_masks", &li::union_masks, pydocs::DOC_LayeredImage_union_masks)
-            .def("union_threshold_masking", &li::union_threshold_masking, pydocs::DOC_LayeredImage_union_threshold_masking)
+            .def("union_threshold_masking", &li::union_threshold_masking,
+                 pydocs::DOC_LayeredImage_union_threshold_masking)
             .def("sub_template", &li::subtract_template, pydocs::DOC_LayeredImage_sub_template)
             .def("save_layers", &li::save_layers, pydocs::DOC_LayeredImage_save_layers)
             .def("get_science", &li::get_science, py::return_value_policy::reference_internal,

src/kbmod/search/layered_image.h

@@ -42,13 +42,13 @@ class LayeredImage {
     // Getter functions for the pixels of the science and variance layers that check
     // the mask layer for any set bits.
     inline float get_science_pixel(const Index& idx) const {
-        return contains(idx) ? (mask.get_pixel(idx) == 0 ? science.get_pixel(idx) : NO_DATA) : NO_DATA;
+        // The get_pixel() functions perform the bounds checking and will return NO_DATA for out of bounds.
+        return mask.get_pixel(idx) == 0 ? science.get_pixel(idx) : NO_DATA;
     }
 
     inline float get_variance_pixel(const Index& idx) const {
-        return contains(idx) ? 
-            (mask.get_pixel(idx) == 0 ? variance.get_pixel(idx) : NO_DATA) :
-             NO_DATA;
+        // The get_pixel() functions perform the bounds checking and will return NO_DATA for out of bounds.
+        return mask.get_pixel(idx) == 0 ? variance.get_pixel(idx) : NO_DATA;
     }
 
     inline bool contains(const Index& idx) const {

src/kbmod/search/psi_phi_array.cpp

@@ -177,11 +177,12 @@ void set_encode_cpu_psi_phi_array(PsiPhiArray& data, const std::vector<RawImage>
         throw std::runtime_error("CPU PsiPhi already allocated.");
     }
     if (debug) {
-        printf("Allocating CPU memory for encoded PsiPhi array using %lu bytes.\n", data.get_total_array_size());
+        printf("Allocating CPU memory for encoded PsiPhi array using %lu bytes.\n",
+               data.get_total_array_size());
     }
     T* encoded = (T*)malloc(data.get_total_array_size());
     if (encoded == nullptr) {
-      throw std::runtime_error("Unable to allocate space for CPU PsiPhi array.");
+        throw std::runtime_error("Unable to allocate space for CPU PsiPhi array.");
     }
 
     // Create a safe maximum that is slightly less than the true max to avoid
@@ -224,7 +225,7 @@ void set_float_cpu_psi_phi_array(PsiPhiArray& data, const std::vector<RawImage>&
     }
     float* encoded = (float*)malloc(data.get_total_array_size());
     if (encoded == nullptr) {
-      throw std::runtime_error("Unable to allocate space for CPU PsiPhi array.");
+        throw std::runtime_error("Unable to allocate space for CPU PsiPhi array.");
     }
 
     int current_index = 0;
@@ -266,12 +267,10 @@ void fill_psi_phi_array(PsiPhiArray& result_data, int num_bytes, const std::vect
         result_data.set_phi_scaling(phi_params[0], phi_params[1], phi_params[2]);
 
         if (debug) {
-            printf("Encoding psi to %i bytes min=%f, max=%f, scale=%f\n",
-                   result_data.get_num_bytes(), psi_params[0], psi_params[1],
-                   psi_params[2]);
-            printf("Encoding phi to %i bytes min=%f, max=%f, scale=%f\n",
-                   result_data.get_num_bytes(), phi_params[0], phi_params[1],
-                   phi_params[2]);
+            printf("Encoding psi to %i bytes min=%f, max=%f, scale=%f\n", result_data.get_num_bytes(),
+                   psi_params[0], psi_params[1], psi_params[2]);
+            printf("Encoding phi to %i bytes min=%f, max=%f, scale=%f\n", result_data.get_num_bytes(),
+                   phi_params[0], phi_params[1], phi_params[2]);
         }
 
         // Do the local encoding.
@@ -281,15 +280,18 @@ void fill_psi_phi_array(PsiPhiArray& result_data, int num_bytes, const std::vect
             set_encode_cpu_psi_phi_array<uint16_t>(result_data, psi_imgs, phi_imgs, debug);
         }
     } else {
-        if (debug) { printf("Encoding psi and phi as floats.\n"); }
+        if (debug) {
+            printf("Encoding psi and phi as floats.\n");
+        }
         // Just interleave psi and phi images.
         set_float_cpu_psi_phi_array(result_data, psi_imgs, phi_imgs, debug);
     }
 
 #ifdef HAVE_CUDA
     // Create a copy of the encoded data in GPU memory.
     if (debug) {
-        printf("Allocating GPU memory for PsiPhi array using %lu bytes.\n", result_data.get_total_array_size());
+        printf("Allocating GPU memory for PsiPhi array using %lu bytes.\n",
+               result_data.get_total_array_size());
     }
     device_allocate_psi_phi_array(&result_data);
     if (result_data.get_gpu_array_ptr() == nullptr) {

src/kbmod/search/psi_phi_array_utils.h

@@ -26,7 +26,7 @@ namespace search {
 std::array<float, 3> compute_scale_params_from_image_vect(const std::vector<RawImage>& imgs, int num_bytes);
 
 void fill_psi_phi_array(PsiPhiArray& result_data, int num_bytes, const std::vector<RawImage>& psi_imgs,
-                        const std::vector<RawImage>& phi_imgs, bool debug=false);
+                        const std::vector<RawImage>& phi_imgs, bool debug = false);
 
 } /* namespace search */
 

src/kbmod/search/pydocs/layered_image_docs.h

@@ -238,11 +238,29 @@ static const auto DOC_LayeredImage_get_variance_pixel = R"doc(
   )doc";
 
 static const auto DOC_LayeredImage_generate_psi_image = R"doc(
-  todo
+  Generates the full psi image where the value of each pixel p in the
+  resulting image is science[p] / variance[p]. To handle masked bits
+  apply_mask() must be called before the psi image is generated. Otherwise,
+  all pixels are used.
+  Convolves the resulting image with the PSF.
+
+  Returns
+  -------
+  result : `kbmod.RawImage`
+      A ``RawImage`` of the same dimensions as the ``LayeredImage``.
   )doc";
 
 static const auto DOC_LayeredImage_generate_phi_image = R"doc(
-  todo
+  Generates the full phi image where the value of each pixel p in the
+  resulting image is 1.0 / variance[p]. To handle masked bits
+  apply_mask() must be called before the phi image is generated. Otherwise,
+  all pixels are used.
+  Convolves the resulting image with the PSF.
+
+  Returns
+  -------
+  result : `kbmod.RawImage`
+      A ``RawImage`` of the same dimensions as the ``LayeredImage``.
   )doc";
 }  // namespace pydocs
 

src/kbmod/search/raw_image.cpp

@@ -110,8 +110,7 @@ float RawImage::interpolate(const Point& p) const {
     return total / sumWeights;
 }
 
-RawImage RawImage::create_stamp(const Point& p, const int radius,
-                                const bool keep_no_data) const {
+RawImage RawImage::create_stamp(const Point& p, const int radius, const bool keep_no_data) const {
     if (radius < 0) throw std::runtime_error("stamp radius must be at least 0");
 
     const int dim = radius * 2 + 1;
@@ -123,8 +122,7 @@ RawImage RawImage::create_stamp(const Point& p, const int radius,
     Image stamp = Image::Constant(dim, dim, NO_DATA);
     stamp.block(anchor.i, anchor.j, h, w) = image.block(corner.i, corner.j, h, w);
 
-    if (!keep_no_data)
-        stamp = (stamp.array() == NO_DATA).select(0.0, stamp);
+    if (!keep_no_data) stamp = (stamp.array() == NO_DATA).select(0.0, stamp);
 
     return RawImage(stamp);
 }
@@ -580,29 +578,25 @@ static void raw_image_bindings(py::module& m) {
             .def("set_pixel", &rie::set_pixel, pydocs::DOC_RawImage_set_pixel)
             .def("set_all", &rie::set_all, pydocs::DOC_RawImage_set_all)
             // python interface adapters (avoids having to construct Index & Point)
-            .def(
-                    "get_pixel",
-                    [](rie& cls, int i, int j) {
-                        return cls.get_pixel({i, j});
-                    })
-            .def(
-                    "pixel_has_data",
-                    [](rie& cls, int i, int j) {
-                        return cls.pixel_has_data({i, j});
-                    })
-            .def(
-                    "set_pixel",
-                    [](rie& cls, int i, int j, double val) {
-                        cls.set_pixel({i, j}, val);
-                    })
+            .def("get_pixel",
+                 [](rie& cls, int i, int j) {
+                     return cls.get_pixel({i, j});
+                 })
+            .def("pixel_has_data",
+                 [](rie& cls, int i, int j) {
+                     return cls.pixel_has_data({i, j});
+                 })
+            .def("set_pixel",
+                 [](rie& cls, int i, int j, double val) {
+                     cls.set_pixel({i, j}, val);
+                 })
             // methods
             .def("l2_allclose", &rie::l2_allclose, pydocs::DOC_RawImage_l2_allclose)
             .def("compute_bounds", &rie::compute_bounds, pydocs::DOC_RawImage_compute_bounds)
             .def("find_peak", &rie::find_peak, pydocs::DOC_RawImage_find_peak)
             .def("find_central_moments", &rie::find_central_moments,
-                pydocs::DOC_RawImage_find_central_moments)
-            .def("center_is_local_max", &rie::center_is_local_max, 
-                pydocs::DOC_RawImage_center_is_local_max)
+                 pydocs::DOC_RawImage_find_central_moments)
+            .def("center_is_local_max", &rie::center_is_local_max, pydocs::DOC_RawImage_center_is_local_max)
             .def("create_stamp", &rie::create_stamp, pydocs::DOC_RawImage_create_stamp)
             .def("interpolate", &rie::interpolate, pydocs::DOC_RawImage_interpolate)
             .def("interpolated_add", &rie::interpolated_add, pydocs::DOC_RawImage_interpolated_add)
@@ -620,21 +614,17 @@ static void raw_image_bindings(py::module& m) {
             .def("append_fits_extension", &rie::append_to_fits, pydocs::DOC_RawImage_append_to_fits)
             .def("load_fits", &rie::from_fits, pydocs::DOC_RawImage_load_fits)
             // python interface adapters
-            .def(
-                    "create_stamp",
-                    [](rie& cls, float x, float y, int radius, bool keep_no_data) {
-                        return cls.create_stamp({x, y}, radius, keep_no_data);
-                    })
-            .def(
-                    "interpolate",
-                    [](rie& cls, float x, float y) {
-                        return cls.interpolate({x, y});
-                    })
-            .def(
-                    "interpolated_add",
-                    [](rie& cls, float x, float y, float val) {
-                        cls.interpolated_add({x, y}, val);
-                    });
+            .def("create_stamp",
+                 [](rie& cls, float x, float y, int radius, bool keep_no_data) {
+                     return cls.create_stamp({x, y}, radius, keep_no_data);
+                 })
+            .def("interpolate",
+                 [](rie& cls, float x, float y) {
+                     return cls.interpolate({x, y});
+                 })
+            .def("interpolated_add", [](rie& cls, float x, float y, float val) {
+                cls.interpolated_add({x, y}, val);
+            });
 }
 #endif
 

tests/test_layered_image.py

@@ -44,6 +44,29 @@ def test_create(self):
                 self.assertGreaterEqual(science.get_pixel(y, x), -100.0)
                 self.assertLessEqual(science.get_pixel(y, x), 100.0)
 
+                # Check direct lookup of pixel values matches the RawImage lookup.
+                self.assertEqual(science.get_pixel(y, x), self.image.get_science_pixel(y, x))
+                self.assertEqual(variance.get_pixel(y, x), self.image.get_variance_pixel(y, x))
+
+        # Check that the LayeredImage pixel lookups work with a masked pixel.
+        # But the the mask was not applied yet to the images themselves.
+        mask.set_pixel(5, 6, 1)
+        self.assertGreater(science.get_pixel(5, 6), KB_NO_DATA)
+        self.assertGreater(variance.get_pixel(5, 6), KB_NO_DATA)
+        self.assertEqual(self.image.get_science_pixel(5, 6), KB_NO_DATA)
+        self.assertEqual(self.image.get_variance_pixel(5, 6), KB_NO_DATA)
+
+        # Test that out of bounds pixel lookups are handled correctly.
+        self.assertEqual(self.image.get_science_pixel(-1, 1), KB_NO_DATA)
+        self.assertEqual(self.image.get_science_pixel(1, -1), KB_NO_DATA)
+        self.assertEqual(self.image.get_science_pixel(self.image.get_height() + 1, 1), KB_NO_DATA)
+        self.assertEqual(self.image.get_science_pixel(1, self.image.get_width() + 1), KB_NO_DATA)
+
+        self.assertEqual(self.image.get_variance_pixel(-1, 1), KB_NO_DATA)
+        self.assertEqual(self.image.get_variance_pixel(1, -1), KB_NO_DATA)
+        self.assertEqual(self.image.get_variance_pixel(self.image.get_height() + 1, 1), KB_NO_DATA)
+        self.assertEqual(self.image.get_variance_pixel(1, self.image.get_width() + 1), KB_NO_DATA)
+
     def test_create_from_layers(self):
         sci = RawImage(30, 40)
         for y in range(40):
@@ -285,7 +308,11 @@ def test_psi_and_phi_image(self):
             for x in range(6):
                 sci.set_pixel(y, x, float(x))
                 var.set_pixel(y, x, float(y + 1))
+
+        # Mask a single pixel and set another to variance of zero.
+        sci.set_pixel(3, 1, KB_NO_DATA)
         var.set_pixel(3, 1, KB_NO_DATA)
+        var.set_pixel(3, 2, 0.0)
 
         # Generate and check psi and phi images.
         psi = img.generate_psi_image()
@@ -298,12 +325,15 @@ def test_psi_and_phi_image(self):
 
         for y in range(5):
             for x in range(6):
-                has_data = not (x == 1 and y == 3)
+                has_data = y != 3 or x == 0 or x > 2
                 self.assertEqual(psi.pixel_has_data(y, x), has_data)
                 self.assertEqual(phi.pixel_has_data(y, x), has_data)
-                if x != 1 or y != 3:
+                if has_data:
                     self.assertAlmostEqual(psi.get_pixel(y, x), x / (y + 1))
                     self.assertAlmostEqual(phi.get_pixel(y, x), 1.0 / (y + 1))
+                else:
+                    self.assertEqual(psi.get_pixel(y, x), KB_NO_DATA)
+                    self.assertEqual(phi.get_pixel(y, x), KB_NO_DATA)
 
     def test_subtract_template(self):
         sci = self.image.get_science()