diff options
Diffstat (limited to 'Python/pywarpx/_libwarpx.py')
| -rwxr-xr-x | Python/pywarpx/_libwarpx.py | 263 |
1 files changed, 168 insertions, 95 deletions
diff --git a/Python/pywarpx/_libwarpx.py b/Python/pywarpx/_libwarpx.py index fcf7f45d6..691288d07 100755 --- a/Python/pywarpx/_libwarpx.py +++ b/Python/pywarpx/_libwarpx.py @@ -300,7 +300,7 @@ class LibWarpX(): self.libwarpx_so.warpx_sett_new.argtypes = [ctypes.c_int, c_real] self.libwarpx_so.warpx_getdt.argtypes = [ctypes.c_int] - def get_boundary_number(self, boundary): + def _get_boundary_number(self, boundary): ''' Utility function to find the boundary number given a boundary name. @@ -308,13 +308,14 @@ class LibWarpX(): Parameters ---------- - boundary : the boundary from which to get the scraped particle data. - In the form x/y/z_hi/lo or eb. + boundary : str + The boundary from which to get the scraped particle data. In the + form x/y/z_hi/lo or eb. Returns ------- - - Integer index in the boundary scraper buffer for the given boundary. + int + Integer index in the boundary scraper buffer for the given boundary. ''' if self.geometry_dim == '3d': dimensions = {'x' : 0, 'y' : 1, 'z' : 2} @@ -348,7 +349,9 @@ class LibWarpX(): @staticmethod def _array1d_from_pointer(pointer, dtype, size): ''' + Function for converting a ctypes pointer to a numpy array + ''' if not pointer: raise Exception(f'_array1d_from_pointer: pointer is a nullptr') @@ -385,7 +388,7 @@ class LibWarpX(): def get_nattr(self): ''' - Get the number of extra attributes. + Get the number of extra particle attributes. ''' # --- The -3 is because the comps include the velocites @@ -393,10 +396,15 @@ class LibWarpX(): def get_nattr_species(self, species_name): ''' - Get the number of real attributes for the given species. + Parameters + ---------- + + species_name: str + Name of the species ''' + return self.libwarpx_so.warpx_nCompsSpecies( ctypes.c_char_p(species_name.encode('utf-8'))) @@ -423,8 +431,7 @@ class LibWarpX(): def initialize(self, argv=None, mpi_comm=None): ''' - Initialize WarpX and AMReX. Must be called before - doing anything else. + Initialize WarpX and AMReX. Must be called before doing anything else. ''' if argv is None: @@ -450,27 +457,33 @@ class LibWarpX(): def getistep(self, level=0): ''' - Get the current time step number for the specified level Parameter --------- - level : the refinement level to reference + level : int + The refinement level to reference ''' + return self.libwarpx_so.warpx_getistep(level) def gett_new(self, level=0): ''' - Get the next time for the specified level + Get the next time for the specified level. + + Parameters + ---------- + level : int + The refinement level to reference ''' + return self.libwarpx_so.warpx_gett_new(level) def evolve(self, num_steps=-1): ''' - Evolve the simulation for num_steps steps. If num_steps=-1, the simulation will be run until the end as specified in the inputs file. @@ -478,21 +491,54 @@ class LibWarpX(): Parameters ---------- - num_steps: int, the number of steps to take - + num_steps: int + The number of steps to take ''' self.libwarpx_so.warpx_evolve(num_steps); def getProbLo(self, direction): + ''' + Get the values of the lower domain boundary. + + Parameters + ---------- + + direction : int + Direction of interest + ''' + assert 0 <= direction < self.dim, 'Inappropriate direction specified' return self.libwarpx_so.warpx_getProbLo(direction) def getProbHi(self, direction): + ''' + Get the values of the upper domain boundary. + + Parameters + ---------- + + direction : int + Direction of interest + ''' + assert 0 <= direction < self.dim, 'Inappropriate direction specified' return self.libwarpx_so.warpx_getProbHi(direction) def getCellSize(self, direction, level=0): + ''' + Get the cell size in the given direction and on the given level. + + Parameters + ---------- + + direction : int + Direction of interest + + level : int + The refinement level to reference + ''' + assert 0 <= direction < 3, 'Inappropriate direction specified' assert 0 <= level and level <= self.libwarpx_so.warpx_finestLevel(), 'Inappropriate level specified' return self.libwarpx_so.warpx_getCellSize(direction, level) @@ -538,25 +584,33 @@ class LibWarpX(): # # self.libwarpx_so.warpx_ComputePMLFactors(lev, dt) - def add_particles(self, species_name, x=None, y=None, z=None, ux=None, uy=None, uz=None, w=None, - unique_particles=True, **kwargs): + def add_particles(self, species_name, x=None, y=None, z=None, ux=None, uy=None, + uz=None, w=None, unique_particles=True, **kwargs): ''' - A function for adding particles to the WarpX simulation. Parameters ---------- - species_name : the species to add the particle to - x, y, z : arrays or scalars of the particle positions (default = 0.) - ux, uy, uz : arrays or scalars of the particle momenta (default = 0.) - w : array or scalar of particle weights (default = 0.) - unique_particles : whether the particles are unique or duplicated on - several processes. (default = True) - kwargs : dictionary containing an entry for all the extra particle - attribute arrays. If an attribute is not given it will be - set to 0. + species_name : str + The type of species for which particles will be added + + x, y, z : arrays or scalars + The particle positions (default = 0.) + + ux, uy, uz : arrays or scalars + The particle momenta (default = 0.) + + w : array or scalars + Particle weights (default = 0.) + + unique_particles : bool + Whether the particles are unique or duplicated on several processes + (default = True) + kwargs : dict + Containing an entry for all the extra particle attribute arrays. If + an attribute is not given it will be set to 0. ''' # --- Get length of arrays, set to one for scalars @@ -646,30 +700,31 @@ class LibWarpX(): def get_particle_count(self, species_name, local=False): ''' - - This returns the number of particles of the specified species in the - simulation. + Get the number of particles of the specified species in the simulation. Parameters ---------- - species_name : the species name that the number will be returned for - local : If True the particle count on this processor will - be returned. + species_name : str + The species name that the number will be returned for + + local : bool + If True the particle count on this processor will be returned. + Default False. Returns ------- + int An integer count of the number of particles - ''' + return self.libwarpx_so.warpx_getNumParticles( ctypes.c_char_p(species_name.encode('utf-8')), local ) def get_particle_structs(self, species_name, level): ''' - This returns a list of numpy arrays containing the particle struct data on each tile for this process. The particle data is represented as a structured numpy array and contains the particle 'x', 'y', 'z', 'id', and 'cpu'. @@ -680,13 +735,17 @@ class LibWarpX(): Parameters ---------- - species_name : the species name that the data will be returned for + species_name : str + The species name that the data will be returned for + + level : int + The refinement level to reference Returns ------- - A List of numpy arrays. - + List of numpy arrays + The requested particle struct data ''' particles_per_tile = _LP_c_int() @@ -709,7 +768,6 @@ class LibWarpX(): def get_particle_arrays(self, species_name, comp_name, level): ''' - This returns a list of numpy arrays containing the particle array data on each tile for this process. @@ -719,14 +777,20 @@ class LibWarpX(): Parameters ---------- - species_name : the species name that the data will be returned for - comp_name : the component of the array data that will be returned. + species_name : str + The species name that the data will be returned for + + comp_name : str + The component of the array data that will be returned + + level : int + The refinement level to reference Returns ------- - A List of numpy arrays. - + List of numpy arrays + The requested particle array data ''' particles_per_tile = _LP_c_int() @@ -887,7 +951,6 @@ class LibWarpX(): def get_particle_comp_index(self, species_name, pid_name): ''' - Get the component index for a given particle attribute. This is useful to get the corrent ordering of attributes when adding new particles using `add_particles()`. @@ -895,15 +958,19 @@ class LibWarpX(): Parameters ---------- - species_name : the species name that the data will be returned for - pid_name : string that is used to identify the new component + species_name : str + The name of the species + + pid_name : str + Name of the component for which the index will be returned Returns ------- + int Integer corresponding to the index of the requested attribute - ''' + return self.libwarpx_so.warpx_getParticleCompIndex( ctypes.c_char_p(species_name.encode('utf-8')), ctypes.c_char_p(pid_name.encode('utf-8')) @@ -911,17 +978,21 @@ class LibWarpX(): def add_real_comp(self, species_name, pid_name, comm=True): ''' - Add a real component to the particle data array. Parameters ---------- - species_name : the species name for which the new component will be added - pid_name : string that is used to identify the new component - comm : should the component be communicated + species_name : str + The species name for which the new component will be added + + pid_name : str + Name that can be used to identify the new component + comm : bool + Should the component be communicated ''' + self.libwarpx_so.warpx_addRealComp( ctypes.c_char_p(species_name.encode('utf-8')), ctypes.c_char_p(pid_name.encode('utf-8')), comm @@ -929,47 +1000,45 @@ class LibWarpX(): def get_species_charge_sum(self, species_name, local=False): ''' - Returns the total charge in the simulation due to the given species. Parameters ---------- - species_name : the species name for which charge will be summed - local : If True return total charge per processor + species_name : str + The species name for which charge will be summed + local : bool + If True return total charge per processor ''' + return self.libwarpx_so.warpx_sumParticleCharge( ctypes.c_char_p(species_name.encode('utf-8')), local ) def get_particle_boundary_buffer_size(self, species_name, boundary): ''' - This returns the number of particles that have been scraped so far in the simulation from the specified boundary and of the specified species. Parameters ---------- - species_name : return the number of scraped particles of this species - boundary : the boundary from which to get the scraped particle data. - In the form x/y/z_hi/lo - - Returns - ------- - - The number of particles scraped so far from a boundary and of a species. + species_name : str + Return the number of scraped particles of this species + boundary : str + The boundary from which to get the scraped particle data in the + form x/y/z_hi/lo ''' + return self.libwarpx_so.warpx_getParticleBoundaryBufferSize( ctypes.c_char_p(species_name.encode('utf-8')), - self.get_boundary_number(boundary) + self._get_boundary_number(boundary) ) def get_particle_boundary_buffer_structs(self, species_name, boundary, level): ''' - This returns a list of numpy arrays containing the particle struct data for a species that has been scraped by a specific simulation boundary. The particle data is represented as a structured numpy array and contains the @@ -981,23 +1050,22 @@ class LibWarpX(): Parameters ---------- - species_name : the species name that the data will be returned for - boundary : the boundary from which to get the scraped particle data. - In the form x/y/z_hi/lo or eb. - level : Which AMR level to retrieve scraped particle data from. + species_name : str + The species name that the data will be returned for - Returns - ------- - - A List of numpy arrays. + boundary : str + The boundary from which to get the scraped particle data in the + form x/y/z_hi/lo or eb. + level : int + Which AMR level to retrieve scraped particle data from. ''' particles_per_tile = _LP_c_int() num_tiles = ctypes.c_int(0) data = self.libwarpx_so.warpx_getParticleBoundaryBufferStructs( ctypes.c_char_p(species_name.encode('utf-8')), - self.get_boundary_number(boundary), level, + self._get_boundary_number(boundary), level, ctypes.byref(num_tiles), ctypes.byref(particles_per_tile) ) @@ -1014,7 +1082,6 @@ class LibWarpX(): def get_particle_boundary_buffer(self, species_name, boundary, comp_name, level): ''' - This returns a list of numpy arrays containing the particle array data for a species that has been scraped by a specific simulation boundary. @@ -1024,33 +1091,34 @@ class LibWarpX(): Parameters ---------- - species_name : the species name that the data will be returned for. - boundary : the boundary from which to get the scraped particle data. - In the form x/y/z_hi/lo or eb. - comp_name : the component of the array data that will be returned. - If "step_scraped" the special attribute holding the - timestep at which a particle was scraped will be - returned. - level : Which AMR level to retrieve scraped particle data from. + species_name : str + The species name that the data will be returned for. - Returns - ------- + boundary : str + The boundary from which to get the scraped particle data in the + form x/y/z_hi/lo or eb. - A List of numpy arrays. + comp_name : str + The component of the array data that will be returned. If + "step_scraped" the special attribute holding the timestep at + which a particle was scraped will be returned. + level : int + Which AMR level to retrieve scraped particle data from. ''' + particles_per_tile = _LP_c_int() num_tiles = ctypes.c_int(0) if comp_name == 'step_scraped': data = self.libwarpx_so.warpx_getParticleBoundaryBufferScrapedSteps( ctypes.c_char_p(species_name.encode('utf-8')), - self.get_boundary_number(boundary), level, + self._get_boundary_number(boundary), level, ctypes.byref(num_tiles), ctypes.byref(particles_per_tile) ) else: data = self.libwarpx_so.warpx_getParticleBoundaryBuffer( ctypes.c_char_p(species_name.encode('utf-8')), - self.get_boundary_number(boundary), level, + self._get_boundary_number(boundary), level, ctypes.byref(num_tiles), ctypes.byref(particles_per_tile), ctypes.c_char_p(comp_name.encode('utf-8')) ) @@ -1083,20 +1151,25 @@ class LibWarpX(): def depositChargeDensity(self, species_name, level, clear_rho=True, sync_rho=True): ''' - Deposit the specified species' charge density in rho_fp in order to - access that data via pywarpx.fields.RhoFPWrapper() + access that data via pywarpx.fields.RhoFPWrapper(). Parameters ---------- - species_name : the species name that will be deposited. - level : Which AMR level to retrieve scraped particle data from. - clear_rho : If True, zero out rho_fp before deposition. - sync_rho : If True, perform MPI exchange and properly set boundary - cells for rho_fp. + species_name : str + The species name that will be deposited. + + level : int + Which AMR level to retrieve scraped particle data from. + clear_rho : bool + If True, zero out rho_fp before deposition. + + sync_rho : bool + If True, perform MPI exchange and properly set boundary cells for rho_fp. ''' + if clear_rho: from . import fields fields.RhoFPWrapper(level, True)[...] = 0.0 |