rbfgpu 1.0.0__py3-none-any.whl

rbfgpu/__init__.py

@@ -0,0 +1,77 @@
+"""
+rbfGPU: GPU/CPU agnostic Radial Basis Function interpolation.
+
+``rbfgpu`` interpolates scattered data with radial basis functions (thin plate
+splines by default) using the `XuPy <https://pypi.org/project/xupy/>`_ array
+namespace, so the very same code runs on CuPy arrays when an NVIDIA GPU is
+available and on NumPy arrays otherwise.  It is the drop-in replacement for
+``scipy.interpolate.RBFInterpolator`` in the regime where the O(n^2) Gram matrix
+and the O(n^3) solve become the bottleneck, i.e. from a few thousand constraint
+points upwards.
+
+Examples
+--------
+>>> import numpy as np
+>>> from rbfgpu import RBFInterpolator
+>>> rng = np.random.default_rng(0)
+>>> y = rng.random((1000, 2))
+>>> d = np.sin(3 * y[:, 0]) * np.cos(3 * y[:, 1])
+>>> rbf = RBFInterpolator(y, d)                  # thin plate spline, degree=1
+>>> values = rbf(rng.random((10000, 2)))         # chunked, bounded memory
+>>> values.shape
+(10000,)
+
+The backend can be inspected (and, through XuPy, switched) at any time:
+
+>>> import rbfgpu
+>>> rbfgpu.backend_info()["backend"] in ("numpy", "cupy")
+True
+"""
+
+from . import __version__ as _version_module
+from ._backend import as_numpy, backend_info, is_on_gpu, require_gpu
+from ._distance import default_epsilon, euclidean_distance
+from ._kernels import (
+    KERNELS,
+    KERNEL_ALIASES,
+    Kernel,
+    available_kernels,
+    resolve_kernel,
+)
+from ._polynomial import monomial_exponents, polynomial_matrix
+from .__version__ import __version__
+from .interpolator import (
+    DEFAULT_MAX_CHUNK_BYTES,
+    DEFAULT_MIN_CHUNK_SIZE,
+    RBFInterpolator,
+)
+
+__title__ = _version_module.__title__
+__description__ = _version_module.__description__
+__author__ = _version_module.__author__
+__author_email__ = _version_module.__author_email__
+__license__ = _version_module.__license__
+__url__ = _version_module.__url__
+
+__all__ = [
+    "__version__",
+    # main API
+    "RBFInterpolator",
+    "available_kernels",
+    "KERNELS",
+    "KERNEL_ALIASES",
+    "Kernel",
+    "resolve_kernel",
+    # backend
+    "backend_info",
+    "is_on_gpu",
+    "require_gpu",
+    "as_numpy",
+    # building blocks (handy for custom interpolants and for testing)
+    "euclidean_distance",
+    "default_epsilon",
+    "monomial_exponents",
+    "polynomial_matrix",
+    "DEFAULT_MAX_CHUNK_BYTES",
+    "DEFAULT_MIN_CHUNK_SIZE",
+]

rbfgpu/__version__.py

@@ -0,0 +1,32 @@
+"""
+Created on October 2026
+
+Author
+------
+Pietro Ferraiuolo
+"""
+
+__title__ = "rbfgpu"
+__long_title__ = "rbfGPU : GPU/CPU agnostic Radial Basis Function interpolation"
+__description__ = (
+    "Thin plate spline and general RBF interpolation that runs on the GPU through "
+    "XuPy (CuPy) and transparently falls back to NumPy on the CPU."
+)
+__version__ = "1.0.0"
+__author__ = "Pietro Ferraiuolo"
+__maintainer__ = "Pietro Ferraiuolo"
+__author_email__ = "pietro.ferraiuolo@inaf.it"
+__license__ = "MIT License"
+__url__ = "https://github.com/pietroferraiuolo/rbfgpu"
+
+__all__ = [
+    "__title__",
+    "__long_title__",
+    "__description__",
+    "__version__",
+    "__author__",
+    "__maintainer__",
+    "__author_email__",
+    "__license__",
+    "__url__",
+]

