Add fake cupyx.scipy and batched xp.optimize.newton

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

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` |

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),
 }
 
 

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`
 

docs/source/reference/index.md

@@ -14,6 +14,7 @@ arguments
 cuda
 rng
 algorithms
+optimize
 mpi
 profiling
 memory

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:
+```

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.<subpackage>.<name>` (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/reduce.cuh>`: `cunumpy_block_sum_to(out, v)`, `cunumpy_block_min/max`, `cunumpy_warp_sum` |

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",

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

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]