diff options
Diffstat (limited to 'Python/pywarpx')
| -rw-r--r-- | Python/pywarpx/AMReX.py | 8 | ||||
| -rw-r--r-- | Python/pywarpx/PGroup.py | 197 | ||||
| -rw-r--r-- | Python/pywarpx/WarpX.py | 15 | ||||
| -rw-r--r-- | Python/pywarpx/__init__.py | 2 | ||||
| -rwxr-xr-x | Python/pywarpx/_libwarpx.py | 524 | ||||
| -rw-r--r-- | Python/pywarpx/timestepper.py | 52 |
6 files changed, 635 insertions, 163 deletions
diff --git a/Python/pywarpx/AMReX.py b/Python/pywarpx/AMReX.py index 444df3490..e84c240af 100644 --- a/Python/pywarpx/AMReX.py +++ b/Python/pywarpx/AMReX.py @@ -8,7 +8,9 @@ from .Langmuirwave import langmuirwave from .Interpolation import interpolation from .Particles import particles -from . import warpxC +import ctypes +from ._libwarpx import libwarpx +from ._libwarpx import amrex_init class AMReX(object): @@ -22,7 +24,7 @@ class AMReX(object): argv += interpolation.attrlist() argv += particles.attrlist() - warpxC.amrex_init(argv) + amrex_init(argv) def finalize(self, finalize_mpi=1): - warpxC.amrex_finalize(finalize_mpi) + libwarpx.amrex_finalize(finalize_mpi) diff --git a/Python/pywarpx/PGroup.py b/Python/pywarpx/PGroup.py index 616141702..bf59d5ce5 100644 --- a/Python/pywarpx/PGroup.py +++ b/Python/pywarpx/PGroup.py @@ -1,15 +1,14 @@ import numpy as np from . import WarpX -from . import warpxC +from . import _libwarpx class PGroup(object): """Implements a class that has the same API as a warp ParticleGroup instance. """ - def __init__(self): + def __init__(self, igroup): + self.igroup = igroup self.ns = 1 # Number of species - self.npmax = 0 # Size of data arrays - self.npid = 0 # number of columns for pid. self.gallot() @@ -24,10 +23,6 @@ class PGroup(object): self.sq = np.zeros(self.ns) # Species charge [C] self.sw = np.zeros(self.ns) # Species weight, (real particles per simulation particles) - self.ins = np.ones(self.ns, dtype=int) # Index of first particle in species - self.nps = np.zeros(self.ns, dtype=int) # Number of particles in species - self.ipmax = np.zeros(self.ns+1, dtype=int) # Max extent within the arrays of each species - 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) @@ -51,152 +46,104 @@ class PGroup(object): self.zshift = np.zeros(self.ns) self.gamma_ebcancel_max = np.ones(self.ns) # maximum value allowed for ExB cancellation - self.gaminv = np.ones(self.npmax) # inverse relativistic gamma factor - self._xp = np.zeros(self.npmax) # X-positions of particles [m] - self._yp = np.zeros(self.npmax) # Y-positions of particles [m] - self._zp = np.zeros(self.npmax) # Z-positions of particles [m] - self._uxp = np.zeros(self.npmax) # gamma * X-velocities of particles [m/s] - self._uyp = np.zeros(self.npmax) # gamma * Y-velocities of particles [m/s] - self._uzp = np.zeros(self.npmax) # gamma * Z-velocities of particles [m/s] - self._ex = np.zeros(self.npmax) # Ex of particles [V/m] - self._ey = np.zeros(self.npmax) # Ey of particles [V/m] - self._ez = np.zeros(self.npmax) # Ez of particles [V/m] - self._bx = np.zeros(self.npmax) # Bx of particles [T] - self._by = np.zeros(self.npmax) # By of particles [T] - self._bz = np.zeros(self.npmax) # Bz of particles [T] - self._pid = np.zeros((self.npmax, self.npid)) # Particle ID - used for various purposes - # --- Temporary fix gchange = gallot def allocated(self, name): - return (getattr(self, name, None) is not None) + return True def addspecies(self): pass - def _updatelocations(self): - warpx = WarpX.warpx.warpx - mypc = warpx.GetPartContainer() - - xplist = [] - yplist = [] - zplist = [] - for ispecie in range(mypc.nSpecies()): - pc = mypc.GetParticleContainer(ispecie) - xx = pc.getLocations() - xplist.append(xx[0,:]) - yplist.append(xx[1,:]) - zplist.append(xx[2,:]) - self.nps[ispecie] = len(xplist[-1]) - if ispecie > 0: - self.ins[ispecie] = self.ins[ispecie-1] + self.nps[ispecie-1] - self.ipmax[ispecie+1] = self.ins[ispecie] + self.nps[ispecie] - 1 - - self._xp = np.concatenate(xplist) - self._yp = np.concatenate(yplist) - self._zp = np.concatenate(zplist) - self.npmax = len(self._xp) - - def _updatevelocities(self): - warpx = WarpX.warpx.warpx - mypc = warpx.GetPartContainer() - - uxplist = [] - uyplist = [] - uzplist = [] - for ispecie in range(mypc.nSpecies()): - pc = mypc.GetParticleContainer(ispecie) - vv = pc.getData(0, 3) - uxplist.append(vv[0,:]) - uyplist.append(vv[1,:]) - uzplist.append(vv[2,:]) - self.nps[ispecie] = len(uxplist[-1]) - if ispecie > 0: - self.ins[ispecie] = self.ins[ispecie-1] + self.nps[ispecie-1] - self.ipmax[ispecie+1] = self.ins[ispecie] + self.nps[ispecie] - 1 - - self._uxp = np.concatenate(uxplist) - self._uyp = np.concatenate(uyplist) - self._uzp = np.concatenate(uzplist) - self.npmax = len(self._xp) - - def _updatepids(self): - warpx = WarpX.warpx.warpx - mypc = warpx.GetPartContainer() - - pidlist = [] - for ispecie in range(mypc.nSpecies()): - pc = mypc.GetParticleContainer(ispecie) - self.npid = pc.nAttribs - 3 - vv = pc.getData(3, self.npid) - pidlist.append(vv) - self.nps[ispecie] = len(uxplist[-1]) - if ispecie > 0: - self.ins[ispecie] = self.ins[ispecie-1] + self.nps[ispecie-1] - self.ipmax[ispecie+1] = self.ins[ispecie] + self.nps[ispecie] - 1 - - self._pid = np.concatenate(pidlist.T, axis=0) - self.npmax = self._pid.shape[0] - - def getxp(self): - self._updatelocations() - return self._xp + + 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, js=0): + return _libwarpx.get_particle_x(js)[self.igroup] xp = property(getxp) - def getyp(self): - self._updatelocations() - return self._yp + def getyp(self, js=0): + return _libwarpx.get_particle_y(js)[self.igroup] yp = property(getyp) - def getzp(self): - self._updatelocations() - return self._zp + def getzp(self, js=0): + return _libwarpx.get_particle_z(js)[self.igroup] zp = property(getzp) - def getuxp(self): - self._updatevelocities() - return self._uxp + def getuxp(self, js=0): + return _libwarpx.get_particle_ux(js)[self.igroup] uxp = property(getuxp) - def getuyp(self): - self._updatevelocities() - return self._uyp + def getuyp(self, js=0): + return _libwarpx.get_particle_uy(js)[self.igroup] uyp = property(getuyp) - def getuzp(self): - self._updatevelocities() - return self._uzp + def getuzp(self, js=0): + return _libwarpx.get_particle_uz(js)[self.igroup] uzp = property(getuzp) - def getpid(self): - self._updatepids() - return self._pid + def getpid(self, js=0): + return _libwarpx.get_particle_id(js)[self.igroup] pid = property(getpid) - def getgaminv(self): - uxp = self.uxp - uyp = self.uyp - uzp = self.uzp + def getgaminv(self, js=0): + uxp = self.getuxp(js) + uyp = self.getuyp(js) + uzp = self.getuzp(js) return sqrt(1. - (uxp**2 + uyp**2 + uzp**2)/warpxC.c**2) gaminv = property(getgaminv) - def getex(self): - return np.zeros(self.npmax) + def getex(self, js=0): + return _libwarpx.get_particle_Ex(js)[self.igroup] ex = property(getex) - def getey(self): - return np.zeros(self.npmax) + + def getey(self, js=0): + return _libwarpx.get_particle_Ey(js)[self.igroup] ey = property(getey) - def getez(self): - return np.zeros(self.npmax) + + def getez(self, js=0): + return _libwarpx.get_particle_Ez(js)[self.igroup] ez = property(getez) - def getbx(self): - return np.zeros(self.npmax) + + def getbx(self, js=0): + return _libwarpx.get_particle_Bx(js)[self.igroup] bx = property(getbx) - def getby(self): - return np.zeros(self.npmax) + + def getby(self, js=0): + return _libwarpx.get_particle_By(js)[self.igroup] by = property(getby) - def getbz(self): - return np.zeros(self.npmax) + + def getbz(self, js=0): + return _libwarpx.get_particle_Bz(js)[self.igroup] bz = property(getbz) +class PGroups(object): + def __init__(self): + xall = _libwarpx.get_particle_x(0) + self.ngroups = len(xall) + + self._pgroups = [] + for igroup in range(self.ngroups): + self._pgroups.append(PGroup(igroup)) + + def __iter__(self): + for igroup in range(self.ngroups): + yield self._pgroups[igroup] + + def __getitem__(self, key): + return self._pgroups[key] + diff --git a/Python/pywarpx/WarpX.py b/Python/pywarpx/WarpX.py index e9ec7cca4..f0a31f3f9 100644 --- a/Python/pywarpx/WarpX.py +++ b/Python/pywarpx/WarpX.py @@ -1,5 +1,5 @@ from .Bucket import Bucket -from . import warpxC +from ._libwarpx import libwarpx class WarpX(Bucket): """ @@ -7,21 +7,18 @@ class WarpX(Bucket): """ def init(self): - self.warpx = warpxC.WarpX.GetInstance() - self.warpx.InitData() + libwarpx.warpx_init() def evolve(self, nsteps=-1): - self.warpx.Evolve(nsteps) + libwarpx.warpx_evolve(nsteps) def finalize(self): - warpxC.WarpX.ResetInstance() + libwarpx.warpx_finalize() def getProbLo(self, direction): - return self.warpx.Geom()[0].ProbLo(direction) - #return warpxC.warpx_getProbLo(direction) + return libwarpx.warpx_getProbLo(direction) def getProbHi(self, direction): - return self.warpx.Geom()[0].ProbHi(direction) - #return warpxC.warpx_getProbHi(direction) + return libwarpx.warpx_getProbHi(direction) warpx = WarpX('warpx') diff --git a/Python/pywarpx/__init__.py b/Python/pywarpx/__init__.py index 4a509419e..cf79c65a6 100644 --- a/Python/pywarpx/__init__.py +++ b/Python/pywarpx/__init__.py @@ -10,4 +10,6 @@ from .AMReX import AMReX from .timestepper import TimeStepper from .PGroup import PGroup +from .PGroup import PGroups +from ._libwarpx import add_particles diff --git a/Python/pywarpx/_libwarpx.py b/Python/pywarpx/_libwarpx.py new file mode 100755 index 000000000..5174e018d --- /dev/null +++ b/Python/pywarpx/_libwarpx.py @@ -0,0 +1,524 @@ +# --- This defines the wrapper functions that directly call the underlying compiled routines +import os +import sys +import ctypes +from ctypes.util import find_library +import numpy as np +from numpy.ctypeslib import ndpointer + + +def get_package_root(): + ''' + Get the path to the installation location (where libwarpx.so would be installed). + ''' + cur = os.path.abspath(__file__) + while True: + name = os.path.basename(cur) + if name == 'pywarpx': + return cur + elif not name: + return '' + cur = os.path.dirname(cur) + +libwarpx = ctypes.CDLL(os.path.join(get_package_root(), "libwarpx.so")) +libc = ctypes.CDLL(find_library('c')) + + + +libwarpx.warpx_getistep.restype = ctypes.c_int +libwarpx.warpx_gett_new.restype = ctypes.c_double +libwarpx.warpx_getdt.restype = ctypes.c_double +libwarpx.warpx_maxStep.restype = ctypes.c_int +libwarpx.warpx_stopTime.restype = ctypes.c_double +libwarpx.warpx_checkInt.restype = ctypes.c_int +libwarpx.warpx_plotInt.restype = ctypes.c_int +libwarpx.warpx_finestLevel.restype = ctypes.c_int + + +libwarpx.warpx_EvolveE.argtypes = [ctypes.c_int, ctypes.c_double] +libwarpx.warpx_EvolveB.argtypes = [ctypes.c_int, ctypes.c_double] +libwarpx.warpx_FillBoundaryE.argtypes = [ctypes.c_int, ctypes.c_bool] +libwarpx.warpx_FillBoundaryB.argtypes = [ctypes.c_int, ctypes.c_bool] +libwarpx.warpx_PushParticlesandDepose.argtypes = [ctypes.c_int, ctypes.c_double] +libwarpx.warpx_getistep.argtypes = [ctypes.c_int] +libwarpx.warpx_setistep.argtypes = [ctypes.c_int, ctypes.c_int] +libwarpx.warpx_gett_new.argtypes = [ctypes.c_int] +libwarpx.warpx_sett_new.argtypes = [ctypes.c_int, ctypes.c_double] +libwarpx.warpx_getdt.argtypes = [ctypes.c_int] + + + +# some useful data structures and typenames +class Particle(ctypes.Structure): + _fields_ = [('x', ctypes.c_double), + ('y', ctypes.c_double), + ('z', ctypes.c_double), + ('id', ctypes.c_int), + ('cpu', ctypes.c_int)] + + +p_dtype = np.dtype([('x', 'f8'), ('y', 'f8'), ('z', 'f8'), + ('id', 'i4'), ('cpu', 'i4')]) + +c_double_p = ctypes.POINTER(ctypes.c_double) +LP_c_char = ctypes.POINTER(ctypes.c_char) +LP_LP_c_char = ctypes.POINTER(LP_c_char) + +# where do I import these? this might only work for CPython... +PyBuf_READ = 0x100 +PyBUF_WRITE = 0x200 + +# this is a function for converting a ctypes pointer to a numpy array +def array1d_from_pointer(pointer, dtype, size): + if sys.version_info.major >= 3: + 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) + + +# set the arg and return types of the wrapped functions +f = libwarpx.amrex_init +f.argtypes = (ctypes.c_int, LP_LP_c_char) + +f = libwarpx.warpx_getParticleStructs +f.restype = ctypes.POINTER(ctypes.POINTER(Particle)) + +f = libwarpx.warpx_getParticleArrays +f.restype = ctypes.POINTER(c_double_p) + +f = libwarpx.warpx_getParticleArrays +f.restype = ctypes.POINTER(c_double_p) + +f = libwarpx.warpx_getEfield +f.restype = ctypes.POINTER(c_double_p) + +f = libwarpx.warpx_getBfield +f.restype = ctypes.POINTER(c_double_p) + +f = libwarpx.warpx_getCurrentDensity +f.restype = ctypes.POINTER(c_double_p) + +f = libwarpx.addNParticles +f.argtypes = (ctypes.c_int, ctypes.c_int, + ndpointer(ctypes.c_double, flags="C_CONTIGUOUS"), + ndpointer(ctypes.c_double, flags="C_CONTIGUOUS"), + ndpointer(ctypes.c_double, flags="C_CONTIGUOUS"), + ndpointer(ctypes.c_double, flags="C_CONTIGUOUS"), + ndpointer(ctypes.c_double, flags="C_CONTIGUOUS"), + ndpointer(ctypes.c_double, flags="C_CONTIGUOUS"), + ctypes.c_int, + ndpointer(ctypes.c_double, flags="C_CONTIGUOUS"), + ctypes.c_int) + +def amrex_init(argv): + # --- Construct the ctype list of strings to pass in + argc = len(argv) + argvC = (LP_c_char * (argc+1))() + for i, arg in enumerate(argv): + enc_arg = arg.encode('utf-8') + argvC[i] = ctypes.create_string_buffer(enc_arg) + + libwarpx.amrex_init(argc, argvC) + +def add_particles(species_number, + x, y, z, ux, uy, uz, attr, unique_particles): + ''' + + A function for adding particles to the WarpX simulation. + + Parameters + ---------- + + species_number : the species to add the particle to + x, y, z : numpy arrays of the particle positions + ux, uy, uz : numpy arrays of the particle momenta + attr : a 2D numpy array with the particle attributes + unique_particles : whether the particles are unique or duplicated on + several processes + + ''' + libwarpx.addNParticles(species_number, x.size, + x, y, z, ux, uy, uz, + attr.shape[1], attr, unique_particles) + +def get_particle_structs(species_number): + ''' + + 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', 'id', and 'cpu'. + + 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_number : the species id that the data will be returned for + + Returns + ------- + + A List of numpy arrays. + + ''' + + particles_per_tile = ctypes.POINTER(ctypes.c_int)() + num_tiles = ctypes.c_int(0) + data = libwarpx.warpx_getParticleStructs(0, ctypes.byref(num_tiles), + ctypes.byref(particles_per_tile)) + + particle_data = [] + for i in range(num_tiles.value): + arr = array1d_from_pointer(data[i], 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(species_number, comp): + ''' + + 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_number : the species id that the data will be returned for + comp : the component of the array data that will be returned. + + Returns + ------- + + A List of numpy arrays. + + ''' + + particles_per_tile = ctypes.POINTER(ctypes.c_int)() + num_tiles = ctypes.c_int(0) + data = libwarpx.warpx_getParticleArrays(0, comp, ctypes.byref(num_tiles), + ctypes.byref(particles_per_tile)) + + particle_data = [] + for i in range(num_tiles.value): + arr = np.ctypeslib.as_array(data[i], (particles_per_tile[i],)) + arr.setflags(write=1) + particle_data.append(arr) + + libc.free(particles_per_tile) + libc.free(data) + return particle_data + + +def get_particle_x(species_number): + ''' + + Return a list of numpy arrays containing the particle 'x' + positions on each tile. + + ''' + structs = get_particle_structs(species_number) + return [struct['x'] for struct in structs] + + +def get_particle_y(species_number): + ''' + + Return a list of numpy arrays containing the particle 'y' + positions on each tile. + + ''' + structs = get_particle_structs(species_number) + return [struct['y'] for struct in structs] + + +def get_particle_z(species_number): + ''' + + Return a list of numpy arrays containing the particle 'z' + positions on each tile. + + ''' + structs = get_particle_structs(species_number) + return [struct['z'] for struct in structs] + + +def get_particle_id(species_number): + ''' + + Return a list of numpy arrays containing the particle 'z' + positions on each tile. + + ''' + structs = get_particle_structs(species_number) + return [struct['id'] for struct in structs] + + +def get_particle_cpu(species_number): + ''' + + Return a list of numpy arrays containing the particle 'z' + positions on each tile. + + ''' + structs = get_particle_structs(species_number) + return [struct['cpu'] for struct in structs] + + +def get_particle_weight(species_number): + ''' + + Return a list of numpy arrays containing the particle + weight on each tile. + + ''' + + return get_particle_arrays(species_number, 0) + + +def get_particle_ux(species_number): + ''' + + Return a list of numpy arrays containing the particle + x momentum on each tile. + + ''' + + return get_particle_arrays(species_number, 1) + + +def get_particle_uy(species_number): + ''' + + Return a list of numpy arrays containing the particle + y momentum on each tile. + + ''' + + return get_particle_arrays(species_number, 2) + + +def get_particle_uz(species_number): + ''' + + Return a list of numpy arrays containing the particle + z momentum on each tile. + + ''' + + return get_particle_arrays(species_number, 3) + + +def get_particle_Ex(species_number): + ''' + + Return a list of numpy arrays containing the particle + x electric field on each tile. + + ''' + + return get_particle_arrays(species_number, 4) + + +def get_particle_Ey(species_number): + ''' + + Return a list of numpy arrays containing the particle + y electric field on each tile. + + ''' + + return get_particle_arrays(species_number, 5) + + +def get_particle_Ez(species_number): + ''' + + Return a list of numpy arrays containing the particle + z electric field on each tile. + + ''' + + return get_particle_arrays(species_number, 6) + + +def get_particle_Bx(species_number): + ''' + + Return a list of numpy arrays containing the particle + x magnetic field on each tile. + + ''' + + return get_particle_arrays(species_number, 7) + + +def get_particle_By(species_number): + ''' + + Return a list of numpy arrays containing the particle + y magnetic field on each tile. + + ''' + + return get_particle_arrays(species_number, 8) + + +def get_particle_Bz(species_number): + ''' + + Return a list of numpy arrays containing the particle + z magnetic field on each tile. + + ''' + + return get_particle_arrays(species_number, 9) + + +def get_mesh_electric_field(level, direction, include_ghosts=True): + ''' + + This returns a list of numpy arrays containing the mesh electric field + 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. + + ''' + + shapes = ctypes.POINTER(ctypes.c_int)() + size = ctypes.c_int(0) + ngrow = ctypes.c_int(0) + data = libwarpx.warpx_getEfield(0, direction, + ctypes.byref(size), ctypes.byref(ngrow), + ctypes.byref(shapes)) + ng = ngrow.value + grid_data = [] + for i in range(size.value): + shape=(shapes[3*i+0], shapes[3*i+1], shapes[3*i+2]) + arr = np.ctypeslib.as_array(data[i], shape) + arr.setflags(write=1) + if include_ghosts: + grid_data.append(arr) + else: + grid_data.append(arr[ng:-ng,ng:-ng,ng:-ng]) + + libc.free(shapes) + libc.free(data) + return grid_data + + +def get_mesh_magnetic_field(level, direction, include_ghosts=True): + ''' + + This returns a list of numpy arrays containing the mesh magnetic field + 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. + + ''' + + shapes = ctypes.POINTER(ctypes.c_int)() + size = ctypes.c_int(0) + ngrow = ctypes.c_int(0) + data = libwarpx.warpx_getBfield(0, direction, + ctypes.byref(size), ctypes.byref(ngrow), + ctypes.byref(shapes)) + ng = ngrow.value + grid_data = [] + for i in range(size.value): + shape=(shapes[3*i+0], shapes[3*i+1], shapes[3*i+2]) + arr = np.ctypeslib.as_array(data[i], shape) + arr.setflags(write=1) + if include_ghosts: + grid_data.append(arr) + else: + grid_data.append(arr[ng:-ng,ng:-ng,ng:-ng]) + + libc.free(shapes) + libc.free(data) + return grid_data + + +def get_mesh_current_density(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. + + ''' + + shapes = ctypes.POINTER(ctypes.c_int)() + size = ctypes.c_int(0) + ngrow = ctypes.c_int(0) + data = libwarpx.warpx_getCurrentDensity(0, direction, + ctypes.byref(size), ctypes.byref(ngrow), + ctypes.byref(shapes)) + ng = ngrow.value + grid_data = [] + for i in range(size.value): + shape=(shapes[3*i+0], shapes[3*i+1], shapes[3*i+2]) + arr = np.ctypeslib.as_array(data[i], shape) + arr.setflags(write=1) + if include_ghosts: + grid_data.append(arr) + else: + grid_data.append(arr[ng:-ng,ng:-ng,ng:-ng]) + + libc.free(shapes) + libc.free(data) + return grid_data + + diff --git a/Python/pywarpx/timestepper.py b/Python/pywarpx/timestepper.py index 71787c7a4..1e3519dcc 100644 --- a/Python/pywarpx/timestepper.py +++ b/Python/pywarpx/timestepper.py @@ -2,18 +2,18 @@ import warp from warp import top from warp import w3d -from . import WarpX +from ._libwarpx import libwarpx class TimeStepper(object): + def step(self, nsteps=1): + for i in range(nsteps): + self.onestep() def onestep(self): - # --- A reference to the instance of the WarpX class - warpx = WarpX.warpx.warpx - mypc = warpx.GetPartContainer() - self.cur_time = warpx.gett_new(0) - self.istep = warpx.getistep(0) + self.cur_time = libwarpx.warpx_gett_new(0) + self.istep = libwarpx.warpx_getistep(0) #if mpi.rank == 0: print "\nSTEP %d starts ..."%(self.istep + 1) @@ -21,53 +21,52 @@ class TimeStepper(object): #if (ParallelDescriptor::NProcs() > 1) # if (okToRegrid(step)) RegridBaseLevel(); - warpx.ComputeDt() - dt = warpx.getdt(0) + libwarpx.warpx_ComputeDt() + dt = libwarpx.warpx_getdt(0) # --- Advance level 0 by dt lev = 0 # --- At the beginning, we have B^{n-1/2} and E^{n}. # --- Particles have p^{n-1/2} and x^{n}. - warpx.EvolveB(lev, 0.5*dt) # We now B^{n} + libwarpx.warpx_EvolveB(lev, 0.5*dt) # We now B^{n} - warpx.FillBoundaryE(lev, False) - warpx.FillBoundaryB(lev, False) + libwarpx.warpx_FillBoundaryE(lev, False) + libwarpx.warpx_FillBoundaryB(lev, False) # --- Evolve particles to p^{n+1/2} and x^{n+1} # --- Depose current, j^{n+1/2} - warpx.PushParticlesandDepose(lev, self.cur_time) + libwarpx.warpx_PushParticlesandDepose(lev, self.cur_time) - mypc.Redistribute() # Redistribute particles + libwarpx.mypc_Redistribute() # Redistribute particles - warpx.EvolveB(lev, 0.5*dt) # We now B^{n+1/2} + libwarpx.warpx_EvolveB(lev, 0.5*dt) # We now B^{n+1/2} - warpx.FillBoundaryB(lev, True) + libwarpx.warpx_FillBoundaryB(lev, True) - warpx.EvolveE(lev, dt) # We now have E^{n+1} + libwarpx.warpx_EvolveE(lev, dt) # We now have E^{n+1} self.istep += 1 self.cur_time += dt - warpx.MoveWindow(); + libwarpx.warpx_MoveWindow(); #if mpi.rank == 0: print "STEP %d ends. TIME = %e DT = %e"%(self.istep, self.cur_time, dt) # --- Sync up time - for i in range(warpx.finestLevel()+1): - warpx.sett_new(i, self.cur_time) - warpx.setistep(i, self.istep) + for i in range(libwarpx.warpx_finestLevel()+1): + libwarpx.warpx_sett_new(i, self.cur_time) + libwarpx.warpx_setistep(i, self.istep) - max_time_reached = ((self.cur_time >= warpx.stopTime() - 1.e-6*dt) or (self.istep >= warpx.maxStep())) + max_time_reached = ((self.cur_time >= libwarpx.warpx_stopTime() - 1.e-6*dt) or (self.istep >= libwarpx.warpx_maxStep())) - if warpx.plotInt() > 0 and (self.istep+1)%warpx.plotInt() == 0 or max_time_reached: - warpx.WritePlotFile() - - if warpx.checkInt() > 0 and (self.istep+1)%warpx.plotInt() == 0 or max_time_reached: - warpx.WriteCheckPointFile() + if libwarpx.warpx_plotInt() > 0 and (self.istep+1)%libwarpx.warpx_plotInt() == 0 or max_time_reached: + libwarpx.warpx_WritePlotFile() + if libwarpx.warpx_checkInt() > 0 and (self.istep+1)%libwarpx.warpx_plotInt() == 0 or max_time_reached: + libwarpx.warpx_WriteCheckPointFile() @@ -80,6 +79,7 @@ class TimeStepper(object): +# --- This is not used class TimeStepperFromPICSAR(object): def __init__(self, package=None, solver=None, l_debug=False): |