From 14510d21e6fc44dd6acc4ea6593d6e4d775aef75 Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 7 Oct 2026 21:11:34 +0200 Subject: [PATCH 1/3] MHD equilibria on the CuPy backend Equilibria with host-only setup (AdhocTorus q_kind 1/2, AdhocTorusQPsi, EQDSKequilibrium) run their __init__ on the NumPy backend (setup_on_host) and hold only host data. SciPy spline evaluations go through host_call, GVEC/DESC evaluations through @evaluate_on_host: device arguments are copied to the host and the result back, once per call; NumPy arguments are evaluated as before. Tokamak no longer builds its default EQDSKequilibrium on the NumPy backend itself. New test_equils_cupy.py compares every equilibrium of equils on CuPy (GPU, or cunumpy's fake CuPy with host geometry kernels) with NumPy. Solves #696. Co-Authored-By: Claude Opus 5.5 --- CUDA_STRATEGY.md | 34 ++- src/struphy/fields_background/base.py | 80 +++++ src/struphy/fields_background/equils.py | 32 +- .../tests/test_equils_cupy.py | 276 ++++++++++++++++++ .../geometry/domains/tokamak/tokamak.py | 7 +- 5 files changed, 410 insertions(+), 19 deletions(-) create mode 100644 src/struphy/fields_background/tests/test_equils_cupy.py diff --git a/CUDA_STRATEGY.md b/CUDA_STRATEGY.md index daf0d25f0..68c596ea3 100644 --- a/CUDA_STRATEGY.md +++ b/CUDA_STRATEGY.md @@ -347,8 +347,8 @@ 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, +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 without host round trips, and array views with more than 4 dimensions in cunumpy (all matrix accumulations write 6D stencil matrix data). @@ -367,11 +367,31 @@ matrix data). - **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. `xp.mpi_is_cuda_aware()` detects it; the exchange in `Particles.mpi_sort_markers` has to be checked for host staging buffers. - **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.** `setup_on_host` (`fields_background/base.py`) 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 `host_call`, and the GVEC/DESC `bv`, `jv`, + `p0`, `n0`, `gradB1` through `@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. ## PR 10 implementation notes @@ -598,7 +618,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)). diff --git a/src/struphy/fields_background/base.py b/src/struphy/fields_background/base.py index 426fd8f28..7e894b366 100644 --- a/src/struphy/fields_background/base.py +++ b/src/struphy/fields_background/base.py @@ -1,9 +1,11 @@ "Base classes for MHD equilibria." +import functools import logging from abc import ABCMeta, abstractmethod import cunumpy as xp +import numpy as np from pyevtk.hl import gridToVTK from struphy.geometry.base import Domain @@ -18,6 +20,84 @@ logger = logging.getLogger("struphy") +def _on_device(arg) -> bool: + """Whether ``arg`` is (or contains, for tuples and lists) a CuPy array.""" + if isinstance(arg, (tuple, list)): + return any(_on_device(a) for a in arg) + return xp.is_gpu(arg) + + +def _to_host(arg): + """Copy CuPy arrays in ``arg`` (an array, or a tuple/list of them) to the host; leave everything else alone.""" + if isinstance(arg, (tuple, list)): + return type(arg)(_to_host(a) for a in arg) + return xp.to_numpy(arg) if xp.is_gpu(arg) else arg + + +def _to_device(arg): + """Copy NumPy arrays in ``arg`` (an array, or a tuple/list of them) to the device; leave scalars alone.""" + if isinstance(arg, (tuple, list)): + return type(arg)(_to_device(a) for a in arg) + return xp.to_cupy(arg) if isinstance(arg, np.ndarray) else arg + + +def host_call(fun, *args, **kwargs): + """Call a host-only function (SciPy spline, external equilibrium code) with arguments of any backend. + + Device (CuPy) arrays among the arguments are copied to the host, ``fun`` runs on the NumPy backend, and its array + results are copied back to the device, once per call. Without device arguments, ``fun`` is called directly on the + NumPy backend (on the NumPy backend, this is a plain call), so the result lives where the arguments live. + + The copies cost one host round trip per call; equilibria are evaluated at setup (initial conditions, projections + onto the FEEC spaces, mappings), not in the time loop, which uses the projected equilibrium. + + Parameters + ---------- + fun : callable + The host-only function. + + *args, **kwargs + Arguments of ``fun``; arrays (or tuples/lists of arrays) of either backend, or scalars. + + Returns + ------- + The result of ``fun`` (an array, a tuple/list of arrays or a scalar), with arrays on the device if any argument + was on the device. + """ + device = _on_device(args) or _on_device(list(kwargs.values())) + if xp.get_backend() == "numpy" and not device: + return fun(*args, **kwargs) + with xp.use_backend("numpy"): + out = fun(*_to_host(args), **{k: _to_host(v) for k, v in kwargs.items()}) + return _to_device(out) if device else out + + +def evaluate_on_host(method): + """Decorator for equilibrium methods that can only be evaluated on the host (see :func:`host_call`).""" + + @functools.wraps(method) + def wrapper(self, *args, **kwargs): + return host_call(method, self, *args, **kwargs) + + return wrapper + + +def setup_on_host(init): + """Decorator for ``__init__`` of equilibria whose setup is host-only (file reading, ODE solves, SciPy fits). + + ``__init__`` runs on the NumPy backend, so the equilibrium holds only host data (NumPy arrays, SciPy splines, + floats) on either backend. Its evaluation follows the backend of the arguments; host-only parts of it go through + :func:`host_call`. + """ + + @functools.wraps(init) + def wrapper(self, *args, **kwargs): + with xp.use_backend("numpy"): + init(self, *args, **kwargs) + + return wrapper + + class FluidEquilibrium(metaclass=ABCMeta): """ Abstract base class for callable fluid equilibria on arbitrary domains. diff --git a/src/struphy/fields_background/equils.py b/src/struphy/fields_background/equils.py index 34dac281f..576c918c0 100644 --- a/src/struphy/fields_background/equils.py +++ b/src/struphy/fields_background/equils.py @@ -27,6 +27,9 @@ NumericalFluidEquilibrium, NumericalFluidEquilibriumWithB, NumericalMHDequilibrium, + evaluate_on_host, + host_call, + setup_on_host, ) from struphy.fields_background.mhd_equil.eqdsk import readeqdsk from struphy.io.options import BaseUnits @@ -874,6 +877,7 @@ def doc_formula(cls): Units are those defned in the parameter file (through :class:`~struphy.io.options.BaseUnits`). """ + @setup_on_host def __init__( self, a: float = 1.0, @@ -1010,7 +1014,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 +1025,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 = host_call(self._psi_i, r, nu=der) # remove all "dimensions" for point-wise evaluation if isinstance(r, (int, float)): @@ -1149,7 +1154,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 = host_call(self._p_i, r) # remove all "dimensions" for point-wise evaluation if isinstance(r, (int, float)): @@ -1365,6 +1370,7 @@ def doc_formula(cls): Units are those defned in the parameter file (through :class:`~struphy.io.options.BaseUnits`). """ + @setup_on_host def __init__( self, a: float = 0.361925, @@ -1470,7 +1476,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 = host_call(self._psi_i, r, nu=der) # remove all "dimensions" for point-wise evaluation if isinstance(r, (int, float)): @@ -1665,6 +1671,7 @@ class EQDSKequilibrium(AxisymmMHDequilibrium): Struphy base units. If None, no rescaling of output is performed. """ + @setup_on_host def __init__( self, rel_path: bool = True, @@ -1864,7 +1871,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 = host_call(self._q_i, psi, nu=der) # remove all "dimensions" for point-wise evaluation if isinstance(psi, (int, float)): @@ -1875,7 +1882,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 = host_call(self._g_i, psi, nu=der) # remove all "dimensions" for point-wise evaluation if isinstance(psi, (int, float)): @@ -1886,7 +1893,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 = host_call(self._p_i, psi, nu=der) # remove all "dimensions" for point-wise evaluation if isinstance(psi, (int, float)): @@ -1922,7 +1929,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 = 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 +2128,7 @@ def units(self) -> Units: """All Struphy units.""" return self._units + @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 +2150,7 @@ def bv(self, *etas, squeeze_out=False): return out + @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 +2180,7 @@ def jv(self, *etas, squeeze_out=False): return out + @evaluate_on_host @profile def p0(self, *etas, squeeze_out=False): """0-form equilibrium pressure on logical cube [0, 1]^3.""" @@ -2190,6 +2200,7 @@ def p0(self, *etas, squeeze_out=False): return self.params["p0"] + tmp / self.units.p + @evaluate_on_host @profile def n0(self, *etas, squeeze_out=False): """0-form equilibrium density on logical cube [0, 1]^3.""" @@ -2407,6 +2418,7 @@ def units(self) -> Units: """All Struphy units.""" return self._units + @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 +2492,7 @@ def _eval_bv(self, *etas, squeeze_out=False): return out + @evaluate_on_host @profile def jv(self, *etas, squeeze_out=False): """Contra-variant (vector field) current density (=curl B) @@ -2555,6 +2568,7 @@ def _eval_jv(self, *etas, squeeze_out=False): return out + @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 +2600,7 @@ def p0(self, *etas, squeeze_out=False): return out + @evaluate_on_host @profile def n0(self, *etas, squeeze_out=False): """0-form equilibrium density on logical cube [0, 1]^3.""" @@ -2610,6 +2625,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 + @evaluate_on_host @profile def gradB1(self, *etas, squeeze_out=False): """1-form gradient of magnetic field strength on logical cube [0, 1]^3.""" diff --git a/src/struphy/fields_background/tests/test_equils_cupy.py b/src/struphy/fields_background/tests/test_equils_cupy.py new file mode 100644 index 000000000..f7ced7543 --- /dev/null +++ b/src/struphy/fields_background/tests/test_equils_cupy.py @@ -0,0 +1,276 @@ +"""Fluid and MHD equilibria on the CuPy backend, compared with the NumPy backend. + +The equilibria are created and evaluated on CuPy (with a GPU) or on cunumpy's fake CuPy (without one, in a +subprocess), which rejects host/device mixing like CuPy, and every evaluation is compared with the NumPy backend. +""" + +import contextlib +import importlib.util +import inspect +import subprocess +import sys + +import cunumpy +import numpy as np +import pytest +from cunumpy.kernel_testing import requires_cupy + +from struphy.geometry.tests.test_domain import _cupy_installed, serial_child_env + +# (equilibrium, its parameters, domain, its parameters); domain None for numerical equilibria (own domain). +# Every equilibrium of `equils` is here (`test_all_equilibria_have_cases`); GVEC and DESC need their packages. +EQUIL_CASES = { + "HomogenSlab": ("HomogenSlab", {}, "Cuboid", {}), + "ShearedSlab": ("ShearedSlab", {}, "Cuboid", {"r1": 1.0, "r2": 2 * np.pi, "r3": 2 * np.pi * 10.0}), + "ShearFluid": ("ShearFluid", {}, "Cuboid", {}), + "ScrewPinch": ("ScrewPinch", {}, "HollowCylinder", {"a1": 0.05, "a2": 1.0, "Lz": 2 * np.pi * 5.0}), + "ScrewPinch-q_inf": ( + "ScrewPinch", + {"q0": "inf", "q1": "inf"}, + "HollowCylinder", + {"a1": 0.05, "a2": 1.0, "Lz": 2 * np.pi * 5.0}, + ), + "AdhocTorus-q0p0": ("AdhocTorus", {"q_kind": 0, "p_kind": 0}, "HollowTorus", {"a1": 0.05, "a2": 1.0}), + "AdhocTorus-q1p0": ("AdhocTorus", {"q_kind": 1, "p_kind": 0}, "HollowTorus", {"a1": 0.05, "a2": 1.0}), + "AdhocTorus-q2p1": ("AdhocTorus", {"q_kind": 2, "p_kind": 1}, "HollowTorus", {"a1": 0.05, "a2": 1.0}), + "AdhocTorus-Tokamak": ("AdhocTorus", {"q_kind": 1, "p_kind": 0}, "Tokamak", {"num_elements": (6, 16)}), + "AdhocTorusQPsi": ("AdhocTorusQPsi", {}, "IGAPolarTorus", {"a": 0.361925, "R0": 1.0, "num_elements": (6, 16)}), + "CircularTokamak": ("CircularTokamak", {}, "HollowTorus", {"a1": 0.05, "a2": 1.0, "R0": 2.0}), + "EQDSKequilibrium": ("EQDSKequilibrium", {}, "Tokamak", {"num_elements": (6, 16)}), + "GVECequilibrium": ( + "GVECequilibrium", + { + "dat_file": "run_01/CIRCTOK_State_0000_00000000.dat", + "param_file": "run_01/parameter.ini", + "num_elements": (6, 6, 4), + }, + None, + {}, + ), + "DESCequilibrium": ("DESCequilibrium", {"num_elements": (6, 6, 6)}, None, {}), + "ConstantVelocity": ("ConstantVelocity", {}, "Cuboid", {}), + "HomogenSlabITG": ("HomogenSlabITG", {}, "Cuboid", {}), + "CurrentSheet": ("CurrentSheet", {}, "Cuboid", {}), + "GenericCartesianFluidEquilibrium": ("GenericCartesianFluidEquilibrium", {}, "Cuboid", {}), + "GenericCartesianFluidEquilibriumWithB": ("GenericCartesianFluidEquilibriumWithB", {}, "Cuboid", {}), +} + +# equilibria that need an optional package (skipped without it) +NEEDS = {"GVECequilibrium": "gvec", "DESCequilibrium": "desc"} + +# methods that models and projections call on the logical domain; those an equilibrium does not provide are skipped +METHODS = ( + "absB0", + "absB3", + "p0", + "p3", + "n0", + "n3", + "t0", + "vth0", + "b1", + "b2", + "bv", + "b_cart", + "unit_b1", + "unit_b2", + "unit_bv", + "gradB1", + "gradB2", + "gradBv", + "j1", + "j2", + "jv", + "absJ0", + "curl_unit_b1", + "curl_unit_b2", + "curl_unit_b_dot_b0", + "u1", + "u2", + "uv", + "a1", + "a2", +) + +# methods of axisymmetric equilibria in (R, Z) (flux and toroidal field function with derivatives) +AXISYMM_DERIVATIVES = ((0, 0), (1, 0), (0, 1), (2, 0), (0, 2), (1, 1)) + + +def _build(case): + """Create the equilibrium of ``case`` with its domain on the active backend.""" + from struphy import domains, equils + + name, params, dom_name, dom_params = EQUIL_CASES[case] + equil = getattr(equils, name)(**params) + if dom_name is None: + return equil + if dom_name == "Tokamak": + domain = domains.Tokamak(equilibrium=equil, **dom_params) + else: + domain = getattr(domains, dom_name)(**dom_params) + equil.domain = domain + return equil + + +def _flatten(out): + """Arrays in ``out`` (an array, or nested tuples/lists of arrays).""" + if isinstance(out, (tuple, list)): + return [a for o in out for a in _flatten(o)] + return [out] + + +def _evaluations(equil, args): + """Evaluate all methods of ``equil`` in ``METHODS`` (and psi, g_tor for axisymmetric ones) at ``args``.""" + from struphy.fields_background.base import AxisymmMHDequilibrium + + out = {} + for name in METHODS: + method = getattr(equil, name, None) + if method is None: + continue + try: + out[name] = method(*args) + except (AssertionError, NotImplementedError) as error: + # not available for this equilibrium (vector potential, GVEC gradB1, ...) + out[name] = type(error) + if isinstance(equil, AxisymmMHDequilibrium): + R = 1 * equil.psi_axis_RZ[0] + 0.3 * args[0] + Z = 1 * equil.psi_axis_RZ[1] + 0.2 * args[0] - 0.1 + for dR, dZ in AXISYMM_DERIVATIVES: + out[f"psi_{dR}{dZ}"] = equil.psi(R, Z, dR=dR, dZ=dZ) + for dR, dZ in AXISYMM_DERIVATIVES[:3]: + out[f"g_tor_{dR}{dZ}"] = equil.g_tor(R, Z, dR=dR, dZ=dZ) + return out + + +@contextlib.contextmanager +def host_geometry_kernels(): + """Fake CuPy only: run the geometry kernels of :class:`~struphy.geometry.base.Domain` with their pyccel versions. + + The fake CuPy cannot launch CUDA kernels. Here the four geometry entry kernels run their pyccel version on the + host buffers of the fake device arrays (shared, so the outputs are written in place), with the pyccel + ``DomainArguments`` built from the same buffers. The equilibria themselves run unchanged on the fake CuPy. The + CUDA geometry kernels are checked against pyccel in ``pic/tests/test_cuda_parity.py`` (GPU) and + ``test_cuda_emulation.py`` (CPU emulation). + """ + import struphy.geometry.base as geometry_base + from struphy.kernel_arguments.pusher_args_cuda import CudaDomainArguments + from struphy.kernel_arguments.pusher_args_kernels import DomainArguments + + def host(arg): + if isinstance(arg, CudaDomainArguments): + return DomainArguments(arg.kind_map, *(host(getattr(arg, name)) for name, _ in arg.fields[1:])) + return arg._a if cunumpy.is_gpu(arg) else arg # the host buffer of a fake CuPy array + + def on_host(kernel): + def launch(*args, **launch_options): + return kernel.host_kernel.kernel(*(host(a) for a in args)) + + return launch + + names = ("kernel_evaluate", "kernel_evaluate_pic", "kernel_pullpush", "kernel_pullpush_pic") + kernels = {name: getattr(geometry_base, name) for name in names} + try: + for name, kernel in kernels.items(): + setattr(geometry_base, name, on_host(kernel)) + yield + finally: + for name, kernel in kernels.items(): + setattr(geometry_base, name, kernel) + + +def check_equil_on_cupy(case): + """Create ``case`` on the CuPy backend, evaluate it there and compare with the NumPy backend. + + Every result is a device array (or a scalar for constant parts like ``psi`` derivatives of analytic profiles) and + agrees with the NumPy result. Evaluated on a meshgrid (three 1d arrays) and at markers (one 2d array). + """ + rng = np.random.default_rng(1234) + e1 = np.sort(rng.uniform(0.1, 0.9, 4)) + e2 = np.sort(rng.uniform(0.0, 1.0, 5)) + e3 = np.sort(rng.uniform(0.0, 1.0, 3)) + markers = np.column_stack([rng.uniform(0.1, 0.9, 7), rng.uniform(0.0, 1.0, 7), rng.uniform(0.0, 1.0, 7)]) + + with cunumpy.use_backend("numpy"): + equil = _build(case) + ref_grid = _evaluations(equil, (e1, e2, e3)) + ref_markers = _evaluations(equil, (markers,)) + + with cunumpy.use_backend("cupy"): + equil = _build(case) + out_grid = _evaluations(equil, tuple(cunumpy.asarray(e) for e in (e1, e2, e3))) + out_markers = _evaluations(equil, (cunumpy.asarray(markers),)) + + for ref, out, where in ((ref_grid, out_grid, "meshgrid"), (ref_markers, out_markers, "markers")): + assert ref.keys() == out.keys() + for name in ref: + if isinstance(ref[name], type): # not available on NumPy, not on CuPy either + assert out[name] is ref[name], (case, where, name) + continue + refs, outs = _flatten(ref[name]), _flatten(out[name]) + assert len(refs) == len(outs), (case, where, name) + for r, o in zip(refs, outs): + if isinstance(r, np.ndarray) and r.ndim > 0: + assert cunumpy.is_gpu(o), (case, where, name, type(o)) + o = cunumpy.to_numpy(o) if cunumpy.is_gpu(o) else np.asarray(o) + assert np.allclose(o, r, rtol=1e-12, atol=1e-12, equal_nan=True), (case, where, name) + + +def _check_needs(case): + package = NEEDS.get(EQUIL_CASES[case][0]) + if package is not None and importlib.util.find_spec(package) is None: + pytest.skip(f"{package} is not installed") + + +def test_all_equilibria_have_cases(): + """Every equilibrium class of ``equils`` is checked here.""" + from struphy import equils + from struphy.fields_background.base import FluidEquilibrium + + classes = { + name + for name, cls in vars(equils).items() + if inspect.isclass(cls) + and issubclass(cls, FluidEquilibrium) + and not inspect.isabstract(cls) + and cls.__module__ == equils.__name__ + } + covered = {case[0] for case in EQUIL_CASES.values()} + assert classes == covered + + +def test_host_call_numpy_is_plain_call(): + """On the NumPy backend, :func:`host_call` calls the function directly and returns its result unchanged.""" + from struphy.fields_background.base import host_call + + a = np.linspace(0.0, 1.0, 5) + with cunumpy.use_backend("numpy"): + out = host_call(lambda x, y=1.0: x * y, a, y=2.0) + assert type(out) is np.ndarray + assert np.array_equal(out, 2 * a) + + +@pytest.mark.skipif(_cupy_installed(), reason="the fake CuPy cannot replace an installed CuPy") +@pytest.mark.parametrize("case", list(EQUIL_CASES)) +def test_equil_fake_cupy(case): + """Without a GPU: :func:`check_equil_on_cupy` with cunumpy's fake CuPy, which rejects host/device mixing. + + Runs in a subprocess because the fake CuPy must be installed before cunumpy is imported. + """ + _check_needs(case) + code = ( + "from struphy.fields_background.tests.test_equils_cupy import check_equil_on_cupy, host_geometry_kernels\n" + f"with host_geometry_kernels(): check_equil_on_cupy({case!r})" + ) + result = subprocess.run( + [sys.executable, "-c", code], env=serial_child_env(CUNUMPY_FAKE_CUPY="1"), capture_output=True, text=True + ) + assert result.returncode == 0, result.stderr[-4000:] + + +@requires_cupy +@pytest.mark.parametrize("case", list(EQUIL_CASES)) +def test_equil_cupy(case): + """On a GPU: the equilibria are created and evaluated on CuPy and agree with NumPy.""" + _check_needs(case) + check_equil_on_cupy(case) diff --git a/src/struphy/geometry/domains/tokamak/tokamak.py b/src/struphy/geometry/domains/tokamak/tokamak.py index c53bb7882..1fa3d7fcc 100644 --- a/src/struphy/geometry/domains/tokamak/tokamak.py +++ b/src/struphy/geometry/domains/tokamak/tokamak.py @@ -67,11 +67,10 @@ def __init__( ): if r_min != 0.0: r0 = r_min - # The equilibrium and the field-line tracing (SciPy) are host-only setup: they run on the NumPy - # backend, and the control points are copied to the active backend once. + # The field-line tracing (SciPy) is host-only setup: it runs on the NumPy backend, and the control points + # are copied to the active backend once. EQDSKequilibrium does its own setup (file, SciPy splines) on the host. if equilibrium is None: - with xp.use_backend("numpy"): - equilibrium = EQDSKequilibrium() + equilibrium = EQDSKequilibrium() else: assert isinstance(equilibrium, AxisymmMHDequilibrium) From b11ee1869896f114955733bed3f583012e07ca74 Mon Sep 17 00:00:00 2001 From: Max Date: Thu, 8 Oct 2026 13:16:00 +0200 Subject: [PATCH 2/3] GVEC: marker evaluation, cunumpy 0.6.2 host helpers, robust fake-CuPy child tests - GVECequilibrium evaluates markers point by point: the coordinates go to gvec as xarray DataArrays with a shared dimension, so bv/jv at N markers no longer return N x N x N arrays and absB0 at markers works (also on NumPy). With use_boozer=True, marker evaluation raises, since gvec computes the Boozer transform per flux surface. New test_gvec_equil.py. - Use xp.host_call, xp.evaluate_on_host and xp.setup_on_host from cunumpy 0.6.2 instead of the copies in fields_background/base.py; pin cunumpy 0.6.2. - run_fake_cupy_child: under mpirun only rank 0 starts the fake-CuPy child (every rank started the same child at once, and on CI one of two concurrent GVEC children died with SIGILL and no output); the child runs with faulthandler and one OpenMP thread, and failures report the signal and the end of stdout/stderr. Used by the four fake-CuPy subprocess tests. Co-Authored-By: Claude Opus 5.5 --- CUDA_STRATEGY.md | 14 +++- pyproject.toml | 2 +- src/struphy/fields_background/base.py | 80 ------------------- src/struphy/fields_background/equils.py | 51 +++++++----- .../tests/test_equils_cupy.py | 15 +--- .../tests/test_gvec_equil.py | 47 +++++++++++ src/struphy/geometry/tests/test_domain.py | 34 +++++++- .../pic/tests/test_accum_matrix_cupy.py | 9 +-- .../tests/test_compile_cuda_kernels.py | 13 +-- 9 files changed, 126 insertions(+), 139 deletions(-) create mode 100644 src/struphy/fields_background/tests/test_gvec_equil.py diff --git a/CUDA_STRATEGY.md b/CUDA_STRATEGY.md index 8fcc4368b..462d18830 100644 --- a/CUDA_STRATEGY.md +++ b/CUDA_STRATEGY.md @@ -382,17 +382,25 @@ views with more than 4 dimensions (all matrix accumulations write 6D stencil mat - **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.** `setup_on_host` (`fields_background/base.py`) runs these `__init__`s on the NumPy backend, so the +- **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 `host_call`, and the GVEC/DESC `bv`, `jv`, - `p0`, `n0`, `gradB1` through `@evaluate_on_host`: device arguments are copied to the host, evaluated on the NumPy +- **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 diff --git a/pyproject.toml b/pyproject.toml index 88547f445..87f42741f 100644 --- a/pyproject.toml +++ b/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", diff --git a/src/struphy/fields_background/base.py b/src/struphy/fields_background/base.py index 7e894b366..426fd8f28 100644 --- a/src/struphy/fields_background/base.py +++ b/src/struphy/fields_background/base.py @@ -1,11 +1,9 @@ "Base classes for MHD equilibria." -import functools import logging from abc import ABCMeta, abstractmethod import cunumpy as xp -import numpy as np from pyevtk.hl import gridToVTK from struphy.geometry.base import Domain @@ -20,84 +18,6 @@ logger = logging.getLogger("struphy") -def _on_device(arg) -> bool: - """Whether ``arg`` is (or contains, for tuples and lists) a CuPy array.""" - if isinstance(arg, (tuple, list)): - return any(_on_device(a) for a in arg) - return xp.is_gpu(arg) - - -def _to_host(arg): - """Copy CuPy arrays in ``arg`` (an array, or a tuple/list of them) to the host; leave everything else alone.""" - if isinstance(arg, (tuple, list)): - return type(arg)(_to_host(a) for a in arg) - return xp.to_numpy(arg) if xp.is_gpu(arg) else arg - - -def _to_device(arg): - """Copy NumPy arrays in ``arg`` (an array, or a tuple/list of them) to the device; leave scalars alone.""" - if isinstance(arg, (tuple, list)): - return type(arg)(_to_device(a) for a in arg) - return xp.to_cupy(arg) if isinstance(arg, np.ndarray) else arg - - -def host_call(fun, *args, **kwargs): - """Call a host-only function (SciPy spline, external equilibrium code) with arguments of any backend. - - Device (CuPy) arrays among the arguments are copied to the host, ``fun`` runs on the NumPy backend, and its array - results are copied back to the device, once per call. Without device arguments, ``fun`` is called directly on the - NumPy backend (on the NumPy backend, this is a plain call), so the result lives where the arguments live. - - The copies cost one host round trip per call; equilibria are evaluated at setup (initial conditions, projections - onto the FEEC spaces, mappings), not in the time loop, which uses the projected equilibrium. - - Parameters - ---------- - fun : callable - The host-only function. - - *args, **kwargs - Arguments of ``fun``; arrays (or tuples/lists of arrays) of either backend, or scalars. - - Returns - ------- - The result of ``fun`` (an array, a tuple/list of arrays or a scalar), with arrays on the device if any argument - was on the device. - """ - device = _on_device(args) or _on_device(list(kwargs.values())) - if xp.get_backend() == "numpy" and not device: - return fun(*args, **kwargs) - with xp.use_backend("numpy"): - out = fun(*_to_host(args), **{k: _to_host(v) for k, v in kwargs.items()}) - return _to_device(out) if device else out - - -def evaluate_on_host(method): - """Decorator for equilibrium methods that can only be evaluated on the host (see :func:`host_call`).""" - - @functools.wraps(method) - def wrapper(self, *args, **kwargs): - return host_call(method, self, *args, **kwargs) - - return wrapper - - -def setup_on_host(init): - """Decorator for ``__init__`` of equilibria whose setup is host-only (file reading, ODE solves, SciPy fits). - - ``__init__`` runs on the NumPy backend, so the equilibrium holds only host data (NumPy arrays, SciPy splines, - floats) on either backend. Its evaluation follows the backend of the arguments; host-only parts of it go through - :func:`host_call`. - """ - - @functools.wraps(init) - def wrapper(self, *args, **kwargs): - with xp.use_backend("numpy"): - init(self, *args, **kwargs) - - return wrapper - - class FluidEquilibrium(metaclass=ABCMeta): """ Abstract base class for callable fluid equilibria on arbitrary domains. diff --git a/src/struphy/fields_background/equils.py b/src/struphy/fields_background/equils.py index 576c918c0..2377cebdc 100644 --- a/src/struphy/fields_background/equils.py +++ b/src/struphy/fields_background/equils.py @@ -27,9 +27,6 @@ NumericalFluidEquilibrium, NumericalFluidEquilibriumWithB, NumericalMHDequilibrium, - evaluate_on_host, - host_call, - setup_on_host, ) from struphy.fields_background.mhd_equil.eqdsk import readeqdsk from struphy.io.options import BaseUnits @@ -877,7 +874,7 @@ def doc_formula(cls): Units are those defned in the parameter file (through :class:`~struphy.io.options.BaseUnits`). """ - @setup_on_host + @xp.setup_on_host def __init__( self, a: float = 1.0, @@ -1025,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 = host_call(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)): @@ -1154,7 +1151,7 @@ def p_r(self, r): # alternative profiles (interpolated) elif self.params["q_kind"] == 1 or self.params["q_kind"] == 2: - pout = host_call(self._p_i, r) + pout = xp.host_call(self._p_i, r) # remove all "dimensions" for point-wise evaluation if isinstance(r, (int, float)): @@ -1370,7 +1367,7 @@ def doc_formula(cls): Units are those defned in the parameter file (through :class:`~struphy.io.options.BaseUnits`). """ - @setup_on_host + @xp.setup_on_host def __init__( self, a: float = 0.361925, @@ -1476,7 +1473,7 @@ def psi_r(self, r, der=0): assert der >= 0 and der <= 2, "Only first and second derivatives available!" - out = host_call(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)): @@ -1671,7 +1668,7 @@ class EQDSKequilibrium(AxisymmMHDequilibrium): Struphy base units. If None, no rescaling of output is performed. """ - @setup_on_host + @xp.setup_on_host def __init__( self, rel_path: bool = True, @@ -1871,7 +1868,7 @@ def psi_axis_RZ(self): def q_psi(self, psi, der=0): """Safety factor q = q(psi).""" - out = host_call(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)): @@ -1882,7 +1879,7 @@ def q_psi(self, psi, der=0): def g_psi(self, psi, der=0): """Toroidal field function g = g(psi).""" - out = host_call(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)): @@ -1893,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 = host_call(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)): @@ -1929,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 = host_call(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: @@ -2128,7 +2125,7 @@ def units(self) -> Units: """All Struphy units.""" return self._units - @evaluate_on_host + @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.""" @@ -2150,7 +2147,7 @@ def bv(self, *etas, squeeze_out=False): return out - @evaluate_on_host + @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.""" @@ -2180,7 +2177,7 @@ def jv(self, *etas, squeeze_out=False): return out - @evaluate_on_host + @xp.evaluate_on_host @profile def p0(self, *etas, squeeze_out=False): """0-form equilibrium pressure on logical cube [0, 1]^3.""" @@ -2200,7 +2197,7 @@ def p0(self, *etas, squeeze_out=False): return self.params["p0"] + tmp / self.units.p - @evaluate_on_host + @xp.evaluate_on_host @profile def n0(self, *etas, squeeze_out=False): """0-form equilibrium density on logical cube [0, 1]^3.""" @@ -2280,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) @@ -2418,7 +2425,7 @@ def units(self) -> Units: """All Struphy units.""" return self._units - @evaluate_on_host + @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.""" @@ -2492,7 +2499,7 @@ def _eval_bv(self, *etas, squeeze_out=False): return out - @evaluate_on_host + @xp.evaluate_on_host @profile def jv(self, *etas, squeeze_out=False): """Contra-variant (vector field) current density (=curl B) @@ -2568,7 +2575,7 @@ def _eval_jv(self, *etas, squeeze_out=False): return out - @evaluate_on_host + @xp.evaluate_on_host @profile def p0(self, *etas, squeeze_out=False): """0-form equilibrium pressure on logical cube [0, 1]^3 in Pascal.""" @@ -2600,7 +2607,7 @@ def p0(self, *etas, squeeze_out=False): return out - @evaluate_on_host + @xp.evaluate_on_host @profile def n0(self, *etas, squeeze_out=False): """0-form equilibrium density on logical cube [0, 1]^3.""" @@ -2625,7 +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 - @evaluate_on_host + @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.""" diff --git a/src/struphy/fields_background/tests/test_equils_cupy.py b/src/struphy/fields_background/tests/test_equils_cupy.py index f7ced7543..e7710570d 100644 --- a/src/struphy/fields_background/tests/test_equils_cupy.py +++ b/src/struphy/fields_background/tests/test_equils_cupy.py @@ -7,15 +7,13 @@ import contextlib import importlib.util import inspect -import subprocess -import sys import cunumpy import numpy as np import pytest from cunumpy.kernel_testing import requires_cupy -from struphy.geometry.tests.test_domain import _cupy_installed, serial_child_env +from struphy.geometry.tests.test_domain import _cupy_installed, run_fake_cupy_child # (equilibrium, its parameters, domain, its parameters); domain None for numerical equilibria (own domain). # Every equilibrium of `equils` is here (`test_all_equilibria_have_cases`); GVEC and DESC need their packages. @@ -240,12 +238,10 @@ def test_all_equilibria_have_cases(): def test_host_call_numpy_is_plain_call(): - """On the NumPy backend, :func:`host_call` calls the function directly and returns its result unchanged.""" - from struphy.fields_background.base import host_call - + """On the NumPy backend, :func:`cunumpy.host_call` calls the function directly and returns its result unchanged.""" a = np.linspace(0.0, 1.0, 5) with cunumpy.use_backend("numpy"): - out = host_call(lambda x, y=1.0: x * y, a, y=2.0) + out = cunumpy.host_call(lambda x, y=1.0: x * y, a, y=2.0) assert type(out) is np.ndarray assert np.array_equal(out, 2 * a) @@ -262,10 +258,7 @@ def test_equil_fake_cupy(case): "from struphy.fields_background.tests.test_equils_cupy import check_equil_on_cupy, host_geometry_kernels\n" f"with host_geometry_kernels(): check_equil_on_cupy({case!r})" ) - result = subprocess.run( - [sys.executable, "-c", code], env=serial_child_env(CUNUMPY_FAKE_CUPY="1"), capture_output=True, text=True - ) - assert result.returncode == 0, result.stderr[-4000:] + run_fake_cupy_child(code) @requires_cupy diff --git a/src/struphy/fields_background/tests/test_gvec_equil.py b/src/struphy/fields_background/tests/test_gvec_equil.py new file mode 100644 index 000000000..40b84e065 --- /dev/null +++ b/src/struphy/fields_background/tests/test_gvec_equil.py @@ -0,0 +1,47 @@ +import importlib.util + +import numpy as np +import pytest + +pytestmark = pytest.mark.skipif(importlib.util.find_spec("gvec") is None, reason="gvec is not installed") + +GVEC_PARAMS = {"dat_file": "run_01/CIRCTOK_State_0000_00000000.dat", "param_file": "run_01/parameter.ini"} + + +def _markers(n=7, seed=3): + """``n`` markers at random logical positions, away from the hole at the axis.""" + rng = np.random.default_rng(seed) + markers = rng.uniform(0.0, 1.0, (n, 3)) + markers[:, 0] = rng.uniform(0.1, 0.9, n) + return markers + + +@pytest.mark.parametrize("name", ["bv", "jv", "p0", "n0", "absB0"]) +def test_gvec_marker_evaluation_matches_grid(name): + """Marker (flat) evaluation gives one value per marker, equal to the grid evaluation at each marker. + + Before, gvec evaluated the markers' rho, theta and zeta as a tensor grid, so ``bv`` and ``jv`` at N markers + returned N x N x N arrays, and ``absB0`` failed. + """ + from struphy.fields_background.equils import GVECequilibrium + + equil = GVECequilibrium(**GVEC_PARAMS) + markers = _markers() + flat = getattr(equil, name)(markers) + flat = flat if isinstance(flat, tuple) else (flat,) + + for ip, (e1, e2, e3) in enumerate(markers): + grid = getattr(equil, name)(np.array([e1]), np.array([e2]), np.array([e3])) + grid = grid if isinstance(grid, tuple) else (grid,) + for f, g in zip(flat, grid): + assert f.shape == (len(markers),) + assert np.isclose(f[ip], g.ravel()[0], rtol=1e-12, atol=1e-14) + + +def test_gvec_marker_evaluation_boozer_raises(): + """gvec computes the Boozer transform per flux surface, so marker evaluation with ``use_boozer=True`` raises.""" + from struphy.fields_background.equils import GVECequilibrium + + equil = GVECequilibrium(**GVEC_PARAMS, use_boozer=True) + with pytest.raises(NotImplementedError, match="use_boozer"): + equil.bv(_markers()) diff --git a/src/struphy/geometry/tests/test_domain.py b/src/struphy/geometry/tests/test_domain.py index 495f9546d..466d138d9 100644 --- a/src/struphy/geometry/tests/test_domain.py +++ b/src/struphy/geometry/tests/test_domain.py @@ -3,6 +3,7 @@ import logging import os import pickle +import signal import subprocess import sys @@ -1199,6 +1200,34 @@ def serial_child_env(**extra): return env +def run_fake_cupy_child(code: str): + """Run ``code`` in a serial child Python process on cunumpy's fake CuPy and fail the test if the child fails. + + The check is serial, so under ``mpirun`` only rank 0 starts the child; the other ranks skip. Before, every rank + started the same child at the same time, and on CI the GVEC case crashed (SIGILL) in one of the two concurrent + children with no output. The child runs with ``faulthandler`` (a crash prints the Python traceback) and one OpenMP + thread, and a failure reports the signal or exit code and the end of both output streams (GVEC writes its Fortran + messages to stdout). + """ + from maybempi import MPI + + if MPI.COMM_WORLD.Get_rank() != 0: + pytest.skip("serial check, runs on MPI rank 0") + result = subprocess.run( + [sys.executable, "-X", "faulthandler", "-c", code], + env=serial_child_env(CUNUMPY_FAKE_CUPY="1", OMP_NUM_THREADS="1"), + capture_output=True, + text=True, + ) + if result.returncode != 0: + rc = result.returncode + how = f"signal {signal.Signals(-rc).name}" if rc < 0 else f"exit code {rc}" + pytest.fail( + f"child process failed with {how}\n" + f"--- stdout (end) ---\n{result.stdout[-2000:]}\n--- stderr (end) ---\n{result.stderr[-4000:]}", + ) + + @pytest.mark.skipif(_cupy_installed(), reason="the fake CuPy cannot replace an installed CuPy") @pytest.mark.parametrize("mapping", CUDA_DOMAIN_MAPPINGS) def test_cuda_args_domain_fake_cupy(mapping): @@ -1207,10 +1236,7 @@ def test_cuda_args_domain_fake_cupy(mapping): Runs in a subprocess because the fake CuPy must be installed before cunumpy is imported. """ code = f"from struphy.geometry.tests.test_domain import check_cuda_args_domain; check_cuda_args_domain({mapping!r})" - result = subprocess.run( - [sys.executable, "-c", code], env=serial_child_env(CUNUMPY_FAKE_CUPY="1"), capture_output=True, text=True - ) - assert result.returncode == 0, result.stderr[-4000:] + run_fake_cupy_child(code) @requires_cupy diff --git a/src/struphy/pic/tests/test_accum_matrix_cupy.py b/src/struphy/pic/tests/test_accum_matrix_cupy.py index cf0aa62c6..4537e8670 100644 --- a/src/struphy/pic/tests/test_accum_matrix_cupy.py +++ b/src/struphy/pic/tests/test_accum_matrix_cupy.py @@ -6,8 +6,6 @@ cunumpy's fake CuPy with CPU-emulated kernel launches (in a subprocess, see :func:`check_accumulator_matrix_fake_cupy`). """ -import subprocess -import sys from dataclasses import dataclass from typing import Any @@ -16,7 +14,7 @@ import pytest from cunumpy.kernel_testing import emulation_compiler, requires_cupy -from struphy.geometry.tests.test_domain import _cupy_installed, serial_child_env +from struphy.geometry.tests.test_domain import _cupy_installed, run_fake_cupy_child ALPHA, KAPPA, VTH, DT = 1.3, 0.7, 0.9, 0.05 N_SCHUR_ITERATIONS = 3 @@ -203,7 +201,4 @@ def test_accumulator_matrix_fake_cupy(): subprocess because the fake CuPy must be installed before cunumpy is imported. """ code = "from struphy.pic.tests.test_accum_matrix_cupy import check_accumulator_matrix_fake_cupy as c; c()" - result = subprocess.run( - [sys.executable, "-c", code], env=serial_child_env(CUNUMPY_FAKE_CUPY="1"), capture_output=True, text=True - ) - assert result.returncode == 0, result.stderr[-4000:] + run_fake_cupy_child(code) diff --git a/src/struphy/simulation/tests/test_compile_cuda_kernels.py b/src/struphy/simulation/tests/test_compile_cuda_kernels.py index 5e7114366..97bec9e23 100644 --- a/src/struphy/simulation/tests/test_compile_cuda_kernels.py +++ b/src/struphy/simulation/tests/test_compile_cuda_kernels.py @@ -5,15 +5,12 @@ (in a subprocess); the test marked ``requires_cupy`` compiles the kernels of a small model on a GPU. """ -import subprocess -import sys - import cunumpy import pytest from cunumpy.kernels import CudaKernel, Kernel from struphy import DerhamOptions, EnvironmentOptions, Simulation, Time, domains, equils, grids, maxwellians -from struphy.geometry.tests.test_domain import _cupy_installed, serial_child_env +from struphy.geometry.tests.test_domain import _cupy_installed, run_fake_cupy_child from struphy.models import Vlasov, VlasovAmpereOneSpecies from struphy.particles.parameters import BoundaryParameters, LoadingParameters from struphy.pic.pushing.kernels.push_bxu_Hdiv import push_bxu_Hdiv @@ -225,13 +222,7 @@ def test_compile_cuda_kernels_on_cupy_backend(tmp_path, monkeypatch): @pytest.mark.skipif(_cupy_installed(), reason="the fake CuPy cannot replace an installed CuPy") def test_compile_cuda_kernels_fake_cupy(): """With cunumpy's fake CuPy the CuPy backend is active: kernels are checked and compiled (compile recorded).""" - result = subprocess.run( - [sys.executable, "-c", FAKE_CUPY_CHECK], - env=serial_child_env(CUNUMPY_FAKE_CUPY="1"), - capture_output=True, - text=True, - ) - assert result.returncode == 0, result.stderr[-4000:] + run_fake_cupy_child(FAKE_CUPY_CHECK) @requires_cupy From 8223abfaff756a844fa20ba7f5736cbb567d40e8 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 9 Oct 2026 15:13:18 +0200 Subject: [PATCH 3/3] Skip the new GVEC tests by default (with_gvec=False), like the other GVEC tests GVEC dies with SIGILL on the CI runners (in gvec's library init), which took down the serial and the 2-rank unit test jobs. Co-Authored-By: Claude Opus 5.5 --- .../fields_background/tests/test_equils_cupy.py | 12 +++++++----- .../fields_background/tests/test_gvec_equil.py | 10 ++++++++-- 2 files changed, 15 insertions(+), 7 deletions(-) diff --git a/src/struphy/fields_background/tests/test_equils_cupy.py b/src/struphy/fields_background/tests/test_equils_cupy.py index e7710570d..28614e65a 100644 --- a/src/struphy/fields_background/tests/test_equils_cupy.py +++ b/src/struphy/fields_background/tests/test_equils_cupy.py @@ -214,7 +214,9 @@ def check_equil_on_cupy(case): assert np.allclose(o, r, rtol=1e-12, atol=1e-12, equal_nan=True), (case, where, name) -def _check_needs(case): +def _check_needs(case, with_gvec=False): + if EQUIL_CASES[case][0] == "GVECequilibrium" and not with_gvec: + pytest.skip("GVEC not tested here (with_gvec=False), like the other GVEC tests") package = NEEDS.get(EQUIL_CASES[case][0]) if package is not None and importlib.util.find_spec(package) is None: pytest.skip(f"{package} is not installed") @@ -248,12 +250,12 @@ def test_host_call_numpy_is_plain_call(): @pytest.mark.skipif(_cupy_installed(), reason="the fake CuPy cannot replace an installed CuPy") @pytest.mark.parametrize("case", list(EQUIL_CASES)) -def test_equil_fake_cupy(case): +def test_equil_fake_cupy(case, with_gvec=False): """Without a GPU: :func:`check_equil_on_cupy` with cunumpy's fake CuPy, which rejects host/device mixing. Runs in a subprocess because the fake CuPy must be installed before cunumpy is imported. """ - _check_needs(case) + _check_needs(case, with_gvec) code = ( "from struphy.fields_background.tests.test_equils_cupy import check_equil_on_cupy, host_geometry_kernels\n" f"with host_geometry_kernels(): check_equil_on_cupy({case!r})" @@ -263,7 +265,7 @@ def test_equil_fake_cupy(case): @requires_cupy @pytest.mark.parametrize("case", list(EQUIL_CASES)) -def test_equil_cupy(case): +def test_equil_cupy(case, with_gvec=False): """On a GPU: the equilibria are created and evaluated on CuPy and agree with NumPy.""" - _check_needs(case) + _check_needs(case, with_gvec) check_equil_on_cupy(case) diff --git a/src/struphy/fields_background/tests/test_gvec_equil.py b/src/struphy/fields_background/tests/test_gvec_equil.py index 40b84e065..d21017ede 100644 --- a/src/struphy/fields_background/tests/test_gvec_equil.py +++ b/src/struphy/fields_background/tests/test_gvec_equil.py @@ -17,12 +17,15 @@ def _markers(n=7, seed=3): @pytest.mark.parametrize("name", ["bv", "jv", "p0", "n0", "absB0"]) -def test_gvec_marker_evaluation_matches_grid(name): +def test_gvec_marker_evaluation_matches_grid(name, with_gvec=False): """Marker (flat) evaluation gives one value per marker, equal to the grid evaluation at each marker. Before, gvec evaluated the markers' rho, theta and zeta as a tensor grid, so ``bv`` and ``jv`` at N markers returned N x N x N arrays, and ``absB0`` failed. """ + if not with_gvec: + pytest.skip("GVEC not tested here (with_gvec=False), like the other GVEC tests") + from struphy.fields_background.equils import GVECequilibrium equil = GVECequilibrium(**GVEC_PARAMS) @@ -38,8 +41,11 @@ def test_gvec_marker_evaluation_matches_grid(name): assert np.isclose(f[ip], g.ravel()[0], rtol=1e-12, atol=1e-14) -def test_gvec_marker_evaluation_boozer_raises(): +def test_gvec_marker_evaluation_boozer_raises(with_gvec=False): """gvec computes the Boozer transform per flux surface, so marker evaluation with ``use_boozer=True`` raises.""" + if not with_gvec: + pytest.skip("GVEC not tested here (with_gvec=False), like the other GVEC tests") + from struphy.fields_background.equils import GVECequilibrium equil = GVECequilibrium(**GVEC_PARAMS, use_boozer=True)