Merge pull request #345 from dirac-institute/prefactor

Cleanup LayeredImage PSF usage

README.md

@@ -111,9 +111,13 @@ velocity = (2, 0)
 # Inject object into images
 for im in imgs:
     dt = im.get_obstime() - t0
-    im.add_object(position[0] + dt * velocity[0], 
-                  position[1] + dt * velocity[1], 
-                  flux)
+    add_fake_object(
+        im,
+        position[0] + dt * velocity[0],
+        position[1] + dt * velocity[1],
+        flux,
+        psf,
+    )
 
 # Create a new image stack with the inserted object.
 stack = kb.image_stack(imgs)

src/kbmod/fake_data_creator.py

@@ -15,6 +15,44 @@
 from kbmod.search import *
 
 
+def add_fake_object(img, x, y, flux, psf=None):
+    """Add a fake object to a LayeredImage or RawImage
+
+    Parameters
+    ----------
+    img : RawImage or LayeredImage
+        The image to modify.
+    x : float
+        The x pixel location of the fake object.
+    y : float
+        The y pixel location of the fake object.
+    flux : float
+        The flux value.
+    psf : PointSpreadFunc
+        The PSF for the image.
+    """
+    if type(img) is layered_image:
+        sci = img.get_science()
+    else:
+        sci = img
+
+    if psf is None:
+        sci.add_pixel_interp(x, y, flux)
+    else:
+        dim = psf.get_dim()
+        initial_x = x - psf.get_radius()
+        initial_y = y - psf.get_radius()
+
+        for i in range(dim):
+            for j in range(dim):
+                sci.add_pixel_interp(initial_x + i, initial_y + j, flux * psf.get_value(i, j))
+
+    # The python/C++ interface requires us to explicitly re-set the science
+    # image in a LayeredImage.
+    if sci is not img:
+        img.set_science(sci)
+
+
 class FakeDataSet:
     """This class creates fake data sets for testing and demo notebooks."""
 
@@ -107,9 +145,9 @@ def insert_object(self, trj):
             # Get the image for the timestep, add the object, and
             # re-set the image. This last step needs to be done
             # explicitly because of how pybind handles references.
-            current_layered_image = self.stack.get_single_image(i)
-            current_layered_image.add_object(px, py, trj.flux)
-            self.stack.set_single_image(i, current_layered_image)
+            current = self.stack.get_single_image(i)
+            add_fake_object(current, px, py, trj.flux, current.get_psf())
+            self.stack.set_single_image(i, current)
 
         # Save the trajectory into the internal list.
         self.trajectories.append(trj)

src/kbmod/search/LayeredImage.cpp

