Update NaN filtering in InterpolatingMap

straxen/itp_map.py

@@ -41,19 +41,21 @@ def __init__(self, points, values, neighbours_to_use=None, array_valued=False):
             self.n_dim = self.values.shape[-1]
 
     def __call__(self, points):
-        distances, indices = self.kdtree.query(points, self.neighbours_to_use)
+        points = np.asarray(points)
 
+        # kdtree doesn't grok NaNs
+        # Start with all Nans, then overwrite for the finite points
         result = np.ones(len(points)) * float("nan")
         if self.array_valued:
             result = np.repeat(result.reshape(-1, 1), self.n_dim, axis=1)
+        valid = np.all(np.isfinite(points), axis=-1)
 
-        # If one of the coordinates is NaN, the neighbour-query fails.
-        # If we don't filter these out, it would result in an IndexError
-        # as the kdtree returns an invalid index if it can't find neighbours.
-        valid = (distances < float("inf")).max(axis=-1)
+        # Get distances to neighbours_to_use nearest neighbours
+        distances, indices = self.kdtree.query(points[valid], self.neighbours_to_use)
 
-        values = self.values[indices[valid]]
-        weights = 1 / np.clip(distances[valid], 1e-6, float("inf"))
+        # Get values and weights for inverse distance weighted interpolation
+        values = self.values[indices]
+        weights = 1 / np.clip(distances, 1e-6, float("inf"))
         if self.array_valued:
             weights = np.repeat(weights, self.n_dim).reshape(values.shape)
 

tests/test_itp_map.py

@@ -1,4 +1,5 @@
 from unittest import TestCase, skipIf
+import numpy as np
 from straxen import utilix_is_configured, get_resource
 from straxen import InterpolatingMap, save_interpolation_formatted_map
 
@@ -77,3 +78,13 @@ def test_array_valued(self):
 
             self.assertAlmostEqual(map_at_random_point[1][0], 2.17815179, places=places)
             self.assertAlmostEqual(map_at_random_point[1][1], 9.47282782, places=places)
+
+        # Test querying nonfinite values gives NaNs, not errors
+        for itp_map in itp_maps:
+            map_at_random_point = itp_map([[0, np.nan, 0], [0, 0, -140]])
+            # Shape is still correct
+            assert map_at_random_point.shape == (2, 2)
+            # First point gives NaN
+            assert np.all(np.isnan(map_at_random_point[0]))
+            # Second point does not
+            assert not np.any(np.isnan(map_at_random_point[1]))