diff --git a/feectools b/feectools index 6e50ffa15..e5c084aee 160000 --- a/feectools +++ b/feectools @@ -1 +1 @@ -Subproject commit 6e50ffa1596565a650d4ccf04e62f40a8e275d70 +Subproject commit e5c084aee1e160754f41f1211b0f1b23ca468a0c diff --git a/params_LinearMHDDriftkineticCC.py b/params_LinearMHDDriftkineticCC.py new file mode 100644 index 000000000..f9f5ecd72 --- /dev/null +++ b/params_LinearMHDDriftkineticCC.py @@ -0,0 +1,116 @@ +# import model, set verbosity +from struphy.models.hybrid import LinearMHDDriftkineticCC + +from struphy import main +from struphy.fields_background import equils +from struphy.geometry import domains +from struphy.initial import perturbations +from struphy.io.options import BaseUnits, DerhamOptions, EnvironmentOptions, FieldsBackground, Time +from struphy.kinetic_background import maxwellians +from struphy.pic.utilities import ( + BinningPlot, + BoundaryParameters, + KernelDensityPlot, + LoadingParameters, + WeightsParameters, +) +from struphy.topology import grids + +# environment options +env = EnvironmentOptions() + +# units +base_units = BaseUnits() + +# time stepping +time_opts = Time() + +# geometry +domain = domains.Cuboid() + +# fluid equilibrium (can be used as part of initial conditions) +equil = equils.HomogenSlab() + +# grid +grid = grids.TensorProductGrid(Nel=(16, 16, 16)) + +# derham options +derham_opts = DerhamOptions() + +# light-weight model instance +model = LinearMHDDriftkineticCC() + +# species parameters +model.mhd.set_phys_params() +model.energetic_ions.set_phys_params() + +loading_params = LoadingParameters(ppc=1000) +weights_params = WeightsParameters() +boundary_params = BoundaryParameters() +model.energetic_ions.set_markers( + loading_params=loading_params, + weights_params=weights_params, + boundary_params=boundary_params, +) +model.energetic_ions.set_sorting_boxes() +model.energetic_ions.set_save_data() + +# propagator options +model.propagators.push_bxe.options = model.propagators.push_bxe.Options( + b_tilde=model.em_fields.b_field, +) +model.propagators.push_parallel.options = model.propagators.push_parallel.Options( + b_tilde=model.em_fields.b_field, +) +model.propagators.shearalfen_cc5d.options = model.propagators.shearalfen_cc5d.Options( + energetic_ions=model.energetic_ions.var, +) +model.propagators.magnetosonic.options = model.propagators.magnetosonic.Options( + b_field=model.em_fields.b_field, +) +model.propagators.cc5d_density.options = model.propagators.cc5d_density.Options( + energetic_ions=model.energetic_ions.var, + b_tilde=model.em_fields.b_field, +) +model.propagators.cc5d_gradb.options = model.propagators.cc5d_gradb.Options( + b_tilde=model.em_fields.b_field, +) +model.propagators.cc5d_curlb.options = model.propagators.cc5d_curlb.Options( + b_tilde=model.em_fields.b_field, +) + +# background, perturbations and initial conditions +model.mhd.velocity.add_background(FieldsBackground()) +model.mhd.velocity.add_perturbation(perturbations.TorusModesCos(given_in_basis="v", comp=0)) +model.mhd.velocity.add_perturbation(perturbations.TorusModesCos(given_in_basis="v", comp=1)) +model.mhd.velocity.add_perturbation(perturbations.TorusModesCos(given_in_basis="v", comp=2)) +maxwellian_1 = maxwellians.GyroMaxwellian2D(n=(1.0, None), equil=equil) +maxwellian_2 = maxwellians.GyroMaxwellian2D(n=(0.1, None), equil=equil) +background = maxwellian_1 + maxwellian_2 +model.energetic_ions.var.add_background(background) + +# if .add_initial_condition is not called, the background is the kinetic initial condition +perturbation = perturbations.TorusModesCos() +maxwellian_1pt = maxwellians.GyroMaxwellian2D(n=(1.0, perturbation), equil=equil) +init = maxwellian_1pt + maxwellian_2 +model.energetic_ions.var.add_initial_condition(init) + +# optional: exclude variables from saving +# model.energetic_ions.var.save_data = False + +if __name__ == "__main__": + # start run + verbose = True + + main.run( + model, + params_path=__file__, + env=env, + base_units=base_units, + time_opts=time_opts, + domain=domain, + equil=equil, + grid=grid, + derham_opts=derham_opts, + verbose=verbose, + ) diff --git a/pyproject.toml b/pyproject.toml index 97c62e82b..39bf76759 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,9 +24,9 @@ classifiers = [ ] dependencies = [ "numpy<=2.5.0", - "cunumpy<=0.1.1", + "cunumpy>=0.1.4, <=0.1.5", "pyccel>=2.2.0, <=2.2.3", - "feectools<=0.1.3", + "feectools<=0.1.8", "scipy<=1.18.0", "h5py<=3.16.0", "matplotlib<=3.11.0", diff --git a/src/struphy/bsplines/bsplines.py b/src/struphy/bsplines/bsplines.py index 09fc2a7de..77690ab90 100644 --- a/src/struphy/bsplines/bsplines.py +++ b/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] diff --git a/src/struphy/console/format.py b/src/struphy/console/format.py index 51544413d..6cc46d378 100644 --- a/src/struphy/console/format.py +++ b/src/struphy/console/format.py @@ -492,9 +492,7 @@ def get_python_files(input_type, path=None): # python_files = [f for f in files if f.endswith(".py") and os.path.isfile(f)] python_files = [ - os.path.join(repopath, f) - for f in files - if (f.endswith(".py") or f.endswith(".ipynb")) and os.path.isfile(os.path.join(repopath, f)) + os.path.join(repopath, f) for f in files if f.endswith(".py") and os.path.isfile(os.path.join(repopath, f)) ] if not python_files: diff --git a/src/struphy/console/test.py b/src/struphy/console/test.py index fc4d190a7..24ad3a879 100644 --- a/src/struphy/console/test.py +++ b/src/struphy/console/test.py @@ -73,6 +73,7 @@ def struphy_test( cmd = [ "pytest", "--testmon", + "-xvs", ] + list_of_tests if with_desc: @@ -156,7 +157,6 @@ def struphy_test( "pytest", "-m", "single", - "-xvs", # "--with-mpi", "--model-name", group, diff --git a/src/struphy/feec/basis_projection_ops.py b/src/struphy/feec/basis_projection_ops.py index b82755484..cf843751e 100644 --- a/src/struphy/feec/basis_projection_ops.py +++ b/src/struphy/feec/basis_projection_ops.py @@ -1,6 +1,8 @@ import logging import cunumpy as xp +import numpy as np +from cunumpy import PyccelKernel from feectools.api.settings import PSYDAC_BACKEND_GPYCCEL from feectools.ddm.mpi import mpi as MPI from feectools.fem.basic import FemSpace @@ -19,7 +21,6 @@ from struphy.polar.basic import PolarDerhamSpace, PolarVector from struphy.polar.linear_operators import PolarExtractionOperator from struphy.utils.docstring_converter import auto_convert_docstring -from struphy.utils.pyccel import Pyccelkernel logger = logging.getLogger("struphy") @@ -53,9 +54,7 @@ def __init__(self, derham, domain, **weights): self._rank = derham.comm.Get_rank() if derham.comm is not None else 0 - if xp.any( - [degree == 1 and num_elements > 1 for degree, num_elements in zip(derham.degree, derham.num_elements)] - ): + if any([degree == 1 and num_elements > 1 for degree, num_elements in zip(derham.degree, derham.num_elements)]): logger.warning( f'WARNING: Class "BasisProjectionOperators" called with degree={derham.degree} (interpolation of piece-wise constants should be avoided).', ) @@ -1057,11 +1056,11 @@ def __init__( if isinstance(V, TensorFemSpace): self._Vspaces = [V.coeff_space] self._V1ds = [V.spaces] - self._VNbasis = xp.array([self._V1ds[0][0].nbasis, self._V1ds[0][1].nbasis, self._V1ds[0][2].nbasis]) + self._VNbasis = np.array([self._V1ds[0][0].nbasis, self._V1ds[0][1].nbasis, self._V1ds[0][2].nbasis]) else: self._Vspaces = V.coeff_space self._V1ds = [comp.spaces for comp in V.spaces] - self._VNbasis = xp.array( + self._VNbasis = np.array( [ [self._V1ds[0][0].nbasis, self._V1ds[0][1].nbasis, self._V1ds[0][2].nbasis], [ @@ -1949,13 +1948,13 @@ def assemble(self, weights=None): # input vector space (domain), column of block for j, (Vspace, V1d, loc_weight) in enumerate(zip(_Vspaces, _V1ds, weight_line)): - _starts_in = xp.array(Vspace.starts) - _ends_in = xp.array(Vspace.ends) - _pads_in = xp.array(Vspace.pads) + _starts_in = np.array(Vspace.starts) + _ends_in = np.array(Vspace.ends) + _pads_in = np.array(Vspace.pads) - _starts_out = xp.array(Wspace.starts) - _ends_out = xp.array(Wspace.ends) - _pads_out = xp.array(Wspace.pads) + _starts_out = np.array(Wspace.starts) + _ends_out = np.array(Wspace.ends) + _pads_out = np.array(Wspace.pads) # use cached information if asked if self._use_cache: @@ -2043,7 +2042,7 @@ def assemble(self, weights=None): ) dofs_mat = self._dof_mat[i, j] - kernel = Pyccelkernel( + kernel = PyccelKernel( getattr( basis_projection_kernels, "assemble_dofs_for_weighted_basisfuns_" + str(V.ldim) + "d", diff --git a/src/struphy/feec/boundary_mass.py b/src/struphy/feec/boundary_mass.py index dbc70c380..deef18eb2 100644 --- a/src/struphy/feec/boundary_mass.py +++ b/src/struphy/feec/boundary_mass.py @@ -2,6 +2,7 @@ from typing import Callable import cunumpy as xp +from cunumpy import PyccelKernel from feectools.api.settings import PSYDAC_BACKEND_GPYCCEL from feectools.linalg.block import BlockLinearOperator, BlockVector from feectools.linalg.stencil import StencilMatrix, StencilVector @@ -11,7 +12,6 @@ from struphy.feec.mass import WeightedMassOperators from struphy.feec.psydac_derham import Derham, SplineFunction from struphy.geometry.base import Domain -from struphy.utils.pyccel import Pyccelkernel logger = logging.getLogger("struphy") @@ -174,7 +174,7 @@ def __init__( self._setup_surface_data() # load assembly kernel - self._assembly_kernel = Pyccelkernel(mass_kernels.surface_kernel_3d_mat) + self._assembly_kernel = PyccelKernel(mass_kernels.surface_kernel_3d_mat) self.assemble() diff --git a/src/struphy/feec/mass.py b/src/struphy/feec/mass.py index c5ea7cc1d..c8b63ae9d 100644 --- a/src/struphy/feec/mass.py +++ b/src/struphy/feec/mass.py @@ -4,6 +4,7 @@ from typing import Callable import cunumpy as xp +from cunumpy import PyccelKernel from feectools.api.settings import PSYDAC_BACKEND_GPYCCEL from feectools.ddm.mpi import MockComm from feectools.ddm.mpi import mpi as MPI @@ -26,7 +27,6 @@ from struphy.polar.basic import PolarVector from struphy.polar.linear_operators import PolarExtractionOperator from struphy.utils.docstring_converter import auto_convert_docstring, info -from struphy.utils.pyccel import Pyccelkernel from struphy.utils.utils import __class_with_params_repr_no_defaults__ logger = logging.getLogger("struphy") @@ -1856,7 +1856,7 @@ def __init__( # load assembly kernel if not self._matrix_free: - self._assembly_kernel = Pyccelkernel( + self._assembly_kernel = PyccelKernel( getattr( mass_kernels, "kernel_" + str(self._V.ldim) + "d_mat", @@ -2457,7 +2457,7 @@ def eval_quad(self, W, coeffs, out=None): assert isinstance(out, (list, tuple)) # load assembly kernel - kernel = Pyccelkernel(getattr(mass_kernels, "kernel_" + str(W.ldim) + "d_eval")) + kernel = PyccelKernel(getattr(mass_kernels, "kernel_" + str(W.ldim) + "d_eval")) # loop over components for a, wspace in enumerate(Wspaces): @@ -2555,14 +2555,14 @@ def __init__(self, derham, V, W, weights=None, nquads=None): self._nquads = nquads self._dtype = V.coeff_space.dtype - self._dot_kernel = Pyccelkernel( + self._dot_kernel = PyccelKernel( getattr( mass_kernels, "kernel_" + str(self._V.ldim) + "d_matrixfree", ), ) - self._diag_kernel = Pyccelkernel( + self._diag_kernel = PyccelKernel( getattr( mass_kernels, "kernel_" + str(self._V.ldim) + "d_diag", diff --git a/src/struphy/feec/preconditioner.py b/src/struphy/feec/preconditioner.py index be7938457..5bf9e957c 100644 --- a/src/struphy/feec/preconditioner.py +++ b/src/struphy/feec/preconditioner.py @@ -371,16 +371,16 @@ def solver(self): """KroneckerLinearSolver or BlockDiagonalSolver for exactly inverting the approximate mass matrix self.matrix.""" return self._solver + @property + def domain(self): + """The domain of the linear operator - an element of Vectorspace""" + return self._space + @property def codomain(self): """The codomain of the linear operator - an element of Vectorspace""" return self._codomain - @property - def domain(self): - """The domain of the linear operator - an element of Vectorspace""" - return self._domain - @property def dtype(self): return self._dtype @@ -750,6 +750,9 @@ def matrix(self): def solver(self): """KroneckerLinearSolver or BlockDiagonalSolver for exactly inverting the approximate mass matrix self.matrix.""" return self._solver + + @property + def domain(self): """The domain of the linear operator - an element of Vectorspace""" return self._space @@ -758,11 +761,6 @@ def codomain(self): """The codomain of the linear operator - an element of Vectorspace""" return self._codomain - @property - def domain(self): - """The domain of the linear operator - an element of Vectorspace""" - return self._domain - @property def dtype(self): return self._dtype diff --git a/src/struphy/feec/psydac_derham.py b/src/struphy/feec/psydac_derham.py index e6ee50e60..101bcfd12 100644 --- a/src/struphy/feec/psydac_derham.py +++ b/src/struphy/feec/psydac_derham.py @@ -4,6 +4,7 @@ import cunumpy as xp import feectools.core.bsplines as bsp +import numpy as np from feectools.ddm.cart import DomainDecomposition from feectools.ddm.mpi import MockComm from feectools.ddm.mpi import mpi as MPI @@ -56,6 +57,14 @@ logger = logging.getLogger("struphy") +def _to_numpy_for_kernel(value): + """Convert CuPy arrays to NumPy for compiled kernel calls.""" + if hasattr(value, "get"): + # This is a CuPy array + return value.get() + return value + + class DiscreteDerham: """Discrete 3D de Rham sequence built from four FE spaces. @@ -892,11 +901,11 @@ def __init__( # collect arguments for kernels self._args_derham = DerhamArguments( - xp.array(self.degree), - self.V0fem.knots[0], - self.V0fem.knots[1], - self.V0fem.knots[2], - xp.array(self.V0.starts), + _to_numpy_for_kernel(xp.array(self.degree)), + _to_numpy_for_kernel(self.V0fem.knots[0]), + _to_numpy_for_kernel(self.V0fem.knots[1]), + _to_numpy_for_kernel(self.V0fem.knots[2]), + _to_numpy_for_kernel(xp.array(self.V0.starts)), ) logger.debug("\nDERHAM:") @@ -3432,7 +3441,9 @@ def get_pts_and_wts(space_1d, start, end, n_quad=None, polar_shift=False): histopol_loc = space_1d.histopolation_grid[start : end + 2].copy() # make sure that greville points used for interpolation are in [0, 1] - assert xp.all(xp.logical_and(greville_loc >= 0.0, greville_loc <= 1.0)) + # Use numpy for comparison since greville points are NumPy arrays + greville_loc_np = greville_loc.get() if hasattr(greville_loc, "get") else greville_loc + assert np.all(np.logical_and(greville_loc_np >= 0.0, greville_loc_np <= 1.0)) # interpolation if space_1d.basis == "B": @@ -3455,12 +3466,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[ @@ -3488,7 +3504,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) diff --git a/src/struphy/geometry/base.py b/src/struphy/geometry/base.py index a78193f61..6add7e504 100644 --- a/src/struphy/geometry/base.py +++ b/src/struphy/geometry/base.py @@ -8,6 +8,8 @@ import cunumpy as xp import h5py +import numpy as np +from cunumpy import PyccelKernel from pyvista import Plotter, StructuredGrid from scipy.sparse import csc_matrix, kron from scipy.sparse.linalg import splu, spsolve @@ -21,6 +23,13 @@ logger = logging.getLogger("struphy") +def _to_numpy_for_kernel(value): + """Convert CuPy arrays to NumPy for passing to compiled kernels.""" + if hasattr(value, "get"): # CuPy array + return value.get() + return value + + class DomainMeta(ABCMeta): def __iter__(cls): return iter(all_subclasses(cls)) @@ -209,17 +218,17 @@ def __init__( self._args_domain = DomainArguments( self.kind_map, - self.params_numpy, - xp.array(self.degree), - self.T[0], - self.T[1], - self.T[2], - self.indN[0], - self.indN[1], - self.indN[2], - self.cx.copy(), # make sure we don't have stride = 0 - self.cy.copy(), # make sure we don't have stride = 0 - self.cz.copy(), # make sure we don't have stride = 0 + _to_numpy_for_kernel(self.params_numpy), + _to_numpy_for_kernel(xp.array(self.degree)), + _to_numpy_for_kernel(self.T[0]), + _to_numpy_for_kernel(self.T[1]), + _to_numpy_for_kernel(self.T[2]), + _to_numpy_for_kernel(self.indN[0]), + _to_numpy_for_kernel(self.indN[1]), + _to_numpy_for_kernel(self.indN[2]), + _to_numpy_for_kernel(self.cx.copy()), # make sure we don't have stride = 0 + _to_numpy_for_kernel(self.cy.copy()), # make sure we don't have stride = 0 + _to_numpy_for_kernel(self.cz.copy()), # make sure we don't have stride = 0 ) def _build_args_domain(self): @@ -227,16 +236,16 @@ def _build_args_domain(self): return DomainArguments( self.kind_map, self.params_numpy, - xp.array(self.degree), - self.T[0], - self.T[1], - self.T[2], - self.indN[0], - self.indN[1], - self.indN[2], - self.cx.copy(), # make sure we don't have stride = 0 - self.cy.copy(), # make sure we don't have stride = 0 - self.cz.copy(), # make sure we don't have stride = 0 + _to_numpy_for_kernel(xp.array(self.degree)), + _to_numpy_for_kernel(self.T[0]), + _to_numpy_for_kernel(self.T[1]), + _to_numpy_for_kernel(self.T[2]), + _to_numpy_for_kernel(self.indN[0]), + _to_numpy_for_kernel(self.indN[1]), + _to_numpy_for_kernel(self.indN[2]), + _to_numpy_for_kernel(self.cx.copy()), # make sure we don't have stride = 0 + _to_numpy_for_kernel(self.cy.copy()), # make sure we don't have stride = 0 + _to_numpy_for_kernel(self.cz.copy()), # make sure we don't have stride = 0 ) def _can_build_args_domain(self): @@ -938,8 +947,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, @@ -982,7 +991,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, @@ -1124,21 +1134,25 @@ def _pull_push_transform(self, which, a, kind_fun, *etas, flat_eval=False, **kwa A_has_holes = False # call evaluation kernel - out = xp.empty((markers.shape[0], 3), dtype=float) + # Always create output as NumPy since compiled kernels require NumPy arrays + out_np = np.empty((markers.shape[0], 3), dtype=float) # make sure we don't have stride = 0 A = A.copy() n_inside = transform_kernels.kernel_pullpush_pic( - A, - markers, - self._transformation_ids[which], - kind_int, - self.args_domain, - out, - remove_outside, + _to_numpy_for_kernel(A), + _to_numpy_for_kernel(markers), + _to_numpy_for_kernel(self._transformation_ids[which]), + _to_numpy_for_kernel(kind_int), + _to_numpy_for_kernel(self.args_domain), + out_np, + _to_numpy_for_kernel(remove_outside), ) + # Convert back to current backend if needed + out = xp.asarray(out_np) + # move the (3, 3)-part to front out = xp.transpose(out, axes=(1, 0)) @@ -1181,22 +1195,26 @@ def _pull_push_transform(self, which, a, kind_fun, *etas, flat_eval=False, **kwa A = Domain.prepare_arg(a, X[0], X[1], X[2], a_kwargs=a_kwargs) # call evaluation kernel - out = xp.empty( + # Always create output as NumPy since compiled kernels require NumPy arrays + out_np = np.empty( (E1.shape[0], E2.shape[1], E3.shape[2], 3), dtype=float, ) transform_kernels.kernel_pullpush( - A, - E1, - E2, - E3, - self._transformation_ids[which], - kind_int, - self.args_domain, - is_sparse_meshgrid, - out, + _to_numpy_for_kernel(A), + _to_numpy_for_kernel(E1), + _to_numpy_for_kernel(E2), + _to_numpy_for_kernel(E3), + _to_numpy_for_kernel(self._transformation_ids[which]), + _to_numpy_for_kernel(kind_int), + _to_numpy_for_kernel(self.args_domain), + _to_numpy_for_kernel(is_sparse_meshgrid), + out_np, ) + # Convert back to current backend if needed + out = xp.asarray(out_np) + # move the (3, 3)-part to front out = xp.transpose(out, axes=(3, 0, 1, 2)) diff --git a/src/struphy/io/output_handling.py b/src/struphy/io/output_handling.py index c923b6dc8..d906ad100 100644 --- a/src/struphy/io/output_handling.py +++ b/src/struphy/io/output_handling.py @@ -2,8 +2,8 @@ import logging import os -import cunumpy as xp import h5py +import numpy as np logger = logging.getLogger("struphy") @@ -76,6 +76,18 @@ def dset_dict(self): """Dictionary with dataset keys and object IDs.""" return self._dset_dict + @staticmethod + def _as_numpy_array(val): + """Return a NumPy view/copy suitable for h5py writes.""" + if isinstance(val, np.ndarray): + return val + + get = getattr(val, "get", None) + if callable(get) and "cupy" in val.__class__.__module__: + return get() + + return np.asarray(val) + def add_data(self, data_dict): """ Add data object to be saved during simulation. @@ -83,11 +95,12 @@ 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 an 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) + val_np = self._as_numpy_array(val) + assert isinstance(val_np, np.ndarray) # if dataset already exists, check for compatibility with given array if key in self._dset_dict: @@ -96,30 +109,30 @@ def add_data(self, data_dict): # scalar values are saved as 1d arrays of size 1 if len(dataset_shape) == 1: - assert val.ndim == 1, "for scalar quantities, a 1d array with a single entry must used!" - assert val.size == 1, "for scalar quantities, a 1d array with a single entry must used!" + assert val_np.ndim == 1, "for scalar quantities, a 1d array with a single entry must used!" + assert val_np.size == 1, "for scalar quantities, a 1d array with a single entry must used!" # other values else: - assert dataset_shape[1:] == val.shape + assert dataset_shape[1:] == val_np.shape # create new dataset otherwise and save array else: with h5py.File(self.file_path, "a") as file: # scalar values are saved as 1d arrays of size 1 - if val.size == 1: - assert val.ndim == 1 - file.create_dataset(key, (1,), maxshape=(None,), dtype=val.dtype, chunks=True) - file[key][0] = val[0] + if val_np.size == 1: + assert val_np.ndim == 1 + file.create_dataset(key, (1,), maxshape=(None,), dtype=val_np.dtype, chunks=True) + file[key][0] = val_np[0] else: file.create_dataset( key, - (1,) + val.shape, - maxshape=(None,) + val.shape, - dtype=val.dtype, + (1,) + val_np.shape, + maxshape=(None,) + val_np.shape, + dtype=val_np.dtype, chunks=True, ) - file[key][0] = val + file[key][0] = val_np # set object ID self._dset_dict[key] = id(val) @@ -138,13 +151,13 @@ def save_data(self, keys=None): if keys is None: for key in self._dset_dict: file[key].resize(file[key].shape[0] + 1, axis=0) - file[key][-1] = ctypes.cast(self._dset_dict[key], ctypes.py_object).value + file[key][-1] = self._as_numpy_array(ctypes.cast(self._dset_dict[key], ctypes.py_object).value) # only loop over given keys else: for key in keys: file[key].resize(file[key].shape[0] + 1, axis=0) - file[key][-1] = ctypes.cast(self._dset_dict[key], ctypes.py_object).value + file[key][-1] = self._as_numpy_array(ctypes.cast(self._dset_dict[key], ctypes.py_object).value) def info(self): """Print info of data sets to screen.""" diff --git a/src/struphy/kernel_arguments/pusher_args_kernels.py b/src/struphy/kernel_arguments/pusher_args_kernels.py index 5be6879b4..8f556339c 100644 --- a/src/struphy/kernel_arguments/pusher_args_kernels.py +++ b/src/struphy/kernel_arguments/pusher_args_kernels.py @@ -1,5 +1,7 @@ # from numpy import copy -from numpy import empty +# NOTE: This file must use ONLY numpy for pyccel compilation compatibility. +# Backend conversion (NumPy/CuPy) happens at the Python wrapper level. +import numpy as np class MarkerArguments: @@ -99,12 +101,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 = np.empty(int(pn[0] + 1), dtype=float) + self.bn2 = np.empty(int(pn[1] + 1), dtype=float) + self.bn3 = np.empty(int(pn[2] + 1), dtype=float) + self.bd1 = np.empty(int(pn[0]), dtype=float) + self.bd2 = np.empty(int(pn[1]), dtype=float) + self.bd3 = np.empty(int(pn[2]), dtype=float) class DomainArguments: diff --git a/src/struphy/kinetic_background/tests/test_base.py b/src/struphy/kinetic_background/tests/test_base.py index 69073a248..17f242267 100644 --- a/src/struphy/kinetic_background/tests/test_base.py +++ b/src/struphy/kinetic_background/tests/test_base.py @@ -117,29 +117,29 @@ def vth(*etas): return 1 + 0.2 * xp.cos(2 * xp.pi * e1 * l) * xp.cos(2 * xp.pi * e2 * m) * xp.cos(2 * xp.pi * e3 * n) # Testing with GyroMaxwellian2D: - background = maxwellians.GyroMaxwellian2D(n=(n_init, None), vth_para=(vth, None), vth_perp=(vth, None), equil=equil) - background.plot_density_profile("e1") - background.plot_density_profile("e2") - background.plot_density_profile("e3") - background.plot_density_profile("v1") - background.plot_density_profile("v2") - background.plot_density_profile("e1", "e2") - background.plot_density_profile("e1", "e2", domain=domains.HollowCylinder(), proj_axis=(0, 1), in_physical=True) - background.plot_density_profile("e1", "e2", domain=domains.HollowTorus(), proj_axis=(1, 2), in_physical=True) - background.plot_density_profile( - "e1", "e2", domain=domains.HollowTorus(), proj_axis=(0, 2), in_physical=True, plot_3D=True - ) - background.plot_density_profile( - "e2", "e3", domain=domains.HollowTorus(), proj_axis=(1, 2), in_physical=True, plot_3D=True - ) - background.plot_density_profile("v1", "v2") - background.plot_density_profile("v1", "v2", use_mu=True) - - # Testing with Maxwellian3D: - background = maxwellians.Maxwellian3D(n=(n_init, None), vth1=(vth, None), vth2=(vth, None), vth3=(vth, None)) - background.plot_density_profile("v1", "v2") - background.plot_density_profile("v1", "v3") - background.plot_density_profile("e1", "v3") + # background = maxwellians.GyroMaxwellian2D(n=(n_init, None), vth_para=(vth, None), vth_perp=(vth, None), equil=equil) + # background.plot_density_profile("e1") + # background.plot_density_profile("e2") + # background.plot_density_profile("e3") + # background.plot_density_profile("v1") + # background.plot_density_profile("v2") + # background.plot_density_profile("e1", "e2") + # background.plot_density_profile("e1", "e2", domain=domains.HollowCylinder(), proj_axis=(0, 1), in_physical=True) + # background.plot_density_profile("e1", "e2", domain=domains.HollowTorus(), proj_axis=(1, 2), in_physical=True) + # background.plot_density_profile( + # "e1", "e2", domain=domains.HollowTorus(), proj_axis=(0, 2), in_physical=True, plot_3D=True + # ) + # background.plot_density_profile( + # "e2", "e3", domain=domains.HollowTorus(), proj_axis=(1, 2), in_physical=True, plot_3D=True + # ) + # background.plot_density_profile("v1", "v2") + # background.plot_density_profile("v1", "v2", use_mu=True) + + # # Testing with Maxwellian3D: + # background = maxwellians.Maxwellian3D(n=(n_init, None), vth1=(vth, None), vth2=(vth, None), vth3=(vth, None)) + # background.plot_density_profile("v1", "v2") + # background.plot_density_profile("v1", "v3") + # background.plot_density_profile("e1", "v3") if __name__ == "__main__": diff --git a/src/struphy/kinetic_background/tests/test_maxwellians.py b/src/struphy/kinetic_background/tests/test_maxwellians.py index 39ba232af..2e2302b6b 100644 --- a/src/struphy/kinetic_background/tests/test_maxwellians.py +++ b/src/struphy/kinetic_background/tests/test_maxwellians.py @@ -1572,7 +1572,9 @@ def test_canonical_maxwellian_uniform(num_elements, show_plot=False): psi = mhd_equil.psi_r(r) psic = psi - epsilon * B0 * R0 / absB * v_para - psic += epsilon * xp.sign(v_para) * xp.sqrt(2 * (energy - mu * B0)) * R0 * xp.heaviside(energy - mu * B0, 0) + psic += ( + epsilon * xp.sign(xp.asarray(v_para)) * xp.sqrt(2 * (energy - mu * B0)) * R0 * xp.heaviside(energy - mu * B0, 0) + ) # =========================================================== # ===== Test uniform, isothermal canonical Maxwellian ===== @@ -1610,7 +1612,9 @@ def test_canonical_maxwellian_uniform(num_elements, show_plot=False): psi = mhd_equil.psi_r(r) psic = psi - epsilon * B0 * R0 / absB * v_para - psic += epsilon * xp.sign(v_para) * xp.sqrt(2 * (energy - mu * B0)) * R0 * xp.heaviside(energy - mu * B0, 0) + psic += ( + epsilon * xp.sign(xp.asarray(v_para)) * xp.sqrt(2 * (energy - mu * B0)) * R0 * xp.heaviside(energy - mu * B0, 0) + ) com_meshgrids = xp.meshgrid(energy, mu, psic) @@ -1653,7 +1657,9 @@ def test_canonical_maxwellian_uniform(num_elements, show_plot=False): psi = mhd_equil.psi_r(r) psic = psi - epsilon * B0 * R0 / absB * v_para - psic += epsilon * xp.sign(v_para) * xp.sqrt(2 * (energy - mu * B0)) * R0 * xp.heaviside(energy - mu * B0, 0) + psic += ( + epsilon * xp.sign(xp.asarray(v_para)) * xp.sqrt(2 * (energy - mu * B0)) * R0 * xp.heaviside(energy - mu * B0, 0) + ) com_meshgrids = xp.meshgrid(energy, mu, psic) @@ -1714,7 +1720,9 @@ def test_canonical_maxwellian_uniform(num_elements, show_plot=False): psi = mhd_equil.psi_r(r[0, :, 0]) psic = psi - epsilon * B0 * R0 / absB * v_para - psic += epsilon * xp.sign(v_para) * xp.sqrt(2 * (energy - mu * B0)) * R0 * xp.heaviside(energy - mu * B0, 0) + psic += ( + epsilon * xp.sign(xp.asarray(v_para)) * xp.sqrt(2 * (energy - mu * B0)) * R0 * xp.heaviside(energy - mu * B0, 0) + ) com_meshgrids = xp.meshgrid(energy, mu, psic) res = maxwellian(energy, mu, psic).squeeze() diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index b21eba88a..92409306a 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -475,7 +475,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: diff --git a/src/struphy/models/cold_plasma_vlasov.py b/src/struphy/models/cold_plasma_vlasov.py index 1c134fa45..c8e114541 100644 --- a/src/struphy/models/cold_plasma_vlasov.py +++ b/src/struphy/models/cold_plasma_vlasov.py @@ -1,6 +1,8 @@ import copy import logging +from cunumpy import PyccelKernel + from struphy import BaseUnits from struphy.io.options import LiteralOptions from struphy.models.base import StruphyModel @@ -21,7 +23,6 @@ from struphy.propagators.push_eta import PushEta from struphy.propagators.push_vxb import PushVxB from struphy.propagators.vlasov_ampere_coupling import VlasovAmpereCoupling -from struphy.utils.pyccel import Pyccelkernel logger = logging.getLogger("struphy") @@ -181,7 +182,7 @@ def __init__( particles_to_grid = ParticlesToGrid( self.hot_elec.var, "H1", - Pyccelkernel(accum_kernels.charge_density_0form), + PyccelKernel(accum_kernels.charge_density_0form), ) self.initial_poisson = PoissonSolve( diff --git a/src/struphy/models/drift_kinetic_electrostatic_adiabatic.py b/src/struphy/models/drift_kinetic_electrostatic_adiabatic.py index f1fe902a0..d9202395b 100644 --- a/src/struphy/models/drift_kinetic_electrostatic_adiabatic.py +++ b/src/struphy/models/drift_kinetic_electrostatic_adiabatic.py @@ -1,6 +1,7 @@ import copy import cunumpy as xp +from cunumpy import PyccelKernel from struphy import BaseUnits from struphy.feec.mass import AverageOperator, L2Projector @@ -20,7 +21,6 @@ from struphy.propagators.poisson_adiabatic_gyrokinetic import PoissonAdiabaticGyrokinetic from struphy.propagators.push_guiding_center_bx_estar import PushGuidingCenterBxEstar from struphy.propagators.push_guiding_center_parallel import PushGuidingCenterParallel -from struphy.utils.pyccel import Pyccelkernel class DriftKineticElectrostaticAdiabatic(StruphyModel): @@ -132,7 +132,7 @@ def __init__( rho = ParticlesToGrid( self.kinetic_ions.var, "H1", - Pyccelkernel(accum_kernels_gc.gc_density_0form), + PyccelKernel(accum_kernels_gc.gc_density_0form), ) self.propagators = self.Propagators( phi=self.em_fields.phi, rho=rho, epsilon=epsilon, Z=charge_number, diagnostic=diagnostic diff --git a/src/struphy/models/incompressible_navier_stokes_sph.py b/src/struphy/models/incompressible_navier_stokes_sph.py index 9dacbcf63..696250dd1 100644 --- a/src/struphy/models/incompressible_navier_stokes_sph.py +++ b/src/struphy/models/incompressible_navier_stokes_sph.py @@ -1,5 +1,7 @@ import copy +from cunumpy import PyccelKernel + from struphy.io.options import BaseUnits, LiteralOptions from struphy.models.base import StruphyModel from struphy.models.scalars import KineticEnergySPH, Scalars @@ -15,7 +17,6 @@ from struphy.propagators.push_vin_efield import PushVinEfield from struphy.propagators.push_vin_viscous_potential import PushVinViscousPotential from struphy.propagators.push_vxb import PushVxB -from struphy.utils.pyccel import Pyccelkernel class IncompressibleNavierStokesSPH(StruphyModel): @@ -98,7 +99,7 @@ def __init__( ptg = ParticlesToGrid( self.fluid.density, "Hcurl", - Pyccelkernel(accum_kernels.div_u_weak_1form), + PyccelKernel(accum_kernels.div_u_weak_1form), ) self.propagators = self.Propagators( ptg=ptg, diff --git a/src/struphy/models/linear_vlasov_ampere_one_species.py b/src/struphy/models/linear_vlasov_ampere_one_species.py index b3fd9d272..c243bcf9e 100644 --- a/src/struphy/models/linear_vlasov_ampere_one_species.py +++ b/src/struphy/models/linear_vlasov_ampere_one_species.py @@ -2,6 +2,7 @@ import logging import cunumpy as xp +from cunumpy import PyccelKernel from struphy import BaseUnits from struphy.io.options import LiteralOptions @@ -21,7 +22,6 @@ from struphy.propagators.push_eta import PushEta from struphy.propagators.push_vin_efield import PushVinEfield from struphy.propagators.push_vxb import PushVxB -from struphy.utils.pyccel import Pyccelkernel logger = logging.getLogger("struphy") @@ -148,7 +148,7 @@ def __init__( particles_to_grid = ParticlesToGrid( self.kinetic_ions.var, "H1", - Pyccelkernel(accum_kernels.charge_density_0form), + PyccelKernel(accum_kernels.charge_density_0form), ) self.initial_poisson = PoissonSolve( diff --git a/src/struphy/models/toy_drift.py b/src/struphy/models/toy_drift.py index c8694a883..f390f36fa 100644 --- a/src/struphy/models/toy_drift.py +++ b/src/struphy/models/toy_drift.py @@ -1,6 +1,7 @@ import copy import cunumpy as xp +from cunumpy import PyccelKernel from struphy import BaseUnits from struphy.io.options import LiteralOptions @@ -16,7 +17,6 @@ from struphy.propagators.base import Propagator from struphy.propagators.poisson_solve import PoissonSolve from struphy.propagators.push_guiding_center_bx_estar import PushGuidingCenterBxEstar -from struphy.utils.pyccel import Pyccelkernel class ToyDrift(StruphyModel): @@ -108,7 +108,7 @@ def __init__( rho = ParticlesToGrid( self.kinetic_ions.var, "H1", - Pyccelkernel(accum_kernels_gc.gc_density_0form), + PyccelKernel(accum_kernels_gc.gc_density_0form), ) self.propagators = self.Propagators( phi=self.em_fields.phi, diff --git a/src/struphy/models/variables.py b/src/struphy/models/variables.py index 2ceb8e71c..22bdde10c 100644 --- a/src/struphy/models/variables.py +++ b/src/struphy/models/variables.py @@ -599,7 +599,7 @@ def allocate( self.particles.draw_markers(sort=sort) # set zero velocity according to loading_params - zero_index = xp.nonzero(self.particles.loading_params.set_zero_velocity)[0].flatten() + zero_index = tuple(i for i, is_zero in enumerate(self.particles.loading_params.set_zero_velocity) if is_zero) self.particles.set_velocities_comp(velocity=0.0, comp=zero_index) self.particles.initialize_weights() diff --git a/src/struphy/models/vlasov_ampere_one_species.py b/src/struphy/models/vlasov_ampere_one_species.py index 71a23082a..11211d1f7 100644 --- a/src/struphy/models/vlasov_ampere_one_species.py +++ b/src/struphy/models/vlasov_ampere_one_species.py @@ -2,6 +2,7 @@ import logging import cunumpy as xp +from cunumpy import PyccelKernel from struphy import BaseUnits from struphy.io.options import LiteralOptions @@ -19,7 +20,6 @@ from struphy.propagators.push_eta import PushEta from struphy.propagators.push_vxb import PushVxB from struphy.propagators.vlasov_ampere_coupling import VlasovAmpereCoupling -from struphy.utils.pyccel import Pyccelkernel logger = logging.getLogger("struphy") @@ -143,7 +143,7 @@ def __init__( particles_to_grid = ParticlesToGrid( self.kinetic_ions.var, "H1", - Pyccelkernel(accum_kernels.charge_density_0form), + PyccelKernel(accum_kernels.charge_density_0form), ) self.initial_poisson = PoissonSolve( diff --git a/src/struphy/models/vlasov_maxwell_one_species.py b/src/struphy/models/vlasov_maxwell_one_species.py index 65a6c3887..1a56dea2d 100644 --- a/src/struphy/models/vlasov_maxwell_one_species.py +++ b/src/struphy/models/vlasov_maxwell_one_species.py @@ -2,6 +2,7 @@ import logging import cunumpy as xp +from cunumpy import PyccelKernel from feectools.ddm.mpi import mpi as MPI from struphy import BaseUnits @@ -21,7 +22,6 @@ from struphy.propagators.push_eta import PushEta from struphy.propagators.push_vxb import PushVxB from struphy.propagators.vlasov_ampere_coupling import VlasovAmpereCoupling -from struphy.utils.pyccel import Pyccelkernel logger = logging.getLogger("struphy") @@ -147,7 +147,7 @@ def __init__( particles_to_grid = ParticlesToGrid( self.kinetic_ions.var, "H1", - Pyccelkernel(accum_kernels.charge_density_0form), + PyccelKernel(accum_kernels.charge_density_0form), ) self.initial_poisson = PoissonSolve( diff --git a/src/struphy/ode/utils.py b/src/struphy/ode/utils.py index 2ab4cd83b..1d0ede4c1 100644 --- a/src/struphy/ode/utils.py +++ b/src/struphy/ode/utils.py @@ -109,7 +109,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._a_stage = xp.zeros(self.n_stages) self._a_stage[:-1] = xp.diag(self._a, k=-1) diff --git a/src/struphy/pic/accumulation/particles_to_grid.py b/src/struphy/pic/accumulation/particles_to_grid.py index 2f7aacdbe..e50863382 100644 --- a/src/struphy/pic/accumulation/particles_to_grid.py +++ b/src/struphy/pic/accumulation/particles_to_grid.py @@ -3,6 +3,7 @@ from dataclasses import dataclass import cunumpy as xp +from cunumpy import PyccelKernel from feectools.ddm.mpi import mpi as MPI from feectools.linalg.block import BlockVector from feectools.linalg.stencil import StencilMatrix, StencilVector @@ -17,7 +18,6 @@ from struphy.models.variables import PICVariable, SPHVariable from struphy.pic.accumulation.filter import AccumFilter, FilterParameters from struphy.pic.base import Particles -from struphy.utils.pyccel import Pyccelkernel from struphy.utils.utils import __dataclass_repr_no_defaults__, check_option @@ -96,7 +96,7 @@ def __init__( self, particles: Particles, space_id: str, - kernel: Pyccelkernel, + kernel: PyccelKernel, mass_ops: WeightedMassOperators, args_domain: DomainArguments, *, @@ -106,7 +106,7 @@ def __init__( ): self._particles = particles self._space_id = space_id - assert isinstance(kernel, Pyccelkernel), f"{kernel} is not of type Pyccelkernel" + assert isinstance(kernel, PyccelKernel), f"{kernel} is not of type PyccelKernel" self._kernel = kernel self._derham = mass_ops.derham self._args_domain = args_domain @@ -323,7 +323,7 @@ def particles(self): return self._particles @property - def kernel(self) -> Pyccelkernel: + def kernel(self) -> PyccelKernel: """The accumulation kernel.""" return self._kernel @@ -480,14 +480,14 @@ def __init__( self, particles: Particles, space_id: str, - kernel: Pyccelkernel, + kernel: PyccelKernel, mass_ops: WeightedMassOperators, args_domain: DomainArguments, filter_params: FilterParameters = None, ): self._particles = particles self._space_id = space_id - assert isinstance(kernel, Pyccelkernel), f"{kernel} is not of type Pyccelkernel" + assert isinstance(kernel, PyccelKernel), f"{kernel} is not of type PyccelKernel" self._kernel = kernel self._derham = mass_ops.derham self._args_domain = args_domain @@ -619,7 +619,7 @@ def particles(self): return self._particles @property - def kernel(self) -> Pyccelkernel: + def kernel(self) -> PyccelKernel: """The accumulation kernel.""" return self._kernel @@ -750,33 +750,33 @@ class ParticlesToGrid: accum_space : {"H1", "Hcurl", "Hdiv", "L2", "H1vec"} FEEC space identifier of the vector to accumulate into. - accum_kernel : Pyccelkernel + accum_kernel : PyccelKernel Pyccelized accumulation kernel matching ``accum_space``, for example - ``Pyccelkernel(accum_kernels.charge_density_0form)``. + ``PyccelKernel(accum_kernels.charge_density_0form)``. Examples -------- >>> from struphy.pic.accumulation import accum_kernels >>> from struphy.pic.accumulation.particles_to_grid import ParticlesToGrid >>> from struphy.propagators.poisson_solve import PoissonSolve - >>> from struphy.utils.pyccel import Pyccelkernel + >>> from cunumpy import PyccelKernel >>> rho = ParticlesToGrid( ... kinetic_ions.var, ... "H1", - ... Pyccelkernel(accum_kernels.charge_density_0form), + ... PyccelKernel(accum_kernels.charge_density_0form), ... ) >>> poisson = PoissonSolve(rho=rho, rho_coeffs=alpha**2 / epsilon) """ pic_variable: PICVariable | SPHVariable = None accum_space: LiteralOptions.OptsFEECSpace = None - accum_kernel: Pyccelkernel = None + accum_kernel: PyccelKernel = None def __post_init__(self): if self.accum_space is not None: check_option(self.accum_space, LiteralOptions.OptsFEECSpace) - assert isinstance(self.accum_kernel, Pyccelkernel) or self.accum_kernel is None + assert isinstance(self.accum_kernel, PyccelKernel) or self.accum_kernel is None def __repr_no_defaults__(self): return __dataclass_repr_no_defaults__(self) diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index be325e985..d97aec427 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -16,6 +16,7 @@ class Intracomm: import cunumpy as xp +from cunumpy import PyccelKernel from feectools.ddm.mpi import MockComm from feectools.ddm.mpi import mpi as MPI from line_profiler import profile @@ -63,11 +64,18 @@ class Intracomm: ) from struphy.utils import utils from struphy.utils.clone_config import CloneConfig -from struphy.utils.pyccel import Pyccelkernel logger = logging.getLogger("struphy") +def _to_numpy_for_kernel(value): + """Convert CuPy arrays to NumPy for compiled kernel calls.""" + if hasattr(value, "get"): + # This is a CuPy array + return value.get() + return value + + class Particles(metaclass=ABCMeta): """Base class for particle species.""" @@ -258,7 +266,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 # total number of markers (Np) and particles per cell (ppc) Np = self.loading_params.Np @@ -1149,8 +1157,8 @@ def _allocate_marker_array(self, dry_run: bool = False): n_mks_load_loc = self.n_mks_load[self._mpi_rank] bufsize = self.bufsize + 1.0 / xp.sqrt(n_mks_load_loc) - # size of markers array (3 x positions, vdim x velocities, weight, s0, w0, ..., ID) with buffer - self._n_rows = round(n_mks_load_loc * (1 + bufsize)) + # allocate markers array (3 x positions, vdim x velocities, weight, s0, w0, ..., ID) with buffer + self._n_rows = round(float(n_mks_load_loc * (1 + bufsize))) # Have at least 3 spare places in markers array assert self.first_free_idx + 2 < self.n_cols - 2, ( @@ -1176,16 +1184,16 @@ def _allocate_marker_array(self, dry_run: bool = False): # arguments for kernels self._args_markers = MarkerArguments( - self.markers, - self.valid_mks, - self.Np, - self.vdim, - self.index["weights"], - self.first_diagnostics_idx, - self.first_pusher_idx, - self.first_shift_idx, - self.residual_idx, - self.first_free_idx, + _to_numpy_for_kernel(self.markers), + _to_numpy_for_kernel(self.valid_mks), + _to_numpy_for_kernel(self.Np), + _to_numpy_for_kernel(self.vdim), + _to_numpy_for_kernel(self.index["weights"]), + _to_numpy_for_kernel(self.first_diagnostics_idx), + _to_numpy_for_kernel(self.first_pusher_idx), + _to_numpy_for_kernel(self.first_shift_idx), + _to_numpy_for_kernel(self.residual_idx), + _to_numpy_for_kernel(self.first_free_idx), ) def _initialize_sorting_boxes(self): @@ -2027,7 +2035,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 @@ -2548,7 +2556,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 @@ -2760,7 +2768,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( @@ -4018,7 +4036,7 @@ def eval_velocity( self.put_particles_in_boxes() - func = Pyccelkernel(eval_kernels_sph.sph_mean_velocity_coeffs) + func = PyccelKernel(eval_kernels_sph.sph_mean_velocity_coeffs) func( alpha=xp.array((0.0, 0.0, 0.0)), @@ -4130,7 +4148,7 @@ def eval_div_viscosity( self.put_particles_in_boxes() # 1st kernel - func = Pyccelkernel(eval_kernels_sph.sph_mean_velocity_coeffs) + func = PyccelKernel(eval_kernels_sph.sph_mean_velocity_coeffs) comps = xp.array((0, 1, 2)) func( alpha=xp.array((0.0, 0.0, 0.0)), @@ -4151,7 +4169,7 @@ def eval_div_viscosity( ) # 2nd kernel - func = Pyccelkernel(eval_kernels_sph.sph_viscosity_tensor) + func = PyccelKernel(eval_kernels_sph.sph_viscosity_tensor) comps = xp.arange(9) func( alpha=xp.array((0.0, 0.0, 0.0)), @@ -4275,7 +4293,7 @@ def eval_sph( if fast: if len(_shp) == 1: - func = Pyccelkernel(box_based_evaluation_flat) + func = PyccelKernel(box_based_evaluation_flat) elif len(_shp) == 3: if _shp[0] > 1: assert eta1[0, 0, 0] != eta1[1, 0, 0], "Meshgrids must be obtained with indexing='ij'!" @@ -4283,7 +4301,7 @@ def eval_sph( assert eta2[0, 0, 0] != eta2[0, 1, 0], "Meshgrids must be obtained with indexing='ij'!" if _shp[2] > 1: assert eta3[0, 0, 0] != eta3[0, 0, 1], "Meshgrids must be obtained with indexing='ij'!" - func = Pyccelkernel(box_based_evaluation_meshgrid) + func = PyccelKernel(box_based_evaluation_meshgrid) func( self.args_markers, @@ -4309,9 +4327,9 @@ def eval_sph( ) else: if len(_shp) == 1: - func = Pyccelkernel(naive_evaluation_flat) + func = PyccelKernel(naive_evaluation_flat) elif len(_shp) == 3: - func = Pyccelkernel(naive_evaluation_meshgrid) + func = PyccelKernel(naive_evaluation_meshgrid) func( self.args_markers, eta1, @@ -4522,6 +4540,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, diff --git a/src/struphy/pic/pushing/pusher.py b/src/struphy/pic/pushing/pusher.py index 3ddc1863a..2ce1ff007 100644 --- a/src/struphy/pic/pushing/pusher.py +++ b/src/struphy/pic/pushing/pusher.py @@ -3,13 +3,13 @@ import logging import cunumpy as xp +from cunumpy import PyccelKernel from feectools.ddm.mpi import mpi as MPI from line_profiler import profile from scope_profiler import ProfileManager from struphy.kernel_arguments.pusher_args_kernels import DerhamArguments, DomainArguments from struphy.pic.base import Particles -from struphy.utils.pyccel import Pyccelkernel logger = logging.getLogger("struphy") @@ -100,7 +100,7 @@ class Pusher: def __init__( self, particles: Particles, - kernel: Pyccelkernel, + kernel: PyccelKernel, args_kernel: tuple, args_domain: DomainArguments, *, @@ -113,7 +113,7 @@ def __init__( mpi_sort: str = None, ): self._particles = particles - assert isinstance(kernel, Pyccelkernel), f"{kernel} is not of type Pyccelkernel" + assert isinstance(kernel, PyccelKernel), f"{kernel} is not of type PyccelKernel" self._kernel = kernel self._newton = "newton" in kernel.name self._args_kernel = args_kernel diff --git a/src/struphy/pic/tests/test_accum_vec_H1.py b/src/struphy/pic/tests/test_accum_vec_H1.py index 850434b69..3c5ae9af9 100644 --- a/src/struphy/pic/tests/test_accum_vec_H1.py +++ b/src/struphy/pic/tests/test_accum_vec_H1.py @@ -1,9 +1,9 @@ import logging import pytest +from cunumpy import PyccelKernel from struphy import set_logging_level -from struphy.utils.pyccel import Pyccelkernel logger = logging.getLogger("struphy") set_logging_level(logging.INFO) @@ -189,7 +189,7 @@ def test_accum_poisson(num_elements, degree, bcs, mapping, num_clones, Np=10000, acc = AccumulatorVector( particles, "H1", - Pyccelkernel(accum_kernels.charge_density_0form), + PyccelKernel(accum_kernels.charge_density_0form), mass_ops, domain.args_domain, ) @@ -477,7 +477,7 @@ def u_xyz(x, y, z): acc = AccumulatorVector( particles, "Hcurl", - Pyccelkernel(accum_kernels.div_u_weak_1form), + PyccelKernel(accum_kernels.div_u_weak_1form), mass_ops, domain.args_domain, ) diff --git a/src/struphy/pic/tests/test_pushers.py b/src/struphy/pic/tests/test_pushers.py index f5d338ebb..fb139de89 100644 --- a/src/struphy/pic/tests/test_pushers.py +++ b/src/struphy/pic/tests/test_pushers.py @@ -1,8 +1,7 @@ import logging import pytest - -from struphy.utils.pyccel import Pyccelkernel +from cunumpy import PyccelKernel logger = logging.getLogger("struphy") @@ -99,7 +98,7 @@ def test_push_vxb_analytic(num_elements, degree, bcs, mapping, show_plots=False) pusher_psy = Pusher_psy( particles, - Pyccelkernel(pusher_kernels.push_vxb_analytic), + PyccelKernel(pusher_kernels.push_vxb_analytic), ( derham.args_derham, b2_eq_psy[0]._data + b2_psy[0]._data, @@ -214,7 +213,7 @@ def test_push_bxu_Hdiv(num_elements, degree, bcs, mapping, show_plots=False): pusher_psy = Pusher_psy( particles, - Pyccelkernel(pusher_kernels.push_bxu_Hdiv), + PyccelKernel(pusher_kernels.push_bxu_Hdiv), ( derham.args_derham, b2_eq_psy[0]._data + b2_psy[0]._data, @@ -333,7 +332,7 @@ def test_push_bxu_Hcurl(num_elements, degree, bcs, mapping, show_plots=False): pusher_psy = Pusher_psy( particles, - Pyccelkernel(pusher_kernels.push_bxu_Hcurl), + PyccelKernel(pusher_kernels.push_bxu_Hcurl), ( derham.args_derham, b2_eq_psy[0]._data + b2_psy[0]._data, @@ -452,7 +451,7 @@ def test_push_bxu_H1vec(num_elements, degree, bcs, mapping, show_plots=False): pusher_psy = Pusher_psy( particles, - Pyccelkernel(pusher_kernels.push_bxu_H1vec), + PyccelKernel(pusher_kernels.push_bxu_H1vec), ( derham.args_derham, b2_eq_psy[0]._data + b2_psy[0]._data, @@ -578,7 +577,7 @@ def test_push_bxu_Hdiv_pauli(num_elements, degree, bcs, mapping, show_plots=Fals pusher_psy = Pusher_psy( particles, - Pyccelkernel(pusher_kernels.push_bxu_Hdiv_pauli), + PyccelKernel(pusher_kernels.push_bxu_Hdiv_pauli), ( derham.args_derham, *derham.degree, @@ -688,7 +687,7 @@ def test_push_eta_rk4(num_elements, degree, bcs, mapping, show_plots=False): pusher_psy = Pusher_psy( particles, - Pyccelkernel(pusher_kernels.push_eta_stage), + PyccelKernel(pusher_kernels.push_eta_stage), (butcher.a_stage, butcher.b, butcher.c), domain.args_domain, alpha_in_kernel=1.0, diff --git a/src/struphy/pic/tests/test_sorting.py b/src/struphy/pic/tests/test_sorting.py index 680331c3b..84f26e0fb 100644 --- a/src/struphy/pic/tests/test_sorting.py +++ b/src/struphy/pic/tests/test_sorting.py @@ -18,7 +18,7 @@ @pytest.mark.parametrize("ny", [16, 80]) @pytest.mark.parametrize("nz", [32, 90]) @pytest.mark.parametrize("algo", ["fortran_ordering", "c_ordering"]) -def test_flattening_1(nx, ny, nz, algo): +def test_flattening_fortran(nx, ny, nz, algo): from struphy.pic.sorting_kernels import flatten_index, unflatten_index n1s = xp.array(xp.random.rand(10) * (nx + 1), dtype=int) @@ -38,7 +38,7 @@ def test_flattening_1(nx, ny, nz, algo): @pytest.mark.parametrize("ny", [16, 80]) @pytest.mark.parametrize("nz", [32, 90]) @pytest.mark.parametrize("algo", ["fortran_ordering", "c_ordering"]) -def test_flattening_2(nx, ny, nz, algo): +def test_flattening_c(nx, ny, nz, algo): from struphy.pic.sorting_kernels import flatten_index, unflatten_index n1s = xp.array(xp.random.rand(10) * (nx + 1), dtype=int) @@ -58,7 +58,7 @@ def test_flattening_2(nx, ny, nz, algo): @pytest.mark.parametrize("ny", [16, 80]) @pytest.mark.parametrize("nz", [32, 90]) @pytest.mark.parametrize("algo", ["fortran_ordering", "c_ordering"]) -def test_flattening_3(nx, ny, nz, algo): +def test_flattening_roundtrip(nx, ny, nz, algo): from struphy.pic.sorting_kernels import flatten_index, unflatten_index n1s = xp.array(xp.random.rand(10) * (nx + 1), dtype=int) @@ -150,7 +150,7 @@ def test_sorting(num_elements, degree, bcs, mapping, Np): if __name__ == "__main__": - test_flattening_1(8, 8, 8, "c_orderwding") + test_flattening_roundtrip(8, 8, 8, "c_ordering") # test_sorting( # [8, 9, 10], # [2, 3, 4], diff --git a/src/struphy/propagators/current_coupling_5d_curlb.py b/src/struphy/propagators/current_coupling_5d_curlb.py index 2014fd27f..28c4a1363 100644 --- a/src/struphy/propagators/current_coupling_5d_curlb.py +++ b/src/struphy/propagators/current_coupling_5d_curlb.py @@ -3,6 +3,7 @@ import logging from dataclasses import dataclass +from cunumpy import PyccelKernel from feectools.ddm.mpi import mpi as MPI from line_profiler import profile @@ -17,7 +18,6 @@ from struphy.pic.pushing import pusher_kernels_gc from struphy.pic.pushing.pusher import Pusher from struphy.propagators.base import Propagator -from struphy.utils.pyccel import Pyccelkernel from struphy.utils.utils import check_option logger = logging.getLogger("struphy") @@ -203,7 +203,7 @@ def allocate(self): self._ACC = Accumulator( self.variables.energetic_ions.particles, self.options.u_space, - Pyccelkernel(accum_kernels_gc.cc_lin_mhd_5d_curlb), + PyccelKernel(accum_kernels_gc.cc_lin_mhd_5d_curlb), self.mass_ops, self.domain.args_domain, add_vector=True, @@ -228,11 +228,11 @@ def allocate(self): # define Pusher if self.options.u_space == "Hcurl": - pusher_kernel = Pyccelkernel(pusher_kernels_gc.push_gc_cc_J1_Hcurl) + pusher_kernel = PyccelKernel(pusher_kernels_gc.push_gc_cc_J1_Hcurl) elif self.options.u_space == "Hdiv": - pusher_kernel = Pyccelkernel(pusher_kernels_gc.push_gc_cc_J1_Hdiv) + pusher_kernel = PyccelKernel(pusher_kernels_gc.push_gc_cc_J1_Hdiv) elif self.options.u_space == "H1vec": - pusher_kernel = Pyccelkernel(pusher_kernels_gc.push_gc_cc_J1_H1vec) + pusher_kernel = PyccelKernel(pusher_kernels_gc.push_gc_cc_J1_H1vec) else: raise ValueError( f'{self.options.u_space =} not valid, choose from "Hcurl", "Hdiv" or "H1vec.', diff --git a/src/struphy/propagators/current_coupling_5d_density.py b/src/struphy/propagators/current_coupling_5d_density.py index 64f4828c6..7f788fe55 100644 --- a/src/struphy/propagators/current_coupling_5d_density.py +++ b/src/struphy/propagators/current_coupling_5d_density.py @@ -1,6 +1,7 @@ import logging from dataclasses import dataclass +from cunumpy import PyccelKernel from feectools.ddm.mpi import mpi as MPI from feectools.linalg.solvers import inverse from line_profiler import profile @@ -13,7 +14,6 @@ from struphy.pic.accumulation.filter import FilterParameters from struphy.pic.accumulation.particles_to_grid import Accumulator from struphy.propagators.base import Propagator -from struphy.utils.pyccel import Pyccelkernel from struphy.utils.utils import check_option logger = logging.getLogger("struphy") @@ -168,7 +168,7 @@ def allocate(self): self._ACC = Accumulator( self.energetic_ions.particles, self.options.u_space, - Pyccelkernel(accum_kernels_gc.cc_lin_mhd_5d_D), + PyccelKernel(accum_kernels_gc.cc_lin_mhd_5d_D), self.mass_ops, self.domain.args_domain, add_vector=False, diff --git a/src/struphy/propagators/current_coupling_5d_gradb.py b/src/struphy/propagators/current_coupling_5d_gradb.py index e5a283345..9e119e71f 100644 --- a/src/struphy/propagators/current_coupling_5d_gradb.py +++ b/src/struphy/propagators/current_coupling_5d_gradb.py @@ -5,6 +5,7 @@ from typing import Literal import cunumpy as xp +from cunumpy import PyccelKernel from feectools.ddm.mpi import mpi as MPI from feectools.linalg.solvers import inverse from line_profiler import profile @@ -20,7 +21,6 @@ from struphy.pic.accumulation.particles_to_grid import Accumulator, AccumulatorVector from struphy.pic.pushing import pusher_kernels_gc from struphy.propagators.base import Propagator -from struphy.utils.pyccel import Pyccelkernel from struphy.utils.utils import check_option logger = logging.getLogger("struphy") @@ -251,7 +251,7 @@ def allocate(self): self._ACC = Accumulator( self.variables.energetic_ions.particles, self.options.u_space, - Pyccelkernel(accum_kernels_gc.cc_lin_mhd_5d_gradB), + PyccelKernel(accum_kernels_gc.cc_lin_mhd_5d_gradB), self.mass_ops, self.domain.args_domain, add_vector=True, diff --git a/src/struphy/propagators/current_coupling_6d_current.py b/src/struphy/propagators/current_coupling_6d_current.py index 1b2b39985..033627f79 100644 --- a/src/struphy/propagators/current_coupling_6d_current.py +++ b/src/struphy/propagators/current_coupling_6d_current.py @@ -3,6 +3,7 @@ import logging from dataclasses import dataclass +from cunumpy import PyccelKernel from feectools.ddm.mpi import mpi as MPI from line_profiler import profile @@ -17,7 +18,6 @@ from struphy.pic.pushing import pusher_kernels from struphy.pic.pushing.pusher import Pusher from struphy.propagators.base import Propagator -from struphy.utils.pyccel import Pyccelkernel from struphy.utils.utils import check_option logger = logging.getLogger("struphy") @@ -173,7 +173,7 @@ def allocate(self): self._accumulator = Accumulator( particles, self.options.u_space, - Pyccelkernel(accum_kernels.cc_lin_mhd_6d_2), + PyccelKernel(accum_kernels.cc_lin_mhd_6d_2), self.mass_ops, self.domain.args_domain, add_vector=True, @@ -197,11 +197,11 @@ def allocate(self): # load particle pusher kernel if self.options.u_space == "Hcurl": - kernel = Pyccelkernel(pusher_kernels.push_bxu_Hcurl) + kernel = PyccelKernel(pusher_kernels.push_bxu_Hcurl) elif self.options.u_space == "Hdiv": - kernel = Pyccelkernel(pusher_kernels.push_bxu_Hdiv) + kernel = PyccelKernel(pusher_kernels.push_bxu_Hdiv) elif self.options.u_space == "H1vec": - kernel = Pyccelkernel(pusher_kernels.push_bxu_H1vec) + kernel = PyccelKernel(pusher_kernels.push_bxu_H1vec) else: raise ValueError( f'{self.options.u_space =} not valid, choose from "Hcurl", "Hdiv" or "H1vec.', diff --git a/src/struphy/propagators/current_coupling_6d_density.py b/src/struphy/propagators/current_coupling_6d_density.py index cfb227c13..19ac34394 100644 --- a/src/struphy/propagators/current_coupling_6d_density.py +++ b/src/struphy/propagators/current_coupling_6d_density.py @@ -1,6 +1,7 @@ import logging from dataclasses import dataclass +from cunumpy import PyccelKernel from feectools.ddm.mpi import mpi as MPI from feectools.linalg.solvers import inverse from line_profiler import profile @@ -13,7 +14,6 @@ from struphy.pic.accumulation.filter import FilterParameters from struphy.pic.accumulation.particles_to_grid import Accumulator from struphy.propagators.base import Propagator -from struphy.utils.pyccel import Pyccelkernel from struphy.utils.utils import check_option logger = logging.getLogger("struphy") @@ -184,7 +184,7 @@ def allocate(self): self._accumulator = Accumulator( particles, self.options.u_space, - Pyccelkernel(accum_kernels.cc_lin_mhd_6d_1), + PyccelKernel(accum_kernels.cc_lin_mhd_6d_1), self.mass_ops, self.domain.args_domain, add_vector=False, diff --git a/src/struphy/propagators/efield_weights_coupling.py b/src/struphy/propagators/efield_weights_coupling.py index 6f07c091b..06127176b 100644 --- a/src/struphy/propagators/efield_weights_coupling.py +++ b/src/struphy/propagators/efield_weights_coupling.py @@ -4,6 +4,7 @@ from dataclasses import dataclass import cunumpy as xp +from cunumpy import PyccelKernel from line_profiler import profile from struphy.feec import preconditioner @@ -17,7 +18,6 @@ from struphy.pic.pushing import pusher_kernels from struphy.pic.pushing.pusher import Pusher from struphy.propagators.base import Propagator -from struphy.utils.pyccel import Pyccelkernel from struphy.utils.utils import check_option logger = logging.getLogger("struphy") @@ -190,7 +190,7 @@ def allocate(self): self._accum = Accumulator( particles, "Hcurl", - Pyccelkernel(accum_kernels.linear_vlasov_ampere), + PyccelKernel(accum_kernels.linear_vlasov_ampere), self.mass_ops, self.domain.args_domain, add_vector=True, @@ -243,7 +243,7 @@ def allocate(self): self._pusher = Pusher( particles, - Pyccelkernel(pusher_kernels.push_weights_with_efield_lin_va), + PyccelKernel(pusher_kernels.push_weights_with_efield_lin_va), args_kernel, self.domain.args_domain, alpha_in_kernel=1.0, diff --git a/src/struphy/propagators/pressure_coupling_6d.py b/src/struphy/propagators/pressure_coupling_6d.py index 5a9e94e33..5e314ec52 100644 --- a/src/struphy/propagators/pressure_coupling_6d.py +++ b/src/struphy/propagators/pressure_coupling_6d.py @@ -3,6 +3,7 @@ import logging from dataclasses import dataclass +from cunumpy import PyccelKernel from feectools.ddm.mpi import mpi as MPI from feectools.linalg.block import BlockVector from line_profiler import profile @@ -19,7 +20,6 @@ from struphy.pic.pushing import pusher_kernels from struphy.pic.pushing.pusher import Pusher from struphy.propagators.base import Propagator -from struphy.utils.pyccel import Pyccelkernel from struphy.utils.utils import check_option logger = logging.getLogger("struphy") @@ -189,11 +189,11 @@ def allocate(self): # Call the accumulation and Pusher class if self.options.use_perp_model: - accum_ker = Pyccelkernel(accum_kernels.pc_lin_mhd_6d) - pusher_ker = Pyccelkernel(pusher_kernels.push_pc_GXu) + accum_ker = PyccelKernel(accum_kernels.pc_lin_mhd_6d) + pusher_ker = PyccelKernel(pusher_kernels.push_pc_GXu) else: - accum_ker = Pyccelkernel(accum_kernels.pc_lin_mhd_6d_full) - pusher_ker = Pyccelkernel(pusher_kernels.push_pc_GXu_full) + accum_ker = PyccelKernel(accum_kernels.pc_lin_mhd_6d_full) + pusher_ker = PyccelKernel(pusher_kernels.push_pc_GXu_full) # define Accumulator and arguments self._ACC = Accumulator( diff --git a/src/struphy/propagators/push_deterministic_diffusion.py b/src/struphy/propagators/push_deterministic_diffusion.py index 592e3341b..e09307438 100644 --- a/src/struphy/propagators/push_deterministic_diffusion.py +++ b/src/struphy/propagators/push_deterministic_diffusion.py @@ -3,6 +3,7 @@ import logging from dataclasses import dataclass +from cunumpy import PyccelKernel from line_profiler import profile from struphy.io.options import OptionsBase @@ -13,7 +14,6 @@ from struphy.pic.pushing import pusher_kernels from struphy.pic.pushing.pusher import Pusher from struphy.propagators.base import Propagator -from struphy.utils.pyccel import Pyccelkernel logger = logging.getLogger("struphy") @@ -119,7 +119,7 @@ def allocate(self): self._u_on_grid = AccumulatorVector( particles, "H1", - Pyccelkernel(accum_kernels.charge_density_0form), + PyccelKernel(accum_kernels.charge_density_0form), self.mass_ops, self.domain.args_domain, ) @@ -139,7 +139,7 @@ def allocate(self): self._pusher = Pusher( particles, - Pyccelkernel(pusher_kernels.push_deterministic_diffusion_stage), + PyccelKernel(pusher_kernels.push_deterministic_diffusion_stage), args_kernel, self.domain.args_domain, alpha_in_kernel=1.0, diff --git a/src/struphy/propagators/push_eta.py b/src/struphy/propagators/push_eta.py index 5c91543b3..de8bfe595 100644 --- a/src/struphy/propagators/push_eta.py +++ b/src/struphy/propagators/push_eta.py @@ -3,6 +3,7 @@ import logging from dataclasses import dataclass +from cunumpy import PyccelKernel from line_profiler import profile from struphy.io.options import OptionsBase @@ -11,7 +12,6 @@ from struphy.pic.pushing import pusher_kernels from struphy.pic.pushing.pusher import Pusher from struphy.propagators.base import Propagator -from struphy.utils.pyccel import Pyccelkernel logger = logging.getLogger("struphy") @@ -91,7 +91,7 @@ def options(self, new): @profile def allocate(self): # get kernel - kernel = Pyccelkernel(pusher_kernels.push_eta_stage) + kernel = PyccelKernel(pusher_kernels.push_eta_stage) # define algorithm butcher = self.options.butcher diff --git a/src/struphy/propagators/push_eta_pc.py b/src/struphy/propagators/push_eta_pc.py index a047f7d48..0be9fdeb7 100644 --- a/src/struphy/propagators/push_eta_pc.py +++ b/src/struphy/propagators/push_eta_pc.py @@ -3,6 +3,7 @@ import logging from dataclasses import dataclass +from cunumpy import PyccelKernel from line_profiler import profile from struphy.io.options import LiteralOptions, OptionsBase @@ -11,7 +12,6 @@ from struphy.pic.pushing import pusher_kernels from struphy.pic.pushing.pusher import Pusher from struphy.propagators.base import Propagator -from struphy.utils.pyccel import Pyccelkernel from struphy.utils.utils import check_option logger = logging.getLogger("struphy") @@ -125,11 +125,11 @@ def allocate(self): # get kernell: if self.options.u_space == "Hcurl": - kernel = Pyccelkernel(pusher_kernels.push_pc_eta_stage_Hcurl) + kernel = PyccelKernel(pusher_kernels.push_pc_eta_stage_Hcurl) elif self.options.u_space == "Hdiv": - kernel = Pyccelkernel(pusher_kernels.push_pc_eta_stage_Hdiv) + kernel = PyccelKernel(pusher_kernels.push_pc_eta_stage_Hdiv) elif self.options.u_space == "H1vec": - kernel = Pyccelkernel(pusher_kernels.push_pc_eta_stage_H1vec) + kernel = PyccelKernel(pusher_kernels.push_pc_eta_stage_H1vec) else: raise ValueError( f'{self.options.u_space =} not valid, choose from "Hcurl", "Hdiv" or "H1vec.', diff --git a/src/struphy/propagators/push_guiding_center_bx_estar.py b/src/struphy/propagators/push_guiding_center_bx_estar.py index 215a65a21..edc770f3f 100644 --- a/src/struphy/propagators/push_guiding_center_bx_estar.py +++ b/src/struphy/propagators/push_guiding_center_bx_estar.py @@ -4,6 +4,7 @@ from dataclasses import dataclass from typing import Literal +from cunumpy import PyccelKernel from line_profiler import profile from struphy.io.options import LiteralOptions, OptionsBase @@ -12,7 +13,6 @@ from struphy.pic.pushing import eval_kernels_gc, pusher_kernels_gc from struphy.pic.pushing.pusher import Pusher from struphy.propagators.base import Propagator -from struphy.utils.pyccel import Pyccelkernel from struphy.utils.utils import check_option logger = logging.getLogger("struphy") @@ -306,7 +306,7 @@ def allocate(self): ) # pusher kernel - kernel = Pyccelkernel(pusher_kernels_gc.push_gc_bxEstar_discrete_gradient_1st_order_newton) + kernel = PyccelKernel(pusher_kernels_gc.push_gc_bxEstar_discrete_gradient_1st_order_newton) alpha_in_kernel = 1.0 # evaluate at eta^{n+1,k} and save args_kernel = ( @@ -340,7 +340,7 @@ def allocate(self): ) # evaluate at eta^{n+1,k} and save # pusher kernel - kernel = Pyccelkernel(pusher_kernels_gc.push_gc_bxEstar_discrete_gradient_1st_order) + kernel = PyccelKernel(pusher_kernels_gc.push_gc_bxEstar_discrete_gradient_1st_order) alpha_in_kernel = 0.5 # evaluate at mid-point args_kernel = ( @@ -386,7 +386,7 @@ def allocate(self): ) # evaluate at eta^{n+1,k} and save) # pusher kernel - kernel = Pyccelkernel(pusher_kernels_gc.push_gc_bxEstar_discrete_gradient_2nd_order) + kernel = PyccelKernel(pusher_kernels_gc.push_gc_bxEstar_discrete_gradient_2nd_order) alpha_in_kernel = 0.5 # evaluate at mid-point args_kernel = ( @@ -427,7 +427,7 @@ def allocate(self): butcher = self.options.butcher # temp fix due to refactoring of ButcherTableau: - kernel = Pyccelkernel(pusher_kernels_gc.push_gc_bxEstar_explicit_multistage) + kernel = PyccelKernel(pusher_kernels_gc.push_gc_bxEstar_explicit_multistage) args_kernel = ( self.derham.args_derham, diff --git a/src/struphy/propagators/push_guiding_center_parallel.py b/src/struphy/propagators/push_guiding_center_parallel.py index 2a21786a8..a10607544 100644 --- a/src/struphy/propagators/push_guiding_center_parallel.py +++ b/src/struphy/propagators/push_guiding_center_parallel.py @@ -4,6 +4,7 @@ from dataclasses import dataclass from typing import Literal +from cunumpy import PyccelKernel from line_profiler import profile from struphy.io.options import LiteralOptions, OptionsBase @@ -12,7 +13,6 @@ from struphy.pic.pushing import eval_kernels_gc, pusher_kernels_gc from struphy.pic.pushing.pusher import Pusher from struphy.propagators.base import Propagator -from struphy.utils.pyccel import Pyccelkernel from struphy.utils.utils import check_option logger = logging.getLogger("struphy") @@ -325,7 +325,7 @@ def allocate(self): ) # pusher kernel - kernel = Pyccelkernel(pusher_kernels_gc.push_gc_Bstar_discrete_gradient_1st_order_newton) + kernel = PyccelKernel(pusher_kernels_gc.push_gc_Bstar_discrete_gradient_1st_order_newton) alpha_in_kernel = 1.0 # evaluate at eta^{n+1,k} and save args_kernel = ( @@ -358,7 +358,7 @@ def allocate(self): ) # evaluate at Z^{n+1,k} and save # pusher kernel - kernel = Pyccelkernel(pusher_kernels_gc.push_gc_Bstar_discrete_gradient_1st_order) + kernel = PyccelKernel(pusher_kernels_gc.push_gc_Bstar_discrete_gradient_1st_order) alpha_in_kernel = 0.5 # evaluate at mid-point args_kernel = ( @@ -404,7 +404,7 @@ def allocate(self): ) # evaluate at Z^{n+1,k} and save # pusher kernel - kernel = Pyccelkernel(pusher_kernels_gc.push_gc_Bstar_discrete_gradient_2nd_order) + kernel = PyccelKernel(pusher_kernels_gc.push_gc_Bstar_discrete_gradient_2nd_order) alpha_in_kernel = 0.5 # evaluate at mid-point args_kernel = ( @@ -448,7 +448,7 @@ def allocate(self): butcher = self.options.butcher # temp fix due to refactoring of ButcherTableau: - kernel = Pyccelkernel(pusher_kernels_gc.push_gc_Bstar_explicit_multistage) + kernel = PyccelKernel(pusher_kernels_gc.push_gc_Bstar_explicit_multistage) args_kernel = ( self.derham.args_derham, diff --git a/src/struphy/propagators/push_random_diffusion.py b/src/struphy/propagators/push_random_diffusion.py index a8ac86f7a..26afa4fb2 100644 --- a/src/struphy/propagators/push_random_diffusion.py +++ b/src/struphy/propagators/push_random_diffusion.py @@ -3,6 +3,7 @@ import logging from dataclasses import dataclass +from cunumpy import PyccelKernel from line_profiler import profile from numpy import array, random @@ -12,7 +13,6 @@ from struphy.pic.pushing import pusher_kernels from struphy.pic.pushing.pusher import Pusher from struphy.propagators.base import Propagator -from struphy.utils.pyccel import Pyccelkernel logger = logging.getLogger("struphy") @@ -125,7 +125,7 @@ def allocate(self): self._pusher = Pusher( particles, - Pyccelkernel(pusher_kernels.push_random_diffusion_stage), + PyccelKernel(pusher_kernels.push_random_diffusion_stage), args_kernel, self.domain.args_domain, alpha_in_kernel=1.0, diff --git a/src/struphy/propagators/push_vin_efield.py b/src/struphy/propagators/push_vin_efield.py index 9269018bd..a49c2c069 100644 --- a/src/struphy/propagators/push_vin_efield.py +++ b/src/struphy/propagators/push_vin_efield.py @@ -4,6 +4,7 @@ from dataclasses import dataclass from typing import Callable +from cunumpy import PyccelKernel from line_profiler import profile from struphy.io.options import OptionsBase @@ -11,7 +12,6 @@ from struphy.pic.pushing import pusher_kernels from struphy.pic.pushing.pusher import Pusher from struphy.propagators.base import Propagator -from struphy.utils.pyccel import Pyccelkernel logger = logging.getLogger("struphy") @@ -136,7 +136,7 @@ def allocate(self): self._pusher = Pusher( self.variables.var.particles, - Pyccelkernel(pusher_kernels.push_v_with_efield), + PyccelKernel(pusher_kernels.push_v_with_efield), args_kernel, self.domain.args_domain, alpha_in_kernel=1.0, diff --git a/src/struphy/propagators/push_vin_sph_pressure.py b/src/struphy/propagators/push_vin_sph_pressure.py index aee1a6c8f..e8b9ef59d 100644 --- a/src/struphy/propagators/push_vin_sph_pressure.py +++ b/src/struphy/propagators/push_vin_sph_pressure.py @@ -5,6 +5,7 @@ from typing import Literal import cunumpy as xp +from cunumpy import PyccelKernel from line_profiler import profile from struphy.io.options import LiteralOptions, OptionsBase @@ -12,7 +13,6 @@ from struphy.pic.pushing import eval_kernels_sph, pusher_kernels_sph from struphy.pic.pushing.pusher import Pusher from struphy.propagators.base import Propagator -from struphy.utils.pyccel import Pyccelkernel from struphy.utils.utils import check_option logger = logging.getLogger("struphy") @@ -158,9 +158,9 @@ def allocate(self): # pusher kernel if self.options.thermodynamics == "isothermal": - kernel = Pyccelkernel(pusher_kernels_sph.push_v_sph_pressure) + kernel = PyccelKernel(pusher_kernels_sph.push_v_sph_pressure) elif self.options.thermodynamics == "polytropic": - kernel = Pyccelkernel(pusher_kernels_sph.push_v_sph_pressure_ideal_gas) + kernel = PyccelKernel(pusher_kernels_sph.push_v_sph_pressure_ideal_gas) gravity = xp.array(self.options.gravity, dtype=float) diff --git a/src/struphy/propagators/push_vin_viscous_potential.py b/src/struphy/propagators/push_vin_viscous_potential.py index f432640ca..849243628 100644 --- a/src/struphy/propagators/push_vin_viscous_potential.py +++ b/src/struphy/propagators/push_vin_viscous_potential.py @@ -4,6 +4,7 @@ from dataclasses import dataclass from typing import Literal +from cunumpy import PyccelKernel from line_profiler import profile from struphy.io.options import LiteralOptions, OptionsBase @@ -11,7 +12,6 @@ from struphy.pic.pushing import eval_kernels_sph, pusher_kernels_sph from struphy.pic.pushing.pusher import Pusher from struphy.propagators.base import Propagator -from struphy.utils.pyccel import Pyccelkernel from struphy.utils.utils import check_option logger = logging.getLogger("struphy") @@ -170,7 +170,7 @@ def allocate(self): # ersetzt init args_init_visc, ) - kernel = Pyccelkernel(pusher_kernels_sph.push_v_viscosity) + kernel = PyccelKernel(pusher_kernels_sph.push_v_viscosity) args_kernel = ( boxes, diff --git a/src/struphy/propagators/push_vxb.py b/src/struphy/propagators/push_vxb.py index 09c02bbf2..fa302b333 100644 --- a/src/struphy/propagators/push_vxb.py +++ b/src/struphy/propagators/push_vxb.py @@ -4,6 +4,7 @@ from dataclasses import dataclass from typing import Literal +from cunumpy import PyccelKernel from feectools.linalg.basic import LinearOperator from feectools.linalg.block import BlockVector from line_profiler import profile @@ -13,7 +14,6 @@ from struphy.pic.pushing import pusher_kernels from struphy.pic.pushing.pusher import Pusher from struphy.propagators.base import Propagator -from struphy.utils.pyccel import Pyccelkernel from struphy.utils.utils import check_option logger = logging.getLogger("struphy") @@ -124,9 +124,9 @@ def allocate(self): # define pusher kernel if self.options.algo == "analytic": - kernel = Pyccelkernel(pusher_kernels.push_vxb_analytic) + kernel = PyccelKernel(pusher_kernels.push_vxb_analytic) elif self.options.algo == "implicit": - kernel = Pyccelkernel(pusher_kernels.push_vxb_implicit) + kernel = PyccelKernel(pusher_kernels.push_vxb_implicit) else: raise ValueError(f"{self.options.algo =} not supported.") diff --git a/src/struphy/propagators/shear_alfven_current_coupling_5d.py b/src/struphy/propagators/shear_alfven_current_coupling_5d.py index f9b5cc69e..004d7a361 100644 --- a/src/struphy/propagators/shear_alfven_current_coupling_5d.py +++ b/src/struphy/propagators/shear_alfven_current_coupling_5d.py @@ -2,6 +2,7 @@ from dataclasses import dataclass from typing import Literal +from cunumpy import PyccelKernel from feectools.ddm.mpi import mpi as MPI from feectools.linalg.block import BlockVector from feectools.linalg.solvers import inverse @@ -22,7 +23,6 @@ from struphy.pic.accumulation.filter import FilterParameters from struphy.pic.accumulation.particles_to_grid import AccumulatorVector from struphy.propagators.base import Propagator -from struphy.utils.pyccel import Pyccelkernel from struphy.utils.utils import check_option logger = logging.getLogger("struphy") @@ -197,7 +197,7 @@ def allocate(self): self._ACC = AccumulatorVector( self.energetic_ions.particles, "H1", - Pyccelkernel(accum_kernels_gc.gc_mag_density_0form), + PyccelKernel(accum_kernels_gc.gc_mag_density_0form), self.mass_ops, self.domain.args_domain, filter_params=self.options.filter_params, diff --git a/src/struphy/propagators/tests/test_curl_curl.py b/src/struphy/propagators/tests/test_curl_curl.py index 06b8b901f..b3205cdcc 100644 --- a/src/struphy/propagators/tests/test_curl_curl.py +++ b/src/struphy/propagators/tests/test_curl_curl.py @@ -3,6 +3,7 @@ import cunumpy as xp import matplotlib.pyplot as plt import pytest +from cunumpy import PyccelKernel from feectools.ddm.mpi import mpi as MPI from struphy import ( @@ -23,7 +24,6 @@ from struphy.propagators.base import Propagator from struphy.propagators.curl_curl_solve import CurlCurlSolve from struphy.topology.grids import TensorProductGrid -from struphy.utils.pyccel import Pyccelkernel logger = logging.getLogger("struphy") set_logging_level(logging.INFO) diff --git a/src/struphy/propagators/tests/test_poisson.py b/src/struphy/propagators/tests/test_poisson.py index 43c8c1b04..55ca649fd 100644 --- a/src/struphy/propagators/tests/test_poisson.py +++ b/src/struphy/propagators/tests/test_poisson.py @@ -3,6 +3,7 @@ import cunumpy as xp import matplotlib.pyplot as plt import pytest +from cunumpy import PyccelKernel from feectools.ddm.mpi import mpi as MPI from struphy import ( @@ -26,7 +27,6 @@ from struphy.propagators.base import Propagator from struphy.propagators.poisson_solve import PoissonSolve from struphy.topology.grids import TensorProductGrid -from struphy.utils.pyccel import Pyccelkernel logger = logging.getLogger("struphy") @@ -344,7 +344,7 @@ def test_poisson_accum_1d(mapping, do_plot=False): particles.initialize_weights() # particle to grid coupling - kernel = Pyccelkernel(charge_density_0form) + kernel = PyccelKernel(charge_density_0form) # control_variate=True, so no PICVariable.species is needed to build the analytical background term pic_var = PICVariable(space="Particles6D") pic_var._particles = particles diff --git a/src/struphy/propagators/vlasov_ampere_coupling.py b/src/struphy/propagators/vlasov_ampere_coupling.py index 134559221..b5c522504 100644 --- a/src/struphy/propagators/vlasov_ampere_coupling.py +++ b/src/struphy/propagators/vlasov_ampere_coupling.py @@ -4,6 +4,7 @@ from dataclasses import dataclass import cunumpy as xp +from cunumpy import PyccelKernel from line_profiler import profile from struphy.feec import preconditioner @@ -16,7 +17,6 @@ from struphy.pic.pushing import pusher_kernels from struphy.pic.pushing.pusher import Pusher from struphy.propagators.base import Propagator -from struphy.utils.pyccel import Pyccelkernel from struphy.utils.utils import check_option logger = logging.getLogger("struphy") @@ -126,7 +126,7 @@ def allocate(self): self._info = self.options.solver_params.info # get accumulation kernel - accum_kernel = Pyccelkernel(accum_kernels.vlasov_maxwell) + accum_kernel = PyccelKernel(accum_kernels.vlasov_maxwell) # Initialize Accumulator object particles = self.variables.ions.particles @@ -181,7 +181,7 @@ def allocate(self): self._pusher = Pusher( particles, - Pyccelkernel(pusher_kernels.push_v_with_efield), + PyccelKernel(pusher_kernels.push_v_with_efield), args_kernel, self.domain.args_domain, alpha_in_kernel=1.0, diff --git a/src/struphy/simulation/sim.py b/src/struphy/simulation/sim.py index 6e1af29f2..730728b57 100644 --- a/src/struphy/simulation/sim.py +++ b/src/struphy/simulation/sim.py @@ -73,6 +73,23 @@ logger = logging.getLogger("struphy") +class CuPyJSONEncoder(json.JSONEncoder): + """JSON encoder that handles CuPy arrays and NumPy arrays.""" + + def default(self, obj): + # Check if it has a .get() method (CuPy array) + if hasattr(obj, "get"): + return obj.get().tolist() if hasattr(obj.get(), "tolist") else obj.get() + # Handle NumPy arrays and scalars + import numpy as np + + if isinstance(obj, np.ndarray): + return obj.tolist() + if isinstance(obj, (np.integer, np.floating)): + return float(obj) if isinstance(obj, np.floating) else int(obj) + return super().default(obj) + + class Simulation(SimulationBase): """Top-level class to configure and run a Struphy simulation. @@ -498,20 +515,24 @@ def save_geometry_and_equil_vtk(self): ] tmp = self.domain(*grids_log) - grids_phy = [tmp[0], tmp[1], tmp[2]] + grids_phy = [ + DataContainer._as_numpy_array(tmp[0]), + DataContainer._as_numpy_array(tmp[1]), + DataContainer._as_numpy_array(tmp[2]), + ] pointData = {} det_df = self.domain.jacobian_det(*grids_log) - pointData["det_df"] = det_df + pointData["det_df"] = DataContainer._as_numpy_array(det_df) if self.equil is not None: p0 = self.equil.p0(*grids_log) - pointData["p0"] = p0 + pointData["p0"] = DataContainer._as_numpy_array(p0) n0 = self.equil.n0(*grids_log) - pointData["n0"] = n0 + pointData["n0"] = DataContainer._as_numpy_array(n0) if isinstance(self.equil, FluidEquilibriumWithB): absB0 = self.equil.absB0(*grids_log) - pointData["absB0"] = absB0 + pointData["absB0"] = DataContainer._as_numpy_array(absB0) gridToVTK(os.path.join(self.env.path_out, "geometry"), *grids_phy, pointData=pointData) @@ -538,21 +559,25 @@ def create_geometry_mesh( ] tmp = self.domain(*grids_log) - grids_phy = [tmp[0], tmp[1], tmp[2]] + grids_phy = [ + DataContainer._as_numpy_array(tmp[0]), + DataContainer._as_numpy_array(tmp[1]), + DataContainer._as_numpy_array(tmp[2]), + ] # Create PyVista structured grid mesh = pv.StructuredGrid(grids_phy[0], grids_phy[1], grids_phy[2]) # Add point data det_df = self.domain.jacobian_det(*grids_log) - mesh["det_df"] = det_df.ravel(order="F") + mesh["det_df"] = DataContainer._as_numpy_array(det_df).ravel(order="F") if self.equil is not None: p0 = self.equil.p0(*grids_log) - mesh["p0"] = p0.ravel(order="F") + mesh["p0"] = DataContainer._as_numpy_array(p0).ravel(order="F") if isinstance(self.equil, FluidEquilibriumWithB): absB0 = self.equil.absB0(*grids_log) - mesh["absB0"] = absB0.ravel(order="F") + mesh["absB0"] = DataContainer._as_numpy_array(absB0).ravel(order="F") return mesh @@ -704,12 +729,12 @@ def run(self, one_time_step: bool = False): self.time_state["index"][0] = file["restart/time/index"][-1] start_step = file["restart/time/index"][-1] - total_steps = int(round((Tend - self.time_state["value"][0]) / dt)) + total_steps = int(round((Tend - float(self.time_state["value"][0])) / dt)) logger.info(f"""\n!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! RESTARTing from: -{self.time_state["value"][0]=} -{self.time_state["value_sec"][0]=} -{self.time_state["index"][0]=} +self.time_state["value"][0]={float(self.time_state["value"][0])} +self.time_state["value_sec"][0]={float(self.time_state["value_sec"][0])} +self.time_state["index"][0]={int(self.time_state["index"][0])} !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! """) else: @@ -748,19 +773,19 @@ def run(self, one_time_step: bool = False): self.Barrier() # stop time loop? - break_cond_1 = self.time_state["value"][0] >= Tend + break_cond_1 = float(self.time_state["value"][0]) >= Tend break_cond_2 = run_time_now > self.env.max_runtime if break_cond_1 or break_cond_2: # save restart data (other data already saved below) self.data.save_data(keys=save_keys_end) end_time = time.time() - logger.info(f"\nTime steps done: {self.time_state['index'][0]}") + logger.info(f"\nTime steps done: {int(self.time_state['index'][0])}") logger.info(f"wall-clock time of simulation [sec]: {end_time - self.start_time}") logger.info("") break - if self.env.sort_step and self.time_state["index"][0] % self.env.sort_step == 0: + if self.env.sort_step and int(self.time_state["index"][0]) % self.env.sort_step == 0: t0 = time.time() for key, val in self.model.pointer.items(): if isinstance(val, Particles): @@ -774,8 +799,10 @@ def run(self, one_time_step: bool = False): logger.info("") # update time and index (round time to 10 decimals for a clean time grid!) - self.time_state["value"][0] = round(self.time_state["value"][0] + dt, 14) - self.time_state["value_sec"][0] = round(self.time_state["value_sec"][0] + dt * self.model.units.t, 14) + self.time_state["value"][0] = round(float(self.time_state["value"][0]) + dt, 14) + self.time_state["value_sec"][0] = round( + float(self.time_state["value_sec"][0]) + dt * self.model.units.t, 14 + ) self.time_state["index"][0] += 1 # perform one time step dt @@ -787,7 +814,7 @@ def run(self, one_time_step: bool = False): run_time_now = (time.time() - self.start_time) / 60 # update diagnostics data and save data - if self.time_state["index"][0] % self.env.save_step == 0: + if int(self.time_state["index"][0]) % self.env.save_step == 0: # compute scalars and kinetic data self.model.update_scalar_quantities() self.model.update_markers_to_be_saved() @@ -807,19 +834,19 @@ def run(self, one_time_step: bool = False): self.data.save_data(keys=save_keys_all) # print current time and scalar quantities to screen - step = str(self.time_state["index"][0]).zfill(len(total_steps_str)) + step = str(int(self.time_state["index"][0])).zfill(len(total_steps_str)) message = "time step:".ljust(25) + f"{step}/{total_steps + start_step}".rjust(25) message += ( "\n" + "normalized time:".ljust(25) - + "{0:4.2e} / {1:4.2e}".format(self.time_state["value"][0], Tend).rjust(25) + + "{0:4.2e} / {1:4.2e}".format(float(self.time_state["value"][0]), Tend).rjust(25) ) message += ( "\n" + "physical time [s]:".ljust(25) + "{0:4.2e} / {1:4.2e}".format( - self.time_state["value_sec"][0], + float(self.time_state["value_sec"][0]), Tend * self.model.units.t, ).rjust(25) ) @@ -1345,9 +1372,9 @@ def _initialize_hdf5_datasets(self, data: DataContainer, size: int): # store grid_info only for runs with 512 ranks or smaller if self.model.scalars.dct and self.derham is not None: if size <= 512: - file["scalar"].attrs["grid_info"] = self.derham.domain_array + file["scalar"].attrs["grid_info"] = DataContainer._as_numpy_array(self.derham.domain_array) else: - file["scalar"].attrs["grid_info"] = self.derham.domain_array[0] + file["scalar"].attrs["grid_info"] = DataContainer._as_numpy_array(self.derham.domain_array[0]) else: pass @@ -1384,9 +1411,9 @@ def _initialize_hdf5_datasets(self, data: DataContainer, size: int): # save field meta data file[key_field].attrs["space_id"] = spline.space_id - file[key_field].attrs["starts"] = spline.starts - file[key_field].attrs["ends"] = spline.ends - file[key_field].attrs["pads"] = spline.pads + file[key_field].attrs["starts"] = DataContainer._as_numpy_array(spline.starts) + file[key_field].attrs["ends"] = DataContainer._as_numpy_array(spline.ends) + file[key_field].attrs["pads"] = DataContainer._as_numpy_array(spline.pads) # save numpy array to be updated only at the end of the simulation for restart. key_field_restart = os.path.join(species_path_restart, variable) @@ -1434,7 +1461,9 @@ def _initialize_hdf5_datasets(self, data: DataContainer, size: int): data.add_data({key_df: bin_plot.df}) for dim, be in enumerate(bin_plot.bin_edges): - file[key_f].attrs["bin_centers" + "_" + str(dim + 1)] = be[:-1] + (be[1] - be[0]) / 2 + file[key_f].attrs["bin_centers" + "_" + str(dim + 1)] = DataContainer._as_numpy_array( + be[:-1] + (be[1] - be[0]) / 2 + ) for i, kd_plot in enumerate(species.saving_params.kernel_density_plots): key_n = os.path.join(key_spec, "n_sph", f"view_{i}") @@ -1444,9 +1473,9 @@ def _initialize_hdf5_datasets(self, data: DataContainer, size: int): eta1 = kd_plot.plot_pts[0][:, 0, 0] eta2 = kd_plot.plot_pts[1][0, :, 0] eta3 = kd_plot.plot_pts[2][0, 0, :] - file[key_n].attrs["eta1"] = eta1 - file[key_n].attrs["eta2"] = eta2 - file[key_n].attrs["eta3"] = eta3 + file[key_n].attrs["eta1"] = DataContainer._as_numpy_array(eta1) + file[key_n].attrs["eta2"] = DataContainer._as_numpy_array(eta2) + file[key_n].attrs["eta3"] = DataContainer._as_numpy_array(eta3) # TODO: maybe add other data # else: @@ -1583,7 +1612,7 @@ def to_run_metadata(self, file_path: str = None, **extra_data) -> str: }, ) - json_str = json.dumps(config, indent=4) + json_str = json.dumps(config, indent=4, cls=CuPyJSONEncoder) if file_path is not None: with open(file_path, "w") as f: f.write(json_str) diff --git a/src/struphy/utils/pyccel.py b/src/struphy/utils/pyccel.py deleted file mode 100644 index 5e62426a3..000000000 --- a/src/struphy/utils/pyccel.py +++ /dev/null @@ -1,25 +0,0 @@ -from typing import Any, Callable - - -class Pyccelkernel: - def __init__(self, kernel: Callable[..., Any], use_cupy: bool = False) -> None: - self._kernel = kernel - self._use_cupy = use_cupy - - def __call__(self, *args: Any, **kwargs: Any) -> Any: - if self.use_cupy: - raise NotImplementedError - else: - return self._kernel(*args, **kwargs) - - @property - def name(self): - return self.kernel.__name__ - - @property - def kernel(self) -> Callable[..., Any]: - return self._kernel - - @property - def use_cupy(self): - return self._use_cupy