Videre
This commit is contained in:
@@ -0,0 +1,518 @@
|
||||
import itertools
|
||||
import functools
|
||||
import operator
|
||||
import numpy as np
|
||||
|
||||
from math import prod
|
||||
from types import GenericAlias
|
||||
|
||||
from . import _dierckx # type: ignore[attr-defined]
|
||||
|
||||
import scipy.sparse.linalg as ssl
|
||||
from scipy.sparse import csr_array
|
||||
from scipy._lib._array_api import array_namespace, xp_capabilities
|
||||
|
||||
from ._bsplines import _not_a_knot, BSpline
|
||||
|
||||
__all__ = ["NdBSpline"]
|
||||
|
||||
|
||||
def _get_dtype(dtype):
|
||||
"""Return np.complex128 for complex dtypes, np.float64 otherwise."""
|
||||
if np.issubdtype(dtype, np.complexfloating):
|
||||
return np.complex128
|
||||
else:
|
||||
return np.float64
|
||||
|
||||
|
||||
@xp_capabilities(
|
||||
cpu_only=True, jax_jit=False,
|
||||
skip_backends=[
|
||||
("dask.array",
|
||||
"https://github.com/data-apis/array-api-extra/issues/488")
|
||||
]
|
||||
)
|
||||
class NdBSpline:
|
||||
"""Tensor product spline object.
|
||||
|
||||
The value at point ``xp = (x1, x2, ..., xN)`` is evaluated as a linear
|
||||
combination of products of one-dimensional b-splines in each of the ``N``
|
||||
dimensions::
|
||||
|
||||
c[i1, i2, ..., iN] * B(x1; i1, t1) * B(x2; i2, t2) * ... * B(xN; iN, tN)
|
||||
|
||||
|
||||
Here ``B(x; i, t)`` is the ``i``-th b-spline defined by the knot vector
|
||||
``t`` evaluated at ``x``.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
t : tuple of 1D ndarrays
|
||||
knot vectors in directions 1, 2, ... N,
|
||||
``len(t[i]) == n[i] + k + 1``
|
||||
c : ndarray, shape (n1, n2, ..., nN, ...)
|
||||
b-spline coefficients
|
||||
k : int or length-d tuple of integers
|
||||
spline degrees.
|
||||
A single integer is interpreted as having this degree for
|
||||
all dimensions.
|
||||
extrapolate : bool, optional
|
||||
Whether to extrapolate out-of-bounds inputs, or return `nan`.
|
||||
Default is to extrapolate.
|
||||
|
||||
Attributes
|
||||
----------
|
||||
t : tuple of ndarrays
|
||||
Knots vectors.
|
||||
c : ndarray
|
||||
Coefficients of the tensor-product spline.
|
||||
k : tuple of integers
|
||||
Degrees for each dimension.
|
||||
extrapolate : bool, optional
|
||||
Whether to extrapolate or return nans for out-of-bounds inputs.
|
||||
Defaults to true.
|
||||
|
||||
Methods
|
||||
-------
|
||||
__call__
|
||||
derivative
|
||||
design_matrix
|
||||
|
||||
See Also
|
||||
--------
|
||||
BSpline : a one-dimensional B-spline object
|
||||
NdPPoly : an N-dimensional piecewise tensor product polynomial
|
||||
|
||||
"""
|
||||
|
||||
# generic type compatibility with scipy-stubs
|
||||
__class_getitem__ = classmethod(GenericAlias)
|
||||
|
||||
def __init__(self, t, c, k, *, extrapolate=None):
|
||||
self._k, self._indices_k1d, (self._t, self._len_t) = _preprocess_inputs(k, t)
|
||||
|
||||
self._asarray = array_namespace(c, *t).asarray
|
||||
|
||||
if extrapolate is None:
|
||||
extrapolate = True
|
||||
self.extrapolate = bool(extrapolate)
|
||||
|
||||
self._c = np.asarray(c)
|
||||
|
||||
ndim = self._t.shape[0] # == len(self.t)
|
||||
if self._c.ndim < ndim:
|
||||
raise ValueError(f"Coefficients must be at least {ndim}-dimensional.")
|
||||
|
||||
for d in range(ndim):
|
||||
td = self.t[d]
|
||||
kd = self.k[d]
|
||||
n = td.shape[0] - kd - 1
|
||||
|
||||
if self._c.shape[d] != n:
|
||||
raise ValueError(f"Knots, coefficients and degree in dimension"
|
||||
f" {d} are inconsistent:"
|
||||
f" got {self._c.shape[d]} coefficients for"
|
||||
f" {len(td)} knots, need at least {n} for"
|
||||
f" k={k}.")
|
||||
|
||||
dt = _get_dtype(self._c.dtype)
|
||||
self._c = np.ascontiguousarray(self._c, dtype=dt)
|
||||
|
||||
@property
|
||||
def k(self):
|
||||
return tuple(self._k)
|
||||
|
||||
@property
|
||||
def t(self):
|
||||
# repack the knots into a tuple
|
||||
return tuple(
|
||||
self._asarray(self._t[d, :self._len_t[d]]) for d in range(self._t.shape[0])
|
||||
)
|
||||
|
||||
@property
|
||||
def c(self):
|
||||
return self._asarray(self._c)
|
||||
|
||||
def __call__(self, xi, *, nu=None, extrapolate=None):
|
||||
"""Evaluate the tensor product b-spline at ``xi``.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
xi : array_like, shape(..., ndim)
|
||||
The coordinates to evaluate the interpolator at.
|
||||
This can be a list or tuple of ndim-dimensional points
|
||||
or an array with the shape (num_points, ndim).
|
||||
nu : sequence of length ``ndim``, optional
|
||||
Orders of derivatives to evaluate. Each must be non-negative.
|
||||
Defaults to the zeroth derivivative.
|
||||
extrapolate : bool, optional
|
||||
Whether to exrapolate based on first and last intervals in each
|
||||
dimension, or return `nan`. Default is to ``self.extrapolate``.
|
||||
|
||||
Returns
|
||||
-------
|
||||
values : ndarray, shape ``xi.shape[:-1] + self.c.shape[ndim:]``
|
||||
Interpolated values at ``xi``
|
||||
"""
|
||||
ndim = self._t.shape[0] # == len(self.t)
|
||||
|
||||
if extrapolate is None:
|
||||
extrapolate = self.extrapolate
|
||||
extrapolate = bool(extrapolate)
|
||||
|
||||
if nu is None:
|
||||
nu = np.zeros((ndim,), dtype=np.int64)
|
||||
else:
|
||||
nu = np.asarray(nu, dtype=np.int64)
|
||||
if nu.ndim != 1 or nu.shape[0] != ndim:
|
||||
raise ValueError(
|
||||
f"invalid number of derivative orders {nu = } for "
|
||||
f"ndim = {len(self.t)}.")
|
||||
if any(nu < 0):
|
||||
raise ValueError(f"derivatives must be positive, got {nu = }")
|
||||
|
||||
# prepare xi : shape (..., m1, ..., md) -> (1, m1, ..., md)
|
||||
xi = np.asarray(xi, dtype=float)
|
||||
xi_shape = xi.shape
|
||||
xi = xi.reshape(-1, xi_shape[-1])
|
||||
xi = np.ascontiguousarray(xi)
|
||||
|
||||
if xi_shape[-1] != ndim:
|
||||
raise ValueError(f"Shapes: xi.shape={xi_shape} and ndim={ndim}")
|
||||
|
||||
# complex -> double
|
||||
was_complex = self._c.dtype.kind == 'c'
|
||||
cc = self._c
|
||||
if was_complex and self._c.ndim == ndim:
|
||||
# make sure that core dimensions are intact, and complex->float
|
||||
# size doubling only adds a trailing dimension
|
||||
cc = self._c[..., None]
|
||||
cc = cc.view(float)
|
||||
|
||||
# prepare the coefficients: flatten the trailing dimensions
|
||||
c1 = cc.reshape(cc.shape[:ndim] + (-1,))
|
||||
c1r = c1.ravel()
|
||||
|
||||
# replacement for np.ravel_multi_index for indexing of `c1`:
|
||||
_strides_c1 = np.asarray([s // c1.dtype.itemsize
|
||||
for s in c1.strides], dtype=np.int64)
|
||||
|
||||
num_c_tr = c1.shape[-1] # # of trailing coefficients
|
||||
out = _dierckx.evaluate_ndbspline(xi,
|
||||
self._t,
|
||||
self._len_t,
|
||||
self._k,
|
||||
nu,
|
||||
extrapolate,
|
||||
c1r,
|
||||
num_c_tr,
|
||||
_strides_c1,
|
||||
self._indices_k1d,
|
||||
)
|
||||
out = out.view(self._c.dtype)
|
||||
out = out.reshape(xi_shape[:-1] + self._c.shape[ndim:])
|
||||
return self._asarray(out)
|
||||
|
||||
@classmethod
|
||||
def design_matrix(cls, xvals, t, k, extrapolate=True):
|
||||
"""Construct the design matrix as a CSR format sparse array.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
xvals : ndarray, shape(npts, ndim)
|
||||
Data points. ``xvals[j, :]`` gives the ``j``-th data point as an
|
||||
``ndim``-dimensional array.
|
||||
t : tuple of 1D ndarrays, length-ndim
|
||||
Knot vectors in directions 1, 2, ... ndim,
|
||||
k : int
|
||||
B-spline degree.
|
||||
extrapolate : bool, optional
|
||||
Whether to extrapolate out-of-bounds values of raise a `ValueError`
|
||||
|
||||
Returns
|
||||
-------
|
||||
design_matrix : a CSR array
|
||||
Each row of the design matrix corresponds to a value in `xvals` and
|
||||
contains values of b-spline basis elements which are non-zero
|
||||
at this value.
|
||||
|
||||
"""
|
||||
xvals = np.asarray(xvals, dtype=float)
|
||||
ndim = xvals.shape[-1]
|
||||
if len(t) != ndim:
|
||||
raise ValueError(
|
||||
f"Data and knots are inconsistent: len(t) = {len(t)} for "
|
||||
f" {ndim = }."
|
||||
)
|
||||
|
||||
# tabulate the flat indices for iterating over the (k+1)**ndim subarray
|
||||
k, _indices_k1d, (_t, len_t) = _preprocess_inputs(k, t)
|
||||
|
||||
# Precompute the shape and strides of the 'coefficients array'.
|
||||
# This would have been the NdBSpline coefficients; in the present context
|
||||
# this is a helper to compute the indices into the colocation matrix.
|
||||
c_shape = tuple(len_t[d] - k[d] - 1 for d in range(ndim))
|
||||
|
||||
# The strides of the coeffs array: the computation is equivalent to
|
||||
# >>> cstrides = [s // 8 for s in np.empty(c_shape).strides]
|
||||
cs = c_shape[1:] + (1,)
|
||||
cstrides = np.cumprod(cs[::-1], dtype=np.int64)[::-1].copy()
|
||||
|
||||
# heavy lifting happens here
|
||||
data, indices, indptr = _dierckx._coloc_nd(xvals,
|
||||
_t, len_t, k, _indices_k1d, cstrides)
|
||||
|
||||
return csr_array((data, indices, indptr))
|
||||
|
||||
def _bspline_derivative_along_axis(self, c, t, k, axis, nu=1):
|
||||
# Move the selected axis to front
|
||||
c = np.moveaxis(c, axis, 0)
|
||||
n = c.shape[0]
|
||||
trailing_shape = c.shape[1:]
|
||||
c_flat = c.reshape(n, -1)
|
||||
|
||||
new_c_list = []
|
||||
new_t = None
|
||||
|
||||
for i in range(c_flat.shape[1]):
|
||||
if k >= nu:
|
||||
b = BSpline.construct_fast(t, c_flat[:, i], k)
|
||||
db = b.derivative(nu)
|
||||
# truncate coefficients to match new knot/degree size
|
||||
db.c = db.c[:len(db.t) - db.k - 1]
|
||||
else:
|
||||
db = BSpline.construct_fast(t, np.zeros(len(t) - 1), 0)
|
||||
|
||||
if new_t is None:
|
||||
new_t = db.t
|
||||
|
||||
new_c_list.append(db.c)
|
||||
|
||||
new_c = np.stack(new_c_list, axis=1).reshape(
|
||||
(len(new_c_list[0]),) + trailing_shape)
|
||||
new_c = np.moveaxis(new_c, 0, axis)
|
||||
|
||||
return new_c, new_t
|
||||
|
||||
def derivative(self, nu):
|
||||
"""
|
||||
Construct a new NdBSpline representing the partial derivative.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
nu : array_like of shape (ndim,)
|
||||
Orders of the partial derivatives to compute along each dimension.
|
||||
|
||||
Returns
|
||||
-------
|
||||
NdBSpline
|
||||
A new NdBSpline representing the partial derivative of the original spline.
|
||||
|
||||
"""
|
||||
nu_arr = np.asarray(nu, dtype=np.int64)
|
||||
ndim = len(self.t)
|
||||
|
||||
if nu_arr.ndim != 1 or nu_arr.shape[0] != ndim:
|
||||
raise ValueError(
|
||||
f"invalid number of derivative orders {nu = } for "
|
||||
f"ndim = {len(self.t)}.")
|
||||
|
||||
if any(nu_arr < 0):
|
||||
raise ValueError(f"derivative orders must be positive, got {nu = }")
|
||||
|
||||
# extract t and c as numpy arrays
|
||||
t_new = [self._t[d, :self._len_t[d]] for d in range(self._t.shape[0])]
|
||||
k_new = list(self.k)
|
||||
c_new = self._c.copy()
|
||||
|
||||
for axis, n in enumerate(nu_arr):
|
||||
if n == 0:
|
||||
continue
|
||||
|
||||
c_new, t_new[axis] = self._bspline_derivative_along_axis(
|
||||
c_new, t_new[axis], k_new[axis], axis, nu=n
|
||||
)
|
||||
k_new[axis] = max(k_new[axis] - n, 0)
|
||||
|
||||
return NdBSpline(tuple(self._asarray(t) for t in t_new),
|
||||
self._asarray(c_new),
|
||||
tuple(k_new),
|
||||
extrapolate=self.extrapolate
|
||||
)
|
||||
|
||||
def _preprocess_inputs(k, t_tpl):
|
||||
"""Helpers: validate and preprocess NdBSpline inputs.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
k : int or tuple
|
||||
Spline orders
|
||||
t_tpl : tuple or array-likes
|
||||
Knots.
|
||||
"""
|
||||
# 1. Make sure t_tpl is a tuple
|
||||
if not isinstance(t_tpl, tuple):
|
||||
raise ValueError(f"Expect `t` to be a tuple of array-likes. "
|
||||
f"Got {t_tpl} instead."
|
||||
)
|
||||
|
||||
# 2. Make ``k`` a tuple of integers
|
||||
ndim = len(t_tpl)
|
||||
try:
|
||||
len(k)
|
||||
except TypeError:
|
||||
# make k a tuple
|
||||
k = (k,)*ndim
|
||||
|
||||
k = np.asarray([operator.index(ki) for ki in k], dtype=np.int64)
|
||||
|
||||
if len(k) != ndim:
|
||||
raise ValueError(f"len(t) = {len(t_tpl)} != {len(k) = }.")
|
||||
|
||||
# 3. Validate inputs
|
||||
ndim = len(t_tpl)
|
||||
for d in range(ndim):
|
||||
td = np.asarray(t_tpl[d])
|
||||
kd = k[d]
|
||||
n = td.shape[0] - kd - 1
|
||||
if kd < 0:
|
||||
raise ValueError(f"Spline degree in dimension {d} cannot be"
|
||||
f" negative.")
|
||||
if td.ndim != 1:
|
||||
raise ValueError(f"Knot vector in dimension {d} must be"
|
||||
f" one-dimensional.")
|
||||
if n < kd + 1:
|
||||
raise ValueError(f"Need at least {2*kd + 2} knots for degree"
|
||||
f" {kd} in dimension {d}.")
|
||||
if (np.diff(td) < 0).any():
|
||||
raise ValueError(f"Knots in dimension {d} must be in a"
|
||||
f" non-decreasing order.")
|
||||
if len(np.unique(td[kd:n + 1])) < 2:
|
||||
raise ValueError(f"Need at least two internal knots in"
|
||||
f" dimension {d}.")
|
||||
if not np.isfinite(td).all():
|
||||
raise ValueError(f"Knots in dimension {d} should not have"
|
||||
f" nans or infs.")
|
||||
|
||||
# 4. tabulate the flat indices for iterating over the (k+1)**ndim subarray
|
||||
# non-zero b-spline elements
|
||||
shape = tuple(kd + 1 for kd in k)
|
||||
indices = np.unravel_index(np.arange(prod(shape)), shape)
|
||||
_indices_k1d = np.asarray(indices, dtype=np.int64).T.copy()
|
||||
|
||||
# 5. pack the knots into a single array:
|
||||
# ([1, 2, 3, 4], [5, 6], (7, 8, 9)) -->
|
||||
# array([[1, 2, 3, 4],
|
||||
# [5, 6, nan, nan],
|
||||
# [7, 8, 9, nan]])
|
||||
t_tpl = [np.asarray(t) for t in t_tpl]
|
||||
ndim = len(t_tpl)
|
||||
len_t = [len(ti) for ti in t_tpl]
|
||||
_t = np.empty((ndim, max(len_t)), dtype=float)
|
||||
_t.fill(np.nan)
|
||||
for d in range(ndim):
|
||||
_t[d, :len(t_tpl[d])] = t_tpl[d]
|
||||
len_t = np.asarray(len_t, dtype=np.int64)
|
||||
|
||||
return k, _indices_k1d, (_t, len_t)
|
||||
|
||||
|
||||
def _iter_solve(a, b, solver=ssl.gcrotmk, **solver_args):
|
||||
# work around iterative solvers not accepting multiple r.h.s.
|
||||
|
||||
# also work around a.dtype == float64 and b.dtype == complex128
|
||||
# cf https://github.com/scipy/scipy/issues/19644
|
||||
if np.issubdtype(b.dtype, np.complexfloating):
|
||||
real = _iter_solve(a, b.real, solver, **solver_args)
|
||||
imag = _iter_solve(a, b.imag, solver, **solver_args)
|
||||
return real + 1j*imag
|
||||
|
||||
if b.ndim == 2 and b.shape[1] !=1:
|
||||
res = np.empty_like(b)
|
||||
for j in range(b.shape[1]):
|
||||
res[:, j], info = solver(a, b[:, j], **solver_args)
|
||||
if info != 0:
|
||||
raise ValueError(f"{solver = } returns {info =} for column {j}.")
|
||||
return res
|
||||
else:
|
||||
res, info = solver(a, b, **solver_args)
|
||||
if info != 0:
|
||||
raise ValueError(f"{solver = } returns {info = }.")
|
||||
return res
|
||||
|
||||
|
||||
def make_ndbspl(points, values, k=3, *, solver=ssl.gcrotmk, **solver_args):
|
||||
"""Construct an interpolating NdBspline.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
points : tuple of ndarrays of float, with shapes (m1,), ... (mN,)
|
||||
The points defining the regular grid in N dimensions. The points in
|
||||
each dimension (i.e. every element of the `points` tuple) must be
|
||||
strictly ascending or descending.
|
||||
values : ndarray of float, shape (m1, ..., mN, ...)
|
||||
The data on the regular grid in n dimensions.
|
||||
k : int, optional
|
||||
The spline degree. Must be odd. Default is cubic, k=3
|
||||
solver : a `scipy.sparse.linalg` solver (iterative or direct), optional.
|
||||
An iterative solver from `scipy.sparse.linalg` or a direct one,
|
||||
`sparse.sparse.linalg.spsolve`.
|
||||
Used to solve the sparse linear system
|
||||
``design_matrix @ coefficients = rhs`` for the coefficients.
|
||||
Default is `scipy.sparse.linalg.gcrotmk`
|
||||
solver_args : dict, optional
|
||||
Additional arguments for the solver. The call signature is
|
||||
``solver(csr_array, rhs_vector, **solver_args)``
|
||||
|
||||
Returns
|
||||
-------
|
||||
spl : NdBSpline object
|
||||
|
||||
Notes
|
||||
-----
|
||||
Boundary conditions are not-a-knot in all dimensions.
|
||||
"""
|
||||
ndim = len(points)
|
||||
xi_shape = tuple(len(x) for x in points)
|
||||
|
||||
try:
|
||||
len(k)
|
||||
except TypeError:
|
||||
# make k a tuple
|
||||
k = (k,)*ndim
|
||||
|
||||
for d, point in enumerate(points):
|
||||
numpts = len(np.atleast_1d(point))
|
||||
if numpts <= k[d]:
|
||||
raise ValueError(f"There are {numpts} points in dimension {d},"
|
||||
f" but order {k[d]} requires at least "
|
||||
f" {k[d]+1} points per dimension.")
|
||||
|
||||
t = tuple(_not_a_knot(np.asarray(points[d], dtype=float), k[d])
|
||||
for d in range(ndim))
|
||||
xvals = np.asarray([xv for xv in itertools.product(*points)], dtype=float)
|
||||
|
||||
# construct the colocation matrix
|
||||
matr = NdBSpline.design_matrix(xvals, t, k)
|
||||
|
||||
# Remove zeros from the sparse matrix
|
||||
# If k=1, then solve() doesn't take long enough for this to help
|
||||
if k[0] >= 3:
|
||||
matr.eliminate_zeros()
|
||||
|
||||
# Solve for the coefficients given `values`.
|
||||
# Trailing dimensions: first ndim dimensions are data, the rest are batch
|
||||
# dimensions, so stack `values` into a 2D array for `spsolve` to undestand.
|
||||
v_shape = values.shape
|
||||
vals_shape = (prod(v_shape[:ndim]), prod(v_shape[ndim:]))
|
||||
vals = values.reshape(vals_shape)
|
||||
|
||||
if solver != ssl.spsolve:
|
||||
solver = functools.partial(_iter_solve, solver=solver)
|
||||
if "atol" not in solver_args:
|
||||
# avoid a DeprecationWarning, grumble grumble
|
||||
solver_args["atol"] = 1e-6
|
||||
|
||||
coef = solver(matr, vals, **solver_args)
|
||||
coef = coef.reshape(xi_shape + v_shape[ndim:])
|
||||
return NdBSpline(t, coef, k)
|
||||
Reference in New Issue
Block a user