diff options
Diffstat (limited to 'Python')
| -rw-r--r-- | Python/pywarpx/PGroup.py | 8 | ||||
| -rwxr-xr-x | Python/pywarpx/_libwarpx.py | 280 | ||||
| -rw-r--r-- | Python/pywarpx/picmi.py | 2 |
3 files changed, 89 insertions, 201 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/_libwarpx.py b/Python/pywarpx/_libwarpx.py index 4c3283b97..5b7e14dbb 100755 --- a/Python/pywarpx/_libwarpx.py +++ b/Python/pywarpx/_libwarpx.py @@ -399,7 +399,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 +413,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 +578,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[[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 +651,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 +679,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 +707,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 +736,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 +764,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 +792,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 +819,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 +847,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 +875,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/picmi.py b/Python/pywarpx/picmi.py index f242f8589..151de388d 100644 --- a/Python/pywarpx/picmi.py +++ b/Python/pywarpx/picmi.py @@ -256,13 +256,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 |