From 91fc64147cf940a4b5d2388610ab694aaac6addb Mon Sep 17 00:00:00 2001
From: Jeremy Kubica <104161096+jeremykubica@users.noreply.github.com>
Date: Fri, 12 Jan 2024 09:58:06 -0500
Subject: [PATCH 1/9] Add some comments, checks, and tests

---
 src/kbmod/search/layered_image.cpp           |  6 ++--
 src/kbmod/search/layered_image.h             |  8 ++---
 src/kbmod/search/pydocs/layered_image_docs.h | 22 +++++++++++--
 tests/test_layered_image.py                  | 34 ++++++++++++++++++--
 4 files changed, 59 insertions(+), 11 deletions(-)

diff --git a/src/kbmod/search/layered_image.cpp b/src/kbmod/search/layered_image.cpp
index efe335fdd..c5f93afca 100644
--- a/src/kbmod/search/layered_image.cpp
+++ b/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,8 +251,8 @@ 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) {
-            result_arr[p] = 1.0 / var_pix;
+        if (var_pix != NO_DATA && var_pix != 0.0) {
+            result_arr[p] = sci_array[p] / var_pix;
         } else {
             result_arr[p] = NO_DATA;
         }
diff --git a/src/kbmod/search/layered_image.h b/src/kbmod/search/layered_image.h
index 11fe159ee..b2e11e421 100644
--- a/src/kbmod/search/layered_image.h
+++ b/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 {
diff --git a/src/kbmod/search/pydocs/layered_image_docs.h b/src/kbmod/search/pydocs/layered_image_docs.h
index cd54decd4..3dd8406f3 100644
--- a/src/kbmod/search/pydocs/layered_image_docs.h
+++ b/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
 
diff --git a/tests/test_layered_image.py b/tests/test_layered_image.py
index bfc45c368..2c739affa 100644
--- a/tests/test_layered_image.py
+++ b/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(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 = (x != 1 or y == 0 or y > 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()

From 74894d662463a51fce69465d352e5063666bd8c1 Mon Sep 17 00:00:00 2001
From: Jeremy Kubica <104161096+jeremykubica@users.noreply.github.com>
Date: Fri, 12 Jan 2024 09:59:49 -0500
Subject: [PATCH 2/9] Apply clang format

---
 src/kbmod/search/bindings.cpp          |  1 -
 src/kbmod/search/layered_image.cpp     | 30 ++++++------
 src/kbmod/search/psi_phi_array.cpp     | 24 +++++-----
 src/kbmod/search/psi_phi_array_utils.h |  2 +-
 src/kbmod/search/raw_image.cpp         | 64 +++++++++++---------------
 5 files changed, 55 insertions(+), 66 deletions(-)

diff --git a/src/kbmod/search/bindings.cpp b/src/kbmod/search/bindings.cpp
index f0d2d1b04..4175c9639 100644
--- a/src/kbmod/search/bindings.cpp
+++ b/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);
diff --git a/src/kbmod/search/layered_image.cpp b/src/kbmod/search/layered_image.cpp
index c5f93afca..95d4cb544 100644
--- a/src/kbmod/search/layered_image.cpp
+++ b/src/kbmod/search/layered_image.cpp
@@ -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,
diff --git a/src/kbmod/search/psi_phi_array.cpp b/src/kbmod/search/psi_phi_array.cpp
index 866fe258f..3f2656680 100644
--- a/src/kbmod/search/psi_phi_array.cpp
+++ b/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,7 +280,9 @@ 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);
     }
@@ -289,7 +290,8 @@ void fill_psi_phi_array(PsiPhiArray& result_data, int num_bytes, const std::vect
 #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) {
diff --git a/src/kbmod/search/psi_phi_array_utils.h b/src/kbmod/search/psi_phi_array_utils.h
index d06a4483c..c363f65d4 100644
--- a/src/kbmod/search/psi_phi_array_utils.h
+++ b/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 */
 
diff --git a/src/kbmod/search/raw_image.cpp b/src/kbmod/search/raw_image.cpp
index 1dcb13214..a2e541355 100644
--- a/src/kbmod/search/raw_image.cpp
+++ b/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
 

From 2ca940080e15840791970b19176ea2896146171d Mon Sep 17 00:00:00 2001
From: Jeremy Kubica <104161096+jeremykubica@users.noreply.github.com>
Date: Fri, 12 Jan 2024 10:30:56 -0500
Subject: [PATCH 3/9] Fix bad copy and paste

---
 src/kbmod/search/layered_image.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/kbmod/search/layered_image.cpp b/src/kbmod/search/layered_image.cpp
index 95d4cb544..24f961e9d 100644
--- a/src/kbmod/search/layered_image.cpp
+++ b/src/kbmod/search/layered_image.cpp
@@ -252,7 +252,7 @@ RawImage LayeredImage::generate_phi_image() {
     for (int p = 0; p < num_pixels; ++p) {
         float var_pix = var_array[p];
         if (var_pix != NO_DATA && var_pix != 0.0) {
-            result_arr[p] = sci_array[p] / var_pix;
+            result_arr[p] = 1.0 / var_pix;
         } else {
             result_arr[p] = NO_DATA;
         }

From e392a1e5ab3583def771f6408d29716078ed826e Mon Sep 17 00:00:00 2001
From: Jeremy Kubica <104161096+jeremykubica@users.noreply.github.com>
Date: Fri, 12 Jan 2024 10:37:38 -0500
Subject: [PATCH 4/9] Fix some typos

---
 tests/test_layered_image.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/tests/test_layered_image.py b/tests/test_layered_image.py
index 2c739affa..4185f9e03 100644
--- a/tests/test_layered_image.py
+++ b/tests/test_layered_image.py
@@ -51,10 +51,10 @@ def test_create(self):
         # 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.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(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)
@@ -325,13 +325,13 @@ def test_psi_and_phi_image(self):
 
         for y in range(5):
             for x in range(6):
-                has_data = (x != 1 or y == 0 or y > 2)
+                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 has_data:
                     self.assertAlmostEqual(psi.get_pixel(y, x), x / (y + 1))
                     self.assertAlmostEqual(phi.get_pixel(y, x), 1.0 / (y + 1))
-                else
+                else:
                     self.assertEqual(psi.get_pixel(y, x), KB_NO_DATA)
                     self.assertEqual(phi.get_pixel(y, x), KB_NO_DATA)
 

From 2cc0848fb93e417ab4bd4179c798c8d0eb98ef6e Mon Sep 17 00:00:00 2001
From: Jeremy Kubica <104161096+jeremykubica@users.noreply.github.com>
Date: Fri, 12 Jan 2024 10:42:23 -0500
Subject: [PATCH 5/9] Fix formatting error

---
 tests/test_layered_image.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/test_layered_image.py b/tests/test_layered_image.py
index 4185f9e03..7ac944844 100644
--- a/tests/test_layered_image.py
+++ b/tests/test_layered_image.py
@@ -325,7 +325,7 @@ def test_psi_and_phi_image(self):
 
         for y in range(5):
             for x in range(6):
-                has_data = (y != 3 or x == 0 or x > 2)
+                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 has_data:

From 049a5448313714640b58eb1698d68d4c0fb74d4c Mon Sep 17 00:00:00 2001
From: Jeremy Kubica <104161096+jeremykubica@users.noreply.github.com>
Date: Fri, 12 Jan 2024 12:14:48 -0500
Subject: [PATCH 6/9] Combine fake WCS generation logic

---
 src/kbmod/fake_data_creator.py | 18 ++---------
 src/kbmod/wcs_utils.py         | 57 ++++++++++++++++++++++++++++++---
 tests/test_regression_test.py  | 21 ++++++------
 tests/test_wcs_utils.py        | 58 ++++++++++++++++++++++++++++------
 tests/test_work_unit.py        | 19 ++---------
 5 files changed, 115 insertions(+), 58 deletions(-)

diff --git a/src/kbmod/fake_data_creator.py b/src/kbmod/fake_data_creator.py
index 7f15b5a45..86e07c7db 100644
--- a/src/kbmod/fake_data_creator.py
+++ b/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()
 
diff --git a/src/kbmod/wcs_utils.py b/src/kbmod/wcs_utils.py
index 0a59d7ab6..2c8687c81 100644
--- a/src/kbmod/wcs_utils.py
+++ b/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
diff --git a/tests/test_regression_test.py b/tests/test_regression_test.py
index 5b8ce73bb..3794b17f8 100644
--- a/tests/test_regression_test.py
+++ b/tests/test_regression_test.py
@@ -17,6 +17,7 @@
 from kbmod.run_search import SearchRunner
 from kbmod.search import *
 from kbmod.trajectory_utils import make_trajectory
+from kbmod.wcs_utils import append_wcs_to_hdu_header, make_fake_wcs_info
 
 
 def ave_trajectory_distance(trjA, trjB, times=[0.0]):
@@ -205,21 +206,19 @@ def add_wcs_header_data(full_file_name):
         The path and filename of the FITS file to modify.
     """
     hdul = fits.open(full_file_name)
-    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
+
+    # Create a fake WCS with some minimal information.
+    wcs_dict = make_fake_wcs_info(200.615, -7.789, 2068, 4088)
+    append_wcs_to_hdu_header(wcs_dict, hdul[1].header)
+
+    # Add rotational and scale parameters that will allow us to use a matching
+    # set of search angles (using KBMOD angle suggestion function).
     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   "
+
+    # Save the augmented header.
     hdul.writeto(full_file_name, overwrite=True)
     hdul.close()
 
diff --git a/tests/test_wcs_utils.py b/tests/test_wcs_utils.py
index 0d0434a33..628a701c4 100644
--- a/tests/test_wcs_utils.py
+++ b/tests/test_wcs_utils.py
@@ -45,16 +45,54 @@ def test_wcs_to_dict(self):
             self.assertAlmostEqual(new_dict[key], self.header_dict[key])
 
     def test_append_wcs_to_hdu_header(self):
-        pri = fits.PrimaryHDU()
-        self.assertFalse("CRVAL1" in pri.header)
-        self.assertFalse("CRVAL2" in pri.header)
-        self.assertFalse("CRPIX1" in pri.header)
-        self.assertFalse("CRPIX2" in pri.header)
-
-        append_wcs_to_hdu_header(self.wcs, pri.header)
-        for key in self.header_dict:
-            self.assertTrue(key in pri.header)
-            self.assertAlmostEqual(pri.header[key], self.header_dict[key])
+        for use_dictionary in [True, False]:
+            if use_dictionary:
+                wcs_info = self.header_dict
+            else:
+                wcs_info = self.wcs
+
+        # Run the dictionary and full WCS tests as separate subtests
+        with self.subTest(i=use_dictionary):
+            pri = fits.PrimaryHDU()
+            self.assertFalse("CRVAL1" in pri.header)
+            self.assertFalse("CRVAL2" in pri.header)
+            self.assertFalse("CRPIX1" in pri.header)
+            self.assertFalse("CRPIX2" in pri.header)
+
+            append_wcs_to_hdu_header(wcs_info, pri.header)
+            for key in self.header_dict:
+                self.assertTrue(key in pri.header)
+                self.assertAlmostEqual(pri.header[key], self.header_dict[key])
+
+    def test_make_fake_wcs_info(self):
+        # Test that we make the dictionary
+        wcs_dict = make_fake_wcs_info(25.0, -10.0, 200, 100, deg_per_pixel=0.01)
+        self.assertEqual(wcs_dict["WCSAXES"], 2)
+        self.assertEqual(wcs_dict["CTYPE1"], "RA---TAN-SIP")
+        self.assertEqual(wcs_dict["CTYPE2"], "DEC--TAN-SIP")
+        self.assertEqual(wcs_dict["CRVAL1"], 25.0)
+        self.assertEqual(wcs_dict["CRVAL2"], -10.0)
+        self.assertEqual(wcs_dict["CRPIX1"], 100.0)
+        self.assertEqual(wcs_dict["CRPIX2"], 50.0)
+        self.assertEqual(wcs_dict["CDELT1"], 0.01)
+        self.assertEqual(wcs_dict["CDELT2"], 0.01)
+        self.assertEqual(wcs_dict["CTYPE1A"], "LINEAR  ")
+        self.assertEqual(wcs_dict["CTYPE2A"], "LINEAR  ")
+        self.assertEqual(wcs_dict["CUNIT1A"], "PIXEL   ")
+        self.assertEqual(wcs_dict["CUNIT2A"], "PIXEL   ")
+
+        # Test the we can convert to a WCS and extra predictions.
+        test_wcs = WCS(wcs_dict)
+
+        # Center position should be at approximately (25.0, -10.0).
+        pos = test_wcs.pixel_to_world(99, 49)
+        self.assertAlmostEqual(pos.ra.degree, 25.0, delta=0.001)
+        self.assertAlmostEqual(pos.dec.degree, -10.0, delta=0.001)
+
+        # One pixel off position should be at approximately (25.01, -10.0).
+        pos = test_wcs.pixel_to_world(100, 48)
+        self.assertAlmostEqual(pos.ra.degree, 25.01, delta=0.01)
+        self.assertAlmostEqual(pos.dec.degree, -10.0, delta=0.01)
 
 
 class test_construct_wcs_tangent_projection(unittest.TestCase):
diff --git a/tests/test_work_unit.py b/tests/test_work_unit.py
index b1c4b7fda..0e1c2f324 100644
--- a/tests/test_work_unit.py
+++ b/tests/test_work_unit.py
@@ -8,6 +8,7 @@
 import warnings
 
 from kbmod.configuration import SearchConfiguration
+from kbmod.fake_data_creator import make_fake_wcs_info
 import kbmod.search as kb
 from kbmod.work_unit import hdu_to_raw_image, raw_image_to_hdu, WorkUnit
 
@@ -44,23 +45,7 @@ def setUp(self):
         self.config.set("repeated_flag_keys", None)
 
         # Create a fake WCS
-        header_dict = {
-            "WCSAXES": 2,
-            "CTYPE1": "RA---TAN-SIP",
-            "CTYPE2": "DEC--TAN-SIP",
-            "CRVAL1": 200.614997245422,
-            "CRVAL2": -7.78878863332778,
-            "CRPIX1": 1033.934327,
-            "CRPIX2": 2043.548284,
-            "CD1_1": -1.13926485986789e-07,
-            "CD1_2": 7.31839748843125e-05,
-            "CD2_1": -7.30064978350695e-05,
-            "CD2_2": -1.27520156332774e-07,
-            "CTYPE1A": "LINEAR  ",
-            "CTYPE2A": "LINEAR  ",
-            "CUNIT1A": "PIXEL   ",
-            "CUNIT2A": "PIXEL   ",
-        }
+        header_dict = make_fake_wcs_info(200.6145, -7.7888, 2000, 4000)
         self.wcs = WCS(header_dict)
         self.per_image_wcs = per_image_wcs = [
             (self.wcs if i % 2 == 0 else None) for i in range(self.num_images)

From f4145d0b374c6e69283e7ae41ddfcc69d34247aa Mon Sep 17 00:00:00 2001
From: Jeremy Kubica <104161096+jeremykubica@users.noreply.github.com>
Date: Fri, 12 Jan 2024 16:02:20 -0500
Subject: [PATCH 7/9] Address PR comments

---
 src/kbmod/search/pydocs/layered_image_docs.h | 2 +-
 tests/test_layered_image.py                  | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/kbmod/search/pydocs/layered_image_docs.h b/src/kbmod/search/pydocs/layered_image_docs.h
index 3dd8406f3..72b924dd7 100644
--- a/src/kbmod/search/pydocs/layered_image_docs.h
+++ b/src/kbmod/search/pydocs/layered_image_docs.h
@@ -253,7 +253,7 @@ static const auto DOC_LayeredImage_generate_psi_image = R"doc(
 static const auto DOC_LayeredImage_generate_phi_image = R"doc(
   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,
+  apply_mask() must be called before the phi image is generated. Otherwise,
   all pixels are used.
   Convolves the resulting image with the PSF.
 
diff --git a/tests/test_layered_image.py b/tests/test_layered_image.py
index 7ac944844..da8c8fe9a 100644
--- a/tests/test_layered_image.py
+++ b/tests/test_layered_image.py
@@ -48,7 +48,7 @@ def test_create(self):
                 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.
+        # 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)

From cbcd81b16f0ed2644edec412c3790dae6dea916e Mon Sep 17 00:00:00 2001
From: Jeremy Kubica <104161096+jeremykubica@users.noreply.github.com>
Date: Fri, 12 Jan 2024 16:04:31 -0500
Subject: [PATCH 8/9] Address PR comment

---
 tests/test_wcs_utils.py | 24 ++++++++++++------------
 1 file changed, 12 insertions(+), 12 deletions(-)

diff --git a/tests/test_wcs_utils.py b/tests/test_wcs_utils.py
index 628a701c4..e12d99371 100644
--- a/tests/test_wcs_utils.py
+++ b/tests/test_wcs_utils.py
@@ -51,18 +51,18 @@ def test_append_wcs_to_hdu_header(self):
             else:
                 wcs_info = self.wcs
 
-        # Run the dictionary and full WCS tests as separate subtests
-        with self.subTest(i=use_dictionary):
-            pri = fits.PrimaryHDU()
-            self.assertFalse("CRVAL1" in pri.header)
-            self.assertFalse("CRVAL2" in pri.header)
-            self.assertFalse("CRPIX1" in pri.header)
-            self.assertFalse("CRPIX2" in pri.header)
-
-            append_wcs_to_hdu_header(wcs_info, pri.header)
-            for key in self.header_dict:
-                self.assertTrue(key in pri.header)
-                self.assertAlmostEqual(pri.header[key], self.header_dict[key])
+            # Run the dictionary and full WCS tests as separate subtests
+            with self.subTest(i=use_dictionary):
+                pri = fits.PrimaryHDU()
+                self.assertFalse("CRVAL1" in pri.header)
+                self.assertFalse("CRVAL2" in pri.header)
+                self.assertFalse("CRPIX1" in pri.header)
+                self.assertFalse("CRPIX2" in pri.header)
+
+                append_wcs_to_hdu_header(wcs_info, pri.header)
+                for key in self.header_dict:
+                    self.assertTrue(key in pri.header)
+                    self.assertAlmostEqual(pri.header[key], self.header_dict[key])
 
     def test_make_fake_wcs_info(self):
         # Test that we make the dictionary

From 44c69b75f83d66ca0f8069f5665e33c726282fee Mon Sep 17 00:00:00 2001
From: Jeremy Kubica <104161096+jeremykubica@users.noreply.github.com>
Date: Fri, 12 Jan 2024 16:09:07 -0500
Subject: [PATCH 9/9] Fix a typo

---
 tests/test_wcs_utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/test_wcs_utils.py b/tests/test_wcs_utils.py
index e12d99371..d95f67c19 100644
--- a/tests/test_wcs_utils.py
+++ b/tests/test_wcs_utils.py
@@ -81,7 +81,7 @@ def test_make_fake_wcs_info(self):
         self.assertEqual(wcs_dict["CUNIT1A"], "PIXEL   ")
         self.assertEqual(wcs_dict["CUNIT2A"], "PIXEL   ")
 
-        # Test the we can convert to a WCS and extra predictions.
+        # Test that we can convert to a WCS and perform predictions.
         test_wcs = WCS(wcs_dict)
 
         # Center position should be at approximately (25.0, -10.0).