Merge branch 'main' into fake_data

src/kbmod/fake_data_creator.py

@@ -14,6 +14,7 @@
 from kbmod.configuration import SearchConfiguration
 from kbmod.file_utils import *
 from kbmod.search import *
+from kbmod.wcs_utils import append_wcs_to_hdu_header, make_fake_wcs_info
 from kbmod.work_unit import WorkUnit
 
 
@@ -204,22 +205,9 @@ def save_fake_data_to_dir(self, data_dir):
             img.save_layers(data_dir + "/")
 
             # Open the file and insert fake WCS data.
+            wcs_dict = make_fake_wcs_info(200.6145, -7.7888, 2000, 4000)
             hdul = fits.open(filename)
-            hdul[1].header["WCSAXES"] = 2
-            hdul[1].header["CTYPE1"] = "RA---TAN-SIP"
-            hdul[1].header["CTYPE2"] = "DEC--TAN-SIP"
-            hdul[1].header["CRVAL1"] = 200.614997245422
-            hdul[1].header["CRVAL2"] = -7.78878863332778
-            hdul[1].header["CRPIX1"] = 1033.934327
-            hdul[1].header["CRPIX2"] = 2043.548284
-            hdul[1].header["CD1_1"] = -1.13926485986789e-07
-            hdul[1].header["CD1_2"] = 7.31839748843125e-05
-            hdul[1].header["CD2_1"] = -7.30064978350695e-05
-            hdul[1].header["CD2_2"] = -1.27520156332774e-07
-            hdul[1].header["CTYPE1A"] = "LINEAR  "
-            hdul[1].header["CTYPE2A"] = "LINEAR  "
-            hdul[1].header["CUNIT1A"] = "PIXEL   "
-            hdul[1].header["CUNIT2A"] = "PIXEL   "
+            append_wcs_to_hdu_header(wcs_dict, hdul[1].header)
             hdul.writeto(filename, overwrite=True)
             hdul.close()
 

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
 

src/kbmod/wcs_utils.py

@@ -344,15 +344,19 @@ def append_wcs_to_hdu_header(wcs, header):
 
     Parameters
     ----------
-    wcs : `astropy.wcs.WCS`
-        The WCS to use.
+    wcs : `astropy.wcs.WCS` or `dict`
+        The WCS to use or a dictionary with the necessary information.
     header : `astropy.io.fits.Header`
         The header to which to append the data.
     """
     if wcs is not None:
-        wcs_header = wcs.to_header()
-        for key in wcs_header:
-            header[key] = wcs_header[key]
+        if type(wcs) is dict:
+            wcs_map = wcs
+        else:
+            wcs_map = wcs.to_header()
+
+        for key in wcs_map:
+            header[key] = wcs_map[key]
 
 
 def wcs_to_dict(wcs):
@@ -374,3 +378,46 @@ def wcs_to_dict(wcs):
         for key in wcs_header:
             result[key] = wcs_header[key]
     return result
+
+
+def make_fake_wcs_info(center_ra, center_dec, height, width, deg_per_pixel=None):
+    """Create a fake WCS given basic information. This is not a realistic
+    WCS in terms of astronomy, but can provide a place holder for many tests.
+
+    Parameters
+    ----------
+    center_ra : `float`
+        The center of the pointing in RA (degrees)
+    center_dev : `float`
+        The center of the pointing in DEC (degrees)
+    height : `int`
+        The image height (pixels).
+    width : `int`
+        The image width (pixels).
+    deg_per_pixel : `float`, optional
+        The angular resolution of a pixel (degrees).
+
+    Returns
+    -------
+    wcs_dict : `dict`
+        A dictionary with the information needed for a WCS.
+    """
+    wcs_dict = {
+        "WCSAXES": 2,
+        "CTYPE1": "RA---TAN-SIP",
+        "CTYPE2": "DEC--TAN-SIP",
+        "CRVAL1": center_ra,
+        "CRVAL2": center_dec,
+        "CRPIX1": height / 2.0,
+        "CRPIX2": width / 2.0,
+        "CTYPE1A": "LINEAR  ",
+        "CTYPE2A": "LINEAR  ",
+        "CUNIT1A": "PIXEL   ",
+        "CUNIT2A": "PIXEL   ",
+    }
+
+    if deg_per_pixel is not None:
+        wcs_dict["CDELT1"] = deg_per_pixel
+        wcs_dict["CDELT2"] = deg_per_pixel
+
+    return wcs_dict