Merge pull request #8175 from ev-br/cubic_spline

Add cupyx.scipy.interpolate.CubicSpline

cupyx/scipy/interpolate/__init__.py

@@ -6,7 +6,7 @@
 from cupyx.scipy.interpolate._interpolate import PPoly, BPoly, NdPPoly  # NOQA
 from cupyx.scipy.interpolate._cubic import (  # NOQA
     CubicHermiteSpline, PchipInterpolator, pchip_interpolate,  # NOQA
-    Akima1DInterpolator)  # NOQA
+    Akima1DInterpolator, CubicSpline)  # NOQA
 
 # 1-D Splines
 from cupyx.scipy.interpolate._bspline import BSpline, splantider, splder  # NOQA

cupyx/scipy/interpolate/_bspline.py

@@ -221,15 +221,16 @@ def _get_typename(dtype):
 def _get_dtype(dtype):
     """Return np.complex128 for complex dtypes, np.float64 otherwise."""
     if cupy.issubdtype(dtype, cupy.complexfloating):
-        return cupy.complex_
+        return cupy.complex128
     else:
-        return cupy.float_
+        return cupy.float64
 
 
 def _as_float_array(x, check_finite=False):
     """Convert the input into a C contiguous float array.
     NB: Upcasts half- and single-precision floats to double precision.
     """
+    x = cupy.asarray(x)
     x = cupy.ascontiguousarray(x)
     dtyp = _get_dtype(x.dtype)
     x = x.astype(dtyp, copy=False)

cupyx/scipy/interpolate/_bspline2.py

@@ -1,50 +1,17 @@
 import operator
+from math import prod
 
 from numpy.core.multiarray import normalize_axis_index
 
 import cupy
 from cupyx.scipy import sparse
 from cupyx.scipy.sparse.linalg import spsolve
+from cupyx.scipy.interpolate._bspline import _get_dtype, _as_float_array
 
 from cupyx.scipy.interpolate._bspline import (
     _get_module_func, INTERVAL_MODULE, D_BOOR_MODULE, BSpline)
 
 
-def _get_dtype(dtype):
-    """Return np.complex128 for complex dtypes, np.float64 otherwise."""
-    if cupy.issubdtype(dtype, cupy.complexfloating):
-        return cupy.complex_
-    else:
-        return cupy.float_
-
-
-def _as_float_array(x, check_finite=False):
-    """Convert the input into a C contiguous float array.
-
-    NB: Upcasts half- and single-precision floats to double precision.
-    """
-    x = cupy.asarray(x)
-    x = cupy.ascontiguousarray(x)
-    dtyp = _get_dtype(x.dtype)
-    x = x.astype(dtyp, copy=False)
-    if check_finite and not cupy.isfinite(x).all():
-        raise ValueError("Array must not contain infs or nans.")
-    return x
-
-
-# vendored from scipy/_lib/_util.py
-def prod(iterable):
-    """
-    Product of a sequence of numbers.
-    Faster than np.prod for short lists like array shapes, and does
-    not overflow if using Python integers.
-    """
-    product = 1
-    for x in iterable:
-        product *= x
-    return product
-
-
 #################################
 #  Interpolating spline helpers #
 #################################

cupyx/scipy/interpolate/_cubic.py

@@ -1,6 +1,9 @@
 
 import cupy
+from cupy.linalg import solve
 from cupyx.scipy.interpolate._interpolate import PPoly
+from cupyx.scipy.interpolate._bspline2 import make_interp_spline
+import cupyx.scipy.special as spec
 
 
 def _isscalar(x):
@@ -412,3 +415,281 @@ def from_spline(cls, tck, extrapolate=None):
     def from_bernstein_basis(cls, bp, extrapolate=None):
         raise NotImplementedError("This method does not make sense for "
                                   "an Akima interpolator.")
