ENH: add max dist to NearestNDInterpolator

scipy/interpolate/_ndgriddata.py

@@ -98,26 +98,55 @@ 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=None):
         """
         Evaluate interpolator at given points.
 
         Parameters
         ----------
+        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. Note
+            that k is restricted to 1 since NearestNDInterpolator has to have k=1 by definition.
+
+            ..versionadded:: 1.12.0
+
         x1, x2, ... xn : array-like of float
             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)``
 
         """
+        if query_options is not None and not isinstance(query_options, dict):
+            raise TypeError("query_options must be a dictionary")
+
+        # gather the query options to pass to cKDTree.query
+        query_options = query_options or {}
+        eps = query_options.get('eps', 0)
+        p = query_options.get('p', 2)
+        distance_upper_bound = query_options.get('distance_upper_bound', np.inf)
+        workers = query_options.get('workers', 1)
+
         # For the sake of enabling subclassing, NDInterpolatorBase._set_xi performs some operations
         # which are not required by NearestNDInterpolator.__call__, hence here we operate
         # on xi directly, without calling a parent class function.
         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]
+        dist, i = self.tree.query(xi, eps=eps, p=p, distance_upper_bound=distance_upper_bound, workers=workers)
+
+        if np.isfinite(distance_upper_bound):
+            # 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.
+            valid_mask = np.isfinite(dist)
+            interp_values = np.zeros(i.shape)
+            interp_values[~valid_mask] = np.nan
+            interp_values[valid_mask] = self.values[i[valid_mask]]
+            return interp_values
+        else:
+            return self.values[i]
 
 
 #------------------------------------------------------------------------------

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, 1, 0, 1],
+                       [0, 0, 1, 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, query_options={"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, query_options={"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, query_options={"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 self.assertRaises(TypeError):
+            NI([0.5, 0.5], query_options="not a dictionary")
+
 
 class TestNDInterpolators:
     @parametrize_interpolators