Merge branch 'main' into suggested_angle

src/kbmod/search/psi_phi_array.cpp

@@ -263,6 +263,24 @@ void fill_psi_phi_array(PsiPhiArray& result_data, int num_bytes, const std::vect
     int height = phi_imgs[0].get_height();
     result_data.set_meta_data(num_bytes, num_times, height, width);
 
+    // Check that we are not processing any bad data (NaNs or infs). We do this once
+    // before allocating space, encoding, etc.
+    for (int t = 0; t < num_times; ++t) {
+        for (int row = 0; row < height; ++row) {
+            for (int col = 0; col < width; ++col) {
+                float psi_val = psi_imgs[t].get_pixel({row, col});
+                if (std::isnan(psi_val) || std::isinf(psi_val)) {
+                    throw std::runtime_error("Invalid value found in PSI array");
+                }
+
+                float phi_val = phi_imgs[t].get_pixel({row, col});
+                if (std::isnan(phi_val) || std::isinf(phi_val)) {
+                    throw std::runtime_error("Invalid value found in PHI array");
+                }
+            }
+        }
+    }
+
     if (result_data.get_num_bytes() == 1 || result_data.get_num_bytes() == 2) {
         // Compute the scaling parameters needed for encoding.
         std::array<float, 3> psi_params =

src/kbmod/search/pydocs/psi_phi_array_docs.h

@@ -158,6 +158,10 @@ static const auto DOC_PsiPhiArray_fill_psi_phi_array = R"doc(
         A list of phi images.
     zeroed_times : `list`
         A list of floating point times starting at zero.
+
+    Raises
+    ------
+    Raises a ``RuntimeError`` if invalid values are found in the psi or phi arrays.
   )doc";
 
 static const auto DOC_PsiPhiArray_fill_psi_phi_array_from_image_stack = R"doc(
@@ -171,6 +175,10 @@ static const auto DOC_PsiPhiArray_fill_psi_phi_array_from_image_stack = R"doc(
         The type of encoding to use (1, 2, or 4).
     stack : `ImageStack`
         The stack of LayeredImages from which to build the psi and phi images.
+
+    Raises
+    ------
+    Raises a ``RuntimeError`` if invalid values are found.
   )doc";
 
 }  // namespace pydocs

tests/test_psi_phi_array.py

@@ -1,6 +1,7 @@
+import math
+import numpy as np
 import unittest
 
-import numpy as np
 
 from kbmod.fake_data.fake_data_creator import make_fake_layered_image
 from kbmod.search import (
@@ -174,6 +175,70 @@ def test_fill_psi_phi_array(self):
                 self.assertFalse(arr.gpu_array_allocated)
                 self.assertFalse(arr.gpu_time_array_allocated)
 
+    def test_fill_invalid(self):
+        arr = PsiPhiArray()
+        old_value = self.psi_1.get_pixel(4, 3)
+
+        # Fails with NaN in psi
+        self.psi_1.set_pixel(4, 3, math.nan)
+        self.assertRaises(
+            RuntimeError,
+            fill_psi_phi_array,
+            arr,
+            4,
+            [self.psi_1, self.psi_2],
+            [self.phi_1, self.phi_2],
+            self.zeroed_times,
+            False,
+        )
+
+        # Fails with inf in psi
+        self.psi_1.set_pixel(4, 3, math.inf)
+        self.assertRaises(
+            RuntimeError,
+            fill_psi_phi_array,
+            arr,
+            4,
+            [self.psi_1, self.psi_2],
+            [self.phi_1, self.phi_2],
+            self.zeroed_times,
+            False,
+        )
+
+        # Reset the psi array and make sure the array builds. Then clear it so we try
+        # to rebuild.
+        self.psi_1.set_pixel(4, 3, old_value)
+        fill_psi_phi_array(
+            arr, 4, [self.psi_1, self.psi_2], [self.phi_1, self.phi_2], self.zeroed_times, False
+        )
+        arr.clear()
+
+        # Fails with NaN in phi
+        self.phi_1.set_pixel(2, 3, math.nan)
+        self.assertRaises(
+            RuntimeError,
+            fill_psi_phi_array,
+            arr,
+            4,
+            [self.psi_1, self.psi_2],
+            [self.phi_1, self.phi_2],
+            self.zeroed_times,
+            False,
+        )
+
+        # Fails with inf in psi
+        self.phi_1.set_pixel(2, 3, math.inf)
+        self.assertRaises(
+            RuntimeError,
+            fill_psi_phi_array,
+            arr,
+            4,
+            [self.psi_1, self.psi_2],
+            [self.phi_1, self.phi_2],
+            self.zeroed_times,
+            False,
+        )
+
     def test_fill_psi_phi_array_from_image_stack(self):
         # Build a fake image stack.
         num_times = 5