rbfgpu/_backend.py

@@ -0,0 +1,121 @@
+"""
+Backend adapter over the XuPy array namespace.
+
+rbfGPU never talks to NumPy or CuPy directly: every array operation goes through
+``xupy``, whose namespace is resolved dynamically against the active backend
+(CuPy when a usable NVIDIA GPU is present, NumPy otherwise).  This module is the
+only place where that duality is dealt with.
+
+Note that, unlike the module this package started from, there is no
+``if xp.on_gpu:`` branch: XuPy already provides the CPU fallback, so the very
+same code path runs (and is tested) on both backends.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import xupy as xp
+
+__all__ = [
+    "LINALG_ERRORS",
+    "as_numpy",
+    "backend_info",
+    "is_on_gpu",
+    "require_gpu",
+    "xp",
+]
+
+
+def _linalg_errors() -> tuple:
+    """
+    Exception types raised by ``xp.linalg.solve`` on a singular matrix.
+
+    NumPy's and CuPy's ``LinAlgError`` are distinct classes, so both are
+    collected: whichever backend is active, the solve step can catch the right
+    one.
+    """
+    errors = [np.linalg.LinAlgError]
+    try:  # pragma: no cover - depends on the machine running the tests
+        import cupy
+
+        errors.append(cupy.linalg.LinAlgError)
+    except Exception:  # CuPy missing, or importable but unusable
+        pass
+    return tuple(errors)
+
+
+LINALG_ERRORS = _linalg_errors()
+
+
+def is_on_gpu() -> bool:
+    """True when XuPy is serving arrays from CuPy (GPU), False for NumPy (CPU)."""
+    return bool(xp.on_gpu)
+
+
+def require_gpu(hint: str = "") -> None:
+    """
+    Raise ``RuntimeError`` unless XuPy is running on the GPU backend.
+
+    Useful at the top of a script that is only worth running on a GPU, so that
+    it fails loudly instead of silently computing on the CPU.
+
+    Parameters
+    ----------
+    hint : str, optional
+        Extra text appended to the error message.
+    """
+    if not is_on_gpu():
+        message = (
+            "rbfGPU is running on the NumPy (CPU) backend: no usable NVIDIA GPU "
+            "or CuPy installation was found. Install a matching CuPy build with "
+            "`pip install 'rbfgpu[cuda12]'` (or `[cuda13]`) and retry."
+        )
+        if hint:
+            message = f"{message} {hint}"
+        raise RuntimeError(message)
+
+
+def as_numpy(array):
+    """
+    Move ``array`` to host memory and return it as a NumPy array.
+
+    A no-op when XuPy is on the CPU backend (the array is returned unchanged if
+    it already is an ``ndarray``).
+    """
+    return xp.asnumpy(array)
+
+
+def backend_info() -> dict:
+    """
+    Describe the array backend rbfGPU is currently running on.
+
+    Returns
+    -------
+    info : dict
+        With keys ``on_gpu``, ``backend`` ("cupy" or "numpy"), ``device``,
+        ``numpy_version``, ``xupy_version`` and ``cupy_version`` (None when
+        CuPy is not importable).
+    """
+    on_gpu = is_on_gpu()
+    info = {
+        "on_gpu": on_gpu,
+        "backend": "cupy" if on_gpu else "numpy",
+        "device": "cpu",
+        "numpy_version": np.__version__,
+        "xupy_version": xp.__version__,
+        "cupy_version": None,
+    }
+    if on_gpu:  # pragma: no cover - needs a GPU
+        try:
+            import cupy
+
+            info["cupy_version"] = cupy.__version__
+            device_id = cupy.cuda.runtime.getDevice()
+            props = cupy.cuda.runtime.getDeviceProperties(device_id)
+            name = props.get("name", b"")
+            info["device"] = f"cuda:{device_id} " + (
+                name.decode() if isinstance(name, bytes) else str(name)
+            )
+        except Exception as error:  # never let an informational call fail
+            info["device"] = f"<unknown: {error}>"
+    return info

rbfgpu/_distance.py

@@ -0,0 +1,138 @@
+"""
+Pairwise Euclidean distances on the active XuPy backend.
+
+Distances are the only O(n*m) quantity in an RBF interpolant, so they are
+computed with the dot-product identity instead of an explicit difference: it
+avoids materialising the (n, m, d) broadcast temporary, which on a GPU is what
+actually limits the problem size.
+"""
+
+from __future__ import annotations
+
+import math
+
+from ._backend import xp
+
+__all__ = ["default_epsilon", "euclidean_distance"]
+
+#: Memory budget (in floats) for the (block, npoints, ndim) temporary that
+#: :func:`default_epsilon` allocates: 2**22 floats is 32 MiB of float64.
+_EPSILON_BLOCK_FLOATS = 2**22
+
+
+def euclidean_distance(x, y, epsilon=None):
+    """
+    Compute the (optionally scaled) Euclidean distance matrix between two clouds.
+
+    Parameters
+    ----------
+    x : array, shape (n, d)
+        First set of points.
+    y : array, shape (m, d)
+        Second set of points.
+    epsilon : float, optional
+        Shape parameter the distances are multiplied by before being fed to the
+        kernel.  ``None`` (default) leaves the distances unscaled.
+
+    Returns
+    -------
+    dist : array, shape (n, m)
+        Pairwise distances, on the active backend (NumPy or CuPy).
+
+    Notes
+    -----
+    The dot-product identity costs one (n, m) temporary instead of the (n, m, d)
+    one an explicit difference would need, which is what makes large problems fit
+    in device memory.  The price is that a distance that is exactly zero (the
+    diagonal of a self-distance matrix, or duplicated points) comes out as
+    round-off instead: up to about ``sqrt(machine epsilon) * ||x||``, i.e. ~1e-8
+    for unit-scale coordinates.  That is harmless for every kernel here (the
+    smallest kernel value it can produce is ``r**2 * log(r)`` ~ 1e-15), but
+    anything that needs *exact* zeros, such as :func:`default_epsilon`, must use
+    the explicit difference instead.
+    """
+    x = xp.asarray(x, dtype=xp.float64)
+    y = xp.asarray(y, dtype=xp.float64)
+
+    # ||x - y||^2 = ||x||^2 + ||y||^2 - 2 x.y
+    x_sqnorm = (x**2).sum(axis=1, keepdims=True)  # (n, 1)
+    y_sqnorm = (y**2).sum(axis=1, keepdims=True)  # (m, 1)
+
+    dist_sq = x_sqnorm + y_sqnorm.T - 2 * (x @ y.T)  # (n, m)
+    # Round-off can make the above slightly negative for coincident points.
+    dist_sq = xp.clip(dist_sq, 0, None)
+    dist = xp.sqrt(dist_sq)
+
+    if epsilon is not None:
+        dist = dist * epsilon
+
+    return dist
+
+
+def default_epsilon(y) -> float:
+    """
+    Data-driven default shape parameter, for the kernels that need one.
+
+    The classic nearest neighbour heuristic: the square root of the mean distance
+    between each point and its closest distinct neighbour.  Small ``epsilon``
+    means wide kernels, large ``epsilon`` means narrow ones.
+
+    Parameters
+    ----------
+    y : array, shape (n, d)
+        Data point coordinates, with ``n >= 2``.
+
+    Returns
+    -------
+    epsilon : float
+        Positive shape parameter.
+
+    Raises
+    ------
+    ValueError
+        If there are fewer than two points, or if the estimate comes out
+        non-finite or non-positive (every point coincident with another).
+
+    Notes
+    -----
+    Distances are computed with the explicit difference, block by block, so that
+    the self-distances are *exactly* zero and can be discarded: the round-off of
+    :func:`euclidean_distance` (~1e-8) would otherwise be mistaken for the
+    nearest neighbour distance.  The blocks bound the temporary memory to a few
+    tens of MiB whatever ``n`` is.  The final reduction to a Python float forces
+    one host synchronisation, paid once at build time.
+    """
+    y = xp.asarray(y, dtype=xp.float64)
+    if y.ndim == 1:
+        y = xp.reshape(y, (y.shape[0], 1))
+    npoints, ndim = y.shape
+    if npoints < 2:
+        raise ValueError(
+            f"`default_epsilon` needs at least two points, got {npoints}: pass "
+            "`epsilon` explicitly."
+        )
+
+    block = max(1, int(_EPSILON_BLOCK_FLOATS // max(npoints * ndim, 1)))
+    nearest_sum = xp.zeros((), dtype=xp.float64)
+
+    for start in range(0, npoints, block):
+        stop = min(start + block, npoints)
+        chunk = y[start:stop]
+
+        difference = chunk[:, None, :] - y[None, :, :]  # (block, npoints, ndim)
+        dist = xp.sqrt((difference**2).sum(axis=-1))
+        del difference
+
+        index = xp.arange(stop - start)
+        dist[index, index + start] = xp.inf  # a point is not its own neighbour
+        dist = xp.where(dist == 0, xp.inf, dist)  # ...nor are its exact copies
+        nearest_sum = nearest_sum + dist.min(axis=1).sum()
+        del dist
+
+    epsilon = float(xp.sqrt(nearest_sum / npoints))
+    if not math.isfinite(epsilon) or epsilon <= 0:
+        raise ValueError(
+            "Cannot estimate a default shape parameter from `y` (all the points "
+            "are coincident). Pass `epsilon` explicitly."
+        )
+    return epsilon

rbfgpu/_kernels.py

@@ -0,0 +1,220 @@
+"""
+Radial basis function kernels.
+
+Every kernel maps the (already ``epsilon``-scaled) distance ``r`` onto a scalar
+weight; an RBF interpolant is a linear combination of those weights centred on
+the data points, plus a polynomial tail.
+
+Names, signs and defaults follow ``scipy.interpolate.RBFInterpolator`` so that
+an interpolant built here and one built by SciPy agree to round-off (see
+``test/test_scipy_parity.py``).
+
+References
+----------
+[1] Fasshauer, G., 2007. Meshfree Approximation Methods with Matlab.
+[2] Wahba, G., 1990. Spline Models for Observational Data.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Any, Callable, NamedTuple, Optional
+
+from ._backend import xp
+
+__all__ = [
+    "KERNELS",
+    "KERNEL_ALIASES",
+    "Kernel",
+    "available_kernels",
+    "resolve_kernel",
+]
+
+
+class Kernel(NamedTuple):
+    """
+    A radial basis function and the polynomial degree it needs.
+
+    Attributes
+    ----------
+    name : str
+        Canonical kernel name (also the key in :data:`KERNELS`).
+    function : callable
+        Maps an array of distances ``r`` to an array of weights.
+    scale_invariant : bool
+        True when ``epsilon`` only rescales the interpolant, so it can safely be
+        left at 1 (``linear``, ``thin_plate_spline``, ``cubic``, ``quintic``).
+    min_degree : int or None
+        Minimum polynomial degree for a well-posed interpolant, or ``None`` for a
+        user supplied callable (then both ``epsilon`` and ``degree`` must be
+        given explicitly).  ``-1`` means "no constraint" (positive definite
+        kernels).
+    """
+
+    name: str
+    function: Callable[[Any], Any]
+    scale_invariant: bool
+    min_degree: Optional[int]
+
+    @property
+    def default_degree(self) -> Optional[int]:
+        """
+        Polynomial degree used when the caller does not pass one.
+
+        The minimum degree required by the kernel, or 0 when the kernel needs no
+        polynomial at all; ``None`` for custom callables (no default).
+        """
+        if self.min_degree is None:
+            return None
+        return max(self.min_degree, 0)
+
+
+def _thin_plate_spline(r):
+    """``r**2 * log(r)``, with the r -> 0 limit (0) applied elementwise."""
+    # log(0) = -inf and 0 * -inf = nan: both are cleaned up to the limit value.
+    with xp.errstate(divide="ignore", invalid="ignore"):
+        result = r**2 * xp.log(r)
+    return xp.nan_to_num(result, nan=0.0, posinf=0.0, neginf=0.0)
+
+
+def _linear(r):
+    """``-r``."""
+    return -r
+
+
+def _cubic(r):
+    """``r**3``."""
+    return r**3
+
+
+def _quintic(r):
+    """``-r**5``."""
+    return -(r**5)
+
+
+def _multiquadric(r):
+    """``-sqrt(1 + r**2)``."""
+    return -xp.sqrt(r**2 + 1)
+
+
+def _inverse_multiquadric(r):
+    """``1/sqrt(1 + r**2)``."""
+    return 1 / xp.sqrt(r**2 + 1)
+
+
+def _inverse_quadratic(r):
+    """``1/(1 + r**2)``."""
+    return 1 / (r**2 + 1)
+
+
+def _gaussian(r):
+    """``exp(-r**2)``."""
+    return xp.exp(-(r**2))
+
+
+def _matern1_2(r):
+    """``exp(-r)``."""
+    return xp.exp(-r)
+
+
+def _matern3_2(r):
+    """``(1 + sqrt(3) r) exp(-sqrt(3) r)``."""
+    s = math.sqrt(3.0) * r
+    return (1 + s) * xp.exp(-s)
+
+
+def _matern5_2(r):
+    """``(1 + sqrt(5) r + 5/3 r**2) exp(-sqrt(5) r)``."""
+    s = math.sqrt(5.0) * r
+    return (1 + s + s**2 / 3) * xp.exp(-s)
+
+
+# name, function, scale invariant, minimum polynomial degree
+_KERNEL_DEFS = (
+    ("linear", _linear, True, 0),
+    ("thin_plate_spline", _thin_plate_spline, True, 1),
+    ("cubic", _cubic, True, 1),
+    ("quintic", _quintic, True, 2),
+    ("multiquadric", _multiquadric, False, 0),
+    ("inverse_multiquadric", _inverse_multiquadric, False, -1),
+    ("inverse_quadratic", _inverse_quadratic, False, -1),
+    ("gaussian", _gaussian, False, -1),
+    ("matern1_2", _matern1_2, False, -1),
+    ("matern3_2", _matern3_2, False, -1),
+    ("matern5_2", _matern5_2, False, -1),
+)
+
+KERNELS: dict = {
+    name: Kernel(name, function, scale_invariant, min_degree)
+    for name, function, scale_invariant, min_degree in _KERNEL_DEFS
+}
+
+#: Accepted spellings that are not canonical names.
+KERNEL_ALIASES: dict = {
+    "inverse_quartic": "inverse_quadratic",  # SciPy < 1.15 spelling
+    "tps": "thin_plate_spline",
+    "thin_plate": "thin_plate_spline",
+    "imq": "inverse_multiquadric",
+}
+
+
+def available_kernels() -> tuple:
+    """
+    Canonical names of every built-in kernel.
+
+    Returns
+    -------
+    names : tuple of str
+        Sorted kernel names, as accepted by ``RBFInterpolator(kernel=...)``.
+    """
+    return tuple(sorted(KERNELS))
+
+
+def resolve_kernel(kernel) -> Kernel:
+    """
+    Turn a kernel specification into a :class:`Kernel`.
+
+    Parameters
+    ----------
+    kernel : str or callable
+        A kernel name (aliases in :data:`KERNEL_ALIASES` are accepted, case
+        insensitive) or a callable mapping an array of distances to an array of
+        weights.
+
+    Returns
+    -------
+    Kernel
+        The resolved kernel.  A callable is wrapped in a ``Kernel`` whose
+        ``min_degree`` is ``None``, i.e. one that requires ``epsilon`` and
+        ``degree`` to be given explicitly.
+
+    Raises
+    ------
+    TypeError
+        If ``kernel`` is neither a string nor a callable.
+    ValueError
+        If the name is unknown.
+    """
+    if isinstance(kernel, str):
+        name = kernel.strip().lower()
+        name = KERNEL_ALIASES.get(name, name)
+        try:
+            return KERNELS[name]
+        except KeyError:
+            raise ValueError(
+                f"Unknown kernel {kernel!r}. Available kernels: "
+                f"{', '.join(available_kernels())}."
+            ) from None
+
+    if callable(kernel):
+        return Kernel(
+            name=getattr(kernel, "__name__", "<callable>"),
+            function=kernel,
+            scale_invariant=False,
+            min_degree=None,
+        )
+
+    raise TypeError(
+        "`kernel` must be a kernel name (str) or a callable taking the distance "
+        f"array, got {type(kernel).__name__}."
+    )

