diff --git a/src/kbmod/search/psi_phi_array.cpp b/src/kbmod/search/psi_phi_array.cpp index e0640e9da..f48c1539e 100644 --- a/src/kbmod/search/psi_phi_array.cpp +++ b/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 psi_params = diff --git a/src/kbmod/search/pydocs/psi_phi_array_docs.h b/src/kbmod/search/pydocs/psi_phi_array_docs.h index 05abe3053..d7239a22c 100644 --- a/src/kbmod/search/pydocs/psi_phi_array_docs.h +++ b/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 diff --git a/tests/test_psi_phi_array.py b/tests/test_psi_phi_array.py index b75a0a463..6de9f8a66 100644 --- a/tests/test_psi_phi_array.py +++ b/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