diff --git a/CHANGELOG.md b/CHANGELOG.md index 298b86e..97bd2cd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] ### Added +- `xp.optimize.newton`: Newton's method or the secant method on many + independent scalar equations at once, with the steps of + `scipy.optimize.newton` for an array `x0`, on the array's backend (no host + copies, fixed shapes, one synchronization per iteration). +- The fake CuPy has `cupyx.scipy`: SciPy behind the CuPy rules (host arrays + rejected, fake device arrays and proxy objects returned), so `xp.scipy` code + runs its CuPy path without a GPU. `cupyx.scipy.interpolate` has only the + names CuPy provides. `cupyx.empty_pinned`/`zeros_pinned` return NumPy arrays. +- Solvers guide: splines on the device with `xp.scipy.interpolate`, and + batched root finding. - Compiled Pyccel regression coverage for F-ordered arrays and column-block copy-back, run in CPU CI with fake CuPy and required compilation. - Publishing now requires CPU and GPU validation of the release commit, checks diff --git a/docs/source/api.md b/docs/source/api.md index fb00929..3771573 100644 --- a/docs/source/api.md +++ b/docs/source/api.md @@ -49,6 +49,7 @@ The submodules are named so that they do not hide a NumPy name (`rng`, not | [`cunumpy.cuda`](reference/cuda.md) | CUDA only | device selection and memory, `stream`, streams/events, `pin_memory`, debug mode, CUDA headers and source tools (`cuda_include_dir`, `parse_cuda_signature`) | | [`cunumpy.rng`](reference/rng.md) | both | `random_streams`, `get_rng`, `philox_*` | | [`cunumpy.algorithms`](reference/algorithms.md) | both | `morton_*`, `sort_by_key`, `segment_sum`, `compact_by_mask` | +| [`cunumpy.optimize`](reference/optimize.md) | both | `newton` (batched Newton/secant), `NewtonResult` | | [`cunumpy.mpi`](reference/mpi.md) | both | `mpi_buffer`, `MPIStaging`, CUDA-aware MPI detection, `synchronize_for_mpi`, and the MPI stand-ins of maybempi (`get_mpi`, `local_rank`, ...) | | [`cunumpy.profiling`](reference/profiling.md) | both | `timed_region`, `nvtx_range`, `count_transfers`, `assert_no_transfers`, `TransferBudget` | | [`cunumpy.memory`](reference/memory.md) | both | `HostStaging`, `HostCopy`, `DeviceMirror` | diff --git a/docs/source/conf.py b/docs/source/conf.py index ffc2437..5f966c0 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -31,6 +31,7 @@ intersphinx_mapping = { "python": ("https://docs.python.org/3", None), "numpy": ("https://numpy.org/doc/stable/", None), + "scipy": ("https://docs.scipy.org/doc/scipy/", None), } diff --git a/docs/source/guides/solvers.md b/docs/source/guides/solvers.md index 4d6f84a..f4b874c 100644 --- a/docs/source/guides/solvers.md +++ b/docs/source/guides/solvers.md @@ -3,8 +3,9 @@ Particle pushes and deposits are only part of a plasma code. Field solves need sparse matrices, iterative solvers and FFTs; fluid (MHD) updates are long chains of pointwise operations; diagnostics reduce over all cells or -particles. This page covers the three tools for that: `xp.scipy`, `xp.kernels.fuse` -and `xp.petsc.petsc_vec`, plus in-kernel reductions with `cunumpy/reduce.cuh`. +particles. This page covers the tools for that: `xp.scipy` (including splines), +`xp.optimize.newton`, `xp.kernels.fuse` and `xp.petsc.petsc_vec`, plus +in-kernel reductions with `cunumpy/reduce.cuh`. ## SciPy on both backends: `xp.scipy` @@ -52,6 +53,56 @@ def periodic_poisson(rho, length): * Special functions for distribution functions and cross sections (`erf`, `erfc`, Bessel functions `i0`, `i1`, `k0`, ...) are in `xp.scipy.special` on both backends. +* On the fake CuPy (tests without a GPU), `xp.scipy` runs SciPy with the CuPy + rules: host arrays are rejected, results are fake device arrays, and + `interpolate` has only the names CuPy has. + +## Splines and interpolation + +`xp.scipy.interpolate` fits and evaluates splines on the active backend. +`cupyx.scipy.interpolate` has `UnivariateSpline` (and its interpolating and +least-squares variants), `BSpline`, `make_interp_spline`, `make_lsq_spline`, +`NdBSpline`, `CubicSpline`, `PchipInterpolator`, `RegularGridInterpolator`, +... but not the FITPACK 2D classes (`RectBivariateSpline`, ...) or +`splrep`/`splev`. An interpolating 2D spline on a grid is +`make_interp_spline` along each axis plus `NdBSpline`: + +```python +interp = xp.scipy.interpolate +bx = interp.make_interp_spline(x, values, k=3) # fit along x (axis 0) +by = interp.make_interp_spline(y, bx.c.T, k=3) # then along y +spline = interp.NdBSpline((bx.t, by.t), by.c.T, (3, 3)) +# RectBivariateSpline clamps points to the grid; NdBSpline extrapolates +points = xp.stack([xp.clip(R, x[0], x[-1]), xp.clip(Z, y[0], y[-1])], axis=-1) +dpsi_dR = spline(points, nu=(1, 0)) +``` + +With `s=0` this is the spline `RectBivariateSpline(x, y, values, s=0)` fits +(same knots, values equal to round-off). + +Fit once at setup and keep the spline object; evaluation then needs no +transfers. + +## Many root-finding problems at once: `xp.optimize.newton` + +SciPy's optimizers run on the host, and `cupyx.scipy` has no `optimize`. +Sequential scalar problems (one `fsolve`, a `quad` integral, an ODE) belong on +the host at setup (see `xp.setup_on_host`). Many independent scalar equations +(a flux surface on every ray, an inverse mapping at every marker) are one +array problem: `xp.optimize.newton` runs Newton's method (with `fprime`) or +the secant method on all of them at once, with the steps of +`scipy.optimize.newton` for an array `x0`: + +```python +# radius where the flux reaches each level along each ray +def residual(r): + return psi(R0 + r * cos_theta, Z0 + r * sin_theta) - levels + +r = xp.optimize.newton(residual, xp.full(levels.shape, 0.3)) +``` + +Each iteration synchronizes once to test convergence. +`full_output=True` returns `root`, `converged` and `zero_der` arrays. ## Fused elementwise updates: `xp.kernels.fuse` diff --git a/docs/source/reference/index.md b/docs/source/reference/index.md index edf425b..5b9c93f 100644 --- a/docs/source/reference/index.md +++ b/docs/source/reference/index.md @@ -14,6 +14,7 @@ arguments cuda rng algorithms +optimize mpi profiling memory diff --git a/docs/source/reference/optimize.md b/docs/source/reference/optimize.md new file mode 100644 index 0000000..abf497f --- /dev/null +++ b/docs/source/reference/optimize.md @@ -0,0 +1,8 @@ +# Root finding: `cunumpy.optimize` + +Batched Newton and secant iterations on either backend. + +```{eval-rst} +.. automodule:: cunumpy.optimize + :members: +``` diff --git a/src/cunumpy/LLM_GUIDE.md b/src/cunumpy/LLM_GUIDE.md index 4064732..83c02bf 100644 --- a/src/cunumpy/LLM_GUIDE.md +++ b/src/cunumpy/LLM_GUIDE.md @@ -99,6 +99,8 @@ https://max-models.github.io/cunumpy/ and in `docs/source/` of the repository. | group arrays/scalars into one kernel argument | `xp.arguments.CudaArguments` (flattened), `xp.arguments.CudaStruct` (C struct), `xp.arguments.CudaStructArguments` (C struct as a class); host kernels take their own argument objects, the caller picks one per backend | | CUDA struct from a Pyccel argument class | `xp.arguments.CudaStruct.from_signature(Cls.__init__, "Name")`, `xp.arguments.write_cuda_header(...)` | | SciPy (sparse, sparse.linalg, fft, special, ndimage, ...) on either backend | `xp.scipy..` (SciPy or `cupyx.scipy`); `xp.scipy.special.available(name)` | +| splines on the device | `xp.scipy.interpolate.UnivariateSpline`, `make_interp_spline`, `NdBSpline` (no `RectBivariateSpline`: `make_interp_spline` per axis + `NdBSpline`) | +| many independent scalar root problems (one per ray, per marker) | `xp.optimize.newton(func, x0_array, fprime=None)` (SciPy's array Newton/secant, on the device) | | chain of elementwise operations as one GPU kernel | `@xp.kernels.fuse` (`cupy.fuse` for CuPy arrays, plain call otherwise) | | PETSc solve on device arrays without copies | `xp.petsc.petsc_vec(array)` (CUDA/HIP petsc4py for CuPy arrays); `xp.synchronize()` around PETSc calls | | reduction inside a CUDA kernel (energy, max velocity) | ``: `cunumpy_block_sum_to(out, v)`, `cunumpy_block_min/max`, `cunumpy_warp_sum` | diff --git a/src/cunumpy/__init__.py b/src/cunumpy/__init__.py index a202070..48106b2 100644 --- a/src/cunumpy/__init__.py +++ b/src/cunumpy/__init__.py @@ -22,6 +22,7 @@ kernels, memory, mpi, + optimize, petsc, profiling, rng, @@ -115,6 +116,7 @@ def require_version(minimum: str) -> None: "memory", "mpi", "numpy_backend", + "optimize", "petsc", "profiling", "require_version", diff --git a/src/cunumpy/__init__.pyi b/src/cunumpy/__init__.pyi index a432d57..8625124 100644 --- a/src/cunumpy/__init__.pyi +++ b/src/cunumpy/__init__.pyi @@ -14,6 +14,7 @@ from cunumpy import cuda as cuda from cunumpy import kernels as kernels from cunumpy import memory as memory from cunumpy import mpi as mpi +from cunumpy import optimize as optimize from cunumpy import petsc as petsc from cunumpy import profiling as profiling from cunumpy import rng as rng diff --git a/src/cunumpy/_fake_cupy.py b/src/cunumpy/_fake_cupy.py index 47d3e09..35a1cc7 100644 --- a/src/cunumpy/_fake_cupy.py +++ b/src/cunumpy/_fake_cupy.py @@ -15,6 +15,8 @@ * arrays have ``data.ptr``, ``device`` and ``__cuda_array_interface__``, so :class:`~cunumpy.arguments.CudaStruct` packing and the argument checks of :class:`~cunumpy.kernels.CudaKernel` work; +* ``cupyx.scipy`` forwards to SciPy under the same rules, with the names of + ``cupyx.scipy.interpolate`` only (see :mod:`cunumpy._fake_cupyx`); * CUDA kernels cannot run: ``RawKernel`` and friends raise ``NotImplementedError`` when called, and :func:`cunumpy.kernel_testing.requires_cupy` skips tests while the fake is active. Inside @@ -37,6 +39,8 @@ import types from pathlib import Path +from cunumpy import _fake_cupyx + __all__ = ["host_buffer", "install", "is_active", "uninstall"] _IMPLEMENTATION = Path(__file__).with_name("_fake_cupy_impl.py") @@ -105,6 +109,7 @@ def install() -> types.ModuleType: sys.modules["cupy"] = module try: exec(code, module.__dict__) # noqa: S102 - our own file, classes named cupy.* + _fake_cupyx.install(module) except BaseException: uninstall() raise @@ -115,7 +120,7 @@ def uninstall() -> None: """Remove the fake ``cupy`` package from ``sys.modules`` (no-op otherwise).""" if not is_active(): return - for name in [n for n in sys.modules if n == "cupy" or n.startswith("cupy.")]: + for name in [n for n in sys.modules if n.split(".")[0] in ("cupy", "cupyx")]: del sys.modules[name] for name in [n for n in sys.modules if n.startswith("array_api_compat.cupy")]: del sys.modules[name] diff --git a/src/cunumpy/_fake_cupyx.py b/src/cunumpy/_fake_cupyx.py new file mode 100644 index 0000000..58b17a0 --- /dev/null +++ b/src/cunumpy/_fake_cupyx.py @@ -0,0 +1,214 @@ +"""The ``cupyx`` package of the fake CuPy: ``cupyx.scipy`` on SciPy, for tests only. + +Installed with the fake ``cupy`` (:func:`cunumpy._fake_cupy.install`), so code +written against ``xp.scipy`` runs its CuPy path without a GPU. Each forwarded +SciPy function, class and method behaves like its CuPy counterpart where it +matters for finding host/device bugs: + +* array arguments must be fake device arrays: NumPy arrays raise, as they do + in CuPy; +* array results are fake device arrays, and objects (a fitted spline, a + factorization) are proxies whose methods follow the same rules; +* ``cupyx.scipy.interpolate`` has only the names CuPy provides + (:data:`INTERPOLATE`), so code using a SciPy name that CuPy lacks (e.g. + ``RectBivariateSpline``) fails here as it would on a GPU. The other + subpackages forward every SciPy name. + +``cupyx.empty_pinned`` and ``cupyx.zeros_pinned`` return NumPy arrays, like +CuPy's (which are backed by pinned memory). +""" + +from __future__ import annotations + +import importlib +import sys +import types +from typing import Any + +import numpy as np + +#: The subpackages of ``cupyx.scipy``; the same as ``_scipy_backend.SUBMODULES``. +SUBPACKAGES = ( + "fft", + "fftpack", + "interpolate", + "linalg", + "ndimage", + "signal", + "sparse", + "sparse.csgraph", + "sparse.linalg", + "spatial", + "special", + "stats", +) + +#: The names of ``cupyx.scipy.interpolate`` (CuPy reference, stable). +INTERPOLATE = frozenset( + { + "Akima1DInterpolator", + "BPoly", + "BSpline", + "BarycentricInterpolator", + "CloughTocher2DInterpolator", + "CubicHermiteSpline", + "CubicSpline", + "InterpolatedUnivariateSpline", + "KroghInterpolator", + "LSQUnivariateSpline", + "LinearNDInterpolator", + "NdBSpline", + "NdPPoly", + "NearestNDInterpolator", + "PPoly", + "PchipInterpolator", + "RBFInterpolator", + "RegularGridInterpolator", + "UnivariateSpline", + "barycentric_interpolate", + "interp1d", + "interpn", + "krogh_interpolate", + "make_interp_spline", + "make_lsq_spline", + "pchip_interpolate", + "splantider", + "splder", + }, +) + +_ALLOWED = {"interpolate": INTERPOLATE} + + +class _Proxy: + """A SciPy object (spline, interpolator, ...) seen as a CuPy object.""" + + def __init__(self, obj: Any, cupy: Any) -> None: + object.__setattr__(self, "_obj", obj) + object.__setattr__(self, "_cupy", cupy) + + def __getattr__(self, name: str) -> Any: + attr = getattr(self._obj, name) + if callable(attr): + return _device_callable( + attr, self._cupy, f"{type(self._obj).__name__}.{name}" + ) + return _to_device(attr, self._cupy) + + def __setattr__(self, name: str, value: Any) -> None: + setattr(self._obj, name, _to_host(value, self._cupy)) + + def __call__(self, *args: Any, **kwargs: Any) -> Any: + where = f"{type(self._obj).__name__}.__call__" + return _device_callable(self._obj, self._cupy, where)(*args, **kwargs) + + def __repr__(self) -> str: + return f"" + + +class _ProxyClass: + """A SciPy class seen as a CuPy class: instances are :class:`_Proxy`.""" + + def __init__(self, cls: type, cupy: Any, where: str) -> None: + self._cls = cls + self._cupy = cupy + self._where = where + self.__name__ = cls.__name__ + self.__doc__ = cls.__doc__ + + def __call__(self, *args: Any, **kwargs: Any) -> Any: + return _device_callable(self._cls, self._cupy, self._where)(*args, **kwargs) + + def __getattr__(self, name: str) -> Any: + attr = getattr(self._cls, name) + if callable(attr): + return _device_callable(attr, self._cupy, f"{self._where}.{name}") + return attr + + def __instancecheck__(self, obj: Any) -> bool: + return isinstance(obj, _Proxy) and isinstance(obj._obj, self._cls) + + def __repr__(self) -> str: + return f"" + + +def _is_scipy_object(value: Any) -> bool: + return type(value).__module__.split(".")[0] == "scipy" + + +def _to_device(value: Any, cupy: Any) -> Any: + if isinstance(value, (tuple, list)): + return type(value)(_to_device(v, cupy) for v in value) + if _is_scipy_object(value) and not isinstance(value, np.ndarray): + return _Proxy(value, cupy) + return cupy._wrap(value) + + +def _to_host(value: Any, cupy: Any) -> Any: + if isinstance(value, _Proxy): + return value._obj + if isinstance(value, (tuple, list)): + return type(value)(_to_host(v, cupy) for v in value) + if isinstance(value, dict): + return {k: _to_host(v, cupy) for k, v in value.items()} + return cupy._unwrap(value) + + +def _device_callable(fun: Any, cupy: Any, where: str) -> Any: + def call(*args: Any, **kwargs: Any) -> Any: + cupy._strict(args, where) + cupy._strict(tuple(kwargs.values()), where) + out = fun(*_to_host(args, cupy), **_to_host(kwargs, cupy)) + return _to_device(out, cupy) + + call.__name__ = getattr(fun, "__name__", where) + call.__doc__ = getattr(fun, "__doc__", None) + return call + + +def _subpackage(path: str, cupy: Any) -> types.ModuleType: + module = types.ModuleType(f"cupyx.scipy.{path}") + allowed = _ALLOWED.get(path) + + def __getattr__(name: str) -> Any: + if name.startswith("__") or (allowed is not None and name not in allowed): + raise AttributeError( + f"module 'cupyx.scipy.{path}' has no attribute {name!r}" + ) + attr = getattr(importlib.import_module(f"scipy.{path}"), name, None) + if attr is None or isinstance(attr, types.ModuleType): + raise AttributeError( + f"module 'cupyx.scipy.{path}' has no attribute {name!r}" + ) + where = f"cupyx.scipy.{path}.{name}" + if isinstance(attr, type): + return _ProxyClass(attr, cupy, where) + if callable(attr): + return _device_callable(attr, cupy, where) + return attr + + module.__getattr__ = __getattr__ + return module + + +def install(cupy: Any) -> None: + """Install ``cupyx`` and ``cupyx.scipy.*`` into ``sys.modules``, on the fake `cupy`.""" + cupyx = types.ModuleType("cupyx") + cupyx.__cunumpy_fake__ = True + cupyx.empty_pinned = lambda shape, dtype=float, order="C": np.empty( + shape, dtype, order + ) + cupyx.zeros_pinned = lambda shape, dtype=float, order="C": np.zeros( + shape, dtype, order + ) + scipy = types.ModuleType("cupyx.scipy") + cupyx.scipy = scipy + modules = {"cupyx": cupyx, "cupyx.scipy": scipy} + for path in SUBPACKAGES: + modules[f"cupyx.scipy.{path}"] = _subpackage(path, cupy) + # parents expose their subpackages as attributes, like real packages + for name, module in modules.items(): + parent, _, leaf = name.rpartition(".") + if parent in modules: + setattr(modules[parent], leaf, module) + sys.modules.update(modules) diff --git a/src/cunumpy/_optimize.py b/src/cunumpy/_optimize.py new file mode 100644 index 0000000..1bde879 --- /dev/null +++ b/src/cunumpy/_optimize.py @@ -0,0 +1,162 @@ +"""Batched root finding on either backend (see :mod:`cunumpy.optimize`).""" + +from __future__ import annotations + +import warnings +from collections.abc import Callable +from typing import Any, NamedTuple + +from cunumpy.xp import array_backend, get_array_backend, get_array_module + + +class NewtonResult(NamedTuple): + """The result of :func:`newton` with ``full_output=True``. + + Attributes + ---------- + root : array + The estimated roots. + converged : array of bool + Where the last step was smaller than the tolerance (or ``func`` was + exactly zero). + zero_der : array of bool + Where the iteration stopped on a zero derivative (Newton) or on equal + function values (secant) without converging. + """ + + root: Any + converged: Any + zero_der: Any + + +def _module(x0: Any) -> Any: + if get_array_backend(x0) == "cupy": + return get_array_module(x0) + return array_backend.xp + + +def newton( + func: Callable[..., Any], + x0: Any, + fprime: Callable[..., Any] | None = None, + args: tuple = (), + tol: float = 1.48e-8, + maxiter: int = 50, + rtol: float = 0.0, + full_output: bool = False, +) -> Any: + """Find a root of each of many independent scalar equations at once. + + Newton's method if `fprime` is given, the secant method otherwise, on an + array of starting points: ``func(x)`` returns ``f_i(x_i)`` for each entry. + The steps of :func:`scipy.optimize.newton` with an array `x0`, but on the + array's own backend (NumPy or CuPy) with fixed shapes (no boolean + indexing), so every step runs on the device. Points that have converged + keep taking (tiny) steps until all have, as in SciPy. Each iteration + synchronizes once, to test for convergence. + + Parameters + ---------- + func : callable + ``func(x, *args)``: the residuals at `x`, an array of `x`'s shape. + x0 : array_like + The starting points. A NumPy or CuPy array stays on its backend; a + scalar or list goes to the active backend. + fprime : callable, optional + ``fprime(x, *args)``: the derivatives ``f_i'(x_i)``. Without it the + secant method is used. + args : tuple, optional + Extra arguments of `func` and `fprime`. + tol : float, optional + Absolute tolerance on the step. + maxiter : int, optional + Maximum number of iterations. + rtol : float, optional + Relative tolerance on the step: a point has converged when + ``abs(step) < tol + rtol * abs(x)``. + full_output : bool, optional + Return a :class:`NewtonResult` instead of only the roots. + + Returns + ------- + array or NewtonResult + The roots, of `x0`'s shape (a 0-d array for a scalar `x0`), or the + :class:`NewtonResult` if `full_output` is true. + + Warns + ----- + RuntimeWarning + Without `full_output`, if some points did not converge. + + See Also + -------- + scipy.optimize.newton : The SciPy version (host only). + + Examples + -------- + >>> s = xp.asarray([2.0, 3.0, 4.0]) + >>> xp.optimize.newton(lambda x: x**2 - s, xp.ones(3), fprime=lambda x: 2 * x) + array([1.41421356, 1.73205081, 2. ]) + >>> xp.optimize.newton(lambda x: x**3 - s, xp.ones(3)) + array([1.25992105, 1.44224957, 1.58740105]) + """ + xpm = _module(x0) + p = xpm.astype(xpm.asarray(x0), xpm.float64, copy=True) + failures = xpm.ones(p.shape, dtype=bool) + nz_der = xpm.ones(p.shape, dtype=bool) + + def tolerance(x: Any) -> Any: + return tol + rtol * xpm.abs(x) + + if fprime is not None: + for _ in range(maxiter): + fval = func(p, *args) + if not bool(xpm.any(fval != 0)): + failures = fval != 0 + break + fder = fprime(p, *args) + nz_der = fder != 0 + if not bool(xpm.any(nz_der)): + break + dp = fval / xpm.where(nz_der, fder, 1.0) + p = xpm.where(nz_der, p - dp, p) + failures = xpm.where(nz_der, xpm.abs(dp) >= tolerance(p), failures) + if not bool(xpm.any(failures & nz_der)): + break + else: + dx = float(xpm.finfo(xpm.float64).eps) ** 0.33 + p1 = p * (1 + dx) + xpm.where(p >= 0, dx, -dx) + q0 = func(p, *args) + q1 = func(p1, *args) + active = xpm.ones(p.shape, dtype=bool) + for _ in range(maxiter): + nz_der = q1 != q0 + if not bool(xpm.any(nz_der)): + p = (p1 + p) / 2.0 + break + dp = q1 * (p1 - p) / xpm.where(nz_der, q1 - q0, 1.0) + new = xpm.where(nz_der, p1 - dp, p) + new = xpm.where(~nz_der & active, (p1 + p) / 2.0, new) + p = new + active = active & nz_der + failures = xpm.where(nz_der, xpm.abs(dp) >= tolerance(p), failures) + if not bool(xpm.any(failures & nz_der)): + break + p1, p = p, p1 + q0 = q1 + q1 = func(p1, *args) + else: + p = p1 # the newest iterate, after the last swap + + zero_der = ~nz_der & failures + if full_output: + return NewtonResult(p, ~failures, zero_der) + if bool(xpm.any(failures)): + n_failed = int(xpm.sum(failures)) + reason = " (zero derivatives)" if bool(xpm.any(zero_der)) else "" + warnings.warn( + f"newton: {n_failed} of {p.size} points did not converge{reason}", + RuntimeWarning, + stacklevel=2, + ) + return p diff --git a/src/cunumpy/optimize.py b/src/cunumpy/optimize.py new file mode 100644 index 0000000..3014989 --- /dev/null +++ b/src/cunumpy/optimize.py @@ -0,0 +1,14 @@ +"""Root finding on either backend. + +SciPy's optimizers run on the host only, and ``cupyx.scipy`` has no +``optimize``. This module has the parts that profit from the device: many +independent problems solved at once, each step one array operation:: + + r = xp.optimize.newton(lambda r: psi(r) - levels, r0) + +See :doc:`/guides/solvers`. +""" + +from cunumpy._optimize import NewtonResult, newton + +__all__ = ["NewtonResult", "newton"] diff --git a/tests/unit/test_fake_cupyx.py b/tests/unit/test_fake_cupyx.py new file mode 100644 index 0000000..3fa26cf --- /dev/null +++ b/tests/unit/test_fake_cupyx.py @@ -0,0 +1,113 @@ +"""The fake cupyx.scipy: SciPy behind the CuPy rules (device arrays in and out).""" + +from pathlib import Path + +import pytest + +from cunumpy import _fake_cupyx +from cunumpy._scipy_backend import SUBMODULES +from cunumpy.kernel_testing import run_in_fake_cupy_subprocess + +SRC = str(Path(__file__).resolve().parents[2] / "src") + + +def run(code): + pytest.importorskip("scipy") + return run_in_fake_cupy_subprocess(code, env={"PYTHONPATH": SRC}).stdout + + +def test_same_subpackages_as_xp_scipy(): + assert _fake_cupyx.SUBPACKAGES == SUBMODULES + + +def test_splines_fit_and_evaluate_on_the_device(): + code = """ +import numpy as np +import scipy.interpolate as si +import cunumpy as xp +xp.set_backend("cupy") +x_h = np.linspace(0.0, 1.0, 11) +x = xp.asarray(x_h) +y = xp.sin(3 * x) +spl = xp.scipy.interpolate.UnivariateSpline(x, y, k=3, s=0.0, ext=3) +ref = si.UnivariateSpline(x_h, np.sin(3 * x_h), k=3, s=0.0, ext=3) +assert isinstance(spl, xp.scipy.interpolate.UnivariateSpline) +for nu in (0, 1, 2): + out = spl(x, nu=nu) + assert xp.is_gpu(out) + assert np.allclose(xp.to_numpy(out), ref(x_h, nu=nu), rtol=0, atol=1e-14) +assert xp.is_gpu(spl(0.3)) and spl(0.3).ndim == 0 +# a returned spline is a proxy too, and its arrays are device arrays +b = xp.scipy.interpolate.make_interp_spline(x, y, k=3) +assert xp.is_gpu(b.t) and xp.is_gpu(b.derivative(1)(x)) +nd = xp.scipy.interpolate.NdBSpline((b.t, b.t), xp.ones((11, 11)), 3) +assert xp.is_gpu(nd(xp.asarray([[0.2, 0.3]]), nu=(1, 0))) +print("ok") +""" + assert "ok" in run(code) + + +def test_host_arrays_are_rejected_like_cupy(): + code = """ +import numpy as np +import cunumpy as xp +xp.set_backend("cupy") +x = xp.linspace(0.0, 1.0, 11) +spl = xp.scipy.interpolate.UnivariateSpline(x, x**2, s=0.0) +for call in (lambda: spl(np.linspace(0.0, 1.0, 3)), + lambda: xp.scipy.interpolate.UnivariateSpline(np.linspace(0, 1, 11), x, s=0.0), + lambda: xp.scipy.special.erf(np.zeros(2))): + try: + call() + except TypeError: + pass + else: + raise AssertionError("host array accepted") +print("ok") +""" + assert "ok" in run(code) + + +def test_interpolate_has_only_the_names_cupy_provides(): + code = """ +import cunumpy as xp +xp.set_backend("cupy") +assert not xp.scipy.interpolate.available("RectBivariateSpline") +assert not xp.scipy.interpolate.available("splrep") +assert xp.scipy.interpolate.available("NdBSpline") +try: + xp.scipy.interpolate.RectBivariateSpline +except AttributeError as error: + assert "not available on the cupy backend" in str(error) +else: + raise AssertionError +print("ok") +""" + assert "ok" in run(code) + + +def test_other_subpackages_and_pinned_memory(): + code = """ +import numpy as np +import cunumpy as xp +import cupyx +xp.set_backend("cupy") +out = xp.scipy.special.erf(xp.zeros(3)) +assert xp.is_gpu(out) +assert type(cupyx.empty_pinned((2, 3))) is np.ndarray +assert np.all(cupyx.zeros_pinned(4) == 0) +print("ok") +""" + assert "ok" in run(code) + + +def test_uninstall_removes_cupyx(): + code = """ +import sys +from cunumpy import _fake_cupy +assert "cupyx.scipy.interpolate" in sys.modules +_fake_cupy.uninstall() +assert not [n for n in sys.modules if n.split(".")[0] in ("cupy", "cupyx")] +print("ok") +""" + assert "ok" in run(code) diff --git a/tests/unit/test_optimize.py b/tests/unit/test_optimize.py new file mode 100644 index 0000000..b9b1e35 --- /dev/null +++ b/tests/unit/test_optimize.py @@ -0,0 +1,103 @@ +"""xp.optimize.newton: batched Newton and secant, compared with scipy.optimize.newton.""" + +from pathlib import Path + +import numpy as np +import pytest + +import cunumpy as xp +from cunumpy.kernel_testing import run_in_fake_cupy_subprocess + +SRC = str(Path(__file__).resolve().parents[2] / "src") + + +def cubic(levels): + return (lambda x: x**3 - 2 * x - levels), (lambda x: 3 * x**2 - 2) + + +@pytest.mark.parametrize("method", ["newton", "secant"]) +def test_matches_scipy_on_an_array(method): + so = pytest.importorskip("scipy.optimize") + levels = np.random.default_rng(0).uniform(1.0, 10.0, 500) + f, fp = cubic(levels) + kw = {"fprime": fp} if method == "newton" else {} + with xp.use_backend("numpy"): + root = xp.optimize.newton(f, np.ones(500), **kw) + np.testing.assert_allclose( + root, so.newton(f, np.ones(500), **kw), rtol=0, atol=1e-14 + ) + np.testing.assert_allclose(f(root), 0.0, atol=1e-12) + + +def test_scalar_and_list_starting_points_go_to_the_active_backend(): + with xp.use_backend("numpy"): + root = xp.optimize.newton(lambda x: x**2 - 2, 1.0, fprime=lambda x: 2 * x) + assert isinstance(root, np.ndarray) and root.shape == () + assert root == pytest.approx(np.sqrt(2)) + roots = xp.optimize.newton(lambda x: x**2 - 2, [1.0, -1.0]) + np.testing.assert_allclose(roots, [np.sqrt(2), -np.sqrt(2)]) + + +def test_x0_is_not_modified(): + x0 = np.ones(3) + xp.optimize.newton(lambda x: x - 2, x0, fprime=lambda x: np.ones_like(x)) + np.testing.assert_array_equal(x0, 1.0) + + +def test_full_output_matches_scipy(): + so = pytest.importorskip("scipy.optimize") + # x**2 + 1 has no real root; x = 0 has a zero derivative from the start, and + # from x = 1 the first step lands on it + x0 = np.array([2.5, 0.0, 1.0, 1.3]) + kw = {"fprime": lambda x: 2 * x, "maxiter": 20, "full_output": True} + f = lambda x: x**2 + 1 + out = xp.optimize.newton(f, x0, **kw) + with pytest.warns(RuntimeWarning): + ref = so.newton(f, x0, **kw) + assert isinstance(out, xp.optimize.NewtonResult) + np.testing.assert_array_equal(out.converged, ref.converged) + np.testing.assert_array_equal(out.zero_der, ref.zero_der) + np.testing.assert_array_equal(out.zero_der, [False, True, True, False]) + + +def test_partial_failure_warns(): + with pytest.warns(RuntimeWarning, match="1 of 2 points did not converge"): + xp.optimize.newton( + lambda x: x**2 - np.array([4.0, -1.0]), + np.array([1.0, 1.0]), + fprime=lambda x: 2 * x, + maxiter=20, + ) + + +def test_tolerances(): + f, fp = cubic(np.array([5.0])) + loose = xp.optimize.newton(f, np.array([10.0]), fprime=fp, tol=1e-2) + tight = xp.optimize.newton(f, np.array([10.0]), fprime=fp, tol=0.0, rtol=1e-15) + assert abs(f(tight)[0]) < abs(f(loose)[0]) + assert abs(f(tight)[0]) < 1e-12 + + +def test_args_are_passed_to_func_and_fprime(): + levels = np.array([1.0, 4.0, 9.0]) + root = xp.optimize.newton( + lambda x, s: x**2 - s, np.ones(3), fprime=lambda x, s: 2 * x, args=(levels,) + ) + np.testing.assert_allclose(root, [1.0, 2.0, 3.0]) + + +def test_runs_on_the_device_on_the_fake_cupy(): + code = """ +import numpy as np +import cunumpy as xp +xp.set_backend("cupy") +levels = xp.asarray(np.linspace(1.0, 10.0, 50)) +for kw in ({"fprime": lambda x: 3 * x**2 - 2}, {}): + out = xp.optimize.newton(lambda x: x**3 - 2 * x - levels, xp.ones(50), full_output=True, **kw) + assert xp.is_gpu(out.root) and xp.is_gpu(out.converged) + assert bool(xp.all(out.converged)) + assert float(xp.max(xp.abs(out.root**3 - 2 * out.root - levels))) < 1e-12 +assert xp.is_gpu(xp.optimize.newton(lambda x: x**2 - 2, 1.0)) +print("ok") +""" + assert "ok" in run_in_fake_cupy_subprocess(code, env={"PYTHONPATH": SRC}).stdout