#!/usr/bin/env python
"""
interpolate.py
Written by Tyler Sutterley (03/2026)
Interpolators for spatial data
PYTHON DEPENDENCIES:
numpy: Scientific Computing Tools For Python
https://numpy.org
https://numpy.org/doc/stable/user/numpy-for-matlab-users.html
scipy: Scientific Tools for Python
https://docs.scipy.org/doc/
xarray: N-D labeled arrays and datasets in Python
https://docs.xarray.dev/en/stable/
UPDATE HISTORY:
Updated 03/2026: break up extrapolation into separate functions to allow
for caching of the kd-tree when interpolating multiple variables
Updated 02/2026: output data from extrapolate as an xarray DataArray
where there are no valid points within the cutoff distance
Updated 08/2025: added a penalized least square inpainting function
Written 12/2022
"""
from __future__ import annotations
import numpy as np
import xarray as xr
import scipy.fftpack
import scipy.spatial
import FirnCorr.spatial
__all__ = [
"inpaint",
"extrapolate",
"_to_cartesian",
"_build_tree",
"_nearest_neighbors",
]
[docs]
def inpaint(
xs: np.ndarray,
ys: np.ndarray,
zs: np.ndarray,
N: int = 0,
s0: int = 3,
power: int = 2,
epsilon: float = 2.0,
**kwargs,
):
"""
Inpaint over missing data in a two-dimensional array using a
penalized least-squares method based on discrete cosine transforms
:cite:p:`Garcia:2010hn,Wang:2012ei`
Parameters
----------
xs: np.ndarray
x-coordinates
ys: np.ndarray
y-coordinates
zs: np.ndarray
Data with masked values
N: int, default 0
Number of iterations (0 for nearest neighbors)
s0: int, default 3
Smoothing factor
power: int, default 2
Power for lambda function
epsilon: float, default 2.0
Relaxation factor
Returns
-------
z0: np.ndarray
Data with inpainted (filled) values
"""
# find masked values
if isinstance(zs, np.ma.MaskedArray):
W = np.logical_not(zs.mask)
else:
W = np.isfinite(zs)
# no valid values can be found
if not np.any(W):
raise ValueError("No valid values found")
# dimensions of input grid
ny, nx = np.shape(zs)
# calculate initial values using nearest neighbors
# computation of distance Matrix
# use scipy spatial KDTree routines
xgrid, ygrid = np.meshgrid(xs, ys)
tree = scipy.spatial.cKDTree(np.c_[xgrid[W], ygrid[W]])
# find nearest neighbors
masked = np.logical_not(W)
_, ii = tree.query(np.c_[xgrid[masked], ygrid[masked]], k=1)
# copy valid original values
z0 = np.zeros((ny, nx), dtype=zs.dtype)
z0[W] = np.copy(zs[W])
# copy nearest neighbors
z0[masked] = zs[W][ii]
# return nearest neighbors interpolation
if N == 0:
return z0
# copy data to new array with 0 values for mask
ZI = np.zeros((ny, nx), dtype=zs.dtype)
ZI[W] = np.copy(z0[W])
# calculate lambda function
L = np.zeros((ny, nx))
L += np.broadcast_to(np.cos(np.pi * np.arange(ny) / ny)[:, None], (ny, nx))
L += np.broadcast_to(np.cos(np.pi * np.arange(nx) / nx)[None, :], (ny, nx))
LAMBDA = np.power(2.0 * (2.0 - L), power)
# smoothness parameters
s = np.logspace(s0, -6, N)
for i in range(N):
# calculate discrete cosine transform
GAMMA = 1.0 / (1.0 + s[i] * LAMBDA)
DISCOS = GAMMA * scipy.fftpack.dctn(W * (ZI - z0) + z0, norm="ortho")
# update interpolated grid
z0 = (
epsilon * scipy.fftpack.idctn(DISCOS, norm="ortho")
+ (1.0 - epsilon) * z0
)
# reset original values
z0[W] = np.copy(zs[W])
# return the inpainted grid
return z0
# PURPOSE: Nearest-neighbor extrapolation of valid data to output data
[docs]
def _to_cartesian(
x: np.ndarray,
y: np.ndarray,
is_geographic: bool = True,
):
"""
Convert input coordinates to an array of points in a
Cartesian coordinate system
Parameters
----------
x: np.ndarray
x-coordinates to be converted
y: np.ndarray
y-coordinates to be converted
is_geographic: bool, default True
Coordinates are geographic
Returns
-------
points: np.ndarray
Output points in Cartesian coordinates
"""
# verify output dimensions
x = np.atleast_1d(x)
y = np.atleast_1d(y)
# calculate coordinates for nearest-neighbors
if is_geographic:
# global or regional equirectangular model
# ellipsoidal major axis in kilometers
a_axis = 6378.137
# calculate Cartesian coordinates of input grid
xi, yi, zi = FirnCorr.spatial.to_cartesian(x, y, a_axis=a_axis)
# calculate Cartesian coordinates of output coordinates
points = np.c_[xi, yi, zi]
else:
points = np.c_[x, y]
# return the output points in Cartesian coordinates
return points
[docs]
def _build_tree(points: np.ndarray, **kwargs):
"""
Build a KD-tree to search for the nearest-neighbors (NN)
Parameters
----------
points: np.ndarray
Input points in Cartesian coordinates
kwargs: dict
Keyword arguments for ``scipy.spatial.cKDTree``
Returns
-------
tree: scipy.spatial.cKDTree
KD-tree from input points
"""
# create KD-tree of points for nearest-neighbor extrapolation
tree = scipy.spatial.cKDTree(points, **kwargs)
return tree
[docs]
def _nearest_neighbors(
tree: scipy.spatial.cKDTree,
points: np.ndarray,
flattened: np.ndarray,
cutoff: int | float = np.inf,
fill_value: float = None,
**kwargs,
):
"""
Nearest-neighbor (NN) extrapolation of valid model data using KD-trees
Parameters
----------
tree: scipy.spatial.cKDTree
KD-tree of valid points for nearest-neighbor extrapolation
points: np.ndarray
Output points in Cartesian coordinates
flattened: np.ndarray
Valid data array to be extrapolated
cutoff: float, default np.inf
Return only neighbors within distance (kilometers)
fill_value: float, default None
Invalid value
dtype: np.dtype, default from input data
Output data type
Returns
-------
data: xr.DataArray
Extrapolated data
"""
# set default data type
dtype = kwargs.get("dtype", flattened.dtype)
# number of data points
npts, _ = points.shape
# query output data points and find nearest neighbor within cutoff
dd, ii = tree.query(points, k=1, distance_upper_bound=cutoff)
# allocate to output extrapolate data array
data = np.ma.zeros((npts), dtype=dtype, fill_value=fill_value)
data.mask = np.ones((npts), dtype=bool)
# initially set all data to fill value
data.data[:] = data.fill_value
# spatially extrapolate using nearest neighbors
if np.any(np.isfinite(dd)):
(ind,) = np.nonzero(np.isfinite(dd))
data.data[ind] = flattened[ii[ind]]
data.mask[ind] = False
# return extrapolated values
return xr.DataArray(data)