diff options
Diffstat (limited to 'Python')
| -rw-r--r-- | Python/pywarpx/PGroup.py | 8 | ||||
| -rw-r--r-- | Python/pywarpx/WarpInterface.py | 171 | ||||
| -rwxr-xr-x | Python/pywarpx/_libwarpx.py | 289 | ||||
| -rw-r--r-- | Python/pywarpx/fields.py | 340 | ||||
| -rw-r--r-- | Python/pywarpx/picmi.py | 2 |
5 files changed, 535 insertions, 275 deletions
diff --git a/Python/pywarpx/PGroup.py b/Python/pywarpx/PGroup.py index 68b37740d..48e68ceb5 100644 --- a/Python/pywarpx/PGroup.py +++ b/Python/pywarpx/PGroup.py @@ -83,6 +83,10 @@ class PGroup(object): return _libwarpx.get_particle_y(self.ispecie)[self.igroup] yp = property(getyp) + def getrp(self): + return _libwarpx.get_particle_r(self.ispecie)[self.igroup] + rp = property(getrp) + def getzp(self): return _libwarpx.get_particle_z(self.ispecie)[self.igroup] zp = property(getzp) @@ -136,6 +140,10 @@ class PGroup(object): return _libwarpx.get_particle_Bz(self.ispecie)[self.igroup] bz = property(getbz) + def gettheta(self): + return _libwarpx.get_particle_theta(self.ispecie)[self.igroup] + theta = property(gettheta) + class PGroups(object): def __init__(self, ispecie=0): self.ispecie = ispecie diff --git a/Python/pywarpx/WarpInterface.py b/Python/pywarpx/WarpInterface.py new file mode 100644 index 000000000..f84c22a5d --- /dev/null +++ b/Python/pywarpx/WarpInterface.py @@ -0,0 +1,171 @@ +import warp +from . import fields +from pywarpx import PGroup + + +def warp_species(warp_type, picmie_species): + """Returns a Warp species that has a reference to the WarpX particles. + """ + elec_pgroups = PGroup.PGroups(ispecie=picmie_species.species_number) + return warp.Species(type=warp_type, pgroups=elec_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/_libwarpx.py b/Python/pywarpx/_libwarpx.py index 4c3283b97..2a37dbd5e 100755 --- a/Python/pywarpx/_libwarpx.py +++ b/Python/pywarpx/_libwarpx.py @@ -1,6 +1,7 @@ # --- This defines the wrapper functions that directly call the underlying compiled routines import os import sys +import atexit import ctypes from ctypes.util import find_library as _find_library import numpy as np @@ -8,6 +9,13 @@ 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 (though don't know why) + from mpi4py import MPI +except ImportError: + pass + # --- Is there a better way of handling constants? clight = 2.99792458e+8 # m/s @@ -184,6 +192,7 @@ def initialize(argv=None): libwarpx.warpx_init() +@atexit.register def finalize(finalize_mpi=1): ''' @@ -399,7 +408,10 @@ def get_particle_x(species_number): ''' structs = get_particle_structs(species_number) - return [struct['x'] for struct in structs] + if geometry_dim == '3d' or geometry_dim == '2d': + return [struct['x'] for struct in structs] + elif geometry_dim == 'rz': + return [struct['x']*np.cos(theta) for struct, theta in zip(structs, get_particle_theta(species_number))] def get_particle_y(species_number): @@ -410,7 +422,26 @@ def get_particle_y(species_number): ''' structs = get_particle_structs(species_number) - return [struct['y'] for struct in structs] + if geometry_dim == '3d' or geometry_dim == '2d': + return [struct['y'] for struct in structs] + elif geometry_dim == 'rz': + return [struct['x']*np.sin(theta) for struct, theta in zip(structs, get_particle_theta(species_number))] + + +def get_particle_r(species_number): + ''' + + Return a list of numpy arrays containing the particle 'r' + positions on each tile. + + ''' + structs = get_particle_structs(species_number) + if geometry_dim == 'rz': + return [struct['x'] for struct in structs] + elif geometry_dim == '3d': + return [np.sqrt(struct['x']**2 + struct['y']**2) for struct in structs] + elif geometry_dim == '2d': + raise Exception('get_particle_r: There is no r coordinate with 2D Cartesian') def get_particle_z(species_number): @@ -556,6 +587,53 @@ def get_particle_Bz(species_number): return get_particle_arrays(species_number, 9) +def get_particle_theta(species_number): + ''' + + Return a list of numpy arrays containing the particle + theta on each tile. + + ''' + + if geometry_dim == 'rz': + return get_particle_arrays(species_number, 10) + elif geometry_dim == '3d': + return [np.arctan2(struct['y'], struct['x']) for struct in structs] + elif geometry_dim == '2d': + raise Exception('get_particle_r: There is no theta coordinate with 2D Cartesian') + + +def _get_mesh_field_list(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) + ngrow = ctypes.c_int(0) + data = warpx_func(level, direction, + ctypes.byref(size), ctypes.byref(ncomps), + ctypes.byref(ngrow), ctypes.byref(shapes)) + ng = ngrow.value + grid_data = [] + shapesize = 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. + arr = np.ctypeslib.as_array(data[i], shape[::-1]).T + arr.setflags(write=1) + if include_ghosts: + grid_data.append(arr) + else: + grid_data.append(arr[tuple([slice(ng, -ng) for _ in range(dim)])]) + + _libc.free(shapes) + _libc.free(data) + return grid_data + + def get_mesh_electric_field(level, direction, include_ghosts=True): ''' @@ -582,28 +660,7 @@ def get_mesh_electric_field(level, direction, include_ghosts=True): ''' assert(level == 0) - - shapes = _LP_c_int() - size = ctypes.c_int(0) - ngrow = ctypes.c_int(0) - data = libwarpx.warpx_getEfield(level, direction, - ctypes.byref(size), ctypes.byref(ngrow), - ctypes.byref(shapes)) - ng = ngrow.value - grid_data = [] - for i in range(size.value): - shape = tuple([shapes[dim*i + d] for d in range(dim)]) - # --- The data is stored in Fortran order, hence shape is reversed and a transpose is taken. - arr = np.ctypeslib.as_array(data[i], shape[::-1]).T - arr.setflags(write=1) - if include_ghosts: - grid_data.append(arr) - else: - grid_data.append(arr[[slice(ng, -ng) for _ in range(dim)]]) - - _libc.free(shapes) - _libc.free(data) - return grid_data + return _get_mesh_field_list(libwarpx.warpx_getEfield, level, direction, include_ghosts) def get_mesh_electric_field_cp(level, direction, include_ghosts=True): @@ -631,28 +688,7 @@ def get_mesh_electric_field_cp(level, direction, include_ghosts=True): ''' assert(level == 0) - - shapes = _LP_c_int() - size = ctypes.c_int(0) - ngrow = ctypes.c_int(0) - data = libwarpx.warpx_getEfieldCP(level, direction, - ctypes.byref(size), ctypes.byref(ngrow), - ctypes.byref(shapes)) - ng = ngrow.value - grid_data = [] - for i in range(size.value): - shape = tuple([shapes[dim*i + d] for d in range(dim)]) - # --- The data is stored in Fortran order, hence shape is reversed and a transpose is taken. - arr = np.ctypeslib.as_array(data[i], shape[::-1]).T - arr.setflags(write=1) - if include_ghosts: - grid_data.append(arr) - else: - grid_data.append(arr[[slice(ng, -ng) for _ in range(dim)]]) - - _libc.free(shapes) - _libc.free(data) - return grid_data + return _get_mesh_field_list(libwarpx.warpx_getEfieldCP, level, direction, include_ghosts) def get_mesh_electric_field_fp(level, direction, include_ghosts=True): @@ -680,28 +716,7 @@ def get_mesh_electric_field_fp(level, direction, include_ghosts=True): ''' assert(level == 0) - - shapes = _LP_c_int() - size = ctypes.c_int(0) - ngrow = ctypes.c_int(0) - data = libwarpx.warpx_getEfieldFP(level, direction, - ctypes.byref(size), ctypes.byref(ngrow), - ctypes.byref(shapes)) - ng = ngrow.value - grid_data = [] - for i in range(size.value): - shape = tuple([shapes[dim*i + d] for d in range(dim)]) - # --- The data is stored in Fortran order, hence shape is reversed and a transpose is taken. - arr = np.ctypeslib.as_array(data[i], shape[::-1]).T - arr.setflags(write=1) - if include_ghosts: - grid_data.append(arr) - else: - grid_data.append(arr[[slice(ng, -ng) for _ in range(dim)]]) - - _libc.free(shapes) - _libc.free(data) - return grid_data + return _get_mesh_field_list(libwarpx.warpx_getEfieldFP, level, direction, include_ghosts) def get_mesh_magnetic_field(level, direction, include_ghosts=True): @@ -730,28 +745,7 @@ def get_mesh_magnetic_field(level, direction, include_ghosts=True): ''' assert(level == 0) - - shapes = _LP_c_int() - size = ctypes.c_int(0) - ngrow = ctypes.c_int(0) - data = libwarpx.warpx_getBfield(level, direction, - ctypes.byref(size), ctypes.byref(ngrow), - ctypes.byref(shapes)) - ng = ngrow.value - grid_data = [] - for i in range(size.value): - shape = tuple([shapes[dim*i + d] for d in range(dim)]) - # --- The data is stored in Fortran order, hence shape is reversed and a transpose is taken. - arr = np.ctypeslib.as_array(data[i], shape[::-1]).T - arr.setflags(write=1) - if include_ghosts: - grid_data.append(arr) - else: - grid_data.append(arr[[slice(ng, -ng) for _ in range(dim)]]) - - _libc.free(shapes) - _libc.free(data) - return grid_data + return _get_mesh_field_list(libwarpx.warpx_getBfield, level, direction, include_ghosts) def get_mesh_magnetic_field_cp(level, direction, include_ghosts=True): @@ -779,28 +773,7 @@ def get_mesh_magnetic_field_cp(level, direction, include_ghosts=True): ''' assert(level == 0) - - shapes = _LP_c_int() - size = ctypes.c_int(0) - ngrow = ctypes.c_int(0) - data = libwarpx.warpx_getBfieldCP(level, direction, - ctypes.byref(size), ctypes.byref(ngrow), - ctypes.byref(shapes)) - ng = ngrow.value - grid_data = [] - for i in range(size.value): - shape = tuple([shapes[dim*i + d] for d in range(dim)]) - # --- The data is stored in Fortran order, hence shape is reversed and a transpose is taken. - arr = np.ctypeslib.as_array(data[i], shape[::-1]).T - arr.setflags(write=1) - if include_ghosts: - grid_data.append(arr) - else: - grid_data.append(arr[[slice(ng, -ng) for _ in range(dim)]]) - - _libc.free(shapes) - _libc.free(data) - return grid_data + return _get_mesh_field_list(libwarpx.warpx_getBfieldCP, level, direction, include_ghosts) def get_mesh_magnetic_field_fp(level, direction, include_ghosts=True): @@ -828,28 +801,7 @@ def get_mesh_magnetic_field_fp(level, direction, include_ghosts=True): ''' assert(level == 0) - - shapes = _LP_c_int() - size = ctypes.c_int(0) - ngrow = ctypes.c_int(0) - data = libwarpx.warpx_getBfieldFP(level, direction, - ctypes.byref(size), ctypes.byref(ngrow), - ctypes.byref(shapes)) - ng = ngrow.value - grid_data = [] - for i in range(size.value): - shape = tuple([shapes[dim*i + d] for d in range(dim)]) - # --- The data is stored in Fortran order, hence shape is reversed and a transpose is taken. - arr = np.ctypeslib.as_array(data[i], shape[::-1]).T - arr.setflags(write=1) - if include_ghosts: - grid_data.append(arr) - else: - grid_data.append(arr[[slice(ng, -ng) for _ in range(dim)]]) - - _libc.free(shapes) - _libc.free(data) - return grid_data + return _get_mesh_field_list(libwarpx.warpx_getBfieldFP, level, direction, include_ghosts) def get_mesh_current_density(level, direction, include_ghosts=True): @@ -876,28 +828,7 @@ def get_mesh_current_density(level, direction, include_ghosts=True): ''' assert(level == 0) - - shapes = _LP_c_int() - size = ctypes.c_int(0) - ngrow = ctypes.c_int(0) - data = libwarpx.warpx_getCurrentDensity(level, direction, - ctypes.byref(size), ctypes.byref(ngrow), - ctypes.byref(shapes)) - ng = ngrow.value - grid_data = [] - for i in range(size.value): - shape = tuple([shapes[dim*i + d] for d in range(dim)]) - # --- The data is stored in Fortran order, hence shape is reversed and a transpose is taken. - arr = np.ctypeslib.as_array(data[i], shape[::-1]).T - arr.setflags(write=1) - if include_ghosts: - grid_data.append(arr) - else: - grid_data.append(arr[[slice(ng, -ng) for _ in range(dim)]]) - - _libc.free(shapes) - _libc.free(data) - return grid_data + return _get_mesh_field_list(libwarpx.warpx_getCurrentDensity, level, direction, include_ghosts) def get_mesh_current_density_cp(level, direction, include_ghosts=True): @@ -925,28 +856,7 @@ def get_mesh_current_density_cp(level, direction, include_ghosts=True): ''' assert(level == 0) - - shapes = _LP_c_int() - size = ctypes.c_int(0) - ngrow = ctypes.c_int(0) - data = libwarpx.warpx_getCurrentDensityCP(level, direction, - ctypes.byref(size), ctypes.byref(ngrow), - ctypes.byref(shapes)) - ng = ngrow.value - grid_data = [] - for i in range(size.value): - shape = tuple([shapes[dim*i + d] for d in range(dim)]) - # --- The data is stored in Fortran order, hence shape is reversed and a transpose is taken. - arr = np.ctypeslib.as_array(data[i], shape[::-1]).T - arr.setflags(write=1) - if include_ghosts: - grid_data.append(arr) - else: - grid_data.append(arr[[slice(ng, -ng) for _ in range(dim)]]) - - _libc.free(shapes) - _libc.free(data) - return grid_data + return _get_mesh_field_list(libwarpx.warpx_getCurrentDensityCP, level, direction, include_ghosts) def get_mesh_current_density_fp(level, direction, include_ghosts=True): @@ -974,28 +884,7 @@ def get_mesh_current_density_fp(level, direction, include_ghosts=True): ''' assert(level == 0) - - shapes = _LP_c_int() - size = ctypes.c_int(0) - ngrow = ctypes.c_int(0) - data = libwarpx.warpx_getCurrentDensityFP(level, direction, - ctypes.byref(size), ctypes.byref(ngrow), - ctypes.byref(shapes)) - ng = ngrow.value - grid_data = [] - for i in range(size.value): - shape = tuple([shapes[dim*i + d] for d in range(dim)]) - # --- The data is stored in Fortran order, hence shape is reversed and a transpose is taken. - arr = np.ctypeslib.as_array(data[i], shape[::-1]).T - arr.setflags(write=1) - if include_ghosts: - grid_data.append(arr) - else: - grid_data.append(arr[[slice(ng, -ng) for _ in range(dim)]]) - - _libc.free(shapes) - _libc.free(data) - return grid_data + return _get_mesh_field_list(libwarpx.warpx_getCurrentDensityFP, level, direction, include_ghosts) def _get_mesh_array_lovects(level, direction, include_ghosts=True, getarrayfunc=None): diff --git a/Python/pywarpx/fields.py b/Python/pywarpx/fields.py index 8237ed867..9068a22ea 100644 --- a/Python/pywarpx/fields.py +++ b/Python/pywarpx/fields.py @@ -11,7 +11,7 @@ import numpy as np from . import _libwarpx -class MultiFABWrapper(object): +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. The indexing is based on global indices. @@ -21,16 +21,23 @@ class MultiFABWrapper(object): - get_fabs: routine that returns the list of FABs - level=0: ignored """ - def __init__(self, direction, overlaps, get_lovects, get_fabs, level=0): + def __init__(self, direction, overlaps, get_lovects, get_fabs, level=0, include_ghosts=False): self.direction = direction self.overlaps = np.array(overlaps) self.get_lovects = get_lovects self.get_fabs = get_fabs self.level = 0 #level - self.include_ghosts = False + self.include_ghosts = include_ghosts + + self.dim = _libwarpx.dim + if self.dim == 2: + # --- Grab the first and last overlaps (for x and z) + self.overlaps = self.overlaps[::2] def _getlovects(self): - return self.get_lovects(self.level, self.direction, self.include_ghosts) + lovects = self.get_lovects(self.level, self.direction, self.include_ghosts) + self.nghosts = -lovects.min() + return lovects def _gethivects(self): lovects = self._getlovects() @@ -38,7 +45,7 @@ class MultiFABWrapper(object): hivects = np.zeros_like(lovects) for i in range(len(fields)): - hivects[:,i] = lovects[:,i] + np.array(fields[i].shape) - self.overlaps + hivects[:,i] = lovects[:,i] + np.array(fields[i].shape[:self.dim]) - self.overlaps return hivects @@ -50,14 +57,29 @@ class MultiFABWrapper(object): 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 - 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 domain. + 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. """ if index == Ellipsis: - index = (slice(None), slice(None), slice(None)) - + index = tuple(self.dim*[slice(None)]) + + if len(index) < self.dim: + # --- Add extra dims to index if needed + index = list(index) + for i in range(len(index), self.dim): + index.append(slice(None)) + index = tuple(index) + + if self.dim == 2: + return self._getitem2d(index) + elif self.dim == 3: + return self._getitem3d(index) + + def _getitem3d(self, index): + """Returns slices of a 3D decomposed array, + """ ix = index[0] iy = index[1] iz = index[2] @@ -66,25 +88,25 @@ class MultiFABWrapper(object): hivects = self._gethivects() fields = self._getfields() - nx = hivects[0,:].max() - ny = hivects[1,:].max() - nz = hivects[2,:].max() + nx = hivects[0,:].max() - self.nghosts + ny = hivects[1,:].max() - self.nghosts + nz = hivects[2,:].max() - self.nghosts if isinstance(ix, slice): - ixstart = max(ix.start, 0) - ixstop = min(ix.stop or nx + 1, nx + self.overlaps[0]) + ixstart = max(ix.start or -self.nghosts, -self.nghosts) + ixstop = min(ix.stop or nx + 1 + self.nghosts, nx + self.overlaps[0] + self.nghosts) else: ixstart = ix ixstop = ix + 1 if isinstance(iy, slice): - iystart = max(iy.start, 0) - iystop = min(iy.stop or ny + 1, ny + self.overlaps[1]) + iystart = max(iy.start or -self.nghosts, -self.nghosts) + iystop = min(iy.stop or ny + 1 + self.nghosts, ny + self.overlaps[1] + self.nghosts) else: iystart = iy iystop = iy + 1 if isinstance(iz, slice): - izstart = max(iz.start, 0) - izstop = min(iz.stop or nz + 1, nz + self.overlaps[2]) + izstart = max(iz.start or -self.nghosts, -self.nghosts) + izstop = min(iz.stop or nz + 1 + self.nghosts, nz + self.overlaps[2] + self.nghosts) else: izstart = iz izstop = iz + 1 @@ -99,11 +121,11 @@ class MultiFABWrapper(object): # --- The ix1, 2 etc are relative to global indexing ix1 = max(ixstart, lovects[0,i]) - ix2 = min((ixstop or nx+1), lovects[0,i] + fields[i].shape[0]) + ix2 = min(ixstop, lovects[0,i] + fields[i].shape[0]) iy1 = max(iystart, lovects[1,i]) - iy2 = min((iystop or ny+1), lovects[1,i] + fields[i].shape[1]) + iy2 = min(iystop, lovects[1,i] + fields[i].shape[1]) iz1 = max(izstart, lovects[2,i]) - iz2 = min((izstop or nz+1), lovects[2,i] + fields[i].shape[2]) + iz2 = min(izstop, lovects[2,i] + fields[i].shape[2]) if ix1 < ix2 and iy1 < iy2 and iz1 < iz2: @@ -126,7 +148,84 @@ class MultiFABWrapper(object): if not isinstance(iz, slice): sss[2] = 0 - return resultglobal[sss] + return resultglobal[tuple(sss)] + + def _getitem2d(self, index): + """Returns slices of a 2D decomposed array, + """ + + lovects = self._getlovects() + hivects = self._gethivects() + fields = self._getfields() + + ix = index[0] + iz = index[1] + + 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[2] + else: + raise Exception('Too many indices given') + else: + ic = None + + nx = hivects[0,:].max() - self.nghosts + nz = hivects[1,:].max() - self.nghosts + + if isinstance(ix, slice): + ixstart = max(ix.start or -self.nghosts, -self.nghosts) + ixstop = min(ix.stop or nx + 1 + self.nghosts, nx + self.overlaps[0] + self.nghosts) + else: + ixstart = ix + ixstop = ix + 1 + if isinstance(iz, slice): + izstart = max(iz.start or -self.nghosts, -self.nghosts) + izstop = min(iz.stop or nz + 1 + self.nghosts, nz + self.overlaps[1] + self.nghosts) + else: + izstart = iz + izstop = iz + 1 + + # --- Setup the size of the array to be returned and create it. + # --- Space is added for multiple components if needed. + sss = (max(0, ixstop - ixstart), + max(0, izstop - izstart)) + if ncomps > 1 and ic is None: + sss = tuple(list(sss) + [ncomps]) + resultglobal = np.zeros(sss) + + for i in range(len(fields)): + + # --- The ix1, 2 etc are relative to global indexing + ix1 = max(ixstart, lovects[0,i]) + ix2 = min(ixstop, lovects[0,i] + fields[i].shape[0]) + iz1 = max(izstart, lovects[1,i]) + iz2 = min(izstop, lovects[1,i] + fields[i].shape[1]) + + if ix1 < ix2 and iz1 < iz2: + + sss = (slice(ix1 - lovects[0,i], ix2 - lovects[0,i]), + slice(iz1 - lovects[1,i], iz2 - lovects[1,i])) + if ic is not None: + sss = tuple(list(sss) + [ic]) + + vslice = (slice(ix1 - ixstart, ix2 - ixstart), + slice(iz1 - izstart, iz2 - izstart)) + + resultglobal[vslice] = fields[i][sss] + + # --- Now remove any of the reduced dimensions. + sss = [slice(None), slice(None)] + if not isinstance(ix, slice): + sss[0] = 0 + if not isinstance(iz, slice): + sss[1] = 0 + + return resultglobal[tuple(sss)] def __setitem__(self, index, value): """Sets slices of a decomposed array. The shape of @@ -135,8 +234,23 @@ class MultiFABWrapper(object): - value: input array (must be supplied) """ if index == Ellipsis: - index = (slice(None), slice(None), slice(None)) - + index = tuple(self.dim*[slice(None)]) + + if len(index) < self.dim: + # --- Add extra dims to index if needed + index = list(index) + for i in range(len(index), self.dim): + index.append(slice(None)) + index = tuple(index) + + if self.dim == 2: + return self._setitem2d(index, value) + elif self.dim == 3: + return self._setitem3d(index, value) + + def _setitem3d(self, index, value): + """Sets slices of a decomposed 3D array. + """ ix = index[0] iy = index[1] iz = index[2] @@ -145,9 +259,9 @@ class MultiFABWrapper(object): hivects = self._gethivects() fields = self._getfields() - nx = hivects[0,:].max() - ny = hivects[1,:].max() - nz = hivects[2,:].max() + nx = hivects[0,:].max() - self.nghosts + ny = hivects[1,:].max() - self.nghosts + nz = hivects[2,:].max() - self.nghosts # --- Add extra dimensions so that the input has the same number of # --- dimensions as array. @@ -160,20 +274,20 @@ class MultiFABWrapper(object): value3d.shape = sss if isinstance(ix, slice): - ixstart = max(ix.start, 0) - ixstop = min(ix.stop or nx + 1, nx + self.overlaps[0]) + ixstart = max(ix.start or -self.nghosts, -self.nghosts) + ixstop = min(ix.stop or nx + 1 + self.nghosts, nx + self.overlaps[0] + self.nghosts) else: ixstart = ix ixstop = ix + 1 if isinstance(iy, slice): - iystart = max(iy.start, 0) - iystop = min(iy.stop or ny + 1, ny + self.overlaps[1]) + iystart = max(iy.start or -self.nghosts, -self.nghosts) + iystop = min(iy.stop or ny + 1 + self.nghosts, ny + self.overlaps[1] + self.nghosts) else: iystart = iy iystop = iy + 1 if isinstance(iz, slice): - izstart = max(iz.start, 0) - izstop = min(iz.stop or nz + 1, nz + self.overlaps[2]) + izstart = max(iz.start or -self.nghosts, -self.nghosts) + izstop = min(iz.stop or nz + 1 + self.nghosts, nz + self.overlaps[2] + self.nghosts) else: izstart = iz izstop = iz + 1 @@ -182,11 +296,11 @@ class MultiFABWrapper(object): # --- The ix1, 2 etc are relative to global indexing ix1 = max(ixstart, lovects[0,i]) - ix2 = min((ixstop or nx+1), lovects[0,i] + fields[i].shape[0]) + ix2 = min(ixstop, lovects[0,i] + fields[i].shape[0]) iy1 = max(iystart, lovects[1,i]) - iy2 = min((iystop or ny+1), lovects[1,i] + fields[i].shape[1]) + iy2 = min(iystop, lovects[1,i] + fields[i].shape[1]) iz1 = max(izstart, lovects[2,i]) - iz2 = min((izstop or nz+1), lovects[2,i] + fields[i].shape[2]) + iz2 = min(izstop, lovects[2,i] + fields[i].shape[2]) if ix1 < ix2 and iy1 < iy2 and iz1 < iz2: @@ -202,49 +316,127 @@ class MultiFABWrapper(object): else: fields[i][sss] = value + def _setitem2d(self, index, value): + """Sets slices of a decomposed 2D array. + """ + ix = index[0] + iz = index[2] -def ExWrapper(level=0): - return MultiFABWrapper(direction=0, overlaps=[0,1,1], - get_lovects=_libwarpx.get_mesh_electric_field_lovects, - get_fabs=_libwarpx.get_mesh_electric_field, level=level) + lovects = self._getlovects() + hivects = self._gethivects() + fields = self._getfields() -def EyWrapper(level=0): - return MultiFABWrapper(direction=1, overlaps=[1,0,1], - get_lovects=_libwarpx.get_mesh_electric_field_lovects, - get_fabs=_libwarpx.get_mesh_electric_field, level=level) + if len(fields[0].shape) > self.dim: + ncomps = fields[0].shape[-1] + else: + ncomps = 1 -def EzWrapper(level=0): - return MultiFABWrapper(direction=2, overlaps=[1,1,0], - get_lovects=_libwarpx.get_mesh_electric_field_lovects, - get_fabs=_libwarpx.get_mesh_electric_field, level=level) + if len(index) > self.dim: + if ncomps > 1: + ic = index[2] + else: + raise Exception('Too many indices given') + else: + ic = None -def BxWrapper(level=0): - return MultiFABWrapper(direction=0, overlaps=[1,0,0], - get_lovects=_libwarpx.get_mesh_magnetic_field_lovects, - get_fabs=_libwarpx.get_mesh_magnetic_field, level=level) + nx = hivects[0,:].max() - self.nghosts + nz = hivects[2,:].max() - self.nghosts -def ByWrapper(level=0): - return MultiFABWrapper(direction=1, overlaps=[0,1,0], - get_lovects=_libwarpx.get_mesh_magnetic_field_lovects, - get_fabs=_libwarpx.get_mesh_magnetic_field, level=level) + # --- Add extra dimensions so that the input has the same number of + # --- dimensions as array. + if isinstance(value, np.ndarray): + value3d = np.array(value, copy=False) + sss = list(value3d.shape) + if not isinstance(ix, slice): sss[0:0] = [1] + if not isinstance(iz, slice): sss[1:1] = [1] + value3d.shape = sss -def BzWrapper(level=0): - return MultiFABWrapper(direction=2, overlaps=[0,0,1], - get_lovects=_libwarpx.get_mesh_magnetic_field_lovects, - get_fabs=_libwarpx.get_mesh_magnetic_field, level=level) + if isinstance(ix, slice): + ixstart = max(ix.start or -self.nghosts, -self.nghosts) + ixstop = min(ix.stop or nx + 1 + self.nghosts, nx + self.overlaps[0] + self.nghosts) + else: + ixstart = ix + ixstop = ix + 1 + if isinstance(iz, slice): + izstart = max(iz.start or -self.nghosts, -self.nghosts) + izstop = min(iz.stop or nz + 1 + self.nghosts, nz + self.overlaps[2] + self.nghosts) + else: + izstart = iz + izstop = iz + 1 + + for i in range(len(fields)): + + # --- The ix1, 2 etc are relative to global indexing + ix1 = max(ixstart, lovects[0,i]) + ix2 = min(ixstop, lovects[0,i] + fields[i].shape[0]) + iz1 = max(izstart, lovects[2,i]) + iz2 = min(izstop, lovects[2,i] + fields[i].shape[2]) + + if ix1 < ix2 and iz1 < iz2: -def JxWrapper(level=0): - return MultiFABWrapper(direction=0, overlaps=[0,1,1], - get_lovects=_libwarpx.get_mesh_current_density_lovects, - get_fabs=_libwarpx.get_mesh_current_density, level=level) + sss = (slice(ix1 - lovects[0,i], ix2 - lovects[0,i]), + slice(iz1 - lovects[2,i], iz2 - lovects[2,i])) + if ic is not None: + sss = tuple(list(sss) + [ic]) -def JyWrapper(level=0): - return MultiFABWrapper(direction=1, overlaps=[1,0,1], - get_lovects=_libwarpx.get_mesh_current_density_lovects, - get_fabs=_libwarpx.get_mesh_current_density, level=level) + if isinstance(value, np.ndarray): + vslice = (slice(ix1 - ixstart, ix2 - ixstart), + slice(iz1 - izstart, iz2 - izstart)) + fields[i][sss] = value3d[vslice] + else: + fields[i][sss] = value -def JzWrapper(level=0): - return MultiFABWrapper(direction=2, overlaps=[1,1,0], - get_lovects=_libwarpx.get_mesh_current_density_lovects, - get_fabs=_libwarpx.get_mesh_current_density, level=level) +def ExWrapper(level=0, include_ghosts=False): + return _MultiFABWrapper(direction=0, overlaps=[0,1,1], + get_lovects=_libwarpx.get_mesh_electric_field_lovects, + get_fabs=_libwarpx.get_mesh_electric_field, + level=level, include_ghosts=include_ghosts) + +def EyWrapper(level=0, include_ghosts=False): + return _MultiFABWrapper(direction=1, overlaps=[1,0,1], + get_lovects=_libwarpx.get_mesh_electric_field_lovects, + get_fabs=_libwarpx.get_mesh_electric_field, + level=level, include_ghosts=include_ghosts) + +def EzWrapper(level=0, include_ghosts=False): + return _MultiFABWrapper(direction=2, overlaps=[1,1,0], + get_lovects=_libwarpx.get_mesh_electric_field_lovects, + get_fabs=_libwarpx.get_mesh_electric_field, + level=level, include_ghosts=include_ghosts) + +def BxWrapper(level=0, include_ghosts=False): + return _MultiFABWrapper(direction=0, overlaps=[1,0,0], + get_lovects=_libwarpx.get_mesh_magnetic_field_lovects, + get_fabs=_libwarpx.get_mesh_magnetic_field, + level=level, include_ghosts=include_ghosts) + +def ByWrapper(level=0, include_ghosts=False): + return _MultiFABWrapper(direction=1, overlaps=[0,1,0], + get_lovects=_libwarpx.get_mesh_magnetic_field_lovects, + get_fabs=_libwarpx.get_mesh_magnetic_field, + level=level, include_ghosts=include_ghosts) + +def BzWrapper(level=0, include_ghosts=False): + return _MultiFABWrapper(direction=2, overlaps=[0,0,1], + get_lovects=_libwarpx.get_mesh_magnetic_field_lovects, + get_fabs=_libwarpx.get_mesh_magnetic_field, + level=level, include_ghosts=include_ghosts) + +def JxWrapper(level=0, include_ghosts=False): + return _MultiFABWrapper(direction=0, overlaps=[0,1,1], + get_lovects=_libwarpx.get_mesh_current_density_lovects, + get_fabs=_libwarpx.get_mesh_current_density, + level=level, include_ghosts=include_ghosts) + +def JyWrapper(level=0, include_ghosts=False): + return _MultiFABWrapper(direction=1, overlaps=[1,0,1], + get_lovects=_libwarpx.get_mesh_current_density_lovects, + get_fabs=_libwarpx.get_mesh_current_density, + level=level, include_ghosts=include_ghosts) + +def JzWrapper(level=0, include_ghosts=False): + return _MultiFABWrapper(direction=2, overlaps=[1,1,0], + get_lovects=_libwarpx.get_mesh_current_density_lovects, + get_fabs=_libwarpx.get_mesh_current_density, + level=level, include_ghosts=include_ghosts) diff --git a/Python/pywarpx/picmi.py b/Python/pywarpx/picmi.py index acf730e5a..ac272089c 100644 --- a/Python/pywarpx/picmi.py +++ b/Python/pywarpx/picmi.py @@ -259,13 +259,13 @@ class CylindricalGrid(picmistandard.PICMI_CylindricalGrid): assert self.lower_bound[0] >= 0., Exception('Lower radial boundary must be >= 0.') assert self.bc_rmin != 'periodic' and self.bc_rmax != 'periodic', Exception('Radial boundaries can not be periodic') - assert self.n_azimuthal_modes is None or self.n_azimuthal_modes == 1, Exception('Only one azimuthal mode supported') # Geometry pywarpx.geometry.coord_sys = 1 # RZ pywarpx.geometry.is_periodic = '0 %d'%(self.bc_zmin=='periodic') # Is periodic? pywarpx.geometry.prob_lo = self.lower_bound # physical domain pywarpx.geometry.prob_hi = self.upper_bound + pywarpx.warpx.nmodes = self.n_azimuthal_modes if self.moving_window_velocity is not None and np.any(np.not_equal(self.moving_window_velocity, 0.)): pywarpx.warpx.do_moving_window = 1 |