image-match/python/py/Lib/site-packages/scipy/interpolate/_ndbspline.py

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)