@@ -9,9 +9,7 @@
 
 namespace search {
 
-LayeredImage::LayeredImage(std::string path, const PointSpreadFunc& psf) : psf(psf), psf_sq(psf) {
-    psf_sq.squarePSF();
-
+LayeredImage::LayeredImage(std::string path, const PointSpreadFunc& psf) : psf(psf) {
     int f_begin = path.find_last_of("/");
     int f_end = path.find_last_of(".fits") - 4;
     filename = path.substr(f_begin, f_end - f_begin);
@@ -35,7 +33,7 @@ LayeredImage::LayeredImage(std::string path, const PointSpreadFunc& psf) : psf(p
 
 LayeredImage::LayeredImage(const RawImage& sci, const RawImage& var, const RawImage& msk,
                            const PointSpreadFunc& psf)
-        : psf(psf), psf_sq(psf) {
+        : psf(psf) {
     // Get the dimensions of the science layer and check for consistency with
     // the other two layers.
     width = sci.getWidth();
@@ -45,9 +43,6 @@ LayeredImage::LayeredImage(const RawImage& sci, const RawImage& var, const RawIm
     if (width != msk.getWidth() or height != msk.getHeight())
         throw std::runtime_error("Science and Mask layers are not the same size.");
 
-    // Set the remaining variables.
-    psf_sq.squarePSF();
-
     // Copy the image layers.
     science = sci;
     mask = msk;
@@ -60,11 +55,10 @@ LayeredImage::LayeredImage(std::string name, int w, int h, float noise_stdev, fl
 
 LayeredImage::LayeredImage(std::string name, int w, int h, float noise_stdev, float pixel_variance, double time,
                            const PointSpreadFunc& psf, int seed)
-        : psf(psf), psf_sq(psf) {
+        : psf(psf) {
     filename = name;
     width = w;
     height = h;
-    psf_sq.squarePSF();
 
     std::vector<float> raw_sci(width * height);
     std::random_device r;
@@ -84,34 +78,24 @@ LayeredImage::LayeredImage(std::string name, int w, int h, float noise_stdev, fl
 
 void LayeredImage::setPSF(const PointSpreadFunc& new_psf) {
     psf = new_psf;
-    psf_sq = new_psf;
-    psf_sq.squarePSF();
-}
-
-void LayeredImage::addObject(float x, float y, float flux) {
-    const std::vector<float>& k = psf.getKernel();
-    int dim = psf.getDim();
-    float initial_x = x - static_cast<float>(psf.getRadius());
-    float initial_y = y - static_cast<float>(psf.getRadius());
-
-    int count = 0;
-    for (int i = 0; i < dim; ++i) {
-        for (int j = 0; j < dim; ++j) {
-            science.addPixelInterp(initial_x + static_cast<float>(i), initial_y + static_cast<float>(j),
-                                   flux * k[count]);
-            count++;
-        }
-    }
 }
 
 void LayeredImage::growMask(int steps) {
     science.growMask(steps);
     variance.growMask(steps);
 }
 
+void LayeredImage::convolveGivenPSF(const PointSpreadFunc& given_psf) {
+    science.convolve(given_psf);
+
+    // Square the PSF use that on the variance image.
+    PointSpreadFunc psfsq = PointSpreadFunc(given_psf);  // Copy
+    psfsq.squarePSF();
+    variance.convolve(psfsq);
+}
+
 void LayeredImage::convolvePSF() {
-    science.convolve(psf);
-    variance.convolve(psf_sq);
+   convolveGivenPSF(psf);
 }
 
 void LayeredImage::applyMaskFlags(int flags, const std::vector<int>& exceptions) {
@@ -217,7 +201,7 @@ RawImage LayeredImage::generatePsiImage() {
     }
 
     // Convolve with the PSF.
-    result.convolve(getPSF());
+    result.convolve(psf);
 
     return result;
 }
@@ -239,7 +223,9 @@ RawImage LayeredImage::generatePhiImage() {
     }
 
     // Convolve with the PSF squared.
-    result.convolve(getPSFSQ());
+    PointSpreadFunc psfsq = PointSpreadFunc(psf);  // Copy
+    psfsq.squarePSF();
+    result.convolve(psfsq);
 
     return result;
 }

src/kbmod/search/LayeredImage.h

@@ -36,7 +36,6 @@ class LayeredImage {
     // Set an image specific point spread function.
     void setPSF(const PointSpreadFunc& psf);
     const PointSpreadFunc& getPSF() const { return psf; }
-    const PointSpreadFunc& getPSFSQ() const { return psf_sq; }
 
     // Basic getter functions for image data.
     std::string getName() const { return filename; }
@@ -65,9 +64,6 @@ class LayeredImage {
     // Subtracts a template image from the science layer.
     void subtractTemplate(const RawImage& sub_template);
 
-    // Adds an (artificial) object to the image (science) data.
-    void addObject(float x, float y, float flux);
-
     // Saves the data in each later to a file.
     void saveLayers(const std::string& path);
 
@@ -76,7 +72,10 @@ class LayeredImage {
     void setMask(RawImage& im);
     void setVariance(RawImage& im);
 
+    // Convolve with a given PSF or the default one.
     void convolvePSF();
+    void convolveGivenPSF(const PointSpreadFunc& psf);
+
     virtual ~LayeredImage(){};
 
     // Generate psi and phi images from the science and variance layers.
@@ -91,7 +90,6 @@ class LayeredImage {
     unsigned height;
 
     PointSpreadFunc psf;
-    PointSpreadFunc psf_sq;
     RawImage science;
     RawImage mask;
     RawImage variance;

src/kbmod/search/PointSpreadFunc.cpp

@@ -9,6 +9,13 @@
 
 namespace search {
 
+PointSpreadFunc::PointSpreadFunc() : kernel(1, 1.0) {
+    dim = 1;
+    radius = 0;
+    width = 1e-12;
+    sum = 1.0;  
+}
+
 PointSpreadFunc::PointSpreadFunc(float stdev) {
     width = stdev;
     float simple_gauss[MAX_KERNEL_RADIUS];

src/kbmod/search/PointSpreadFunc.h

@@ -26,6 +26,7 @@ namespace search {
 
 class PointSpreadFunc {
 public:
+    PointSpreadFunc();  // Create a no-op PSF.
     PointSpreadFunc(float stdev);
     PointSpreadFunc(const PointSpreadFunc& other);  // Copy constructor
     PointSpreadFunc(PointSpreadFunc&& other);       // Move constructor

src/kbmod/search/bindings.cpp

@@ -55,6 +55,7 @@ PYBIND11_MODULE(search, m) {
                                        2, {m.getDim(), m.getDim()},
                                        {sizeof(float) * m.getDim(), sizeof(float)});
             })
+            .def(py::init<>())
             .def(py::init<float>())
             .def(py::init<py::array_t<float>>())
             .def(py::init<pf &>())
@@ -100,6 +101,7 @@ PYBIND11_MODULE(search, m) {
             .def("create_stamp", &ri::createStamp)
             .def("set_pixel", &ri::setPixel, "Set the value of a given pixel.")
             .def("add_pixel", &ri::addToPixel, "Add to the value of a given pixel.")
+            .def("add_pixel_interp", &ri::addPixelInterp, "Add to the interpolated value of a given pixel.")
             .def("apply_mask", &ri::applyMask)
             .def("grow_mask", &ri::growMask)
             .def("pixel_has_data", &ri::pixelHasData,
@@ -139,7 +141,6 @@ PYBIND11_MODULE(search, m) {
             .def(py::init<std::string, int, int, double, float, float, pf &, int>())
             .def("set_psf", &li::setPSF, "Sets the PSF object.")
             .def("get_psf", &li::getPSF, "Returns the PSF object.")
-            .def("get_psfsq", &li::getPSFSQ)
             .def("apply_mask_flags", &li::applyMaskFlags)
             .def("apply_mask_threshold", &li::applyMaskThreshold)
             .def("sub_template", &li::subtractTemplate)
@@ -150,8 +151,8 @@ PYBIND11_MODULE(search, m) {
             .def("set_science", &li::setScience)
             .def("set_mask", &li::setMask)
             .def("set_variance", &li::setVariance)
-            .def("convolve_psf", &li::convolvePSF)
-            .def("add_object", &li::addObject)
+            .def("convolve_psf", &li::convolvePSF, "Convolve each layer with the object's PSF.")
+            .def("convolve_given_psf", &li::convolveGivenPSF, "Convolve each layer with a given PSF.")
             .def("grow_mask", &li::growMask)
             .def("get_name", &li::getName, "Returns the name of the layered image.")
             .def("get_width", &li::getWidth, "Returns the image's width in pixels.")

tests/regression_test.py

@@ -12,6 +12,7 @@
 import numpy as np
 from astropy.io import fits
 
+from kbmod.fake_data_creator import add_fake_object
 from kbmod.file_utils import *
 from kbmod.run_search import run_search
 from kbmod.search import *
@@ -213,7 +214,7 @@ def make_fake_image_stack(times, trjs, psf_vals):
         for trj in trjs:
             px = trj.x + time * trj.x_v + 0.5
             py = trj.y + time * trj.y_v + 0.5
-            img.add_object(px, py, trj.flux)
+            add_fake_object(img, px, py, trj.flux, p)
 
         imlist.append(img)
     stack = image_stack(imlist)

tests/test_analysis_utils.py

@@ -1,6 +1,7 @@
 import unittest
 
 from kbmod.analysis_utils import *
+from kbmod.fake_data_creator import add_fake_object
 from kbmod.result_list import *
 from kbmod.search import *
 
@@ -179,10 +180,12 @@ def test_apply_stamp_filter(self):
 
         for i in range(self.img_count):
             time = i / self.img_count
-            self.imlist[i].add_object(
+            add_fake_object(
+                self.imlist[i],
                 self.start_x + time * self.x_vel + 0.5,
                 self.start_y + time * self.y_vel + 0.5,
                 self.object_flux,
+                self.p,
             )
 
         stack = image_stack(self.imlist)
@@ -228,10 +231,12 @@ def test_apply_stamp_filter_2(self):
 
         for i in range(self.img_count):
             time = i / self.img_count
-            self.imlist[i].add_object(
+            add_fake_object(
+                self.imlist[i],
                 self.start_x + time * self.x_vel,
                 self.start_y + time * self.y_vel,
                 self.object_flux,
+                self.p,
             )
 
         stack = image_stack(self.imlist)
@@ -346,7 +351,7 @@ def test_load_and_filter_results_lh(self):
 
             # Add the objects.
             for j, trj in enumerate(trjs):
-                im.add_object(trj.x, trj.y, fluxes[j])
+                add_fake_object(im, trj.x, trj.y, fluxes[j], self.p)
 
             # Append the image.
             imlist.append(im)

tests/test_bilinear_interp.py

@@ -2,6 +2,7 @@
 
 import numpy
 
+from kbmod.fake_data_creator import add_fake_object
 import kbmod.search as kb
 
 
@@ -12,7 +13,7 @@ def setUp(self):
         self.images = []
         for c in range(self.im_count):
             im = kb.layered_image(str(c), 10, 10, 0.0, 1.0, c, p)
-            im.add_object(2 + c * 0.5 + 0.5, 2 + c * 0.5 + 0.5, 1)
+            add_fake_object(im, 2 + c * 0.5 + 0.5, 2 + c * 0.5 + 0.5, 1, p)
             self.images.append(im)
 
     def test_pixels(self):

tests/test_image_stack.py

@@ -1,6 +1,7 @@
 import tempfile
 import unittest
 
+from kbmod.fake_data_creator import add_fake_object
 from kbmod.search import *
 
 
@@ -150,7 +151,7 @@ def test_different_psfs(self):
         last_val = -100.0
         for i in range(self.num_images):
             img = self.im_stack.get_single_image(i)
-            img.add_object(10, 20, 500.0)
+            add_fake_object(img, 10, 20, 500.0, self.p[i])
 
             sci = img.get_science()
             pix_val = sci.get_pixel(10, 20)