diff options
| -rwxr-xr-x | Python/pywarpx/_libwarpx.py | 35 | ||||
| -rw-r--r-- | Python/pywarpx/fields.py | 97 | ||||
| -rw-r--r-- | Source/Python/WarpXWrappers.cpp | 39 |
3 files changed, 94 insertions, 77 deletions
diff --git a/Python/pywarpx/_libwarpx.py b/Python/pywarpx/_libwarpx.py index 953de0698..2560da235 100755 --- a/Python/pywarpx/_libwarpx.py +++ b/Python/pywarpx/_libwarpx.py @@ -533,7 +533,10 @@ def get_particle_z(species_number, level=0): ''' structs = get_particle_structs(species_number, level) - return [struct['z'] for struct in structs] + if geometry_dim == '3d': + return [struct['z'] for struct in structs] + elif geometry_dim == 'rz' or geometry_dim == '2d': + return [struct['y'] for struct in structs] def get_particle_id(species_number, level=0): @@ -625,16 +628,16 @@ def _get_mesh_field_list(warpx_func, level, direction, include_ghosts): shapes = _LP_c_int() size = ctypes.c_int(0) ncomps = ctypes.c_int(0) - ngrow = ctypes.c_int(0) + ngrowvect = _LP_c_int() if direction is None: data = warpx_func(level, ctypes.byref(size), ctypes.byref(ncomps), - ctypes.byref(ngrow), ctypes.byref(shapes)) + ctypes.byref(ngrowvect), ctypes.byref(shapes)) else: data = warpx_func(level, direction, ctypes.byref(size), ctypes.byref(ncomps), - ctypes.byref(ngrow), ctypes.byref(shapes)) - ng = ngrow.value + ctypes.byref(ngrowvect), ctypes.byref(shapes)) + ngvect = [ngrowvect[i] for i in range(dim)] grid_data = [] shapesize = dim if ncomps.value > 1: @@ -651,7 +654,7 @@ def _get_mesh_field_list(warpx_func, level, direction, include_ghosts): if include_ghosts: grid_data.append(arr) else: - grid_data.append(arr[tuple([slice(ng, -ng) for _ in range(dim)])]) + grid_data.append(arr[tuple([slice(ngvect[d], -ngvect[d]) for d in range(dim)])]) _libc.free(shapes) _libc.free(data) @@ -1132,15 +1135,15 @@ def get_mesh_charge_density_fp(level, include_ghosts=True): return _get_mesh_field_list(libwarpx.warpx_getChargeDensityFP, level, None, include_ghosts) -def _get_mesh_array_lovects(level, direction, include_ghosts=True, getarrayfunc=None): +def _get_mesh_array_lovects(level, direction, include_ghosts=True, getlovectsfunc=None): assert(0 <= level and level <= libwarpx.warpx_finestLevel()) size = ctypes.c_int(0) - ngrow = ctypes.c_int(0) + ngrowvect = _LP_c_int() if direction is None: - data = getarrayfunc(level, ctypes.byref(size), ctypes.byref(ngrow)) + data = getlovectsfunc(level, ctypes.byref(size), ctypes.byref(ngrowvect)) else: - data = getarrayfunc(level, direction, ctypes.byref(size), ctypes.byref(ngrow)) + data = getlovectsfunc(level, direction, ctypes.byref(size), ctypes.byref(ngrowvect)) lovects_ref = np.ctypeslib.as_array(data, (size.value, dim)) @@ -1148,12 +1151,18 @@ def _get_mesh_array_lovects(level, direction, include_ghosts=True, getarrayfunc= # --- Also, take the transpose to give shape (dims, number of grids) lovects = lovects_ref.copy().T - if not include_ghosts: - lovects += ngrow.value + ng = [] + if include_ghosts: + for d in range(dim): + ng.append(ngrowvect[d]) + else: + for d in range(dim): + ng.append(0) + lovects[d,:] += ngrowvect[d] del lovects_ref _libc.free(data) - return lovects + return lovects, ng def get_mesh_electric_field_lovects(level, direction, include_ghosts=True): diff --git a/Python/pywarpx/fields.py b/Python/pywarpx/fields.py index eeed3283d..e36f4707a 100644 --- a/Python/pywarpx/fields.py +++ b/Python/pywarpx/fields.py @@ -45,26 +45,26 @@ class _MultiFABWrapper(object): self.dim = _libwarpx.dim # overlaps is one along the axes where the grid boundaries overlap the neighboring grid, - # which is the case with node centering + # 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 = self.get_lovects(self.level, self.include_ghosts) + lovects, ngrow = self.get_lovects(self.level, self.include_ghosts) else: - lovects = self.get_lovects(self.level, self.direction, self.include_ghosts) - self.nghosts = -lovects.min() - return lovects + lovects, ngrow = self.get_lovects(self.level, self.direction, self.include_ghosts) + return lovects, ngrow def _gethivects(self): - lovects = self._getlovects() + 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 + return hivects, ngrow def _getfields(self): if self.direction is None: @@ -73,7 +73,8 @@ class _MultiFABWrapper(object): return self.get_fabs(self.level, self.direction, self.include_ghosts) def __len__(self): - return lend(self._getlovects()) + lovects, ngrow = self._getlovects() + return lend(lovects) def mesh(self, direction): """Returns the mesh along the specified direction with the appropriate centering. @@ -96,8 +97,8 @@ class _MultiFABWrapper(object): raise Exception('Inappropriate direction given') # --- Get the total number of cells along the direction - hivects = self._gethivects() - nn = hivects[idir,:].max() - self.nghosts + self.overlaps[idir] + hivects, ngrow = self._gethivects() + nn = hivects[idir,:].max() - ngrow[idir] + self.overlaps[idir] if npes > 1: nn = comm_world.allreduce(nn, op=mpi.MAX) @@ -143,8 +144,8 @@ class _MultiFABWrapper(object): """Returns slices of a 3D decomposed array, """ - lovects = self._getlovects() - hivects = self._gethivects() + lovects, ngrow = self._getlovects() + hivects, ngrow = self._gethivects() fields = self._getfields() ix = index[0] @@ -164,9 +165,9 @@ class _MultiFABWrapper(object): else: ic = None - nx = hivects[0,:].max() - self.nghosts - ny = hivects[1,:].max() - self.nghosts - nz = hivects[2,:].max() - self.nghosts + nx = hivects[0,:].max() - ngrow[0] + ny = hivects[1,:].max() - ngrow[1] + nz = hivects[2,:].max() - ngrow[2] if npes > 1: nx = comm_world.allreduce(nx, op=mpi.MAX) @@ -174,20 +175,20 @@ class _MultiFABWrapper(object): nz = comm_world.allreduce(nz, op=mpi.MAX) 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) + ixstart = max(ix.start or -ngrow[0], -ngrow[0]) + ixstop = min(ix.stop or nx + 1 + ngrow[0], nx + self.overlaps[0] + ngrow[0]) else: ixstart = ix ixstop = ix + 1 if isinstance(iy, slice): - iystart = max(iy.start or -self.nghosts, -self.nghosts) - iystop = min(iy.stop or ny + 1 + self.nghosts, ny + self.overlaps[1] + self.nghosts) + iystart = max(iy.start or -ngrow[1], -ngrow[1]) + iystop = min(iy.stop or ny + 1 + ngrow[1], ny + self.overlaps[1] + ngrow[1]) else: iystart = iy iystop = iy + 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) + izstart = max(iz.start or -ngrow[2], -ngrow[2]) + izstop = min(iz.stop or nz + 1 + ngrow[2], nz + self.overlaps[2] + ngrow[2]) else: izstart = iz izstop = iz + 1 @@ -250,8 +251,8 @@ class _MultiFABWrapper(object): """Returns slices of a 2D decomposed array, """ - lovects = self._getlovects() - hivects = self._gethivects() + lovects, ngrow = self._getlovects() + hivects, ngrow = self._gethivects() fields = self._getfields() ix = index[0] @@ -270,22 +271,22 @@ class _MultiFABWrapper(object): else: ic = None - nx = hivects[0,:].max() - self.nghosts - nz = hivects[1,:].max() - self.nghosts + nx = hivects[0,:].max() - ngrow[0] + nz = hivects[1,:].max() - ngrow[1] if npes > 1: nx = comm_world.allreduce(nx, op=mpi.MAX) nz = comm_world.allreduce(nz, op=mpi.MAX) 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) + ixstart = max(ix.start or -ngrow[0], -ngrow[0]) + ixstop = min(ix.stop or nx + 1 + ngrow[0], nx + self.overlaps[0] + ngrow[0]) 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) + izstart = max(iz.start or -ngrow[1], -ngrow[1]) + izstop = min(iz.stop or nz + 1 + ngrow[1], nz + self.overlaps[1] + ngrow[1]) else: izstart = iz izstop = iz + 1 @@ -365,8 +366,8 @@ class _MultiFABWrapper(object): iy = index[1] iz = index[2] - lovects = self._getlovects() - hivects = self._gethivects() + lovects, ngrow = self._getlovects() + hivects, ngrow = self._gethivects() fields = self._getfields() if len(index) > self.dim: @@ -377,9 +378,9 @@ class _MultiFABWrapper(object): else: ic = None - nx = hivects[0,:].max() - self.nghosts - ny = hivects[1,:].max() - self.nghosts - nz = hivects[2,:].max() - self.nghosts + nx = hivects[0,:].max() - ngrow[0] + ny = hivects[1,:].max() - ngrow[1] + nz = hivects[2,:].max() - ngrow[2] # --- Add extra dimensions so that the input has the same number of # --- dimensions as array. @@ -392,20 +393,20 @@ class _MultiFABWrapper(object): value3d.shape = sss 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) + ixstart = max(ix.start or -ngrow[0], -ngrow[0]) + ixstop = min(ix.stop or nx + 1 + ngrow[0], nx + self.overlaps[0] + ngrow[0]) else: ixstart = ix ixstop = ix + 1 if isinstance(iy, slice): - iystart = max(iy.start or -self.nghosts, -self.nghosts) - iystop = min(iy.stop or ny + 1 + self.nghosts, ny + self.overlaps[1] + self.nghosts) + iystart = max(iy.start or -ngrow[1], -ngrow[1]) + iystop = min(iy.stop or ny + 1 + ngrow[1], ny + self.overlaps[1] + ngrow[1]) else: iystart = iy iystop = iy + 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) + izstart = max(iz.start or -ngrow[2], -ngrow[2]) + izstop = min(iz.stop or nz + 1 + ngrow[2], nz + self.overlaps[2] + ngrow[2]) else: izstart = iz izstop = iz + 1 @@ -442,8 +443,8 @@ class _MultiFABWrapper(object): ix = index[0] iz = index[2] - lovects = self._getlovects() - hivects = self._gethivects() + lovects, ngrow = self._getlovects() + hivects, ngrow = self._gethivects() fields = self._getfields() if len(fields[0].shape) > self.dim: @@ -459,8 +460,8 @@ class _MultiFABWrapper(object): else: ic = None - nx = hivects[0,:].max() - self.nghosts - nz = hivects[2,:].max() - self.nghosts + nx = hivects[0,:].max() - ngrow[0] + nz = hivects[2,:].max() - ngrow[1] # --- Add extra dimensions so that the input has the same number of # --- dimensions as array. @@ -472,14 +473,14 @@ class _MultiFABWrapper(object): value3d.shape = sss 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) + ixstart = max(ix.start or -ngrow[0], -ngrow[0]) + ixstop = min(ix.stop or nx + 1 + ngrow[0], nx + self.overlaps[0] + ngrow[0]) 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) + izstart = max(iz.start or -ngrow[1], -ngrow[1]) + izstop = min(iz.stop or nz + 1 + ngrow[1], nz + self.overlaps[2] + ngrow[1]) else: izstart = iz izstop = iz + 1 diff --git a/Source/Python/WarpXWrappers.cpp b/Source/Python/WarpXWrappers.cpp index ac0f60449..cc42f8cd7 100644 --- a/Source/Python/WarpXWrappers.cpp +++ b/Source/Python/WarpXWrappers.cpp @@ -19,12 +19,15 @@ namespace { - amrex::Real** getMultiFabPointers(const amrex::MultiFab& mf, int *num_boxes, int *ncomps, int *ngrow, int **shapes) + amrex::Real** getMultiFabPointers(const amrex::MultiFab& mf, int *num_boxes, int *ncomps, int **ngrowvect, int **shapes) { *ncomps = mf.nComp(); - *ngrow = mf.nGrow(); *num_boxes = mf.local_size(); int shapesize = AMREX_SPACEDIM; + *ngrowvect = static_cast<int*>(malloc(sizeof(int)*shapesize)); + for (int j = 0; j < AMREX_SPACEDIM; ++j) { + (*ngrowvect)[j] = mf.nGrow(j); + } if (mf.nComp() > 1) shapesize += 1; *shapes = static_cast<int*>(malloc(sizeof(int)*shapesize * (*num_boxes))); auto data = @@ -43,9 +46,13 @@ namespace } return data; } - int* getMultiFabLoVects(const amrex::MultiFab& mf, int *num_boxes, int *ngrow) + int* getMultiFabLoVects(const amrex::MultiFab& mf, int *num_boxes, int **ngrowvect) { - *ngrow = mf.nGrow(); + int shapesize = AMREX_SPACEDIM; + *ngrowvect = static_cast<int*>(malloc(sizeof(int)*shapesize)); + for (int j = 0; j < AMREX_SPACEDIM; ++j) { + (*ngrowvect)[j] = mf.nGrow(j); + } *num_boxes = mf.local_size(); int *loVects = (int*) malloc((*num_boxes)*AMREX_SPACEDIM * sizeof(int)); @@ -246,16 +253,16 @@ extern "C" #define WARPX_GET_FIELD(FIELD, GETTER) \ amrex::Real** FIELD(int lev, int direction, \ - int *return_size, int *ncomps, int *ngrow, int **shapes) { \ + int *return_size, int *ncomps, int **ngrowvect, int **shapes) { \ auto & mf = GETTER(lev, direction); \ - return getMultiFabPointers(mf, return_size, ncomps, ngrow, shapes); \ + return getMultiFabPointers(mf, return_size, ncomps, ngrowvect, shapes); \ } #define WARPX_GET_LOVECTS(FIELD, GETTER) \ int* FIELD(int lev, int direction, \ - int *return_size, int *ngrow) { \ + int *return_size, int **ngrowvect) { \ auto & mf = GETTER(lev, direction); \ - return getMultiFabLoVects(mf, return_size, ngrow); \ + return getMultiFabLoVects(mf, return_size, ngrowvect); \ } WARPX_GET_FIELD(warpx_getEfield, WarpX::GetInstance().getEfield) @@ -295,16 +302,16 @@ extern "C" #define WARPX_GET_SCALAR(SCALAR, GETTER) \ amrex::Real** SCALAR(int lev, \ - int *return_size, int *ncomps, int *ngrow, int **shapes) { \ + int *return_size, int *ncomps, int **ngrowvect, int **shapes) { \ auto & mf = GETTER(lev); \ - return getMultiFabPointers(mf, return_size, ncomps, ngrow, shapes); \ + return getMultiFabPointers(mf, return_size, ncomps, ngrowvect, shapes); \ } #define WARPX_GET_LOVECTS_SCALAR(SCALAR, GETTER) \ int* SCALAR(int lev, \ - int *return_size, int *ngrow) { \ + int *return_size, int **ngrowvect) { \ auto & mf = GETTER(lev); \ - return getMultiFabLoVects(mf, return_size, ngrow); \ + return getMultiFabLoVects(mf, return_size, ngrowvect); \ } WARPX_GET_SCALAR(warpx_getChargeDensityCP, WarpX::GetInstance().getrho_cp) @@ -315,11 +322,11 @@ extern "C" #define WARPX_GET_FIELD_PML(FIELD, GETTER) \ amrex::Real** FIELD(int lev, int direction, \ - int *return_size, int *ncomps, int *ngrow, int **shapes) { \ + int *return_size, int *ncomps, int **ngrowvect, int **shapes) { \ auto * pml = WarpX::GetInstance().GetPML(lev); \ if (pml) { \ auto & mf = *(pml->GETTER()[direction]); \ - return getMultiFabPointers(mf, return_size, ncomps, ngrow, shapes); \ + return getMultiFabPointers(mf, return_size, ncomps, ngrowvect, shapes); \ } else { \ return nullptr; \ } \ @@ -327,11 +334,11 @@ extern "C" #define WARPX_GET_LOVECTS_PML(FIELD, GETTER) \ int* FIELD(int lev, int direction, \ - int *return_size, int *ngrow) { \ + int *return_size, int **ngrowvect) { \ auto * pml = WarpX::GetInstance().GetPML(lev); \ if (pml) { \ auto & mf = *(pml->GETTER()[direction]); \ - return getMultiFabLoVects(mf, return_size, ngrow); \ + return getMultiFabLoVects(mf, return_size, ngrowvect); \ } else { \ return nullptr; \ } \ |