MHD equilibria on the CuPy backend

CUDA_STRATEGY.md

@@ -348,10 +348,10 @@ The three spaces (H1vec/Hcurl/Hdiv) differ only in the basis; port one, then the
 `push_v_sph_pressure`, `push_v_sph_pressure_ideal_gas`, `push_v_viscosity`, `div_u_weak_1form`.
 
 Infrastructure that gates the steps, independent of the kernels: mappings are no gate any more (every CUDA kernel
-accepts every mapping since PR 19), except for what is still host-only around them: polar splines in `Derham` and MHD
-equilibria such as `EQDSKequilibrium` cannot be created on CuPy (see [Open questions](#open-questions)); multi-rank marker sorting without host round trips. Array views with more than 4 dimensions
-(all matrix accumulations write 6D stencil matrix data) are in cunumpy since 0.6.1 (`Array6D`); `linear_vlasov_ampere` is
-the first kernel that uses them.
+accepts every mapping since PR 19), except for what is still host-only around them: polar splines in `Derham` cannot be
+created on CuPy (see [Open questions](#open-questions)). Multi-rank marker sorting stays on the device since #712. Array
+views with more than 4 dimensions (all matrix accumulations write 6D stencil matrix data) are in cunumpy since 0.6.1
+(`Array6D`); `linear_vlasov_ampere` is the first kernel that uses them.
 
 ## Testing
 
@@ -368,11 +368,39 @@ the first kernel that uses them.
 - **Marker layout.** The markers array is row-major (`n_markers × n_cols`). With one thread per marker, the memory accesses are strided. This is fine for now (each thread reads a few neighbouring columns), but a column-major or struct-of-arrays layout may be faster later. This would affect the CPU code too, so it is out of scope here. The array view represents strides explicitly on the CUDA side.
 - **MPI + GPUs.** One GPU per MPI rank (`xp.bind_local_device()` before `MPI_Init`, with feectools#86/#87), and GPU-aware MPI for the marker exchange, so markers do not go through the host. The marker exchange in `Particles.mpi_sort_markers` uses device buffers since #698 (see [Marker exchange implementation notes](#marker-exchange-implementation-notes-698)); the SPH ghost-box exchange (`_sendrecv_markers_boxes`) does not yet.
 - **Single-source alternatives.** Hand-written CUDA stays the default. Generating whole kernels from the Python source (`cupyx.jit`, numba-cuda, or a pyccel CUDA backend) is worth a look before the guiding-center kernels (the largest ones) are ported. Those tools take flat arguments, which `fields` also provides.
-- **Polar splines and MHD equilibria on the GPU** (left after PR 19): spline mappings run on the device, but
+- **Polar splines on the GPU** (left after PR 19): spline mappings run on the device, but
   `Derham` with `polar_splines=True` raises on CuPy (`PolarExtractionBlocksC1` builds SciPy sparse matrices, and the
-  polar extraction operators would apply them to device stencil data; needs `cupyx.scipy.sparse` or kernels), and
-  `EQDSKequilibrium` cannot be created on CuPy (its SciPy splines get device arrays). A `Tokamak` on CuPy builds its
-  default equilibrium on the host. Needed once a model with a polar domain or an EQDSK equilibrium runs on the GPU.
+  polar extraction operators would apply them to device stencil data; needs `cupyx.scipy.sparse` or kernels). Needed
+  once a model with a polar domain runs on the GPU. (MHD equilibria run on CuPy since #696, see
+  [MHD equilibria](#mhd-equilibria-on-cupy-696).)
+- **Equilibrium splines on the device.** The SciPy splines of `EQDSKequilibrium`, `AdhocTorus` (`q_kind` 1, 2) and
+  `AdhocTorusQPsi`, and GVEC/DESC, are evaluated on the host with one copy per call (#696). Fine while equilibria are
+  evaluated only at setup; a device B-spline evaluation of their knots and coefficients would remove the copies.
+
+## MHD equilibria on CuPy (#696)
+
+- **What failed.** Creating `AdhocTorus` (`q_kind` 1, 2: SciPy `quad`/`UnivariateSpline`), `AdhocTorusQPsi` (`odeint`,
+  `fsolve`) and `EQDSKequilibrium` (`RectBivariateSpline` on `xp.linspace`) on CuPy; evaluating `GVECequilibrium` and
+  `DESCequilibrium` (gvec/DESC got device arrays). The analytic equilibria already worked.
+- **Host setup.** `xp.setup_on_host` (cunumpy >= 0.6.2, which also provides `xp.host_call` and `xp.evaluate_on_host`) runs these `__init__`s on the NumPy backend, so the
+  equilibria hold only host data (floats, NumPy arrays, SciPy splines) on either backend. `Tokamak` no longer builds
+  its default `EQDSKequilibrium` on the NumPy backend itself; the field-line tracing still runs there.
+- **Evaluation follows the arguments.** SciPy spline evaluations go through `xp.host_call`, and the GVEC/DESC `bv`,
+  `jv`, `p0`, `n0`, `gradB1` through `@xp.evaluate_on_host`: device arguments are copied to the host, evaluated on the NumPy
+  backend, and the result is copied back, once per call (equilibria are evaluated at setup; the time loop uses the
+  projected equilibrium). NumPy arguments are evaluated as before.
+- **Tests.** `fields_background/tests/test_equils_cupy.py` creates every equilibrium of `equils` on CuPy, evaluates
+  the methods models call (meshgrid and markers, plus `psi`/`g_tor` with derivatives) and compares with NumPy: on a
+  GPU, and without one on the fake CuPy, where the four geometry kernels run their pyccel version on the fake arrays'
+  host buffers (`host_geometry_kernels`), since the fake CuPy cannot launch CUDA kernels.
+- **GVEC at markers (#715).** gvec evaluated the markers' `rho`, `theta`, `zeta` as a tensor grid, so `bv`/`jv` at N
+  markers returned N x N x N arrays and `absB0` failed (on NumPy too). The coordinates are now passed as
+  `xarray.DataArray`s with one shared dimension, which gvec evaluates point by point. Boozer coordinates
+  (`use_boozer=True`) are computed per flux surface, so marker evaluation raises there. Test: `test_gvec_equil.py`.
+- **Fake-CuPy child processes under MPI.** `run_fake_cupy_child` (`geometry/tests/test_domain.py`) starts the child
+  only on rank 0 (under `mpirun` every rank used to start the same child at once; on CI one of the two concurrent GVEC
+  children died with SIGILL and no output), with `faulthandler` and one OpenMP thread, and a failure reports the
+  signal and the end of stdout and stderr (gvec writes its Fortran messages to stdout).
 
 ## PR 10 implementation notes
 
@@ -617,7 +645,7 @@ Every mapping runs on the GPU: the spline mappings (`kind_map` 0–2) join the a
   NumPy backend, and the control points are copied to the active backend once. Found and tested without a GPU with
   cunumpy's fake CuPy (`CUNUMPY_FAKE_CUPY=1`), which rejects host/device mixing like CuPy.
 - **Still host-only.** `EQDSKequilibrium` cannot be created on CuPy (a `Tokamak` on CuPy takes an equilibrium created
-  on the NumPy backend, or builds its default one there). Polar splines: `PolarExtractionBlocksC1` builds SciPy sparse
+  on the NumPy backend, or builds its default one there; fixed in #696). Polar splines: `PolarExtractionBlocksC1` builds SciPy sparse
   matrices from the control points, and the polar extraction operators apply them to the stencil data, which lives on
   the device on CuPy; `Derham` raises `NotImplementedError` for `polar_splines=True` on CuPy (see
   [Open questions](#open-questions)).

pyproject.toml

@@ -24,7 +24,7 @@ classifiers = [
 ]
 dependencies = [
     "numpy<=2.5.0",
-    "cunumpy >= 0.6.1, <= 0.6.1",
+    "cunumpy >= 0.6.2, <= 0.6.2",
     "pyccel>=2.2.0, <=2.2.3",
     "maybempi>=0.1.2",
     "feectools>=0.5.0, <=0.5.0",

src/struphy/fields_background/equils.py

@@ -874,6 +874,7 @@ def doc_formula(cls):
         Units are those defned in the parameter file (through :class:`~struphy.io.options.BaseUnits`).
         """
 
+    @xp.setup_on_host
     def __init__(
         self,
         a: float = 1.0,
@@ -1010,7 +1011,8 @@ def psi_r(self, r, der=0):
 
             if der == 0:
                 out = -self.params["B0"] * self.params["a"] ** 2 / xp.sqrt(dq * q0 * eps**2 + dq**2)
-                out *= xp.arctanh(
+                # not in place: on CuPy, the prefactor is a 0-d device array
+                out = out * xp.arctanh(
                     xp.sqrt((dq - dq * (r / self.params["R0"]) ** 2) / (q0 * eps**2 + dq)),
                 )
             elif der == 1:
@@ -1020,7 +1022,7 @@ def psi_r(self, r, der=0):
 
         # alternative profile (interpolated)
         elif self.params["q_kind"] == 1 or self.params["q_kind"] == 2:
-            out = self._psi_i(r, nu=der)
+            out = xp.host_call(self._psi_i, r, nu=der)
 
             # remove all "dimensions" for point-wise evaluation
             if isinstance(r, (int, float)):
@@ -1149,7 +1151,7 @@ def p_r(self, r):
 
             # alternative profiles (interpolated)
             elif self.params["q_kind"] == 1 or self.params["q_kind"] == 2:
-                pout = self._p_i(r)
+                pout = xp.host_call(self._p_i, r)
 
                 # remove all "dimensions" for point-wise evaluation
                 if isinstance(r, (int, float)):
@@ -1365,6 +1367,7 @@ def doc_formula(cls):
         Units are those defned in the parameter file (through :class:`~struphy.io.options.BaseUnits`).
         """
 
+    @xp.setup_on_host
     def __init__(
         self,
         a: float = 0.361925,
@@ -1470,7 +1473,7 @@ def psi_r(self, r, der=0):
 
         assert der >= 0 and der <= 2, "Only first and second derivatives available!"
 
-        out = self._psi_i(r, nu=der)
+        out = xp.host_call(self._psi_i, r, nu=der)
 
         # remove all "dimensions" for point-wise evaluation
         if isinstance(r, (int, float)):
@@ -1665,6 +1668,7 @@ class EQDSKequilibrium(AxisymmMHDequilibrium):
         Struphy base units. If None, no rescaling of output is performed.
     """
 
+    @xp.setup_on_host
     def __init__(
         self,
         rel_path: bool = True,
@@ -1864,7 +1868,7 @@ def psi_axis_RZ(self):
 
     def q_psi(self, psi, der=0):
         """Safety factor q = q(psi)."""
-        out = self._q_i(psi, nu=der)
+        out = xp.host_call(self._q_i, psi, nu=der)
 
         # remove all "dimensions" for point-wise evaluation
         if isinstance(psi, (int, float)):
@@ -1875,7 +1879,7 @@ def q_psi(self, psi, der=0):
 
     def g_psi(self, psi, der=0):
         """Toroidal field function g = g(psi)."""
-        out = self._g_i(psi, nu=der)
+        out = xp.host_call(self._g_i, psi, nu=der)
 
         # remove all "dimensions" for point-wise evaluation
         if isinstance(psi, (int, float)):
@@ -1886,7 +1890,7 @@ def g_psi(self, psi, der=0):
 
     def p_psi(self, psi, der=0):
         """Pressure profile p = p(psi) in units Pa (as in the EQDSK file)."""
-        out = self._p_i(psi, nu=der)
+        out = xp.host_call(self._p_i, psi, nu=der)
 
         # remove all "dimensions" for point-wise evaluation
         if isinstance(psi, (int, float)):
@@ -1922,7 +1926,7 @@ def psi(self, R, Z, dR=0, dZ=0):
 
         is_float = all(isinstance(v, (int, float)) for v in [R, Z])
 
-        out = self._psi_i(R, Z, dx=dR, dy=dZ, grid=False)
+        out = xp.host_call(self._psi_i, R, Z, dx=dR, dy=dZ, grid=False)
 
         # remove all "dimensions" for point-wise evaluation
         if is_float:
@@ -2121,6 +2125,7 @@ def units(self) -> Units:
         """All Struphy units."""
         return self._units
 
+    @xp.evaluate_on_host
     @profile
     def bv(self, *etas, squeeze_out=False):
         """Contra-variant (vector field) magnetic field on logical cube [0, 1]^3 in Tesla / meter."""
@@ -2142,6 +2147,7 @@ def bv(self, *etas, squeeze_out=False):
 
         return out
 
+    @xp.evaluate_on_host
     @profile
     def jv(self, *etas, squeeze_out=False):
         """Contra-variant (vector field) current density (=curl B) on logical cube [0, 1]^3 in Ampere / meter^3."""
@@ -2171,6 +2177,7 @@ def jv(self, *etas, squeeze_out=False):
 
         return out
 
+    @xp.evaluate_on_host
     @profile
     def p0(self, *etas, squeeze_out=False):
         """0-form equilibrium pressure on logical cube [0, 1]^3."""
@@ -2190,6 +2197,7 @@ def p0(self, *etas, squeeze_out=False):
 
         return self.params["p0"] + tmp / self.units.p
 
+    @xp.evaluate_on_host
     @profile
     def n0(self, *etas, squeeze_out=False):
         """0-form equilibrium density on logical cube [0, 1]^3."""
@@ -2269,7 +2277,17 @@ def _gvec_evaluations(self, *etas):
 
         # evaluate
         if self.params["use_boozer"]:
+            if flat_eval:
+                # the Boozer transform is computed per flux surface, so gvec cannot evaluate it at scattered points
+                raise NotImplementedError("GVECequilibrium: marker evaluation is not available with use_boozer=True.")
             ev = gvec.EvaluationsBoozer(rho=rho, theta_B=theta, zeta_B=zeta, state=self.state)
+        elif flat_eval:
+            import xarray as xr
+
+            # coordinates sharing one dimension make gvec evaluate point by point instead of on the tensor grid
+            # rho x theta x zeta (one value per marker, not n_markers**3)
+            rho, theta, zeta = (xr.DataArray(c, dims="marker") for c in (rho, theta, zeta))
+            ev = gvec.Evaluations(rho=rho, theta=theta, zeta=zeta, state=self.state)
         else:
             ev = gvec.Evaluations(rho=rho, theta=theta, zeta=zeta, state=self.state)
 
@@ -2407,6 +2425,7 @@ def units(self) -> Units:
         """All Struphy units."""
         return self._units
 
+    @xp.evaluate_on_host
     @profile
     def bv(self, *etas, squeeze_out=False):
         """Contra-variant (vector field) magnetic field on logical cube [0, 1]^3 in Tesla / meter."""
@@ -2480,6 +2499,7 @@ def _eval_bv(self, *etas, squeeze_out=False):
 
         return out
 
+    @xp.evaluate_on_host
     @profile
     def jv(self, *etas, squeeze_out=False):
         """Contra-variant (vector field) current density (=curl B)
@@ -2555,6 +2575,7 @@ def _eval_jv(self, *etas, squeeze_out=False):
 
         return out
 
+    @xp.evaluate_on_host
     @profile
     def p0(self, *etas, squeeze_out=False):
         """0-form equilibrium pressure on logical cube [0, 1]^3 in Pascal."""
@@ -2586,6 +2607,7 @@ def p0(self, *etas, squeeze_out=False):
 
         return out
 
+    @xp.evaluate_on_host
     @profile
     def n0(self, *etas, squeeze_out=False):
         """0-form equilibrium density on logical cube [0, 1]^3."""
@@ -2610,6 +2632,7 @@ def n0(self, *etas, squeeze_out=False):
         # density in default units, n=1 --> 10^20 m^(-3)
         return p0_pascal / (self.params["T_kelvin"] * k_Boltzmann) / self.units.n
 
+    @xp.evaluate_on_host
     @profile
     def gradB1(self, *etas, squeeze_out=False):
         """1-form gradient of magnetic field strength on logical cube [0, 1]^3."""