MHD equilibria on the CuPy backend

[max-models]

Stack: part 7 of 14, based on #713 (merge that first), next: #718. Full order: #709 → #708 → #711 → #712 → #705 → #713 → #714 → #718 → #720 → #721 → #724 → #725 → #723 → #722.

---

Solves the following issue(s):

Solves #696, part of #650 (see CUDA_STRATEGY.md). Every equilibrium in fields_background/equils.py can now be created and evaluated on the CuPy backend. The time loop does not change: it uses the projected equilibrium.

Core changes:

• What failed on CuPy. I checked this with cunumpy's fake CuPy in a clean serial subprocess. Every method models call was evaluated on a meshgrid and at markers.
  • Construction failed for:
    • AdhocTorus with q_kind 1 or 2: SciPy quad/UnivariateSpline got device arrays.
    • AdhocTorusQPsi: odeint/fsolve got device arrays.
    • EQDSKequilibrium: RectBivariateSpline got the device arrays from xp.linspace.
  • Evaluation failed for every field of GVECequilibrium and DESCequilibrium: gvec and DESC/JAX got device arrays.
  • Already working: the analytic ones (HomogenSlab, ShearedSlab, ShearFluid, ScrewPinch, AdhocTorus with q_kind=0, CircularTokamak, ConstantVelocity, HomogenSlabITG, CurrentSheet, the generic ones).
  • Also found: AdhocTorus.psi_r (q_kind=0) did out *= ... on a 0-d device array, which fails on CuPy when psi is called with arrays. It is now out = out * ....
• New helpers in fields_background/base.py:
  • setup_on_host: a decorator for __init__. It runs the setup on the NumPy backend (file reading, ODE solves, SciPy fits), so these equilibria hold only host data on either backend. Used on AdhocTorus, AdhocTorusQPsi and EQDSKequilibrium.
  • host_call(fun, *args, **kwargs): calls a host-only function, copying device arguments to the host and array results back once per call. Without device arguments it is a plain call, so the NumPy path does not change. Used for every SciPy spline evaluation: psi_r/p_r of AdhocTorus, psi_r of AdhocTorusQPsi, and psi/q_psi/g_psi/p_psi of EQDSK.
  • evaluate_on_host: the same as a method decorator. Used on the GVEC bv, jv, p0, n0 and the DESC bv, jv, p0, n0, gradB1.
• Cost of the host evaluation. Each call costs one host round trip. For example, gradB_xyz of EQDSK makes about 10 psi calls, so 10 round trips. This is fine because equilibria are only evaluated at setup: initial conditions, projections onto the FEEC spaces and mappings. A device B-spline evaluation of the SciPy knots and coefficients would remove these copies. It is listed as an open question in CUDA_STRATEGY.md.
• Tokamak: the workaround from PR 19 is removed. It no longer builds its default EQDSKequilibrium inside use_backend("numpy"). The field-line tracing (SciPy) still runs on the NumPy backend.
• Tests: new file fields_background/tests/test_equils_cupy.py.

  • What it checks: 19 cases that cover every equilibrium class of equils, including EQDSK, GVEC and DESC with the files shipped in the repo or with the packages. test_all_equilibria_have_cases checks that no class is missing. Each case is created on NumPy and on CuPy, and 30 logical-domain methods (absB0, p0, n0, b2, unit_b1, gradB1, curl_unit_b2, j2, …) are evaluated:

    • on a meshgrid and at markers,
    • plus psi/g_tor with all derivatives for the axisymmetric equilibria.

    Every array result must be a device array and match NumPy to 1e-12.

  • test_equil_fake_cupy runs without a GPU, in a subprocess on the fake CuPy. The fake CuPy cannot launch CUDA kernels, so host_geometry_kernels() runs the four geometry entry kernels with their pyccel version on the host buffers of the fake arrays. The CUDA geometry kernels have their own parity and emulation tests. The emulated_launches() of CUDA version of the linear_vlasov_ampere accumulation #705 would also work here, but it compiles every launch, and these cases make hundreds of domain evaluations.

  • test_equil_cupy runs the same check on a GPU.

Model-specific changes:

None.

Documentation changes:

CUDA_STRATEGY.md:

• New section "MHD equilibria on CuPy (MHD equilibria on the CuPy backend #696)".
• The open question "Polar splines and MHD equilibria" now covers polar splines only.
• New open question: equilibrium splines on the device.

Testing (macOS, no GPU; pyccel kernels compiled with GNU/Fortran; GPU tests not run):

	before (devel)	after
fields_background/tests geometry/tests	92 passed, 12 skipped	113 passed, 31 skipped
kinetic_background/tests post_processing/tests/test_orbits_tools.py feec/tests/test_field_init.py	50 passed	50 passed

The 21 new passes are the 19 fake-CuPy cases, test_all_equilibria_have_cases and test_host_call_numpy_is_plain_call. The 19 new skips are the GPU tests (test_equil_cupy). The DESC case takes about 35 s, because it loads W7-X from the desc package.

🤖 Generated with Claude Code