diff options
Diffstat (limited to 'Python/pywarpx')
| -rw-r--r-- | Python/pywarpx/PGroup.py | 182 | ||||
| -rw-r--r-- | Python/pywarpx/WarpInterface.py | 190 | ||||
| -rw-r--r-- | Python/pywarpx/WarpX.py | 5 | ||||
| -rw-r--r-- | Python/pywarpx/__init__.py | 28 | ||||
| -rwxr-xr-x | Python/pywarpx/_libwarpx.py | 2734 | ||||
| -rw-r--r-- | Python/pywarpx/callbacks.py | 410 | ||||
| -rw-r--r-- | Python/pywarpx/fields.py | 1286 | ||||
| -rw-r--r-- | Python/pywarpx/particle_containers.py | 601 |
8 files changed, 1404 insertions, 4032 deletions
diff --git a/Python/pywarpx/PGroup.py b/Python/pywarpx/PGroup.py deleted file mode 100644 index 0dcb664e5..000000000 --- a/Python/pywarpx/PGroup.py +++ /dev/null @@ -1,182 +0,0 @@ -# Copyright 2017-2019 David Grote -# -# This file is part of WarpX. -# -# License: BSD-3-Clause-LBNL - -import numpy as np - -from ._libwarpx import libwarpx -from ._picmi import constants - - -class PGroup(object): - """Implements a class that has the same API as a warp ParticleGroup instance. - """ - - def __init__(self, igroup, ispecie, level=0): - self.igroup = igroup - self.ispecie = ispecie - self.level = level - self.ns = 1 # Number of species - - self.gallot() - - def name(self): - return 'WarpXParticleGroup' - - def gallot(self): - self.lebcancel_pusher = 0 # turns on/off cancellation of E+VxB within V push (logical type) - self.lebcancel = 0 # turns on/off cancellation of E+VxB before V push (logical type) - - self.sm = np.zeros(self.ns) # Species mass [kg] - self.sq = np.zeros(self.ns) # Species charge [C] - self.sw = np.ones(self.ns) # Species weight, (real particles per simulation particles) - - self.sid = np.arange(self.ns, dtype=int) # Global species index for each species - self.ndts = np.ones(self.ns, dtype=int) # Stride for time step advance for each species - self.ldts = np.ones(self.ns, dtype=int) # (logical type) - self.lvdts = np.ones(self.ns, dtype=int) # (logical type) - self.iselfb = np.zeros(self.ns, dtype=int) # Group number for particles that are affected by - # their own magnetic field, using the 1/gamma**2 - # approximation. The correction is not applied to - # group number -1. - self.fselfb = np.zeros(self.ns) # The scaling factor, vz. - self.l_maps = np.zeros(self.ns) - self.dtscale = np.ones(self.ns) # Scale factor applied to time step size for each - # species. Only makes sense in steaday and and - # transverse slice modes. - - self.limplicit = np.zeros(self.ns, dtype=int) # Flags implicit particle species (logical type) - self.iimplicit = np.full(self.ns, -1, dtype=int) # Group number for implicit particles - self.ldoadvance = np.ones(self.ns, dtype=int) # Flags whether particles are time advanced (logical type) - self.lboundaries = np.ones(self.ns, dtype=int) # Flags whether boundary conditions need to be applied (logical type) - self.lparaxial = np.zeros(self.ns, dtype=int) # Flags to turn on/off paraxial approximation (logical type) - - self.zshift = np.zeros(self.ns) - self.gamma_ebcancel_max = np.ones(self.ns) # maximum value allowed for ExB cancellation - - # --- Temporary fix - gchange = gallot - - def allocated(self, name): - return True - - def addspecies(self): - pass - - def getnpid(self): - return libwarpx.get_nattr() - npid = property(getnpid) - - def getnps(self): - return np.array([len(self.xp)], dtype='l') - nps = property(getnps) - - def getins(self): - return np.ones(self.ns, dtype='l') - ins = property(getins) - - def getipmax(self): - return np.array([0, len(self.xp)], dtype='l') - ipmax = property(getipmax) - - def getnpmax(self): - return self.nps.sum() - npmax = property(getnpmax) - - def getxp(self): - return libwarpx.get_particle_x(self.ispecie, self.level)[self.igroup] - xp = property(getxp) - - def getyp(self): - return libwarpx.get_particle_y(self.ispecie, self.level)[self.igroup] - yp = property(getyp) - - def getrp(self): - return libwarpx.get_particle_r(self.ispecie, self.level)[self.igroup] - rp = property(getrp) - - def getzp(self): - return libwarpx.get_particle_z(self.ispecie, self.level)[self.igroup] - zp = property(getzp) - - def getuxp(self): - return libwarpx.get_particle_ux(self.ispecie, self.level)[self.igroup] - uxp = property(getuxp) - - def getuyp(self): - return libwarpx.get_particle_uy(self.ispecie, self.level)[self.igroup] - uyp = property(getuyp) - - def getuzp(self): - return libwarpx.get_particle_uz(self.ispecie, self.level)[self.igroup] - uzp = property(getuzp) - - def getw(self): - return libwarpx.get_particle_weight(self.ispecie, self.level)[self.igroup] - - def getpid(self, id): - pid = libwarpx.get_particle_arrays(self.ispecie, id, self.level)[self.igroup] - return np.array([pid]).T - - def getgaminv(self): - uxp = self.getuxp() - uyp = self.getuyp() - uzp = self.getuzp() - return np.sqrt(1. - (uxp**2 + uyp**2 + uzp**2)/constants.c**2) - gaminv = property(getgaminv) - - def getex(self): - raise Exception('Particle E fields not supported') - ex = property(getex) - - def getey(self): - raise Exception('Particle E fields not supported') - ey = property(getey) - - def getez(self): - raise Exception('Particle E fields not supported') - ez = property(getez) - - def getbx(self): - raise Exception('Particle B fields not supported') - bx = property(getbx) - - def getby(self): - raise Exception('Particle B fields not supported') - by = property(getby) - - def getbz(self): - raise Exception('Particle B fields not supported') - bz = property(getbz) - - def gettheta(self): - return libwarpx.get_particle_theta(self.ispecie, self.level)[self.igroup] - theta = property(gettheta) - -class PGroups(object): - def __init__(self, ispecie=0, level=0): - self.ispecie = ispecie - self.level = level - - def setuppgroups(self): - xall = libwarpx.get_particle_x(self.ispecie, self.level) - self.ngroups = len(xall) - - self._pgroups = [] - for igroup in range(self.ngroups): - self._pgroups.append(PGroup(igroup, self.ispecie, self.level)) - - def __iter__(self): - self.setuppgroups() - for igroup in range(self.ngroups): - yield self._pgroups[igroup] - - def __getitem__(self, key): - self.setuppgroups() - return self._pgroups[key] - - def __len__(self): - self.setuppgroups() - return len(self._pgroups) diff --git a/Python/pywarpx/WarpInterface.py b/Python/pywarpx/WarpInterface.py deleted file mode 100644 index dc80e4bc9..000000000 --- a/Python/pywarpx/WarpInterface.py +++ /dev/null @@ -1,190 +0,0 @@ -# Copyright 2019 David Grote -# -# This file is part of WarpX. -# -# License: BSD-3-Clause-LBNL - -# This sets up an interface between Warp and WarpX, allowing access to WarpX data using -# classes from Warp. - -# The routine warp_species will return an instance of the Species class from Warp, -# giving nearly all of the capability of that class, including accessing the data -# using down selection and all of the plots. - -# The class WarpX_EM3D inherits from Warp's EM3D class. It primarily provides -# access to the field plotting routines. - -import warp - -from pywarpx import PGroup - -from . import fields - -# The particle weight is always the first pid -warp.top.wpid = 1 - -def warp_species(warp_type, picmi_species, level=0): - """Returns a Warp species that has a reference to the WarpX particles. - """ - pgroups = PGroup.PGroups(ispecie=picmi_species.species_number, level=level) - return warp.Species(type=warp_type, pgroups=pgroups) - - -class _WarpX_FIELDtype(object): - """Mirrors part of the EM3D_FIELDtype type from Warp - yf - """ - def __init__(self, yf): - self.yf = yf - - -class _WarpX_BLOCKtype(object): - """Mirrors part of the EM3D_BLOCKtype type from Warp - core - xmin, ymin, zmin - xmax, ymax, zmax - dx, dy, dz - xrbnd, yrbnd, zrbnd - """ - def __init__(self, fields, picmi_grid): - self.core = _WarpX_FIELDtype(fields) - self.picmi_grid = picmi_grid - - self.xmin = self.picmi_grid.lower_bound[0] - if self.picmi_grid.number_of_dimensions == 3: - self.ymin = self.picmi_grid.lower_bound[1] - else: - self.ymin = 0. - self.zmin = self.picmi_grid.lower_bound[-1] - - self.xmax = self.picmi_grid.upper_bound[0] - if self.picmi_grid.number_of_dimensions == 3: - self.ymax = self.picmi_grid.upper_bound[1] - else: - self.ymax = 0. - self.zmax = self.picmi_grid.upper_bound[-1] - - self.dx = (self.xmax - self.xmin)/self.picmi_grid.number_of_cells[0] - if self.picmi_grid.number_of_dimensions == 3: - self.dy = (self.ymax - self.ymin)/self.picmi_grid.number_of_cells[1] - else: - self.dy = 1. - self.dz = (self.zmax - self.zmin)/self.picmi_grid.number_of_cells[-1] - - self.xrbnd = 0 - self.yrbnd = 0 - self.zrbnd = 0 - - -class _WarpX_YEEFIELDtype(object): - """Mirrors part of the EM3D_YEEFIELDtype type from Warp - Exp, Eyp, Ezp - Bxp, Byp, Bzp - Ex, Ey, Ez - Bx, By, Bz - Jx, Jy, Jz - """ - def __init__(self, level=0): - self.level = level - self._Ex_wrap = fields.ExWrapper(level, include_ghosts=True) - self._Ey_wrap = fields.EyWrapper(level, include_ghosts=True) - self._Ez_wrap = fields.EzWrapper(level, include_ghosts=True) - self._Bx_wrap = fields.BxWrapper(level, include_ghosts=True) - self._By_wrap = fields.ByWrapper(level, include_ghosts=True) - self._Bz_wrap = fields.BzWrapper(level, include_ghosts=True) - self._Jx_wrap = fields.JxWrapper(level, include_ghosts=True) - self._Jy_wrap = fields.JyWrapper(level, include_ghosts=True) - self._Jz_wrap = fields.JzWrapper(level, include_ghosts=True) - - self._Ex_wrap._getlovects() # --- Calculated nghosts - self.nxguard = self._Ex_wrap.nghosts - self.nyguard = self._Ex_wrap.nghosts - self.nzguard = self._Ex_wrap.nghosts - - def _get_wrapped_array(self, wrapper): - result = wrapper[...] - if len(result.shape) == 2: - # --- Add the middle dimension that Warp is expecting - result.shape = [result.shape[0], 1, result.shape[1]] - return result - - @property - def Exp(self): - return self._get_wrapped_array(self._Ex_wrap) - @property - def Eyp(self): - return self._get_wrapped_array(self._Ey_wrap) - @property - def Ezp(self): - return self._get_wrapped_array(self._Ez_wrap) - @property - def Bxp(self): - return self._get_wrapped_array(self._Bx_wrap) - @property - def Byp(self): - return self._get_wrapped_array(self._By_wrap) - @property - def Bzp(self): - return self._get_wrapped_array(self._Bz_wrap) - @property - def Ex(self): - return self._get_wrapped_array(self._Ex_wrap) - @property - def Ey(self): - return self._get_wrapped_array(self._Ey_wrap) - @property - def Ez(self): - return self._get_wrapped_array(self._Ez_wrap) - @property - def Bx(self): - return self._get_wrapped_array(self._Bx_wrap) - @property - def By(self): - return self._get_wrapped_array(self._By_wrap) - @property - def Bz(self): - return self._get_wrapped_array(self._Bz_wrap) - @property - def Jx(self): - return self._get_wrapped_array(self._Jx_wrap) - @property - def Jy(self): - return self._get_wrapped_array(self._Jy_wrap) - @property - def Jz(self): - return self._get_wrapped_array(self._Jz_wrap) - - -class WarpX_EM3D(warp.EM3D): - """Mirrors part of the Warp EM3D class, mostly diagnostics. - """ - def __init__(self, picmi_grid, level=0): - self.picmi_grid = picmi_grid - self.level = level - - # --- Only define what is necessary for the diagnostics - self.fields = _WarpX_YEEFIELDtype(level) - self.block = _WarpX_BLOCKtype(self.fields, picmi_grid) - - self.isactive = True - - self.l_1dz = (picmi_grid.number_of_dimensions == 1) - self.l_2dxz = (picmi_grid.number_of_dimensions == 2) - try: - picmi_grid.nr - except AttributeError: - self.l_2drz = False - else: - self.l_2drz = True - - self.l4symtry = False - self.l2symtry = False - - self.nx = picmi_grid.number_of_cells[0] - if not self.l_2dxz: - self.ny = picmi_grid.number_of_cells[1] - else: - self.ny = 0 - self.nz = picmi_grid.number_of_cells[-1] - - self.zgrid = 0. # --- This should be obtained from WarpX diff --git a/Python/pywarpx/WarpX.py b/Python/pywarpx/WarpX.py index ef35f28ce..cfa556414 100644 --- a/Python/pywarpx/WarpX.py +++ b/Python/pywarpx/WarpX.py @@ -6,6 +6,7 @@ # License: BSD-3-Clause-LBNL import re +import sys from . import Particles from .Algo import algo @@ -91,7 +92,9 @@ class WarpX(Bucket): return argv def init(self, mpi_comm=None, **kw): - argv = ['warpx'] + self.create_argv_list(**kw) + # note: argv[0] needs to be an absolute path so it works with AMReX backtraces + # https://github.com/AMReX-Codes/amrex/issues/3435 + argv = [sys.executable] + self.create_argv_list(**kw) libwarpx.initialize(argv, mpi_comm=mpi_comm) def evolve(self, nsteps=-1): diff --git a/Python/pywarpx/__init__.py b/Python/pywarpx/__init__.py index c74405acf..f0f4177d4 100644 --- a/Python/pywarpx/__init__.py +++ b/Python/pywarpx/__init__.py @@ -1,9 +1,27 @@ -# Copyright 2016-2022 Andrew Myers, David Grote, Lorenzo Giacomel +# Copyright 2016-2023 The WarpX Community # # This file is part of WarpX. # +# Authors: Andrew Myers, David Grote, Lorenzo Giacomel, Axel Huebl # License: BSD-3-Clause-LBNL +import os + +# Python 3.8+ on Windows: DLL search paths for dependent +# shared libraries +# Refs.: +# - https://github.com/python/cpython/issues/80266 +# - https://docs.python.org/3.8/library/os.html#os.add_dll_directory +if os.name == "nt": + # add anything in the current directory + pwd = __file__.rsplit(os.sep, 1)[0] + os.sep + os.add_dll_directory(pwd) + # add anything in PATH + paths = os.environ.get("PATH", "") + for p in paths.split(";"): + if os.path.exists(p): + os.add_dll_directory(p) + from .Algo import algo from .Amr import amr from .Amrex import amrex @@ -21,5 +39,11 @@ from .Particles import newspecies, particles from .WarpX import warpx from ._libwarpx import libwarpx -# This is a circulor import and must happen after the import of libwarpx +# This is a circular import and must happen after the import of libwarpx from . import picmi # isort:skip + +# TODO +#__version__ = cxx.__version__ +#__doc__ = cxx.__doc__ +#__license__ = cxx.__license__ +#__author__ = cxx.__author__ diff --git a/Python/pywarpx/_libwarpx.py b/Python/pywarpx/_libwarpx.py index 2a297b908..592459132 100755 --- a/Python/pywarpx/_libwarpx.py +++ b/Python/pywarpx/_libwarpx.py @@ -1,52 +1,31 @@ -# Copyright 2017-2019 Andrew Myers, David Grote, Remi Lehe -# Weiqun Zhang +# Copyright 2017-2022 The WarpX Community # -# This file is part of WarpX. +# This file is part of WarpX. It defines the wrapper functions that directly +# call the underlying compiled routines through pybind11. +# +# NOTE: We will reduce the libwarpx.py level of abstraction eventually! +# Please add new functionality directly to pybind11-bound modules +# and call them via sim.extension.libwarpx_so. ... and sim.extension.warpx. +# ... from user code. +# +# Authors: Axel Huebl, Andrew Myers, David Grote, Remi Lehe, Weiqun Zhang # # License: BSD-3-Clause-LBNL import atexit -import ctypes -from ctypes.util import find_library as _find_library -# --- This defines the wrapper functions that directly call the underlying compiled routines import os -import platform import sys import numpy as np -from numpy.ctypeslib import ndpointer as _ndpointer from .Geometry import geometry -try: - # --- If mpi4py is going to be used, this needs to be imported - # --- before libwarpx is loaded, because mpi4py calls MPI_Init - from mpi4py import MPI - - # --- Change MPI Comm type depending on MPICH (int) or OpenMPI (void*) - if MPI._sizeof(MPI.Comm) == ctypes.sizeof(ctypes.c_int): - _MPI_Comm_type = ctypes.c_int - else: - _MPI_Comm_type = ctypes.c_void_p -except ImportError: - MPI = None - _MPI_Comm_type = ctypes.c_void_p - -if platform.system() == 'Windows': - _path_libc = _find_library('msvcrt') -else: - _path_libc = _find_library('c') -_libc = ctypes.CDLL(_path_libc) -_LP_c_int = ctypes.POINTER(ctypes.c_int) -_LP_c_char = ctypes.c_char_p - class LibWarpX(): - - """This class manages the warpx shared object, the library from the compiled C++ code. + """This class manages the WarpX classes, part of the Python module from the compiled C++ code. It will only load the library when it is referenced, and this can only be done after the geometry is defined so that the version of the library that is needed can be determined. - Once loaded, all of the setting of function call interfaces is setup. + Once loaded, all the settings of function call interfaces are set up. """ def __init__(self): @@ -82,7 +61,9 @@ class LibWarpX(): if 'libwarpx_so' in self.__dict__: raise RuntimeError( - "Invalid attempt to load libwarpx_so library multiple times." + "Invalid attempt to load the pybind11 bindings library multiple times. " + "Note that multiple AMReX/WarpX geometries cannot be loaded yet into the same Python process. " + "Please write separate scripts for each geometry." ) # --- Use geometry to determine whether to import the 1D, 2D, 3D or RZ version. @@ -91,7 +72,7 @@ class LibWarpX(): _prob_lo = geometry.prob_lo _dims = geometry.dims except AttributeError: - raise Exception('The shared object could not be loaded. The geometry must be setup before the WarpX shared object can be accessesd. The geometry determines which version of the shared object to load.') + raise Exception('The shared object could not be loaded. The geometry must be setup before the WarpX pybind11 module can be accessesd. The geometry determines which version of the shared object to load.') if _dims == 'RZ': self.geometry_dim = 'rz' @@ -100,287 +81,33 @@ class LibWarpX(): else: raise Exception('Undefined geometry %d'%_dims) - # this is a plain C/C++ shared library, not a Python module - if os.name == 'nt': - mod_ext = 'dll' - else: - mod_ext = 'so' - libname = f'libwarpx.{self.geometry_dim}.{mod_ext}' - try: - self.libwarpx_so = ctypes.CDLL(os.path.join(self._get_package_root(), libname)) - except OSError as e: - value = e.args[0] - if f'{libname}: cannot open shared object file: No such file or directory' in value: - raise Exception(f'"{libname}" was not installed. Installation instructions can be found here https://warpx.readthedocs.io/en/latest/install/users.html') from e - else: - print("Failed to load the libwarpx shared object library") - raise - - # WarpX can be compiled using either double or float - self.libwarpx_so.warpx_Real_size.restype = ctypes.c_int - self.libwarpx_so.warpx_ParticleReal_size.restype = ctypes.c_int - - _Real_size = self.libwarpx_so.warpx_Real_size() - _ParticleReal_size = self.libwarpx_so.warpx_ParticleReal_size() - - if _Real_size == 8: - c_real = ctypes.c_double - self._numpy_real_dtype = 'f8' - else: - c_real = ctypes.c_float - self._numpy_real_dtype = 'f4' - - if _ParticleReal_size == 8: - c_particlereal = ctypes.c_double - self._numpy_particlereal_dtype = 'f8' - else: - c_particlereal = ctypes.c_float - self._numpy_particlereal_dtype = 'f4' - - self.dim = self.libwarpx_so.warpx_SpaceDim() - - # our particle data type, depends on _ParticleReal_size - _p_struct = ( - [(d, self._numpy_particlereal_dtype) for d in 'xyz'[:self.dim]] - + [('idcpu', 'u8')] - ) - self._p_dtype = np.dtype(_p_struct, align=True) - - _numpy_to_ctypes = {} - _numpy_to_ctypes[self._numpy_particlereal_dtype] = c_particlereal - _numpy_to_ctypes['u8'] = ctypes.c_uint64 - - class Particle(ctypes.Structure): - _fields_ = [(v[0], _numpy_to_ctypes[v[1]]) for v in _p_struct] - - # some useful typenames - _LP_particle_p = ctypes.POINTER(ctypes.POINTER(Particle)) - _LP_LP_c_int = ctypes.POINTER(_LP_c_int) - #_LP_c_void_p = ctypes.POINTER(ctypes.c_void_p) - _LP_c_real = ctypes.POINTER(c_real) - _LP_LP_c_real = ctypes.POINTER(_LP_c_real) - _LP_c_particlereal = ctypes.POINTER(c_particlereal) - _LP_LP_c_particlereal = ctypes.POINTER(_LP_c_particlereal) - _LP_LP_c_char = ctypes.POINTER(_LP_c_char) - - # set the arg and return types of the wrapped functions - self.libwarpx_so.amrex_init.argtypes = (ctypes.c_int, _LP_LP_c_char) - self.libwarpx_so.amrex_init_with_inited_mpi.argtypes = (ctypes.c_int, _LP_LP_c_char, _MPI_Comm_type) - self.libwarpx_so.warpx_getParticleStructs.restype = _LP_particle_p - self.libwarpx_so.warpx_getParticleArrays.restype = _LP_LP_c_particlereal - self.libwarpx_so.warpx_getParticleCompIndex.restype = ctypes.c_int - self.libwarpx_so.warpx_getEfield.restype = _LP_LP_c_real - self.libwarpx_so.warpx_getEfieldLoVects.restype = _LP_c_int - self.libwarpx_so.warpx_getEfieldCP.restype = _LP_LP_c_real - self.libwarpx_so.warpx_getEfieldCPLoVects.restype = _LP_c_int - self.libwarpx_so.warpx_getEfieldFP.restype = _LP_LP_c_real - self.libwarpx_so.warpx_getEfieldFPLoVects.restype = _LP_c_int - self.libwarpx_so.warpx_getEfieldCP_PML.restype = _LP_LP_c_real - self.libwarpx_so.warpx_getEfieldCPLoVects_PML.restype = _LP_c_int - self.libwarpx_so.warpx_getEfieldFP_PML.restype = _LP_LP_c_real - self.libwarpx_so.warpx_getEfieldFPLoVects_PML.restype = _LP_c_int - self.libwarpx_so.warpx_getBfield.restype = _LP_LP_c_real - self.libwarpx_so.warpx_getBfieldLoVects.restype = _LP_c_int - self.libwarpx_so.warpx_getBfieldCP.restype = _LP_LP_c_real - self.libwarpx_so.warpx_getBfieldCPLoVects.restype = _LP_c_int - self.libwarpx_so.warpx_getBfieldFP.restype = _LP_LP_c_real - self.libwarpx_so.warpx_getBfieldFPLoVects.restype = _LP_c_int - self.libwarpx_so.warpx_getBfieldCP_PML.restype = _LP_LP_c_real - self.libwarpx_so.warpx_getBfieldCPLoVects_PML.restype = _LP_c_int - self.libwarpx_so.warpx_getBfieldFP_PML.restype = _LP_LP_c_real - self.libwarpx_so.warpx_getBfieldFPLoVects_PML.restype = _LP_c_int - self.libwarpx_so.warpx_getCurrentDensity.restype = _LP_LP_c_real - self.libwarpx_so.warpx_getCurrentDensityLoVects.restype = _LP_c_int - self.libwarpx_so.warpx_getCurrentDensityCP.restype = _LP_LP_c_real - self.libwarpx_so.warpx_getCurrentDensityCPLoVects.restype = _LP_c_int - self.libwarpx_so.warpx_getCurrentDensityFP.restype = _LP_LP_c_real - self.libwarpx_so.warpx_getCurrentDensityFPLoVects.restype = _LP_c_int - self.libwarpx_so.warpx_getCurrentDensityCP_PML.restype = _LP_LP_c_real - self.libwarpx_so.warpx_getCurrentDensityCPLoVects_PML.restype = _LP_c_int - self.libwarpx_so.warpx_getCurrentDensityFP_PML.restype = _LP_LP_c_real - self.libwarpx_so.warpx_getCurrentDensityFPLoVects_PML.restype = _LP_c_int - self.libwarpx_so.warpx_getCurrentDensityFP_Ampere.restype = _LP_LP_c_real - self.libwarpx_so.warpx_getChargeDensityCP.restype = _LP_LP_c_real - self.libwarpx_so.warpx_getChargeDensityCPLoVects.restype = _LP_c_int - self.libwarpx_so.warpx_getChargeDensityFP.restype = _LP_LP_c_real - self.libwarpx_so.warpx_getChargeDensityFPLoVects.restype = _LP_c_int - self.libwarpx_so.warpx_getPhiFP.restype = _LP_LP_c_real - self.libwarpx_so.warpx_getPhiFPLoVects.restype = _LP_c_int - self.libwarpx_so.warpx_getVectorPotentialFP.restype = _LP_LP_c_real - self.libwarpx_so.warpx_getVectorPotentialFPLoVects.restype = _LP_c_int - self.libwarpx_so.warpx_getFfieldCP.restype = _LP_LP_c_real - self.libwarpx_so.warpx_getFfieldCPLoVects.restype = _LP_c_int - self.libwarpx_so.warpx_getFfieldFP.restype = _LP_LP_c_real - self.libwarpx_so.warpx_getFfieldFPLoVects.restype = _LP_c_int - self.libwarpx_so.warpx_getFfieldCP_PML.restype = _LP_LP_c_real - self.libwarpx_so.warpx_getFfieldCPLoVects_PML.restype = _LP_c_int - self.libwarpx_so.warpx_getFfieldFP_PML.restype = _LP_LP_c_real - self.libwarpx_so.warpx_getFfieldFPLoVects_PML.restype = _LP_c_int - self.libwarpx_so.warpx_getGfieldCP.restype = _LP_LP_c_real - self.libwarpx_so.warpx_getGfieldCPLoVects.restype = _LP_c_int - self.libwarpx_so.warpx_getGfieldFP.restype = _LP_LP_c_real - self.libwarpx_so.warpx_getGfieldFPLoVects.restype = _LP_c_int - self.libwarpx_so.warpx_getGfieldCP_PML.restype = _LP_LP_c_real - self.libwarpx_so.warpx_getGfieldCPLoVects_PML.restype = _LP_c_int - self.libwarpx_so.warpx_getGfieldFP_PML.restype = _LP_LP_c_real - self.libwarpx_so.warpx_getGfieldFPLoVects_PML.restype = _LP_c_int - self.libwarpx_so.warpx_getEdgeLengths.restype = _LP_LP_c_real - self.libwarpx_so.warpx_getEdgeLengthsLoVects.restype = _LP_c_int - self.libwarpx_so.warpx_getFaceAreas.restype = _LP_LP_c_real - self.libwarpx_so.warpx_getFaceAreasLoVects.restype = _LP_c_int - - self.libwarpx_so.warpx_sumParticleCharge.restype = c_real - self.libwarpx_so.warpx_getParticleBoundaryBufferSize.restype = ctypes.c_int - self.libwarpx_so.warpx_getParticleBoundaryBufferStructs.restype = _LP_LP_c_particlereal - self.libwarpx_so.warpx_getParticleBoundaryBuffer.restype = _LP_LP_c_particlereal - self.libwarpx_so.warpx_getParticleBoundaryBufferScrapedSteps.restype = _LP_LP_c_int - - self.libwarpx_so.warpx_getEx_nodal_flag.restype = _LP_c_int - self.libwarpx_so.warpx_getEy_nodal_flag.restype = _LP_c_int - self.libwarpx_so.warpx_getEz_nodal_flag.restype = _LP_c_int - self.libwarpx_so.warpx_getBx_nodal_flag.restype = _LP_c_int - self.libwarpx_so.warpx_getBy_nodal_flag.restype = _LP_c_int - self.libwarpx_so.warpx_getBz_nodal_flag.restype = _LP_c_int - self.libwarpx_so.warpx_getJx_nodal_flag.restype = _LP_c_int - self.libwarpx_so.warpx_getJy_nodal_flag.restype = _LP_c_int - self.libwarpx_so.warpx_getJz_nodal_flag.restype = _LP_c_int - self.libwarpx_so.warpx_getAx_nodal_flag.restype = _LP_c_int - self.libwarpx_so.warpx_getAy_nodal_flag.restype = _LP_c_int - self.libwarpx_so.warpx_getAz_nodal_flag.restype = _LP_c_int - self.libwarpx_so.warpx_getRho_nodal_flag.restype = _LP_c_int - self.libwarpx_so.warpx_getPhi_nodal_flag.restype = _LP_c_int - self.libwarpx_so.warpx_getF_nodal_flag.restype = _LP_c_int - self.libwarpx_so.warpx_getG_nodal_flag.restype = _LP_c_int - self.libwarpx_so.warpx_getF_pml_nodal_flag.restype = _LP_c_int - self.libwarpx_so.warpx_getG_pml_nodal_flag.restype = _LP_c_int - - self.libwarpx_so.warpx_get_edge_lengths_x_nodal_flag.restype = _LP_c_int - self.libwarpx_so.warpx_get_edge_lengths_y_nodal_flag.restype = _LP_c_int - self.libwarpx_so.warpx_get_edge_lengths_z_nodal_flag.restype = _LP_c_int - self.libwarpx_so.warpx_get_face_areas_x_nodal_flag.restype = _LP_c_int - self.libwarpx_so.warpx_get_face_areas_y_nodal_flag.restype = _LP_c_int - self.libwarpx_so.warpx_get_face_areas_z_nodal_flag.restype = _LP_c_int - - #self.libwarpx_so.warpx_getPMLSigma.restype = _LP_c_real - #self.libwarpx_so.warpx_getPMLSigmaStar.restype = _LP_c_real - #self.libwarpx_so.warpx_ComputePMLFactors.argtypes = (ctypes.c_int, c_real) - self.libwarpx_so.warpx_addNParticles.argtypes = (ctypes.c_char_p, ctypes.c_int, - _ndpointer(c_particlereal, flags="C_CONTIGUOUS"), - _ndpointer(c_particlereal, flags="C_CONTIGUOUS"), - _ndpointer(c_particlereal, flags="C_CONTIGUOUS"), - _ndpointer(c_particlereal, flags="C_CONTIGUOUS"), - _ndpointer(c_particlereal, flags="C_CONTIGUOUS"), - _ndpointer(c_particlereal, flags="C_CONTIGUOUS"), - ctypes.c_int, - _ndpointer(c_particlereal, flags="C_CONTIGUOUS"), - ctypes.c_int, - _ndpointer(ctypes.c_int, flags="C_CONTIGUOUS"), - ctypes.c_int) - self.libwarpx_so.warpx_convert_id_to_long.argtypes = (_ndpointer(ctypes.c_int64, flags="C_CONTIGUOUS"), - _ndpointer(Particle, flags="C_CONTIGUOUS"), - ctypes.c_int) - self.libwarpx_so.warpx_convert_cpu_to_int.argtypes = (_ndpointer(ctypes.c_int32, flags="C_CONTIGUOUS"), - _ndpointer(Particle, flags="C_CONTIGUOUS"), - ctypes.c_int) - self.libwarpx_so.warpx_getProbLo.restype = c_real - self.libwarpx_so.warpx_getProbHi.restype = c_real - self.libwarpx_so.warpx_getCellSize.restype = c_real - self.libwarpx_so.warpx_getistep.restype = ctypes.c_int - self.libwarpx_so.warpx_gett_new.restype = c_real - self.libwarpx_so.warpx_getdt.restype = c_real - self.libwarpx_so.warpx_maxStep.restype = ctypes.c_int - self.libwarpx_so.warpx_stopTime.restype = c_real - self.libwarpx_so.warpx_finestLevel.restype = ctypes.c_int - self.libwarpx_so.warpx_getMyProc.restype = ctypes.c_int - self.libwarpx_so.warpx_getNProcs.restype = ctypes.c_int - - self.libwarpx_so.warpx_EvolveE.argtypes = [c_real] - self.libwarpx_so.warpx_EvolveB.argtypes = [c_real] - self.libwarpx_so.warpx_FillBoundaryE.argtypes = [] - self.libwarpx_so.warpx_FillBoundaryB.argtypes = [] - self.libwarpx_so.warpx_UpdateAuxilaryData.argtypes = [] - self.libwarpx_so.warpx_SyncCurrent.argtypes = [] - self.libwarpx_so.warpx_PushParticlesandDepose.argtypes = [c_real] - self.libwarpx_so.warpx_getProbLo.argtypes = [ctypes.c_int] - self.libwarpx_so.warpx_getProbHi.argtypes = [ctypes.c_int] - self.libwarpx_so.warpx_getCellSize.argtypes = [ctypes.c_int, ctypes.c_int] - self.libwarpx_so.warpx_getistep.argtypes = [ctypes.c_int] - self.libwarpx_so.warpx_setistep.argtypes = [ctypes.c_int, ctypes.c_int] - self.libwarpx_so.warpx_gett_new.argtypes = [ctypes.c_int] - self.libwarpx_so.warpx_sett_new.argtypes = [ctypes.c_int, c_real] - self.libwarpx_so.warpx_getdt.argtypes = [ctypes.c_int] - self.libwarpx_so.warpx_setPotentialEB.argtypes = [ctypes.c_char_p] - self.libwarpx_so.warpx_setPlasmaLensStrength.argtypes = [ctypes.c_int, c_real, c_real] - - def _get_boundary_number(self, boundary): - ''' - - Utility function to find the boundary number given a boundary name. - - Parameters - ---------- - - boundary : str - The boundary from which to get the scraped particle data. In the - form x/y/z_hi/lo or eb. - - Returns - ------- - int - Integer index in the boundary scraper buffer for the given boundary. - ''' - if self.geometry_dim == '3d': - dimensions = {'x' : 0, 'y' : 1, 'z' : 2} - elif self.geometry_dim == '2d' or self.geometry_dim == 'rz': - dimensions = {'x' : 0, 'z' : 1} - elif self.geometry_dim == '1d': - dimensions = {'z' : 0} - else: - raise RuntimeError(f"Unknown simulation geometry: {self.geometry_dim}") - - if boundary != 'eb': - boundary_parts = boundary.split("_") - dim_num = dimensions[boundary_parts[0]] - if boundary_parts[1] == 'lo': - side = 0 - elif boundary_parts[1] == 'hi': - side = 1 - else: - raise RuntimeError(f'Unknown boundary specified: {boundary}') - boundary_num = 2 * dim_num + side - else: - if self.geometry_dim == '3d': - boundary_num = 6 - elif self.geometry_dim == '2d' or self.geometry_dim == 'rz': - boundary_num = 4 - elif self.geometry_dim == '1d': - boundary_num = 2 - - return boundary_num - - @staticmethod - def _array1d_from_pointer(pointer, dtype, size): - ''' - - Function for converting a ctypes pointer to a numpy array - - ''' - if not pointer: - raise Exception(f'_array1d_from_pointer: pointer is a nullptr') - if sys.version_info.major >= 3: - # from where do I import these? this might only work for CPython... - #PyBuf_READ = 0x100 - PyBUF_WRITE = 0x200 - buffer_from_memory = ctypes.pythonapi.PyMemoryView_FromMemory - buffer_from_memory.argtypes = (ctypes.c_void_p, ctypes.c_int, ctypes.c_int) - buffer_from_memory.restype = ctypes.py_object - buf = buffer_from_memory(pointer, dtype.itemsize*size, PyBUF_WRITE) - else: - buffer_from_memory = ctypes.pythonapi.PyBuffer_FromReadWriteMemory - buffer_from_memory.restype = ctypes.py_object - buf = buffer_from_memory(pointer, dtype.itemsize*size) - return np.frombuffer(buf, dtype=dtype, count=size) + if self.geometry_dim == "1d": + import amrex.space1d as amr + self.amr = amr + from . import warpx_pybind_1d as cxx_1d + self.libwarpx_so = cxx_1d + self.dim = 1 + elif self.geometry_dim == "2d": + import amrex.space2d as amr + self.amr = amr + from . import warpx_pybind_2d as cxx_2d + self.libwarpx_so = cxx_2d + self.dim = 2 + elif self.geometry_dim == "rz": + import amrex.space2d as amr + self.amr = amr + from . import warpx_pybind_rz as cxx_rz + self.libwarpx_so = cxx_rz + self.dim = 2 + elif self.geometry_dim == "3d": + import amrex.space3d as amr + self.amr = amr + from . import warpx_pybind_3d as cxx_3d + self.libwarpx_so = cxx_3d + self.dim = 3 + except ImportError: + raise Exception(f"Dimensionality '{self.geometry_dim}' was not compiled in this Python install. Please recompile with -DWarpX_DIMS={_dims}") def getNProcs(self): ''' @@ -388,7 +115,7 @@ class LibWarpX(): Get the number of processors ''' - return self.libwarpx_so.warpx_getNProcs() + return self.libwarpx_so.getNProcs() def getMyProc(self): ''' @@ -396,7 +123,7 @@ class LibWarpX(): Get the number of the processor ''' - return self.libwarpx_so.warpx_getMyProc() + return self.libwarpx_so.getMyProc() def get_nattr(self): ''' @@ -417,29 +144,16 @@ class LibWarpX(): species_name: str Name of the species ''' - - return self.libwarpx_so.warpx_nCompsSpecies( - ctypes.c_char_p(species_name.encode('utf-8'))) + warpx = self.libwarpx_so.get_instance() + mpc = warpx.multi_particle_container() + pc = mpc.get_particle_container_from_name(species_name) + return pc.num_real_comps() def amrex_init(self, argv, mpi_comm=None): - # --- Construct the ctype list of strings to pass in - argc = len(argv) - - # note: +1 since there is an extra char-string array element, - # that ANSII C requires to be a simple NULL entry - # https://stackoverflow.com/a/39096006/2719194 - argvC = (_LP_c_char * (argc+1))() - for i, arg in enumerate(argv): - enc_arg = arg.encode('utf-8') - argvC[i] = _LP_c_char(enc_arg) - argvC[argc] = _LP_c_char(b"\0") # +1 element must be NULL - - if mpi_comm is None or MPI is None: - self.libwarpx_so.amrex_init(argc, argvC) + if mpi_comm is None: # or MPI is None: + self.libwarpx_so.amrex_init(argv) else: - comm_ptr = MPI._addressof(mpi_comm) - comm_val = _MPI_Comm_type.from_address(comm_ptr) - self.libwarpx_so.amrex_init_with_inited_mpi(argc, argvC, comm_val) + raise Exception('mpi_comm argument not yet supported') def initialize(self, argv=None, mpi_comm=None): ''' @@ -450,11 +164,16 @@ class LibWarpX(): if argv is None: argv = sys.argv self.amrex_init(argv, mpi_comm) - self.libwarpx_so.warpx_ConvertLabParamsToBoost() - self.libwarpx_so.warpx_ReadBCParams() + self.libwarpx_so.convert_lab_params_to_boost() + self.libwarpx_so.read_BC_params() if self.geometry_dim == 'rz': - self.libwarpx_so.warpx_CheckGriddingForRZSpectral() - self.libwarpx_so.warpx_init() + self.libwarpx_so.check_gridding_for_RZ_spectral() + self.warpx = self.libwarpx_so.get_instance() + self.warpx.initialize_data() + self.libwarpx_so.execute_python_callback("afterinit") + self.libwarpx_so.execute_python_callback("particleloader") + + #self.libwarpx_so.warpx_init() self.initialized = True @@ -464,9 +183,15 @@ class LibWarpX(): Call finalize for WarpX and AMReX. Registered to run at program exit. ''' + # TODO: simplify, part of pyAMReX already if self.initialized: - self.libwarpx_so.warpx_finalize() - self.libwarpx_so.amrex_finalize(finalize_mpi) + del self.warpx + # The call to warpx_finalize causes a crash - don't know why + #self.libwarpx_so.warpx_finalize() + self.libwarpx_so.amrex_finalize() + + from pywarpx import callbacks + callbacks.clear_all() def getistep(self, level=0): ''' @@ -479,7 +204,7 @@ class LibWarpX(): The refinement level to reference ''' - return self.libwarpx_so.warpx_getistep(level) + return self.warpx.getistep(level) def gett_new(self, level=0): ''' @@ -493,7 +218,7 @@ class LibWarpX(): The refinement level to reference ''' - return self.libwarpx_so.warpx_gett_new(level) + return self.warpx.gett_new(level) def evolve(self, num_steps=-1): ''' @@ -508,9 +233,9 @@ class LibWarpX(): The number of steps to take ''' - self.libwarpx_so.warpx_evolve(num_steps); + self.warpx.evolve(num_steps) - def getProbLo(self, direction): + def getProbLo(self, direction, level=0): ''' Get the values of the lower domain boundary. @@ -522,9 +247,9 @@ class LibWarpX(): ''' assert 0 <= direction < self.dim, 'Inappropriate direction specified' - return self.libwarpx_so.warpx_getProbLo(direction) + return self.warpx.Geom(level).ProbLo(direction) - def getProbHi(self, direction): + def getProbHi(self, direction, level=0): ''' Get the values of the upper domain boundary. @@ -536,7 +261,7 @@ class LibWarpX(): ''' assert 0 <= direction < self.dim, 'Inappropriate direction specified' - return self.libwarpx_so.warpx_getProbHi(direction) + return self.warpx.Geom(level).ProbHi(direction) def getCellSize(self, direction, level=0): ''' @@ -597,598 +322,7 @@ class LibWarpX(): # # self.libwarpx_so.warpx_ComputePMLFactors(lev, dt) - def add_particles(self, species_name, x=None, y=None, z=None, ux=None, uy=None, - uz=None, w=None, unique_particles=True, **kwargs): - ''' - A function for adding particles to the WarpX simulation. - - Parameters - ---------- - - species_name : str - The type of species for which particles will be added - - x, y, z : arrays or scalars - The particle positions (default = 0.) - - ux, uy, uz : arrays or scalars - The particle momenta (default = 0.) - - w : array or scalars - Particle weights (default = 0.) - - unique_particles : bool - Whether the particles are unique or duplicated on several processes - (default = True) - kwargs : dict - Containing an entry for all the extra particle attribute arrays. If - an attribute is not given it will be set to 0. - ''' - - # --- Get length of arrays, set to one for scalars - lenx = np.size(x) - leny = np.size(y) - lenz = np.size(z) - lenux = np.size(ux) - lenuy = np.size(uy) - lenuz = np.size(uz) - lenw = np.size(w) - - # --- Find the max length of the parameters supplied - maxlen = 0 - if x is not None: - maxlen = max(maxlen, lenx) - if y is not None: - maxlen = max(maxlen, leny) - if z is not None: - maxlen = max(maxlen, lenz) - if ux is not None: - maxlen = max(maxlen, lenux) - if uy is not None: - maxlen = max(maxlen, lenuy) - if uz is not None: - maxlen = max(maxlen, lenuz) - if w is not None: - maxlen = max(maxlen, lenw) - - # --- Make sure that the lengths of the input parameters are consistent - assert x is None or lenx==maxlen or lenx==1, "Length of x doesn't match len of others" - assert y is None or leny==maxlen or leny==1, "Length of y doesn't match len of others" - assert z is None or lenz==maxlen or lenz==1, "Length of z doesn't match len of others" - assert ux is None or lenux==maxlen or lenux==1, "Length of ux doesn't match len of others" - assert uy is None or lenuy==maxlen or lenuy==1, "Length of uy doesn't match len of others" - assert uz is None or lenuz==maxlen or lenuz==1, "Length of uz doesn't match len of others" - assert w is None or lenw==maxlen or lenw==1, "Length of w doesn't match len of others" - for key, val in kwargs.items(): - assert np.size(val)==1 or len(val)==maxlen, f"Length of {key} doesn't match len of others" - - # --- Broadcast scalars into appropriate length arrays - # --- If the parameter was not supplied, use the default value - if lenx == 1: - x = np.full(maxlen, (x or 0.), self._numpy_particlereal_dtype) - if leny == 1: - y = np.full(maxlen, (y or 0.), self._numpy_particlereal_dtype) - if lenz == 1: - z = np.full(maxlen, (z or 0.), self._numpy_particlereal_dtype) - if lenux == 1: - ux = np.full(maxlen, (ux or 0.), self._numpy_particlereal_dtype) - if lenuy == 1: - uy = np.full(maxlen, (uy or 0.), self._numpy_particlereal_dtype) - if lenuz == 1: - uz = np.full(maxlen, (uz or 0.), self._numpy_particlereal_dtype) - if lenw == 1: - w = np.full(maxlen, (w or 0.), self._numpy_particlereal_dtype) - for key, val in kwargs.items(): - if np.size(val) == 1: - kwargs[key] = np.full( - maxlen, val, self._numpy_particlereal_dtype - ) - - # --- The number of built in attributes - # --- The three velocities - built_in_attrs = 3 - if self.geometry_dim == 'rz': - # --- With RZ, there is also theta - built_in_attrs += 1 - - # --- The number of extra attributes (including the weight) - nattr = self.get_nattr_species(species_name) - built_in_attrs - attr = np.zeros((maxlen, nattr), self._numpy_particlereal_dtype) - attr[:,0] = w - - # --- Note that the velocities are handled separately and not included in attr - # --- (even though they are stored as attributes in the C++) - for key, vals in kwargs.items(): - attr[:,self.get_particle_comp_index(species_name, key) - built_in_attrs] = vals - - nattr_int = 0 - attr_int = np.empty([0], ctypes.c_int) - - # Iff x/y/z/ux/uy/uz are not numpy arrays of the correct dtype, new - # array copies are made with the correct dtype - x = x.astype(self._numpy_particlereal_dtype, copy=False) - y = y.astype(self._numpy_particlereal_dtype, copy=False) - z = z.astype(self._numpy_particlereal_dtype, copy=False) - ux = ux.astype(self._numpy_particlereal_dtype, copy=False) - uy = uy.astype(self._numpy_particlereal_dtype, copy=False) - uz = uz.astype(self._numpy_particlereal_dtype, copy=False) - - self.libwarpx_so.warpx_addNParticles( - ctypes.c_char_p(species_name.encode('utf-8')), x.size, - x, y, z, ux, uy, uz, nattr, attr, nattr_int, attr_int, unique_particles - ) - - def get_particle_count(self, species_name, local=False): - ''' - Get the number of particles of the specified species in the simulation. - - Parameters - ---------- - - species_name : str - The species name that the number will be returned for - - local : bool - If True the particle count on this processor will be returned. - Default False. - - Returns - ------- - - int - An integer count of the number of particles - ''' - - return self.libwarpx_so.warpx_getNumParticles( - ctypes.c_char_p(species_name.encode('utf-8')), local - ) - - def get_particle_structs(self, species_name, level): - ''' - This returns a list of numpy arrays containing the particle struct data - on each tile for this process. The particle data is represented as a structured - numpy array and contains the particle 'x', 'y', 'z', and 'idcpu'. - - The data for the numpy arrays are not copied, but share the underlying - memory buffer with WarpX. The numpy arrays are fully writeable. - - Parameters - ---------- - - species_name : str - The species name that the data will be returned for - - level : int - The refinement level to reference - - Returns - ------- - - List of numpy arrays - The requested particle struct data - ''' - - particles_per_tile = _LP_c_int() - num_tiles = ctypes.c_int(0) - data = self.libwarpx_so.warpx_getParticleStructs( - ctypes.c_char_p(species_name.encode('utf-8')), level, - ctypes.byref(num_tiles), ctypes.byref(particles_per_tile) - ) - - particle_data = [] - for i in range(num_tiles.value): - if particles_per_tile[i] == 0: - continue - arr = self._array1d_from_pointer(data[i], self._p_dtype, particles_per_tile[i]) - particle_data.append(arr) - - _libc.free(particles_per_tile) - _libc.free(data) - return particle_data - - def get_particle_arrays(self, species_name, comp_name, level): - ''' - This returns a list of numpy arrays containing the particle array data - on each tile for this process. - - The data for the numpy arrays are not copied, but share the underlying - memory buffer with WarpX. The numpy arrays are fully writeable. - - Parameters - ---------- - - species_name : str - The species name that the data will be returned for - - comp_name : str - The component of the array data that will be returned - - level : int - The refinement level to reference - - Returns - ------- - - List of numpy arrays - The requested particle array data - ''' - - particles_per_tile = _LP_c_int() - num_tiles = ctypes.c_int(0) - data = self.libwarpx_so.warpx_getParticleArrays( - ctypes.c_char_p(species_name.encode('utf-8')), - ctypes.c_char_p(comp_name.encode('utf-8')), - level, ctypes.byref(num_tiles), ctypes.byref(particles_per_tile) - ) - - particle_data = [] - for i in range(num_tiles.value): - if particles_per_tile[i] == 0: - continue - if not data[i]: - raise Exception(f'get_particle_arrays: data[i] for i={i} was not initialized') - arr = np.ctypeslib.as_array(data[i], (particles_per_tile[i],)) - try: - # This fails on some versions of numpy - arr.setflags(write=1) - except ValueError: - pass - particle_data.append(arr) - - _libc.free(particles_per_tile) - _libc.free(data) - return particle_data - - def get_particle_x(self, species_name, level=0): - ''' - - Return a list of numpy arrays containing the particle 'x' - positions on each tile. - - ''' - structs = self.get_particle_structs(species_name, level) - if self.geometry_dim == '3d' or self.geometry_dim == '2d': - return [struct['x'] for struct in structs] - elif self.geometry_dim == 'rz': - return [struct['x']*np.cos(theta) for struct, theta in zip(structs, self.get_particle_theta(species_name))] - elif self.geometry_dim == '1d': - raise Exception('get_particle_x: There is no x coordinate with 1D Cartesian') - - def get_particle_y(self, species_name, level=0): - ''' - - Return a list of numpy arrays containing the particle 'y' - positions on each tile. - - ''' - structs = self.get_particle_structs(species_name, level) - if self.geometry_dim == '3d': - return [struct['y'] for struct in structs] - elif self.geometry_dim == 'rz': - return [struct['x']*np.sin(theta) for struct, theta in zip(structs, self.get_particle_theta(species_name))] - elif self.geometry_dim == '1d' or self.geometry_dim == '2d': - raise Exception('get_particle_y: There is no y coordinate with 1D or 2D Cartesian') - - def get_particle_r(self, species_name, level=0): - ''' - - Return a list of numpy arrays containing the particle 'r' - positions on each tile. - - ''' - structs = self.get_particle_structs(species_name, level) - if self.geometry_dim == 'rz': - return [struct['x'] for struct in structs] - elif self.geometry_dim == '3d': - return [np.sqrt(struct['x']**2 + struct['y']**2) for struct in structs] - elif self.geometry_dim == '2d' or self.geometry_dim == '1d': - raise Exception('get_particle_r: There is no r coordinate with 1D or 2D Cartesian') - - def get_particle_z(self, species_name, level=0): - ''' - - Return a list of numpy arrays containing the particle 'z' - positions on each tile. - - ''' - structs = self.get_particle_structs(species_name, level) - if self.geometry_dim == '3d': - return [struct['z'] for struct in structs] - elif self.geometry_dim == 'rz' or self.geometry_dim == '2d': - return [struct['y'] for struct in structs] - elif self.geometry_dim == '1d': - return [struct['x'] for struct in structs] - - def get_particle_id(self, species_name, level=0): - ''' - - Return a list of numpy arrays containing the particle 'id' - numbers on each tile. - - ''' - ids = [] - structs = self.get_particle_structs(species_name, level) - for ptile_of_structs in structs: - arr = np.empty(ptile_of_structs.shape, np.int64) - self.libwarpx_so.warpx_convert_id_to_long(arr, ptile_of_structs, arr.size) - ids.append(arr) - return ids - - def get_particle_cpu(self, species_name, level=0): - ''' - - Return a list of numpy arrays containing the particle 'cpu' - numbers on each tile. - - ''' - cpus = [] - structs = self.get_particle_structs(species_name, level) - for ptile_of_structs in structs: - arr = np.empty(ptile_of_structs.shape, np.int32) - self.libwarpx_so.warpx_convert_cpu_to_int(arr, ptile_of_structs, arr.size) - cpus.append(arr) - return cpus - - def get_particle_weight(self, species_name, level=0): - ''' - - Return a list of numpy arrays containing the particle - weight on each tile. - - ''' - - return self.get_particle_arrays(species_name, 'w', level) - - def get_particle_ux(self, species_name, level=0): - ''' - - Return a list of numpy arrays containing the particle - x momentum on each tile. - - ''' - - return self.get_particle_arrays(species_name, 'ux', level) - - def get_particle_uy(self, species_name, level=0): - ''' - - Return a list of numpy arrays containing the particle - y momentum on each tile. - - ''' - - return self.get_particle_arrays(species_name, 'uy', level) - - def get_particle_uz(self, species_name, level=0): - ''' - - Return a list of numpy arrays containing the particle - z momentum on each tile. - - ''' - - return self.get_particle_arrays(species_name, 'uz', level) - - def get_particle_theta(self, species_name, level=0): - ''' - - Return a list of numpy arrays containing the particle - theta on each tile. - - ''' - - if self.geometry_dim == 'rz': - return self.get_particle_arrays(species_name, 'theta', level) - elif self.geometry_dim == '3d': - structs = self.get_particle_structs(species_name, level) - return [np.arctan2(struct['y'], struct['x']) for struct in structs] - elif self.geometry_dim == '2d' or self.geometry_dim == '1d': - raise Exception('get_particle_theta: There is no theta coordinate with 1D or 2D Cartesian') - - def get_particle_comp_index(self, species_name, pid_name): - ''' - Get the component index for a given particle attribute. This is useful - to get the corrent ordering of attributes when adding new particles using - `add_particles()`. - - Parameters - ---------- - - species_name : str - The name of the species - - pid_name : str - Name of the component for which the index will be returned - - Returns - ------- - - int - Integer corresponding to the index of the requested attribute - ''' - - return self.libwarpx_so.warpx_getParticleCompIndex( - ctypes.c_char_p(species_name.encode('utf-8')), - ctypes.c_char_p(pid_name.encode('utf-8')) - ) - - def add_real_comp(self, species_name, pid_name, comm=True): - ''' - Add a real component to the particle data array. - - Parameters - ---------- - - species_name : str - The species name for which the new component will be added - - pid_name : str - Name that can be used to identify the new component - - comm : bool - Should the component be communicated - ''' - - self.libwarpx_so.warpx_addRealComp( - ctypes.c_char_p(species_name.encode('utf-8')), - ctypes.c_char_p(pid_name.encode('utf-8')), comm - ) - - def get_species_charge_sum(self, species_name, local=False): - ''' - Returns the total charge in the simulation due to the given species. - - Parameters - ---------- - - species_name : str - The species name for which charge will be summed - - local : bool - If True return total charge per processor - ''' - - return self.libwarpx_so.warpx_sumParticleCharge( - ctypes.c_char_p(species_name.encode('utf-8')), local - ) - - def get_particle_boundary_buffer_size(self, species_name, boundary, local=False): - ''' - This returns the number of particles that have been scraped so far in the simulation - from the specified boundary and of the specified species. - - Parameters - ---------- - - species_name : str - Return the number of scraped particles of this species - - boundary : str - The boundary from which to get the scraped particle data in the - form x/y/z_hi/lo - - local : bool - Whether to only return the number of particles in the current - processor's buffer - ''' - - return self.libwarpx_so.warpx_getParticleBoundaryBufferSize( - ctypes.c_char_p(species_name.encode('utf-8')), - self._get_boundary_number(boundary), local - ) - - def get_particle_boundary_buffer_structs(self, species_name, boundary, level): - ''' - This returns a list of numpy arrays containing the particle struct data - for a species that has been scraped by a specific simulation boundary. The - particle data is represented as a structured numpy array and contains the - particle 'x', 'y', 'z', and 'idcpu'. - - The data for the numpy arrays are not copied, but share the underlying - memory buffer with WarpX. The numpy arrays are fully writeable. - - Parameters - ---------- - - species_name : str - The species name that the data will be returned for - - boundary : str - The boundary from which to get the scraped particle data in the - form x/y/z_hi/lo or eb. - - level : int - Which AMR level to retrieve scraped particle data from. - ''' - - particles_per_tile = _LP_c_int() - num_tiles = ctypes.c_int(0) - data = self.libwarpx_so.warpx_getParticleBoundaryBufferStructs( - ctypes.c_char_p(species_name.encode('utf-8')), - self._get_boundary_number(boundary), level, - ctypes.byref(num_tiles), ctypes.byref(particles_per_tile) - ) - - particle_data = [] - for i in range(num_tiles.value): - if particles_per_tile[i] == 0: - continue - arr = self._array1d_from_pointer(data[i], self._p_dtype, particles_per_tile[i]) - particle_data.append(arr) - - _libc.free(particles_per_tile) - _libc.free(data) - return particle_data - - def get_particle_boundary_buffer(self, species_name, boundary, comp_name, level): - ''' - This returns a list of numpy arrays containing the particle array data - for a species that has been scraped by a specific simulation boundary. - - The data for the numpy arrays are not copied, but share the underlying - memory buffer with WarpX. The numpy arrays are fully writeable. - - Parameters - ---------- - - species_name : str - The species name that the data will be returned for. - - boundary : str - The boundary from which to get the scraped particle data in the - form x/y/z_hi/lo or eb. - - comp_name : str - The component of the array data that will be returned. If - "step_scraped" the special attribute holding the timestep at - which a particle was scraped will be returned. - - level : int - Which AMR level to retrieve scraped particle data from. - ''' - - particles_per_tile = _LP_c_int() - num_tiles = ctypes.c_int(0) - if comp_name == 'step_scraped': - data = self.libwarpx_so.warpx_getParticleBoundaryBufferScrapedSteps( - ctypes.c_char_p(species_name.encode('utf-8')), - self._get_boundary_number(boundary), level, - ctypes.byref(num_tiles), ctypes.byref(particles_per_tile) - ) - else: - data = self.libwarpx_so.warpx_getParticleBoundaryBuffer( - ctypes.c_char_p(species_name.encode('utf-8')), - self._get_boundary_number(boundary), level, - ctypes.byref(num_tiles), ctypes.byref(particles_per_tile), - ctypes.c_char_p(comp_name.encode('utf-8')) - ) - - particle_data = [] - for i in range(num_tiles.value): - if particles_per_tile[i] == 0: - continue - if not data[i]: - raise Exception(f'get_particle_arrays: data[i] for i={i} was not initialized') - arr = np.ctypeslib.as_array(data[i], (particles_per_tile[i],)) - try: - # This fails on some versions of numpy - arr.setflags(write=1) - except ValueError: - pass - particle_data.append(arr) - - _libc.free(particles_per_tile) - _libc.free(data) - return particle_data - - def clearParticleBoundaryBuffer(self): - ''' - - Clear the buffer that holds the particles lost at the boundaries. - - ''' - self.libwarpx_so.warpx_clearParticleBoundaryBuffer() def depositChargeDensity(self, species_name, level, clear_rho=True, sync_rho=True): ''' @@ -1210,15 +344,30 @@ class LibWarpX(): sync_rho : bool If True, perform MPI exchange and properly set boundary cells for rho_fp. ''' + rho_fp = self.warpx.multifab(f'rho_fp[level={level}]') + + if rho_fp is None: + raise RuntimeError("Multifab `rho_fp` is not allocated.") + # ablastr::warn_manager::WMRecordWarning( + # "WarpXWrappers", "rho_fp is not allocated", + # ablastr::warn_manager::WarnPriority::low + # ); + # return if clear_rho: - from . import fields - fields.RhoFPWrapper(level, True)[...] = 0.0 - self.libwarpx_so.warpx_depositChargeDensity( - ctypes.c_char_p(species_name.encode('utf-8')), level - ) + rho_fp.set_val(0.0) + + # deposit the charge density from the desired species + mypc = self.warpx.multi_particle_container() + myspc = mypc.get_particle_container_from_name(species_name) + myspc.deposit_charge(rho_fp, level) + + if self.geometry_dim == 'rz': + self.warpx.apply_inverse_volume_scaling_to_charge_density(rho_fp, level) + if sync_rho: - self.libwarpx_so.warpx_SyncRho() + self.warpx.sync_rho() + def set_potential_EB(self, potential): """ @@ -1230,1696 +379,7 @@ class LibWarpX(): potential : str The expression string """ - self.libwarpx_so.warpx_setPotentialEB(ctypes.c_char_p(potential.encode('utf-8'))) - - def set_plasma_lens_strength( self, i_lens, strength_E, strength_B ): - """ - Set the strength of the `i_lens`-th lens - - Parameters - ---------- - i_lens: int - Index of the lens to be modified - - strength_E, strength_B: floats - The electric and magnetic focusing strength of the lens - """ - if self._numpy_real_dtype == 'f8': - c_real = ctypes.c_double - else: - c_real = ctypes.c_float - - self.libwarpx_so.warpx_setPlasmaLensStrength( - ctypes.c_int(i_lens), c_real(strength_E), c_real(strength_B) ) - - - def _get_mesh_field_list(self, warpx_func, level, direction, include_ghosts): - """ - Generic routine to fetch the list of field data arrays. - """ - shapes = _LP_c_int() - size = ctypes.c_int(0) - ncomps = ctypes.c_int(0) - ngrowvect = _LP_c_int() - if direction is None: - data = warpx_func(level, - ctypes.byref(size), ctypes.byref(ncomps), - ctypes.byref(ngrowvect), ctypes.byref(shapes)) - else: - data = warpx_func(level, direction, - ctypes.byref(size), ctypes.byref(ncomps), - ctypes.byref(ngrowvect), ctypes.byref(shapes)) - if not data: - raise Exception('object was not initialized') - - ngvect = [ngrowvect[i] for i in range(self.dim)] - grid_data = [] - shapesize = self.dim - if ncomps.value > 1: - shapesize += 1 - for i in range(size.value): - shape = tuple([shapes[shapesize*i + d] for d in range(shapesize)]) - # --- The data is stored in Fortran order, hence shape is reversed and a transpose is taken. - if shape[::-1] == 0: - continue - if not data[i]: - raise Exception(f'get_particle_arrays: data[i] for i={i} was not initialized') - arr = np.ctypeslib.as_array(data[i], shape[::-1]).T - try: - # This fails on some versions of numpy - arr.setflags(write=1) - except ValueError: - pass - if include_ghosts: - grid_data.append(arr) - else: - grid_data.append(arr[tuple([slice(ngvect[d], -ngvect[d]) for d in range(self.dim)])]) - - _libc.free(shapes) - _libc.free(data) - return grid_data - - def get_mesh_electric_field(self, level, direction, include_ghosts=True): - ''' - - This returns a list of numpy arrays containing the mesh electric field - data on each grid for this process. - - This version is for the full "auxiliary" solution on the given level. - - The data for the numpy arrays are not copied, but share the underlying - memory buffer with WarpX. The numpy arrays are fully writeable. - - Parameters - ---------- - - level : the AMR level to get the data for - direction : the component of the data you want - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A List of numpy arrays. - - ''' - - return self._get_mesh_field_list(self.libwarpx_so.warpx_getEfield, level, direction, include_ghosts) - - def get_mesh_electric_field_cp(self, level, direction, include_ghosts=True): - ''' - - This returns a list of numpy arrays containing the mesh electric field - data on each grid for this process. This version returns the field on - the coarse patch for the given level. - - The data for the numpy arrays are not copied, but share the underlying - memory buffer with WarpX. The numpy arrays are fully writeable. - - Parameters - ---------- - - level : the AMR level to get the data for - direction : the component of the data you want - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A List of numpy arrays. - - ''' - - return self._get_mesh_field_list(self.libwarpx_so.warpx_getEfieldCP, level, direction, include_ghosts) - - def get_mesh_electric_field_fp(self, level, direction, include_ghosts=True): - ''' - - This returns a list of numpy arrays containing the mesh electric field - data on each grid for this process. This version returns the field on - the fine patch for the given level. - - The data for the numpy arrays are not copied, but share the underlying - memory buffer with WarpX. The numpy arrays are fully writeable. - - Parameters - ---------- - - level : the AMR level to get the data for - direction : the component of the data you want - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A List of numpy arrays. - - ''' - - return self._get_mesh_field_list(self.libwarpx_so.warpx_getEfieldFP, level, direction, include_ghosts) - - def get_mesh_electric_field_cp_pml(self, level, direction, include_ghosts=True): - ''' - - This returns a list of numpy arrays containing the mesh electric field - data on each grid for this process. This version returns the field on - the coarse patch for the PML for the given level. - - The data for the numpy arrays are not copied, but share the underlying - memory buffer with WarpX. The numpy arrays are fully writeable. - - Parameters - ---------- - - level : the AMR level to get the data for - direction : the component of the data you want - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A List of numpy arrays. - - ''' - - try: - return self._get_mesh_field_list(self.libwarpx_so.warpx_getEfieldCP_PML, level, direction, include_ghosts) - except ValueError: - raise Exception('PML not initialized') - - def get_mesh_electric_field_fp_pml(self, level, direction, include_ghosts=True): - ''' - - This returns a list of numpy arrays containing the mesh electric field - data on each grid for this process. This version returns the field on - the fine patch for the PML for the given level. - - The data for the numpy arrays are not copied, but share the underlying - memory buffer with WarpX. The numpy arrays are fully writeable. - - Parameters - ---------- - - level : the AMR level to get the data for - direction : the component of the data you want - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A List of numpy arrays. - - ''' - - try: - return self._get_mesh_field_list(self.libwarpx_so.warpx_getEfieldFP_PML, level, direction, include_ghosts) - except ValueError: - raise Exception('PML not initialized') - - def get_mesh_magnetic_field(self, level, direction, include_ghosts=True): - ''' - - This returns a list of numpy arrays containing the mesh magnetic field - data on each grid for this process. - - This version is for the full "auxiliary" solution on the given level. - - The data for the numpy arrays are not copied, but share the underlying - memory buffer with WarpX. The numpy arrays are fully writeable. - - Parameters - ---------- - - level : the AMR level to get the data for - direction : the component of the data you want - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A List of numpy arrays. - - ''' - - return self._get_mesh_field_list(self.libwarpx_so.warpx_getBfield, level, direction, include_ghosts) - - def get_mesh_magnetic_field_cp(self, level, direction, include_ghosts=True): - ''' - - This returns a list of numpy arrays containing the mesh magnetic field - data on each grid for this process. This version returns the field on - the coarse patch for the given level. - - The data for the numpy arrays are not copied, but share the underlying - memory buffer with WarpX. The numpy arrays are fully writeable. - - Parameters - ---------- - - level : the AMR level to get the data for - direction : the component of the data you want - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A List of numpy arrays. - - ''' - - return self._get_mesh_field_list(self.libwarpx_so.warpx_getBfieldCP, level, direction, include_ghosts) - - def get_mesh_magnetic_field_fp(self, level, direction, include_ghosts=True): - ''' - - This returns a list of numpy arrays containing the mesh magnetic field - data on each grid for this process. This version returns the field on - the fine patch for the given level. - - The data for the numpy arrays are not copied, but share the underlying - memory buffer with WarpX. The numpy arrays are fully writeable. - - Parameters - ---------- - - level : the AMR level to get the data for - direction : the component of the data you want - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A List of numpy arrays. - - ''' - - return self._get_mesh_field_list(self.libwarpx_so.warpx_getBfieldFP, level, direction, include_ghosts) - - def get_mesh_vector_potential_fp(self, level, direction, include_ghosts=True): - ''' - - This returns a list of numpy arrays containing the mesh vector potential - data on each grid for this process. This version returns the field on - the fine patch for the given level. - - The data for the numpy arrays are not copied, but share the underlying - memory buffer with WarpX. The numpy arrays are fully writeable. - - Parameters - ---------- - - level : the AMR level to get the data for - direction : the component of the data you want - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A List of numpy arrays. - - ''' - - return self._get_mesh_field_list(self.libwarpx_so.warpx_getVectorPotentialFP, level, direction, include_ghosts) - - def get_mesh_magnetic_field_cp_pml(self, level, direction, include_ghosts=True): - ''' - - This returns a list of numpy arrays containing the mesh magnetic field - data on each grid for this process. This version returns the field on - the coarse patch for the PML for the given level. - - The data for the numpy arrays are not copied, but share the underlying - memory buffer with WarpX. The numpy arrays are fully writeable. - - Parameters - ---------- - - level : the AMR level to get the data for - direction : the component of the data you want - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A List of numpy arrays. - - ''' - - try: - return self._get_mesh_field_list(self.libwarpx_so.warpx_getBfieldCP_PML, level, direction, include_ghosts) - except ValueError: - raise Exception('PML not initialized') - - def get_mesh_magnetic_field_fp_pml(self, level, direction, include_ghosts=True): - ''' - - This returns a list of numpy arrays containing the mesh magnetic field - data on each grid for this process. This version returns the field on - the fine patch for the PML for the given level. - - The data for the numpy arrays are not copied, but share the underlying - memory buffer with WarpX. The numpy arrays are fully writeable. - - Parameters - ---------- - - level : the AMR level to get the data for - direction : the component of the data you want - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A List of numpy arrays. - - ''' - - try: - return self._get_mesh_field_list(self.libwarpx_so.warpx_getBfieldFP_PML, level, direction, include_ghosts) - except ValueError: - raise Exception('PML not initialized') - - def get_mesh_current_density(self, level, direction, include_ghosts=True): - ''' - - This returns a list of numpy arrays containing the mesh current density - data on each grid for this process. - - The data for the numpy arrays are not copied, but share the underlying - memory buffer with WarpX. The numpy arrays are fully writeable. - - Parameters - ---------- - - level : the AMR level to get the data for - direction : the component of the data you want - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A List of numpy arrays. - - ''' - - return self._get_mesh_field_list(self.libwarpx_so.warpx_getCurrentDensity, level, direction, include_ghosts) - - def get_mesh_current_density_cp(self, level, direction, include_ghosts=True): - ''' - - This returns a list of numpy arrays containing the mesh current density - data on each grid for this process. This version returns the density for - the coarse patch on the given level. - - The data for the numpy arrays are not copied, but share the underlying - memory buffer with WarpX. The numpy arrays are fully writeable. - - Parameters - ---------- - - level : the AMR level to get the data for - direction : the component of the data you want - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A List of numpy arrays. - - ''' - - return self._get_mesh_field_list(self.libwarpx_so.warpx_getCurrentDensityCP, level, direction, include_ghosts) - - def get_mesh_current_density_fp(self, level, direction, include_ghosts=True): - ''' - - This returns a list of numpy arrays containing the mesh current density - data on each grid for this process. This version returns the density on - the fine patch for the given level. - - The data for the numpy arrays are not copied, but share the underlying - memory buffer with WarpX. The numpy arrays are fully writeable. - - Parameters - ---------- - - level : the AMR level to get the data for - direction : the component of the data you want - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A List of numpy arrays. - - ''' - - return self._get_mesh_field_list(self.libwarpx_so.warpx_getCurrentDensityFP, level, direction, include_ghosts) - - def get_mesh_current_density_cp_pml(self, level, direction, include_ghosts=True): - ''' - - This returns a list of numpy arrays containing the mesh current density - data on each grid for this process. This version returns the density for - the coarse patch for the PML for the given level. - - The data for the numpy arrays are not copied, but share the underlying - memory buffer with WarpX. The numpy arrays are fully writeable. - - Parameters - ---------- - - level : the AMR level to get the data for - direction : the component of the data you want - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A List of numpy arrays. - - ''' - - try: - return self._get_mesh_field_list(self.libwarpx_so.warpx_getCurrentDensityCP_PML, level, direction, include_ghosts) - except ValueError: - raise Exception('PML not initialized') - - def get_mesh_current_density_fp_pml(self, level, direction, include_ghosts=True): - ''' - - This returns a list of numpy arrays containing the mesh current density - data on each grid for this process. This version returns the density on - the fine patch for the PML for the given level. - - The data for the numpy arrays are not copied, but share the underlying - memory buffer with WarpX. The numpy arrays are fully writeable. - - Parameters - ---------- - - level : the AMR level to get the data for - direction : the component of the data you want - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A List of numpy arrays. - - ''' - - try: - return self._get_mesh_field_list(self.libwarpx_so.warpx_getCurrentDensityFP_PML, level, direction, include_ghosts) - except ValueError: - raise Exception('PML not initialized') - - def get_mesh_current_density_fp_ampere(self, level, direction, include_ghosts=True): - ''' - - This returns a list of numpy arrays containing the mesh current density - data on each grid for this process calculated from the curl of B. This - quantity is calculated in the kinetic-fluid hybrid model to get the - electron current. This function returns the current density on the fine - patch for the given level. - - The data for the numpy arrays are not copied, but share the underlying - memory buffer with WarpX. The numpy arrays are fully writeable. - - Parameters - ---------- - - level : the AMR level to get the data for - direction : the component of the data you want - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A List of numpy arrays. - - ''' - - try: - return self._get_mesh_field_list(self.libwarpx_so.warpx_getCurrentDensityFP_Ampere, level, direction, include_ghosts) - except ValueError: - raise Exception('Current multifab not allocated.') - - def get_mesh_charge_density_cp(self, level, include_ghosts=True): - ''' - - This returns a list of numpy arrays containing the mesh charge density - data on each grid for this process. This version returns the density for - the coarse patch on the given level. - - The data for the numpy arrays are not copied, but share the underlying - memory buffer with WarpX. The numpy arrays are fully writeable. - - Parameters - ---------- - - level : the AMR level to get the data for - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A List of numpy arrays. - - ''' - - return self._get_mesh_field_list(self.libwarpx_so.warpx_getChargeDensityCP, level, None, include_ghosts) - - def get_mesh_charge_density_fp(self, level, include_ghosts=True): - ''' - - This returns a list of numpy arrays containing the mesh charge density - data on each grid for this process. This version returns the density on - the fine patch for the given level. - - The data for the numpy arrays are not copied, but share the underlying - memory buffer with WarpX. The numpy arrays are fully writeable. - - Parameters - ---------- - - level : the AMR level to get the data for - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A List of numpy arrays. - - ''' - - return self._get_mesh_field_list(self.libwarpx_so.warpx_getChargeDensityFP, level, None, include_ghosts) - - def get_mesh_phi_fp(self, level, include_ghosts=True): - ''' - - This returns a list of numpy arrays containing the mesh electrostatic - potential data on each grid for this process. This version returns the - potential on the fine patch for the given level. - - The data for the numpy arrays are not copied, but share the underlying - memory buffer with WarpX. The numpy arrays are fully writeable. - - Parameters - ---------- - - level : the AMR level to get the data for - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A List of numpy arrays. - - ''' - return self._get_mesh_field_list(self.libwarpx_so.warpx_getPhiFP, level, None, include_ghosts) - - def get_mesh_F_cp(self, level, include_ghosts=True): - ''' - - This returns a list of numpy arrays containing the mesh F field - data on each grid for this process. This version returns the F field for - the coarse patch on the given level. - - The data for the numpy arrays are not copied, but share the underlying - memory buffer with WarpX. The numpy arrays are fully writeable. - - Parameters - ---------- - - level : the AMR level to get the data for - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A List of numpy arrays. - - ''' - - return self._get_mesh_field_list(self.libwarpx_so.warpx_getFfieldCP, level, None, include_ghosts) - - def get_mesh_F_fp(self, level, include_ghosts=True): - ''' - - This returns a list of numpy arrays containing the mesh F field - data on each grid for this process. This version returns the F field on - the fine patch for the given level. - - The data for the numpy arrays are not copied, but share the underlying - memory buffer with WarpX. The numpy arrays are fully writeable. - - Parameters - ---------- - - level : the AMR level to get the data for - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A List of numpy arrays. - - ''' - - return self._get_mesh_field_list(self.libwarpx_so.warpx_getFfieldFP, level, None, include_ghosts) - - def get_mesh_F_fp_pml(self, level, include_ghosts=True): - ''' - - This returns a list of numpy arrays containing the mesh F field - data on each grid for this process. This version returns the F field on - the fine patch for the PML on the given level. - - The data for the numpy arrays are not copied, but share the underlying - memory buffer with WarpX. The numpy arrays are fully writeable. - - Parameters - ---------- - - level : the AMR level to get the data for - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A List of numpy arrays. - - ''' - try: - return self._get_mesh_field_list(self.libwarpx_so.warpx_getFfieldFP_PML, level, None, include_ghosts) - except ValueError: - raise Exception('PML not initialized') - - def get_mesh_F_cp_pml(self, level, include_ghosts=True): - ''' - - This returns a list of numpy arrays containing the mesh F field - data on each grid for this process. This version returns the F field on - the coarse patch for the PML on the given level. - - The data for the numpy arrays are not copied, but share the underlying - memory buffer with WarpX. The numpy arrays are fully writeable. - - Parameters - ---------- - - level : the AMR level to get the data for - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A List of numpy arrays. - - ''' - try: - return self._get_mesh_field_list(self.libwarpx_so.warpx_getFfieldCP_PML, level, None, include_ghosts) - except ValueError: - raise Exception('PML not initialized') - - def get_mesh_G_cp(self, level, include_ghosts=True): - ''' - - This returns a list of numpy arrays containing the mesh G field - data on each grid for this process. This version returns the G field for - the coarse patch on the given level. - - The data for the numpy arrays are not copied, but share the underlying - memory buffer with WarpX. The numpy arrays are fully writeable. - - Parameters - ---------- - - level : the AMR level to get the data for - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A List of numpy arrays. - - ''' - - return self._get_mesh_field_list(self.libwarpx_so.warpx_getGfieldCP, level, None, include_ghosts) - - def get_mesh_G_fp(self, level, include_ghosts=True): - ''' - - This returns a list of numpy arrays containing the mesh G field - data on each grid for this process. This version returns the G field on - the fine patch for the given level. - - The data for the numpy arrays are not copied, but share the underlying - memory buffer with WarpX. The numpy arrays are fully writeable. - - Parameters - ---------- - - level : the AMR level to get the data for - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A List of numpy arrays. - - ''' - - return self._get_mesh_field_list(self.libwarpx_so.warpx_getGfieldFP, level, None, include_ghosts) - - def get_mesh_G_cp_pml(self, level, include_ghosts=True): - ''' - - This returns a list of numpy arrays containing the mesh G field - data on each grid for this process. This version returns the G field on - the coarse patch for the PML on the given level. - - The data for the numpy arrays are not copied, but share the underlying - memory buffer with WarpX. The numpy arrays are fully writeable. - - Parameters - ---------- - - level : the AMR level to get the data for - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A List of numpy arrays. - - ''' - try: - return self._get_mesh_field_list(self.libwarpx_so.warpx_getGfieldCP_PML, level, None, include_ghosts) - except ValueError: - raise Exception('PML not initialized') - - def get_mesh_G_fp_pml(self, level, include_ghosts=True): - ''' - - This returns a list of numpy arrays containing the mesh G field - data on each grid for this process. This version returns the G field on - the fine patch for the PML on the given level. - - The data for the numpy arrays are not copied, but share the underlying - memory buffer with WarpX. The numpy arrays are fully writeable. - - Parameters - ---------- - - level : the AMR level to get the data for - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A List of numpy arrays. - - ''' - try: - return self._get_mesh_field_list(self.libwarpx_so.warpx_getGfieldFP_PML, level, None, include_ghosts) - except ValueError: - raise Exception('PML not initialized') - - def get_mesh_edge_lengths(self, level, direction, include_ghosts=True): - ''' - - This returns a list of numpy arrays containing the mesh edge lengths - data on each grid for this process. This version returns the density on - the fine patch for the given level. - - The data for the numpy arrays are not copied, but share the underlying - memory buffer with WarpX. The numpy arrays are fully writeable. - - Parameters - ---------- - - level : the AMR level to get the data for - direction : the component of the data you want - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A List of numpy arrays. - - ''' - - return self._get_mesh_field_list(self.libwarpx_so.warpx_getEdgeLengths, level, direction, include_ghosts) - - def get_mesh_face_areas(self, level, direction, include_ghosts=True): - ''' - - This returns a list of numpy arrays containing the mesh face areas - data on each grid for this process. This version returns the density on - the fine patch for the given level. - - The data for the numpy arrays are not copied, but share the underlying - memory buffer with WarpX. The numpy arrays are fully writeable. - - Parameters - ---------- - - level : the AMR level to get the data for - direction : the component of the data you want - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A List of numpy arrays. - - ''' - - return self._get_mesh_field_list(self.libwarpx_so.warpx_getFaceAreas, level, direction, include_ghosts) - - def _get_mesh_array_lovects(self, level, direction, include_ghosts=True, getlovectsfunc=None): - assert(0 <= level and level <= self.libwarpx_so.warpx_finestLevel()) - - size = ctypes.c_int(0) - ngrowvect = _LP_c_int() - if direction is None: - data = getlovectsfunc(level, ctypes.byref(size), ctypes.byref(ngrowvect)) - else: - data = getlovectsfunc(level, direction, ctypes.byref(size), ctypes.byref(ngrowvect)) - - if not data: - raise Exception('object was not initialized') - - lovects_ref = np.ctypeslib.as_array(data, (size.value, self.dim)) - - # --- Make a copy of the data to avoid memory problems - # --- Also, take the transpose to give shape (dims, number of grids) - lovects = lovects_ref.copy().T - - ng = [] - if include_ghosts: - for d in range(self.dim): - ng.append(ngrowvect[d]) - else: - for d in range(self.dim): - ng.append(0) - lovects[d,:] += ngrowvect[d] - - del lovects_ref - _libc.free(data) - return lovects, ng - - def get_mesh_electric_field_lovects(self, level, direction, include_ghosts=True): - ''' - - This returns a list of the lo vectors of the arrays containing the mesh electric field - data on each grid for this process. - - This version is for the full "auxiliary" solution on the given level. - - Parameters - ---------- - - level : the AMR level to get the data for - direction : the component of the data you want - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A 2d numpy array of the lo vector for each grid with the shape (dims, number of grids) - - ''' - return self._get_mesh_array_lovects(level, direction, include_ghosts, self.libwarpx_so.warpx_getEfieldLoVects) - - def get_mesh_electric_field_cp_lovects(self, level, direction, include_ghosts=True): - ''' - - This returns a list of the lo vectors of the arrays containing the mesh electric field - data on each grid for this process. - - Parameters - ---------- - - level : the AMR level to get the data for - direction : the component of the data you want - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A 2d numpy array of the lo vector for each grid with the shape (dims, number of grids) - - ''' - return self._get_mesh_array_lovects(level, direction, include_ghosts, self.libwarpx_so.warpx_getEfieldCPLoVects) - - def get_mesh_electric_field_fp_lovects(self, level, direction, include_ghosts=True): - ''' - - This returns a list of the lo vectors of the arrays containing the mesh electric field - data on each grid for this process. - - Parameters - ---------- - - level : the AMR level to get the data for - direction : the component of the data you want - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A 2d numpy array of the lo vector for each grid with the shape (dims, number of grids) - - ''' - return self._get_mesh_array_lovects(level, direction, include_ghosts, self.libwarpx_so.warpx_getEfieldFPLoVects) - - def get_mesh_electric_field_cp_lovects_pml(self, level, direction, include_ghosts=True): - ''' - - This returns a list of the lo vectors of the arrays containing the mesh electric field - data on each PML grid for this process. - - Parameters - ---------- - - level : the AMR level to get the data for - direction : the component of the data you want - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A 2d numpy array of the lo vector for each grid with the shape (dims, number of grids) - - ''' - try: - return self._get_mesh_array_lovects(level, direction, include_ghosts, self.libwarpx_so.warpx_getEfieldCPLoVects_PML) - except ValueError: - raise Exception('PML not initialized') - - def get_mesh_electric_field_fp_lovects_pml(self, level, direction, include_ghosts=True): - ''' - - This returns a list of the lo vectors of the arrays containing the mesh electric field - data on each PML grid for this process. - - Parameters - ---------- - - level : the AMR level to get the data for - direction : the component of the data you want - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A 2d numpy array of the lo vector for each grid with the shape (dims, number of grids) - - ''' - try: - return self._get_mesh_array_lovects(level, direction, include_ghosts, self.libwarpx_so.warpx_getEfieldFPLoVects_PML) - except ValueError: - raise Exception('PML not initialized') - - def get_mesh_magnetic_field_lovects(self, level, direction, include_ghosts=True): - ''' - - This returns a list of the lo vectors of the arrays containing the mesh electric field - data on each grid for this process. - - This version is for the full "auxiliary" solution on the given level. - - Parameters - ---------- - - level : the AMR level to get the data for - direction : the component of the data you want - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A 2d numpy array of the lo vector for each grid with the shape (dims, number of grids) - - ''' - return self._get_mesh_array_lovects(level, direction, include_ghosts, self.libwarpx_so.warpx_getBfieldLoVects) - - def get_mesh_magnetic_field_cp_lovects(self, level, direction, include_ghosts=True): - ''' - - This returns a list of the lo vectors of the arrays containing the mesh electric field - data on each grid for this process. - - Parameters - ---------- - - level : the AMR level to get the data for - direction : the component of the data you want - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A 2d numpy array of the lo vector for each grid with the shape (dims, number of grids) - - ''' - return self._get_mesh_array_lovects(level, direction, include_ghosts, self.libwarpx_so.warpx_getBfieldCPLoVects) - - def get_mesh_magnetic_field_fp_lovects(self, level, direction, include_ghosts=True): - ''' - - This returns a list of the lo vectors of the arrays containing the mesh electric field - data on each grid for this process. - - Parameters - ---------- - - level : the AMR level to get the data for - direction : the component of the data you want - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A 2d numpy array of the lo vector for each grid with the shape (dims, number of grids) - - ''' - return self._get_mesh_array_lovects(level, direction, include_ghosts, self.libwarpx_so.warpx_getBfieldFPLoVects) - - def get_mesh_vector_potential_fp_lovects(self, level, direction, include_ghosts=True): - ''' - - This returns a list of the lo vectors of the arrays containing the mesh vector potential field - data on each grid for this process. - - Parameters - ---------- - - level : the AMR level to get the data for - direction : the component of the data you want - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A 2d numpy array of the lo vector for each grid with the shape (dims, number of grids) - - ''' - return self._get_mesh_array_lovects(level, direction, include_ghosts, self.libwarpx_so.warpx_getVectorPotentialFPLoVects) - - def get_mesh_magnetic_field_cp_lovects_pml(self, level, direction, include_ghosts=True): - ''' - - This returns a list of the lo vectors of the arrays containing the mesh electric field - data on each PML grid for this process. - - Parameters - ---------- - - level : the AMR level to get the data for - direction : the component of the data you want - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A 2d numpy array of the lo vector for each grid with the shape (dims, number of grids) - - ''' - try: - return self._get_mesh_array_lovects(level, direction, include_ghosts, self.libwarpx_so.warpx_getBfieldCPLoVects_PML) - except ValueError: - raise Exception('PML not initialized') - - def get_mesh_magnetic_field_fp_lovects_pml(self, level, direction, include_ghosts=True): - ''' - - This returns a list of the lo vectors of the arrays containing the mesh electric field - data on each PML grid for this process. - - Parameters - ---------- - - level : the AMR level to get the data for - direction : the component of the data you want - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A 2d numpy array of the lo vector for each grid with the shape (dims, number of grids) - - ''' - try: - return self._get_mesh_array_lovects(level, direction, include_ghosts, self.libwarpx_so.warpx_getBfieldFPLoVects_PML) - except ValueError: - raise Exception('PML not initialized') - - def get_mesh_current_density_lovects(self, level, direction, include_ghosts=True): - ''' - - This returns a list of the lo vectors of the arrays containing the mesh electric field - data on each grid for this process. - - Parameters - ---------- - - level : the AMR level to get the data for - direction : the component of the data you want - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A 2d numpy array of the lo vector for each grid with the shape (dims, number of grids) - - ''' - return self._get_mesh_array_lovects(level, direction, include_ghosts, self.libwarpx_so.warpx_getCurrentDensityLoVects) - - def get_mesh_current_density_cp_lovects(self, level, direction, include_ghosts=True): - ''' - - This returns a list of the lo vectors of the arrays containing the mesh electric field - data on each grid for this process. - - Parameters - ---------- - - level : the AMR level to get the data for - direction : the component of the data you want - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A 2d numpy array of the lo vector for each grid with the shape (dims, number of grids) - - ''' - return self._get_mesh_array_lovects(level, direction, include_ghosts, self.libwarpx_so.warpx_getCurrentDensityCPLoVects) - - def get_mesh_current_density_fp_lovects(self, level, direction, include_ghosts=True): - ''' - - This returns a list of the lo vectors of the arrays containing the mesh electric field - data on each grid for this process. - - Parameters - ---------- - - level : the AMR level to get the data for - direction : the component of the data you want - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A 2d numpy array of the lo vector for each grid with the shape (dims, number of grids) - - ''' - return self._get_mesh_array_lovects(level, direction, include_ghosts, self.libwarpx_so.warpx_getCurrentDensityFPLoVects) - - def get_mesh_current_density_cp_lovects_pml(self, level, direction, include_ghosts=True): - ''' - - This returns a list of the lo vectors of the arrays containing the mesh electric field - data on each PML grid for this process. - - Parameters - ---------- - - level : the AMR level to get the data for - direction : the component of the data you want - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A 2d numpy array of the lo vector for each grid with the shape (dims, number of grids) - - ''' - try: - return self._get_mesh_array_lovects(level, direction, include_ghosts, self.libwarpx_so.warpx_getCurrentDensityCPLoVects_PML) - except ValueError: - raise Exception('PML not initialized') - - def get_mesh_current_density_fp_lovects_pml(self, level, direction, include_ghosts=True): - ''' - - This returns a list of the lo vectors of the arrays containing the mesh electric field - data on each PML grid for this process. - - Parameters - ---------- - - level : the AMR level to get the data for - direction : the component of the data you want - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A 2d numpy array of the lo vector for each grid with the shape (dims, number of grids) - - ''' - try: - return self._get_mesh_array_lovects(level, direction, include_ghosts, self.libwarpx_so.warpx_getCurrentDensityFPLoVects_PML) - except ValueError: - raise Exception('PML not initialized') - - def get_mesh_charge_density_cp_lovects(self, level, include_ghosts=True): - ''' - - This returns a list of the lo vectors of the arrays containing the mesh electric field - data on each grid for this process. - - Parameters - ---------- - - level : the AMR level to get the data for - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A 2d numpy array of the lo vector for each grid with the shape (dims, number of grids) - - ''' - return self._get_mesh_array_lovects(level, None, include_ghosts, self.libwarpx_so.warpx_getChargeDensityCPLoVects) - - def get_mesh_charge_density_fp_lovects(self, level, include_ghosts=True): - ''' - - This returns a list of the lo vectors of the arrays containing the mesh - charge density data on each grid for this process. - - Parameters - ---------- - - level : the AMR level to get the data for - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A 2d numpy array of the lo vector for each grid with the shape (dims, number of grids) - - ''' - return self._get_mesh_array_lovects(level, None, include_ghosts, self.libwarpx_so.warpx_getChargeDensityFPLoVects) - - def get_mesh_phi_fp_lovects(self, level, include_ghosts=True): - ''' - - This returns a list of the lo vectors of the arrays containing the mesh - electrostatic potential data on each grid for this process. - - Parameters - ---------- - - level : the AMR level to get the data for - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A 2d numpy array of the lo vector for each grid with the shape (dims, number of grids) - - ''' - return self._get_mesh_array_lovects(level, None, include_ghosts, self.libwarpx_so.warpx_getPhiFPLoVects) - - def get_mesh_F_cp_lovects(self, level, include_ghosts=True): - ''' - - This returns a list of the lo vectors of the arrays containing the mesh F field - data on each grid for this process. - - Parameters - ---------- - - level : the AMR level to get the data for - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A 2d numpy array of the lo vector for each grid with the shape (dims, number of grids) - - ''' - return self._get_mesh_array_lovects(level, None, include_ghosts, self.libwarpx_so.warpx_getFfieldCPLoVects) - - def get_mesh_F_fp_lovects(self, level, include_ghosts=True): - ''' - - This returns a list of the lo vectors of the arrays containing the mesh F field - data on each grid for this process. - - Parameters - ---------- - - level : the AMR level to get the data for - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A 2d numpy array of the lo vector for each grid with the shape (dims, number of grids) - - ''' - return self._get_mesh_array_lovects(level, None, include_ghosts, self.libwarpx_so.warpx_getFfieldFPLoVects) - - def get_mesh_F_cp_lovects_pml(self, level, include_ghosts=True): - ''' - - This returns a list of the lo vectors of the arrays containing the mesh F field - data on each PML grid for this process. - - Parameters - ---------- - - level : the AMR level to get the data for - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A 2d numpy array of the lo vector for each grid with the shape (dims, number of grids) - - ''' - try: - return self._get_mesh_array_lovects(level, None, include_ghosts, self.libwarpx_so.warpx_getFfieldCPLoVects_PML) - except ValueError: - raise Exception('PML not initialized') - - def get_mesh_F_fp_lovects_pml(self, level, include_ghosts=True): - ''' - - This returns a list of the lo vectors of the arrays containing the mesh F field - data on each PML grid for this process. - - Parameters - ---------- - - level : the AMR level to get the data for - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A 2d numpy array of the lo vector for each grid with the shape (dims, number of grids) - - ''' - try: - return self._get_mesh_array_lovects(level, None, include_ghosts, self.libwarpx_so.warpx_getFfieldFPLoVects_PML) - except ValueError: - raise Exception('PML not initialized') - - def get_mesh_G_cp_lovects(self, level, include_ghosts=True): - ''' - - This returns a list of the lo vectors of the arrays containing the mesh G field - data on each grid for this process. - - Parameters - ---------- - - level : the AMR level to get the data for - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A 2d numpy array of the lo vector for each grid with the shape (dims, number of grids) - - ''' - return self._get_mesh_array_lovects(level, None, include_ghosts, self.libwarpx_so.warpx_getGfieldCPLoVects) - - def get_mesh_G_fp_lovects(self, level, include_ghosts=True): - ''' - - This returns a list of the lo vectors of the arrays containing the mesh G field - data on each grid for this process. - - Parameters - ---------- - - level : the AMR level to get the data for - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A 2d numpy array of the lo vector for each grid with the shape (dims, number of grids) - - ''' - return self._get_mesh_array_lovects(level, None, include_ghosts, self.libwarpx_so.warpx_getGfieldFPLoVects) - - def get_mesh_G_cp_lovects_pml(self, level, include_ghosts=True): - ''' - - This returns a list of the lo vectors of the arrays containing the mesh G field - data on each PML grid for this process. - - Parameters - ---------- - - level : the AMR level to get the data for - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A 2d numpy array of the lo vector for each grid with the shape (dims, number of grids) - - ''' - try: - return self._get_mesh_array_lovects(level, None, include_ghosts, self.libwarpx_so.warpx_getGfieldCPLoVects_PML) - except ValueError: - raise Exception('PML not initialized') - - def get_mesh_G_fp_lovects_pml(self, level, include_ghosts=True): - ''' - - This returns a list of the lo vectors of the arrays containing the mesh G field - data on each PML grid for this process. - - Parameters - ---------- - - level : the AMR level to get the data for - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A 2d numpy array of the lo vector for each grid with the shape (dims, number of grids) - - ''' - try: - return self._get_mesh_array_lovects(level, None, include_ghosts, self.libwarpx_so.warpx_getGfieldFPLoVects_PML) - except ValueError: - raise Exception('PML not initialized') - - def get_mesh_edge_lengths_lovects(self, level, direction, include_ghosts=True): - ''' - - This returns a list of the lo vectors of the arrays containing the mesh edge lengths - data on each grid for this process. - - Parameters - ---------- - - level : the AMR level to get the data for - direction : the component of the data you want - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A 2d numpy array of the lo vector for each grid with the shape (dims, number of grids) - - ''' - return self._get_mesh_array_lovects(level, direction, include_ghosts, self.libwarpx_so.warpx_getEdgeLengthsLoVects) - - def get_mesh_face_areas_lovects(self, level, direction, include_ghosts=True): - ''' - - This returns a list of the lo vectors of the arrays containing the mesh face areas - data on each grid for this process. - - Parameters - ---------- - - level : the AMR level to get the data for - direction : the component of the data you want - include_ghosts : whether to include ghost zones or not - - Returns - ------- - - A 2d numpy array of the lo vector for each grid with the shape (dims, number of grids) - - ''' - return self._get_mesh_array_lovects(level, direction, include_ghosts, self.libwarpx_so.warpx_getFaceAreasLoVects) - - def _get_nodal_flag(self, getdatafunc): - data = getdatafunc() - if not data: - raise Exception('object was not initialized') - - nodal_flag_ref = np.ctypeslib.as_array(data, (self.dim,)) - - # --- Make a copy of the data to avoid memory problems - nodal_flag = nodal_flag_ref.copy() - - del nodal_flag_ref - _libc.free(data) - return nodal_flag - - def get_Ex_nodal_flag(self): - ''' - This returns a 1d array of the nodal flags for Ex along each direction. A 1 means node centered, and 0 cell centered. - ''' - return self._get_nodal_flag(self.libwarpx_so.warpx_getEx_nodal_flag) - - def get_Ey_nodal_flag(self): - ''' - This returns a 1d array of the nodal flags for Ey along each direction. A 1 means node centered, and 0 cell centered. - ''' - return self._get_nodal_flag(self.libwarpx_so.warpx_getEy_nodal_flag) - - def get_Ez_nodal_flag(self): - ''' - This returns a 1d array of the nodal flags for Ez along each direction. A 1 means node centered, and 0 cell centered. - ''' - return self._get_nodal_flag(self.libwarpx_so.warpx_getEz_nodal_flag) - - def get_Bx_nodal_flag(self): - ''' - This returns a 1d array of the nodal flags for Bx along each direction. A 1 means node centered, and 0 cell centered. - ''' - return self._get_nodal_flag(self.libwarpx_so.warpx_getBx_nodal_flag) - - def get_By_nodal_flag(self): - ''' - This returns a 1d array of the nodal flags for By along each direction. A 1 means node centered, and 0 cell centered. - ''' - return self._get_nodal_flag(self.libwarpx_so.warpx_getBy_nodal_flag) - - def get_Bz_nodal_flag(self): - ''' - This returns a 1d array of the nodal flags for Bz along each direction. A 1 means node centered, and 0 cell centered. - ''' - return self._get_nodal_flag(self.libwarpx_so.warpx_getBz_nodal_flag) - - def get_Jx_nodal_flag(self): - ''' - This returns a 1d array of the nodal flags for Jx along each direction. A 1 means node centered, and 0 cell centered. - ''' - return self._get_nodal_flag(self.libwarpx_so.warpx_getJx_nodal_flag) - - def get_Jy_nodal_flag(self): - ''' - This returns a 1d array of the nodal flags for Jy along each direction. A 1 means node centered, and 0 cell centered. - ''' - return self._get_nodal_flag(self.libwarpx_so.warpx_getJy_nodal_flag) - - def get_Jz_nodal_flag(self): - ''' - This returns a 1d array of the nodal flags for Jz along each direction. A 1 means node centered, and 0 cell centered. - ''' - return self._get_nodal_flag(self.libwarpx_so.warpx_getJz_nodal_flag) - - def get_Ax_nodal_flag(self): - ''' - This returns a 1d array of the nodal flags for Ax along each direction. A 1 means node centered, and 0 cell centered. - ''' - return self._get_nodal_flag(self.libwarpx_so.warpx_getAx_nodal_flag) - - def get_Ay_nodal_flag(self): - ''' - This returns a 1d array of the nodal flags for Ay along each direction. A 1 means node centered, and 0 cell centered. - ''' - return self._get_nodal_flag(self.libwarpx_so.warpx_getAy_nodal_flag) - - def get_Az_nodal_flag(self): - ''' - This returns a 1d array of the nodal flags for Az along each direction. A 1 means node centered, and 0 cell centered. - ''' - return self._get_nodal_flag(self.libwarpx_so.warpx_getAz_nodal_flag) - - def get_Rho_nodal_flag(self): - ''' - This returns a 1d array of the nodal flags for Rho along each direction. A 1 means node centered, and 0 cell centered. - ''' - return self._get_nodal_flag(self.libwarpx_so.warpx_getRho_nodal_flag) - - def get_Phi_nodal_flag(self): - ''' - This returns a 1d array of the nodal flags for Phi along each direction. A 1 means node centered, and 0 cell centered. - ''' - return self._get_nodal_flag(self.libwarpx_so.warpx_getPhi_nodal_flag) - - def get_F_nodal_flag(self): - ''' - This returns a 1d array of the nodal flags for F along each direction. A 1 means node centered, and 0 cell centered. - ''' - return self._get_nodal_flag(self.libwarpx_so.warpx_getF_nodal_flag) - - def get_G_nodal_flag(self): - ''' - This returns a 1d array of the nodal flags for G along each direction. A 1 means node centered, and 0 cell centered. - ''' - return self._get_nodal_flag(self.libwarpx_so.warpx_getG_nodal_flag) - - def get_edge_lengths_x_nodal_flag(self): - ''' - This returns a 1d array of the nodal flags for the x edge lengths along each direction. A 1 means node centered, and 0 cell centered. - ''' - return self._get_nodal_flag(self.libwarpx_so.warpx_get_edge_lengths_x_nodal_flag) - - def get_edge_lengths_y_nodal_flag(self): - ''' - This returns a 1d array of the nodal flags for the y edge lengths along each direction. A 1 means node centered, and 0 cell centered. - ''' - return self._get_nodal_flag(self.libwarpx_so.warpx_get_edge_lengths_y_nodal_flag) - - def get_edge_lengths_z_nodal_flag(self): - ''' - This returns a 1d array of the nodal flags for the z edge lengths along each direction. A 1 means node centered, and 0 cell centered. - ''' - return self._get_nodal_flag(self.libwarpx_so.warpx_get_edge_lengths_z_nodal_flag) - - def get_face_areas_x_nodal_flag(self): - ''' - This returns a 1d array of the nodal flags for the x face areas along each direction. A 1 means node centered, and 0 cell centered. - ''' - return self._get_nodal_flag(self.libwarpx_so.warpx_get_face_areas_x_nodal_flag) - - def get_face_areas_y_nodal_flag(self): - ''' - This returns a 1d array of the nodal flags for the y face areas along each direction. A 1 means node centered, and 0 cell centered. - ''' - return self._get_nodal_flag(self.libwarpx_so.warpx_get_face_areas_y_nodal_flag) - - def get_face_areas_z_nodal_flag(self): - ''' - This returns a 1d array of the nodal flags for the z face areas along each direction. A 1 means node centered, and 0 cell centered. - ''' - return self._get_nodal_flag(self.libwarpx_so.warpx_get_face_areas_z_nodal_flag) - - def get_F_pml_nodal_flag(self): - ''' - This returns a 1d array of the nodal flags for F in the PML along each direction. A 1 means node centered, and 0 cell centered. - ''' - return self._get_nodal_flag(self.libwarpx_so.warpx_getF_pml_nodal_flag) - - def get_G_pml_nodal_flag(self): - ''' - This returns a 1d array of the nodal flags for G in the PML along each direction. A 1 means node centered, and 0 cell centered. - ''' - return self._get_nodal_flag(self.libwarpx_so.warpx_getG_pml_nodal_flag) + self.warpx.set_potential_on_eb(potential) libwarpx = LibWarpX() diff --git a/Python/pywarpx/callbacks.py b/Python/pywarpx/callbacks.py index 02a8f5aad..b8824d0ef 100644 --- a/Python/pywarpx/callbacks.py +++ b/Python/pywarpx/callbacks.py @@ -1,7 +1,9 @@ -# Copyright 2017 David Grote +# Copyright 2017-2022 The WarpX Community # # This file is part of WarpX. # +# Authors: David Grote, Roelof Groenewald +# # License: BSD-3-Clause-LBNL """call back operations @@ -10,41 +12,43 @@ These are the functions which allow installing user created functions so that they are called at various places along the time step. -For each call back type, the following three functions are defined. - - install___: Installs a function to be called at that specified time - - uninstall___: Uninstalls the function (so it won't be called anymore) - - isinstalled___: Checks if the function is installed +The following three functions allow the user to install, uninstall and verify +the different call back types. + - installcallback: Installs a function to be called at that specified time + - uninstallcallback: Uninstalls the function (so it won't be called anymore) + - isinstalled: Checks if the function is installed -These functions all take a function or instance method as an argument. Note that -if an instance method is used, an extra reference to the method's object is saved. +These functions all take a callback location name (string) and function or +instance method as an argument. Note that if an instance method is used, an +extra reference to the method's object is saved. -The install can be done using a decorator, which has the prefix "callfrom". See example below. +The install can be done using a decorator, which has the prefix "callfrom". See +example below. Functions can be called at the following times: - - beforeInitEsolve <installbeforeInitEsolve>: before the initial solve for the E fields (i.e. before the PIC loop starts) - - afterInitEsolve <installafterInitEsolve>: after the initial solve for the E fields (i.e. before the PIC loop starts) - - afterinit <installafterinit>: immediately after the init is complete - - beforeEsolve <installbeforeEsolve>: before the solve for E fields - - poissonsolver <installpoissonsolver>: In place of the computePhi call but only in an electrostatic simulation - - afterEsolve <installafterEsolve>: after the solve for E fields - - beforedeposition <installbeforedeposition>: before the particle deposition (for charge and/or current) - - afterdeposition <installafterdeposition>: after particle deposition (for charge and/or current) - - beforestep <installbeforestep>: before the time step - - afterstep <installafterstep>: after the time step - - afterdiagnostics <installafterdiagnostics>: after diagnostic output - - oncheckpointsignal <installoncheckpointsignal>: on a checkpoint signal - - onbreaksignal <installonbreaksignal>: on a break signal. These callbacks will be the last ones executed - before the simulation ends. - - particlescraper <installparticlescraper>: just after the particle boundary conditions are applied - but before lost particles are processed - - particleloader <installparticleloader>: at the time that the standard particle loader is called - - particleinjection <installparticleinjection>: called when particle injection happens, after the position - advance and before deposition is called, allowing a user defined - particle distribution to be injected each time step - - appliedfields <installappliedfields>: allows directly specifying any fields to be applied to the particles - during the advance - -To use a decorator, the syntax is as follows. This will install the function myplots to be called after each step. + - beforeInitEsolve: before the initial solve for the E fields (i.e. before the PIC loop starts) + - afterinit: immediately after the init is complete + - beforeEsolve: before the solve for E fields + - poissonsolver: In place of the computePhi call but only in an electrostatic simulation + - afterEsolve: after the solve for E fields + - beforedeposition: before the particle deposition (for charge and/or current) + - afterdeposition: after particle deposition (for charge and/or current) + - beforestep: before the time step + - afterstep: after the time step + - afterdiagnostics: after diagnostic output + - oncheckpointsignal: on a checkpoint signal + - onbreaksignal: on a break signal. These callbacks will be the last ones executed before the simulation ends. + - particlescraper: just after the particle boundary conditions are applied + but before lost particles are processed + - particleloader: at the time that the standard particle loader is called + - particleinjection: called when particle injection happens, after the position + advance and before deposition is called, allowing a user + defined particle distribution to be injected each time step + - appliedfields: allows directly specifying any fields to be applied to the particles + during the advance + +To use a decorator, the syntax is as follows. This will install the function +``myplots`` to be called after each step. @callfromafterstep def myplots(): @@ -54,13 +58,10 @@ This is equivalent to the following: def myplots(): ppzx() -installafterstep(myplots) +installcallback('afterstep', myplots) """ -from __future__ import generators - import copy -import ctypes import sys import time import types @@ -84,12 +85,13 @@ class CallbackFunctions(object): installed a method in one of the call back lists. """ - def __init__(self,name=None,lcallonce=0): + def __init__(self,name=None,lcallonce=0,singlefunconly=False): self.funcs = [] self.time = 0. self.timers = {} self.name = name self.lcallonce = lcallonce + self.singlefunconly = singlefunconly def __call__(self,*args,**kw): """Call all of the functions in the list""" @@ -100,9 +102,7 @@ class CallbackFunctions(object): def clearlist(self): """Unregister/clear out all registered C callbacks""" self.funcs = [] - libwarpx.libwarpx_so.warpx_clear_callback_py( - ctypes.c_char_p(self.name.encode('utf-8')) - ) + libwarpx.libwarpx_so.remove_python_callback(self.name) def __bool__(self): """Returns True if functions are installed, otherwise False""" @@ -158,15 +158,15 @@ class CallbackFunctions(object): def installfuncinlist(self,f): """Check if the specified function is installed""" + if self.singlefunconly and self.hasfuncsinstalled(): + raise RuntimeError( + f"Only one function can be installed for callback {self.name}." + ) + if len(self.funcs) == 0: # If this is the first function installed, set the callback in the C++ # to call this class instance. - # Note that the _c_func must be saved. - _CALLBACK_FUNC_0 = ctypes.CFUNCTYPE(None) - self._c_func = _CALLBACK_FUNC_0(self) - libwarpx.libwarpx_so.warpx_set_callback_py( - ctypes.c_char_p(self.name.encode('utf-8')), self._c_func - ) + libwarpx.libwarpx_so.add_python_callback(self.name, self) if isinstance(f,types.MethodType): # --- If the function is a method of a class instance, then save a full # --- reference to that instance and the method name. @@ -254,30 +254,51 @@ class CallbackFunctions(object): #============================================================================= +callback_instances = { + "beforeInitEsolve": {}, + "afterInitEsolve": {}, + "afterinit": {}, + "beforecollisions": {}, + "aftercollisions": {}, + "beforeEsolve": {}, + "poissonsolver": {'singlefunconly': True}, # external Poisson solver + "afterEsolve": {}, + "beforedeposition": {}, + "afterdeposition": {}, + "particlescraper": {}, + "particleloader": {}, + "beforestep": {}, + "afterstep": {}, + "afterdiagnostics": {}, + "afterrestart": {}, + "oncheckpointsignal": {}, + "onbreaksignal": {}, + "particleinjection": {}, + "appliedfields": {} +} + # --- Now create the actual instances. -_beforeInitEsolve = CallbackFunctions('beforeInitEsolve') -_afterInitEsolve = CallbackFunctions('afterInitEsolve') -_afterinit = CallbackFunctions('afterinit') -_beforecollisions = CallbackFunctions('beforecollisions') -_aftercollisions = CallbackFunctions('aftercollisions') -_beforeEsolve = CallbackFunctions('beforeEsolve') -_poissonsolver = CallbackFunctions('poissonsolver') -_afterEsolve = CallbackFunctions('afterEsolve') -_beforedeposition = CallbackFunctions('beforedeposition') -_afterdeposition = CallbackFunctions('afterdeposition') -_particlescraper = CallbackFunctions('particlescraper') -_particleloader = CallbackFunctions('particleloader') -_beforestep = CallbackFunctions('beforestep') -_afterstep = CallbackFunctions('afterstep') -_afterdiagnostics = CallbackFunctions('afterdiagnostics') -_afterrestart = CallbackFunctions('afterrestart',lcallonce=1) -_oncheckpointsignal = CallbackFunctions('oncheckpointsignal') -_onbreaksignal = CallbackFunctions('onbreaksignal') -_particleinjection = CallbackFunctions('particleinjection') -_appliedfields = CallbackFunctions('appliedfields') +for key, val in callback_instances.items(): + callback_instances[key] = CallbackFunctions(name=key, **val) + +def installcallback(name, f): + "Adds a function to the list of functions called by this callback" + callback_instances[name].installfuncinlist(f) + +def uninstallcallback(name, f): + "Removes the function from the list of functions called by this callback" + callback_instances[name].uninstallfuncinlist(f) +def isinstalled(name, f): + "Checks if the function is called by this callback" + return callback_instances[name].isinstalledfuncinlist(f) + +def clear_all(): + for key, val in callback_instances.items(): + val.clearlist() #============================================================================= + def printcallbacktimers(tmin=1.,lminmax=False,ff=None): """Prints timings of installed functions. - tmin=1.: only functions with time greater than tmin will be printed @@ -285,18 +306,7 @@ def printcallbacktimers(tmin=1.,lminmax=False,ff=None): - ff=None: If given, timings will be written to the file object instead of stdout """ if ff is None: ff = sys.stdout - for c in [ _beforeInitEsolve, _afterInitEsolve, - _afterinit,_beforeEsolve,_poissonsolver,_afterEsolve, - _beforedeposition,_afterdeposition, - _particlescraper, - _particleloader, - _beforestep,_afterstep, - _afterdiagnostics, - _afterrestart, - _oncheckpointsignal, - _onbreaksignal, - _particleinjection, - _appliedfields]: + for c in callback_instances.values(): for fname, time in c.timers.items(): #vlist = numpy.array(gather(time)) vlist = numpy.array([time]) @@ -316,299 +326,129 @@ def printcallbacktimers(tmin=1.,lminmax=False,ff=None): ff.write('\n') #============================================================================= + # ---------------------------------------------------------------------------- def callfrombeforeInitEsolve(f): - installbeforeInitEsolve(f) + installcallback('beforeInitEsolve', f) return f def installbeforeInitEsolve(f): - "Adds a function to the list of functions called before the initial E solve" - _beforeInitEsolve.installfuncinlist(f) -def uninstallbeforeInitEsolve(f): - "Removes the function from the list of functions called before the initial E solve" - _beforeInitEsolve.uninstallfuncinlist(f) -def isinstalledbeforeInitEsolve(f): - "Checks if the function is called before the initial E solve" - return _beforeInitEsolve.isinstalledfuncinlist(f) + installcallback('beforeInitEsolve', f) # ---------------------------------------------------------------------------- def callfromafterInitEsolve(f): - installafterInitEsolve(f) + installcallback('afterInitEsolve', f) return f def installafterInitEsolve(f): - "Adds a function to the list of functions called after the initial E solve" - _afterInitEsolve.installfuncinlist(f) -def uninstallafterInitEsolve(f): - "Removes the function from the list of functions called after the initial E solve" - _afterInitEsolve.uninstallfuncinlist(f) -def isinstalledafterInitEsolve(f): - "Checks if the function is called after the initial E solve" - return _afterInitEsolve.isinstalledfuncinlist(f) + installcallback('afterInitEsolve', f) # ---------------------------------------------------------------------------- def callfromafterinit(f): - installafterinit(f) + installcallback('afterinit', f) return f def installafterinit(f): - "Adds a function to the list of functions called after the init" - _afterinit.installfuncinlist(f) -def uninstallafterinit(f): - "Removes the function from the list of functions called after the init" - _afterinit.uninstallfuncinlist(f) -def isinstalledafterinit(f): - "Checks if the function is called after a init" - return _afterinit.isinstalledfuncinlist(f) + installcallback('afterinit', f) # ---------------------------------------------------------------------------- def callfrombeforecollisions(f): - installbeforecollisions(f) + installcallback('beforecollisions', f) return f def installbeforecollisions(f): - "Adds a function to the list of functions called before collisions" - _beforecollisions.installfuncinlist(f) -def uninstallbeforecollisions(f): - "Removes the function from the list of functions called before collisions" - _beforecollisions.uninstallfuncinlist(f) -def isinstalledbeforecollisions(f): - "Checks if the function is called before collisions" - return _beforecollisions.isinstalledfuncinlist(f) + installcallback('beforecollisions', f) # ---------------------------------------------------------------------------- def callfromaftercollisions(f): - installaftercollisions(f) + installcallback('aftercollisions', f) return f def installaftercollisions(f): - "Adds a function to the list of functions called after collisions" - _aftercollisions.installfuncinlist(f) -def uninstallaftercollisions(f): - "Removes the function from the list of functions called after collisions" - _aftercollisions.uninstallfuncinlist(f) -def isinstalledaftercollisions(f): - "Checks if the function is called after collisions" - return _aftercollisions.isinstalledfuncinlist(f) + installcallback('aftercollisions', f) # ---------------------------------------------------------------------------- def callfrombeforeEsolve(f): - installbeforeEsolve(f) + installcallback('beforeEsolve', f) return f def installbeforeEsolve(f): - "Adds a function to the list of functions called before an E solve" - _beforeEsolve.installfuncinlist(f) -def uninstallbeforeEsolve(f): - "Removes the function from the list of functions called before an E solve" - _beforeEsolve.uninstallfuncinlist(f) -def isinstalledbeforeEsolve(f): - "Checks if the function is called before an E solve" - return _beforeEsolve.isinstalledfuncinlist(f) + installcallback('beforeEsolve', f) # ---------------------------------------------------------------------------- def callfrompoissonsolver(f): - installpoissonsolver(f) + installcallback('poissonsolver', f) return f def installpoissonsolver(f): - """Installs an external function to solve Poisson's equation""" - if _poissonsolver.hasfuncsinstalled(): - raise RuntimeError("Only one external Poisson solver can be installed.") - _poissonsolver.installfuncinlist(f) -def uninstallpoissonsolver(f): - """Removes the external function to solve Poisson's equation""" - _poissonsolver.uninstallfuncinlist(f) -def isinstalledpoissonsolver(f): - """Checks if the function is called to solve Poisson's equation""" - return _poissonsolver.isinstalledfuncinlist(f) + installcallback('poissonsolver', f) # ---------------------------------------------------------------------------- def callfromafterEsolve(f): - installafterEsolve(f) + installcallback('afterEsolve', f) return f def installafterEsolve(f): - "Adds a function to the list of functions called after an E solve" - _afterEsolve.installfuncinlist(f) -def uninstallafterEsolve(f): - "Removes the function from the list of functions called after an E solve" - _afterEsolve.uninstallfuncinlist(f) -def isinstalledafterEsolve(f): - "Checks if the function is called after an E solve" - return _afterEsolve.isinstalledfuncinlist(f) + installcallback('afterEsolve', f) # ---------------------------------------------------------------------------- def callfrombeforedeposition(f): - installbeforedeposition(f) + installcallback('beforedeposition', f) return f def installbeforedeposition(f): - "Adds a function to the list of functions called before a particle deposition" - _beforedeposition.installfuncinlist(f) -def uninstallbeforedeposition(f): - "Removes the function from the list of functions called before a particle deposition" - _beforedeposition.uninstallfuncinlist(f) -def isinstalledbeforedeposition(f): - "Checks if the function is called before a particle deposition" - return _beforedeposition.isinstalledfuncinlist(f) + installcallback('beforedeposition', f) # ---------------------------------------------------------------------------- def callfromafterdeposition(f): - installafterdeposition(f) + installcallback('afterdeposition', f) return f def installafterdeposition(f): - "Adds a function to the list of functions called after a particle deposition" - _afterdeposition.installfuncinlist(f) -def uninstallafterdeposition(f): - "Removes the function from the list of functions called after a particle deposition" - _afterdeposition.uninstallfuncinlist(f) -def isinstalledafterdeposition(f): - "Checks if the function is called after a particle deposition" - return _afterdeposition.isinstalledfuncinlist(f) + installcallback('afterdeposition', f) # ---------------------------------------------------------------------------- def callfromparticlescraper(f): - installparticlescraper(f) + installcallback('particlescraper', f) return f def installparticlescraper(f): - "Adds a function to the list of functions called to scrape particles" - _particlescraper.installfuncinlist(f) -def uninstallparticlescraper(f): - "Removes the function from the list of functions called to scrape particles" - _particlescraper.uninstallfuncinlist(f) -def isinstalledparticlescraper(f): - "Checks if the function is called to scrape particles" - return _particlescraper.isinstalledfuncinlist(f) + installcallback('particlescraper', f) # ---------------------------------------------------------------------------- def callfromparticleloader(f): - installparticleloader(f) + installcallback('particleloader', f) return f def installparticleloader(f): - "Adds a function to the list of functions called to load particles" - _particleloader.installfuncinlist(f) -def uninstallparticleloader(f): - "Removes the function from the list of functions called to load particles" - _particleloader.uninstallfuncinlist(f) -def isinstalledparticleloader(f): - "Checks if the function is called to load particles" - return _particleloader.isinstalledfuncinlist(f) + installcallback('particleloader', f) # ---------------------------------------------------------------------------- def callfrombeforestep(f): - installbeforestep(f) + installcallback('beforestep', f) return f def installbeforestep(f): - "Adds a function to the list of functions called before a step" - _beforestep.installfuncinlist(f) -def uninstallbeforestep(f): - "Removes the function from the list of functions called before a step" - _beforestep.uninstallfuncinlist(f) -def isinstalledbeforestep(f): - "Checks if the function is called before a step" - return _beforestep.isinstalledfuncinlist(f) + installcallback('beforestep', f) # ---------------------------------------------------------------------------- def callfromafterstep(f): - installafterstep(f) + installcallback('afterstep', f) return f def installafterstep(f): - "Adds a function to the list of functions called after a step" - _afterstep.installfuncinlist(f) -def uninstallafterstep(f): - "Removes the function from the list of functions called after a step" - _afterstep.uninstallfuncinlist(f) -def isinstalledafterstep(f): - "Checks if the function is called after a step" - return _afterstep.isinstalledfuncinlist(f) + installcallback('afterstep', f) # ---------------------------------------------------------------------------- def callfromafterdiagnostics(f): - installafterdiagnostics(f) + installcallback('afterdiagnostics', f) return f def installafterdiagnostics(f): - "Adds a function to the list of functions called after diagnostic output" - _afterdiagnostics.installfuncinlist(f) -def uninstallafterdiagnostics(f): - "Removes the function from the list of functions called after diagnostic output" - _afterdiagnostics.uninstallfuncinlist(f) -def isinstalledafterdiagnostics(f): - "Checks if the function is called after diagnostic output" - return _afterdiagnostics.isinstalledfuncinlist(f) - -# ---------------------------------------------------------------------------- -def callfromafterrestart(f): - raise Exception('restart call back not implemented yet') - installafterrestart(f) - return f -def installafterrestart(f): - "Adds a function to the list of functions called immediately after a restart" - raise Exception('restart call back not implemented yet') - _afterrestart.installfuncinlist(f) -def uninstallafterrestart(f): - "Removes the function from the list of functions called immediately after a restart" - raise Exception('restart call back not implemented yet') - _afterrestart.uninstallfuncinlist(f) -def isinstalledafterrestart(f): - "Checks if the function is called immediately after a restart" - raise Exception('restart call back not implemented yet') - return _afterrestart.isinstalledfuncinlist(f) + installcallback('afterdiagnostics', f) # ---------------------------------------------------------------------------- def oncheckpointsignal(f): - installoncheckpointsignal(f) + installcallback('oncheckpointsignal', f) return f def installoncheckpointsignal(f): - "Adds a function to the list of functions called on checkpoint signal" - _oncheckpointsignal.installfuncinlist(f) -def uninstalloncheckpointsignal(f): - "Removes the function from the list of functions called on checkpoint signal" - _oncheckpointsignal.uninstallfuncinlist(f) -def isinstalledoncheckpointsignal(f): - "Checks if the function is called on checkpoint signal" - return _oncheckpointsignal.isinstalledfuncinlist(f) + installcallback('oncheckpointsignal', f) # ---------------------------------------------------------------------------- def onbreaksignal(f): - installonbreaksignal(f) + installcallback('onbreaksignal', f) return f def installonbreaksignal(f): - "Adds a function to the list of functions called on a break signal" - _onbreaksignal.installfuncinlist(f) -def uninstallonbreaksignal(f): - "Removes the function from the list of functions called on a break signal" - _onbreaksignal.uninstallfuncinlist(f) -def isinstalledonbreaksignal(f): - "Checks if the function is called on a break signal" - return _onbreaksignal.isinstalledfuncinlist(f) + installcallback('onbreaksignal', f) # ---------------------------------------------------------------------------- def callfromparticleinjection(f): - installparticleinjection(f) + installcallback('particleinjection', f) return f def installparticleinjection(f): - """ - Adds a user defined function that is to be called when particle - injection happens, after the position advance and before deposition is - called, allowing a user defined particle distribution to be injected - each time step""" - _particleinjection.installfuncinlist(f) -def uninstallparticleinjection(f): - "Removes the function installed by installparticleinjection" - _particleinjection.uninstallfuncinlist(f) -def isinstalledparticleinjection(f): - "Checks if the function is called when particles injection happens" - return _particleinjection.isinstalledfuncinlist(f) - -# ---------------------------------------------------------------------------- -def callfromappliedfields(f): - raise Exception('applied fields call back not implemented yet') - installappliedfields(f) - return f -def installappliedfields(f): - """ - Adds a user defined function which can specify E and B fields which are applied - to the particles during the particle advance. - """ - raise Exception('applied fields call back not implemented yet') - _appliedfields.installfuncinlist(f) -def uninstallappliedfields(f): - "Removes the function installed by installappliedfields" - raise Exception('applied fields call back not implemented yet') - _appliedfields.uninstallfuncinlist(f) -def isinstalledappliedfields(f): - "Checks if the function is called when which applies fields" - raise Exception('applied fields call back not implemented yet') - return _appliedfields.isinstalledfuncinlist(f) + installcallback('particleinjection', f) diff --git a/Python/pywarpx/fields.py b/Python/pywarpx/fields.py index 40ce6aa4a..ab66a4646 100644 --- a/Python/pywarpx/fields.py +++ b/Python/pywarpx/fields.py @@ -1,10 +1,10 @@ -# Copyright 2017-2019 David Grote +# Copyright 2017-2023 David Grote # # This file is part of WarpX. # # License: BSD-3-Clause-LBNL -"""Provides wrappers around field and current density on multiFABs +"""Provides wrappers around multiFABs Available routines: @@ -12,10 +12,43 @@ ExWrapper, EyWrapper, EzWrapper BxWrapper, ByWrapper, BzWrapper JxWrapper, JyWrapper, JzWrapper +ExFPWrapper, EyFPWrapper, EzFPWrapper +BxFPWrapper, ByFPWrapper, BzFPWrapper +JxFPWrapper, JyFPWrapper, JzFPWrapper + +RhoFPWrapper, PhiFPWrapper +FFPWrapper, GFPWrapper +AxFPWrapper, AyFPWrapper, AzFPWrapper + +ExCPWrapper, EyCPWrapper, EzCPWrapper +BxCPWrapper, ByCPWrapper, BzCPWrapper +JxCPWrapper, JyCPWrapper, JzCPWrapper + +RhoCPWrapper +FCPWrapper, GCPWrapper + +EdgeLengthsxWrapper, EdgeLengthsyWrapper, EdgeLengthszWrapper +FaceAreasxWrapper, FaceAreasyWrapper, FaceAreaszWrapper + +ExFPPMLWrapper, EyFPPMLWrapper, EzFPPMLWrapper +BxFPPMLWrapper, ByFPPMLWrapper, BzFPPMLWrapper +JxFPPMLWrapper, JyFPPMLWrapper, JzFPPMLWrapper +JxFPAmpereWrapper, JyFPAmpereWrapper, JzFPAmpereWrapper +FFPPMLWrapper, GFPPMLWrapper + +ExCPPMLWrapper, EyCPPMLWrapper, EzCPPMLWrapper +BxCPPMLWrapper, ByCPPMLWrapper, BzCPPMLWrapper +JxCPPMLWrapper, JyCPPMLWrapper, JzCPPMLWrapper +FCPPMLWrapper, GCPPMLWrapper """ import numpy as np try: + import cupy as cp +except ImportError: + cp = None + +try: from mpi4py import MPI as mpi comm_world = mpi.COMM_WORLD npes = comm_world.Get_size() @@ -26,63 +59,73 @@ from ._libwarpx import libwarpx class _MultiFABWrapper(object): - """Wrapper around field arrays at level 0 - This provides a convenient way to query and set fields that are broken up into FABs. + """Wrapper around MultiFabs + This provides a convenient way to query and set data in the MultiFabs. The indexing is based on global indices. - - direction: component to access, one of the values (0, 1, 2) or None - - get_lovects: routine that returns the list of lo vectors - - get_fabs: routine that returns the list of FABs - - get_nodal_flag: routine that returns the list of nodal flag - - level: refinement level + + Parameters + ---------- + mf_name: string + The name of the MultiFab to be accessed, the tag specified when the + MultiFab is allocated + + level: int + The refinement level + + include_ghosts: bool, default=False + Whether to include the ghost cells. + Note that when True, the first n-ghost negative indices will refer to the lower + ghost cells. """ - def __init__(self, direction, get_lovects, get_fabs, get_nodal_flag, level, include_ghosts=False): - self.direction = direction - self.get_lovects = get_lovects - self.get_fabs = get_fabs - self.get_nodal_flag = get_nodal_flag + def __init__(self, mf_name, level, include_ghosts=False): + self.mf_name = mf_name self.level = level self.include_ghosts = include_ghosts self.dim = libwarpx.dim - # overlaps is one along the axes where the grid boundaries overlap the neighboring grid, + # The overlaps list is one along the axes where the grid boundaries overlap the neighboring grid, # which is the case with node centering. - # This presumably will never change during a calculation. - self.overlaps = self.get_nodal_flag() - - def _getlovects(self): - if self.direction is None: - lovects, ngrow = self.get_lovects(self.level, self.include_ghosts) - else: - lovects, ngrow = self.get_lovects(self.level, self.direction, self.include_ghosts) - return lovects, ngrow - - def _gethivects(self): - lovects, ngrow = self._getlovects() - fields = self._getfields() - - hivects = np.zeros_like(lovects) - for i in range(len(fields)): - hivects[:,i] = lovects[:,i] + np.array(fields[i].shape[:self.dim]) - self.overlaps - - return hivects, ngrow - - def _getfields(self): - if self.direction is None: - return self.get_fabs(self.level, self.include_ghosts) - else: - return self.get_fabs(self.level, self.direction, self.include_ghosts) + ix_type = self.mf.box_array().ix_type() + self.overlaps = self._get_indices([int(ix_type.node_centered(i)) for i in range(self.dim)], 0) def __len__(self): - lovects, ngrow = self._getlovects() - return len(lovects) + "Returns the number of blocks" + return self.mf.size + + def __iter__(self): + "The iteration is over the MultiFab" + return self.mf.__iter__() + + @property + def mf(self): + # Always fetch this anew in case the C++ MultiFab is recreated + warpx = libwarpx.libwarpx_so.get_instance() + # All MultiFab names have the level suffix + return warpx.multifab(f'{self.mf_name}[level={self.level}]') + + @property + def shape(self): + """Returns the shape of the global array + """ + min_box = self.mf.box_array().minimal_box() + shape = list(min_box.size - min_box.small_end) + if self.include_ghosts: + nghosts = self.mf.n_grow_vect() + shape = [shape[i] + 2*nghosts[i] for i in range(self.dim)] + shape.append(self.mf.nComp) + return tuple(shape) def mesh(self, direction): """Returns the mesh along the specified direction with the appropriate centering. - - direction: In 3d, one of 'x', 'y', or 'z'. - In 2d, Cartesian, one of 'x', or 'z'. - In RZ, one of 'r', or 'z' - In Z, 'z'. + + Parameters + ---------- + direction: string + In 3d, one of 'x', 'y', or 'z'. + In 2d, Cartesian, one of 'x', or 'z'. + In RZ, one of 'r', or 'z' + In Z, 'z'. """ try: @@ -101,833 +144,606 @@ class _MultiFABWrapper(object): except ValueError: raise Exception('Inappropriate direction given') - # --- Get the total number of cells along the direction - hivects, ngrow = self._gethivects() - nn = hivects[idir,:].max() - ngrow[idir] + self.overlaps[idir] - if npes > 1: - nn = comm_world.allreduce(nn, op=mpi.MAX) - - # --- Cell size in the direction - dd = libwarpx.getCellSize(celldir, self.level) - - # --- Get the nodal flag along direction - nodal_flag = self.get_nodal_flag()[idir] - - # --- The centering shift - if nodal_flag == 1: + min_box = self.mf.box_array().minimal_box() + ilo = min_box.small_end[idir] + ihi = min_box.big_end[idir] + + if self.include_ghosts: + # The ghost cells are added to the upper and lower end of the global domain. + nghosts = self.mf.n_grow_vect() + ilo = list(ilo) + ihi = list(ihi) + min_box = self.mf.box_array().minimal_box() + imax = min_box.big_end + for i in range(self.dim): + if ilo[i] == 0: + ilo[i] -= nghosts[i] + if ihi[i] == imax[i]: + ihi[i] += nghosts[i] + + # Cell size in the direction + warpx = libwarpx.libwarpx_so.get_instance() + dd = warpx.Geom(self.level).data().CellSize(idir) + + # The centering shift + ix_type = self.mf.box_array().ix_type() + if ix_type.node_centered(idir): # node centered shift = 0. else: # cell centered shift = 0.5*dd - return np.arange(nn)*dd + shift + lo = warpx.Geom(self.level).ProbLo(idir) + return lo + np.arange(ilo,ihi+1)*dd + shift + + def _get_indices(self, index, missing): + """Expand the index list to length three. + + Parameters + ---------- + index: sequence of length dims + The indices for each dim + + missing: + The value used to fill in the extra dimensions added + """ + result = [] + for i in range(self.dim): + result.append(index[i]) + for i in range(self.dim, 3): + result.append(missing) + return result + + def _get_n_ghosts(self): + """Return the list of number of ghosts. This includes the component dimension.""" + nghosts = list(self._get_indices(self.mf.n_grow_vect(), 0)) + # The components always has nghosts = 0 + nghosts.append(0) + return nghosts + + def _get_min_indices(self): + """Returns the minimum indices, expanded to length 3""" + min_box = self.mf.box_array().minimal_box() + if self.include_ghosts: + min_box.grow(self.mf.n_grow_vect()) + imin = self._get_indices(min_box.small_end, 0) + return imin + + def _get_max_indices(self): + """Returns the maximum indices, expanded to length 3. + """ + min_box = self.mf.box_array().minimal_box() + if self.include_ghosts: + min_box.grow(self.mf.n_grow_vect()) + imax = self._get_indices(min_box.big_end, 0) + return imax + + def _fix_index(self, ii, imax, d): + """Handle negative index, wrapping them as needed. + This depends on whether ghost cells are included. When true, the indices are + shifted by the number of ghost cells before being wrapped. + """ + nghosts = self._get_n_ghosts() + if self.include_ghosts: + ii += nghosts[d] + if ii < 0: + ii += imax + if self.include_ghosts: + ii -= nghosts[d] + return ii def _find_start_stop(self, ii, imin, imax, d): """Given the input index, calculate the start and stop range of the indices. - - ii: input index, either a slice object or an integer - - imin: the global lowest lovect value in the specified direction - - imax: the global highest hivect value in the specified direction - - d: the direction, an integer, 0, 1, or 2 + + Parameters + ---------- + ii: None, slice, integer + Input index, either None, a slice object, or an integer. + Note that ii can be negative. + + imin: integer + The global lowest lower bound in the specified direction. + This can include the ghost cells. + + imax: integer + The global highest upper bound in the specified direction. + This can include the ghost cells. + This should be the max index + 1. + + d: integer + The dimension number, 0, 1, 2, or 3 (3 being the components) + If ii is a slice, the start and stop values are used directly, unless they are None, then the lower or upper bound is used. An assertion checks if the indices are within the bounds. """ - if isinstance(ii, slice): + if ii is None: + iistart = imin + iistop = imax + elif isinstance(ii, slice): if ii.start is None: iistart = imin else: - iistart = ii.start + iistart = self._fix_index(ii.start, imax, d) if ii.stop is None: - iistop = imax + self.overlaps[d] + iistop = imax else: - iistop = ii.stop + iistop = self._fix_index(ii.stop, imax, d) else: + ii = self._fix_index(ii, imax, d) iistart = ii iistop = ii + 1 - assert imin <= iistart <= imax + self.overlaps[d], Exception(f'Dimension {d} lower index is out of bounds') - assert imin <= iistop <= imax + self.overlaps[d], Exception(f'Dimension {d} upper index is out of bounds') + assert imin <= iistart <= imax, Exception(f'Dimension {d+1} lower index is out of bounds') + assert imin <= iistop <= imax, Exception(f'Dimension {d+1} upper index is out of bounds') return iistart, iistop - def _get_indices(self, index): - if self.dim == 1: - return None, None, index[0] - elif self.dim == 2: - return index[0], None, index[1] - elif self.dim == 3: - return index[0], index[1], index[2] - - def _get_min_indices(self, lovects): - if self.dim == 1: - return 0, 0, lovects[0,:].min() - elif self.dim == 2: - return lovects[0,:].min(), 0, lovects[1,:].min() - elif self.dim == 3: - return lovects[0,:].min(), lovects[1,:].min(), lovects[2,:].min() - - def _get_max_indices(self, hivects): - if self.dim == 1: - return 0, 0, hivects[0,:].max() - elif self.dim == 2: - return hivects[0,:].max(), 0, hivects[1,:].max() - elif self.dim == 3: - return hivects[0,:].max(), hivects[1,:].max(), hivects[2,:].max() - - def _get_vslice(self, lovects, fields_shape, ixstart, ixstop, iystart, - iystop, izstart, izstop, ic): - # --- The ix1, 2 etc are relative to global indexing - if self.dim == 1: - ix1, ix2 = 0, 1 - iy1, iy2 = 0, 1 - iz1 = max(izstart, lovects[0]) - iz2 = min(izstop, lovects[0] + fields_shape[0]) - elif self.dim == 2: - ix1 = max(ixstart, lovects[0]) - ix2 = min(ixstop, lovects[0] + fields_shape[0]) - iy1, iy2 = 0, 1 - iz1 = max(izstart, lovects[1]) - iz2 = min(izstop, lovects[1] + fields_shape[1]) - elif self.dim == 3: - ix1 = max(ixstart, lovects[0]) - ix2 = min(ixstop, lovects[0] + fields_shape[0]) - iy1 = max(iystart, lovects[1]) - iy2 = min(iystop, lovects[1] + fields_shape[1]) - iz1 = max(izstart, lovects[2]) - iz2 = min(izstop, lovects[2] + fields_shape[2]) - - if ix1 < ix2 and iy1 < iy2 and iz1 < iz2: - - if self.dim == 1: - sss = (slice(iz1 - lovects[0], iz2 - lovects[0])) - vslice = (slice(iz1 - izstart, iz2 - izstart)) - - elif self.dim == 2: - sss = (slice(ix1 - lovects[0], ix2 - lovects[0]), - slice(iz1 - lovects[1], iz2 - lovects[1])) - vslice = (slice(ix1 - ixstart, ix2 - ixstart), - slice(iz1 - izstart, iz2 - izstart)) - - elif self.dim == 3: - sss = (slice(ix1 - lovects[0], ix2 - lovects[0]), - slice(iy1 - lovects[1], iy2 - lovects[1]), - slice(iz1 - lovects[2], iz2 - lovects[2])) - vslice = (slice(ix1 - ixstart, ix2 - ixstart), - slice(iy1 - iystart, iy2 - iystart), - slice(iz1 - izstart, iz2 - izstart)) - - if ic is not None: - sss = tuple(list(sss) + [ic]) - - return sss, vslice + def _get_field(self, mfi): + """Return the field at the given mfi. + If include ghosts is true, return the whole array, otherwise + return the interior slice that does not include the ghosts. + """ + # Note that the array will always have 4 dimensions. + # even when self.dim < 3. + # The transpose is taken since the Python array interface to Array4 in + # self.mf.array(mfi) is in C ordering. + # Note: transposing creates a view and not a copy. + device_arr4 = self.mf.array(mfi) + if cp is not None: + device_arr = cp.array(device_arr4, copy=False).T + else: + device_arr = np.array(device_arr4, copy=False).T + if not self.include_ghosts: + nghosts = self._get_n_ghosts() + device_arr = device_arr[tuple([slice(ng, -ng) for ng in nghosts[:self.dim]])] + return device_arr + + def _get_intersect_slice(self, mfi, starts, stops, icstart, icstop): + """Return the slices where the block intersects with the global slice. + If the block does not intersect, return None. + This also shifts the block slices by the number of ghost cells in the + MultiFab arrays since the arrays include the ghost cells. + + Parameters + ---------- + mfi: MFIter + The MFIter instance for the current block, + + starts: sequence + The minimum indices of the global slice. + These can be negative. + + stops: sequence + The maximum indices of the global slice. + These can be negative. + + icstart: integer + The minimum component index of the global slice. + These can be negative. + + icstops: integer + The maximum component index of the global slice. + These can be negative. + + Returns + ------- + block_slices: + The slice of the intersection relative to the block + + global_slices: + The slice of the intersection relative to the global array where the data from individual block will go + """ + box = mfi.tilebox() + if self.include_ghosts: + box.grow(self.mf.n_grow_vect()) + + ilo = self._get_indices(box.small_end, 0) + ihi = self._get_indices(box.big_end, 0) + + # Add 1 to the upper end to be consistent with the slicing notation + ihi_p1 = [i + 1 for i in ihi] + i1 = np.maximum(starts, ilo) + i2 = np.minimum(stops, ihi_p1) + + if np.all(i1 < i2): + + block_slices = [] + global_slices = [] + for i in range(3): + block_slices.append(slice(i1[i] - ilo[i], i2[i] - ilo[i])) + global_slices.append(slice(i1[i] - starts[i], i2[i] - starts[i])) + + block_slices.append(slice(icstart, icstop)) + global_slices.append(slice(0, icstop - icstart)) + + return tuple(block_slices), tuple(global_slices) else: return None, None def __getitem__(self, index): - """Returns slices of a decomposed array, The shape of - the object returned depends on the number of ix, iy and iz specified, which + """Returns slice of the MultiFab using global indexing. + The shape of the object returned depends on the number of ix, iy and iz specified, which can be from none to all three. Note that the values of ix, iy and iz are relative to the fortran indexing, meaning that 0 is the lower boundary - of the whole domain. + of the whole domain, and in fortran ordering, i.e. [ix,iy,iz]. + This allows negative indexing, though with ghosts cells included, the first n-ghost negative + indices will refer to the lower guard cells. + + Parameters + ---------- + index: integer, or sequence of integers or slices, or Ellipsis + Index of the slice to return """ + # Note that the index can have negative values (which wrap around) and has 1 added to the upper + # limit using python style slicing if index == Ellipsis: - index = tuple(self.dim*[slice(None)]) - - if len(index) < self.dim: - # --- Add extra dims to index if needed + index = self.dim*[slice(None)] + elif isinstance(index, slice): + # If only one slice passed in, it was not wrapped in a list + index = [index] + + if len(index) < self.dim+1: + # Add extra dims to index, including for the component. + # These are the dims left out and assumed to extend over the full size of the dim index = list(index) - for i in range(len(index), self.dim): + while len(index) < self.dim+1: index.append(slice(None)) - index = tuple(index) - - lovects, ngrow = self._getlovects() - hivects, ngrow = self._gethivects() - fields = self._getfields() - - ix, iy, iz = self._get_indices(index) - - if len(fields[0].shape) > self.dim: - ncomps = fields[0].shape[-1] - else: - ncomps = 1 - - if len(index) > self.dim: - if ncomps > 1: - ic = index[self.dim] - else: - raise Exception('Too many indices given') - else: - ic = None - - ixmin, iymin, izmin = self._get_min_indices(lovects) - ixmax, iymax, izmax = self._get_max_indices(hivects) - - if npes > 1: - izmin = comm_world.allreduce(izmin, op=mpi.MIN) - izmax = comm_world.allreduce(izmax, op=mpi.MAX) - if self.dim > 1: - ixmin = comm_world.allreduce(ixmin, op=mpi.MIN) - ixmax = comm_world.allreduce(ixmax, op=mpi.MAX) - if self.dim == 3: - iymin = comm_world.allreduce(iymin, op=mpi.MIN) - iymax = comm_world.allreduce(iymax, op=mpi.MAX) - - # --- Setup the size of the array to be returned. - if self.dim == 1: - ixstart, ixstop = None, None - iystart, iystop = None, None - izstart, izstop = self._find_start_stop(iz, izmin, izmax, 0) - - sss = [max(0, izstop - izstart)] - - elif self.dim == 2: - ixstart, ixstop = self._find_start_stop(ix, ixmin, ixmax, 0) - iystart, iystop = None, None - izstart, izstop = self._find_start_stop(iz, izmin, izmax, 1) - - sss = (max(0, ixstop - ixstart), - max(0, izstop - izstart)) - - elif self.dim == 3: - ixstart, ixstop = self._find_start_stop(ix, ixmin, ixmax, 0) - iystart, iystop = self._find_start_stop(iy, iymin, iymax, 1) - izstart, izstop = self._find_start_stop(iz, izmin, izmax, 2) - - sss = (max(0, ixstop - ixstart), - max(0, iystop - iystart), - max(0, izstop - izstart)) - - # --- Space is added for multiple components if needed. - if ncomps > 1 and ic is None: - sss = tuple(list(sss) + [ncomps]) - # --- Create the array to be returned. - resultglobal = np.zeros(sss, dtype=libwarpx._numpy_real_dtype) - + elif len(index) > self.dim+1: + raise Exception('Too many indices given') + + # Expand the indices to length 3 + ii = self._get_indices(index, None) + ic = index[-1] + + # Global extent. These include the ghost cells when include_ghosts is True + ixmin, iymin, izmin = self._get_min_indices() + ixmax, iymax, izmax = self._get_max_indices() + + # Setup the size of the array to be returned + ixstart, ixstop = self._find_start_stop(ii[0], ixmin, ixmax+1, 0) + iystart, iystop = self._find_start_stop(ii[1], iymin, iymax+1, 1) + izstart, izstop = self._find_start_stop(ii[2], izmin, izmax+1, 2) + icstart, icstop = self._find_start_stop(ic, 0, self.mf.n_comp(), 3) + + # Gather the data to be included in a list to be sent to other processes + starts = [ixstart, iystart, izstart] + stops = [ixstop, iystop, izstop] datalist = [] - for i in range(len(fields)): - sss, vslice = self._get_vslice( - lovects[:,i], fields[i].shape, ixstart, ixstop, iystart, - iystop, izstart, izstop, ic - ) - if vslice is not None: - datalist.append((vslice, fields[i][sss])) + for mfi in self.mf: + block_slices, global_slices = self._get_intersect_slice(mfi, starts, stops, icstart, icstop) + if global_slices is not None: + # Note that the array will always have 4 dimensions. + device_arr = self._get_field(mfi) + if cp is not None: + # Copy the data from the device to the host + slice_arr = cp.asnumpy(device_arr[block_slices]) + else: + slice_arr = device_arr[block_slices] + datalist.append((global_slices, slice_arr)) + # Gather the data from all processors if npes == 1: all_datalist = [datalist] else: all_datalist = comm_world.allgather(datalist) + # Create the array to be returned + result_shape = (max(0, ixstop - ixstart), + max(0, iystop - iystart), + max(0, izstop - izstart), + max(0, icstop - icstart)) + + # Now, copy the data into the result array + result_global = None for datalist in all_datalist: - for vslice, ff in datalist: - resultglobal[vslice] = ff - - # --- Now remove any of the reduced dimensions. - if self.dim == 1: - sss = [slice(None)] - if not isinstance(iz, slice): - sss[0] = 0 - elif self.dim == 2: - sss = [slice(None), slice(None)] - if not isinstance(ix, slice): - sss[0] = 0 - if not isinstance(iz, slice): - sss[1] = 0 - elif self.dim == 3: - sss = [slice(None), slice(None), slice(None)] - if not isinstance(ix, slice): - sss[0] = 0 - if not isinstance(iy, slice): - sss[1] = 0 - if not isinstance(iz, slice): - sss[2] = 0 - - return resultglobal[tuple(sss)] + for global_slices, f_arr in datalist: + if result_global is None: + # Delay allocation to here so that the type can be obtained + result_global = np.zeros(result_shape, dtype=f_arr.dtype) + result_global[global_slices] = f_arr + + if result_global is None: + # Something went wrong with the index and no data was found. Return an empty array. + result_global = np.zeros(0) + + # Remove dimensions of length 1, and if all dimensions + # are removed, return a scalar (that's what the [()] does) + return result_global.squeeze()[()] def __setitem__(self, index, value): - """Sets slices of a decomposed array. The shape of - the input object depends on the number of arguments specified, which can + """Sets slices of a decomposed array. + The shape of the input object depends on the number of arguments specified, which can be from none to all three. - - value: input array (must be supplied) + This allows negative indexing, though with ghosts cells included, the first n-ghost negative + indices will refer to the lower guard cells. + + Parameters + ---------- + index: integer, or sequence of integers or slices, or Ellipsis + The slice to set + + value: scalar or array + Input value to assign to the specified slice of the MultiFab """ + # Note that the index can have negative values (which wrap around) and has 1 added to the upper + # limit using python style slicing if index == Ellipsis: index = tuple(self.dim*[slice(None)]) + elif isinstance(index, slice): + # If only one slice passed in, it was not wrapped in a list + index = [index] - if len(index) < self.dim: - # --- Add extra dims to index if needed + if len(index) < self.dim+1: + # Add extra dims to index, including for the component. + # These are the dims left out and assumed to extend over the full size of the dim. index = list(index) - for i in range(len(index), self.dim): + while len(index) < self.dim+1: index.append(slice(None)) - index = tuple(index) + elif len(index) > self.dim+1: + raise Exception('Too many indices given') - lovects, ngrow = self._getlovects() - hivects, ngrow = self._gethivects() - fields = self._getfields() + # Expand the indices to length 3 + ii = self._get_indices(index, None) + ic = index[-1] - ix, iy, iz = self._get_indices(index) + # Global extent. These include the ghost cells when include_ghosts is True + ixmin, iymin, izmin = self._get_min_indices() + ixmax, iymax, izmax = self._get_max_indices() - if len(fields[0].shape) > self.dim: - ncomps = fields[0].shape[-1] - else: - ncomps = 1 - - if len(index) > self.dim: - if ncomps > 1: - ic = index[self.dim] - else: - raise Exception('Too many indices given') - else: - ic = None - - ixmin, iymin, izmin = self._get_min_indices(lovects) - ixmax, iymax, izmax = self._get_max_indices(hivects) - - if npes > 1: - izmin = comm_world.allreduce(izmin, op=mpi.MIN) - izmax = comm_world.allreduce(izmax, op=mpi.MAX) - if self.dim > 1: - ixmin = comm_world.allreduce(ixmin, op=mpi.MIN) - ixmax = comm_world.allreduce(ixmax, op=mpi.MAX) - if self.dim == 3: - iymin = comm_world.allreduce(iymin, op=mpi.MIN) - iymax = comm_world.allreduce(iymax, op=mpi.MAX) - - # --- Add extra dimensions so that the input has the same number of - # --- dimensions as array. - if self.dim == 1: - ixstart, ixstop = None, None - iystart, iystop = None, None - izstart, izstop = self._find_start_stop(iz, izmin, izmax, 0) - - elif self.dim == 2: - ixstart, ixstop = self._find_start_stop(ix, ixmin, ixmax, 0) - iystart, iystop = None, None - izstart, izstop = self._find_start_stop(iz, izmin, izmax, 1) - - elif self.dim == 3: - ixstart, ixstop = self._find_start_stop(ix, ixmin, ixmax, 0) - iystart, iystop = self._find_start_stop(iy, iymin, iymax, 1) - izstart, izstop = self._find_start_stop(iz, izmin, izmax, 2) + # Setup the size of the global array to be set + ixstart, ixstop = self._find_start_stop(ii[0], ixmin, ixmax+1, 0) + iystart, iystop = self._find_start_stop(ii[1], iymin, iymax+1, 1) + izstart, izstop = self._find_start_stop(ii[2], izmin, izmax+1, 2) + icstart, icstop = self._find_start_stop(ic, 0, self.mf.n_comp(), 3) if isinstance(value, np.ndarray): - value3d = np.array(value, copy=False) - sss = list(value3d.shape) - if self.dim == 1: - if not isinstance(iz, slice): sss[0:0] = [1] - elif self.dim == 2: - if not isinstance(ix, slice): sss[0:0] = [1] - if not isinstance(iz, slice): sss[1:1] = [1] - elif self.dim == 3: - if not isinstance(ix, slice): sss[0:0] = [1] - if not isinstance(iy, slice): sss[1:1] = [1] - if not isinstance(iz, slice): sss[2:2] = [1] - value3d.shape = sss - - for i in range(len(fields)): - sss, vslice = self._get_vslice( - lovects[:,i], fields[i].shape, ixstart, ixstop, iystart, - iystop, izstart, izstop, ic - ) - if vslice is not None: + # Expand the shape of the input array to match the shape of the global array + # (it needs to be 4-D). + # This converts value to an array if needed, and the [...] grabs a view so + # that the shape change below doesn't affect value. + value3d = np.array(value)[...] + global_shape = list(value3d.shape) + # The shape of 1 is added for the extra dimensions and when index is an integer + # (in which case the dimension was not in the input array). + if not isinstance(ii[0], slice): global_shape[0:0] = [1] + if not isinstance(ii[1], slice): global_shape[1:1] = [1] + if not isinstance(ii[2], slice): global_shape[2:2] = [1] + if not isinstance(ic , slice) or len(global_shape) < 4: global_shape[3:3] = [1] + value3d.shape = global_shape + + starts = [ixstart, iystart, izstart] + stops = [ixstop, iystop, izstop] + for mfi in self.mf: + block_slices, global_slices = self._get_intersect_slice(mfi, starts, stops, icstart, icstop) + if global_slices is not None: + mf_arr = self._get_field(mfi) if isinstance(value, np.ndarray): - fields[i][sss] = value3d[vslice] + slice_value = value3d[global_slices] + if cp is not None: + # Copy data from host to device + slice_value = cp.asarray(value3d[global_slices]) + mf_arr[block_slices] = slice_value else: - fields[i][sss] = value + mf_arr[block_slices] = value + + def min(self, *args): + return self.mf.min(*args) + + def max(self, *args): + return self.mf.max(*args) + + def sum(self, *args): + return self.mf.sum(*args) + def min_index(self, *args): + return self.mf.minIndex(*args) + + def max_index(self, *args): + return self.mf.maxIndex(*args) + + def norm0(self, *args): + return self.mf.norm0(*args) def ExWrapper(level=0, include_ghosts=False): - return _MultiFABWrapper(direction=0, - get_lovects=libwarpx.get_mesh_electric_field_lovects, - get_fabs=libwarpx.get_mesh_electric_field, - get_nodal_flag=libwarpx.get_Ex_nodal_flag, - level=level, include_ghosts=include_ghosts) + return _MultiFABWrapper(mf_name=f'Efield_aux[x]', level=level, include_ghosts=include_ghosts) def EyWrapper(level=0, include_ghosts=False): - return _MultiFABWrapper(direction=1, - get_lovects=libwarpx.get_mesh_electric_field_lovects, - get_fabs=libwarpx.get_mesh_electric_field, - get_nodal_flag=libwarpx.get_Ey_nodal_flag, - level=level, include_ghosts=include_ghosts) + return _MultiFABWrapper(mf_name=f'Efield_aux[y]', level=level, include_ghosts=include_ghosts) def EzWrapper(level=0, include_ghosts=False): - return _MultiFABWrapper(direction=2, - get_lovects=libwarpx.get_mesh_electric_field_lovects, - get_fabs=libwarpx.get_mesh_electric_field, - get_nodal_flag=libwarpx.get_Ez_nodal_flag, - level=level, include_ghosts=include_ghosts) + return _MultiFABWrapper(mf_name=f'Efield_aux[z]', level=level, include_ghosts=include_ghosts) def BxWrapper(level=0, include_ghosts=False): - return _MultiFABWrapper(direction=0, - get_lovects=libwarpx.get_mesh_magnetic_field_lovects, - get_fabs=libwarpx.get_mesh_magnetic_field, - get_nodal_flag=libwarpx.get_Bx_nodal_flag, - level=level, include_ghosts=include_ghosts) + return _MultiFABWrapper(mf_name=f'Bfield_aux[x]', level=level, include_ghosts=include_ghosts) def ByWrapper(level=0, include_ghosts=False): - return _MultiFABWrapper(direction=1, - get_lovects=libwarpx.get_mesh_magnetic_field_lovects, - get_fabs=libwarpx.get_mesh_magnetic_field, - get_nodal_flag=libwarpx.get_By_nodal_flag, - level=level, include_ghosts=include_ghosts) + return _MultiFABWrapper(mf_name=f'Bfield_aux[y]', level=level, include_ghosts=include_ghosts) def BzWrapper(level=0, include_ghosts=False): - return _MultiFABWrapper(direction=2, - get_lovects=libwarpx.get_mesh_magnetic_field_lovects, - get_fabs=libwarpx.get_mesh_magnetic_field, - get_nodal_flag=libwarpx.get_Bz_nodal_flag, - level=level, include_ghosts=include_ghosts) + return _MultiFABWrapper(mf_name=f'Bfield_aux[z]', level=level, include_ghosts=include_ghosts) def JxWrapper(level=0, include_ghosts=False): - return _MultiFABWrapper(direction=0, - get_lovects=libwarpx.get_mesh_current_density_lovects, - get_fabs=libwarpx.get_mesh_current_density, - get_nodal_flag=libwarpx.get_Jx_nodal_flag, - level=level, include_ghosts=include_ghosts) + return _MultiFABWrapper(mf_name=f'current_fp[x]', level=level, include_ghosts=include_ghosts) def JyWrapper(level=0, include_ghosts=False): - return _MultiFABWrapper(direction=1, - get_lovects=libwarpx.get_mesh_current_density_lovects, - get_fabs=libwarpx.get_mesh_current_density, - get_nodal_flag=libwarpx.get_Jy_nodal_flag, - level=level, include_ghosts=include_ghosts) + return _MultiFABWrapper(mf_name=f'current_fp[y]', level=level, include_ghosts=include_ghosts) def JzWrapper(level=0, include_ghosts=False): - return _MultiFABWrapper(direction=2, - get_lovects=libwarpx.get_mesh_current_density_lovects, - get_fabs=libwarpx.get_mesh_current_density, - get_nodal_flag=libwarpx.get_Jz_nodal_flag, - level=level, include_ghosts=include_ghosts) - -def ExCPWrapper(level=1, include_ghosts=False): - assert level>0, Exception('Coarse patch only available on levels > 0') - return _MultiFABWrapper(direction=0, - get_lovects=libwarpx.get_mesh_electric_field_cp_lovects, - get_fabs=libwarpx.get_mesh_electric_field_cp, - get_nodal_flag=libwarpx.get_Ex_nodal_flag, - level=level, include_ghosts=include_ghosts) - -def EyCPWrapper(level=1, include_ghosts=False): - assert level>0, Exception('Coarse patch only available on levels > 0') - return _MultiFABWrapper(direction=1, - get_lovects=libwarpx.get_mesh_electric_field_cp_lovects, - get_fabs=libwarpx.get_mesh_electric_field_cp, - get_nodal_flag=libwarpx.get_Ey_nodal_flag, - level=level, include_ghosts=include_ghosts) - -def EzCPWrapper(level=1, include_ghosts=False): - assert level>0, Exception('Coarse patch only available on levels > 0') - return _MultiFABWrapper(direction=2, - get_lovects=libwarpx.get_mesh_electric_field_cp_lovects, - get_fabs=libwarpx.get_mesh_electric_field_cp, - get_nodal_flag=libwarpx.get_Ez_nodal_flag, - level=level, include_ghosts=include_ghosts) - -def BxCPWrapper(level=1, include_ghosts=False): - assert level>0, Exception('Coarse patch only available on levels > 0') - return _MultiFABWrapper(direction=0, - get_lovects=libwarpx.get_mesh_magnetic_field_cp_lovects, - get_fabs=libwarpx.get_mesh_magnetic_field_cp, - get_nodal_flag=libwarpx.get_Bx_nodal_flag, - level=level, include_ghosts=include_ghosts) - -def ByCPWrapper(level=1, include_ghosts=False): - assert level>0, Exception('Coarse patch only available on levels > 0') - return _MultiFABWrapper(direction=1, - get_lovects=libwarpx.get_mesh_magnetic_field_cp_lovects, - get_fabs=libwarpx.get_mesh_magnetic_field_cp, - get_nodal_flag=libwarpx.get_By_nodal_flag, - level=level, include_ghosts=include_ghosts) - -def BzCPWrapper(level=1, include_ghosts=False): - assert level>0, Exception('Coarse patch only available on levels > 0') - return _MultiFABWrapper(direction=2, - get_lovects=libwarpx.get_mesh_magnetic_field_cp_lovects, - get_fabs=libwarpx.get_mesh_magnetic_field_cp, - get_nodal_flag=libwarpx.get_Bz_nodal_flag, - level=level, include_ghosts=include_ghosts) - -def JxCPWrapper(level=1, include_ghosts=False): - assert level>0, Exception('Coarse patch only available on levels > 0') - return _MultiFABWrapper(direction=0, - get_lovects=libwarpx.get_mesh_current_density_cp_lovects, - get_fabs=libwarpx.get_mesh_current_density_cp, - get_nodal_flag=libwarpx.get_Jx_nodal_flag, - level=level, include_ghosts=include_ghosts) - -def JyCPWrapper(level=1, include_ghosts=False): - assert level>0, Exception('Coarse patch only available on levels > 0') - return _MultiFABWrapper(direction=1, - get_lovects=libwarpx.get_mesh_current_density_cp_lovects, - get_fabs=libwarpx.get_mesh_current_density_cp, - get_nodal_flag=libwarpx.get_Jy_nodal_flag, - level=level, include_ghosts=include_ghosts) - -def JzCPWrapper(level=1, include_ghosts=False): - assert level>0, Exception('Coarse patch only available on levels > 0') - return _MultiFABWrapper(direction=2, - get_lovects=libwarpx.get_mesh_current_density_cp_lovects, - get_fabs=libwarpx.get_mesh_current_density_cp, - get_nodal_flag=libwarpx.get_Jz_nodal_flag, - level=level, include_ghosts=include_ghosts) - -def RhoCPWrapper(level=1, include_ghosts=False): - assert level>0, Exception('Coarse patch only available on levels > 0') - return _MultiFABWrapper(direction=None, - get_lovects=libwarpx.get_mesh_charge_density_cp_lovects, - get_fabs=libwarpx.get_mesh_charge_density_cp, - get_nodal_flag=libwarpx.get_Rho_nodal_flag, - level=level, include_ghosts=include_ghosts) - -def FCPWrapper(level=1, include_ghosts=False): - assert level>0, Exception('Coarse patch only available on levels > 0') - return _MultiFABWrapper(direction=None, - get_lovects=libwarpx.get_mesh_F_cp_lovects, - get_fabs=libwarpx.get_mesh_F_cp, - get_nodal_flag=libwarpx.get_F_nodal_flag, - level=level, include_ghosts=include_ghosts) - -def GCPWrapper(level=1, include_ghosts=False): - assert level>0, Exception('Coarse patch only available on levels > 0') - return _MultiFABWrapper(direction=None, - get_lovects=libwarpx.get_mesh_G_cp_lovects, - get_fabs=libwarpx.get_mesh_G_cp, - get_nodal_flag=libwarpx.get_G_nodal_flag, - level=level, include_ghosts=include_ghosts) + return _MultiFABWrapper(mf_name=f'current_fp[z]', level=level, include_ghosts=include_ghosts) def ExFPWrapper(level=0, include_ghosts=False): - return _MultiFABWrapper(direction=0, - get_lovects=libwarpx.get_mesh_electric_field_fp_lovects, - get_fabs=libwarpx.get_mesh_electric_field_fp, - get_nodal_flag=libwarpx.get_Ex_nodal_flag, - level=level, include_ghosts=include_ghosts) + return _MultiFABWrapper(mf_name=f'Efield_fp[x]', level=level, include_ghosts=include_ghosts) def EyFPWrapper(level=0, include_ghosts=False): - return _MultiFABWrapper(direction=1, - get_lovects=libwarpx.get_mesh_electric_field_fp_lovects, - get_fabs=libwarpx.get_mesh_electric_field_fp, - get_nodal_flag=libwarpx.get_Ey_nodal_flag, - level=level, include_ghosts=include_ghosts) + return _MultiFABWrapper(mf_name=f'Efield_fp[y]', level=level, include_ghosts=include_ghosts) def EzFPWrapper(level=0, include_ghosts=False): - return _MultiFABWrapper(direction=2, - get_lovects=libwarpx.get_mesh_electric_field_fp_lovects, - get_fabs=libwarpx.get_mesh_electric_field_fp, - get_nodal_flag=libwarpx.get_Ez_nodal_flag, - level=level, include_ghosts=include_ghosts) + return _MultiFABWrapper(mf_name=f'Efield_fp[z]', level=level, include_ghosts=include_ghosts) def BxFPWrapper(level=0, include_ghosts=False): - return _MultiFABWrapper(direction=0, - get_lovects=libwarpx.get_mesh_magnetic_field_fp_lovects, - get_fabs=libwarpx.get_mesh_magnetic_field_fp, - get_nodal_flag=libwarpx.get_Bx_nodal_flag, - level=level, include_ghosts=include_ghosts) + return _MultiFABWrapper(mf_name=f'Bfield_fp[x]', level=level, include_ghosts=include_ghosts) def ByFPWrapper(level=0, include_ghosts=False): - return _MultiFABWrapper(direction=1, - get_lovects=libwarpx.get_mesh_magnetic_field_fp_lovects, - get_fabs=libwarpx.get_mesh_magnetic_field_fp, - get_nodal_flag=libwarpx.get_By_nodal_flag, - level=level, include_ghosts=include_ghosts) + return _MultiFABWrapper(mf_name=f'Bfield_fp[y]', level=level, include_ghosts=include_ghosts) def BzFPWrapper(level=0, include_ghosts=False): - return _MultiFABWrapper(direction=2, - get_lovects=libwarpx.get_mesh_magnetic_field_fp_lovects, - get_fabs=libwarpx.get_mesh_magnetic_field_fp, - get_nodal_flag=libwarpx.get_Bz_nodal_flag, - level=level, include_ghosts=include_ghosts) + return _MultiFABWrapper(mf_name=f'Bfield_fp[z]', level=level, include_ghosts=include_ghosts) def JxFPWrapper(level=0, include_ghosts=False): - return _MultiFABWrapper(direction=0, - get_lovects=libwarpx.get_mesh_current_density_fp_lovects, - get_fabs=libwarpx.get_mesh_current_density_fp, - get_nodal_flag=libwarpx.get_Jx_nodal_flag, - level=level, include_ghosts=include_ghosts) + return _MultiFABWrapper(mf_name=f'current_fp[x]', level=level, include_ghosts=include_ghosts) def JyFPWrapper(level=0, include_ghosts=False): - return _MultiFABWrapper(direction=1, - get_lovects=libwarpx.get_mesh_current_density_fp_lovects, - get_fabs=libwarpx.get_mesh_current_density_fp, - get_nodal_flag=libwarpx.get_Jy_nodal_flag, - level=level, include_ghosts=include_ghosts) + return _MultiFABWrapper(mf_name=f'current_fp[y]', level=level, include_ghosts=include_ghosts) def JzFPWrapper(level=0, include_ghosts=False): - return _MultiFABWrapper(direction=2, - get_lovects=libwarpx.get_mesh_current_density_fp_lovects, - get_fabs=libwarpx.get_mesh_current_density_fp, - get_nodal_flag=libwarpx.get_Jz_nodal_flag, - level=level, include_ghosts=include_ghosts) + return _MultiFABWrapper(mf_name=f'current_fp[z]', level=level, include_ghosts=include_ghosts) def RhoFPWrapper(level=0, include_ghosts=False): - return _MultiFABWrapper(direction=None, - get_lovects=libwarpx.get_mesh_charge_density_fp_lovects, - get_fabs=libwarpx.get_mesh_charge_density_fp, - get_nodal_flag=libwarpx.get_Rho_nodal_flag, - level=level, include_ghosts=include_ghosts) + return _MultiFABWrapper(mf_name=f'rho_fp', level=level, include_ghosts=include_ghosts) def PhiFPWrapper(level=0, include_ghosts=False): - return _MultiFABWrapper(direction=None, - get_lovects=libwarpx.get_mesh_phi_fp_lovects, - get_fabs=libwarpx.get_mesh_phi_fp, - get_nodal_flag=libwarpx.get_Phi_nodal_flag, - level=level, include_ghosts=include_ghosts) + return _MultiFABWrapper(mf_name=f'phi_fp', level=level, include_ghosts=include_ghosts) + +def FFPWrapper(level=0, include_ghosts=False): + return _MultiFABWrapper(mf_name=f'F_fp', level=level, include_ghosts=include_ghosts) + +def GFPWrapper(level=0, include_ghosts=False): + return _MultiFABWrapper(mf_name=f'G_fp', level=level, include_ghosts=include_ghosts) def AxFPWrapper(level=0, include_ghosts=False): - return _MultiFABWrapper(direction=0, - get_lovects=libwarpx.get_mesh_vector_potential_fp_lovects, - get_fabs=libwarpx.get_mesh_vector_potential_fp, - get_nodal_flag=libwarpx.get_Ax_nodal_flag, - level=level, include_ghosts=include_ghosts) + return _MultiFABWrapper(mf_name=f'vector_potential_fp_nodal[x]', level=level, include_ghosts=include_ghosts) def AyFPWrapper(level=0, include_ghosts=False): - return _MultiFABWrapper(direction=1, - get_lovects=libwarpx.get_mesh_vector_potential_fp_lovects, - get_fabs=libwarpx.get_mesh_vector_potential_fp, - get_nodal_flag=libwarpx.get_Ay_nodal_flag, - level=level, include_ghosts=include_ghosts) + return _MultiFABWrapper(mf_name=f'vector_potential_fp_nodal[y]', level=level, include_ghosts=include_ghosts) def AzFPWrapper(level=0, include_ghosts=False): - return _MultiFABWrapper(direction=2, - get_lovects=libwarpx.get_mesh_vector_potential_fp_lovects, - get_fabs=libwarpx.get_mesh_vector_potential_fp, - get_nodal_flag=libwarpx.get_Az_nodal_flag, - level=level, include_ghosts=include_ghosts) + return _MultiFABWrapper(mf_name=f'vector_potential_fp_nodal[z]', level=level, include_ghosts=include_ghosts) -def FFPWrapper(level=0, include_ghosts=False): - return _MultiFABWrapper(direction=None, - get_lovects=libwarpx.get_mesh_F_fp_lovects, - get_fabs=libwarpx.get_mesh_F_fp, - get_nodal_flag=libwarpx.get_F_nodal_flag, - level=level, include_ghosts=include_ghosts) +def ExCPWrapper(level=0, include_ghosts=False): + return _MultiFABWrapper(mf_name=f'Efield_cp[x]', level=level, include_ghosts=include_ghosts) -def GFPWrapper(level=0, include_ghosts=False): - return _MultiFABWrapper(direction=None, - get_lovects=libwarpx.get_mesh_G_fp_lovects, - get_fabs=libwarpx.get_mesh_G_fp, - get_nodal_flag=libwarpx.get_G_nodal_flag, - level=level, include_ghosts=include_ghosts) +def EyCPWrapper(level=0, include_ghosts=False): + return _MultiFABWrapper(mf_name=f'Efield_cp[y]', level=level, include_ghosts=include_ghosts) + +def EzCPWrapper(level=0, include_ghosts=False): + return _MultiFABWrapper(mf_name=f'Efield_cp[z]', level=level, include_ghosts=include_ghosts) + +def BxCPWrapper(level=0, include_ghosts=False): + return _MultiFABWrapper(mf_name=f'Bfield_cp[x]', level=level, include_ghosts=include_ghosts) + +def ByCPWrapper(level=0, include_ghosts=False): + return _MultiFABWrapper(mf_name=f'Bfield_cp[y]', level=level, include_ghosts=include_ghosts) + +def BzCPWrapper(level=0, include_ghosts=False): + return _MultiFABWrapper(mf_name=f'Bfield_cp[z]', level=level, include_ghosts=include_ghosts) + +def JxCPWrapper(level=0, include_ghosts=False): + return _MultiFABWrapper(mf_name=f'current_cp[x]', level=level, include_ghosts=include_ghosts) + +def JyCPWrapper(level=0, include_ghosts=False): + return _MultiFABWrapper(mf_name=f'current_cp[y]', level=level, include_ghosts=include_ghosts) + +def JzCPWrapper(level=0, include_ghosts=False): + return _MultiFABWrapper(mf_name=f'current_cp[z]', level=level, include_ghosts=include_ghosts) + +def RhoCPWrapper(level=0, include_ghosts=False): + return _MultiFABWrapper(mf_name=f'rho_cp', level=level, include_ghosts=include_ghosts) + +def FCPWrapper(level=0, include_ghosts=False): + return _MultiFABWrapper(mf_name=f'F_cp', level=level, include_ghosts=include_ghosts) + +def GCPWrapper(level=0, include_ghosts=False): + return _MultiFABWrapper(mf_name=f'G_cp', level=level, include_ghosts=include_ghosts) def EdgeLengthsxWrapper(level=0, include_ghosts=False): - return _MultiFABWrapper(direction=0, - get_lovects=libwarpx.get_mesh_edge_lengths_lovects, - get_fabs=libwarpx.get_mesh_edge_lengths, - get_nodal_flag=libwarpx.get_edge_lengths_x_nodal_flag, - level=level, include_ghosts=include_ghosts) + return _MultiFABWrapper(mf_name=f'm_edge_lengths[x]', level=level, include_ghosts=include_ghosts) def EdgeLengthsyWrapper(level=0, include_ghosts=False): - return _MultiFABWrapper(direction=1, - get_lovects=libwarpx.get_mesh_edge_lengths_lovects, - get_fabs=libwarpx.get_mesh_edge_lengths, - get_nodal_flag=libwarpx.get_edge_lengths_y_nodal_flag, - level=level, include_ghosts=include_ghosts) + return _MultiFABWrapper(mf_name=f'm_edge_lengths[y]', level=level, include_ghosts=include_ghosts) def EdgeLengthszWrapper(level=0, include_ghosts=False): - return _MultiFABWrapper(direction=2, - get_lovects=libwarpx.get_mesh_edge_lengths_lovects, - get_fabs=libwarpx.get_mesh_edge_lengths, - get_nodal_flag=libwarpx.get_edge_lengths_z_nodal_flag, - level=level, include_ghosts=include_ghosts) + return _MultiFABWrapper(mf_name=f'm_edge_lengths[z]', level=level, include_ghosts=include_ghosts) def FaceAreasxWrapper(level=0, include_ghosts=False): - return _MultiFABWrapper(direction=0, - get_lovects=libwarpx.get_mesh_face_areas_lovects, - get_fabs=libwarpx.get_mesh_face_areas, - get_nodal_flag=libwarpx.get_face_areas_x_nodal_flag, - level=level, include_ghosts=include_ghosts) + return _MultiFABWrapper(mf_name=f'm_face_areas[x]', level=level, include_ghosts=include_ghosts) def FaceAreasyWrapper(level=0, include_ghosts=False): - return _MultiFABWrapper(direction=1, - get_lovects=libwarpx.get_mesh_face_areas_lovects, - get_fabs=libwarpx.get_mesh_face_areas, - get_nodal_flag=libwarpx.get_face_areas_y_nodal_flag, - level=level, include_ghosts=include_ghosts) + return _MultiFABWrapper(mf_name=f'm_face_areas[y]', level=level, include_ghosts=include_ghosts) def FaceAreaszWrapper(level=0, include_ghosts=False): - return _MultiFABWrapper(direction=2, - get_lovects=libwarpx.get_mesh_face_areas_lovects, - get_fabs=libwarpx.get_mesh_face_areas, - get_nodal_flag=libwarpx.get_face_areas_z_nodal_flag, - level=level, include_ghosts=include_ghosts) - -def ExCPPMLWrapper(level=1, include_ghosts=False): - assert level>0, Exception('Coarse patch only available on levels > 0') - return _MultiFABWrapper(direction=0, - get_lovects=libwarpx.get_mesh_electric_field_cp_lovects_pml, - get_fabs=libwarpx.get_mesh_electric_field_cp_pml, - get_nodal_flag=libwarpx.get_Ex_nodal_flag, - level=level, include_ghosts=include_ghosts) - -def EyCPPMLWrapper(level=1, include_ghosts=False): - assert level>0, Exception('Coarse patch only available on levels > 0') - return _MultiFABWrapper(direction=1, - get_lovects=libwarpx.get_mesh_electric_field_cp_lovects_pml, - get_fabs=libwarpx.get_mesh_electric_field_cp_pml, - get_nodal_flag=libwarpx.get_Ey_nodal_flag, - level=level, include_ghosts=include_ghosts) - -def EzCPPMLWrapper(level=1, include_ghosts=False): - assert level>0, Exception('Coarse patch only available on levels > 0') - return _MultiFABWrapper(direction=2, - get_lovects=libwarpx.get_mesh_electric_field_cp_lovects_pml, - get_fabs=libwarpx.get_mesh_electric_field_cp_pml, - get_nodal_flag=libwarpx.get_Ez_nodal_flag, - level=level, include_ghosts=include_ghosts) - -def BxCPPMLWrapper(level=1, include_ghosts=False): - assert level>0, Exception('Coarse patch only available on levels > 0') - return _MultiFABWrapper(direction=0, - get_lovects=libwarpx.get_mesh_magnetic_field_cp_lovects_pml, - get_fabs=libwarpx.get_mesh_magnetic_field_cp_pml, - get_nodal_flag=libwarpx.get_Bx_nodal_flag, - level=level, include_ghosts=include_ghosts) - -def ByCPPMLWrapper(level=1, include_ghosts=False): - assert level>0, Exception('Coarse patch only available on levels > 0') - return _MultiFABWrapper(direction=1, - get_lovects=libwarpx.get_mesh_magnetic_field_cp_lovects_pml, - get_fabs=libwarpx.get_mesh_magnetic_field_cp_pml, - get_nodal_flag=libwarpx.get_By_nodal_flag, - level=level, include_ghosts=include_ghosts) - -def BzCPPMLWrapper(level=1, include_ghosts=False): - assert level>0, Exception('Coarse patch only available on levels > 0') - return _MultiFABWrapper(direction=2, - get_lovects=libwarpx.get_mesh_magnetic_field_cp_lovects_pml, - get_fabs=libwarpx.get_mesh_magnetic_field_cp_pml, - get_nodal_flag=libwarpx.get_Bz_nodal_flag, - level=level, include_ghosts=include_ghosts) - -def JxCPPMLWrapper(level=1, include_ghosts=False): - assert level>0, Exception('Coarse patch only available on levels > 0') - return _MultiFABWrapper(direction=0, - get_lovects=libwarpx.get_mesh_current_density_cp_lovects_pml, - get_fabs=libwarpx.get_mesh_current_density_cp_pml, - get_nodal_flag=libwarpx.get_Jx_nodal_flag, - level=level, include_ghosts=include_ghosts) - -def JyCPPMLWrapper(level=1, include_ghosts=False): - assert level>0, Exception('Coarse patch only available on levels > 0') - return _MultiFABWrapper(direction=1, - get_lovects=libwarpx.get_mesh_current_density_cp_lovects_pml, - get_fabs=libwarpx.get_mesh_current_density_cp_pml, - get_nodal_flag=libwarpx.get_Jy_nodal_flag, - level=level, include_ghosts=include_ghosts) - -def JzCPPMLWrapper(level=1, include_ghosts=False): - assert level>0, Exception('Coarse patch only available on levels > 0') - return _MultiFABWrapper(direction=2, - get_lovects=libwarpx.get_mesh_current_density_cp_lovects_pml, - get_fabs=libwarpx.get_mesh_current_density_cp_pml, - get_nodal_flag=libwarpx.get_Jz_nodal_flag, - level=level, include_ghosts=include_ghosts) - -def FCPPMLWrapper(level=1, include_ghosts=False): - assert level>0, Exception('Coarse patch only available on levels > 0') - return _MultiFABWrapper(direction=None, - get_lovects=libwarpx.get_mesh_F_cp_lovects_pml, - get_fabs=libwarpx.get_mesh_F_cp_pml, - get_nodal_flag=libwarpx.get_F_nodal_flag, - level=level, include_ghosts=include_ghosts) - -def GCPPMLWrapper(level=1, include_ghosts=False): - assert level>0, Exception('Coarse patch only available on levels > 0') - return _MultiFABWrapper(direction=None, - get_lovects=libwarpx.get_mesh_G_cp_lovects_pml, - get_fabs=libwarpx.get_mesh_G_cp_pml, - get_nodal_flag=libwarpx.get_G_nodal_flag, - level=level, include_ghosts=include_ghosts) + return _MultiFABWrapper(mf_name=f'm_face_areas[z]', level=level, include_ghosts=include_ghosts) + +def JxFPAmpereWrapper(level=0, include_ghosts=False): + return _MultiFABWrapper(mf_name=f'current_fp_ampere[x]', level=level, include_ghosts=include_ghosts) + +def JyFPAmpereWrapper(level=0, include_ghosts=False): + return _MultiFABWrapper(mf_name=f'current_fp_ampere[y]', level=level, include_ghosts=include_ghosts) + +def JzFPAmpereWrapper(level=0, include_ghosts=False): + return _MultiFABWrapper(mf_name=f'current_fp_ampere[z]', level=level, include_ghosts=include_ghosts) def ExFPPMLWrapper(level=0, include_ghosts=False): - return _MultiFABWrapper(direction=0, - get_lovects=libwarpx.get_mesh_electric_field_fp_lovects_pml, - get_fabs=libwarpx.get_mesh_electric_field_fp_pml, - get_nodal_flag=libwarpx.get_Ex_nodal_flag, - level=level, include_ghosts=include_ghosts) + return _MultiFABWrapper(mf_name='pml_E_fp[x]', level=level, include_ghosts=include_ghosts) def EyFPPMLWrapper(level=0, include_ghosts=False): - return _MultiFABWrapper(direction=1, - get_lovects=libwarpx.get_mesh_electric_field_fp_lovects_pml, - get_fabs=libwarpx.get_mesh_electric_field_fp_pml, - get_nodal_flag=libwarpx.get_Ey_nodal_flag, - level=level, include_ghosts=include_ghosts) + return _MultiFABWrapper(mf_name='pml_E_fp[y]', level=level, include_ghosts=include_ghosts) def EzFPPMLWrapper(level=0, include_ghosts=False): - return _MultiFABWrapper(direction=2, - get_lovects=libwarpx.get_mesh_electric_field_fp_lovects_pml, - get_fabs=libwarpx.get_mesh_electric_field_fp_pml, - get_nodal_flag=libwarpx.get_Ez_nodal_flag, - level=level, include_ghosts=include_ghosts) + return _MultiFABWrapper(mf_name='pml_E_fp[z]', level=level, include_ghosts=include_ghosts) def BxFPPMLWrapper(level=0, include_ghosts=False): - return _MultiFABWrapper(direction=0, - get_lovects=libwarpx.get_mesh_magnetic_field_fp_lovects_pml, - get_fabs=libwarpx.get_mesh_magnetic_field_fp_pml, - get_nodal_flag=libwarpx.get_Bx_nodal_flag, - level=level, include_ghosts=include_ghosts) + return _MultiFABWrapper(mf_name='pml_B_fp[x]', level=level, include_ghosts=include_ghosts) def ByFPPMLWrapper(level=0, include_ghosts=False): - return _MultiFABWrapper(direction=1, - get_lovects=libwarpx.get_mesh_magnetic_field_fp_lovects_pml, - get_fabs=libwarpx.get_mesh_magnetic_field_fp_pml, - get_nodal_flag=libwarpx.get_By_nodal_flag, - level=level, include_ghosts=include_ghosts) + return _MultiFABWrapper(mf_name='pml_B_fp[y]', level=level, include_ghosts=include_ghosts) def BzFPPMLWrapper(level=0, include_ghosts=False): - return _MultiFABWrapper(direction=2, - get_lovects=libwarpx.get_mesh_magnetic_field_fp_lovects_pml, - get_fabs=libwarpx.get_mesh_magnetic_field_fp_pml, - get_nodal_flag=libwarpx.get_Bz_nodal_flag, - level=level, include_ghosts=include_ghosts) + return _MultiFABWrapper(mf_name='pml_B_fp[z]', level=level, include_ghosts=include_ghosts) def JxFPPMLWrapper(level=0, include_ghosts=False): - return _MultiFABWrapper(direction=0, - get_lovects=libwarpx.get_mesh_current_density_fp_lovects_pml, - get_fabs=libwarpx.get_mesh_current_density_fp_pml, - get_nodal_flag=libwarpx.get_Jx_nodal_flag, - level=level, include_ghosts=include_ghosts) + return _MultiFABWrapper(mf_name='pml_j_fp[x]', level=level, include_ghosts=include_ghosts) def JyFPPMLWrapper(level=0, include_ghosts=False): - return _MultiFABWrapper(direction=1, - get_lovects=libwarpx.get_mesh_current_density_fp_lovects_pml, - get_fabs=libwarpx.get_mesh_current_density_fp_pml, - get_nodal_flag=libwarpx.get_Jy_nodal_flag, - level=level, include_ghosts=include_ghosts) + return _MultiFABWrapper(mf_name='pml_j_fp[y]', level=level, include_ghosts=include_ghosts) def JzFPPMLWrapper(level=0, include_ghosts=False): - return _MultiFABWrapper(direction=2, - get_lovects=libwarpx.get_mesh_current_density_fp_lovects_pml, - get_fabs=libwarpx.get_mesh_current_density_fp_pml, - get_nodal_flag=libwarpx.get_Jz_nodal_flag, - level=level, include_ghosts=include_ghosts) + return _MultiFABWrapper(mf_name='pml_j_fp[z]', level=level, include_ghosts=include_ghosts) -def JxFPAmpereWrapper(level=0, include_ghosts=False): - return _MultiFABWrapper(direction=0, - get_lovects=libwarpx.get_mesh_current_density_fp_lovects, - get_fabs=libwarpx.get_mesh_current_density_fp_ampere, - get_nodal_flag=libwarpx.get_Jx_nodal_flag, - level=level, include_ghosts=include_ghosts) +def FFPPMLWrapper(level=0, include_ghosts=False): + return _MultiFABWrapper(mf_name='pml_F_fp', level=level, include_ghosts=include_ghosts) -def JyFPAmpereWrapper(level=0, include_ghosts=False): - return _MultiFABWrapper(direction=1, - get_lovects=libwarpx.get_mesh_current_density_fp_lovects, - get_fabs=libwarpx.get_mesh_current_density_fp_ampere, - get_nodal_flag=libwarpx.get_Jy_nodal_flag, - level=level, include_ghosts=include_ghosts) +def GFPPMLWrapper(level=0, include_ghosts=False): + return _MultiFABWrapper(mf_name='pml_G_fp', level=level, include_ghosts=include_ghosts) -def JzFPAmpereWrapper(level=0, include_ghosts=False): - return _MultiFABWrapper(direction=2, - get_lovects=libwarpx.get_mesh_current_density_fp_lovects, - get_fabs=libwarpx.get_mesh_current_density_fp_ampere, - get_nodal_flag=libwarpx.get_Jz_nodal_flag, - level=level, include_ghosts=include_ghosts) +def ExCPPMLWrapper(level=0, include_ghosts=False): + return _MultiFABWrapper(mf_name='pml_E_cp[x]', level=level, include_ghosts=include_ghosts) -def FFPPMLWrapper(level=0, include_ghosts=False): - return _MultiFABWrapper(direction=None, - get_lovects=libwarpx.get_mesh_F_fp_lovects_pml, - get_fabs=libwarpx.get_mesh_F_fp_pml, - get_nodal_flag=libwarpx.get_F_pml_nodal_flag, - level=level, include_ghosts=include_ghosts) +def EyCPPMLWrapper(level=0, include_ghosts=False): + return _MultiFABWrapper(mf_name='pml_E_cp[y]', level=level, include_ghosts=include_ghosts) -def GFPPMLWrapper(level=0, include_ghosts=False): - return _MultiFABWrapper(direction=None, - get_lovects=libwarpx.get_mesh_G_fp_lovects_pml, - get_fabs=libwarpx.get_mesh_G_fp_pml, - get_nodal_flag=libwarpx.get_G_pml_nodal_flag, - level=level, include_ghosts=include_ghosts) +def EzCPPMLWrapper(level=0, include_ghosts=False): + return _MultiFABWrapper(mf_name='pml_E_cp[z]', level=level, include_ghosts=include_ghosts) + +def BxCPPMLWrapper(level=0, include_ghosts=False): + return _MultiFABWrapper(mf_name='pml_B_cp[x]', level=level, include_ghosts=include_ghosts) + +def ByCPPMLWrapper(level=0, include_ghosts=False): + return _MultiFABWrapper(mf_name='pml_B_cp[y]', level=level, include_ghosts=include_ghosts) + +def BzCPPMLWrapper(level=0, include_ghosts=False): + return _MultiFABWrapper(mf_name='pml_B_cp[z]', level=level, include_ghosts=include_ghosts) + +def JxCPPMLWrapper(level=0, include_ghosts=False): + return _MultiFABWrapper(mf_name='pml_j_cp[x]', level=level, include_ghosts=include_ghosts) + +def JyCPPMLWrapper(level=0, include_ghosts=False): + return _MultiFABWrapper(mf_name='pml_j_cp[y]', level=level, include_ghosts=include_ghosts) + +def JzCPPMLWrapper(level=0, include_ghosts=False): + return _MultiFABWrapper(mf_name='pml_j_cp[z]', level=level, include_ghosts=include_ghosts) + +def FCPPMLWrapper(level=0, include_ghosts=False): + return _MultiFABWrapper(mf_name='pml_F_cp', level=level, include_ghosts=include_ghosts) + +def GCPPMLWrapper(level=0, include_ghosts=False): + return _MultiFABWrapper(mf_name='pml_G_cp', level=level, include_ghosts=include_ghosts) diff --git a/Python/pywarpx/particle_containers.py b/Python/pywarpx/particle_containers.py new file mode 100644 index 000000000..69e6c4f8e --- /dev/null +++ b/Python/pywarpx/particle_containers.py @@ -0,0 +1,601 @@ +# Copyright 2017-2023 The WarpX Community +# +# This file is part of WarpX. +# +# Authors: David Grote, Roelof Groenewald +# +# License: BSD-3-Clause-LBNL + +import numpy as np + +from ._libwarpx import libwarpx + + +class ParticleContainerWrapper(object): + """Wrapper around particle containers. + This provides a convenient way to query and set data in the particle containers. + + Parameters + ---------- + species_name: string + The name of the species to be accessed. + """ + + def __init__(self, species_name): + self.name = species_name + + # grab the desired particle container + mypc = libwarpx.warpx.multi_particle_container() + self.particle_container = mypc.get_particle_container_from_name(self.name) + + def add_particles(self, x=None, y=None, z=None, ux=None, uy=None, + uz=None, w=None, unique_particles=True, **kwargs): + ''' + A function for adding particles to the WarpX simulation. + + Parameters + ---------- + + species_name : str + The type of species for which particles will be added + + x, y, z : arrays or scalars + The particle positions (default = 0.) + + ux, uy, uz : arrays or scalars + The particle momenta (default = 0.) + + w : array or scalars + Particle weights (default = 0.) + + unique_particles : bool + Whether the particles are unique or duplicated on several processes + (default = True) + + kwargs : dict + Containing an entry for all the extra particle attribute arrays. If + an attribute is not given it will be set to 0. + ''' + + # --- Get length of arrays, set to one for scalars + lenx = np.size(x) + leny = np.size(y) + lenz = np.size(z) + lenux = np.size(ux) + lenuy = np.size(uy) + lenuz = np.size(uz) + lenw = np.size(w) + + # --- Find the max length of the parameters supplied + maxlen = 0 + if x is not None: + maxlen = max(maxlen, lenx) + if y is not None: + maxlen = max(maxlen, leny) + if z is not None: + maxlen = max(maxlen, lenz) + if ux is not None: + maxlen = max(maxlen, lenux) + if uy is not None: + maxlen = max(maxlen, lenuy) + if uz is not None: + maxlen = max(maxlen, lenuz) + if w is not None: + maxlen = max(maxlen, lenw) + + # --- Make sure that the lengths of the input parameters are consistent + assert x is None or lenx==maxlen or lenx==1, "Length of x doesn't match len of others" + assert y is None or leny==maxlen or leny==1, "Length of y doesn't match len of others" + assert z is None or lenz==maxlen or lenz==1, "Length of z doesn't match len of others" + assert ux is None or lenux==maxlen or lenux==1, "Length of ux doesn't match len of others" + assert uy is None or lenuy==maxlen or lenuy==1, "Length of uy doesn't match len of others" + assert uz is None or lenuz==maxlen or lenuz==1, "Length of uz doesn't match len of others" + assert w is None or lenw==maxlen or lenw==1, "Length of w doesn't match len of others" + for key, val in kwargs.items(): + assert np.size(val)==1 or len(val)==maxlen, f"Length of {key} doesn't match len of others" + + # --- Broadcast scalars into appropriate length arrays + # --- If the parameter was not supplied, use the default value + if lenx == 1: + x = np.full(maxlen, (x or 0.)) + if leny == 1: + y = np.full(maxlen, (y or 0.)) + if lenz == 1: + z = np.full(maxlen, (z or 0.)) + if lenux == 1: + ux = np.full(maxlen, (ux or 0.)) + if lenuy == 1: + uy = np.full(maxlen, (uy or 0.)) + if lenuz == 1: + uz = np.full(maxlen, (uz or 0.)) + if lenw == 1: + w = np.full(maxlen, (w or 0.)) + for key, val in kwargs.items(): + if np.size(val) == 1: + kwargs[key] = np.full(maxlen, val) + + # --- The number of built in attributes + # --- The three velocities + built_in_attrs = 3 + if libwarpx.geometry_dim == 'rz': + # --- With RZ, there is also theta + built_in_attrs += 1 + + # --- The number of extra attributes (including the weight) + nattr = self.particle_container.num_real_comps() - built_in_attrs + attr = np.zeros((maxlen, nattr)) + attr[:,0] = w + + # --- Note that the velocities are handled separately and not included in attr + # --- (even though they are stored as attributes in the C++) + for key, vals in kwargs.items(): + attr[:,self.particle_container.get_comp_index(key) - built_in_attrs] = vals + + nattr_int = 0 + attr_int = np.empty([0], dtype=np.int32) + + # TODO: expose ParticleReal through pyAMReX + # and cast arrays to the correct types, before calling add_n_particles + # x = x.astype(self._numpy_particlereal_dtype, copy=False) + # y = y.astype(self._numpy_particlereal_dtype, copy=False) + # z = z.astype(self._numpy_particlereal_dtype, copy=False) + # ux = ux.astype(self._numpy_particlereal_dtype, copy=False) + # uy = uy.astype(self._numpy_particlereal_dtype, copy=False) + # uz = uz.astype(self._numpy_particlereal_dtype, copy=False) + + self.particle_container.add_n_particles( + 0, x.size, x, y, z, ux, uy, uz, + nattr, attr, nattr_int, attr_int, unique_particles + ) + + def get_particle_count(self, local=False): + ''' + Get the number of particles of this species in the simulation. + + Parameters + ---------- + + local : bool + If True the particle count on this processor will be returned. + Default False. + + Returns + ------- + + int + An integer count of the number of particles + ''' + return self.particle_container.total_number_of_particles(True, local) + nps = property(get_particle_count) + + def add_real_comp(self, pid_name, comm=True): + ''' + Add a real component to the particle data array. + + Parameters + ---------- + + pid_name : str + Name that can be used to identify the new component + + comm : bool + Should the component be communicated + ''' + self.particle_container.add_real_comp(pid_name, comm) + + def get_particle_structs(self, level): + ''' + This returns a list of numpy arrays containing the particle struct data + on each tile for this process. The particle data is represented as a structured + numpy array and contains the particle 'x', 'y', 'z', and 'idcpu'. + + The data for the numpy arrays are not copied, but share the underlying + memory buffer with WarpX. The numpy arrays are fully writeable. + + Parameters + ---------- + + level : int + The refinement level to reference + + Returns + ------- + + List of numpy arrays + The requested particle struct data + ''' + particle_data = [] + for pti in libwarpx.libwarpx_so.WarpXParIter(self.particle_container, level): + aos_arr = np.array(pti.aos(), copy=False) + particle_data.append(aos_arr) + return particle_data + + def get_particle_arrays(self, comp_name, level): + ''' + This returns a list of numpy arrays containing the particle array data + on each tile for this process. + + The data for the numpy arrays are not copied, but share the underlying + memory buffer with WarpX. The numpy arrays are fully writeable. + + Parameters + ---------- + + comp_name : str + The component of the array data that will be returned + + level : int + The refinement level to reference + + Returns + ------- + + List of numpy arrays + The requested particle array data + ''' + comp_idx = self.particle_container.get_comp_index(comp_name) + + data_array = [] + for pti in libwarpx.libwarpx_so.WarpXParIter(self.particle_container, level): + soa = pti.soa() + data_array.append(np.array(soa.GetRealData(comp_idx), copy=False)) + return data_array + + def get_particle_id(self, level=0): + ''' + + Return a list of numpy arrays containing the particle 'id' + numbers on each tile. + + ''' + structs = self.get_particle_structs(level) + return [libwarpx.amr.unpack_ids(struct['cpuid']) for struct in structs] + + def get_particle_cpu(self, level=0): + ''' + + Return a list of numpy arrays containing the particle 'cpu' + numbers on each tile. + + ''' + structs = self.get_particle_structs(level) + return [libwarpx.amr.unpack_cpus(struct['cpuid']) for struct in structs] + + def get_particle_x(self, level=0): + ''' + + Return a list of numpy arrays containing the particle 'x' + positions on each tile. + + ''' + structs = self.get_particle_structs(level) + if libwarpx.geometry_dim == '3d' or libwarpx.geometry_dim == '2d': + return [struct['x'] for struct in structs] + elif libwarpx.geometry_dim == 'rz': + return [struct['x']*np.cos(theta) for struct, theta in zip(structs, self.get_particle_theta(species_name))] + elif libwarpx.geometry_dim == '1d': + raise Exception('get_particle_x: There is no x coordinate with 1D Cartesian') + xp = property(get_particle_x) + + def get_particle_y(self, level=0): + ''' + + Return a list of numpy arrays containing the particle 'y' + positions on each tile. + + ''' + structs = self.get_particle_structs(level) + if libwarpx.geometry_dim == '3d': + return [struct['y'] for struct in structs] + elif libwarpx.geometry_dim == 'rz': + return [struct['x']*np.sin(theta) for struct, theta in zip(structs, self.get_particle_theta(species_name))] + elif libwarpx.geometry_dim == '1d' or libwarpx.geometry_dim == '2d': + raise Exception('get_particle_y: There is no y coordinate with 1D or 2D Cartesian') + yp = property(get_particle_y) + + def get_particle_r(self, level=0): + ''' + + Return a list of numpy arrays containing the particle 'r' + positions on each tile. + + ''' + structs = self.get_particle_structs(level) + if libwarpx.geometry_dim == 'rz': + return [struct['x'] for struct in structs] + elif libwarpx.geometry_dim == '3d': + return [np.sqrt(struct['x']**2 + struct['y']**2) for struct in structs] + elif libwarpx.geometry_dim == '2d' or libwarpx.geometry_dim == '1d': + raise Exception('get_particle_r: There is no r coordinate with 1D or 2D Cartesian') + rp = property(get_particle_r) + + def get_particle_theta(self, species_name, level=0): + ''' + + Return a list of numpy arrays containing the particle + theta on each tile. + + ''' + if libwarpx.geometry_dim == 'rz': + return self.get_particle_arrays('theta', level) + elif libwarpx.geometry_dim == '3d': + structs = self.get_particle_structs(level) + return [np.arctan2(struct['y'], struct['x']) for struct in structs] + elif libwarpx.geometry_dim == '2d' or libwarpx.geometry_dim == '1d': + raise Exception('get_particle_theta: There is no theta coordinate with 1D or 2D Cartesian') + thetap = property(get_particle_theta) + + def get_particle_z(self, level=0): + ''' + + Return a list of numpy arrays containing the particle 'z' + positions on each tile. + + ''' + structs = self.get_particle_structs(level) + if libwarpx.geometry_dim == '3d': + return [struct['z'] for struct in structs] + elif libwarpx.geometry_dim == 'rz' or libwarpx.geometry_dim == '2d': + return [struct['y'] for struct in structs] + elif libwarpx.geometry_dim == '1d': + return [struct['x'] for struct in structs] + zp = property(get_particle_z) + + def get_particle_weight(self, level=0): + ''' + + Return a list of numpy arrays containing the particle + weight on each tile. + + ''' + return self.get_particle_arrays('w', level) + wp = property(get_particle_weight) + + def get_particle_ux(self, level=0): + ''' + + Return a list of numpy arrays containing the particle + x momentum on each tile. + + ''' + return self.get_particle_arrays('ux', level) + uxp = property(get_particle_ux) + + def get_particle_uy(self, level=0): + ''' + + Return a list of numpy arrays containing the particle + y momentum on each tile. + + ''' + return self.get_particle_arrays('uy', level) + uyp = property(get_particle_uy) + + def get_particle_uz(self, level=0): + ''' + + Return a list of numpy arrays containing the particle + z momentum on each tile. + + ''' + + return self.get_particle_arrays('uz', level) + uzp = property(get_particle_uz) + + def get_species_charge_sum(self, local=False): + ''' + Returns the total charge in the simulation due to the given species. + + Parameters + ---------- + + local : bool + If True return total charge per processor + ''' + raise NotImplementedError() + return self.libwarpx_so.warpx_sumParticleCharge( + ctypes.c_char_p(species_name.encode('utf-8')), local + ) + + def getex(self): + raise NotImplementedError('Particle E fields not supported') + ex = property(getex) + + def getey(self): + raise NotImplementedError('Particle E fields not supported') + ey = property(getey) + + def getez(self): + raise NotImplementedError('Particle E fields not supported') + ez = property(getez) + + def getbx(self): + raise NotImplementedError('Particle B fields not supported') + bx = property(getbx) + + def getby(self): + raise NotImplementedError('Particle B fields not supported') + by = property(getby) + + def getbz(self): + raise NotImplementedError('Particle B fields not supported') + bz = property(getbz) + + +class ParticleBoundaryBufferWrapper(object): + """Wrapper around particle boundary buffer containers. + This provides a convenient way to query data in the particle boundary + buffer containers. + """ + + def __init__(self): + self.particle_buffer = libwarpx.warpx.get_particle_boundary_buffer() + + def get_particle_boundary_buffer_size(self, species_name, boundary, local=False): + ''' + This returns the number of particles that have been scraped so far in the simulation + from the specified boundary and of the specified species. + + Parameters + ---------- + + species_name : str + Return the number of scraped particles of this species + + boundary : str + The boundary from which to get the scraped particle data in the + form x/y/z_hi/lo + + local : bool + Whether to only return the number of particles in the current + processor's buffer + ''' + return self.particle_buffer.get_num_particles_in_container( + species_name, self._get_boundary_number(boundary), + local=local + ) + + def get_particle_boundary_buffer_structs(self, species_name, boundary, level): + ''' + This returns a list of numpy arrays containing the particle struct data + for a species that has been scraped by a specific simulation boundary. The + particle data is represented as a structured numpy array and contains the + particle 'x', 'y', 'z', and 'idcpu'. + + The data for the numpy arrays are not copied, but share the underlying + memory buffer with WarpX. The numpy arrays are fully writeable. + + Parameters + ---------- + + species_name : str + The species name that the data will be returned for + + boundary : str + The boundary from which to get the scraped particle data in the + form x/y/z_hi/lo or eb. + + level : int + Which AMR level to retrieve scraped particle data from. + ''' + + particles_per_tile = _LP_c_int() + num_tiles = ctypes.c_int(0) + data = self.libwarpx_so.warpx_getParticleBoundaryBufferStructs( + ctypes.c_char_p(species_name.encode('utf-8')), + self._get_boundary_number(boundary), level, + ctypes.byref(num_tiles), ctypes.byref(particles_per_tile) + ) + + particle_data = [] + for i in range(num_tiles.value): + if particles_per_tile[i] == 0: + continue + arr = self._array1d_from_pointer(data[i], self._p_dtype, particles_per_tile[i]) + particle_data.append(arr) + + _libc.free(particles_per_tile) + _libc.free(data) + return particle_data + + def get_particle_boundary_buffer(self, species_name, boundary, comp_name, level): + ''' + This returns a list of numpy arrays containing the particle array data + for a species that has been scraped by a specific simulation boundary. + + The data for the numpy arrays are not copied, but share the underlying + memory buffer with WarpX. The numpy arrays are fully writeable. + + Parameters + ---------- + + species_name : str + The species name that the data will be returned for. + + boundary : str + The boundary from which to get the scraped particle data in the + form x/y/z_hi/lo or eb. + + comp_name : str + The component of the array data that will be returned. If + "step_scraped" the special attribute holding the timestep at + which a particle was scraped will be returned. + + level : int + Which AMR level to retrieve scraped particle data from. + ''' + part_container = self.particle_buffer.get_particle_container( + species_name, self._get_boundary_number(boundary) + ) + data_array = [] + if comp_name == 'step_scraped': + # the step scraped is always the final integer component + comp_idx = part_container.num_int_comps() - 1 + for ii, pti in enumerate(libwarpx.libwarpx_so.BoundaryBufferParIter(part_container, level)): + soa = pti.soa() + data_array.append(np.array(soa.GetIntData(comp_idx), copy=False)) + else: + mypc = libwarpx.warpx.multi_particle_container() + sim_part_container_wrapper = mypc.get_particle_container_from_name(species_name) + comp_idx = sim_part_container_wrapper.get_comp_index(comp_name) + for ii, pti in enumerate(libwarpx.libwarpx_so.BoundaryBufferParIter(part_container, level)): + soa = pti.soa() + data_array.append(np.array(soa.GetRealData(comp_idx), copy=False)) + + return data_array + + def clear_buffer(self): + ''' + + Clear the buffer that holds the particles lost at the boundaries. + + ''' + self.particle_buffer.clear_particles() + + def _get_boundary_number(self, boundary): + ''' + + Utility function to find the boundary number given a boundary name. + + Parameters + ---------- + + boundary : str + The boundary from which to get the scraped particle data. In the + form x/y/z_hi/lo or eb. + + Returns + ------- + int + Integer index in the boundary scraper buffer for the given boundary. + ''' + if libwarpx.geometry_dim == '3d': + dimensions = {'x' : 0, 'y' : 1, 'z' : 2} + elif libwarpx.geometry_dim == '2d' or libwarpx.geometry_dim == 'rz': + dimensions = {'x' : 0, 'z' : 1} + elif libwarpx.geometry_dim == '1d': + dimensions = {'z' : 0} + else: + raise RuntimeError(f"Unknown simulation geometry: {libwarpx.geometry_dim}") + + if boundary != 'eb': + boundary_parts = boundary.split("_") + dim_num = dimensions[boundary_parts[0]] + if boundary_parts[1] == 'lo': + side = 0 + elif boundary_parts[1] == 'hi': + side = 1 + else: + raise RuntimeError(f'Unknown boundary specified: {boundary}') + boundary_num = 2 * dim_num + side + else: + if libwarpx.geometry_dim == '3d': + boundary_num = 6 + elif libwarpx.geometry_dim == '2d' or libwarpx.geometry_dim == 'rz': + boundary_num = 4 + elif libwarpx.geometry_dim == '1d': + boundary_num = 2 + else: + raise RuntimeError(f"Unknown simulation geometry: {libwarpx.geometry_dim}") + + return boundary_num |