Coverage for src/cosmic_toolbox/NearestWeightedNDInterpolator.py: 94%
35 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-07-10 10:31 +0000
1"""
2Convenience interface to N-D interpolation.
4Provides a weighted nearest-neighbor interpolator using BallTree for efficient
5N-dimensional interpolation.
7author: Tomasz Kacprzak
8"""
10import numpy as np
12try:
13 from scipy.interpolate import NDInterpolatorBase
14except ImportError:
15 try:
16 from scipy.interpolate._ndgriddata import NDInterpolatorBase
17 except ImportError:
18 from scipy.interpolate.interpnd import NDInterpolatorBase
20from sklearn.neighbors import BallTree
21from sklearn.preprocessing import MinMaxScaler
24class NearestWeightedNDInterpolator(NDInterpolatorBase):
25 """
26 Weighted nearest-neighbor interpolation in N dimensions.
28 This interpolator uses BallTree for efficient nearest-neighbor queries
29 and computes weighted averages based on inverse distance.
31 :param x: Training points, shape (n_samples, n_dims).
32 :type x: numpy.ndarray
33 :param y: Training values, shape (n_samples,).
34 :type y: numpy.ndarray
35 :param k: Number of nearest neighbors to use. Defaults to n_dims + 1
36 (number of vertices of an n_dims dimensional simplex).
37 :type k: int or None
38 :param tree_options: Options passed to sklearn's BallTree constructor.
39 :type tree_options: dict or None
41 Example
42 -------
43 >>> import numpy as np
44 >>> from cosmic_toolbox import NearestWeightedNDInterpolator
45 >>> x = np.array([[0, 0], [0, 1], [1, 0], [1, 1]])
46 >>> y = np.array([0, 1, 1, 2])
47 >>> interp = NearestWeightedNDInterpolator(x, y)
48 >>> interp(np.array([[0.5, 0.5]]))
49 array([1.])
50 """
52 def __init__(self, x, y, k=None, tree_options=None):
53 if tree_options is None:
54 tree_options = {}
55 self.x = x
56 self.y = y
57 self.ndim = self.x.shape[1]
58 if k is None:
59 k = self.ndim + 1 # number of vertices of ndim dimensional simplex
60 self.k = k
61 self.scaler = MinMaxScaler()
62 self.x = self.scaler.fit_transform(self.x)
63 self.tree = BallTree(self.x, **tree_options)
65 def __call__(self, xi):
66 """
67 Evaluate the interpolator at given points.
69 :param xi: Points at which to interpolate, shape (n_points, n_dims).
70 :type xi: numpy.ndarray
71 :return: Interpolated values, shape (n_points,).
72 :rtype: numpy.ndarray
73 :raises AssertionError: If xi has wrong shape or dimensionality.
74 """
75 assert len(xi.shape) == 2
76 assert self.ndim == xi.shape[1]
78 xi = self.scaler.transform(xi)
80 dist, i = self.tree.query(xi, self.k)
81 vi = self.y[i].reshape((xi.shape[0], self.k))
82 if self.k > 1:
83 weight = 1.0 / dist
84 weight[~np.isfinite(weight)] = 0
85 weight = weight.reshape((xi.shape[0], self.k))
86 vi = np.average(vi, weights=weight, axis=1)
88 return vi
91# class NearestWeightedNDInterpolator(NDInterpolatorBase):
92# """
93# NearestWeightedNDInterpolator(x, y)
94# NN interpolation in N dimensions with weighted interpolation.
95# Uses BallTree instead of cKDTree
96# .. versionadded:: 0.9
97# Methods
98# -------
99# __call__
100# Parameters
101# ----------
102# x : (Npoints, Ndims) ndarray of floats
103# Data point coordinates.
104# y : (Npoints,) ndarray of float or complex
105# Data values.
106# rescale : boolean, optional
107# Rescale points to unit cube before performing interpolation.
108# This is useful if some of the input dimensions have
109# incommensurable units and differ by many orders of magnitude.
110# .. versionadded:: 0.14.0
111# tree_options : dict, optional
112# Options passed to the underlying ``cKDTree``.
113# .. versionadded:: 0.17.0
114# Notes
115# -----
116# Uses ``scipy.spatial.cKDTree``
117# Examples
118# --------
119# We can interpolate values on a 2D plane:
120# >>> from scipy.interpolate import NearestNDInterpolator
121# >>> import matplotlib.pyplot as plt
122# >>> np.random.seed(0)
123# >>> x = np.random.random(10) - 0.5
124# >>> y = np.random.random(10) - 0.5
125# >>> z = np.hypot(x, y)
126# >>> X = np.linspace(min(x), max(x))
127# >>> Y = np.linspace(min(y), max(y))
128# >>> X, Y = np.meshgrid(X, Y) # 2D grid for interpolation
129# >>> interp = NearestNDInterpolator(list(zip(x, y)), z)
130# >>> Z = interp(X, Y)
131# >>> plt.pcolormesh(X, Y, Z, shading='auto')
132# >>> plt.plot(x, y, "ok", label="input point")
133# >>> plt.legend()
134# >>> plt.colorbar()
135# >>> plt.axis("equal")
136# >>> plt.show()
137# See also
138# --------
139# griddata :
140# Interpolate unstructured D-D data.
141# LinearNDInterpolator :
142# Piecewise linear interpolant in N dimensions.
143# CloughTocher2DInterpolator :
144# Piecewise cubic, C1 smooth, curvature-minimizing interpolant in 2D.
145# """
147# def __init__(self, x, y, k=None, tree_options={}):
148# NDInterpolatorBase.__init__(self, x, y, rescale=True,
149# need_contiguous=False,
150# need_values=False)
151# ndim = self.points.shape[1]
152# if k==None:
153# k = ndim + 1 # number of vertices of ndim dimensional simplex
154# self.k = k
155# self.scaler = MinMaxScaler(copy=False)
156# self.scaler.fit_transform(self.points)
157# self.tree = BallTree(self.points, **tree_options)
158# self.values = np.asarray(y)
160# def __call__(self, *args):
161# """
162# Evaluate interpolator at given points.
163# Parameters
164# ----------
165# x1, x2, ... xn: array-like of float
166# Points where to interpolate data at.
167# x1, x2, ... xn can be array-like of float with broadcastable
168# shape or x1 can be array-like of float with shape ``(..., ndim)``
169# """
171# ndim = self.points.shape[1]
172# xi = _ndim_coords_from_arrays(args, ndim=ndim)
173# xi = self._check_call_shape(xi)
174# # xi = self._scale_x(xi)
175# orig_shape = xi.shape
177# if len(xi.shape) > 2:
178# xi = xi.reshape(-1, xi.shape[-1])
179# self.scaler.transform(xi)
181# dist, i = self.tree.query(xi, self.k)
182# weight = 1./dist
183# print(weight.shape)
184# weight[~np.isfinite(weight)] = 0
185# import ipdb; ipdb.set_trace()
186# try:
187# vi = np.average(self.values[i].squeeze(), weights=weight, axis=1)
188# except Exception as err:
189# print(err)
190# import ipdb; ipdb.set_trace()
192# return np.reshape(vi, orig_shape[:-1])