diff options
Diffstat (limited to 'Python/pywarpx/_libwarpx.py')
| -rwxr-xr-x | Python/pywarpx/_libwarpx.py | 280 |
1 files changed, 80 insertions, 200 deletions
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): |