diff --git a/CHANGELOG.md b/CHANGELOG.md index 7dd14f5..652f241 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,14 +8,22 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] ### Added +- MIT license (`LICENSE.txt`, referenced by `pyproject.toml`). - `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. + runs its CuPy path without a GPU. Every `cupyx.scipy` subpackage has only + the names CuPy v14.2.0 provides, so e.g. `special.jv` fails as on a GPU. + `cupyx.empty_pinned`/`zeros_pinned` return NumPy arrays. +- `xp.optimize.fsolve`, `root`, `minimize` and `xp.integrate.quad`, `odeint` + (`cupyx.scipy` has neither): SciPy on the host with SciPy's arguments, the + callbacks get their arrays on the backend of `x0`/`y0`, and results come + back on that backend. +- `xp.scipy.linalg.solve_circulant` on the CuPy backend, which `cupyx.scipy` + lacks: SciPy's FFT solve and arguments, on the device. - 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 diff --git a/LICENSE.txt b/LICENSE.txt new file mode 100644 index 0000000..4b4eeec --- /dev/null +++ b/LICENSE.txt @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2025 Max Lindqvist + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/docs/source/api.md b/docs/source/api.md index 3771573..fdb68db 100644 --- a/docs/source/api.md +++ b/docs/source/api.md @@ -49,7 +49,8 @@ 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.optimize`](reference/optimize.md) | both | `newton` (batched Newton/secant), `NewtonResult`; `fsolve`, `root`, `minimize` (SciPy on the host, callbacks on the device) | +| [`cunumpy.integrate`](reference/integrate.md) | both | `quad`, `odeint` (SciPy on the host, callbacks on the device) | | [`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/guides/solvers.md b/docs/source/guides/solvers.md index f4b874c..ee2d948 100644 --- a/docs/source/guides/solvers.md +++ b/docs/source/guides/solvers.md @@ -44,6 +44,12 @@ def periodic_poisson(rho, length): * `cupyx.scipy` has only part of SciPy. A missing name raises `AttributeError` saying which backend lacks it; check with `xp.scipy.special.available("erfcx")` where a fallback is possible. + Commonly missing: `special.jv` (only `j0`, `j1`, `yn`), `special.erfi`, + `linalg.solve`/`inv` (use `xp.linalg`), `RectBivariateSpline`, and all of + `integrate` and `optimize` (see `xp.optimize` and `xp.integrate` below). +* cunumpy fills in `xp.scipy.linalg.solve_circulant` on the CuPy backend: an + FFT solve on the device, with SciPy's arguments (a circulant preconditioner + per iteration stays on the GPU). * Keyword arguments can differ between the two: SciPy's `cg` takes `rtol` (since SciPy 1.12), CuPy's takes `tol`. Pass tolerances only after checking both, or through a small wrapper. @@ -54,8 +60,8 @@ def periodic_poisson(rho, length): (`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. + rules: host arrays are rejected, results are fake device arrays, and each + subpackage has only the names CuPy has. ## Splines and interpolation @@ -87,7 +93,7 @@ transfers. 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 +the host at setup (see the next section and `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 @@ -104,6 +110,30 @@ 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. +## SciPy solvers with device callbacks + +For a single problem, `xp.optimize.fsolve`, `root` and `minimize` and +`xp.integrate.quad` and `odeint` run SciPy on the host with SciPy's +arguments, so they work on both backends. With `x0` (`y0` for `odeint`) on +the device, the callbacks get `x` on the device and the result comes back on +the device, so a residual written with `xp` that uses device data needs no +host copy of that data: + +```python +B = equilibrium.b_field_on_grid() # device array + +def residual(x): + return interpolate(B, x) - b_target + +x = xp.optimize.fsolve(residual, xp.asarray([0.5, 0.0])) # device array +``` + +Every evaluation copies `x` to the device and the result to the host, so use +them at setup, not in a time loop. `quad` calls its integrand with a Python +float; a 0-d device array returned is copied to the host. Bounds and +constraint matrices of `minimize` are host data; constraint functions get +device arrays like `fun`. + ## Fused elementwise updates: `xp.kernels.fuse` On the GPU, `(gamma - 1) * (E - 0.5 * rho * u**2)` runs as five kernels, each diff --git a/docs/source/reference/index.md b/docs/source/reference/index.md index 5b9c93f..2c74f22 100644 --- a/docs/source/reference/index.md +++ b/docs/source/reference/index.md @@ -15,6 +15,7 @@ cuda rng algorithms optimize +integrate mpi profiling memory diff --git a/docs/source/reference/integrate.md b/docs/source/reference/integrate.md new file mode 100644 index 0000000..94a36ce --- /dev/null +++ b/docs/source/reference/integrate.md @@ -0,0 +1,8 @@ +# Integration: `cunumpy.integrate` + +SciPy's `quad` and `odeint` with callbacks on the device. + +```{eval-rst} +.. automodule:: cunumpy.integrate + :members: +``` diff --git a/docs/source/reference/optimize.md b/docs/source/reference/optimize.md index abf497f..e1a2f7e 100644 --- a/docs/source/reference/optimize.md +++ b/docs/source/reference/optimize.md @@ -1,6 +1,7 @@ -# Root finding: `cunumpy.optimize` +# Optimization: `cunumpy.optimize` -Batched Newton and secant iterations on either backend. +Batched Newton and secant iterations on the device, and SciPy's `fsolve`, +`root` and `minimize` with callbacks on the device. ```{eval-rst} .. automodule:: cunumpy.optimize diff --git a/src/cunumpy/LLM_GUIDE.md b/src/cunumpy/LLM_GUIDE.md index 83c02bf..0fea615 100644 --- a/src/cunumpy/LLM_GUIDE.md +++ b/src/cunumpy/LLM_GUIDE.md @@ -100,6 +100,8 @@ https://max-models.github.io/cunumpy/ and in `docs/source/` of the repository. | 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`) | +| circulant solve (FFT preconditioner) on the device | `xp.scipy.linalg.solve_circulant(c, b)` (cunumpy fill-in on CuPy, SciPy's arguments) | +| one root/minimization/integral/ODE at setup with an `xp` callback using device data | `xp.optimize.fsolve`/`root`/`minimize`, `xp.integrate.quad`/`odeint` (SciPy on the host, callback args and results on `x0`'s backend) | | 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 | diff --git a/src/cunumpy/__init__.py b/src/cunumpy/__init__.py index 48106b2..1d5247b 100644 --- a/src/cunumpy/__init__.py +++ b/src/cunumpy/__init__.py @@ -19,6 +19,7 @@ algorithms, arguments, cuda, + integrate, kernels, memory, mpi, @@ -110,6 +111,7 @@ def require_version(minimum: str) -> None: "get_array_module", "get_backend", "host_call", + "integrate", "is_cpu", "is_gpu", "kernels", diff --git a/src/cunumpy/__init__.pyi b/src/cunumpy/__init__.pyi index 8625124..9e23036 100644 --- a/src/cunumpy/__init__.pyi +++ b/src/cunumpy/__init__.pyi @@ -11,6 +11,7 @@ from numpy import * from cunumpy import algorithms as algorithms from cunumpy import arguments as arguments from cunumpy import cuda as cuda +from cunumpy import integrate as integrate from cunumpy import kernels as kernels from cunumpy import memory as memory from cunumpy import mpi as mpi diff --git a/src/cunumpy/_fake_cupy.py b/src/cunumpy/_fake_cupy.py index 35a1cc7..a96b1be 100644 --- a/src/cunumpy/_fake_cupy.py +++ b/src/cunumpy/_fake_cupy.py @@ -16,7 +16,7 @@ :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`); + ``cupyx.scipy`` (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 diff --git a/src/cunumpy/_fake_cupyx.py b/src/cunumpy/_fake_cupyx.py index 58b17a0..900b0ef 100644 --- a/src/cunumpy/_fake_cupyx.py +++ b/src/cunumpy/_fake_cupyx.py @@ -9,10 +9,10 @@ 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. +* each subpackage has only the names CuPy provides (:data:`NAMES`), so code + using a SciPy name that CuPy lacks (e.g. ``RectBivariateSpline``, + ``special.jv``, ``linalg.solve_circulant``) fails here as it would on a + GPU. A CuPy name that the installed SciPy lacks is missing too. ``cupyx.empty_pinned`` and ``cupyx.zeros_pinned`` return NumPy arrays, like CuPy's (which are backed by pinned memory). @@ -43,41 +43,139 @@ "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} +def _names(text: str) -> frozenset[str]: + return frozenset(text.split()) + + +#: The public names of each ``cupyx.scipy`` subpackage (CuPy v14.2.0). +NAMES = { + "fft": _names( + """ + dct dctn dst dstn fft fft2 fftfreq fftn fftshift fht get_fft_plan hfft + hfft2 hfftn idct idctn idst idstn ifft ifft2 ifftn ifftshift ifht ihfft + ihfft2 ihfftn irfft irfft2 irfftn next_fast_len rfft rfft2 rfftfreq + rfftn + """, + ), + "fftpack": _names( + """ + fft fft2 fftn get_fft_plan ifft ifft2 ifftn irfft rfft + """, + ), + "interpolate": _names( + """ + 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 + pchip_interpolate splantider splder + """, + ), + "linalg": _names( + """ + bandwidth block_diag circulant companion convolution_matrix dft expm + fiedler fiedler_companion hadamard hankel helmert hilbert khatri_rao + kron leslie lu lu_factor lu_solve solve_triangular toeplitz + """, + ), + "ndimage": _names( + """ + affine_transform binary_closing binary_dilation binary_erosion + binary_fill_holes binary_hit_or_miss binary_opening binary_propagation + black_tophat center_of_mass convolve convolve1d correlate correlate1d + distance_transform_edt extrema find_objects fourier_ellipsoid + fourier_gaussian fourier_shift fourier_uniform gaussian_filter + gaussian_filter1d gaussian_gradient_magnitude gaussian_laplace + generate_binary_structure generic_filter generic_filter1d + generic_gradient_magnitude generic_laplace grey_closing grey_dilation + grey_erosion grey_opening histogram iterate_structure label + labeled_comprehension laplace map_coordinates maximum maximum_filter + maximum_filter1d maximum_position mean median median_filter minimum + minimum_filter minimum_filter1d minimum_position morphological_gradient + morphological_laplace percentile_filter prewitt rank_filter rotate shift + sobel spline_filter spline_filter1d standard_deviation sum sum_labels + uniform_filter uniform_filter1d value_indices variance white_tophat zoom + """, + ), + "signal": _names( + """ + BadCoefficients CZT StateSpace TransferFunction ZerosPolesGain ZoomFFT + abcd_normalize argrelextrema argrelmax argrelmin band_stop_obj bilinear + bilinear_zpk bode buttap butter buttord cheb1ap cheb1ord cheb2ap + cheb2ord cheby1 cheby2 check_COLA check_NOLA chirp choose_conv_method + coherence cont2discrete convolve convolve2d correlate correlate2d + correlation_lags csd cspline1d cspline1d_eval cspline2d cwt czt + czt_points dbode decimate deconvolve detrend dfreqresp dimpulse dlsim + dlti dstep ellip ellipap ellipord fftconvolve filtfilt find_peaks + findfreqs firls firwin firwin2 freqresp freqs freqs_zpk freqz freqz_sos + freqz_zpk gammatone gauss_spline gausspulse get_window group_delay + hilbert hilbert2 iircomb iirdesign iirfilter iirnotch iirpeak impulse + invres invresz istft kaiser_atten kaiser_beta kaiserord lfilter + lfilter_zi lfiltic lombscargle lp2bp lp2bp_zpk lp2bs lp2bs_zpk lp2hp + lp2hp_zpk lp2lp lp2lp_zpk lsim lti max_len_seq medfilt medfilt2d + minimum_phase morlet morlet2 normalize oaconvolve order_filter + peak_prominences peak_widths periodogram place_poles qmf qspline1d + qspline1d_eval qspline2d resample resample_poly residue residuez ricker + savgol_coeffs savgol_filter sawtooth sepfir2d sos2tf sos2zpk sosfilt + sosfilt_zi sosfiltfilt sosfreqz spectrogram spline_filter square ss2tf + ss2zpk step stft sweep_poly symiirorder1 symiirorder2 tf2sos tf2ss + tf2zpk unique_roots unit_impulse upfirdn vectorstrength welch wiener + zoom_fft zpk2sos zpk2ss zpk2tf + """, + ), + "sparse": _names( + """ + SparseEfficiencyWarning SparseWarning block_array block_diag bmat + coo_array coo_matrix csc_array csc_matrix csr_array csr_matrix dia_array + dia_matrix diags diags_array eye eye_array find get_index_dtype hstack + identity issparse isspmatrix isspmatrix_coo isspmatrix_csc + isspmatrix_csr isspmatrix_dia kron kronsum matrix_transpose permute_dims + rand random random_array safely_cast_index_arrays sparray spdiags + spmatrix swapaxes tril triu vstack + """, + ), + "sparse.csgraph": _names( + """ + connected_components + """, + ), + "sparse.linalg": _names( + """ + LinearOperator SuperLU aslinearoperator bicgstab cg cgs eigsh factorized + gmres lobpcg lsmr lsqr matrix_power minres norm spilu splu spsolve + spsolve_triangular svds + """, + ), + "spatial": _names( + """ + Delaunay KDTree distance_matrix + """, + ), + "special": _names( + """ + bdtr bdtrc bdtri beta betainc betaincinv betaln binom boxcox boxcox1p + btdtr btdtri cbrt chdtr chdtrc chdtri cosdg cosm1 cotdg digamma + ellipeinc ellipj ellipk ellipkinc ellipkm1 entr erf erfc erfcinv erfcx + erfinv exp1 exp10 exp2 expi expit expm1 expn exprel fdtr fdtrc fdtri + gamma gammainc gammaincc gammainccinv gammaincinv gammaln gammasgn gdtr + gdtrc huber i0 i0e i1 i1e inv_boxcox inv_boxcox1p j0 j1 k0 k0e k1 k1e + kl_div lambertw log1p log_expit log_ndtr log_softmax loggamma logit + logsumexp lpmv multigammaln nbdtr nbdtrc nbdtri ndtr ndtri pdtr pdtrc + pdtri poch polygamma pseudo_huber psi radian rel_entr rgamma round + shichi sici sinc sindg softmax sph_harm sph_harm_y spherical_yn tandg + wright_bessel xlog1py xlogy y0 y1 yn zeta zetac + """, + ), + "stats": _names( + """ + boxcox_llf entropy trim_mean zmap zscore + """, + ), +} class _Proxy: @@ -168,10 +266,10 @@ def call(*args: Any, **kwargs: Any) -> Any: def _subpackage(path: str, cupy: Any) -> types.ModuleType: module = types.ModuleType(f"cupyx.scipy.{path}") - allowed = _ALLOWED.get(path) + allowed = NAMES[path] def __getattr__(name: str) -> Any: - if name.startswith("__") or (allowed is not None and name not in allowed): + if name.startswith("__") or name not in allowed: raise AttributeError( f"module 'cupyx.scipy.{path}' has no attribute {name!r}" ) diff --git a/src/cunumpy/_linalg.py b/src/cunumpy/_linalg.py new file mode 100644 index 0000000..41e8ffb --- /dev/null +++ b/src/cunumpy/_linalg.py @@ -0,0 +1,103 @@ +"""SciPy linear algebra that ``cupyx.scipy.linalg`` lacks, on the array's backend.""" + +from __future__ import annotations + +from typing import Any + +import numpy as np + +from cunumpy.xp import get_array_module + + +def _moveaxis(xpm: Any, a: Any, source: int, destination: int) -> Any: + if source % a.ndim == destination % a.ndim: + return a + return xpm.moveaxis(a, source, destination) + + +def _as_inexact(xpm: Any, a: Any) -> Any: + a = xpm.asarray(a) + if a.ndim == 0: + a = xpm.reshape(a, (1,)) + if not xpm.isdtype(a.dtype, ("real floating", "complex floating")): + a = xpm.astype(a, xpm.float64) + return a + + +def solve_circulant( + c: Any, + b: Any, + singular: str = "raise", + tol: float | None = None, + caxis: int = -1, + baxis: int = 0, + outaxis: int = 0, +) -> Any: + """Solve ``C x = b`` for a circulant matrix ``C`` with first column `c`. + + The steps and arguments of :func:`scipy.linalg.solve_circulant`, with FFTs + on the backend of `c`: on the CuPy backend nothing leaves the device, and + the singularity check synchronizes once. Available as + ``xp.scipy.linalg.solve_circulant`` on both backends. + + Parameters + ---------- + c : array + First column of ``C``; a batch of columns along the other axes. + b : array + Right-hand side(s), of length ``len(c)`` along `baxis`. + singular : {"raise", "lstsq"}, optional + On a (near) singular ``C``: raise, or return the least-squares + solution. + tol : float, optional + Eigenvalues of ``C`` with magnitude at most `tol` count as zero. + Default ``max(abs(fft(c))) * len(c) * eps``. + caxis, baxis, outaxis : int, optional + The axis of `c`, `b` and the result along which the vectors lie. + + Returns + ------- + array + The solution ``x``, real unless `c` or `b` is complex. + + Raises + ------ + numpy.linalg.LinAlgError + If ``C`` is near singular and `singular` is ``"raise"``. + ValueError + If the lengths of `c` and `b` differ. + + Examples + -------- + >>> c = xp.asarray([4.0, 1.0, 0.0, 1.0]) + >>> xp.scipy.linalg.solve_circulant(c, xp.asarray([6.0, 6.0, 6.0, 6.0])) + array([1., 1., 1., 1.]) + """ + xpm = get_array_module(c) + c = _as_inexact(xpm, c) + b = _as_inexact(xpm, b) + n = c.shape[caxis] + if b.shape[baxis] != n: + raise ValueError(f"Shapes of c {c.shape} and b {b.shape} are incompatible") + fc = xpm.fft.fft(_moveaxis(xpm, c, caxis, -1), axis=-1) + abs_fc = xpm.abs(fc) + if tol is None: + # the tolerance of numpy.linalg.matrix_rank, as in SciPy + tol = xpm.max(abs_fc, axis=-1, keepdims=True) * n * np.finfo(np.float64).eps + near_zeros = abs_fc <= tol + is_near_singular = bool(xpm.any(near_zeros)) + if is_near_singular: + if singular == "raise": + raise np.linalg.LinAlgError("near singular circulant matrix.") + fc = xpm.where(near_zeros, xpm.ones_like(fc), fc) + fb = xpm.fft.fft(_moveaxis(xpm, b, baxis, -1), axis=-1) + q = fb / fc + if is_near_singular: + q = xpm.where(near_zeros, xpm.zeros_like(q), q) + x = xpm.fft.ifft(q, axis=-1) + if not ( + xpm.isdtype(c.dtype, "complex floating") + or xpm.isdtype(b.dtype, "complex floating") + ): + x = xpm.real(x) + return _moveaxis(xpm, x, -1, outaxis) diff --git a/src/cunumpy/_scipy_backend.py b/src/cunumpy/_scipy_backend.py index 3e00026..cc9f527 100644 --- a/src/cunumpy/_scipy_backend.py +++ b/src/cunumpy/_scipy_backend.py @@ -11,7 +11,9 @@ Names are resolved at every access, so a backend switch takes effect at once; nothing is imported until a name is used, and SciPy is not a dependency of cunumpy. ``cupyx.scipy`` covers only part of SciPy, and keyword arguments can -differ (SciPy's ``cg`` takes ``rtol``, CuPy's ``tol``). A SciPy sparse matrix +differ (SciPy's ``cg`` takes ``rtol``, CuPy's ``tol``). A few missing +functions are filled in by cunumpy on the device (``FILLS``, e.g. +``xp.scipy.linalg.solve_circulant``). A SciPy sparse matrix moves to the device once with ``xp.scipy.sparse.csr_matrix(host_matrix)`` on the CuPy backend, and back with ``matrix.get()``. See :doc:`/guides/solvers`. """ @@ -22,9 +24,10 @@ from types import ModuleType from typing import Any +from cunumpy import _linalg from cunumpy.xp import get_backend -__all__ = ["SUBMODULES", "ScipyNamespace", "scipy"] +__all__ = ["FILLS", "SUBMODULES", "ScipyNamespace", "scipy"] #: The SciPy subpackages forwarded by ``xp.scipy``: those that ``cupyx.scipy`` #: provides as well. @@ -43,6 +46,10 @@ "stats", ) +#: Functions ``cupyx.scipy`` lacks, implemented by cunumpy on the device; used +#: on the CuPy backend only, where the ``cupyx`` module has no such name. +FILLS = {"linalg": {"solve_circulant": _linalg.solve_circulant}} + _ROOTS = {"numpy": "scipy", "cupy": "cupyx.scipy"} _INSTALL = { "numpy": "SciPy is not installed (pip install scipy)", @@ -54,7 +61,8 @@ class ScipyNamespace: """A SciPy (sub)package of the active backend; see :data:`cunumpy.scipy`. Attribute access forwards to the module of the active backend. A name the - backend's module lacks raises ``AttributeError`` naming the backend; a + backend's module lacks raises ``AttributeError`` naming the backend, unless + cunumpy fills it in (``FILLS``, CuPy backend only); a missing SciPy (NumPy backend) or CuPy raises ``ImportError``. Parameters @@ -124,6 +132,9 @@ def __getattr__(self, name: str) -> Any: try: return getattr(module, name) except AttributeError: + fill = self._fill(name) + if fill is not None: + return fill other = "cupy" if get_backend() == "numpy" else "numpy" raise AttributeError( f"{module.__name__} has no attribute {name!r}: it is not available " @@ -131,6 +142,11 @@ def __getattr__(self, name: str) -> Any: f"{self._name(_ROOTS[other])})", ) from None + def _fill(self, name: str) -> Any: + if get_backend() != "cupy": + return None + return FILLS.get(self._path, {}).get(name) + def available(self, name: str) -> bool: """Return whether `name` exists in this namespace on the active backend. @@ -157,7 +173,7 @@ def available(self, name: str) -> bool: return False return True try: - return hasattr(self.resolve(), name) + return hasattr(self.resolve(), name) or self._fill(name) is not None except ImportError: return False @@ -172,6 +188,8 @@ def __dir__(self) -> list[str]: names = dir(self.resolve()) except ImportError: names = [] + if get_backend() == "cupy": + names = [*names, *FILLS.get(self._path, {})] return sorted(set(names) | set(children) | {"available", "resolve"}) diff --git a/src/cunumpy/_scipy_host.py b/src/cunumpy/_scipy_host.py new file mode 100644 index 0000000..0a62dff --- /dev/null +++ b/src/cunumpy/_scipy_host.py @@ -0,0 +1,370 @@ +"""SciPy solvers that ``cupyx.scipy`` lacks, run on the host with callbacks on the device.""" + +from __future__ import annotations + +import copy +import functools +import importlib +from collections.abc import Callable +from typing import Any + +from cunumpy import xp +from cunumpy._host import _to_device + + +def _scipy(name: str, function: str) -> Any: + try: + return importlib.import_module(f"scipy.{name}") + except ImportError as error: + raise ImportError( + f"cunumpy.{name}.{function} needs SciPy (pip install scipy)", + ) from error + + +def _host(value: Any) -> Any: + """Copy device arrays in `value` (or a tuple/list of them) to the host.""" + if isinstance(value, (tuple, list)): + return type(value)(_host(v) for v in value) + return xp.to_numpy(value) if xp.is_gpu(value) else value + + +def _device(value: Any) -> Any: + """Copy NumPy arrays in `value` (a tuple, list or dict of them) to the device.""" + if isinstance(value, dict): # OptimizeResult, infodict: a copy of the same type + out = copy.copy(value) + for key, item in value.items(): + out[key] = _device(item) + return out + if isinstance(value, (tuple, list)): + return type(value)(_device(v) for v in value) + return _to_device(value) + + +class _Bridge: + """Callbacks get arrays on the backend of `x0`; SciPy sees host arrays.""" + + def __init__(self, x0: Any) -> None: + self.device = xp.is_gpu(x0) + + def callback(self, fun: Any, n_arrays: int = 1) -> Any: + # the first `n_arrays` positional and all keyword arguments of `fun` + # arrive on the device + if not (self.device and callable(fun)): + return fun + + # wraps: SciPy reads the signature (a callback named `intermediate_result`) + @functools.wraps(fun) + def call(*args: Any, **kwargs: Any) -> Any: + args = (*_device(list(args[:n_arrays])), *args[n_arrays:]) + kwargs = {k: _device(v) for k, v in kwargs.items()} # intermediate_result + return _host(fun(*args, **kwargs)) + + return call + + def result(self, value: Any) -> Any: + return _device(value) if self.device else value + + def constraints(self, constraints: Any) -> Any: + if isinstance(constraints, (tuple, list)): + return type(constraints)(self.constraints(c) for c in constraints) + if isinstance(constraints, dict): # {"type": "ineq", "fun": ..., "jac": ...} + out = dict(constraints) + for key in ("fun", "jac"): + out[key] = self.callback(out.get(key)) + return {k: v for k, v in out.items() if v is not None} + out = copy.copy(constraints) # NonlinearConstraint, LinearConstraint + for name, n_arrays in (("fun", 1), ("jac", 1), ("hess", 2)): + if callable(getattr(out, name, None)): + setattr(out, name, self.callback(getattr(out, name), n_arrays)) + return out + + +def fsolve( + func: Callable[..., Any], + x0: Any, + args: tuple = (), + fprime: Callable[..., Any] | None = None, + **kwargs: Any, +) -> Any: + """Find a root of a system of equations with :func:`scipy.optimize.fsolve`. + + SciPy (MINPACK) runs on the host. With `x0` on the device, `func` and + `fprime` get ``x`` on the device, their results are copied to the host, + and the root comes back on the device: two transfers per evaluation, so + use it for setup, not in a time loop. + + Parameters + ---------- + func : callable + ``func(x, *args)``, the residual, an array of the shape of `x0`. + x0 : array + Initial guess; its backend is the backend of the callbacks and result. + args : tuple, optional + Extra arguments of `func` and `fprime`, passed unchanged. + fprime : callable, optional + ``fprime(x, *args)``, the Jacobian. + **kwargs + The other arguments of :func:`scipy.optimize.fsolve`. + + Returns + ------- + array or tuple + The root, on the backend of `x0`; with ``full_output=True`` SciPy's + tuple, with its arrays on that backend. + + See Also + -------- + root : The same with a choice of method. + newton : Many independent scalar equations at once, on the device. + + Examples + -------- + >>> xp.optimize.fsolve(lambda x: x**2 - xp.asarray([4.0, 9.0]), xp.ones(2)) + array([2., 3.]) + """ + optimize = _scipy("optimize", "fsolve") + bridge = _Bridge(x0) + out = optimize.fsolve( + bridge.callback(func), + _host(x0), + args=args, + fprime=bridge.callback(fprime), + **kwargs, + ) + return bridge.result(out) + + +def root( + fun: Callable[..., Any], + x0: Any, + args: tuple = (), + jac: Any = None, + callback: Callable[..., Any] | None = None, + **kwargs: Any, +) -> Any: + """Find a root of a system of equations with :func:`scipy.optimize.root`. + + SciPy runs on the host; with `x0` on the device, `fun`, `jac` and + `callback` get ``x`` on the device and the arrays of the result are + copied back to the device. See :func:`fsolve`. + + Parameters + ---------- + fun : callable + ``fun(x, *args)``, the residual (or ``(residual, jacobian)`` with + ``jac=True``). + x0 : array + Initial guess; its backend is the backend of the callbacks and result. + args : tuple, optional + Extra arguments of `fun` and `jac`, passed unchanged. + jac : callable or bool, optional + ``jac(x, *args)``, the Jacobian; ``True`` if `fun` returns it too. + callback : callable, optional + ``callback(x, f)``, called after each iteration (some methods only). + **kwargs + The other arguments of :func:`scipy.optimize.root` (``method``, + ``tol``, ``options``). + + Returns + ------- + OptimizeResult + SciPy's result; ``x``, ``fun`` and the other arrays on the backend of + `x0`. + + Examples + -------- + >>> res = xp.optimize.root(lambda x: x**3 - 8.0, xp.ones(1)) + >>> res.success, res.x + (True, array([2.])) + """ + optimize = _scipy("optimize", "root") + bridge = _Bridge(x0) + out = optimize.root( + bridge.callback(fun), + _host(x0), + args=args, + jac=bridge.callback(jac), + callback=bridge.callback(callback, 2), + **kwargs, + ) + return bridge.result(out) + + +def minimize( + fun: Callable[..., Any], + x0: Any, + args: tuple = (), + jac: Any = None, + hess: Any = None, + hessp: Callable[..., Any] | None = None, + bounds: Any = None, + constraints: Any = (), + callback: Callable[..., Any] | None = None, + **kwargs: Any, +) -> Any: + """Minimize a scalar function with :func:`scipy.optimize.minimize`. + + SciPy runs on the host; with `x0` on the device, every callback (`fun`, + `jac`, `hess`, `hessp`, `callback` and the functions of `constraints`) + gets its arrays on the device, and the arrays of the result are copied + back to the device. See :func:`fsolve`. + + Parameters + ---------- + fun : callable + ``fun(x, *args)``, a scalar (or ``(value, gradient)`` with + ``jac=True``). + x0 : array + Initial guess; its backend is the backend of the callbacks and result. + args : tuple, optional + Extra arguments of the callbacks, passed unchanged. + jac, hess : callable, str or bool, optional + Gradient and Hessian, as in SciPy (a callable, a finite-difference + scheme such as ``"2-point"``, or a ``HessianUpdateStrategy``). + hessp : callable, optional + ``hessp(x, p, *args)``, the Hessian times ``p``. + bounds : sequence or Bounds, optional + Bounds on ``x``, as in SciPy (host data). + constraints : dict, constraint object or sequence of them, optional + As in SciPy; their functions get device arrays like `fun`, their + bounds and matrices are host data. + callback : callable, optional + Called after each iteration, as in SciPy (``callback(intermediate_result)`` + or ``callback(xk)``). + **kwargs + The other arguments of :func:`scipy.optimize.minimize` (``method``, + ``tol``, ``options``). + + Returns + ------- + OptimizeResult + SciPy's result; ``x``, ``jac`` and the other arrays on the backend of + `x0`. + + Examples + -------- + >>> res = xp.optimize.minimize(lambda x: xp.sum((x - 1.0) ** 2), xp.zeros(2)) + >>> bool(xp.allclose(res.x, 1.0, atol=1e-6)) + True + """ + optimize = _scipy("optimize", "minimize") + bridge = _Bridge(x0) + if bridge.device: + constraints = bridge.constraints(constraints) + out = optimize.minimize( + bridge.callback(fun), + _host(x0), + args=args, + jac=bridge.callback(jac), + hess=bridge.callback(hess), + hessp=bridge.callback(hessp, 2), + bounds=bounds, + constraints=constraints, + callback=bridge.callback(callback, 2), # trust-constr: (xk, state) + **kwargs, + ) + return bridge.result(out) + + +def quad( + func: Callable[..., Any], + a: float, + b: float, + args: tuple = (), + **kwargs: Any, +) -> Any: + """Integrate a scalar function with :func:`scipy.integrate.quad`. + + SciPy (QUADPACK) runs on the host and calls `func` with a Python float. + `func` may compute on the device and return a 0-d device array: it is + copied to the host, one synchronization per evaluation. For many + integrals at once, evaluate the integrand on a fixed grid instead. + + Parameters + ---------- + func : callable + ``func(x, *args)``, a scalar. + a, b : float + Integration limits (``numpy.inf`` allowed). + args : tuple, optional + Extra arguments of `func`, passed unchanged. + **kwargs + The other arguments of :func:`scipy.integrate.quad`. + + Returns + ------- + tuple + SciPy's result, ``(value, abserr)`` plus the outputs that + ``full_output`` adds, on the host. + + Examples + -------- + >>> value, error = xp.integrate.quad(lambda x: x**2, 0.0, 3.0) + >>> round(value, 12) + 9.0 + """ + integrate = _scipy("integrate", "quad") + return integrate.quad( + lambda x, *a: _host(func(x, *a)), + _host(a), + _host(b), + args=args, + **kwargs, + ) + + +def odeint( + func: Callable[..., Any], + y0: Any, + t: Any, + args: tuple = (), + Dfun: Callable[..., Any] | None = None, + **kwargs: Any, +) -> Any: + """Integrate a system of ODEs with :func:`scipy.integrate.odeint`. + + SciPy (LSODA) runs on the host; with `y0` on the device, `func` and + `Dfun` get ``y`` on the device, and the solution comes back on the + device. See :func:`cunumpy.optimize.fsolve` for the cost. + + Parameters + ---------- + func : callable + ``func(y, t, *args)`` (``func(t, y, *args)`` with ``tfirst=True``), + the derivative of ``y``. + y0 : array + Initial condition; its backend is the backend of the callbacks and + result. + t : array + Times at which to return ``y``, starting with the time of `y0`; either + backend. + args : tuple, optional + Extra arguments of `func` and `Dfun`, passed unchanged. + Dfun : callable, optional + The Jacobian of `func`, with the same arguments. + **kwargs + The other arguments of :func:`scipy.integrate.odeint`. + + Returns + ------- + array or tuple + ``y`` at the times `t`, shape ``(len(t), len(y0))``, on the backend of + `y0`; with ``full_output=True`` also SciPy's info dict. + + Examples + -------- + >>> y = xp.integrate.odeint(lambda y, t: -y, xp.ones(1), xp.asarray([0.0, 1.0])) + >>> bool(xp.allclose(y[-1], np.exp(-1.0), rtol=1e-6)) + True + """ + integrate = _scipy("integrate", "odeint") + bridge = _Bridge(y0) + out = integrate.odeint( + bridge.callback(func, 2), + _host(y0), + _host(t), + args=args, + Dfun=bridge.callback(Dfun, 2), + **kwargs, + ) + return bridge.result(out) diff --git a/src/cunumpy/integrate.py b/src/cunumpy/integrate.py new file mode 100644 index 0000000..fd9b5a2 --- /dev/null +++ b/src/cunumpy/integrate.py @@ -0,0 +1,10 @@ +"""Integration on either backend. + +``cupyx.scipy`` has no ``integrate``. :func:`quad` and :func:`odeint` run +SciPy on the host, with callbacks that may compute on the device and results +on the backend of the input. See :doc:`/guides/solvers`. +""" + +from cunumpy._scipy_host import odeint, quad + +__all__ = ["odeint", "quad"] diff --git a/src/cunumpy/optimize.py b/src/cunumpy/optimize.py index 3014989..86b612d 100644 --- a/src/cunumpy/optimize.py +++ b/src/cunumpy/optimize.py @@ -1,14 +1,19 @@ """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:: +``optimize``. :func:`newton` runs on the device: many independent problems +solved at once, each step one array operation:: r = xp.optimize.newton(lambda r: psi(r) - levels, r0) +:func:`fsolve`, :func:`root` and :func:`minimize` run SciPy on the host, with +callbacks that get their arrays on the backend of ``x0`` (so code written with +``xp`` that uses device data works unchanged) and results on that backend. + See :doc:`/guides/solvers`. """ from cunumpy._optimize import NewtonResult, newton +from cunumpy._scipy_host import fsolve, minimize, root -__all__ = ["NewtonResult", "newton"] +__all__ = ["NewtonResult", "fsolve", "minimize", "newton", "root"] diff --git a/tests/unit/test_fake_cupyx.py b/tests/unit/test_fake_cupyx.py index 3fa26cf..f392ad2 100644 --- a/tests/unit/test_fake_cupyx.py +++ b/tests/unit/test_fake_cupyx.py @@ -18,6 +18,7 @@ def run(code): def test_same_subpackages_as_xp_scipy(): assert _fake_cupyx.SUBPACKAGES == SUBMODULES + assert tuple(_fake_cupyx.NAMES) == SUBMODULES def test_splines_fit_and_evaluate_on_the_device(): @@ -86,6 +87,61 @@ def test_interpolate_has_only_the_names_cupy_provides(): assert "ok" in run(code) +def test_every_subpackage_has_only_the_names_cupy_provides(): + code = """ +import cunumpy as xp +import cupyx.scipy.linalg +xp.set_backend("cupy") +missing = { + "special": ("jv", "erfi"), + "linalg": ("solve_circulant", "inv", "eigh"), + "sparse.linalg": ("inv",), +} +present = { + "special": ("yn", "erf", "erfinv", "ndtri"), + "linalg": ("circulant", "lu_factor", "lu_solve"), + "sparse": ("csr_matrix", "kron", "bmat", "identity"), + "sparse.linalg": ("splu", "spsolve", "cg"), + "signal": ("argrelextrema",), + "fft": ("rfft", "irfft"), +} +ns = lambda path: eval("xp.scipy." + path) +for path, names in missing.items(): + for name in names: + assert not hasattr(ns(path).resolve(), name), (path, name) +for path, names in present.items(): + for name in names: + assert ns(path).available(name), (path, name) +# cunumpy fills in solve_circulant, which raw cupyx lacks +assert not hasattr(cupyx.scipy.linalg, "solve_circulant") +assert xp.scipy.linalg.available("solve_circulant") +print("ok") +""" + assert "ok" in run(code) + + +def test_solve_circulant_stays_on_the_device(): + code = """ +import numpy as np +import scipy.linalg +import cunumpy as xp +xp.set_backend("cupy") +rng = np.random.default_rng(0) +c_h, b_h = rng.random(6) + 2.0, rng.random((6, 3)) +x = xp.scipy.linalg.solve_circulant(xp.asarray(c_h), xp.asarray(b_h)) +assert xp.is_gpu(x) +assert np.allclose(xp.to_numpy(x), scipy.linalg.solve_circulant(c_h, b_h)) +try: + xp.scipy.linalg.solve_circulant(xp.asarray([1.0, 1.0, 0.0, 0.0]), xp.ones(4)) +except np.linalg.LinAlgError: + pass +else: + raise AssertionError("singular matrix accepted") +print("ok") +""" + assert "ok" in run(code) + + def test_other_subpackages_and_pinned_memory(): code = """ import numpy as np diff --git a/tests/unit/test_linalg.py b/tests/unit/test_linalg.py new file mode 100644 index 0000000..e17df92 --- /dev/null +++ b/tests/unit/test_linalg.py @@ -0,0 +1,44 @@ +"""Tests for `solve_circulant`, the cunumpy fill-in for `cupyx.scipy.linalg`.""" + +import numpy as np +import pytest + +from cunumpy._linalg import solve_circulant + +scipy_linalg = pytest.importorskip("scipy.linalg") + +RNG = np.random.default_rng(0) + + +@pytest.mark.parametrize( + ("c", "b", "kwargs"), + [ + (RNG.random(5), RNG.random(5), {}), + (RNG.random(5), RNG.random((5, 3)), {}), + (RNG.random((2, 5)), RNG.random((5, 2)), {}), + (RNG.random(6) + 1j * RNG.random(6), RNG.random(6), {}), + ( + RNG.random((5, 3)), + RNG.random((3, 5)), + {"caxis": 0, "baxis": 1, "outaxis": -1}, + ), + (np.array([1.0, 1.0, 0.0, 0.0]), np.ones(4), {"singular": "lstsq"}), + (np.array([4, 1, 0, 1]), np.array([6, 6, 6, 6]), {}), + (np.float32([4, 1, 0, 1]), np.float32([1, 2, 3, 4]), {}), + ], +) +def test_matches_scipy(c, b, kwargs): + x = solve_circulant(c, b, **kwargs) + ref = scipy_linalg.solve_circulant(c, b, **kwargs) + assert x.dtype == ref.dtype + np.testing.assert_allclose(x, ref, rtol=1e-6, atol=1e-12) + + +def test_singular_raises(): + with pytest.raises(np.linalg.LinAlgError, match="near singular"): + solve_circulant(np.array([1.0, 1.0, 0.0, 0.0]), np.ones(4)) + + +def test_incompatible_shapes(): + with pytest.raises(ValueError, match="incompatible"): + solve_circulant(np.ones(4), np.ones(5)) diff --git a/tests/unit/test_scipy_backend.py b/tests/unit/test_scipy_backend.py index 01f9777..ed928b7 100644 --- a/tests/unit/test_scipy_backend.py +++ b/tests/unit/test_scipy_backend.py @@ -7,6 +7,7 @@ import pytest import cunumpy as xp +from cunumpy import _linalg from cunumpy import _scipy_backend as scipy_backend from cunumpy._scipy_backend import SUBMODULES, ScipyNamespace @@ -85,6 +86,21 @@ def test_cupy_backend_forwards_to_cupyx(fake_cupyx): xp.scipy.ndimage.resolve() +def test_cupy_backend_fills_in_missing_functions(fake_cupyx): + modules = {"cupyx.scipy.linalg": types.ModuleType("cupyx.scipy.linalg")} + fake_cupyx["cupyx.scipy"].linalg = modules["cupyx.scipy.linalg"] + sys.modules.update(modules) + try: + assert xp.scipy.linalg.solve_circulant is _linalg.solve_circulant + assert xp.scipy.linalg.available("solve_circulant") + assert "solve_circulant" in dir(xp.scipy.linalg) + # a cupyx version would win over the fill + modules["cupyx.scipy.linalg"].solve_circulant = "device solve_circulant" + assert xp.scipy.linalg.solve_circulant == "device solve_circulant" + finally: + del sys.modules["cupyx.scipy.linalg"] + + def test_backend_switch_takes_effect_immediately(fake_cupyx, monkeypatch): pytest.importorskip("scipy") import scipy.special diff --git a/tests/unit/test_scipy_host.py b/tests/unit/test_scipy_host.py new file mode 100644 index 0000000..5a20d3d --- /dev/null +++ b/tests/unit/test_scipy_host.py @@ -0,0 +1,136 @@ +"""SciPy solvers cupyx lacks: `xp.optimize.fsolve/root/minimize`, `xp.integrate`.""" + +from pathlib import Path + +import numpy as np +import pytest + +import cunumpy as xp +from cunumpy.kernel_testing import run_in_fake_cupy_subprocess + +optimize = pytest.importorskip("scipy.optimize") +integrate = pytest.importorskip("scipy.integrate") + +SRC = str(Path(__file__).resolve().parents[2] / "src") + + +def residual(x): + return np.array( + [x[0] + 0.5 * (x[0] - x[1]) ** 3 - 1.0, 0.5 * (x[1] - x[0]) ** 3 + x[1]] + ) + + +def test_exported(): + assert set(xp.optimize.__all__) >= {"fsolve", "root", "minimize", "newton"} + assert xp.integrate.__all__ == ["odeint", "quad"] + assert "integrate" in xp.__all__ + + +def test_host_arrays_give_scipy_results(): + x0 = np.zeros(2) + np.testing.assert_allclose( + xp.optimize.fsolve(residual, x0), optimize.fsolve(residual, x0) + ) + _, info, ier, _ = xp.optimize.fsolve(residual, x0, full_output=True) + assert ier == 1 and info["nfev"] > 0 + res = xp.optimize.root(residual, x0, method="lm") + np.testing.assert_allclose(res.x, optimize.root(residual, x0, method="lm").x) + rosen = optimize.rosen + res = xp.optimize.minimize( + rosen, np.zeros(3), jac=optimize.rosen_der, method="BFGS" + ) + np.testing.assert_allclose(res.x, 1.0, atol=1e-5) + value, _ = xp.integrate.quad(np.exp, 0.0, 1.0) + assert value == pytest.approx(np.e - 1.0) + t = np.linspace(0.0, 1.0, 5) + np.testing.assert_allclose( + xp.integrate.odeint(lambda y, t: -2.0 * y, np.ones(2), t), + integrate.odeint(lambda y, t: -2.0 * y, np.ones(2), t), + ) + + +def test_missing_scipy(monkeypatch): + import sys + + monkeypatch.setitem(sys.modules, "scipy.optimize", None) + with pytest.raises(ImportError, match=r"cunumpy.optimize.fsolve needs SciPy"): + xp.optimize.fsolve(residual, np.zeros(2)) + + +def test_callbacks_and_results_on_the_device(): + code = """ +import numpy as np +import scipy.optimize as so +import cunumpy as xp +xp.set_backend("cupy") + +target = xp.asarray([4.0, 9.0]) # device data captured by the callbacks + +def on_device(*arrays): + for a in arrays: + assert xp.is_gpu(a), type(a) + +def f(x): + on_device(x) + return x**2 - target + +def jac(x): + on_device(x) + return xp.diag(2 * x) + +x = xp.optimize.fsolve(f, xp.ones(2), fprime=jac) +assert xp.is_gpu(x) and np.allclose(xp.to_numpy(x), [2.0, 3.0]) +x, info, ier, msg = xp.optimize.fsolve(f, xp.ones(2), full_output=True) +assert xp.is_gpu(x) and xp.is_gpu(info["fvec"]) + +seen = [] +res = xp.optimize.root(f, xp.ones(2), jac=jac, method="hybr") +assert res.success and xp.is_gpu(res.x) and xp.is_gpu(res.fun) +res = xp.optimize.root(lambda x: (f(x), jac(x)), xp.ones(2), jac=True) +assert np.allclose(xp.to_numpy(res.x), [2.0, 3.0]) + +def fun(x): + on_device(x) + return xp.sum((x - target) ** 2) + +def grad(x): + on_device(x) + return 2 * (x - target) + +def hessp(x, p): + on_device(x, p) + return 2 * p + +def callback(intermediate_result): + on_device(intermediate_result.x) + seen.append(1) + +res = xp.optimize.minimize(fun, xp.zeros(2), jac=grad, hessp=hessp, method="trust-ncg", callback=callback) +assert seen and xp.is_gpu(res.x) and np.allclose(xp.to_numpy(res.x), [4.0, 9.0]) +cons = [ + {"type": "ineq", "fun": lambda x: (on_device(x), 5.0 - x[1])[1]}, + so.NonlinearConstraint(lambda x: (on_device(x), x[0])[1], -np.inf, 3.0), + so.LinearConstraint(np.asarray([[1.0, 1.0]]), -np.inf, 7.0), +] +bounds = so.Bounds(np.zeros(2), np.full(2, 10.0)) +res = xp.optimize.minimize(fun, xp.ones(2), method="SLSQP", constraints=cons, bounds=bounds) +assert res.success and np.allclose(xp.to_numpy(res.x), [2.0, 5.0], atol=1e-6) +res = xp.optimize.minimize(fun, xp.ones(2), bounds=[(0.0, 3.0), (0.0, 3.0)], method="L-BFGS-B") +assert np.allclose(xp.to_numpy(res.x), [3.0, 3.0]) + +rate = xp.asarray([1.0, 2.0]) +def rhs(y, t): + on_device(y) + return -rate * y +t = xp.linspace(0.0, 1.0, 5) +y = xp.integrate.odeint(rhs, xp.ones(2), t) +assert xp.is_gpu(y) and y.shape == (5, 2) +assert np.allclose(xp.to_numpy(y[-1]), np.exp([-1.0, -2.0]), rtol=1e-6) +y = xp.integrate.odeint(lambda t, y: rhs(y, t), xp.ones(2), t, tfirst=True) +assert np.allclose(xp.to_numpy(y[-1]), np.exp([-1.0, -2.0]), rtol=1e-6) + +value, error = xp.integrate.quad(lambda s: xp.sum(rate) * s, 0.0, 1.0) +assert abs(value - 1.5) < 1e-12 +print("ok") +""" + assert "ok" in run_in_fake_cupy_subprocess(code, env={"PYTHONPATH": SRC}).stdout