GPU stack 2/8: host/device conversions in post-processing; math.sqrt in Species

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
@@ -203,7 +203,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

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: