Struphy with cupy

src/struphy/bsplines/bsplines.py

@@ -611,8 +611,8 @@ def make_knots(breaks, degree, periodic):
 
     if periodic:
         period = breaks[-1] - breaks[0]
-        T[0:p] = [xi - period for xi in breaks[-p - 1 : -1]]
-        T[-p:] = [xi + period for xi in breaks[1 : p + 1]]
+        T[0:p] = xp.asarray([xi - period for xi in breaks[-p - 1 : -1]])
+        T[-p:] = xp.asarray([xi + period for xi in breaks[1 : p + 1]])
     else:
         T[0:p] = breaks[0]
         T[-p:] = breaks[-1]

src/struphy/feec/basis_projection_ops.py

@@ -51,7 +51,7 @@ def __init__(self, derham, domain, verbose=True, **weights):
 
         self._rank = derham.comm.Get_rank() if derham.comm is not None else 0
 
-        if xp.any([p == 1 and Nel > 1 for p, Nel in zip(derham.p, derham.Nel)]):
+        if xp.any(xp.array([p == 1 and Nel > 1 for p, Nel in zip(derham.p, derham.Nel)])):
             if MPI.COMM_WORLD.Get_rank() == 0:
                 print(
                     f'\nWARNING: Class "BasisProjectionOperators" called with p={derham.p} (interpolation of piece-wise constants should be avoided).',

src/struphy/feec/psydac_derham.py

@@ -2,6 +2,7 @@
 import importlib.metadata
 
 import cunumpy as xp
+import numpy as np
 import psydac.core.bsplines as bsp
 from psydac.ddm.cart import DomainDecomposition
 from psydac.ddm.mpi import MockComm, MockMPI
@@ -117,7 +118,7 @@ def __init__(
         if dirichlet_bc is not None:
             assert len(dirichlet_bc) == 3
             # make sure that boundary conditions are compatible with spline space
-            assert xp.all([bc == (False, False) for i, bc in enumerate(dirichlet_bc) if spl_kind[i]])
+            assert xp.all(xp.array([bc == (False, False) for i, bc in enumerate(dirichlet_bc) if spl_kind[i]]))
 
         self._dirichlet_bc = dirichlet_bc
 
@@ -2789,12 +2790,17 @@ def get_pts_and_wts(space_1d, start, end, n_quad=None, polar_shift=False):
             union_breaks = space_1d.breaks[:-1]
 
         # Make union of Greville and break points
-        tmp = set(xp.round(space_1d.histopolation_grid, decimals=14)).union(
-            xp.round(union_breaks, decimals=14),
+        # tmp = set(xp.round(space_1d.histopolation_grid, decimals=14)).union(
+        #     xp.round(union_breaks, decimals=14),
+        # )
+        # tmp = list(tmp)
+        # tmp.sort()
+        # tmp_a = xp.array(tmp)
+
+        tmp = set(xp.round(space_1d.histopolation_grid, decimals=14).tolist()).union(
+            xp.round(union_breaks, decimals=14).tolist()
         )
-
-        tmp = list(tmp)
-        tmp.sort()
+        tmp = sorted(tmp)
         tmp_a = xp.array(tmp)
 
         x_grid = tmp_a[
@@ -2822,7 +2828,13 @@ def get_pts_and_wts(space_1d, start, end, n_quad=None, polar_shift=False):
             # products of basis functions are integrated exactly
             n_quad = space_1d.degree + 1
 
-        pts_loc, wts_loc = xp.polynomial.legendre.leggauss(n_quad)
+        pts_loc, wts_loc = np.polynomial.legendre.leggauss(n_quad)
+
+        if "cupy" in xp.__name__:
+            import cupy as cp
+
+            pts_loc = cp.array(pts_loc)
+            wts_loc = cp.array(wts_loc)
 
         x, wts = bsp.quadrature_grid(x_grid, pts_loc, wts_loc)
 

src/struphy/geometry/base.py

@@ -12,6 +12,7 @@
 from struphy.geometry import evaluation_kernels, transform_kernels
 from struphy.kernel_arguments.pusher_args_kernels import DomainArguments
 from struphy.linear_algebra import linalg_kron
+from struphy.utils.pyccel import Pyccelkernel
 
 
 class Domain(metaclass=ABCMeta):
@@ -769,8 +770,8 @@ def _evaluate_metric_coefficient(self, *etas, which=0, **kwargs):
 
             # to keep C-ordering the (3, 3)-part is in the last indices
             out = xp.empty((markers.shape[0], 3, 3), dtype=float)
-
-            n_inside = evaluation_kernels.kernel_evaluate_pic(
+            kernel = Pyccelkernel(evaluation_kernels.kernel_evaluate_pic)
+            n_inside = kernel(
                 markers,
                 which,
                 self.args_domain,
@@ -813,7 +814,8 @@ def _evaluate_metric_coefficient(self, *etas, which=0, **kwargs):
                 (E1.shape[0], E2.shape[1], E3.shape[2], 3, 3),
                 dtype=float,
             )
-            evaluation_kernels.kernel_evaluate(
+            kernel = Pyccelkernel(evaluation_kernels.kernel_evaluate)
+            kernel(
                 E1,
                 E2,
                 E3,

src/struphy/geometry/evaluation_kernels.py

@@ -3,7 +3,7 @@
 corresponding to mappings (x, y, z) = F(eta_1, eta_2, eta_3).
 """
 
-from numpy import empty, shape, zeros
+from cupy import empty, shape, zeros
 from pyccel.decorators import stack_array
 
 import struphy.geometry.mappings_kernels as mappings_kernels

src/struphy/geometry/transform_kernels.py

@@ -43,7 +43,14 @@
 - 2-form --> vector : (a_1, a_2, a_3) =          (a^2_1, a^2_2, a^2_3) / |det(DF)|
 """
 
-from numpy import empty, shape, sqrt, zeros
+import cunumpy as np
+from cunumpy.xp import array_backend
+
+if array_backend.backend == "cupy":
+    from cupy import empty, shape, sqrt, zeros
+else:
+    from numpy import empty, shape, sqrt, zeros
+
 from pyccel.decorators import stack_array
 
 import struphy.geometry.evaluation_kernels as evaluation_kernels

src/struphy/io/output_handling.py

@@ -1,8 +1,8 @@
 import ctypes
 import os
 
-import cunumpy as xp
 import h5py
+import numpy as np
 
 
 class DataContainer:
@@ -82,11 +82,11 @@ def add_data(self, data_dict):
         Parameters
         ----------
         data_dict : dict
-            Name-object pairs to save during time stepping, e.g. {key : val}. key must be a string and val must be a xp.array of fixed shape. Scalar values (floats) must therefore be passed as 1d arrays of size 1.
+            Name-object pairs to save during time stepping, e.g. {key : val}. key must be a string and val must be a np.array of fixed shape. Scalar values (floats) must therefore be passed as 1d arrays of size 1.
         """
 
         for key, val in data_dict.items():
-            assert isinstance(val, xp.ndarray)
+            assert isinstance(val, np.ndarray)
 
             # if dataset already exists, check for compatibility with given array
             if key in self._dset_dict:

src/struphy/kernel_arguments/pusher_args_kernels.py

@@ -1,5 +1,5 @@
 # from numpy import copy
-from numpy import empty
+import cunumpy as xp
 
 
 class MarkerArguments:
@@ -99,12 +99,12 @@ def __init__(
         self.tn3 = tn3
         self.starts = starts
 
-        self.bn1 = empty(pn[0] + 1, dtype=float)
-        self.bn2 = empty(pn[1] + 1, dtype=float)
-        self.bn3 = empty(pn[2] + 1, dtype=float)
-        self.bd1 = empty(pn[0], dtype=float)
-        self.bd2 = empty(pn[1], dtype=float)
-        self.bd3 = empty(pn[2], dtype=float)
+        self.bn1 = xp.empty(int(pn[0] + 1), dtype=float)
+        self.bn2 = xp.empty(int(pn[1] + 1), dtype=float)
+        self.bn3 = xp.empty(int(pn[2] + 1), dtype=float)
+        self.bd1 = xp.empty(int(pn[0]), dtype=float)
+        self.bd2 = xp.empty(int(pn[1]), dtype=float)
+        self.bd3 = xp.empty(int(pn[2]), dtype=float)
 
 
 class DomainArguments:

src/struphy/main.py

@@ -254,12 +254,25 @@ def run(
                 absB0 = model.equil.absB0(*grids_log)
                 pointData["absB0"] = absB0
 
-        gridToVTK(os.path.join(path_out, "geometry"), *grids_phy, pointData=pointData)
+        from cunumpy.xp import array_backend
+
+        print("calling gridToVTK")
+        if array_backend.backend == "numpy":
+            gridToVTK(os.path.join(path_out, "geometry"), *grids_phy, pointData=pointData)
+        else:
+            # cupy
+            grids_phy_cpu = [g.get() if isinstance(g, xp.ndarray) else g for g in grids_phy]
+
+            # Convert pointData values to NumPy
+            pointData_cpu = {k: (v.get() if isinstance(v, xp.ndarray) else v) for k, v in pointData.items()}
+
+            # Now call gridToVTK safely
+            gridToVTK(os.path.join(path_out, "geometry"), *grids_phy_cpu, pointData=pointData_cpu)
 
     # data object for saving (will either create new hdf5 files if restart==False or open existing files if restart==True)
     # use MPI.COMM_WORLD as communicator when storing the outputs
     data = DataContainer(path_out, comm=comm)
-
+    print("ccc")
     # time quantities (current time value, value in seconds and index)
     time_state = {}
     time_state["value"] = xp.zeros(1, dtype=float)
@@ -270,8 +283,17 @@ def run(
     for key, val in time_state.items():
         key_time = "time/" + key
         key_time_restart = "restart/time/" + key
-        data.add_data({key_time: val})
-        data.add_data({key_time_restart: val})
+        if array_backend.backend == "numpy":
+            data.add_data({key_time: val})
+            data.add_data({key_time_restart: val})
+        else:
+            val_cpu = val.get()  #  if isinstance(val, xp.ndarray) else val
+
+            # Then assign
+            print(f"{val_cpu = } {type(val_cpu) = }")
+            data.add_data({key_time: val_cpu})
+            # self._file[key][0] = val_cpu[0]
+            data.add_data({key_time_restart: val_cpu})
 
     # retrieve time parameters
     dt = time_opts.dt

src/struphy/models/base.py

@@ -5,6 +5,7 @@
 from functools import reduce
 from textwrap import indent
 
+import numpy as np
 import cunumpy as xp
 import yaml
 from line_profiler import profile
@@ -521,11 +522,12 @@ def update_scalar(self, name, value=None):
 
         if summands is None:
             # Ensure the value is a float if there are no summands
-            assert isinstance(value, float)
-
-            # Create a numpy array to hold the scalar value
-            value_array = xp.array([value], dtype=xp.float64)
-
+            if isinstance(value, float):
+                # Create a numpy array to hold the scalar value
+                value_array = np.array([value])
+            else:
+                value_array = np.asarray(value)
+            value_array = np.array(value_array)
             # Perform MPI operations based on the compute flags
             if "sum_world" in compute_operations and not isinstance(MPI, MockMPI):
                 MPI.COMM_WORLD.Allreduce(
@@ -562,11 +564,14 @@ def update_scalar(self, name, value=None):
 
             if "divide_n_mks" in compute_operations:
                 # Initialize the total number of markers
-                n_mks_tot = xp.array([variable.particles.Np])
+                n_mks_tot = np.array([variable.particles.Np])
                 value_array /= n_mks_tot
 
             # Update the scalar value
-            self._scalar_quantities[name]["value"][0] = value_array[0]
+            if value_array.ndim == 0:
+                self._scalar_quantities[name]["value"][0] = value_array.item()
+            else:
+                self._scalar_quantities[name]["value"][0] = value_array[0]
 
         else:
             # Sum the values of the summands
@@ -776,7 +781,7 @@ def update_distr_functions(self):
                     h2 = 1 / obj.boxes_per_dim[1]
                     h3 = 1 / obj.boxes_per_dim[2]
 
-                    ndim = xp.count_nonzero([d > 1 for d in obj.boxes_per_dim])
+                    ndim = xp.count_nonzero(xp.array([d > 1 for d in obj.boxes_per_dim]))
                     if ndim == 0:
                         kernel_type = "gaussian_3d"
                     else:

src/struphy/models/toy.py

@@ -1,3 +1,4 @@
+import numpy as np
 import cunumpy as xp
 from psydac.ddm.mpi import mpi as MPI
 
@@ -157,7 +158,7 @@ def velocity_scale(self):
         return "cyclotron"
 
     def allocate_helpers(self):
-        self._tmp = xp.empty(1, dtype=float)
+        self._tmp = np.empty(1, dtype=float)
 
     def update_scalar_quantities(self):
         particles = self.kinetic_ions.var.particles

src/struphy/ode/utils.py

@@ -77,7 +77,8 @@ def __post_init__(self):
         self._a = xp.tri(self.n_stages, k=-1)
         for l, st in enumerate(a):
             assert len(st) == l + 1
-            self._a[l + 1, : l + 1] = st
+
+            self._a[l + 1, : l + 1] = xp.array(st)
 
         self._conv_rate = conv_rate
 

src/struphy/pic/base.py

@@ -235,7 +235,7 @@ def __init__(
             assert all([nboxes % nproc == 0 for nboxes, nproc in zip(self.boxes_per_dim, self.nprocs)]), (
                 f"Number of boxes {self.boxes_per_dim =} must be divisible by number of processes {self.nprocs =} in each direction."
             )
-            n_boxes = xp.prod(self.boxes_per_dim, dtype=int) * self.num_clones
+            n_boxes = xp.prod(xp.array(self.boxes_per_dim), dtype=int) * self.num_clones
 
         # if verbose:
         #     print(f"\n{self.mpi_rank = }, {n_boxes = }")
@@ -1095,7 +1095,7 @@ def _allocate_marker_array(self):
         bufsize = self.bufsize + 1.0 / xp.sqrt(n_mks_load_loc)
 
         # allocate markers array (3 x positions, vdim x velocities, weight, s0, w0, ..., ID) with buffer
-        self._n_rows = round(n_mks_load_loc * (1 + bufsize))
+        self._n_rows = round(float(n_mks_load_loc * (1 + bufsize)))
         self._markers = xp.zeros((self.n_rows, self.n_cols), dtype=float)
 
         # allocate auxiliary arrays
@@ -1917,7 +1917,7 @@ def binning(
             The reconstructed delta-f distribution function.
         """
 
-        assert xp.count_nonzero(components) == len(bin_edges)
+        assert xp.count_nonzero(xp.array(components)) == len(bin_edges)
 
         # volume of a bin
         bin_vol = 1.0
@@ -2424,7 +2424,7 @@ def _set_boxes(self):
             n_particles = self._markers_shape[0]
             n_mkr = int(n_particles / n_box_in) + 1
             n_cols = round(
-                n_mkr * (1 + 1 / xp.sqrt(n_mkr) + self._box_bufsize),
+                float(n_mkr) * (1 + 1 / float(xp.sqrt(n_mkr)) + self._box_bufsize),
             )
 
             # cartesian boxes
@@ -2643,7 +2643,17 @@ def check_and_assign_particles_to_boxes(self):
         """Check whether the box array has enough columns (detect load imbalance wrt to sorting boxes),
         and then assigne the particles to boxes."""
 
-        bcount = xp.bincount(xp.int64(self.markers_wo_holes[:, -2]))
+        from cunumpy.xp import array_backend
+
+        if array_backend.backend == "numpy":
+            bcount = xp.bincount(xp.int64(self.markers_wo_holes[:, -2]))
+        else:
+            import cupy as cp
+
+            indices = self.markers_wo_holes[:, -2]
+            indices = indices.astype(cp.int64)
+            bcount = cp.bincount(indices)
+
         max_in_box = xp.max(bcount)
         if max_in_box > self._sorting_boxes.boxes.shape[1]:
             warnings.warn(
@@ -4049,6 +4059,7 @@ def _gather_scalar_in_subcomm_array(self, scalar: int, out: xp.ndarray = None):
         _tmp[self.mpi_rank] = scalar
 
         if self.mpi_comm is not None:
+            print(f"{self.mpi_comm = }")
             self.mpi_comm.Allgather(
                 _tmp[self.mpi_rank],
                 _tmp,

src/struphy/propagators/propagators_markers.py

@@ -104,7 +104,7 @@ def allocate(self):
         import cunumpy as xp
 
         butcher._a = xp.diag(butcher.a, k=-1)
-        butcher._a = xp.array(list(butcher.a) + [0.0])
+        butcher._a = xp.concatenate([xp.asarray(butcher.a), xp.array([0.0])])
 
         args_kernel = (
             butcher.a,
@@ -867,7 +867,7 @@ def allocate(self):
             import cunumpy as xp
 
             butcher._a = xp.diag(butcher.a, k=-1)
-            butcher._a = xp.array(list(butcher.a) + [0.0])
+            butcher._a = xp.concatenate([xp.asarray(butcher.a), xp.array([0.0])])
 
             kernel = Pyccelkernel(pusher_kernels_gc.push_gc_bxEstar_explicit_multistage)
 
@@ -1315,7 +1315,7 @@ def allocate(self):
             import cunumpy as xp
 
             butcher._a = xp.diag(butcher.a, k=-1)
-            butcher._a = xp.array(list(butcher.a) + [0.0])
+            butcher._a = xp.concatenate([xp.asarray(butcher.a), xp.array([0.0])])
 
             kernel = Pyccelkernel(pusher_kernels_gc.push_gc_Bstar_explicit_multistage)