diff options
Diffstat (limited to 'Python')
| -rwxr-xr-x | Python/pywarpx/_libwarpx.py | 12 | ||||
| -rw-r--r-- | Python/pywarpx/fields.py | 40 | ||||
| -rw-r--r-- | Python/pywarpx/picmi.py | 21 |
3 files changed, 66 insertions, 7 deletions
diff --git a/Python/pywarpx/_libwarpx.py b/Python/pywarpx/_libwarpx.py index 2a37dbd5e..5df9aa0f9 100755 --- a/Python/pywarpx/_libwarpx.py +++ b/Python/pywarpx/_libwarpx.py @@ -392,7 +392,11 @@ def get_particle_arrays(species_number, comp): particle_data = [] for i in range(num_tiles.value): arr = np.ctypeslib.as_array(data[i], (particles_per_tile[i],)) - arr.setflags(write=1) + try: + # This fails on some versions of numpy + arr.setflags(write=1) + except ValueError: + pass particle_data.append(arr) _libc.free(particles_per_tile) @@ -623,7 +627,11 @@ def _get_mesh_field_list(warpx_func, level, direction, include_ghosts): 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) + try: + # This fails on some versions of numpy + arr.setflags(write=1) + except ValueError: + pass if include_ghosts: grid_data.append(arr) else: diff --git a/Python/pywarpx/fields.py b/Python/pywarpx/fields.py index 9068a22ea..fcb04f865 100644 --- a/Python/pywarpx/fields.py +++ b/Python/pywarpx/fields.py @@ -8,6 +8,13 @@ JxWrapper, JyWrapper, JzWrapper """ import numpy as np +try: + from mpi4py import MPI as mpi + comm_world = mpi.COMM_WORLD + npes = comm_world.Get_size() +except ImportError: + npes = 1 + from . import _libwarpx @@ -92,6 +99,11 @@ class _MultiFABWrapper(object): ny = hivects[1,:].max() - self.nghosts nz = hivects[2,:].max() - self.nghosts + if npes > 1: + nx = comm_world.allreduce(nx, op=mpi.MAX) + ny = comm_world.allreduce(ny, 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) @@ -117,6 +129,7 @@ class _MultiFABWrapper(object): max(0, izstop - izstart)) resultglobal = np.zeros(sss) + datalist = [] for i in range(len(fields)): # --- The ix1, 2 etc are relative to global indexing @@ -137,7 +150,16 @@ class _MultiFABWrapper(object): slice(iy1 - iystart, iy2 - iystart), slice(iz1 - izstart, iz2 - izstart)) - resultglobal[vslice] = fields[i][sss] + datalist.append((vslice, fields[i][sss])) + + if npes == 1: + all_datalist = [datalist] + else: + all_datalist = comm_world.allgather(datalist) + + for datalist in all_datalist: + for vslice, ff in datalist: + resultglobal[vslice] = ff # --- Now remove any of the reduced dimensions. sss = [slice(None), slice(None), slice(None)] @@ -177,6 +199,10 @@ class _MultiFABWrapper(object): nx = hivects[0,:].max() - self.nghosts nz = hivects[1,:].max() - self.nghosts + 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) @@ -198,6 +224,7 @@ class _MultiFABWrapper(object): sss = tuple(list(sss) + [ncomps]) resultglobal = np.zeros(sss) + datalist = [] for i in range(len(fields)): # --- The ix1, 2 etc are relative to global indexing @@ -216,7 +243,16 @@ class _MultiFABWrapper(object): vslice = (slice(ix1 - ixstart, ix2 - ixstart), slice(iz1 - izstart, iz2 - izstart)) - resultglobal[vslice] = fields[i][sss] + datalist.append((vslice, fields[i][sss])) + + if npes == 1: + all_datalist = [datalist] + else: + all_datalist = comm_world.allgather(datalist) + + for datalist in all_datalist: + for vslice, ff in datalist: + resultglobal[vslice] = ff # --- Now remove any of the reduced dimensions. sss = [slice(None), slice(None)] diff --git a/Python/pywarpx/picmi.py b/Python/pywarpx/picmi.py index 47f8c562f..c33700278 100644 --- a/Python/pywarpx/picmi.py +++ b/Python/pywarpx/picmi.py @@ -207,7 +207,6 @@ class AnalyticDistribution(picmistandard.PICMI_AnalyticDistribution): species.zmin = self.lower_bound[2] species.zmax = self.upper_bound[2] - # --- Only constant density is supported at this time species.profile = "parse_density_function" if density_scale is None: species.__setattr__('density_function(x,y,z)', self.density_expression) @@ -218,7 +217,10 @@ class AnalyticDistribution(picmistandard.PICMI_AnalyticDistribution): setattr(pywarpx.my_constants, k, v) # --- Note that WarpX takes gamma*beta as input - if np.any(np.not_equal(self.rms_velocity, 0.)): + if np.any(np.not_equal(self.momentum_expressions, None)): + species.momentum_distribution_type = 'parse_momentum_function' + self.setup_parse_momentum_functions(species) + elif np.any(np.not_equal(self.rms_velocity, 0.)): species.momentum_distribution_type = "gaussian" species.ux_m = self.directed_velocity[0]/constants.c species.uy_m = self.directed_velocity[1]/constants.c @@ -235,6 +237,19 @@ class AnalyticDistribution(picmistandard.PICMI_AnalyticDistribution): if self.fill_in: species.do_continuous_injection = 1 + def setup_parse_momentum_functions(self, species): + if self.momentum_expressions[0] is not None: + species.__setattr__('momentum_function_ux(x,y,z)', '({0})/{1}'.format(self.momentum_expressions[0], constants.c)) + else: + species.__setattr__('momentum_function_ux(x,y,z)', '({0})/{1}'.format(self.directed_velocity[0], constants.c)) + if self.momentum_expressions[1] is not None: + species.__setattr__('momentum_function_uy(x,y,z)', '({0})/{1}'.format(self.momentum_expressions[1], constants.c)) + else: + species.__setattr__('momentum_function_uy(x,y,z)', '({0})/{1}'.format(self.directed_velocity[1], constants.c)) + if self.momentum_expressions[2] is not None: + species.__setattr__('momentum_function_uz(x,y,z)', '({0})/{1}'.format(self.momentum_expressions[2], constants.c)) + else: + species.__setattr__('momentum_function_uz(x,y,z)', '({0})/{1}'.format(self.directed_velocity[2], constants.c)) class ParticleListDistribution(picmistandard.PICMI_ParticleListDistribution): def init(self, kw): @@ -308,7 +323,7 @@ class CylindricalGrid(picmistandard.PICMI_CylindricalGrid): 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 + pywarpx.warpx.n_rz_azimuthal_modes = 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 |