rbfgpu/_polynomial.py

@@ -0,0 +1,104 @@
+"""
+Polynomial tail of the RBF interpolant.
+
+The RBF part alone cannot extrapolate and, for conditionally positive definite
+kernels, is not even well posed: a polynomial of degree ``degree`` is added to
+the interpolant and constrained to be orthogonal to the RBF coefficients
+(``P(y)^T a = 0``).  The basis used here is the plain monomial basis
+``x**p`` with ``|p| <= degree``.
+
+Any basis spanning the same polynomial space gives the same interpolant, so the
+column order below is arbitrary (it differs from SciPy's, which also shifts and
+scales the coordinates): the coefficients are consistent because the very same
+``powers`` are used to build and to evaluate the system.
+"""
+
+from __future__ import annotations
+
+import itertools
+
+from ._backend import xp
+
+__all__ = ["monomial_exponents", "polynomial_matrix"]
+
+
+def monomial_exponents(ndim: int, degree: int) -> tuple:
+    """
+    Exponent tuples of every monomial of total degree <= ``degree``.
+
+    Parameters
+    ----------
+    ndim : int
+        Number of coordinates of the points.
+    degree : int
+        Polynomial degree.  ``-1`` means "no polynomial terms at all", ``0``
+        only the constant term.
+
+    Returns
+    -------
+    powers : tuple of tuple of int
+        ``len(powers)`` is the number of polynomial terms, i.e.
+        ``C(ndim + degree, degree)``.  Empty for ``degree == -1``.
+
+    Raises
+    ------
+    ValueError
+        If ``degree < -1``, or ``ndim < 1``.
+    """
+    if not isinstance(degree, (int,)) or isinstance(degree, bool):
+        degree = int(degree)
+    if degree < -1:
+        raise ValueError(f"`degree` must be >= -1 (-1 means no polynomial), got {degree}.")
+    if ndim < 1:
+        raise ValueError(f"`ndim` must be >= 1, got {ndim}.")
+    if degree == -1:
+        return ()
+
+    return tuple(
+        powers
+        for powers in itertools.product(range(degree + 1), repeat=ndim)
+        if sum(powers) <= degree
+    )
+
+
+def polynomial_matrix(x, powers) -> object:
+    """
+    Evaluate a monomial basis at the points ``x``.
+
+    Parameters
+    ----------
+    x : array, shape (n, d)
+        Point coordinates.
+    powers : sequence of tuple of int
+        Exponents of the monomials, as returned by :func:`monomial_exponents`.
+
+    Returns
+    -------
+    P : array, shape (n, len(powers))
+        ``P[i, j]`` is ``prod_k x[i, k] ** powers[j][k]``.  For an empty
+        ``powers`` an ``(n, 0)`` array is returned.
+    """
+    x = xp.asarray(x)
+    n, ndim = x.shape
+
+    if len(powers) == 0:
+        return xp.empty((n, 0), dtype=x.dtype)
+
+    columns = []
+    for exponents in powers:
+        if len(exponents) != ndim:
+            raise ValueError(
+                f"`powers` holds {len(exponents)}-dimensional exponents but `x` "
+                f"has {ndim} columns."
+            )
+        column = xp.ones(n, dtype=x.dtype)
+        for axis, exponent in enumerate(exponents):
+            if exponent == 0:
+                continue
+            if exponent == 1:
+                column = column * x[:, axis]
+            else:
+                column = column * x[:, axis] ** exponent
+        columns.append(column)
+
+    return xp.column_stack(columns)