ENH: add max dist to NearestNDInterpolator

scipy/interpolate/_ndgriddata.py

@@ -30,9 +30,9 @@ class NearestNDInterpolator(NDInterpolatorBase):
 
     Parameters
     ----------
-    x : (npoints, ndims) 2-D ndarray of floats
+    x : (npoints, ndims) n-D ndarray of floats
         Data point coordinates.
-    y : (npoints, ) 1-D ndarray of float or complex
+    y : (npoints, ...) n-D ndarray of floats or complex
         Data values.
     rescale : boolean, optional
         Rescale points to unit cube before performing interpolation.
@@ -98,7 +98,7 @@ def __init__(self, x, y, rescale=False, tree_options=None):
         self.tree = cKDTree(self.points, **tree_options)
         self.values = np.asarray(y)
 
-    def __call__(self, *args):
+    def __call__(self, *args, **query_options):
         """
         Evaluate interpolator at given points.
 
@@ -108,6 +108,11 @@ def __call__(self, *args):
             Points where to interpolate data at.
             x1, x2, ... xn can be array-like of float with broadcastable shape.
             or x1 can be array-like of float with shape ``(..., ndim)``
+        query_options : dict, optional
+            This allows `eps`, `p`, `distance_upper_bound`, and `workers` being passed to the cKDTree's query function
+            to be explicitly set. See the `scipy.spatial.cKDTree.query` for an overview of the different options.
+
+            ..versionadded:: 1.12.0
 
         """
         # For the sake of enabling subclassing, NDInterpolatorBase._set_xi performs some operations
@@ -116,14 +121,50 @@ def __call__(self, *args):
         xi = _ndim_coords_from_arrays(args, ndim=self.points.shape[1])
         xi = self._check_call_shape(xi)
         xi = self._scale_x(xi)
-        dist, i = self.tree.query(xi)
-        return self.values[i]
+
+        # We need to handle two important cases for compatibility with a flexible griddata:
+        # (1) the case where xi is of some dimension (n, m, ..., D), where D is the coordinate dimension, and
+        # (2) the case where y is multidimensional (npoints, k, l, ...).
+        # We will first flatten xi to deal with case (1) and build an intermediate return array with shape
+        # (n*m*.., k, l, ...) and then reshape back to (n, m, ..., k, l, ...).
+
+        # Flatten xi for the query
+        xi_flat = xi.reshape(-1, xi.shape[-1])
+        original_shape = xi.shape
+        flattened_shape = xi_flat.shape
+
+        # if distance_upper_bound is set to not be infinite, then we need to consider the case where cKDtree
+        # does not find any points within distance_upper_bound to return. It marks those points as having infinte
+        # distance, which is what will be used below to mask the array and return only the points that were deemed
+        # to have a close enough neighbor to return something useful.
+        dist, i = self.tree.query(xi_flat, **query_options)
+        valid_mask = np.isfinite(dist)
+
+        # create a holder interp_values array with shape (n*m*.., k, l, ...) and fill with nans.
+        interp_shape = flattened_shape[:-1] + self.values.shape[1:] if self.values.ndim > 1 else flattened_shape[:-1]
+
+        if np.issubdtype(self.values.dtype, np.complexfloating):
+            interp_values = np.full(interp_shape, np.nan, dtype=self.values.dtype)
+        else:
+            interp_values = np.full(interp_shape, np.nan)
+
+        if self.values.ndim == 1:
+            interp_values[valid_mask] = self.values[i[valid_mask]]
+        else:
+            interp_values[valid_mask] = self.values[i[valid_mask], ...]
+
+        # (n*m*.., k, l, ...) -> (n, m, ..., k, l, ...)
+        new_shape = original_shape[:-1] + self.values.shape[1:] if self.values.ndim > 1 else original_shape[:-1]
+        interp_values = interp_values.reshape(new_shape)
+
+        return interp_values
 
 
 #------------------------------------------------------------------------------
 # Convenience interface function
 #------------------------------------------------------------------------------
 
+
 def griddata(points, values, xi, method='linear', fill_value=np.nan,
              rescale=False):
     """

scipy/interpolate/tests/test_ndgriddata.py

@@ -196,6 +196,38 @@ def test_nearest_list_argument(self):
         NI = NearestNDInterpolator((d[0], d[1]), list(d[2]))
         assert_array_equal(NI([0.1, 0.9], [0.1, 0.9]), [0, 2])
 
+    def test_nearest_query_options(self):
+        nd = np.array([[0, 0.5, 0, 1],
+                       [0, 0, 0.5, 1],
+                       [0, 1, 1, 2]])
+        delta = 0.1
+        query_points = [0 + delta, 1 + delta], [0 + delta, 1 + delta]
+
+        # case 1 - query max_dist is smaller than the query points' nearest distance to nd.
+        NI = NearestNDInterpolator((nd[0], nd[1]), nd[2])
+        distance_upper_bound = np.sqrt(delta ** 2 + delta ** 2) - 1e-7
+        assert_array_equal(NI(query_points, distance_upper_bound=distance_upper_bound),
+                           [np.nan, np.nan])
+
+        # case 2 - query p is inf, will return [0, 2]
+        distance_upper_bound = np.sqrt(delta ** 2 + delta ** 2) - 1e-7
+        p = np.inf
+        assert_array_equal(NI(query_points, distance_upper_bound=distance_upper_bound, p=p),
+                           [0, 2])
+
+        # case 3 - query max_dist is larger, so should return non np.nan
+        distance_upper_bound = np.sqrt(delta ** 2 + delta ** 2) + 1e-7
+        assert_array_equal(NI(query_points, distance_upper_bound=distance_upper_bound),
+                           [0, 2])
+
+    def test_nearest_query_valid_inputs(self):
+        nd = np.array([[0, 1, 0, 1],
+                       [0, 0, 1, 1],
+                       [0, 1, 1, 2]])
+        NI = NearestNDInterpolator((nd[0], nd[1]), nd[2])
+        with assert_raises(TypeError):
+            NI([0.5, 0.5], query_options="not a dictionary")
+
 
 class TestNDInterpolators:
     @parametrize_interpolators