From 43fe80900b702b859438e014af0d695bba23c741 Mon Sep 17 00:00:00 2001
From: Jeremy Kubica <104161096+jeremykubica@users.noreply.github.com>
Date: Tue, 27 Feb 2024 13:07:48 -0500
Subject: [PATCH 1/4] Add support for NaNs in kernel functions

---
 src/kbmod/search/image_kernels.cu |  8 ++---
 src/kbmod/search/kernels.cu       |  7 +++--
 src/kbmod/search/raw_image.cpp    | 29 +++++++++---------
 src/kbmod/search/raw_image.h      |  5 +++-
 tests/test_raw_image.py           | 34 ++++++++++++++++++++-
 tests/test_search.py              | 50 +++++++++++++++++++++++--------
 6 files changed, 97 insertions(+), 36 deletions(-)

diff --git a/src/kbmod/search/image_kernels.cu b/src/kbmod/search/image_kernels.cu
index ad3436bf4..4f094db01 100644
--- a/src/kbmod/search/image_kernels.cu
+++ b/src/kbmod/search/image_kernels.cu
@@ -30,13 +30,13 @@ __global__ void convolve_psf(int width, int height, float *source_img, float *re
     float sum = 0.0;
     float psf_portion = 0.0;
     float center = source_img[y * width + x];
-    if (center != NO_DATA) {
+    if ((center != NO_DATA) && !isnan(center)) {
         for (int j = -psf_radius; j <= psf_radius; j++) {
             // #pragma unroll
             for (int i = -psf_radius; i <= psf_radius; i++) {
                 if ((x + i >= 0) && (x + i < width) && (y + j >= 0) && (y + j < height)) {
                     float current_pix = source_img[(y + j) * width + (x + i)];
-                    if (current_pix != NO_DATA) {
+                    if ((current_pix != NO_DATA) && !isnan(current_pix)) {
                         float current_psf = psf[(j + psf_radius) * psf_dim + (i + psf_radius)];
                         psf_portion += current_psf;
                         sum += current_pix * current_psf;
@@ -47,8 +47,8 @@ __global__ void convolve_psf(int width, int height, float *source_img, float *re
 
         result_img[y * width + x] = (sum * psf_sum) / psf_portion;
     } else {
-        // Leave masked pixel alone (these could be replaced here with zero)
-        result_img[y * width + x] = NO_DATA;  // 0.0
+        // Leave masked and NaN pixels alone (these could be replaced here with zero)
+        result_img[y * width + x] = center;  // 0.0
     }
 }
 
diff --git a/src/kbmod/search/kernels.cu b/src/kbmod/search/kernels.cu
index a2eb45950..871aab974 100644
--- a/src/kbmod/search/kernels.cu
+++ b/src/kbmod/search/kernels.cu
@@ -179,9 +179,10 @@ extern "C" __device__ __host__ void evaluateTrajectory(PsiPhiArrayMeta psi_phi_m
         int current_x = candidate->x + int(candidate->vx * curr_time + 0.5);
         int current_y = candidate->y + int(candidate->vy * curr_time + 0.5);
 
-        // Get the Psi and Phi pixel values.
+        // Get the Psi and Phi pixel values. Skip values that are NaN or marked with the NO_DATA.
         PsiPhi pixel_vals = read_encoded_psi_phi(psi_phi_meta, psi_phi_vect, i, current_y, current_x);
-        if (pixel_vals.psi != NO_DATA && pixel_vals.phi != NO_DATA) {
+        if (pixel_vals.psi != NO_DATA && pixel_vals.phi != NO_DATA && !isnan(pixel_vals.psi) &&
+            !isnan(pixel_vals.phi)) {
             psi_sum += pixel_vals.psi;
             phi_sum += pixel_vals.phi;
             psi_array[num_seen] = pixel_vals.psi;
@@ -414,7 +415,7 @@ __global__ void deviceGetCoaddStamp(int num_images, int width, int height, float
         int img_y = current_y - params.radius + stamp_y;
         if ((img_x >= 0) && (img_x < width) && (img_y >= 0) && (img_y < height)) {
             int pixel_index = width * height * t + img_y * width + img_x;
-            if (image_vect[pixel_index] != NO_DATA) {
+            if ((image_vect[pixel_index] != NO_DATA) && !isnan(image_vect[pixel_index])) {
                 values[num_values] = image_vect[pixel_index];
                 ++num_values;
             }
diff --git a/src/kbmod/search/raw_image.cpp b/src/kbmod/search/raw_image.cpp
index 4334061c8..9d21b300a 100644
--- a/src/kbmod/search/raw_image.cpp
+++ b/src/kbmod/search/raw_image.cpp
@@ -173,9 +173,9 @@ void RawImage::convolve_cpu(PSF& psf) {
 
     for (int y = 0; y < height; ++y) {
         for (int x = 0; x < width; ++x) {
-            // Pixels with NO_DATA remain NO_DATA.
-            if (image(y, x) == NO_DATA) {
-                result(y, x) = NO_DATA;
+            // Pixels with NO_DATA or NaN remain NO_DATA or NaN.
+            if ((image(y, x) == NO_DATA) || std::isnan(image(y, x))) {
+                result(y, x) = image(y, x);
                 continue;
             }
 
@@ -186,7 +186,7 @@ void RawImage::convolve_cpu(PSF& psf) {
                     if ((x + i >= 0) && (x + i < width) && (y + j >= 0) && (y + j < height)) {
                         float current_pixel = image(y + j, x + i);
                         // note that convention for index access is flipped for PSF
-                        if (current_pixel != NO_DATA) {
+                        if ((current_pixel != NO_DATA) && !std::isnan(current_pixel)) {
                             float current_psf = psf.get_value(i + psf_rad, j + psf_rad);
                             psf_portion += current_psf;
                             sum += current_pixel * current_psf;
@@ -354,11 +354,8 @@ RawImage create_median_image(const std::vector<RawImage>& images) {
             int num_unmasked = 0;
             for (auto& img : images) {
                 // Only used the unmasked array.
-                // we have a get_pixel and pixel_has_data, but we don't use them
-                // here in the original code, so I go to get_image()() too...
-                if ((img.get_image()(y, x) != NO_DATA) &&
-                    (!std::isnan(img.get_image()(y, x)))) {  // why are we suddenly checking nans?!
-                    pix_array[num_unmasked] = img.get_image()(y, x);
+                if (img.pixel_has_data({y, x})) {
+                    pix_array[num_unmasked] = img.get_pixel({y, x});
                     num_unmasked += 1;
                 }
             }
@@ -395,7 +392,13 @@ RawImage create_summed_image(const std::vector<RawImage>& images) {
 
     Image result = Image::Zero(height, width);
     for (auto& img : images) {
-        result += (img.get_image().array() == NO_DATA).select(0, img.get_image());
+        for (int y = 0; y < height; ++y) {
+            for (int x = 0; x < width; ++x) {
+                if (img.pixel_has_data({y, x})) {
+                    result(y, x) += img.get_pixel({y, x});
+                }
+            }
+        }
     }
     return RawImage(result);
 }
@@ -416,11 +419,9 @@ RawImage create_mean_image(const std::vector<RawImage>& images) {
             float sum = 0.0;
             float count = 0.0;
             for (auto& img : images) {
-                // we have a get_pixel and pixel_has_data, but we don't use them
-                // here in the original code, so I go to get_image()() too...
-                if ((img.get_image()(y, x) != NO_DATA) && (!std::isnan(img.get_image()(y, x)))) {
+                if (img.pixel_has_data({y, x})) {
                     count += 1.0;
-                    sum += img.get_image()(y, x);
+                    sum += img.get_pixel({y, x});
                 }
             }
 
diff --git a/src/kbmod/search/raw_image.h b/src/kbmod/search/raw_image.h
index 5c24a4382..8f8898ae1 100644
--- a/src/kbmod/search/raw_image.h
+++ b/src/kbmod/search/raw_image.h
@@ -51,7 +51,10 @@ class RawImage {
 
     inline float get_pixel(const Index& idx) const { return contains(idx) ? image(idx.i, idx.j) : NO_DATA; }
 
-    inline bool pixel_has_data(const Index& idx) const { return get_pixel(idx) != NO_DATA ? true : false; }
+    inline bool pixel_has_data(const Index& idx) const {
+        float pix_val = get_pixel(idx);
+        return ((pix_val != NO_DATA) && !std::isnan(pix_val)) ? true : false;
+    }
 
     inline void set_pixel(const Index& idx, float value) {
         if (!contains(idx)) throw std::runtime_error("Index out of bounds!");
diff --git a/tests/test_raw_image.py b/tests/test_raw_image.py
index 2c7a9bc7a..5ec57d970 100644
--- a/tests/test_raw_image.py
+++ b/tests/test_raw_image.py
@@ -1,3 +1,4 @@
+import math
 import numpy as np
 import os
 import tempfile
@@ -246,7 +247,38 @@ def test_convolve_psf_mask_cpu(self):
     @unittest.skipIf(not HAS_GPU, "Skipping test (no GPU detected)")
     def test_convolve_psf_mask_gpu(self):
         """Test masked convolution with a identity kernel on GPU"""
-        self.convolve_psf_mask("CPU")
+        self.convolve_psf_mask("GPU")
+
+    def convolve_psf_nan(self, device):
+        p = PSF(1.0)
+
+        # Mask out three pixels.
+        img = RawImage(self.array)
+        img.set_pixel(0, 3, math.nan)
+        img.set_pixel(5, 6, np.nan)
+        img.set_pixel(5, 7, np.nan)
+
+        if device.upper() == "CPU":
+            img.convolve_cpu(p)
+        elif device.upper() == "GPU":
+            img.convolve_gpu(p)
+        else:
+            raise ValueError(f"Unknown device. Expected GPU or CPU got {device}")
+
+        # Check that the same pixels are NaN (we ignore those pixels).
+        for y in range(self.height):
+            for x in range(self.width):
+                if (y == 5 and x == 6) or (y == 0 and x == 3) or (y == 5 and x == 7):
+                    self.assertTrue(math.isnan(img.get_pixel(y, x)))
+                else:
+                    self.assertFalse(math.isnan(img.get_pixel(y, x)))
+
+    def test_convolve_psf_nan_cpu(self):
+        self.convolve_psf_nan("CPU")
+
+    @unittest.skipIf(not HAS_GPU, "Skipping test (no GPU detected)")
+    def test_convolve_psf_nan_gpu(self):
+        self.convolve_psf_nan("GPU")
 
     # confused, sort out later
     def convolve_psf_average(self, device):
diff --git a/tests/test_search.py b/tests/test_search.py
index 19b2d1965..67c899156 100644
--- a/tests/test_search.py
+++ b/tests/test_search.py
@@ -483,24 +483,30 @@ def test_coadd_cpu_simple(self):
         # Create an image set with three images.
         imlist = []
         for i in range(3):
-            time = i
-            im = make_fake_layered_image(3, 3, 0.1, 0.01, i, self.p, seed=i)
+            # Set the first column to i, the middle column to i + 1, and the last column to 0.5.
+            # In the very first image, make pixel (0, 2) NaN
+            sci = np.array([[i, i + 1, 0.5], [i, i + 1, 0.5], [i, i + 1, 0.5]])
+            if i == 0:
+                sci[0][2] = np.nan
 
-            # Overwrite the middle row to be i + 1.
-            sci = im.get_science()
-            for x in range(3):
-                sci.set_pixel(x, 1, i + 1)
+            var = np.full((3, 3), 0.1)
 
-            # Mask out the row's first pixel twice and second pixel once.
-            mask = im.get_mask()
+            # Mask out the column's first pixel twice and second pixel once.
+            msk = np.zeros((3, 3))
             if i == 0:
-                mask.set_pixel(0, 1, 1)
-                mask.set_pixel(1, 1, 1)
+                msk[0][1] = 1
+                msk[1][1] = 1
             if i == 1:
-                mask.set_pixel(0, 1, 1)
-            im.apply_mask(1)
+                msk[0][1] = 1
 
-            imlist.append(im)
+            img = LayeredImage(
+                RawImage(sci.astype(np.float32)),
+                RawImage(var.astype(np.float32)),
+                RawImage(msk.astype(np.float32)),
+                self.p,
+            )
+            img.apply_mask(1)
+            imlist.append(img)
         stack = ImageStack(imlist)
         search = StackSearch(stack)
         all_valid = [True, True, True]  # convenience array
@@ -516,23 +522,41 @@ def test_coadd_cpu_simple(self):
         # Test summed.
         params.stamp_type = StampType.STAMP_SUM
         stamps = StampCreator.get_coadded_stamps(search.get_imagestack(), [trj], [all_valid], params, False)
+        self.assertAlmostEqual(stamps[0].get_pixel(0, 0), 3.0)
+        self.assertAlmostEqual(stamps[0].get_pixel(1, 0), 3.0)
+        self.assertAlmostEqual(stamps[0].get_pixel(2, 0), 3.0)
         self.assertAlmostEqual(stamps[0].get_pixel(0, 1), 3.0)
         self.assertAlmostEqual(stamps[0].get_pixel(1, 1), 5.0)
         self.assertAlmostEqual(stamps[0].get_pixel(2, 1), 6.0)
+        self.assertAlmostEqual(stamps[0].get_pixel(0, 2), 1.0)
+        self.assertAlmostEqual(stamps[0].get_pixel(1, 2), 1.5)
+        self.assertAlmostEqual(stamps[0].get_pixel(2, 2), 1.5)
 
         # Test mean.
         params.stamp_type = StampType.STAMP_MEAN
         stamps = StampCreator.get_coadded_stamps(search.get_imagestack(), [trj], [all_valid], params, False)
+        self.assertAlmostEqual(stamps[0].get_pixel(0, 0), 1.0)
+        self.assertAlmostEqual(stamps[0].get_pixel(1, 0), 1.0)
+        self.assertAlmostEqual(stamps[0].get_pixel(2, 0), 1.0)
         self.assertAlmostEqual(stamps[0].get_pixel(0, 1), 3.0)
         self.assertAlmostEqual(stamps[0].get_pixel(1, 1), 2.5)
         self.assertAlmostEqual(stamps[0].get_pixel(2, 1), 2.0)
+        self.assertAlmostEqual(stamps[0].get_pixel(0, 2), 0.5)
+        self.assertAlmostEqual(stamps[0].get_pixel(1, 2), 0.5)
+        self.assertAlmostEqual(stamps[0].get_pixel(2, 2), 0.5)
 
         # Test median.
         params.stamp_type = StampType.STAMP_MEDIAN
         stamps = StampCreator.get_coadded_stamps(search.get_imagestack(), [trj], [all_valid], params, False)
+        self.assertAlmostEqual(stamps[0].get_pixel(0, 0), 1.0)
+        self.assertAlmostEqual(stamps[0].get_pixel(1, 0), 1.0)
+        self.assertAlmostEqual(stamps[0].get_pixel(2, 0), 1.0)
         self.assertAlmostEqual(stamps[0].get_pixel(0, 1), 3.0)
         self.assertAlmostEqual(stamps[0].get_pixel(1, 1), 2.5)
         self.assertAlmostEqual(stamps[0].get_pixel(2, 1), 2.0)
+        self.assertAlmostEqual(stamps[0].get_pixel(0, 2), 0.5)
+        self.assertAlmostEqual(stamps[0].get_pixel(1, 2), 0.5)
+        self.assertAlmostEqual(stamps[0].get_pixel(2, 2), 0.5)
 
     @unittest.skipIf(not HAS_GPU, "Skipping test (no GPU detected)")
     def test_coadd_gpu_simple(self):

From db4b41feeb711c5f56b8e2ea6fa44187e0ac3964 Mon Sep 17 00:00:00 2001
From: Jeremy Kubica <104161096+jeremykubica@users.noreply.github.com>
Date: Fri, 1 Mar 2024 09:21:11 -0500
Subject: [PATCH 2/4] Refactor change to use a inlined testing function

---
 src/kbmod/search/bindings.cpp      |  4 +--
 src/kbmod/search/common.h          |  5 +++
 src/kbmod/search/image_kernels.cu  |  7 +++--
 src/kbmod/search/kernels.cu        |  9 +++---
 src/kbmod/search/layered_image.cpp |  8 +++--
 src/kbmod/search/raw_image.cpp     | 27 ++++++++--------
 src/kbmod/search/raw_image.h       |  5 ++-
 tests/test_layered_image.py        | 49 ++++++++++++++++++++----------
 tests/test_raw_image.py            | 19 ++++++++++++
 9 files changed, 90 insertions(+), 43 deletions(-)

diff --git a/src/kbmod/search/bindings.cpp b/src/kbmod/search/bindings.cpp
index 3362a32e2..244660437 100644
--- a/src/kbmod/search/bindings.cpp
+++ b/src/kbmod/search/bindings.cpp
@@ -6,7 +6,6 @@
 
 namespace py = pybind11;
 
-#include "logging.h"
 #include "common.h"
 #include "geom.h"
 
@@ -28,7 +27,6 @@ PYBIND11_MODULE(search, m) {
             .value("STAMP_MEAN", search::StampType::STAMP_MEAN)
             .value("STAMP_MEDIAN", search::StampType::STAMP_MEDIAN)
             .export_values();
-    logging::logging_bindings(m);
     indexing::index_bindings(m);
     indexing::point_bindings(m);
     indexing::rectangle_bindings(m);
@@ -44,6 +42,8 @@ PYBIND11_MODULE(search, m) {
     search::stamp_parameters_bindings(m);
     search::psi_phi_array_binding(m);
     search::debug_timer_binding(m);
+    // Helper function from common.h
+    m.def("pixel_value_valid", &search::pixel_value_valid);
     // Functions from raw_image.cpp
     m.def("create_median_image", &search::create_median_image);
     m.def("create_summed_image", &search::create_summed_image);
diff --git a/src/kbmod/search/common.h b/src/kbmod/search/common.h
index 285d51876..2625be590 100644
--- a/src/kbmod/search/common.h
+++ b/src/kbmod/search/common.h
@@ -26,6 +26,11 @@ constexpr float NO_DATA = -9999.0;
 
 enum StampType { STAMP_SUM = 0, STAMP_MEAN, STAMP_MEDIAN };
 
+// A helper function to check that a pixel value is valid.
+inline bool pixel_value_valid(float value) {
+    return ((value != NO_DATA) && !std::isnan(value));
+}
+
 /*
  * Data structure to represent an objects trajectory
  * through a stack of images
diff --git a/src/kbmod/search/image_kernels.cu b/src/kbmod/search/image_kernels.cu
index 4f094db01..3a53a32fb 100644
--- a/src/kbmod/search/image_kernels.cu
+++ b/src/kbmod/search/image_kernels.cu
@@ -16,6 +16,9 @@
 
 namespace search {
 
+// This is defined in kernels.cu.
+__host__ __device__ bool device_pixel_valid(float value);
+
 /*
  * Device kernel that convolves the provided image with the psf
  */
@@ -30,13 +33,13 @@ __global__ void convolve_psf(int width, int height, float *source_img, float *re
     float sum = 0.0;
     float psf_portion = 0.0;
     float center = source_img[y * width + x];
-    if ((center != NO_DATA) && !isnan(center)) {
+    if (device_pixel_valid(center)) {
         for (int j = -psf_radius; j <= psf_radius; j++) {
             // #pragma unroll
             for (int i = -psf_radius; i <= psf_radius; i++) {
                 if ((x + i >= 0) && (x + i < width) && (y + j >= 0) && (y + j < height)) {
                     float current_pix = source_img[(y + j) * width + (x + i)];
-                    if ((current_pix != NO_DATA) && !isnan(current_pix)) {
+                    if (device_pixel_valid(current_pix)) {
                         float current_psf = psf[(j + psf_radius) * psf_dim + (i + psf_radius)];
                         psf_portion += current_psf;
                         sum += current_pix * current_psf;
diff --git a/src/kbmod/search/kernels.cu b/src/kbmod/search/kernels.cu
index 871aab974..5226ddadd 100644
--- a/src/kbmod/search/kernels.cu
+++ b/src/kbmod/search/kernels.cu
@@ -23,6 +23,8 @@
 
 namespace search {
 
+__host__ __device__ bool device_pixel_valid(float value) { return ((value != NO_DATA) && !isnan(value)); }
+
 extern "C" void device_allocate_psi_phi_arrays(PsiPhiArray *data) {
     if (data == nullptr) {
         throw std::runtime_error("No data given.");
@@ -179,10 +181,9 @@ extern "C" __device__ __host__ void evaluateTrajectory(PsiPhiArrayMeta psi_phi_m
         int current_x = candidate->x + int(candidate->vx * curr_time + 0.5);
         int current_y = candidate->y + int(candidate->vy * curr_time + 0.5);
 
-        // Get the Psi and Phi pixel values. Skip values that are NaN or marked with the NO_DATA.
+        // Get the Psi and Phi pixel values. Skip invalid values, such as those marked NaN or NO_DATA.
         PsiPhi pixel_vals = read_encoded_psi_phi(psi_phi_meta, psi_phi_vect, i, current_y, current_x);
-        if (pixel_vals.psi != NO_DATA && pixel_vals.phi != NO_DATA && !isnan(pixel_vals.psi) &&
-            !isnan(pixel_vals.phi)) {
+        if (device_pixel_valid(pixel_vals.psi) && device_pixel_valid(pixel_vals.phi)) {
             psi_sum += pixel_vals.psi;
             phi_sum += pixel_vals.phi;
             psi_array[num_seen] = pixel_vals.psi;
@@ -415,7 +416,7 @@ __global__ void deviceGetCoaddStamp(int num_images, int width, int height, float
         int img_y = current_y - params.radius + stamp_y;
         if ((img_x >= 0) && (img_x < width) && (img_y >= 0) && (img_y < height)) {
             int pixel_index = width * height * t + img_y * width + img_x;
-            if ((image_vect[pixel_index] != NO_DATA) && !isnan(image_vect[pixel_index])) {
+            if (device_pixel_valid(image_vect[pixel_index])) {
                 values[num_values] = image_vect[pixel_index];
                 ++num_values;
             }
diff --git a/src/kbmod/search/layered_image.cpp b/src/kbmod/search/layered_image.cpp
index efe20474d..b6d84f74b 100644
--- a/src/kbmod/search/layered_image.cpp
+++ b/src/kbmod/search/layered_image.cpp
@@ -118,8 +118,10 @@ void LayeredImage::subtract_template(RawImage& sub_template) {
     float* sci_pixels = science.data();
     float* tem_pixels = sub_template.data();
     for (unsigned i = 0; i < num_pixels; ++i) {
-        if ((sci_pixels[i] != NO_DATA) && (tem_pixels[i] != NO_DATA)) {
+        if (pixel_value_valid(sci_pixels[i]) && pixel_value_valid(tem_pixels[i])) {
             sci_pixels[i] -= tem_pixels[i];
+        } else {
+            sci_pixels[i] = NO_DATA;
         }
     }
 }
@@ -154,7 +156,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 && var_pix != 0.0 && sci_array[p] != NO_DATA) {
+        if (pixel_value_valid(var_pix) && var_pix != 0.0 && pixel_value_valid(sci_array[p])) {
             result_arr[p] = sci_array[p] / var_pix;
         } else {
             result_arr[p] = NO_DATA;
@@ -176,7 +178,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 && var_pix != 0.0) {
+        if (pixel_value_valid(var_pix) && var_pix != 0.0) {
             result_arr[p] = 1.0 / var_pix;
         } else {
             result_arr[p] = NO_DATA;
diff --git a/src/kbmod/search/raw_image.cpp b/src/kbmod/search/raw_image.cpp
index 9d21b300a..dee177ada 100644
--- a/src/kbmod/search/raw_image.cpp
+++ b/src/kbmod/search/raw_image.cpp
@@ -153,7 +153,7 @@ std::array<float, 2> RawImage::compute_bounds() const {
     float max_val = -FLT_MAX;
 
     for (auto elem : image.reshaped())
-        if (elem != NO_DATA) {
+        if (pixel_value_valid(elem)) {
             min_val = std::min(min_val, elem);
             max_val = std::max(max_val, elem);
         }
@@ -173,8 +173,8 @@ void RawImage::convolve_cpu(PSF& psf) {
 
     for (int y = 0; y < height; ++y) {
         for (int x = 0; x < width; ++x) {
-            // Pixels with NO_DATA or NaN remain NO_DATA or NaN.
-            if ((image(y, x) == NO_DATA) || std::isnan(image(y, x))) {
+            // Pixels with invalid data (e.g. NO_DATA or NaN) do not change.
+            if (!pixel_value_valid(image(y, x))) {
                 result(y, x) = image(y, x);
                 continue;
             }
@@ -186,7 +186,7 @@ void RawImage::convolve_cpu(PSF& psf) {
                     if ((x + i >= 0) && (x + i < width) && (y + j >= 0) && (y + j < height)) {
                         float current_pixel = image(y + j, x + i);
                         // note that convention for index access is flipped for PSF
-                        if ((current_pixel != NO_DATA) && !std::isnan(current_pixel)) {
+                        if (pixel_value_valid(current_pixel)) {
                             float current_psf = psf.get_value(i + psf_rad, j + psf_rad);
                             psf_portion += current_psf;
                             sum += current_pixel * current_psf;
@@ -240,22 +240,23 @@ Index RawImage::find_peak(bool furthest_from_center) const {
 
     // Initialize the variables for tracking the peak's location.
     Index result = {0, 0};
-    float max_val = image(0, 0);
+    float max_val = NO_DATA;
     float dist2 = c_x * c_x + c_y * c_y;
 
     // Search each pixel for the peak.
     for (int y = 0; y < height; ++y) {
         for (int x = 0; x < width; ++x) {
-            if (image(y, x) > max_val) {
-                max_val = image(y, x);
+            float pix_val = image(y, x);
+            if (pixel_value_valid(pix_val) && (pix_val > max_val)) {
+                max_val = pix_val;
                 result.i = y;
                 result.j = x;
                 dist2 = (c_x - x) * (c_x - x) + (c_y - y) * (c_y - y);
-            } else if (image(y, x) == max_val) {
+            } else if (pixel_value_valid(pix_val) && (pix_val == max_val)) {
                 int new_dist2 = (c_x - x) * (c_x - x) + (c_y - y) * (c_y - y);
                 if ((furthest_from_center && (new_dist2 > dist2)) ||
                     (!furthest_from_center && (new_dist2 < dist2))) {
-                    max_val = image(y, x);
+                    max_val = pix_val;
                     result.i = y;
                     result.j = x;
                     dist2 = new_dist2;
@@ -282,13 +283,13 @@ ImageMoments RawImage::find_central_moments() const {
     // Find the min (non-NO_DATA) value to subtract off.
     float min_val = FLT_MAX;
     for (int p = 0; p < num_pixels; ++p) {
-        min_val = ((pixels[p] != NO_DATA) && (pixels[p] < min_val)) ? pixels[p] : min_val;
+        min_val = (pixel_value_valid(pixels[p]) && (pixels[p] < min_val)) ? pixels[p] : min_val;
     }
 
     // Find the sum of the zero-shifted (non-NO_DATA) pixels.
     double sum = 0.0;
     for (int p = 0; p < num_pixels; ++p) {
-        sum += (pixels[p] != NO_DATA) ? (pixels[p] - min_val) : 0.0;
+        sum += pixel_value_valid(pixels[p]) ? (pixels[p] - min_val) : 0.0;
     }
     if (sum == 0.0) return res;
 
@@ -296,7 +297,7 @@ ImageMoments RawImage::find_central_moments() const {
     for (int y = 0; y < height; ++y) {
         for (int x = 0; x < width; ++x) {
             int ind = y * width + x;
-            float pix_val = (pixels[ind] != NO_DATA) ? (pixels[ind] - min_val) / sum : 0.0;
+            float pix_val = pixel_value_valid(pixels[ind]) ? (pixels[ind] - min_val) / sum : 0.0;
 
             res.m00 += pix_val;
             res.m10 += (x - c_x) * pix_val;
@@ -326,7 +327,7 @@ bool RawImage::center_is_local_max(double flux_thresh, bool local_max) const {
         if (p != c_ind && local_max && pix_val >= center_val) {
             return false;
         }
-        sum += (pix_val != NO_DATA) ? pix_val : 0.0;
+        sum += pixel_value_valid(pixels[p]) ? pix_val : 0.0;
     }
     if (sum == 0.0) return false;
     return center_val / sum >= flux_thresh;
diff --git a/src/kbmod/search/raw_image.h b/src/kbmod/search/raw_image.h
index 8f8898ae1..a12b6ae48 100644
--- a/src/kbmod/search/raw_image.h
+++ b/src/kbmod/search/raw_image.h
@@ -52,8 +52,7 @@ class RawImage {
     inline float get_pixel(const Index& idx) const { return contains(idx) ? image(idx.i, idx.j) : NO_DATA; }
 
     inline bool pixel_has_data(const Index& idx) const {
-        float pix_val = get_pixel(idx);
-        return ((pix_val != NO_DATA) && !std::isnan(pix_val)) ? true : false;
+        return pixel_value_valid(get_pixel(idx)) ? true : false;
     }
 
     inline void set_pixel(const Index& idx, float value) {
@@ -107,7 +106,7 @@ class RawImage {
     // Find the basic image moments in order to test if stamps have a gaussian shape.
     // It computes the moments on the "normalized" image where the minimum
     // value has been shifted to zero and the sum of all elements is 1.0.
-    // Elements with NO_DATA are treated as zero.
+    // Elements with NO_DATA, NaN, etc. are treated as zero.
     ImageMoments find_central_moments() const;
 
     bool center_is_local_max(double flux_thresh, bool local_max) const;
diff --git a/tests/test_layered_image.py b/tests/test_layered_image.py
index f0d511a34..0fbae27e1 100644
--- a/tests/test_layered_image.py
+++ b/tests/test_layered_image.py
@@ -1,3 +1,5 @@
+import math
+import numpy as np
 import os
 import tempfile
 import unittest
@@ -309,10 +311,14 @@ def test_psi_and_phi_image(self):
                 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.
+        # Mask a single pixel, set another to variance of zero,
+        # and mark two as NaN.
         sci.set_pixel(3, 1, KB_NO_DATA)
         var.set_pixel(3, 1, KB_NO_DATA)
         var.set_pixel(3, 2, 0.0)
+        var.set_pixel(3, 0, np.nan)
+        sci.set_pixel(3, 3, math.nan)
+        sci.set_pixel(3, 4, np.nan)
 
         # Generate and check psi and phi images.
         psi = img.generate_psi_image()
@@ -325,36 +331,47 @@ 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
-                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))
+                psi_has_data = y != 3 or x > 4
+                self.assertEqual(psi.pixel_has_data(y, x), psi_has_data)
+                if psi_has_data:
+                    self.assertAlmostEqual(psi.get_pixel(y, x), x / (y + 1), delta=1e-5)
                 else:
                     self.assertEqual(psi.get_pixel(y, x), KB_NO_DATA)
+                    
+                phi_has_data = y != 3 or x > 2
+                self.assertEqual(phi.pixel_has_data(y, x), phi_has_data)
+                if phi_has_data:
+                    self.assertAlmostEqual(phi.get_pixel(y, x), 1.0 / (y + 1), delta=1e-5)
+                else:
                     self.assertEqual(phi.get_pixel(y, x), KB_NO_DATA)
 
     def test_subtract_template(self):
         sci = self.image.get_science()
         sci.set_pixel(7, 10, KB_NO_DATA)
-        sci.set_pixel(21, 10, KB_NO_DATA)
+        sci.set_pixel(7, 11, KB_NO_DATA)
+        sci.set_pixel(7, 12, math.nan)
+        sci.set_pixel(7, 13, np.nan)
         old_sci = RawImage(sci.image.copy())  # Make a copy.
 
         template = RawImage(self.image.get_width(), self.image.get_height())
         template.set_all(0.0)
         for h in range(sci.height):
-            template.set_pixel(h, 10, 0.01 * h)
+            for w in range(4, sci.width):
+                template.set_pixel(h, w, 0.01 * h)
         self.image.sub_template(template)
 
-        for x in range(sci.width):
-            for y in range(sci.height):
-                val1 = old_sci.get_pixel(y, x)
-                val2 = sci.get_pixel(y, x)
-                if x == 10 and y != 7 and y != 21:
-                    self.assertAlmostEqual(val2, val1 - 0.01 * y, delta=1e-6)
-                else:
+        for y in range(sci.height):
+            for x in range(sci.width):
+                if y == 7 and (x >= 10 and x <= 13):
+                    self.assertFalse(sci.pixel_has_data(y, x))
+                elif x < 4:
+                    val1 = old_sci.get_pixel(y, x)
+                    val2 = sci.get_pixel(y, x)
                     self.assertEqual(val1, val2)
+                else:
+                    val1 = old_sci.get_pixel(y, x) - 0.01 * y
+                    val2 = sci.get_pixel(y, x)
+                    self.assertAlmostEqual(val1, val2, delta=1e-5)
 
     def test_read_write_files(self):
         with tempfile.TemporaryDirectory() as dir_name:
diff --git a/tests/test_raw_image.py b/tests/test_raw_image.py
index 5ec57d970..bd52685d0 100644
--- a/tests/test_raw_image.py
+++ b/tests/test_raw_image.py
@@ -134,6 +134,13 @@ def test_get_bounds(self):
         self.assertAlmostEqual(lower, 0.1, delta=1e-6)
         self.assertAlmostEqual(upper, 100.0, delta=1e-6)
 
+        # Insert a NaN and make sure that does not mess up the computation.
+        img.set_pixel(2, 3, math.nan)
+        img.set_pixel(3, 2, np.nan)
+        lower, upper = img.compute_bounds()
+        self.assertAlmostEqual(lower, 0.1, delta=1e-6)
+        self.assertAlmostEqual(upper, 100.0, delta=1e-6)
+
     def test_find_peak(self):
         "Test RawImage find_peak"
         img = RawImage(self.masked_array)
@@ -146,6 +153,13 @@ def test_find_peak(self):
         self.assertEqual(idx.i, 3)
         self.assertEqual(idx.j, 1)
 
+        # We are okay when the data includes NaNs.
+        img.set_pixel(2, 3, math.nan)
+        img.set_pixel(3, 2, np.nan)
+        idx = img.find_peak(False)
+        self.assertEqual(idx.i, 5)
+        self.assertEqual(idx.j, 5)
+
     def test_find_central_moments(self):
         """Test RawImage central moments."""
         img = RawImage(5, 5, value=0.1)
@@ -191,6 +205,11 @@ def test_find_central_moments(self):
         self.assertAlmostEqual(img_mom.m02, 1.01695, delta=1e-4)
         self.assertAlmostEqual(img_mom.m20, 1.57627, delta=1e-4)
 
+        # Check that nothing fails with NaNs.
+        img.set_pixel(2, 3, math.nan)
+        img.set_pixel(3, 2, np.nan)
+        img_mom = img.find_central_moments()
+
     def convolve_psf_identity(self, device):
         psf_data = np.zeros((3, 3), dtype=np.single)
         psf_data[1, 1] = 1.0

From 3ec1179ecc927e872ab3b3c87fb81494c9487d8c Mon Sep 17 00:00:00 2001
From: Jeremy Kubica <104161096+jeremykubica@users.noreply.github.com>
Date: Fri, 1 Mar 2024 09:36:47 -0500
Subject: [PATCH 3/4] Use the pixel_value_valid function in a few more places

---
 src/kbmod/search/psi_phi_array_ds.h | 2 +-
 src/kbmod/search/stack_search.cpp   | 2 +-
 tests/test_psi_phi_array.py         | 4 ++++
 3 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/src/kbmod/search/psi_phi_array_ds.h b/src/kbmod/search/psi_phi_array_ds.h
index 477b17041..a4caacc3b 100644
--- a/src/kbmod/search/psi_phi_array_ds.h
+++ b/src/kbmod/search/psi_phi_array_ds.h
@@ -37,7 +37,7 @@ struct PsiPhi {
 
 // Helper utility functions.
 inline float encode_uint_scalar(float value, float min_val, float max_val, float scale) {
-    return (value == NO_DATA) ? 0 : (std::max(std::min(value, max_val), min_val) - min_val) / scale + 1.0;
+    return !pixel_value_valid(value) ? 0 : (std::max(std::min(value, max_val), min_val) - min_val) / scale + 1.0;
 }
 
 inline float decode_uint_scalar(float value, float min_val, float scale) {
diff --git a/src/kbmod/search/stack_search.cpp b/src/kbmod/search/stack_search.cpp
index 8d7076775..5062cd106 100644
--- a/src/kbmod/search/stack_search.cpp
+++ b/src/kbmod/search/stack_search.cpp
@@ -202,7 +202,7 @@ std::vector<float> StackSearch::extract_psi_or_phi_curve(Trajectory& trj, bool e
         PsiPhi psi_phi_val = psi_phi_array.read_psi_phi(i, pred_idx.i, pred_idx.j);
 
         float value = (extract_psi) ? psi_phi_val.psi : psi_phi_val.phi;
-        if (value != NO_DATA) {
+        if (pixel_value_valid(value)) {
             result[i] = value;
         }
     }
diff --git a/tests/test_psi_phi_array.py b/tests/test_psi_phi_array.py
index 6de9f8a66..93b395fb3 100644
--- a/tests/test_psi_phi_array.py
+++ b/tests/test_psi_phi_array.py
@@ -105,6 +105,10 @@ def encode_uint_scalar(self):
         # NO_DATA always encodes to 0.0.
         self.assertAlmostEqual(encode_uint_scalar(KB_NO_DATA, 0.0, 10.0, 0.1), 0.0)
 
+        # NAN always encodes to 0.0.
+        self.assertAlmostEqual(encode_uint_scalar(math.nan, 0.0, 10.0, 0.1), 0.0)
+        self.assertAlmostEqual(encode_uint_scalar(np.nan, 0.0, 10.0, 0.1), 0.0)
+        
         # Test clipping.
         self.assertAlmostEqual(encode_uint_scalar(11.0, 0.0, 10.0, 0.1), 100.0)
         self.assertAlmostEqual(encode_uint_scalar(-100.0, 0.0, 10.0, 0.1), 1.0)

From be591f374e979ff65ddbe51fe7e880db40977d42 Mon Sep 17 00:00:00 2001
From: Jeremy Kubica <104161096+jeremykubica@users.noreply.github.com>
Date: Fri, 1 Mar 2024 11:39:48 -0500
Subject: [PATCH 4/4] Fix linting errors

---
 tests/test_layered_image.py | 2 +-
 tests/test_psi_phi_array.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/test_layered_image.py b/tests/test_layered_image.py
index 0fbae27e1..d40a04129 100644
--- a/tests/test_layered_image.py
+++ b/tests/test_layered_image.py
@@ -337,7 +337,7 @@ def test_psi_and_phi_image(self):
                     self.assertAlmostEqual(psi.get_pixel(y, x), x / (y + 1), delta=1e-5)
                 else:
                     self.assertEqual(psi.get_pixel(y, x), KB_NO_DATA)
-                    
+
                 phi_has_data = y != 3 or x > 2
                 self.assertEqual(phi.pixel_has_data(y, x), phi_has_data)
                 if phi_has_data:
diff --git a/tests/test_psi_phi_array.py b/tests/test_psi_phi_array.py
index 93b395fb3..5fc1542c6 100644
--- a/tests/test_psi_phi_array.py
+++ b/tests/test_psi_phi_array.py
@@ -108,7 +108,7 @@ def encode_uint_scalar(self):
         # NAN always encodes to 0.0.
         self.assertAlmostEqual(encode_uint_scalar(math.nan, 0.0, 10.0, 0.1), 0.0)
         self.assertAlmostEqual(encode_uint_scalar(np.nan, 0.0, 10.0, 0.1), 0.0)
-        
+
         # Test clipping.
         self.assertAlmostEqual(encode_uint_scalar(11.0, 0.0, 10.0, 0.1), 100.0)
         self.assertAlmostEqual(encode_uint_scalar(-100.0, 0.0, 10.0, 0.1), 1.0)