From 8b7d737fd46cdd27a9cdaa707da6e686e1d4cf0c Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 10:00:05 +0200 Subject: [PATCH] CuPy backend: host/device conversions in post-processing; math.sqrt for host scalars in Species Post-processing evaluates fields on the device and writes them from the host: the evaluation grids, point data and coordinates are converted where they cross, and HDF5 data is moved to the device before it is written into the spline coefficients. Co-Authored-By: Claude Opus 5.5 --- src/struphy/models/species.py | 4 ++-- src/struphy/post_processing/output.py | 31 ++++++++++++++++++--------- 2 files changed, 23 insertions(+), 12 deletions(-) diff --git a/src/struphy/models/species.py b/src/struphy/models/species.py index b43464713..0ad27067d 100644 --- a/src/struphy/models/species.py +++ b/src/struphy/models/species.py @@ -1,9 +1,9 @@ import logging +import math import warnings from abc import ABCMeta, abstractmethod from dataclasses import fields, is_dataclass -import cunumpy as xp from maybempi import MPI from struphy.models.variables import Variable @@ -220,7 +220,7 @@ def __init__( con = ConstantsOfNature() # relevant frequencies - om_p = xp.sqrt(units.n * (Z * con.e) ** 2 / (con.eps0 * A * con.mH)) + om_p = math.sqrt(units.n * (Z * con.e) ** 2 / (con.eps0 * A * con.mH)) om_c = Z * con.e * units.B / (A * con.mH) # compute equation parameters diff --git a/src/struphy/post_processing/output.py b/src/struphy/post_processing/output.py index bfcf5fd58..cc65d56f8 100644 --- a/src/struphy/post_processing/output.py +++ b/src/struphy/post_processing/output.py @@ -85,6 +85,13 @@ def load_plasma_plots() -> bool: return True +def _to_host(value): + """`value` (an array, or a list/tuple of arrays for vector-valued fields) as NumPy array(s).""" + if isinstance(value, (list, tuple)): + return type(value)(_to_host(v) for v in value) + return xp.to_numpy(value) + + class ProductMapping(Mapping[str, xr.DataArray]): """A discoverable mapping whose products are loaded on first access.""" @@ -568,13 +575,14 @@ def _attach_physical_coords( self, array: xr.DataArray, etas: tuple[Any, Any, Any], dims: tuple[str, ...] ) -> xr.DataArray: """Map a raw FEEC evaluation grid and attach its Cartesian coordinates.""" - grids = [np.atleast_1d(np.asarray(eta, dtype=float)) for eta in etas] - mapped = self.domain(*grids) + # the logical grids may be device arrays (CuPy backend): map them where they live, keep the coordinates on the host + grids = [np.atleast_1d(np.asarray(xp.to_numpy(eta), dtype=float)) for eta in etas] + mapped = self.domain(*(xp.asarray(grid) for grid in grids)) shape = tuple(len(grid) for grid in grids) keep = tuple(slice(None) if dim in dims else 0 for dim in ("eta1", "eta2", "eta3")) coordinates = {} for name, values in zip(("X", "Y", "Z"), mapped): - values = np.asarray(values).reshape(shape)[keep] + values = np.asarray(xp.to_numpy(values)).reshape(shape)[keep] coordinates[name] = (dims, values) if dims else values.item() return array.assign_coords(coordinates) @@ -1295,7 +1303,9 @@ def _process_fields( # the physical grid is only needed for output, hence it is only built on rank 0 if self._pproc_rank == 0: - grids_phy = list(self.domain(*grids_log)) + # output coordinates, on the host (the logical grids stay on the device for the field evaluation) + grids_phy = [xp.to_numpy(grid) for grid in self.domain(*grids_log)] + grids_log_host = [xp.to_numpy(grid) for grid in grids_log] else: grids_phy = None @@ -1327,8 +1337,9 @@ def _process_fields( if self._pproc_rank == 0: for species, vars in vals.items(): for name, val in vars.items(): - point_data[species][name][t] = val - point_data_phy[species][name][t] = vals_phy[species][name] + # evaluated on the device on the CuPy backend; written to files from the host + point_data[species][name][t] = _to_host(val) + point_data_phy[species][name][t] = _to_host(vals_phy[species][name]) # directory for the vtk files path_fields = os.path.join(self.path_pproc, "fields_data") @@ -1338,10 +1349,10 @@ def _process_fields( for species, vars in point_data.items(): variables = {} for name, val in vars.items(): - variables[name] = wrap_field_data(val, grids_log, grids_phy=grids_phy, name=name) + variables[name] = wrap_field_data(val, grids_log_host, grids_phy=grids_phy, name=name) if physical: variables[name + "_xyz"] = wrap_field_data( - point_data_phy[species][name], grids_log, grids_phy=grids_phy, name=name + "_xyz" + point_data_phy[species][name], grids_log_host, grids_phy=grids_phy, name=name + "_xyz" ) store.write_group(store.store_path(self.path_pproc), f"/{species}", xr.Dataset(variables)) @@ -1517,7 +1528,7 @@ def _load_femfields(self, fields: dict, files: list, n: int, step: int = 1): s1 : e1 + 1, s2 : e2 + 1, s3 : e3 + 1, - ] = ddset[n * step, p1:-p1, p2:-p2, p3:-p3] + ] = xp.asarray(ddset[n * step, p1:-p1, p2:-p2, p3:-p3]) # HDF5 data is on the host # vector-valued field else: @@ -1530,7 +1541,7 @@ def _load_femfields(self, fields: dict, files: list, n: int, step: int = 1): s1 : e1 + 1, s2 : e2 + 1, s3 : e3 + 1, - ] = ddset[str(comp + 1)][n * step, p1:-p1, p2:-p2, p3:-p3] + ] = xp.asarray(ddset[str(comp + 1)][n * step, p1:-p1, p2:-p2, p3:-p3]) vector.update_ghost_regions() if is_polar: