diff options
25 files changed, 1506 insertions, 1021 deletions
diff --git a/Example/Langmuir/langmuir.py b/Example/Langmuir/langmuir.py new file mode 100644 index 000000000..da9e3d8ad --- /dev/null +++ b/Example/Langmuir/langmuir.py @@ -0,0 +1,78 @@ +import numpy as np +from pywarpx import * + +nx = 64 +ny = 64 +nz = 64 + +xmin = -20.e-6 +ymin = -20.e-6 +zmin = -20.e-6 +xmax = +20.e-6 +ymax = +20.e-6 +zmax = +20.e-6 + +dx = (xmax - xmin)/nx +dy = (ymax - ymin)/ny +dz = (zmax - zmin)/nz + +# Maximum number of time steps +max_step = 40 + +# number of grid points +amr.n_cell = [nx, ny, nz] + +# Maximum allowable size of each subdomain in the problem domain; +# this is used to decompose the domain for parallel calculations. +amr.max_grid_size = 32 + +# Maximum level in hierarchy (for now must be 0, i.e., one level in total) +amr.max_level = 0 + +amr.plot_int = 1 # How often to write plotfiles. "<= 0" means no plotfiles. + +# Geometry +geometry.coord_sys = 0 # 0: Cartesian +geometry.is_periodic = [1, 1, 1] # Is periodic? +geometry.prob_lo = [xmin, ymin, zmin] # physical domain +geometry.prob_hi = [xmax, ymax, zmax] + +# Verbosity +warpx.verbose = 1 +warpx.do_moving_window = 0 +warpx.moving_window_dir = 'x' +warpx.moving_window_v = 0.0 # in units of the speed of light + +# Algorithms +algo.current_deposition = 3 +algo.charge_deposition = 0 +algo.field_gathering = 0 +algo.particle_pusher = 0 + +# CFL +warpx.cfl = 1.0 + +particles.nspecies = 1 +particles.species_names = 'electrons' + +electrons.charge = '-q_e' +electrons.mass = 'm_e' +electrons.injection_style = "NUniformPerCell" +electrons.num_particles_per_cell_each_dim = [2, 2, 2] + +electrons.xmin = xmin +electrons.xmax = 0.e-6 +electrons.ymin = ymin +electrons.ymax = ymax +electrons.zmin = zmin +electrons.zmax = zmax + +electrons.profile = 'constant' +electrons.density = 1.e25 # number of electrons per m^3 + +electrons.momentum_distribution_type = "constant" +electrons.ux = 0.01 # ux = gamma*beta_x + +# --- Initialize the simulation +warpx.write_inputs('inputs_from_python', max_step=max_step) + diff --git a/Example/Langmuir/langmuir_PICMI.py b/Example/Langmuir/langmuir_PICMI.py new file mode 100644 index 000000000..0a1200ffa --- /dev/null +++ b/Example/Langmuir/langmuir_PICMI.py @@ -0,0 +1,46 @@ +import numpy as np +from pywarpx import picmi + +nx = 64 +ny = 64 +nz = 64 + +xmin = -20.e-6 +ymin = -20.e-6 +zmin = -20.e-6 +xmax = +20.e-6 +ymax = +20.e-6 +zmax = +20.e-6 + +uniform_plasma = picmi.UniformDistribution(density=1.e25, upper_bound=[0., None, None], directed_velocity=[0.1, 0., 0.]) + +electrons = picmi.Species(particle_type='electron', name='electrons', initial_distribution=uniform_plasma) + +grid = picmi.Cartesian3DGrid(number_of_cells = [nx, ny, nz], + lower_bound = [xmin, ymin, zmin], + upper_bound = [xmax, ymax, zmax], + lower_boundary_conditions = ['periodic', 'periodic', 'periodic'], + upper_boundary_conditions = ['periodic', 'periodic', 'periodic'], + moving_window_velocity = [0., 0., 0.], + warpx_max_grid_size=32, warpx_max_level=0, warpx_coord_sys=0) + +solver = picmi.ElectromagneticSolver(grid=grid, cfl=1.) + +sim = picmi.Simulation(solver = solver, + verbose = 1, + max_steps = 40, + plot_int = 1, + warpx_current_deposition_algo = 3, + warpx_charge_deposition_algo = 0, + warpx_field_gathering_algo = 0, + warpx_particle_pusher_algo = 0) + +sim.add_species(electrons, layout=picmi.GriddedLayout(n_macroparticle_per_cell=[2,2,2], grid=grid)) + +# write_inputs will create an inputs file that can be used to run +# with the compiled version. +sim.write_input_file(file_name='inputs_from_PICMI') + +# Alternatively, sim.step will run WarpX, controlling it from Python +#sim.step() + diff --git a/Example/gaussian_beam/gaussian_beam_PICMI.py b/Example/gaussian_beam/gaussian_beam_PICMI.py new file mode 100644 index 000000000..ab8d22c4b --- /dev/null +++ b/Example/gaussian_beam/gaussian_beam_PICMI.py @@ -0,0 +1,59 @@ +import numpy as np +from pywarpx import PICMI +#from warp import PICMI + +nx = 32 +ny = 32 +nz = 32 + +xmin = -2. +xmax = +2. +ymin = -2. +ymax = +2. +zmin = -2. +zmax = +2. + +number_sim_particles = 32678 +total_charge = 8.010883097437485e-07 + +beam_rms_size = 0.25 +electron_beam_divergence = -0.04 + +em_order = 3 + +grid = PICMI.Grid(nx=nx, ny=ny, nz=nz, xmin=xmin, xmax=xmax, ymin=ymin, ymax=ymax, zmin=zmin, zmax=zmax, + bcxmin='periodic', bcxmax='periodic', bcymin='periodic', bcymax='periodic', bczmin='open', bczmax='open', + max_grid_size=16, max_level=0, coord_sys=0) + +solver = PICMI.EM_solver(current_deposition_algo = 0, + charge_deposition_algo = 0, + field_gathering_algo = 0, + particle_pusher_algo = 0, + norderx = em_order, nordery = em_order, norderz = em_order) + +electrons = PICMI.Species(type='electron', name='electrons') +protons = PICMI.Species(type='electron', name='protons') + +electron_beam = PICMI.GaussianBeam(electrons, + number_sim_particles = number_sim_particles, + number_real_particles = total_charge/PICMI.q_e, + Xrms = beam_rms_size, Yrms = beam_rms_size, Zrms = beam_rms_size, + UXdiv = electron_beam_divergence, UYdiv = electron_beam_divergence, UZdiv = electron_beam_divergence) + +proton_beam = PICMI.GaussianBeam(protons, + number_sim_particles = number_sim_particles, + number_real_particles = total_charge/PICMI.q_e, + Xrms = beam_rms_size, Yrms = beam_rms_size, Zrms = beam_rms_size) + +sim = PICMI.Simulation(plot_int = 8, + verbose = 1, + cfl = 1.0, + max_step = 1000) + +# write_inputs will create an inputs file that can be used to run +# with the compiled version. +sim.write_inputs(inputs_name = 'inputs_from_PICMI') + +# Alternatively, sim.step will run WarpX, controlling it from Python +#sim.step() + diff --git a/Example/laser_acceleration/laser_acceleration_PICMI.py b/Example/laser_acceleration/laser_acceleration_PICMI.py index 761ff7cfe..4cb906378 100644 --- a/Example/laser_acceleration/laser_acceleration_PICMI.py +++ b/Example/laser_acceleration/laser_acceleration_PICMI.py @@ -1,168 +1,77 @@ import numpy as np -from pywarpx import PICMI -#from warp import PICMI +from pywarpx import picmi +#from warp import picmi -from mpi4py import MPI -comm = MPI.COMM_WORLD +nx = 64 +ny = 64 +nz = 480 -def get_parallel_indices(Np, rank, size): - ''' +xmin = -30.e-6 +xmax = +30.e-6 +ymin = -30.e-6 +ymax = +30.e-6 +zmin = -56.e-6 +zmax = +12.e-6 - This decomposes the arrays of particle attributes into subarrays - for parallel processing. - - ''' - Navg = Np / size - Nleft = Np - Navg * size - if (rank < Nleft): - lo = rank*(Navg + 1) - hi = lo + Navg + 1 - else: - lo = rank * Navg + Nleft - hi = lo + Navg - return lo, hi - - -def load_plasma(ncells, domain_min, domain_max, injected_plasma_density, injected_plasma_ppc, plasma_min, plasma_max): - ''' - - Sets initial conditions for the plasma acceleration setup. - - ''' - - comm.Barrier() - - dx = (domain_max - domain_min) / ncells - - nplasma_cells = ((plasma_max - plasma_min)/dx + 0.5).astype('l') + 1 - iplasma_min = ((plasma_min - domain_min)/dx).astype('l')*dx + domain_min - - # Species 0 - the plasma - plasma_weight = injected_plasma_density * dx[0] * dx[1] * dx[2] / injected_plasma_ppc - - # Fill the entire domain with particles. Only a subset of these points - # will be selected for each species - Z, Y, X, P = np.mgrid[0:nplasma_cells[2], 0:nplasma_cells[1], 0:nplasma_cells[0], 0:injected_plasma_ppc] - Z = Z.flatten() - Y = Y.flatten() - X = X.flatten() - P = P.flatten() - - particle_shift = (0.5 + P) / injected_plasma_ppc - - xp = (X + particle_shift)*dx[0] + iplasma_min[0] - yp = (Y + particle_shift)*dx[1] + iplasma_min[1] - zp = (Z + particle_shift)*dx[2] + iplasma_min[2] - - # now do the plasma species - plasma_locs = np.logical_and(xp >= plasma_min[0], xp < plasma_max[0]) - plasma_locs = np.logical_and(plasma_locs, yp >= plasma_min[1]) - plasma_locs = np.logical_and(plasma_locs, zp >= plasma_min[2]) - plasma_locs = np.logical_and(plasma_locs, yp < plasma_max[1]) - plasma_locs = np.logical_and(plasma_locs, zp < plasma_max[2]) - - plasma_xp = xp[plasma_locs] - plasma_yp = yp[plasma_locs] - plasma_zp = zp[plasma_locs] - - plasma_uxp = np.zeros_like(plasma_xp) - plasma_uyp = np.zeros_like(plasma_xp) - plasma_uzp = np.zeros_like(plasma_xp) - - # --- and weights - Np = plasma_xp.shape[0] - plasma_wp = np.full(Np, plasma_weight) - - lo, hi = get_parallel_indices(Np, comm.rank, comm.size) - - electrons.add_particles(x=plasma_xp[lo:hi], y=plasma_yp[lo:hi], z=plasma_zp[lo:hi], - ux=plasma_uxp[lo:hi], uy=plasma_uyp[lo:hi], uz=plasma_uzp[lo:hi], - w=plasma_wp[lo:hi], unique_particles=True) - - comm.Barrier() - -def inject_new_plasma(ncells, domain_min, domain_max, num_shift, direction, injected_plasma_density, injected_plasma_ppc, plasma_min, plasma_max): - ''' - - This injects fresh plasma into the domain after the moving window has been upated. - - ''' - - comm.Barrier() - - dx = (domain_max - domain_min) / ncells - - plasma_min_inject = plasma_min.copy() - plasma_max_inject = plasma_max.copy() - if (num_shift > 0): - plasma_min_inject[direction] = domain_max[direction] - plasma_max_inject[direction] = domain_max[direction] + num_shift*dx[direction] - else: - plasma_min_inject[direction] = domain_min[direction] - num_shift*dx[direction] - plasma_max_inject[direction] = domain_min[direction] - - load_plasma(ncells, domain_min, domain_max, injected_plasma_density, injected_plasma_ppc, plasma_min_inject, plasma_max_inject) - -ncells = np.array([64, 64, 480])/2 -domain_min = np.array([-30.e-6, -30.e-6, -56.e-6]) -domain_max = np.array([ 30.e-6, 30.e-6, 12.e-6]) -dx = (domain_max - domain_min) / ncells - -moving_window_velocity = np.array([0., 0., PICMI.clight]) +moving_window_velocity = [0., 0., picmi.c] injected_plasma_density = 1.e23 -injected_plasma_ppc = 4 -plasma_min = np.array([-20.e-6, -20.e-6, 0.0e-6]) -plasma_max = np.array([ 20.e-6, 20.e-6, domain_max[2]]) - -grid = PICMI.Grid(nx=ncells[0], ny=ncells[1], nz=ncells[2], - xmin=domain_min[0], xmax=domain_max[0], ymin=domain_min[1], ymax=domain_max[1], zmin=domain_min[2], zmax=domain_max[2], - bcxmin='periodic', bcxmax='periodic', bcymin='periodic', bcymax='periodic', bczmin='open', bczmax='open', - moving_window_velocity = moving_window_velocity, - max_grid_size=32, max_level=0, coord_sys=0) - -solver = PICMI.EM_solver(current_deposition = 3, - charge_deposition = 0, - field_gathering = 0, - particle_pusher = 0) - -laser = PICMI.Gaussian_laser(antenna_z0 = 9.e-6, # This point is on the laser plane - antenna_zvec = 1., # The plane normal direction - pol_angle = PICMI.pi/2., # The main polarization vector - E0 = 16.e12, # Maximum amplitude of the laser field (in V/m) - waist = 5.e-6, # The waist of the laser (in meters) - duration = 15.e-15, # The duration of the laser (in seconds) - t_peak = 30.e-15, # The time at which the laser reaches its peak (in seconds) - focal_position = 109.e-6, # Focal position (m) - wavelength = 0.8e-6, # The wavelength of the laser (in meters) - em_solver = solver) - -electrons = PICMI.Species(type=PICMI.Electron, name='electrons') - -# Maximum number of time steps -max_step = 1000 - -sim = PICMI.Simulation(plot_int = 200, +number_per_cell_each_dim = [2, 2, 1] +plasma_min = [-20.e-6, -20.e-6, 0.0e-6] +plasma_max = [ 20.e-6, 20.e-6, zmax] + +grid = picmi.Cartesian3DGrid(number_of_cells = [nx, ny, nz], + lower_bound = [xmin, ymin, zmin], + upper_bound = [xmax, ymax, zmax], + lower_boundary_conditions = ['periodic', 'periodic', 'open'], + upper_boundary_conditions = ['periodic', 'periodic', 'open'], + moving_window_velocity = moving_window_velocity, + warpx_max_grid_size=32, warpx_max_level=0, warpx_coord_sys=0) + +solver = picmi.ElectromagneticSolver(grid=grid, cfl=1.) + +t_peak = 30.e-15 # The time at which the laser reaches its peak at the antenna (in seconds) +focal_distance = 100.e-6 # Focal distance from the antenna (in meters) +antenna_z0 = 9.e-6 # This point is on the laser plane +laser = picmi.GaussianLaser(wavelength = 0.8e-6, # The wavelength of the laser (in meters) + waist = 5.e-6, # The waist of the laser (in meters) + duration = 15.e-15, # The duration of the laser (in seconds) + polarization_angle = np.pi/2., # The main polarization vector + focal_position = [0., 0., focal_distance + antenna_z0], # Focal position (m) + E0 = 16.e12, # Maximum amplitude of the laser field (in V/m) + centroid_position = [0., 0., antenna_z0 - picmi.c*t_peak], # Position of the laser centroid in Z at time 0 + propagation_direction = [0,0,1]) + +laser_antenna = picmi.LaserAntenna(position = [0., 0., antenna_z0], # This point is on the laser plane + normal_vector = [0., 0., 1.]) # The plane normal direction + +uniform_plasma = picmi.UniformDistribution(density = injected_plasma_density, + lower_bound = plasma_min, + upper_bound = plasma_max, + fill_in = True) + +electrons = picmi.Species(particle_type = 'electron', + name = 'electrons', + initial_distribution = uniform_plasma) + +sim = picmi.Simulation(solver = solver, + plot_int = 100, verbose = 1, - cfl = 1.0) - -load_plasma(ncells, domain_min, domain_max, injected_plasma_density, injected_plasma_ppc, plasma_min, plasma_max) + cfl = 1.0, + max_steps = 1000, + warpx_current_deposition_algo = 3, + warpx_charge_deposition_algo = 0, + warpx_field_gathering_algo = 0, + warpx_particle_pusher_algo = 0) -direction = np.argmax(abs(moving_window_velocity)) -old_x = grid.getmins()[direction] -new_x = old_x -for i in range(1, max_step + 1): +sim.add_species(electrons, layout=picmi.GriddedLayout(grid=grid, n_macroparticle_per_cell=number_per_cell_each_dim)) - # check whether the moving window has updated - num_shift = int( 0.5 + (new_x - old_x) / dx[direction]) - if (num_shift != 0): - inject_new_plasma(ncells, domain_min, domain_max, num_shift, direction, injected_plasma_density, injected_plasma_ppc, plasma_min, plasma_max) - domain_min[direction] += num_shift*dx[direction] - domain_max[direction] += num_shift*dx[direction] +sim.add_laser(laser, injection_method=laser_antenna) - sim.step(1) +# write_inputs will create an inputs file that can be used to run +# with the compiled version. +sim.write_input_file(file_name = 'inputs_from_PICMI') - old_x = new_x - new_x = grid.getmins()[direction] +# Alternatively, sim.step will run WarpX, controlling it from Python +#sim.step() -sim.finalize() diff --git a/Example/plasma_acceleration/plasma_acceleration_PICMI.py b/Example/plasma_acceleration/plasma_acceleration_PICMI.py new file mode 100644 index 000000000..0d831f3e1 --- /dev/null +++ b/Example/plasma_acceleration/plasma_acceleration_PICMI.py @@ -0,0 +1,62 @@ +import numpy as np +from pywarpx import picmi +#from warp import picmi + +nx = 64 +ny = 64 +nz = 64 + +xmin = -200.e-6 +xmax = +200.e-6 +ymin = -200.e-6 +ymax = +200.e-6 +zmin = -200.e-6 +zmax = +200.e-6 + +moving_window_velocity = [0., 0., picmi.c] + +number_per_cell_each_dim = [2, 2, 1] + +grid = picmi.Cartesian3DGrid(number_of_cells = [nx, ny, nz], + lower_bound = [xmin, ymin, zmin], + upper_bound = [xmax, ymax, zmax], + lower_boundary_conditions = ['periodic', 'periodic', 'open'], + upper_boundary_conditions = ['periodic', 'periodic', 'open'], + moving_window_velocity = moving_window_velocity, + warpx_max_grid_size=32, warpx_max_level=0, warpx_coord_sys=0) + +solver = picmi.ElectromagneticSolver(grid=grid, cfl=1) + +beam_distribution = picmi.UniformDistribution(density = 1.e23, + lower_bound = [-20.e-6, -20.e-6, -150.e-6], + upper_bound = [+20.e-6, +20.e-6, -100.e-6], + directed_velocity = [0., 0., 1.e9]) + +plasma_distribution = picmi.UniformDistribution(density = 1.e22, + lower_bound = [-200.e-6, -200.e-6, 0.], + upper_bound = [+200.e-6, +200.e-6, None], + fill_in = True) + +beam = picmi.Species(particle_type='electron', name='beam', initial_distribution=beam_distribution) +plasma = picmi.Species(particle_type='electron', name='plasma', initial_distribution=plasma_distribution) + +sim = picmi.Simulation(solver = solver, + plot_int = 2, + verbose = 1, + cfl = 1.0, + max_steps = 1000, + warpx_current_deposition_algo = 3, + warpx_charge_deposition_algo = 0, + warpx_field_gathering_algo = 0, + warpx_particle_pusher_algo = 0) + +sim.add_species(beam, layout=picmi.GriddedLayout(grid=grid, n_macroparticle_per_cell=number_per_cell_each_dim)) +sim.add_species(plasma, layout=picmi.GriddedLayout(grid=grid, n_macroparticle_per_cell=number_per_cell_each_dim)) + +# write_inputs will create an inputs file that can be used to run +# with the compiled version. +sim.write_input_file(file_name = 'inputs_from_PICMI') + +# Alternatively, sim.step will run WarpX, controlling it from Python +#sim.step() + diff --git a/Python/pywarpx/AMReX.py b/Python/pywarpx/AMReX.py deleted file mode 100644 index 4ce7ac61c..000000000 --- a/Python/pywarpx/AMReX.py +++ /dev/null @@ -1,42 +0,0 @@ -from .Bucket import Bucket - -from .WarpX import warpx -from .Amr import amr -from .Geometry import geometry -from .Algo import algo -from .Langmuirwave import langmuirwave -from .Interpolation import interpolation -from .Laser import laser -from . import Particles -from .Particles import particles, particles_list - -import ctypes -from ._libwarpx import libwarpx -from ._libwarpx import amrex_init - -class AMReX(object): - - def init(self): - argv = ['warpx'] - argv += warpx.attrlist() - argv += amr.attrlist() - argv += geometry.attrlist() - argv += algo.attrlist() - argv += langmuirwave.attrlist() - argv += interpolation.attrlist() - argv += particles.attrlist() - argv += laser.attrlist() - - if not particles_list: - # --- This is needed in case only species_names has been set, - # --- assuming that only the built in particle types are being used. - for pstring in particles.species_names.split(' '): - particles_list.append(getattr(Particles, pstring)) - - for particle in particles_list: - argv += particle.attrlist() - - amrex_init(argv) - - def finalize(self, finalize_mpi=1): - libwarpx.amrex_finalize(finalize_mpi) diff --git a/Python/pywarpx/Bucket.py b/Python/pywarpx/Bucket.py index 5c429f530..c73b3dac9 100644 --- a/Python/pywarpx/Bucket.py +++ b/Python/pywarpx/Bucket.py @@ -1,5 +1,3 @@ -from six import iteritems - class Bucket(object): """ The purpose of this class is to be a named bucket for holding attributes. @@ -25,10 +23,19 @@ class Bucket(object): def attrlist(self): "Concatenate the attributes into a string" result = [] - for attr, value in iteritems(self.argvattrs): + for attr, value in self.argvattrs.items(): + if value is None: + continue # --- repr is applied to value so that for floats, all of the digits are included. # --- The strip is then needed when value is a string. - attrstring = '{0}.{1}={2} '.format(self.instancename, attr, repr(value).strip("'\"")) + if isinstance(value, str): + rhs = value + elif hasattr(value, '__iter__'): + # --- For lists, tuples, and arrays make a space delimited string of the values + rhs = ' '.join(map(repr, value)) + else: + rhs = value + attrstring = '{0}.{1} = {2}'.format(self.instancename, attr, repr(rhs).strip("'\"")) result += [attrstring] return result diff --git a/Python/pywarpx/PGroup.py b/Python/pywarpx/PGroup.py index fa8572c0a..64e396d38 100644 --- a/Python/pywarpx/PGroup.py +++ b/Python/pywarpx/PGroup.py @@ -148,7 +148,7 @@ class PGroups(object): self._pgroups = [] for igroup in range(self.ngroups): self._pgroups.append(PGroup(igroup, self.ispecie)) - + def __iter__(self): self.setuppgroups() for igroup in range(self.ngroups): @@ -158,3 +158,6 @@ class PGroups(object): self.setuppgroups() return self._pgroups[key] + def __len__(self): + self.setuppgroups() + return len(self._pgroups) diff --git a/Python/pywarpx/PICMI.py b/Python/pywarpx/PICMI.py deleted file mode 100644 index 1bc36e5ca..000000000 --- a/Python/pywarpx/PICMI.py +++ /dev/null @@ -1,133 +0,0 @@ -"""Classes following the PICMI standard -""" -from PICMI_Base import * -import numpy as np -from pywarpx import * - -codename = 'WarpX' - -def _args_to_string(*args): - # --- Converts of sequence of number to a string that is appropriate for input. - return ' '.join(map(repr, args)) - -class Grid(PICMI_Grid): - def init(self, **kw): - - amr.n_cell = _args_to_string(self.nx, self.ny, self.nz) - - # Maximum allowable size of each subdomain in the problem domain; - # this is used to decompose the domain for parallel calculations. - amr.max_grid_size = kw.get('max_grid_size', 32) - - # Maximum level in hierarchy (for now must be 0, i.e., one level in total) - amr.max_level = kw.get('max_level', 0) - - # Geometry - geometry.coord_sys = kw.get('coord_sys', 0) # 0: Cartesian - geometry.is_periodic = '%d %d %d'%(self.bcxmin=='periodic', self.bcymin=='periodic', self.bczmin=='periodic') # Is periodic? - geometry.prob_lo = _args_to_string(self.xmin, self.ymin, self.zmin) # physical domain - geometry.prob_hi = _args_to_string(self.xmax, self.ymax, self.zmax) - - if self.moving_window_velocity is not None and np.any(self.moving_window_velocity != 0): - warpx.do_moving_window = 1 - if self.moving_window_velocity[0] != 0.: - warpx.moving_window_dir = 'x' - warpx.moving_window_v = self.moving_window_velocity[0]/clight # in units of the speed of light - if self.moving_window_velocity[1] != 0.: - warpx.moving_window_dir = 'y' - warpx.moving_window_v = self.moving_window_velocity[1]/clight # in units of the speed of light - if self.moving_window_velocity[2] != 0.: - warpx.moving_window_dir = 'z' - warpx.moving_window_v = self.moving_window_velocity[2]/clight # in units of the speed of light - - def getmins(self, **kw): - return np.array([warpx.getProbLo(0), warpx.getProbLo(1), warpx.getProbLo(2)]) - - def getmaxs(self, **kw): - return np.array([warpx.getProbHi(0), warpx.getProbHi(1), warpx.getProbHi(2)]) - - def getxmin(self): - return warpx.getProbLo(0) - def getxmax(self): - return warpx.getProbHi(0) - def getymin(self): - return warpx.getProbLo(1) - def getymax(self): - return warpx.getProbHi(1) - def getzmin(self): - return warpx.getProbLo(2) - def getzmax(self): - return warpx.getProbHi(2) - - -class EM_solver(PICMI_EM_solver): - def init(self, **kw): - - if self.current_deposition_algo is not None: - algo.current_deposition = self.current_deposition_algo - if self.charge_deposition_algo is not None: - algo.charge_deposition = self.charge_deposition_algo - if self.field_gathering_algo is not None: - algo.field_gathering = self.field_gathering_algo - if self.particle_pusher_algo is not None: - algo.particle_pusher = self.particle_pusher_algo - - -class Gaussian_laser(PICMI_Gaussian_laser): - def init(self, **kw): - - warpx.use_laser = 1 - laser.profile = "Gaussian" - laser.position = _args_to_string(self.antenna_x0, self.antenna_y0, self.antenna_z0) # This point is on the laser plane - laser.direction = _args_to_string(self.antenna_xvec, self.antenna_yvec, self.antenna_zvec) # The plane normal direction - laser.polarization = _args_to_string(np.cos(self.pol_angle), np.sin(self.pol_angle), 0.) # The main polarization vector - laser.e_max = self.E0 # Maximum amplitude of the laser field (in V/m) - laser.profile_waist = self.waist # The waist of the laser (in meters) - laser.profile_duration = self.duration # The duration of the laser (in seconds) - laser.profile_t_peak = self.t_peak # The time at which the laser reaches its peak (in seconds) - laser.profile_focal_distance = self.focal_position - self.antenna_z0 # Focal distance from the antenna (in meters) - laser.wavelength = self.wavelength # The wavelength of the laser (in meters) - - -class Species(PICMI_Species): - def init(self, **kw): - - self.species_number = particles.nspecies - particles.nspecies = particles.nspecies + 1 - particles.species_names = particles.species_names + ' ' + self.name - - self.bucket = Bucket.Bucket(self.name, mass=self.mass, charge=self.charge, injection_style = 'python') - Particles.particles_list.append(self.bucket) - - def add_particles(self, n=None, - x=None, y=None, z=None, - ux=None, uy=None, uz=None, w=None, - unique_particles=None, **kw): - pid = np.array([w]).T - add_particles(self.species_number, x, y, z, ux, uy, uz, pid, unique_particles) - - -class Simulation(PICMI_Simulation): - def set_warpx_attr(self, warpx_obj, attr, kw): - value = kw.get(attr, None) - if value is not None: - setattr(warpx_obj, attr, value) - setattr(self, attr, value) - - def init(self, **kw): - - warpx.verbose = self.verbose - warpx.cfl = self.cfl - amr.plot_int = self.plot_int - - self.amrex = AMReX() - self.amrex.init() - warpx.init() - - def step(self, nsteps=-1): - warpx.evolve(nsteps) - - def finalize(self): - warpx.finalize() - self.amrex.finalize() - diff --git a/Python/pywarpx/Particles.py b/Python/pywarpx/Particles.py index e7bac7e10..757a49d00 100644 --- a/Python/pywarpx/Particles.py +++ b/Python/pywarpx/Particles.py @@ -1,6 +1,6 @@ from .Bucket import Bucket -particles = Bucket('particles', nspecies=0, species_names='') +particles = Bucket('particles', nspecies=0, species_names=None) particles_list = [] electrons = Bucket('electrons') @@ -22,3 +22,8 @@ particle_dict = {'electrons':electrons, 'positrons':positrons, 'protons':protons } + +def newspecies(name): + result = Bucket(name) + particles_list.append(result) + return result diff --git a/Python/pywarpx/WarpX.py b/Python/pywarpx/WarpX.py index f0a31f3f9..f1fb9918a 100644 --- a/Python/pywarpx/WarpX.py +++ b/Python/pywarpx/WarpX.py @@ -1,19 +1,63 @@ from .Bucket import Bucket + +from .Amr import amr +from .Geometry import geometry +from .Algo import algo +from .Langmuirwave import langmuirwave +from .Interpolation import interpolation +from .Laser import laser +from . import Particles +from .Particles import particles, particles_list + +import ctypes from ._libwarpx import libwarpx +from ._libwarpx import amrex_init class WarpX(Bucket): """ A Python wrapper for the WarpX C++ class """ + def create_argv_list(self): + argv = [] + argv += warpx.attrlist() + argv += amr.attrlist() + argv += geometry.attrlist() + argv += algo.attrlist() + argv += langmuirwave.attrlist() + argv += interpolation.attrlist() + argv += particles.attrlist() + argv += laser.attrlist() + + # --- Search through species_names and add any predefined particle objects in the list. + particles_list_names = [p.instancename for p in particles_list] + for pstring in particles.species_names.split(' '): + if pstring in particles_list_names: + # --- The species is already included in particles_list + continue + elif hasattr(Particles, pstring): + # --- Add the predefined species to particles_list + particles_list.append(getattr(Particles, pstring)) + particles_list_names.append(pstring) + else: + raise Exception('Species %s listed in species_names not defined'%pstring) + + for particle in particles_list: + argv += particle.attrlist() + + return argv + def init(self): + argv = ['warpx'] + self.create_argv_list() + amrex_init(argv) libwarpx.warpx_init() def evolve(self, nsteps=-1): libwarpx.warpx_evolve(nsteps) - def finalize(self): + def finalize(self, finalize_mpi=1): libwarpx.warpx_finalize() + libwarpx.amrex_finalize(finalize_mpi) def getProbLo(self, direction): return libwarpx.warpx_getProbLo(direction) @@ -21,4 +65,15 @@ class WarpX(Bucket): def getProbHi(self, direction): return libwarpx.warpx_getProbHi(direction) + def write_inputs(self, filename='inputs', **kw): + argv = self.create_argv_list() + with open(filename, 'w') as ff: + + for k, v in kw.iteritems(): + ff.write('{0} = {1}\n'.format(k, v)) + + for arg in argv: + ff.write('{0}\n'.format(arg)) + warpx = WarpX('warpx') + diff --git a/Python/pywarpx/WarpXPIC.py b/Python/pywarpx/WarpXPIC.py new file mode 100644 index 000000000..77ab8464f --- /dev/null +++ b/Python/pywarpx/WarpXPIC.py @@ -0,0 +1,50 @@ +from warp.run_modes.timestepper import PICAPI +from ._libwarpx import libwarpx + +class WarpXPIC(PICAPI): + + def get_time(self): + return libwarpx.warpx_gett_new(0) + + def set_time(self, time): + for i in range(libwarpx.warpx_finestLevel()+1): + libwarpx.warpx_sett_new(i, time) + + def get_step_size(self): + libwarpx.warpx_ComputeDt() + return libwarpx.warpx_getdt(0) + + def get_step_number(self): + return libwarpx.warpx_getistep(0) + + def set_step_number(self, it): + for i in range(libwarpx.warpx_finestLevel()+1): + libwarpx.warpx_setistep(i, it) + + def push_positions(self, dt): + libwarpx.warpx_PushX(0, dt) + + def push_velocities_withE(self, dt): + libwarpx.warpx_EPushV(0, dt) + + def push_velocities_withB(self, dt): + libwarpx.warpx_BPushV(0, dt) + + def get_self_fields(self): + libwarpx.warpx_FieldGather(0) + + def calculate_source(self): + libwarpx.warpx_CurrentDeposition(0) + + def push_Efields(self, dt): + libwarpx.warpx_EvolveE(0, dt) + libwarpx.warpx_FillBoundaryE(0, True) + + def push_Bfields(self, dt): + libwarpx.warpx_EvolveB(0, dt) + libwarpx.warpx_FillBoundaryB(0, True) + + def apply_particle_boundary_conditions(self): + libwarpx.mypc_Redistribute() # Redistribute particles + libwarpx.warpx_MoveWindow() # !!! not the correct place yet + diff --git a/Python/pywarpx/__init__.py b/Python/pywarpx/__init__.py index bddbb78b7..c3f824e8f 100644 --- a/Python/pywarpx/__init__.py +++ b/Python/pywarpx/__init__.py @@ -5,12 +5,14 @@ from .Geometry import geometry from .Algo import algo from .Langmuirwave import langmuirwave from .Interpolation import interpolation -from .Particles import particles +from .Particles import particles, electrons, positrons, protons, newspecies from .Laser import laser -from .AMReX import AMReX #from .timestepper import TimeStepper from .PGroup import PGroup from .PGroup import PGroups +#from .WarpXPIC import WarpXPIC from ._libwarpx import add_particles + +from .callbacks import * diff --git a/Python/pywarpx/callbacks.py b/Python/pywarpx/callbacks.py new file mode 100644 index 000000000..68629d164 --- /dev/null +++ b/Python/pywarpx/callbacks.py @@ -0,0 +1,481 @@ +"""call back operations +===================== + +These are the functions which allow installing user created functions so that +they are called at various places along the time step. + +For each call back type, the following three functions are defined. + - install___: Installs a function to be called at that specified time + - uninstall___: Uninstalls the function (so it won't be called anymore) + - isinstalled___: Checks if the function is installed + +These functions all take a function or instance method as an argument. Note that +if an instance method is used, an extra reference to the method's object is saved. + +The install can be done using a decorator, which has the prefix "callfrom". See example below. + +Functions can be called at the following times: + - afterinit <installafterinit>: immediately after the init is complete + - beforeEsolve <installbeforeEsolve>: before the solve for E fields + - afterEsolve <installafterEsolve>: after the solve for E fields + - beforedeposition <installbeforedeposition>: before the particle deposition (for charge and/or current) + - afterdeposition <installafterdeposition>: after particle deposition (for charge and/or current) + - beforestep <installbeforestep>: before the time step + - afterstep <installafterstep>: after the time step + - particlescraper <installparticlescraper>: just after the particle boundary conditions are applied + but before lost particles are processed + - particleloader <installparticleloader>: at the time that the standard particle loader is called + - particleinjection <installparticleinjection>: called when particle injection happens, after the position + advance and before deposition is called, allowing a user defined + particle distribution to be injected each time step + - appliedfields <installappliedfields>: allows directly specifying any fields to be applied to the particles + during the advance + +To use a decorator, the syntax is as follows. This will install the function myplots to be called after each step. + +@callfromafterstep +def myplots(): + ppzx() + +This is equivalent to the following: + +def myplots(): + ppzx() +installafterstep(myplots) + +""" +from __future__ import generators + +import types +import copy +import time +import ctypes +import sys +import numpy + +from ._libwarpx import libwarpx + +class CallbackFunctions(object): + """ + Class to handle call back function lists. + + Note that for functions passed in that are methods of a class instance, + a full reference of the instance is saved. This extra reference means + that the object will not actually deleted if the user deletes the + original reference. This is good since the user does not need to keep + the reference to the object (for example it can be created using a local + variable in a function). It may be bad if the user thinks an object was + deleted, but it actually isn't since it had (unkown to the user) + installed a method in one of the call back lists + """ + + def __init__(self,name=None,lcallonce=0): + self.funcs = [] + self.time = 0. + self.timers = {} + self.name = name + self.lcallonce = lcallonce + + def __call__(self,*args,**kw): + "Call all of the functions in the list" + tt = self.callfuncsinlist(*args,**kw) + self.time = self.time + tt + if self.lcallonce: self.funcs = [] + + def clearlist(self): + self.funcs = [] + + def __nonzero__(self): + "Returns True if functions are installed, otherwise False" + return self.hasfuncsinstalled() + + def __len__(self): + "Returns number of functions installed" + return len(self.funcs) + + def hasfuncsinstalled(self): + "Checks if there are any functions installed" + return len(self.funcs) > 0 + + def _getmethodobject(self,func): + "For call backs that are methods, returns the method's instance" + return func[0] + + def callbackfunclist(self): + "Generator returning callable functions from the list" + funclistcopy = copy.copy(self.funcs) + for f in funclistcopy: + if isinstance(f,list): + object = self._getmethodobject(f) + if object is None: + self.funcs.remove(f) + continue + result = getattr(object,f[1]) + elif isinstance(f,basestring): + import __main__ + if f in __main__.__dict__: + result = __main__.__dict__[f] + # --- If the function with the name is found, then replace the + # --- name in the list with the function. + self.funcs[self.funcs.index(f)] = result + else: + continue + else: + result = f + if not callable(result): + print("\n\nWarning: a call back was found that is not callable.") + if self.name is not None: + print("For %s"%self.name) + print("Only callable objects can be installed.") + print("It is possible that the callable's name has been overwritten") + print("by something not callable. This can happen during restart") + print("if a function name had later been used as a variable name.") + print(self.name) + if isinstance(f,basestring): + print("The name of the call back is %s"%f) + print("\n\n") + continue + yield result + + def installfuncinlist(self,f): + "Check if the specified function is installed" + if isinstance(f,types.MethodType): + # --- If the function is a method of a class instance, then save a full + # --- reference to that instance and the method name. + finstance = f.im_self + fname = f.__name__ + self.funcs.append([finstance,fname]) + elif callable(f): + # --- If a function had already been installed by name, then skip the install. + # --- This is problematic, since no warning message is given, but it is unlikely + # --- to arise under normal circumstances. + # --- The purpose of this check is to avoid redundant installation of functions + # --- during a restore from a dump file. Without the check, functions that had been + # --- installed via a decorator would be installed an extra time since the source + # --- of the function contains the decoration (which is activated when the source + # --- is exec'd). + if f.__name__ not in self.funcs: + self.funcs.append(f) + else: + self.funcs.append(f) + + def uninstallfuncinlist(self,f): + "Uninstall the specified function" + # --- An element by element search is needed + # --- f can be a function or method object, or a name (string). + # --- Note that method objects can not be removed by name. + funclistcopy = copy.copy(self.funcs) + for func in funclistcopy: + if f == func: + self.funcs.remove(f) + return + elif isinstance(func,list) and isinstance(f,types.MethodType): + object = self._getmethodobject(func) + if f.im_self is object and f.__name__ == func[1]: + self.funcs.remove(func) + return + elif isinstance(func,basestring): + if f.__name__ == func: + self.funcs.remove(func) + return + elif isinstance(f,basestring): + if isinstance(func,basestring): funcname = func + elif isinstance(func,list): funcname = None + else: funcname = func.__name__ + if f == funcname: + self.funcs.remove(func) + return + raise Exception('Warning: no such function had been installed') + + def isinstalledfuncinlist(self,f): + "Checks if the specified function is installed" + # --- An element by element search is needed + funclistcopy = copy.copy(self.funcs) + for func in funclistcopy: + if f == func: + return 1 + elif isinstance(func,list) and isinstance(f,types.MethodType): + object = self._getmethodobject(func) + if f.im_self is object and f.__name__ == func[1]: + return 1 + elif isinstance(func,basestring): + if f.__name__ == func: + return 1 + return 0 + + def callfuncsinlist(self,*args,**kw): + "Call the functions in the list" + bb = time.time() + for f in self.callbackfunclist(): + #barrier() + t1 = time.time() + f(*args,**kw) + #barrier() + t2 = time.time() + # --- For the timers, use the function (or method) name as the key. + self.timers[f.__name__] = self.timers.get(f.__name__,0.) + (t2 - t1) + aa = time.time() + return aa - bb + +#============================================================================= + +# --- Now create the actual instances. +_afterinit = CallbackFunctions('afterinit') +_beforeEsolve = CallbackFunctions('beforeEsolve') +_afterEsolve = CallbackFunctions('afterEsolve') +_beforedeposition = CallbackFunctions('beforedeposition') +_afterdeposition = CallbackFunctions('afterdeposition') +_particlescraper = CallbackFunctions('particlescraper') +_particleloader = CallbackFunctions('particleloader') +_beforestep = CallbackFunctions('beforestep') +_afterstep = CallbackFunctions('afterstep') +_afterrestart = CallbackFunctions('afterrestart',lcallonce=1) +_particleinjection = CallbackFunctions('particleinjection') +_appliedfields = CallbackFunctions('appliedfields') + +# --- Create the objects that can be called from C. +# --- Note that each of the CFUNCTYPE instances need to be saved +_CALLBACK_FUNC_0 = ctypes.CFUNCTYPE(None) +_c_afterinit = _CALLBACK_FUNC_0(_afterinit) +libwarpx.warpx_set_callback_py_afterinit(_c_afterinit) +_c_beforeEsolve = _CALLBACK_FUNC_0(_beforeEsolve) +libwarpx.warpx_set_callback_py_beforeEsolve(_c_beforeEsolve) +_c_afterEsolve = _CALLBACK_FUNC_0(_afterEsolve) +libwarpx.warpx_set_callback_py_afterEsolve(_c_afterEsolve) +_c_beforedeposition = _CALLBACK_FUNC_0(_beforedeposition) +libwarpx.warpx_set_callback_py_beforedeposition(_c_beforedeposition) +_c_afterdeposition = _CALLBACK_FUNC_0(_afterdeposition) +libwarpx.warpx_set_callback_py_afterdeposition(_c_afterdeposition) +_c_particlescraper = _CALLBACK_FUNC_0(_particlescraper) +libwarpx.warpx_set_callback_py_particlescraper(_c_particlescraper) +_c_particleloader = _CALLBACK_FUNC_0(_particleloader) +libwarpx.warpx_set_callback_py_particleloader(_c_particleloader) +_c_beforestep = _CALLBACK_FUNC_0(_beforestep) +libwarpx.warpx_set_callback_py_beforestep(_c_beforestep) +_c_afterstep = _CALLBACK_FUNC_0(_afterstep) +libwarpx.warpx_set_callback_py_afterstep(_c_afterstep) +_c_afterrestart = _CALLBACK_FUNC_0(_afterrestart) +libwarpx.warpx_set_callback_py_afterrestart(_c_afterrestart) +_c_particleinjection = _CALLBACK_FUNC_0(_particleinjection) +libwarpx.warpx_set_callback_py_particleinjection(_c_particleinjection) +_c_appliedfields = _CALLBACK_FUNC_0(_appliedfields) +libwarpx.warpx_set_callback_py_appliedfields(_c_appliedfields) + +#============================================================================= +def printcallbacktimers(tmin=1.,lminmax=False,ff=None): + """Prints timings of installed functions. + - tmin=1.: only functions with time greater than tmin will be printed + - lminmax=False: If True, prints the min and max times over all processors + - ff=None: If given, timings will be written to the file object instead of stdout + """ + if ff is None: ff = sys.stdout + for c in [_afterinit,_beforeEsolve,_afterEsolve, + _beforedeposition,_afterdeposition, + _particlescraper, + _particleloader, + _beforestep,_afterstep, + _afterrestart, + _particleinjection, + _appliedfields]: + for fname, time in c.timers.items(): + #vlist = numpy.array(gather(time)) + vlist = numpy.array([time]) + #if me > 0: continue + vsum = numpy.sum(vlist) + if vsum <= tmin: continue + vrms = numpy.sqrt(max(0.,numpy.sum(vlist**2)/len(vlist) - (numpy.sum(vlist)/len(vlist))**2)) + npes = 1. # Only works for one processor + ff.write('%20s %s %10.4f %10.4f %10.4f'%(c.name,fname,vsum,vsum/npes,vrms)) + if lminmax: + vmin = numpy.min(vlist) + vmax = numpy.max(vlist) + ff.write(' %10.4f %10.4f'%(vmin,vmax)) + it = libwarpx.warpx_getistep(0) + if it > 0: + ff.write(' %10.4f'%(vsum/npes/(it))) + ff.write('\n') + +#============================================================================= +# ---------------------------------------------------------------------------- +def callfromafterinit(f): + installafterinit(f) + return f +def installafterinit(f): + "Adds a function to the list of functions called after the init" + _afterinit.installfuncinlist(f) +def uninstallafterinit(f): + "Removes the function from the list of functions called after the init" + _afterinit.uninstallfuncinlist(f) +def isinstalledafterinit(f): + "Checks if the function is called after a init" + return _afterinit.isinstalledfuncinlist(f) + +# ---------------------------------------------------------------------------- +def callfrombeforeEsolve(f): + installbeforeEsolve(f) + return f +def installbeforeEsolve(f): + "Adds a function to the list of functions called before an E solve" + _beforeEsolve.installfuncinlist(f) +def uninstallbeforeEsolve(f): + "Removes the function from the list of functions called before an E solve" + _beforeEsolve.uninstallfuncinlist(f) +def isinstalledbeforeEsolve(f): + "Checks if the function is called before an E solve" + return _beforeEsolve.isinstalledfuncinlist(f) + +# ---------------------------------------------------------------------------- +def callfromafterEsolve(f): + installafterEsolve(f) + return f +def installafterEsolve(f): + "Adds a function to the list of functions called after an E solve" + _afterEsolve.installfuncinlist(f) +def uninstallafterEsolve(f): + "Removes the function from the list of functions called after an E solve" + _afterEsolve.uninstallfuncinlist(f) +def isinstalledafterEsolve(f): + "Checks if the function is called after an E solve" + return _afterEsolve.isinstalledfuncinlist(f) + +# ---------------------------------------------------------------------------- +def callfrombeforedeposition(f): + installbeforedeposition(f) + return f +def installbeforedeposition(f): + "Adds a function to the list of functions called before a particle deposition" + _beforedeposition.installfuncinlist(f) +def uninstallbeforedeposition(f): + "Removes the function from the list of functions called before a particle deposition" + _beforedeposition.uninstallfuncinlist(f) +def isinstalledbeforedeposition(f): + "Checks if the function is called before a particle deposition" + return _beforedeposition.isinstalledfuncinlist(f) + +# ---------------------------------------------------------------------------- +def callfromafterdeposition(f): + installafterdeposition(f) + return f +def installafterdeposition(f): + "Adds a function to the list of functions called after a particle deposition" + _afterdeposition.installfuncinlist(f) +def uninstallafterdeposition(f): + "Removes the function from the list of functions called after a particle deposition" + _afterdeposition.uninstallfuncinlist(f) +def isinstalledafterdeposition(f): + "Checks if the function is called after a particle deposition" + return _afterdeposition.isinstalledfuncinlist(f) + +# ---------------------------------------------------------------------------- +def callfromparticlescraper(f): + installparticlescraper(f) + return f +def installparticlescraper(f): + "Adds a function to the list of functions called to scrape particles" + _particlescraper.installfuncinlist(f) +def uninstallparticlescraper(f): + "Removes the function from the list of functions called to scrape particles" + _particlescraper.uninstallfuncinlist(f) +def isinstalledparticlescraper(f): + "Checks if the function is called to scrape particles" + return _particlescraper.isinstalledfuncinlist(f) + +# ---------------------------------------------------------------------------- +def callfromparticleloader(f): + installparticleloader(f) + return f +def installparticleloader(f): + "Adds a function to the list of functions called to load particles" + _particleloader.installfuncinlist(f) +def uninstallparticleloader(f): + "Removes the function from the list of functions called to load particles" + _particleloader.uninstallfuncinlist(f) +def isinstalledparticleloader(f): + "Checks if the function is called to load particles" + return _particleloader.isinstalledfuncinlist(f) + +# ---------------------------------------------------------------------------- +def callfrombeforestep(f): + installbeforestep(f) + return f +def installbeforestep(f): + "Adds a function to the list of functions called before a step" + _beforestep.installfuncinlist(f) +def uninstallbeforestep(f): + "Removes the function from the list of functions called before a step" + _beforestep.uninstallfuncinlist(f) +def isinstalledbeforestep(f): + "Checks if the function is called before a step" + return _beforestep.isinstalledfuncinlist(f) + +# ---------------------------------------------------------------------------- +def callfromafterstep(f): + installafterstep(f) + return f +def installafterstep(f): + "Adds a function to the list of functions called after a step" + _afterstep.installfuncinlist(f) +def uninstallafterstep(f): + "Removes the function from the list of functions called after a step" + _afterstep.uninstallfuncinlist(f) +def isinstalledafterstep(f): + "Checks if the function is called after a step" + return _afterstep.isinstalledfuncinlist(f) + +# ---------------------------------------------------------------------------- +def callfromafterrestart(f): + raise Exception('restart call back not implemented yet') + installafterrestart(f) + return f +def installafterrestart(f): + "Adds a function to the list of functions called immediately after a restart" + raise Exception('restart call back not implemented yet') + _afterrestart.installfuncinlist(f) +def uninstallafterrestart(f): + "Removes the function from the list of functions called immediately after a restart" + raise Exception('restart call back not implemented yet') + _afterrestart.uninstallfuncinlist(f) +def isinstalledafterrestart(f): + "Checks if the function is called immediately after a restart" + raise Exception('restart call back not implemented yet') + return _afterrestart.isinstalledfuncinlist(f) + +# ---------------------------------------------------------------------------- +def callfromparticleinjection(f): + installparticleinjection(f) + return f +def installparticleinjection(f): + """ + Adds a user defined function that is to be called when particle + injection happens, after the position advance and before deposition is + called, allowing a user defined particle distribution to be injected + each time step""" + _particleinjection.installfuncinlist(f) +def uninstallparticleinjection(f): + "Removes the function installed by installparticleinjection" + _particleinjection.uninstallfuncinlist(f) +def isinstalledparticleinjection(f): + "Checks if the function is called when particles injection happens" + return _particleinjection.isinstalledfuncinlist(f) + +# ---------------------------------------------------------------------------- +def callfromappliedfields(f): + raise Exception('applied fields call back not implemented yet') + installappliedfields(f) + return f +def installappliedfields(f): + """ + Adds a user defined function which can specify E and B fields which are applied + to the particles during the particle advance. + """ + raise Exception('applied fields call back not implemented yet') + _appliedfields.installfuncinlist(f) +def uninstallappliedfields(f): + "Removes the function installed by installappliedfields" + raise Exception('applied fields call back not implemented yet') + _appliedfields.uninstallfuncinlist(f) +def isinstalledappliedfields(f): + "Checks if the function is called when which applies fields" + raise Exception('applied fields call back not implemented yet') + return _appliedfields.isinstalledfuncinlist(f) + diff --git a/Python/pywarpx/picmi.py b/Python/pywarpx/picmi.py new file mode 100644 index 000000000..11e1df6a2 --- /dev/null +++ b/Python/pywarpx/picmi.py @@ -0,0 +1,420 @@ +"""Classes following the PICMI standard +""" +import PICMI_Base +import numpy as np +import pywarpx + +codename = 'warpx' +PICMI_Base.register_codename(codename) + +# --- Values from WarpXConst.H +c = 299792458. +ep0 = 8.854187817e-12 +mu0 = 1.2566370614359173e-06 +q_e = 1.602176462e-19 +m_e = 9.10938291e-31 +m_p = 1.6726231e-27 + + +class Species(PICMI_Base.PICMI_Species): + def init(self, kw): + + if self.particle_type == 'electron': + if self.charge is None: self.charge = '-q_e' + if self.mass is None: self.mass = 'm_e' + elif self.particle_type == 'positron': + if self.charge is None: self.charge = 'q_e' + if self.mass is None: self.mass = 'm_e' + elif self.particle_type == 'proton': + if self.charge is None: self.charge = 'q_e' + if self.mass is None: self.mass = 'm_p' + elif self.particle_type == 'anti-proton': + if self.charge is None: self.charge = '-q_e' + if self.mass is None: self.mass = 'm_p' + + def initialize_inputs(self, layout): + self.species_number = pywarpx.particles.nspecies + pywarpx.particles.nspecies += 1 + + if self.name is None: + self.name = 'species{}'.format(self.species_number) + + if pywarpx.particles.species_names is None: + pywarpx.particles.species_names = self.name + else: + pywarpx.particles.species_names += ' ' + self.name + + self.species = pywarpx.Bucket.Bucket(self.name, mass=self.mass, charge=self.charge, injection_style = 'python') + pywarpx.Particles.particles_list.append(self.species) + + if self.initial_distribution is not None: + self.initial_distribution.initialize_inputs(self.species_number, layout, self.species) + + +PICMI_Base.PICMI_MultiSpecies.Species_class = Species +class MultiSpecies(PICMI_Base.PICMI_MultiSpecies): + pass + + +class GaussianBunchDistribution(PICMI_Base.PICMI_GaussianBunchDistribution): + def init(self, kw): + if self.seed is not None: + print('Warning: WarpX does not support specifying the random number seed') + + def initialize_inputs(self, species_number, layout, species): + species.injection_style = "gaussian_beam" + species.x_m = self.centroid_position[0] + species.y_m = self.centroid_position[1] + species.z_m = self.centroid_position[2] + species.x_rms = self.rms_bunch_size[0] + species.y_rms = self.rms_bunch_size[1] + species.z_rms = self.rms_bunch_size[2] + + # --- Only PseudoRandomLayout is supported + species.npart = layout.n_macroparticles + + # --- Calculate the total charge. Note that charge might be a string instead of a number. + charge = species.charge + if charge == 'q_e' or charge == '+q_e': + charge = q_e + elif charge == '-q_e': + charge = -q_e + species.q_tot = self.number_real_particles*charge + + # --- These need to be defined even though they are not used + species.profile = "constant" + species.density = 1 + + # --- The PICMI standard doesn't yet have a way of specifying these values. + # --- They should default to the size of the domain. They are not typically + # --- necessary though since any particles outside the domain are rejected. + #species.xmin + #species.xmax + #species.ymin + #species.ymax + #species.zmin + #species.zmax + + if np.any(np.not_equal(self.velocity_divergence, 0.)): + species.momentum_distribution_type = "radial_expansion" + species.u_over_r = self.velocity_divergence[0] + #species.u_over_y = self.velocity_divergence[1] + #species.u_over_z = self.velocity_divergence[2] + elif np.any(np.not_equal(self.rms_velocity, 0.)): + species.momentum_distribution_type = "gaussian" + species.ux_m = self.centroid_velocity[0] + species.uy_m = self.centroid_velocity[1] + species.uz_m = self.centroid_velocity[2] + species.ux_th = self.rms_velocity[0] + species.uy_th = self.rms_velocity[1] + species.uz_th = self.rms_velocity[2] + else: + species.momentum_distribution_type = "constant" + species.ux = self.centroid_velocity[0] + species.uy = self.centroid_velocity[1] + species.uz = self.centroid_velocity[2] + + +class UniformDistribution(PICMI_Base.PICMI_UniformDistribution): + def initialize_inputs(self, species_number, layout, species): + + if isinstance(layout, GriddedLayout): + # --- Note that the grid attribute of GriddedLayout is ignored + species.injection_style = "nuniformpercell" + species.num_particles_per_cell_each_dim = layout.n_macroparticle_per_cell + elif isinstance(layout, PseudoRandomLayout): + assert (layout.n_macroparticles_per_cell is not None), Exception('WarpX only supports n_macroparticles_per_cell for the PseudoRandomLayout with UniformDistribution') + species.injection_style = "nrandompercell" + species.num_particles_per_cell = layout.n_macroparticles_per_cell + else: + raise Exception('WarpX does not support the specified layout for UniformDistribution') + + species.xmin = self.lower_bound[0] + species.xmax = self.upper_bound[0] + species.ymin = self.lower_bound[1] + species.ymax = self.upper_bound[1] + species.zmin = self.lower_bound[2] + species.zmax = self.upper_bound[2] + + # --- Only constant density is supported at this time + species.profile = "constant" + species.density = self.density + + if np.any(np.not_equal(self.rms_velocity, 0.)): + species.momentum_distribution_type = "gaussian" + species.ux_m = self.directed_velocity[0] + species.uy_m = self.directed_velocity[1] + species.uz_m = self.directed_velocity[2] + species.ux_th = self.rms_velocity[0] + species.uy_th = self.rms_velocity[1] + species.uz_th = self.rms_velocity[2] + else: + species.momentum_distribution_type = "constant" + species.ux = self.directed_velocity[0] + species.uy = self.directed_velocity[1] + species.uz = self.directed_velocity[2] + + if self.fill_in: + pywarpx.warpx.do_plasma_injection = 1 + if not hasattr(pywarpx.warpx, 'injected_plasma_species'): + pywarpx.warpx.injected_plasma_species = [] + + pywarpx.warpx.injected_plasma_species.append(species_number) + pywarpx.warpx.num_injected_species = len(pywarpx.warpx.injected_plasma_species) + + +class AnalyticDistribution(PICMI_Base.PICMI_AnalyticDistribution): + def initialize_inputs(self, species_number, layout, species): + raise Exception('WarpX does not support AnalyticDistribution') + + +class ParticleList(PICMI_Base.PICMI_ParticleList): + def init(self, kw): + + if len(x) > 1: + raise Exception('Only a single particle can be loaded') + + def initialize_inputs(self, species_number, layout, species): + + species.injection_style = "singleparticle" + species.single_particle_pos = [self.x[0], self.y[0], self.z[0]] + species.single_particle_vel = [self.ux[0]/c, self.uy[0]/c, self.uz[0]/c] + species.single_particle_weight = self.weight + + # --- These need to be defined even though they are not used + species.profile = "constant" + species.density = 1 + species.momentum_distribution_type = 'constant' + + +class ParticleDistributionPlanarInjector(PICMI_Base.PICMI_ParticleDistributionPlanarInjector): + pass + + +class GriddedLayout(PICMI_Base.PICMI_GriddedLayout): + pass + + +class PseudoRandomLayout(PICMI_Base.PICMI_PseudoRandomLayout): + pass + + +class BinomialSmoother(PICMI_Base.PICMI_BinomialSmoother): + pass + + +class CylindricalGrid(PICMI_Base.PICMI_CylindricalGrid): + def init(self, kw): + raise Exception('WarpX does not support CylindricalGrid') + + +class Cartesian2DGrid(PICMI_Base.PICMI_Cartesian2DGrid): + def init(self, kw): + self.max_grid_size = kw.pop('warpx_max_grid_size', 32) + self.max_level = kw.pop('warpx_max_level', 0) + self.coord_sys = kw.pop('warpx_coord_sys', 0) + + def initialize_inputs(self): + pywarpx.amr.n_cell = self.number_of_cells + + # Maximum allowable size of each subdomain in the problem domain; + # this is used to decompose the domain for parallel calculations. + pywarpx.amr.max_grid_size = self.max_grid_size + + # Maximum level in hierarchy (for now must be 0, i.e., one level in total) + pywarpx.amr.max_level = self.max_level + + # Geometry + pywarpx.geometry.coord_sys = self.coord_sys + pywarpx.geometry.is_periodic = '%d %d %d'%(self.bc_xmin=='periodic', self.bc_ymin=='periodic') # Is periodic? + pywarpx.geometry.prob_lo = self.lower_bound # physical domain + pywarpx.geometry.prob_hi = self.upper_bound + + 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 + if self.moving_window_velocity[0] != 0.: + pywarpx.warpx.moving_window_dir = 'x' + pywarpx.warpx.moving_window_v = self.moving_window_velocity[0]/c # in units of the speed of light + if self.moving_window_velocity[1] != 0.: + pywarpx.warpx.moving_window_dir = 'y' + pywarpx.warpx.moving_window_v = self.moving_window_velocity[1]/c # in units of the speed of light + + +class Cartesian3DGrid(PICMI_Base.PICMI_Cartesian3DGrid): + def init(self, kw): + self.max_grid_size = kw.pop('warpx_max_grid_size', 32) + self.max_level = kw.pop('warpx_max_level', 0) + self.coord_sys = kw.pop('warpx_coord_sys', 0) + + def initialize_inputs(self): + pywarpx.amr.n_cell = self.number_of_cells + + # Maximum allowable size of each subdomain in the problem domain; + # this is used to decompose the domain for parallel calculations. + pywarpx.amr.max_grid_size = self.max_grid_size + + # Maximum level in hierarchy (for now must be 0, i.e., one level in total) + pywarpx.amr.max_level = self.max_level + + # Geometry + pywarpx.geometry.coord_sys = self.coord_sys + pywarpx.geometry.is_periodic = '%d %d %d'%(self.bc_xmin=='periodic', self.bc_ymin=='periodic', self.bc_zmin=='periodic') # Is periodic? + pywarpx.geometry.prob_lo = self.lower_bound # physical domain + pywarpx.geometry.prob_hi = self.upper_bound + + 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 + if self.moving_window_velocity[0] != 0.: + pywarpx.warpx.moving_window_dir = 'x' + pywarpx.warpx.moving_window_v = self.moving_window_velocity[0]/c # in units of the speed of light + if self.moving_window_velocity[1] != 0.: + pywarpx.warpx.moving_window_dir = 'y' + pywarpx.warpx.moving_window_v = self.moving_window_velocity[1]/c # in units of the speed of light + if self.moving_window_velocity[2] != 0.: + pywarpx.warpx.moving_window_dir = 'z' + pywarpx.warpx.moving_window_v = self.moving_window_velocity[2]/c # in units of the speed of light + + +class ElectromagneticSolver(PICMI_Base.PICMI_ElectromagneticSolver): + def init(self, kw): + assert self.method is None or self.method in ['Yee'], Exception("Only 'Yee' FDTD is supported") + + self.do_pml = kw.pop('warpx_do_pml', None) + + def initialize_inputs(self): + + self.grid.initialize_inputs() + + pywarpx.warpx.do_pml = self.do_pml + + if self.cfl is not None: + pywarpx.warpx.cfl = self.cfl + + if self.stencil_order is not None: + pywarpx.interpolation.nox = self.stencil_order[0] + pywarpx.interpolation.noy = self.stencil_order[1] + pywarpx.interpolation.noz = self.stencil_order[2] + + +class Electrostatic_solver(PICMI_Base.PICMI_Electrostatic_solver): + def initialize_inputs(self): + pass + + +class GaussianLaser(PICMI_Base.PICMI_GaussianLaser): + def initialize_inputs(self): + pywarpx.warpx.use_laser = 1 + pywarpx.laser.profile = "Gaussian" + pywarpx.laser.wavelength = self.wavelength # The wavelength of the laser (in meters) + pywarpx.laser.e_max = self.E0 # Maximum amplitude of the laser field (in V/m) + pywarpx.laser.polarization = [np.cos(self.polarization_angle), np.sin(self.polarization_angle), 0.] # The main polarization vector + pywarpx.laser.profile_waist = self.waist # The waist of the laser (in meters) + pywarpx.laser.profile_duration = self.duration # The duration of the laser (in seconds) + + +class LaserAntenna(PICMI_Base.PICMI_LaserAntenna): + def initialize_inputs(self, laser): + pywarpx.laser.position = self.position # This point is on the laser plane + pywarpx.laser.direction = self.normal_vector # The plane normal direction + pywarpx.laser.profile_focal_distance = laser.focal_position[2] - self.position[2] # Focal distance from the antenna (in meters) + pywarpx.laser.profile_t_peak = (self.position[2] - laser.centroid_position[2])/c # The time at which the laser reaches its peak (in seconds) + + +class Simulation(PICMI_Base.PICMI_Simulation): + def init(self, kw): + + self.plot_int = kw.pop('warpx_plot_int', None) + self.plot_file = kw.pop('warpx_plot_file', None) + self.current_deposition_algo = kw.pop('warpx_current_deposition_algo', None) + self.charge_deposition_algo = kw.pop('warpx_charge_deposition_algo', None) + self.field_gathering_algo = kw.pop('warpx_field_gathering_algo', None) + self.particle_pusher_algo = kw.pop('warpx_particle_pusher_algo', None) + self.use_filter = kw.pop('warpx_use_filter', None) + self.serialize_ics = kw.pop('warpx_serialize_ics', None) + self.do_dynamic_scheduling = kw.pop('warpx_do_dynamic_scheduling', None) + self.load_balance_int = kw.pop('warpx_load_balance_int', None) + self.load_balance_with_sfc = kw.pop('warpx_load_balance_with_sfc', None) + + self.inputs_initialized = False + self.warpx_initialized = False + + def initialize_inputs(self): + if self.inputs_initialized: + return + + self.inputs_initialized = True + + pywarpx.warpx.verbose = self.verbose + if self.time_step_size is not None: + pywarpx.warpx.const_dt = self.timestep + + pywarpx.amr.plot_int = self.plot_int + pywarpx.amr.plot_file = self.plot_file + pywarpx.algo.current_deposition = self.current_deposition_algo + pywarpx.algo.charge_deposition = self.charge_deposition_algo + pywarpx.algo.field_gathering = self.field_gathering_algo + pywarpx.algo.particle_pusher = self.particle_pusher_algo + + pywarpx.warpx.use_filter = self.use_filter + pywarpx.warpx.serialize_ics = self.serialize_ics + + pywarpx.warpx.do_dynamic_scheduling = self.do_dynamic_scheduling + pywarpx.warpx.load_balance_int = self.load_balance_int + pywarpx.warpx.load_balance_with_sfc = self.load_balance_with_sfc + + particle_shape = self.particle_shape + for s in self.species: + if s.particle_shape is not None: + assert particle_shape is None or particle_shape == s.particle_shape, Exception('WarpX only supports one particle shape for all species') + # --- If this was set for any species, use that value. + particle_shape = s.particle_shape + + if particle_shape is not None: + if isinstance(particle_shape, str): + interpolation_order = {'NGP':0, 'linear':1, 'quadratic':2, 'cubic':3}[particle_shape] + else: + interpolation_order = particle_shape + pywarpx.interpolation.nox = interpolation_order + pywarpx.interpolation.noy = interpolation_order + pywarpx.interpolation.noz = interpolation_order + + self.solver.initialize_inputs() + + for i in range(len(self.species)): + assert not self.initialize_self_fields[i], Exception('WarpX does not support initializing self fields') + self.species[i].initialize_inputs(self.layouts[i]) + + for i in range(len(self.lasers)): + self.lasers[i].initialize_inputs() + self.laser_injection_methods[i].initialize_inputs(self.lasers[i]) + + def initialize_warpx(self): + if self.warpx_initialized: + return + + self.warpx_initialized = True + pywarpx.warpx.init() + + def write_input_file(self, file_name='inputs'): + self.initialize_inputs() + kw = {} + if self.max_steps is not None: + kw['max_step'] = self.max_steps + if self.max_time is not None: + kw['stop_time'] = self.max_time + pywarpx.warpx.write_inputs(file_name, **kw) + + def step(self, nsteps=None): + self.initialize_inputs() + self.initialize_warpx() + if nsteps is None: + if self.max_steps is not None: + nsteps = self.max_steps + else: + nsteps = -1 + pywarpx.warpx.evolve(nsteps) + + def finalize(self): + if self.warpx_initialized: + self.warpx_initialized = False + pywarpx.warpx.finalize() diff --git a/Python/pywarpx/timestepper.py b/Python/pywarpx/timestepper.py index f72b1669f..3068cbc82 100644 --- a/Python/pywarpx/timestepper.py +++ b/Python/pywarpx/timestepper.py @@ -1,8 +1,5 @@ -import warp -from warp import top -from warp import w3d - from ._libwarpx import libwarpx +from . import callbacks class TimeStepper(object): @@ -12,6 +9,8 @@ class TimeStepper(object): def onestep(self): + callbacks._beforestep() + self.cur_time = libwarpx.warpx_gett_new(0) self.istep = libwarpx.warpx_getistep(0) @@ -33,7 +32,11 @@ class TimeStepper(object): # --- Evolve particles to p^{n+1/2} and x^{n+1} # --- Depose current, j^{n+1/2} + callbacks._particleinjection() + callbacks._particlescraper() + callbacks._beforedeposition() libwarpx.warpx_PushParticlesandDepose(self.cur_time) + callbacks._afterdeposition() libwarpx.mypc_Redistribute() # Redistribute particles @@ -43,7 +46,9 @@ class TimeStepper(object): libwarpx.warpx_SyncCurrent() libwarpx.warpx_FillBoundaryB() + callbacks._beforeEsolve() libwarpx.warpx_EvolveE(dt,0) # We now have E^{n+1} + callbacks._afterEsolve() self.istep += 1 @@ -67,318 +72,4 @@ class TimeStepper(object): if libwarpx.warpx_checkInt() > 0 and (self.istep+1)%libwarpx.warpx_plotInt() == 0 or max_time_reached: libwarpx.warpx_WriteCheckPointFile() - - - - - - - - - - - -# --- This is not used -class TimeStepperFromPICSAR(object): - - def __init__(self, package=None, solver=None, l_debug=False): - self.package = package - self.solver = solver - self.l_debug = l_debug - - def setpackagestepnumber(self, it): - if self.package is not None: - self.package.setstepnumber(it) - - def step(self, n=1, freq_print=10, lallspecl=0): - """ - This function performs a range of Particle-In-Cell iterations - - Inputs: - - n: number of iterations - - freq_print: print frequency - """ - - if (self.l_debug): print("Call step") - - for i in range(n): - if(me == 0): - if top.it%freq_print == 0: - print('it = %g time = %g'%(top.it, top.time)) - - l_first = (lallspecl or (i == 0)) - l_last = (lallspecl or (i == n-1)) - - self.onestep(l_first, l_last) - - if (self.l_debug): print("End step") - - def onestep(self, l_first, l_last): - """ - Perform a single particle-in-cell step - """ - - if (self.l_debug): print("Call onestep") - - # --- Iteration number - self.setpackagestepnumber(top.it) - - # --- call beforestep functions - if (self.l_debug): print("Call beforestep functions") - warp.callbeforestepfuncs.callfuncsinlist() - - # --- gather fields from grid to particles - if (self.l_debug): print("Call Field gathering and particle push") - - # --- push - if l_first: - if self.package is None: - # --- Standard Warp advance - for specie in warp.listofallspecies: - for pg in specie.flatten(specie.pgroups): - for js in range(pg.ns): - self.push_velocity_second_part(js, pg) - self.push_positions(js, pg) - warp.particleboundaries3d(pg, -1, False) - - else: - # --- Particle pusher - if (self.l_debug): print("Call package particle push") - pxr.pxrpush_particles_part2() - - # --- Particle boundary consitions - if (self.l_debug): print("Call package particle boundary conditions") - pxr.particle_bcs() - - if (self.l_debug): print("Call aliasparticlearrays()") - self.aliasparticlearrays() - - else: - if self.package is None: - # --- Standard Warp advance - - for specie in warp.listofallspecies: - for pg in specie.flatten(specie.pgroups): - for js in range(pg.ns): - self.push_velocity_full(js, pg) - self.push_positions(js, pg) - - warp.particleboundaries3d(pg, -1, False) - - else: - # --- Particle pusher - if (self.l_debug): print("Call package particle pusher") - pxr.field_gathering_plus_particle_pusher() - - # --- Particle boundary conditions - if (self.l_debug): print("Call package particle boundary conditions") - pxr.particle_bcs() - - self.aliasparticlearrays() - - # --- Particle sorting - if (self.l_debug): print("Call Particle Sorting") - if self.package is not None: - # --- This should be a function installed before load rho - pxr.particle_sorting_sub() - - # --- call beforeloadrho functions - if (self.l_debug): print("Call beforeloadrho functions") - warp.beforeloadrho.callfuncsinlist() - - pgroups = [] - for specie in warp.listofallspecies: - pgroups += specie.flatten(specie.pgroups) - - self.pgroups = pgroups - - # --- Call user-defined injection routines - if (self.l_debug): print("Call user-defined injection routines") - warp.userinjection.callfuncsinlist() - - xgriddiff = w3d.xmmin - pxr.xmin - ygriddiff = w3d.ymmin - pxr.ymin - zgriddiff = w3d.zmmin - pxr.zmin - - if (xgriddiff != 0 or ygriddiff != 0 or zgriddiff != 0): - pxr.pxr_move_sim_boundaries(xgriddiff, ygriddiff, zgriddiff) - - - if (self.l_debug): print("Call loadrho") - self.solver.loadrho(pgroups = pgroups) - if (self.l_debug): print("Call loadj") - self.solver.loadj(pgroups = pgroups) - - if (self.l_debug): print("Call dosolve") - self.solver.dosolve() - - if self.package is None: - for specie in warp.listofallspecies: - for pg in specie.flatten(specie.pgroups): - for js in range(pg.ns): - self.fetcheb(js, pg) - if l_last: - self.push_velocity_first_part(js, pg) - - else: - if l_last: - if (self.l_debug): print("Call package push particles 1") - pxr.pxrpush_particles_part1() - - # --- update time, time counter - top.time += top.dt - if top.it%top.nhist == 0: -# zmmnt() - minidiag(top.it, top.time, top.lspecial) - top.it += 1 - - # --- Load balance function should be installed after step - - # --- call afterstep functions - if (self.l_debug): print("Call callafterstepfuncs.callfuncsinlist()") - warp.callafterstepfuncs.callfuncsinlist() - - - # --- The following methods use the standard Warp routines - def fetcheb(self, js, pg=None): - if self.l_verbose:print(me, 'enter fetcheb') - if pg is None: - pg = top.pgroup - np = pg.nps[js] - if np == 0: return - il = pg.ins[js] - 1 - iu = il + pg.nps[js] - w3d.pgroupfsapi = pg - w3d.ipminfsapi = pg.ins[js] - w3d.npfsapi = pg.nps[js] - pg.ex[il:iu] = 0. - pg.ey[il:iu] = 0. - pg.ez[il:iu] = 0. - pg.bx[il:iu] = 0. - pg.by[il:iu] = 0. - pg.bz[il:iu] = 0. - self.fetche() - self.fetchb() - w3d.pgroupfsapi = None - - def push_velocity_full(self, js, pg=None): - if self.l_verbose:print(me, 'enter push_velocity_full') - if pg is None: - pg = top.pgroup - np = pg.nps[js] - if np == 0: return - il = pg.ins[js] - 1 - iu = il + pg.nps[js] - if pg.lebcancel_pusher: - warp.ebcancelpush3d(np, pg.uxp[il:iu], pg.uyp[il:iu], pg.uzp[il:iu], pg.gaminv[il:iu], - pg.ex[il:iu], pg.ey[il:iu], pg.ez[il:iu], - pg.bx[il:iu], pg.by[il:iu], pg.bz[il:iu], - pg.sq[js], pg.sm[js], top.dt, 0) - else: - # --- push velocity from electric field (half step) - warp.epush3d(np, pg.uxp[il:iu], pg.uyp[il:iu], pg.uzp[il:iu], - pg.ex[il:iu], pg.ey[il:iu], pg.ez[il:iu], - pg.sq[js], pg.sm[js], 0.5*top.dt) - # --- update gamma - self.set_gamma(js, pg) - # --- push velocity from magnetic field - warp.bpush3d(np, pg.uxp[il:iu], pg.uyp[il:iu], pg.uzp[il:iu], pg.gaminv[il:iu], - pg.bx[il:iu], pg.by[il:iu], pg.bz[il:iu], - pg.sq[js], pg.sm[js], top.dt, top.ibpush) - # --- push velocity from electric field (half step) - warp.epush3d(np, pg.uxp[il:iu], pg.uyp[il:iu], pg.uzp[il:iu], - pg.ex[il:iu], pg.ey[il:iu], pg.ez[il:iu], - pg.sq[js], pg.sm[js], 0.5*top.dt) - # --- update gamma - self.set_gamma(js, pg) - - if self.l_verbose:print(me, 'exit push_velocity_first_part') - - def push_velocity_first_part(self, js, pg=None): - if self.l_verbose:print(me, 'enter push_velocity_first_part') - if pg is None: - pg = top.pgroup - np = pg.nps[js] - if np == 0: return - il = pg.ins[js] - 1 - iu = il + pg.nps[js] - if pg.lebcancel_pusher: - warp.ebcancelpush3d(np, pg.uxp[il:iu], pg.uyp[il:iu], pg.uzp[il:iu], pg.gaminv[il:iu], - pg.ex[il:iu], pg.ey[il:iu], pg.ez[il:iu], - pg.bx[il:iu], pg.by[il:iu], pg.bz[il:iu], - pg.sq[js], pg.sm[js], top.dt, 1) - else: - # --- push velocity from electric field (half step) - warp.epush3d(np, pg.uxp[il:iu], pg.uyp[il:iu], pg.uzp[il:iu], - pg.ex[il:iu], pg.ey[il:iu], pg.ez[il:iu], - pg.sq[js], pg.sm[js], 0.5*top.dt) - # --- update gamma - self.set_gamma(js,pg) - # --- push velocity from magnetic field - warp.bpush3d(np, pg.uxp[il:iu], pg.uyp[il:iu], pg.uzp[il:iu], pg.gaminv[il:iu], - pg.bx[il:iu], pg.by[il:iu], pg.bz[il:iu], - pg.sq[js], pg.sm[js], 0.5*top.dt, top.ibpush) - - if self.l_verbose:print(me, 'exit push_velocity_first_part') - - def push_velocity_second_part(self, js, pg=None): - if self.l_verbose:print(me, 'enter push_velocity_second_part') - if pg is None: - pg = top.pgroup - np = pg.nps[js] - if np == 0: return - il = pg.ins[js] - 1 - iu = il + pg.nps[js] - if pg.lebcancel_pusher: - warp.ebcancelpush3d(np, pg.uxp[il:iu], pg.uyp[il:iu], pg.uzp[il:iu], pg.gaminv[il:iu], - pg.ex[il:iu], pg.ey[il:iu], pg.ez[il:iu], - pg.bx[il:iu], pg.by[il:iu], pg.bz[il:iu], - pg.sq[js], pg.sm[js], top.dt, 2) - else: - # --- push velocity from magnetic field - warp.bpush3d(np, pg.uxp[il:iu], pg.uyp[il:iu], pg.uzp[il:iu], pg.gaminv[il:iu], - pg.bx[il:iu], pg.by[il:iu], pg.bz[il:iu], - pg.sq[js], pg.sm[js], 0.5*top.dt, top.ibpush) - # --- push velocity from electric field (half step) - warp.epush3d(np, pg.uxp[il:iu], pg.uyp[il:iu], pg.uzp[il:iu], - pg.ex[il:iu], pg.ey[il:iu], pg.ez[il:iu], - pg.sq[js], pg.sm[js], 0.5*top.dt) - # --- update gamma - self.set_gamma(js, pg) - - if self.l_verbose:print(me, 'exit push_velocity_second_part') - - def set_gamma(self, js, pg=None): - if self.l_verbose:print(me, 'enter set_gamma') - if pg is None: - pg = top.pgroup - np = pg.nps[js] - if np == 0: return - il = pg.ins[js] - 1 - iu = il + pg.nps[js] - # --- update gamma - warp.gammaadv(np, pg.gaminv[il:iu], pg.uxp[il:iu], pg.uyp[il:iu], pg.uzp[il:iu], top.gamadv, top.lrelativ) - - if self.l_verbose:print(me, 'exit set_gamma') - - def push_positions(self, js, pg=None): - if self.l_verbose:print(me, 'enter push_positions') - if pg is None: - pg = top.pgroup - np = pg.nps[js] - if np == 0: return - il = pg.ins[js] - 1 - iu = il + pg.nps[js] - warp.xpush3d(np, pg.xp[il:iu], pg.yp[il:iu], pg.zp[il:iu], - pg.uxp[il:iu], pg.uyp[il:iu], pg.uzp[il:iu], - pg.gaminv[il:iu], top.dt) - - if self.l_verbose:print(me, 'exit push_positions') - - - - - - - - + callbacks._afterstep() diff --git a/Python/setup.py b/Python/setup.py index 20e7c5d9f..71f82dbc7 100644 --- a/Python/setup.py +++ b/Python/setup.py @@ -6,9 +6,16 @@ setup.py file for WarpX from distutils.core import setup +try: + from distutils.command.build_py import build_py_2to3 as build_py +except ImportError: + from distutils.command.build_py import build_py + + setup (name = 'pywarpx', packages = ['pywarpx'], package_dir = {'pywarpx':'pywarpx'}, description = """Wrapper of WarpX""", package_data = {'pywarpx' : ['libwarpx.so']}, + cmdclass={'build_py': build_py} ) diff --git a/Source/WarpXElectrostatic.cpp b/Source/WarpXElectrostatic.cpp index 277dc2887..c4e37ce1a 100644 --- a/Source/WarpXElectrostatic.cpp +++ b/Source/WarpXElectrostatic.cpp @@ -63,6 +63,9 @@ WarpX::EvolveES (int numsteps) { // Start loop on time steps amrex::Print() << "\nSTEP " << step+1 << " starts ...\n"; +#ifdef WARPX_USE_PY + if (warpx_py_beforestep) warpx_py_beforestep(); +#endif // At initialization, particles have p^{n-1/2} and x^{n-1/2}. // Beyond one step, particles have p^{n-1/2} and x^{n}. @@ -86,10 +89,22 @@ WarpX::EvolveES (int numsteps) { // Evolve particles to p^{n+1/2} and x^{n+1} mypc->EvolveES(eFieldNodal, rhoNodal, cur_time, dt[lev]); +#ifdef WARPX_USE_PY + if (warpx_py_particleinjection) warpx_py_particleinjection(); + if (warpx_py_particlescraper) warpx_py_particlescraper(); + if (warpx_py_beforedeposition) warpx_py_beforedeposition(); +#endif mypc->DepositCharge(rhoNodal); +#ifdef WARPX_USE_PY + if (warpx_py_afterdeposition) warpx_py_afterdeposition(); + if (warpx_py_beforeEsolve) warpx_py_beforeEsolve(); +#endif computePhi(rhoNodal, phiNodal); computeE(eFieldNodal, phiNodal); +#ifdef WARPX_USE_PY + if (warpx_py_afterEsolve) warpx_py_afterEsolve(); +#endif if (cur_time + dt[0] >= stop_time - 1.e-3*dt[0] || step == numsteps_max-1) { // on last step, push by only 0.5*dt to synchronize all at n+1/2 @@ -135,6 +150,9 @@ WarpX::EvolveES (int numsteps) { break; } +#ifdef WARPX_USE_PY + if (warpx_py_afterstep) warpx_py_afterstep(); +#endif // End loop on time steps } diff --git a/Source/WarpXEvolve.cpp b/Source/WarpXEvolve.cpp index a92f033b5..b42a781d1 100644 --- a/Source/WarpXEvolve.cpp +++ b/Source/WarpXEvolve.cpp @@ -5,7 +5,9 @@ #include <WarpX.H> #include <WarpXConst.H> #include <WarpX_f.H> +#ifdef WARPX_USE_PY #include <WarpX_py.H> +#endif using namespace amrex; @@ -45,12 +47,12 @@ WarpX::EvolveEM (int numsteps) Real walltime, walltime_start = ParallelDescriptor::second(); for (int step = istep[0]; step < numsteps_max && cur_time < stop_time; ++step) { - if (warpx_py_print_step) { - warpx_py_print_step(step); - } // Start loop on time steps amrex::Print() << "\nSTEP " << step+1 << " starts ...\n"; +#ifdef WARPX_USE_PY + if (warpx_py_beforestep) warpx_py_beforestep(); +#endif if (costs[0] != nullptr) { @@ -97,7 +99,15 @@ WarpX::EvolveEM (int numsteps) // from p^{n-1/2} to p^{n+1/2} // Deposit current j^{n+1/2} // Deposit charge density rho^{n} +#ifdef WARPX_USE_PY + if (warpx_py_particleinjection) warpx_py_particleinjection(); + if (warpx_py_particlescraper) warpx_py_particlescraper(); + if (warpx_py_beforedeposition) warpx_py_beforedeposition(); +#endif PushParticlesandDepose(cur_time); +#ifdef WARPX_USE_PY + if (warpx_py_afterdeposition) warpx_py_afterdeposition(); +#endif SyncCurrent(); @@ -122,6 +132,9 @@ WarpX::EvolveEM (int numsteps) FillBoundaryB(); #endif +#ifdef WARPX_USE_PY + if (warpx_py_beforeEsolve) warpx_py_beforeEsolve(); +#endif if (cur_time + dt[0] >= stop_time - 1.e-3*dt[0] || step == numsteps_max-1) { // At the end of last step, push p by 0.5*dt to synchronize UpdateAuxilaryData(); @@ -132,6 +145,9 @@ WarpX::EvolveEM (int numsteps) } is_synchronized = true; } +#ifdef WARPX_USE_PY + if (warpx_py_afterEsolve) warpx_py_afterEsolve(); +#endif for (int lev = 0; lev <= max_level; ++lev) { ++istep[lev]; @@ -196,6 +212,9 @@ WarpX::EvolveEM (int numsteps) break; } +#ifdef WARPX_USE_PY + if (warpx_py_afterstep) warpx_py_afterstep(); +#endif // End loop on time steps } diff --git a/Source/WarpXWrappers.cpp b/Source/WarpXWrappers.cpp index e65c5cae6..87faac93b 100644 --- a/Source/WarpXWrappers.cpp +++ b/Source/WarpXWrappers.cpp @@ -97,6 +97,8 @@ extern "C" { WarpX& warpx = WarpX::GetInstance(); warpx.InitData(); + if (warpx_py_afterinit) warpx_py_afterinit(); + if (warpx_py_particleloader) warpx_py_particleloader(); } void warpx_finalize () @@ -104,9 +106,53 @@ extern "C" WarpX::ResetInstance(); } - void warpx_set_callback_py_funcs (WARPX_CALLBACK_PY_FUNC_1 print_step) + void warpx_set_callback_py_afterinit (WARPX_CALLBACK_PY_FUNC_0 callback) { - warpx_py_print_step = print_step; + warpx_py_afterinit = callback; + } + void warpx_set_callback_py_beforeEsolve (WARPX_CALLBACK_PY_FUNC_0 callback) + { + warpx_py_beforeEsolve = callback; + } + void warpx_set_callback_py_afterEsolve (WARPX_CALLBACK_PY_FUNC_0 callback) + { + warpx_py_afterEsolve = callback; + } + void warpx_set_callback_py_beforedeposition (WARPX_CALLBACK_PY_FUNC_0 callback) + { + warpx_py_beforedeposition = callback; + } + void warpx_set_callback_py_afterdeposition (WARPX_CALLBACK_PY_FUNC_0 callback) + { + warpx_py_afterdeposition = callback; + } + void warpx_set_callback_py_particlescraper (WARPX_CALLBACK_PY_FUNC_0 callback) + { + warpx_py_particlescraper = callback; + } + void warpx_set_callback_py_particleloader (WARPX_CALLBACK_PY_FUNC_0 callback) + { + warpx_py_particleloader = callback; + } + void warpx_set_callback_py_beforestep (WARPX_CALLBACK_PY_FUNC_0 callback) + { + warpx_py_beforestep = callback; + } + void warpx_set_callback_py_afterstep (WARPX_CALLBACK_PY_FUNC_0 callback) + { + warpx_py_afterstep = callback; + } + void warpx_set_callback_py_afterrestart (WARPX_CALLBACK_PY_FUNC_0 callback) + { + warpx_py_afterrestart = callback; + } + void warpx_set_callback_py_particleinjection (WARPX_CALLBACK_PY_FUNC_0 callback) + { + warpx_py_particleinjection = callback; + } + void warpx_set_callback_py_appliedfields (WARPX_CALLBACK_PY_FUNC_0 callback) + { + warpx_py_appliedfields = callback; } void warpx_evolve (int numsteps) @@ -260,11 +306,16 @@ extern "C" auto & myspc = mypc.GetParticleContainer(speciesnumber); const int level = 0; - *num_tiles = myspc.numLocalTilesAtLevel(level); + + int i = 0; + for (WarpXParIter pti(myspc, level); pti.isValid(); ++pti, ++i) {} + + // *num_tiles = myspc.numLocalTilesAtLevel(level); + *num_tiles = i; *particles_per_tile = (int*) malloc(*num_tiles*sizeof(int)); - + double** data = (double**) malloc(*num_tiles*sizeof(typename WarpXParticleContainer::ParticleType*)); - int i = 0; + i = 0; for (WarpXParIter pti(myspc, level); pti.isValid(); ++pti, ++i) { auto& aos = pti.GetArrayOfStructs(); data[i] = (double*) aos.data(); @@ -279,11 +330,16 @@ extern "C" auto & myspc = mypc.GetParticleContainer(speciesnumber); const int level = 0; - *num_tiles = myspc.numLocalTilesAtLevel(level); + + int i = 0; + for (WarpXParIter pti(myspc, level); pti.isValid(); ++pti, ++i) {} + + // *num_tiles = myspc.numLocalTilesAtLevel(level); + *num_tiles = i; *particles_per_tile = (int*) malloc(*num_tiles*sizeof(int)); - + double** data = (double**) malloc(*num_tiles*sizeof(double*)); - int i = 0; + i = 0; for (WarpXParIter pti(myspc, level); pti.isValid(); ++pti, ++i) { auto& soa = pti.GetStructOfArrays(); data[i] = (double*) soa.GetRealData(comp).dataPtr(); diff --git a/Source/WarpXWrappers.h b/Source/WarpXWrappers.h index 7862b0fc3..07a573cc8 100644 --- a/Source/WarpXWrappers.h +++ b/Source/WarpXWrappers.h @@ -31,9 +31,21 @@ extern "C" { void warpx_finalize (); - typedef void(*WARPX_CALLBACK_PY_FUNC_1)(int); - void warpx_set_callback_py_funcs (WARPX_CALLBACK_PY_FUNC_1); - + typedef void(*WARPX_CALLBACK_PY_FUNC_0)(); + + void warpx_set_callback_py_afterinit (WARPX_CALLBACK_PY_FUNC_0); + void warpx_set_callback_py_beforeEsolve (WARPX_CALLBACK_PY_FUNC_0); + void warpx_set_callback_py_afterEsolve (WARPX_CALLBACK_PY_FUNC_0); + void warpx_set_callback_py_beforedeposition (WARPX_CALLBACK_PY_FUNC_0); + void warpx_set_callback_py_afterdeposition (WARPX_CALLBACK_PY_FUNC_0); + void warpx_set_callback_py_particlescraper (WARPX_CALLBACK_PY_FUNC_0); + void warpx_set_callback_py_particleloader (WARPX_CALLBACK_PY_FUNC_0); + void warpx_set_callback_py_beforestep (WARPX_CALLBACK_PY_FUNC_0); + void warpx_set_callback_py_afterstep (WARPX_CALLBACK_PY_FUNC_0); + void warpx_set_callback_py_afterrestart (WARPX_CALLBACK_PY_FUNC_0); + void warpx_set_callback_py_particleinjection (WARPX_CALLBACK_PY_FUNC_0); + void warpx_set_callback_py_appliedfields (WARPX_CALLBACK_PY_FUNC_0); + void warpx_evolve (int numsteps); // -1 means the inputs parameter will be used. void warpx_addNParticles(int speciesnumber, int lenx, diff --git a/Source/WarpX_py.H b/Source/WarpX_py.H index 75977acf2..d8cf22155 100644 --- a/Source/WarpX_py.H +++ b/Source/WarpX_py.H @@ -4,7 +4,20 @@ #include <WarpXWrappers.h> extern "C" { - extern WARPX_CALLBACK_PY_FUNC_1 warpx_py_print_step; + + extern WARPX_CALLBACK_PY_FUNC_0 warpx_py_afterinit; + extern WARPX_CALLBACK_PY_FUNC_0 warpx_py_beforeEsolve; + extern WARPX_CALLBACK_PY_FUNC_0 warpx_py_afterEsolve; + extern WARPX_CALLBACK_PY_FUNC_0 warpx_py_beforedeposition; + extern WARPX_CALLBACK_PY_FUNC_0 warpx_py_afterdeposition; + extern WARPX_CALLBACK_PY_FUNC_0 warpx_py_particlescraper; + extern WARPX_CALLBACK_PY_FUNC_0 warpx_py_particleloader; + extern WARPX_CALLBACK_PY_FUNC_0 warpx_py_beforestep; + extern WARPX_CALLBACK_PY_FUNC_0 warpx_py_afterstep; + extern WARPX_CALLBACK_PY_FUNC_0 warpx_py_afterrestart; + extern WARPX_CALLBACK_PY_FUNC_0 warpx_py_particleinjection; + extern WARPX_CALLBACK_PY_FUNC_0 warpx_py_appliedfields; + } #endif diff --git a/Source/WarpX_py.cpp b/Source/WarpX_py.cpp index 8f7b36e45..276d637d7 100644 --- a/Source/WarpX_py.cpp +++ b/Source/WarpX_py.cpp @@ -1,6 +1,19 @@ #include <WarpX_py.H> extern "C" { - WARPX_CALLBACK_PY_FUNC_1 warpx_py_print_step = nullptr; + + WARPX_CALLBACK_PY_FUNC_0 warpx_py_afterinit = nullptr; + WARPX_CALLBACK_PY_FUNC_0 warpx_py_beforeEsolve = nullptr; + WARPX_CALLBACK_PY_FUNC_0 warpx_py_afterEsolve = nullptr; + WARPX_CALLBACK_PY_FUNC_0 warpx_py_beforedeposition = nullptr; + WARPX_CALLBACK_PY_FUNC_0 warpx_py_afterdeposition = nullptr; + WARPX_CALLBACK_PY_FUNC_0 warpx_py_particlescraper = nullptr; + WARPX_CALLBACK_PY_FUNC_0 warpx_py_particleloader = nullptr; + WARPX_CALLBACK_PY_FUNC_0 warpx_py_beforestep = nullptr; + WARPX_CALLBACK_PY_FUNC_0 warpx_py_afterstep = nullptr; + WARPX_CALLBACK_PY_FUNC_0 warpx_py_afterrestart = nullptr; + WARPX_CALLBACK_PY_FUNC_0 warpx_py_particleinjection = nullptr; + WARPX_CALLBACK_PY_FUNC_0 warpx_py_appliedfields = nullptr; + } diff --git a/tests/Langmuir/langmuir_PICMI.py b/tests/Langmuir/langmuir_PICMI.py deleted file mode 100644 index 976cd90b1..000000000 --- a/tests/Langmuir/langmuir_PICMI.py +++ /dev/null @@ -1,117 +0,0 @@ -import numpy as np -from pywarpx import PICMI -#from warp import PICMI - -from mpi4py import MPI -comm = MPI.COMM_WORLD - -def get_parallel_indices(Np, rank, size): - ''' - - This decomposes the arrays of particle attributes into subarrays - for parallel processing. - - ''' - Navg = Np / size - Nleft = Np - Navg * size - if (rank < Nleft): - lo = rank*(Navg + 1) - hi = lo + Navg + 1 - else: - lo = rank * Navg + Nleft - hi = lo + Navg - return lo, hi - - -def set_initial_conditions(): - ''' - - Sets initial conditions for the Langmuir wave setup. - - - ''' - - comm.Barrier() - - n_e = 1.0e25 - num_particles_per_cell = 10 - ux = 0.01 - uy = 0.0 - uz = 0.0 - - weight = n_e * dx * dy * dz / num_particles_per_cell - c = 299792458.0 - ux *= c - uy *= c - uz *= c - - # --- This creates particles only the left half of the domain - Z, Y, X, P = np.mgrid[0:64, 0:64, 0:32, 0:num_particles_per_cell] - Z = Z.flatten() - Y = Y.flatten() - X = X.flatten() - P = P.flatten() - - particle_shift = (0.5 + P) / num_particles_per_cell - - # -- particle positions - xp = (X + particle_shift)*dx + xmin - yp = (Y + particle_shift)*dy + ymin - zp = (Z + particle_shift)*dz + zmin - - # --- velocities - uxp = ux * np.ones_like(xp) - uyp = uy * np.ones_like(xp) - uzp = uz * np.ones_like(xp) - - # --- and weights - Np = xp.shape[0] - wp = np.full(Np, weight) - - lo, hi = get_parallel_indices(Np, comm.rank, comm.size) - - electrons.add_particles(x=xp[lo:hi], y=yp[lo:hi], z=zp[lo:hi], - ux=uxp[lo:hi], uy=uyp[lo:hi], uz=uzp[lo:hi], - w=wp[lo:hi], unique_particles=1) - - comm.Barrier() - -nx = 64 -ny = 64 -nz = 64 - -xmin = -20.e-6 -ymin = -20.e-6 -zmin = -20.e-6 -xmax = +20.e-6 -ymax = +20.e-6 -zmax = +20.e-6 - -dx = (xmax - xmin)/nx -dy = (ymax - ymin)/ny -dz = (zmax - zmin)/nz - -grid = PICMI.Grid(nx=nx, ny=ny, nz=nz, - xmin=xmin, xmax=xmax, ymin=ymin, ymax=ymax, zmin=zmin, zmax=zmax, - bcxmin='periodic', bcxmax='periodic', bcymin='periodic', bcymax='periodic', bczmin='periodic', bczmax='periodic', - max_grid_size=32, max_level=0, coord_sys=0) - -solver = PICMI.EM_solver(current_deposition = 3, - charge_deposition = 0, - field_gathering = 1, - particle_pusher = 0) - -electrons = PICMI.Species(type=PICMI.Electron, name='electrons', lvariableweights=True) - -sim = PICMI.Simulation(plot_int = 10, - verbose = 1, - cfl = 1.0) - -set_initial_conditions() - -# Maximum number of time steps -max_step = 40 - -sim.step(max_step) - -sim.finalize() diff --git a/tests/plasma_acceleration/plasma_acceleration_PICMI.py b/tests/plasma_acceleration/plasma_acceleration_PICMI.py deleted file mode 100644 index e309773fc..000000000 --- a/tests/plasma_acceleration/plasma_acceleration_PICMI.py +++ /dev/null @@ -1,229 +0,0 @@ -import numpy as np -from pywarpx import PICMI -#from warp import PICMI -if PICMI.codename == 'WarpX': - from pywarpx import PGroups - -from mpi4py import MPI -comm = MPI.COMM_WORLD - -def get_parallel_indices(Np, rank, size): - ''' - - This decomposes the arrays of particle attributes into subarrays - for parallel processing. - - ''' - Navg = Np / size - Nleft = Np - Navg * size - if (rank < Nleft): - lo = rank*(Navg + 1) - hi = lo + Navg + 1 - else: - lo = rank * Navg + Nleft - hi = lo + Navg - return lo, hi - - -def set_initial_conditions(ncells, domain_min, domain_max): - ''' - - Sets initial conditions for the plasma acceleration setup. - - - ''' - - comm.Barrier() - - num_ppc = 4 - c = 299792458.0 - dx = (domain_max - domain_min) / ncells - - # Species 0 - the beam - beam_min = np.array([ -20e-6, -20e-6, -150e-6]) - beam_max = np.array([ 20e-6, 20e-6, -100e-6]) - beam_density = 1.e23 - beam_gamma = 1.e9 - beam_uz = np.sqrt(beam_gamma**2 - 1)*c - beam_weight = beam_density * dx[0] * dx[1] * dx[2] / num_ppc - - # Species 1 - the plasma - plasma_min = np.array([-200e-6, -200e-6, 0.0e-6]) - plasma_max = np.array([ 200e-6, 200e-6, 200e-6]) - plasma_density = 1.e22 - plasma_weight = plasma_density * dx[0] * dx[1] * dx[2] / num_ppc - - # Fill the entire domain with particles. Only a subset of these points - # will be selected for each species - Z, Y, X, P = np.mgrid[0:ncells[0], 0:ncells[1], 0:ncells[2], 0:num_ppc] - Z = Z.flatten() - Y = Y.flatten() - X = X.flatten() - P = P.flatten() - - particle_shift = (0.5 + P) / num_ppc - - xp = (X + particle_shift)*dx[0] + domain_min[0] - yp = (Y + particle_shift)*dx[1] + domain_min[1] - zp = (Z + particle_shift)*dx[2] + domain_min[2] - - # Do the beam species first - beam_locs = np.logical_and(xp >= beam_min[0], xp < beam_max[0]) - beam_locs = np.logical_and(beam_locs, yp >= beam_min[1]) - beam_locs = np.logical_and(beam_locs, zp >= beam_min[2]) - beam_locs = np.logical_and(beam_locs, yp < beam_max[1]) - beam_locs = np.logical_and(beam_locs, zp < beam_max[2]) - - beam_xp = xp[beam_locs] - beam_yp = yp[beam_locs] - beam_zp = zp[beam_locs] - - beam_uxp = np.zeros_like(beam_xp) - beam_uyp = np.zeros_like(beam_xp) - beam_uzp = beam_uz * np.ones_like(beam_xp) - - # --- and weights - Np = beam_xp.shape[0] - beam_wp = np.full(Np, beam_weight) - - lo, hi = get_parallel_indices(Np, comm.rank, comm.size) - - beam.add_particles(x=beam_xp[lo:hi], y=beam_yp[lo:hi], z=beam_zp[lo:hi], - ux=beam_uxp[lo:hi], uy=beam_uyp[lo:hi], uz=beam_uzp[lo:hi], - w=beam_wp[lo:hi], unique_particles=1) - - # now do the plasma species - plasma_locs = np.logical_and(xp >= plasma_min[0], xp < plasma_max[0]) - plasma_locs = np.logical_and(plasma_locs, yp >= plasma_min[1]) - plasma_locs = np.logical_and(plasma_locs, zp >= plasma_min[2]) - plasma_locs = np.logical_and(plasma_locs, yp < plasma_max[1]) - plasma_locs = np.logical_and(plasma_locs, zp < plasma_max[2]) - - plasma_xp = xp[plasma_locs] - plasma_yp = yp[plasma_locs] - plasma_zp = zp[plasma_locs] - - plasma_uxp = np.zeros_like(plasma_xp) - plasma_uyp = np.zeros_like(plasma_xp) - plasma_uzp = np.zeros_like(plasma_xp) - - # --- and weights - Np = plasma_xp.shape[0] - plasma_wp = np.full(Np, plasma_weight) - - lo, hi = get_parallel_indices(Np, comm.rank, comm.size) - - plasma.add_particles(x=plasma_xp[lo:hi], y=plasma_yp[lo:hi], z=plasma_zp[lo:hi], - ux=plasma_uxp[lo:hi], uy=plasma_uyp[lo:hi], uz=plasma_uzp[lo:hi], - w=plasma_wp[lo:hi], unique_particles=1) - - comm.Barrier() - - -def inject_plasma(num_shift, direction): - ''' - - This injects fresh plasma into the domain after the moving window has been upated. - - ''' - - comm.Barrier() - - num_ppc = 4 - density = 1.e22 - weight = density * dx[0]*dx[1]*dx[2] / num_ppc - - shift_box = np.array(ncells) - shift_box[direction] = abs(num_shift) - - Z, Y, X, P = np.mgrid[0:shift_box[2], 0:shift_box[1], 0:shift_box[0], 0:num_ppc] - X = X.flatten() - Y = Y.flatten() - Z = Z.flatten() - P = P.flatten() - particle_shift = (0.5 + P) / num_ppc - - pos = [[],[],[]] - - if (num_shift > 0): - pos[0] = (X + particle_shift)*dx[0] + domain_min[0] - pos[1] = (Y + particle_shift)*dx[1] + domain_min[1] - pos[2] = (Z + particle_shift)*dx[2] + domain_min[2] - pos[direction] -= domain_min[direction] - pos[direction] += domain_max[direction] - - if (num_shift < 0): - pos[0] = (X + particle_shift)*dx[0] + domain_min[0] - pos[1] = (Y + particle_shift)*dx[1] + domain_min[1] - pos[2] = (Z + particle_shift)*dx[2] + domain_min[2] - pos[direction] += num_shift*dx[direction] - - xp = pos[0] - yp = pos[1] - zp = pos[2] - - uxp = np.zeros_like(xp) - uyp = np.zeros_like(xp) - uzp = np.zeros_like(xp) - - Np = xp.shape[0] - wp = np.full(Np, weight) - - lo, hi = get_parallel_indices(Np, comm.rank, comm.size) - - plasma.add_particles(x=xp[lo:hi], y=yp[lo:hi], z=zp[lo:hi], - ux=uxp[lo:hi], uy=uyp[lo:hi], uz=uzp[lo:hi], - w=wp[lo:hi], unique_particles=1) - - comm.Barrier() - -ncells = np.array([64, 64, 64]) -domain_min = np.array([-200e-6, -200e-6, -200e-6]) -domain_max = np.array([ 200e-6, 200e-6, 200e-6]) -dx = (domain_max - domain_min) / ncells - -moving_window_velocity = np.array([0., 0., PICMI.clight]) - -grid = PICMI.Grid(nx=ncells[0], ny=ncells[1], nz=ncells[2], - xmin=domain_min[0], xmax=domain_max[0], ymin=domain_min[1], ymax=domain_max[1], zmin=domain_min[2], zmax=domain_max[2], - bcxmin='periodic', bcxmax='periodic', bcymin='periodic', bcymax='periodic', bczmin='open', bczmax='open', - moving_window_velocity = moving_window_velocity, - max_grid_size=32, max_level=0, coord_sys=0) - - -solver = PICMI.EM_solver(current_deposition = 3, - charge_deposition = 0, - field_gathering = 0, - particle_pusher = 0) - -beam = PICMI.Species(type=PICMI.Electron, name='beam') -plasma = PICMI.Species(type=PICMI.Electron, name='plasma') - - -# Maximum number of time steps -max_step = 160 - -sim = PICMI.Simulation(plot_int = 2, - verbose = 1, - cfl = 1.0) - -set_initial_conditions(ncells, domain_min, domain_max) - -direction = np.argmax(abs(moving_window_velocity)) -old_x = grid.getmins()[direction] -new_x = old_x -for i in range(1, max_step + 1): - - # check whether the moving window has updated - num_shift = int( 0.5 + (new_x - old_x) / dx[direction]) - if (num_shift != 0): - inject_plasma(num_shift, direction) - domain_min[direction] += num_shift*dx[direction] - domain_max[direction] += num_shift*dx[direction] - - sim.step(1) - - old_x = new_x - new_x = grid.getmins()[direction] - -sim.finalize() |