+
+
+def _validate_bc(bc_type, y, expected_deriv_shape, axis):
+    """Validate and prepare boundary conditions.
+
+    Returns
+    -------
+    validated_bc : 2-tuple
+        Boundary conditions for a curve start and end.
+    y : ndarray
+        y casted to complex dtype if one of the boundary conditions has
+        complex dtype.
+    """
+    if isinstance(bc_type, str):
+        if bc_type == 'periodic':
+            if not cupy.allclose(y[0], y[-1], rtol=1e-15, atol=1e-15):
+                raise ValueError(
+                    f"The first and last `y` point along axis {axis} must "
+                    "be identical (within machine precision) when "
+                    "bc_type='periodic'.")
+
+        bc_type = (bc_type, bc_type)
+
+    else:
+        if len(bc_type) != 2:
+            raise ValueError("`bc_type` must contain 2 elements to "
+                             "specify start and end conditions.")
+
+        if 'periodic' in bc_type:
+            raise ValueError("'periodic' `bc_type` is defined for both "
+                             "curve ends and cannot be used with other "
+                             "boundary conditions.")
+
+    validated_bc = []
+    for bc in bc_type:
+        if isinstance(bc, str):
+            if bc == 'clamped':
+                validated_bc.append((1, cupy.zeros(expected_deriv_shape)))
+            elif bc == 'natural':
+                validated_bc.append((2, cupy.zeros(expected_deriv_shape)))
+            elif bc in ['not-a-knot', 'periodic']:
+                validated_bc.append(bc)
+            else:
+                raise ValueError(f"bc_type={bc} is not allowed.")
+        else:
+            try:
+                deriv_order, deriv_value = bc
+            except Exception as e:
+                raise ValueError(
+                    "A specified derivative value must be "
+                    "given in the form (order, value)."
+                ) from e
+
+            if deriv_order not in [1, 2]:
+                raise ValueError("The specified derivative order must "
+                                 "be 1 or 2.")
+
+            deriv_value = cupy.asarray(deriv_value)
+            if deriv_value.shape != expected_deriv_shape:
+                raise ValueError(
+                    "`deriv_value` shape {} is not the expected one {}."
+                    .format(deriv_value.shape, expected_deriv_shape))
+
+            if cupy.issubdtype(deriv_value.dtype, cupy.complexfloating):
+                y = y.astype(complex, copy=False)
+
+            validated_bc.append((deriv_order, deriv_value))
+
+    return validated_bc, y
+
+# XXX: upstream to scipy? (careful with splPrep)
+
+
+def _from_spline(spl):
+    """PPoly.from_spline replacement which handles y.ndim > 1."""
+    t, c, k = spl.tck
+    axis = spl.axis
+    cvals = cupy.empty((k+1, len(t)-1) + c.shape[1:], dtype=c.dtype)
+
+    # convert: here axis=0 because spl(x) rolls the interpolation axis back
+    for m in range(k, -1, -1):
+        ym = spl(t[:-1], nu=m)
+        ym = cupy.moveaxis(ym, axis, 0)
+        cvals[k - m, ...] = ym / spec.gamma(m+1)
+
+    # redo the axis reshuffle in _PPolyBase.__init__:
+    # https://github.com/scipy/scipy/blob/v1.12.0/scipy/interpolate/_interpolate.py#L826
+    cvals_ = cupy.moveaxis(cvals, 0, axis+1)
+    cvals_ = cupy.moveaxis(cvals_, 0, axis+1)
+
+    pp = PPoly(cvals_, t, axis=axis)
+    return pp
+
+
+class CubicSpline(CubicHermiteSpline):
+    """Cubic spline data interpolator.
+
+    Interpolate data with a piecewise cubic polynomial which is twice
+    continuously differentiable. The result is represented as a `PPoly`
+    instance with breakpoints matching the given data.
+
+    Parameters
+    ----------
+    x : array_like, shape (n,)
+        1-D array containing values of the independent variable.
+        Values must be real, finite and in strictly increasing order.
+    y : array_like, shape (n,)
+        Array containing values of the dependent variable. It can have
+        arbitrary number of dimensions, but the length along ``axis``
+        (see below) must match the length of ``x``. Values must be finite.
+    axis : int, optional
+        Axis along which `y` is assumed to be varying. Meaning that for
+        ``x[i]`` the corresponding values are ``np.take(y, i, axis=axis)``.
+        Default is 0.
+    bc_type : string or 2-tuple, optional
+        Boundary condition type. Two additional equations, given by the
+        boundary conditions, are required to determine all coefficients of
+        polynomials on each segment.
+
+        If `bc_type` is a string, then the specified condition will be applied
+        at both ends of a spline. Available conditions are:
+
+        * 'not-a-knot' (default): The first and second segment at a curve end
+          are the same polynomial. It is a good default when there is no
+          information on boundary conditions.
+        * 'periodic': The interpolated functions is assumed to be periodic
+          of period ``x[-1] - x[0]``. The first and last value of `y` must be
+          identical: ``y[0] == y[-1]``. This boundary condition will result in
+          ``y'[0] == y'[-1]`` and ``y''[0] == y''[-1]``.
+        * 'clamped': The first derivative at curves ends are zero. Assuming
+          a 1D `y`, ``bc_type=((1, 0.0), (1, 0.0))`` is the same condition.
+        * 'natural': The second derivative at curve ends are zero. Assuming
+          a 1D `y`, ``bc_type=((2, 0.0), (2, 0.0))`` is the same condition.
+
+        If `bc_type` is a 2-tuple, the first and the second value will be
+        applied at the curve start and end respectively. The tuple values can
+        be one of the previously mentioned strings (except 'periodic') or a
+        tuple `(order, deriv_values)` allowing to specify arbitrary
+        derivatives at curve ends:
+
+        * ``order``: the derivative order, 1 or 2.
+        * ``deriv_value``: array_like containing derivative values, shape must
+          be the same as `y`, excluding ``axis`` dimension. For example, if
+          `y` is 1-D, then `deriv_value` must be a scalar. If `y` is 3-D with
+          the shape (n0, n1, n2) and axis=2, then `deriv_value` must be 2-D
+          and have the shape (n0, n1).
+    extrapolate : {bool, 'periodic', None}, optional
+        If bool, determines whether to extrapolate to out-of-bounds points
+        based on first and last intervals, or to return NaNs. If 'periodic',
+        periodic extrapolation is used. If None (default), ``extrapolate`` is
+        set to 'periodic' for ``bc_type='periodic'`` and to True otherwise.
+
+    Attributes
+    ----------
+    x : ndarray, shape (n,)
+        Breakpoints. The same ``x`` which was passed to the constructor.
+    c : ndarray, shape (4, n-1, ...)
+        Coefficients of the polynomials on each segment. The trailing
+        dimensions match the dimensions of `y`, excluding ``axis``.
+        For example, if `y` is 1-d, then ``c[k, i]`` is a coefficient for
+        ``(x-x[i])**(3-k)`` on the segment between ``x[i]`` and ``x[i+1]``.
+    axis : int
+        Interpolation axis. The same axis which was passed to the
+        constructor.
+
+    See Also
+    --------
+    scipy.interpolate.CubicSpline
+
+    Notes
+    -----
+    Parameters `bc_type` and ``extrapolate`` work independently, i.e. the
+    former controls only construction of a spline, and the latter only
+    evaluation.
+
+    When a boundary condition is 'not-a-knot' and n = 2, it is replaced by
+    a condition that the first derivative is equal to the linear interpolant
+    slope. When both boundary conditions are 'not-a-knot' and n = 3, the
+    solution is sought as a parabola passing through given points.
+
+    """
+
+    def __init__(self, x, y, axis=0, bc_type='not-a-knot', extrapolate=None):
+        x = cupy.asarray(x)
+        y = cupy.asarray(y)
+
+        if y.size == 0:
+            # bail out early for zero-sized arrays
+            s = cupy.zeros_like(y)
+            super().__init__(x, y, s, axis=axis, extrapolate=extrapolate)
+            self.axis = axis
+        elif len(x) <= 3:
+            # special cases: make_interp_spline requires >k data points
+            # vendor what scipy.interpolate.CubicSpline does
+            x, dx, y, axis, _ = prepare_input(x, y, axis)
+            n = len(x)
+
+            bc, y = _validate_bc(bc_type, y, y.shape[1:], axis)
+
+            if extrapolate is None:
+                if bc[0] == 'periodic':
+                    extrapolate = 'periodic'
+                else:
+                    extrapolate = True
+
+            dxr = dx.reshape([dx.shape[0]] + [1] * (y.ndim - 1))
+            slope = cupy.diff(y, axis=0) / dxr
+
+            # If bc is 'not-a-knot' this change is just a convention.
+            # If bc is 'periodic' then we already checked that y[0] == y[-1],
+            # and the spline is just a constant, we handle this case in the
+            # same way by setting the first derivatives to slope, which is 0.
+            if n == 2:
+                if bc[0] in ['not-a-knot', 'periodic']:
+                    bc[0] = (1, slope[0])
+                if bc[1] in ['not-a-knot', 'periodic']:
+                    bc[1] = (1, slope[0])
+                s = cupy.r_[slope, slope]
+
+            # This is a special case, when both conditions are 'not-a-knot'
+            # and n == 3. In this case 'not-a-knot' can't be handled regularly
+            # as the both conditions are identical. We handle this case by
+            # constructing a parabola passing through given points.
+            if n == 3 and bc[0] == 'not-a-knot' and bc[1] == 'not-a-knot':
+                A = cupy.zeros((3, 3))  # This is a standard matrix.
+                b = cupy.empty((3,) + y.shape[1:], dtype=y.dtype)
+
+                A[0, 0] = 1
+                A[0, 1] = 1
+                A[1, 0] = dx[1]
+                A[1, 1] = 2 * (dx[0] + dx[1])
+                A[1, 2] = dx[0]
+                A[2, 1] = 1
+                A[2, 2] = 1
+
+                b[0] = 2 * slope[0]
+                b[1] = 3 * (dxr[0] * slope[1] + dxr[1] * slope[0])
+                b[2] = 2 * slope[1]
+
+                s = solve(A, b)
+            elif n == 3 and bc[0] == 'periodic':
+                # In case when number of points is 3 we compute the derivatives
+                # manually
+                t = (slope / dxr).sum(0) / (1. / dxr).sum(0)
+                s = cupy.broadcast_to(t, (n,) + y.shape[1:])
+
+            # finally, construct the object
+            super().__init__(x, y, s, axis=0, extrapolate=extrapolate)
+            self.axis = axis
+        else:
+            # general case: delegate to make_interp_spline and convert to PPoly
+            # first, repackage the boundary conditions.
+            # Also account for that complex b.c. upcast y in CubicSpline
+            need_complex = False
+            if not isinstance(bc_type, str):
+                # must be a 2-tuple
+                bc_0, bc_1 = bc_type
+                if not isinstance(bc_0, str):
+                    need_complex = cupy.iscomplexobj(bc_0[1])
+                    bc_0 = [bc_0]
+                if not isinstance(bc_1, str):
+                    need_complex = cupy.iscomplexobj(bc_1[1])
+                    bc_1 = [bc_1]
+                bc_type = (bc_0, bc_1)
+
+            if need_complex:
+                y = y.astype(complex)
+
+            # actually do the work
+            spl = make_interp_spline(x, y, k=3, axis=axis, bc_type=bc_type)
+            pp = _from_spline(spl)
+            self.x = pp.x
+            self.c = pp.c
+            self.axis = pp.axis
+
+            if extrapolate is None:
+                extrapolate = True
+            self.extrapolate = extrapolate

docs/source/reference/scipy_interpolate.rst

@@ -21,6 +21,7 @@ Univariate interpolation
    Akima1DInterpolator
    PPoly
    BPoly
+   CubicSpline
 
 1-D Splines
 -----------