# Copyright 2018-2022 Andrew Myers, David Grote, Ligia Diana Amorim
# Maxence Thevenet, Remi Lehe, Revathi Jambunathan, Lorenzo Giacomel
#
#
# This file is part of WarpX.
#
# License: BSD-3-Clause-LBNL
"""Classes following the PICMI standard
"""
import os
import re
import numpy as np
import periodictable
import picmistandard
import pywarpx
codename = 'warpx'
picmistandard.register_codename(codename)
# dictionary to map field boundary conditions from picmistandard to WarpX
BC_map = {
'open':'pml', 'dirichlet':'pec', 'periodic':'periodic', 'damped':'damped',
'neumann':'neumann', 'none':'none', None:'none'
}
class constants:
# --- Put the constants in their own namespace
# --- Values from WarpXConst.H
c = 299792458.
ep0 = 8.8541878128e-12
mu0 = 1.25663706212e-06
q_e = 1.602176634e-19
m_e = 9.1093837015e-31
m_p = 1.67262192369e-27
hbar = 1.054571817e-34
kb = 1.380649e-23
picmistandard.register_constants(constants)
[docs]class Species(picmistandard.PICMI_Species):
"""
See `Input Parameters <https://warpx.readthedocs.io/en/latest/usage/parameters.html>`_ for more information.
Parameters
----------
warpx_boost_adjust_transverse_positions: bool, default=False
Whether to adjust transverse positions when apply the boost
to the simulation frame
warpx_self_fields_required_precision: float, default=1.e-11
Relative precision on the electrostatic solver
(when using the relativistic solver)
warpx_self_fields_absolute_tolerance: float, default=0.
Absolute precision on the electrostatic solver
(when using the relativistic solver)
warpx_self_fields_max_iters: integer, default=200
Maximum number of iterations for the electrostatic
solver for the species
warpx_self_fields_verbosity: integer, default=2
Level of verbosity for the electrostatic solver
warpx_save_previous_position: bool, default=False
Whether to save the old particle positions
warpx_do_not_deposit: bool, default=False
Whether or not to deposit the charge and current density for
for this species
warpx_reflection_model_xlo: string, default='0.'
Expression (in terms of the velocity "v") specifying the probability
that the particle will reflect on the lower x boundary
warpx_reflection_model_xhi: string, default='0.'
Expression (in terms of the velocity "v") specifying the probability
that the particle will reflect on the upper x boundary
warpx_reflection_model_ylo: string, default='0.'
Expression (in terms of the velocity "v") specifying the probability
that the particle will reflect on the lower y boundary
warpx_reflection_model_yhi: string, default='0.'
Expression (in terms of the velocity "v") specifying the probability
that the particle will reflect on the upper y boundary
warpx_reflection_model_zlo: string, default='0.'
Expression (in terms of the velocity "v") specifying the probability
that the particle will reflect on the lower z boundary
warpx_reflection_model_zhi: string, default='0.'
Expression (in terms of the velocity "v") specifying the probability
that the particle will reflect on the upper z boundary
warpx_save_particles_at_xlo: bool, default=False
Whether to save particles lost at the lower x boundary
warpx_save_particles_at_xhi: bool, default=False
Whether to save particles lost at the upper x boundary
warpx_save_particles_at_ylo: bool, default=False
Whether to save particles lost at the lower y boundary
warpx_save_particles_at_yhi: bool, default=False
Whether to save particles lost at the upper y boundary
warpx_save_particles_at_zlo: bool, default=False
Whether to save particles lost at the lower z boundary
warpx_save_particles_at_zhi: bool, default=False
Whether to save particles lost at the upper z boundary
warpx_save_particles_at_eb: bool, default=False
Whether to save particles lost at the embedded boundary
"""
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'
else:
if self.charge is None and self.charge_state is not None:
if self.charge_state == +1.:
self.charge = 'q_e'
elif self.charge_state == -1.:
self.charge = '-q_e'
else:
self.charge = self.charge_state*constants.q_e
if self.particle_type is not None:
# Match a string of the format '#nXx', with the '#n' optional isotope number.
m = re.match(r'(?P<iso>#[\d+])*(?P<sym>[A-Za-z]+)', self.particle_type)
if m is not None:
element = periodictable.elements.symbol(m['sym'])
if m['iso'] is not None:
element = element[m['iso'][1:]]
if self.charge_state is not None:
assert self.charge_state <= element.number, Exception('%s charge state not valid'%self.particle_type)
try:
element = element.ion[self.charge_state]
except ValueError:
# Note that not all valid charge states are defined in elements,
# so this value error can be ignored.
pass
self.element = element
if self.mass is None:
self.mass = element.mass*periodictable.constants.atomic_mass_constant
else:
raise Exception('The species "particle_type" is not known')
self.boost_adjust_transverse_positions = kw.pop('warpx_boost_adjust_transverse_positions', None)
# For the relativistic electrostatic solver
self.self_fields_required_precision = kw.pop('warpx_self_fields_required_precision', None)
self.self_fields_absolute_tolerance = kw.pop('warpx_self_fields_absolute_tolerance', None)
self.self_fields_max_iters = kw.pop('warpx_self_fields_max_iters', None)
self.self_fields_verbosity = kw.pop('warpx_self_fields_verbosity', None)
self.save_previous_position = kw.pop('warpx_save_previous_position', None)
self.do_not_deposit = kw.pop('warpx_do_not_deposit', None)
# For particle reflection
self.reflection_model_xlo = kw.pop('warpx_reflection_model_xlo', None)
self.reflection_model_xhi = kw.pop('warpx_reflection_model_xhi', None)
self.reflection_model_ylo = kw.pop('warpx_reflection_model_ylo', None)
self.reflection_model_yhi = kw.pop('warpx_reflection_model_yhi', None)
self.reflection_model_zlo = kw.pop('warpx_reflection_model_zlo', None)
self.reflection_model_zhi = kw.pop('warpx_reflection_model_zhi', None)
# self.reflection_model_eb = kw.pop('warpx_reflection_model_eb', None)
# For the scraper buffer
self.save_particles_at_xlo = kw.pop('warpx_save_particles_at_xlo', None)
self.save_particles_at_xhi = kw.pop('warpx_save_particles_at_xhi', None)
self.save_particles_at_ylo = kw.pop('warpx_save_particles_at_ylo', None)
self.save_particles_at_yhi = kw.pop('warpx_save_particles_at_yhi', None)
self.save_particles_at_zlo = kw.pop('warpx_save_particles_at_zlo', None)
self.save_particles_at_zhi = kw.pop('warpx_save_particles_at_zhi', None)
self.save_particles_at_eb = kw.pop('warpx_save_particles_at_eb', None)
def initialize_inputs(self, layout,
initialize_self_fields = False,
injection_plane_position = None,
injection_plane_normal_vector = None):
self.species_number = len(pywarpx.particles.species_names)
if self.name is None:
self.name = 'species{}'.format(self.species_number)
pywarpx.particles.species_names.append(self.name)
if initialize_self_fields is None:
initialize_self_fields = False
self.species = pywarpx.Bucket.Bucket(self.name,
mass = self.mass,
charge = self.charge,
injection_style = None,
initialize_self_fields = int(initialize_self_fields),
boost_adjust_transverse_positions = self.boost_adjust_transverse_positions,
self_fields_required_precision = self.self_fields_required_precision,
self_fields_absolute_tolerance = self.self_fields_absolute_tolerance,
self_fields_max_iters = self.self_fields_max_iters,
self_fields_verbosity = self.self_fields_verbosity,
save_particles_at_xlo = self.save_particles_at_xlo,
save_particles_at_xhi = self.save_particles_at_xhi,
save_particles_at_ylo = self.save_particles_at_ylo,
save_particles_at_yhi = self.save_particles_at_yhi,
save_particles_at_zlo = self.save_particles_at_zlo,
save_particles_at_zhi = self.save_particles_at_zhi,
save_particles_at_eb = self.save_particles_at_eb,
save_previous_position = self.save_previous_position,
do_not_deposit = self.do_not_deposit)
# add reflection models
self.species.add_new_attr("reflection_model_xlo(E)", self.reflection_model_xlo)
self.species.add_new_attr("reflection_model_xhi(E)", self.reflection_model_xhi)
self.species.add_new_attr("reflection_model_ylo(E)", self.reflection_model_ylo)
self.species.add_new_attr("reflection_model_yhi(E)", self.reflection_model_yhi)
self.species.add_new_attr("reflection_model_zlo(E)", self.reflection_model_zlo)
self.species.add_new_attr("reflection_model_zhi(E)", self.reflection_model_zhi)
# self.species.add_new_attr("reflection_model_eb(E)", self.reflection_model_eb)
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, self.density_scale)
if injection_plane_position is not None:
if injection_plane_normal_vector is not None:
assert injection_plane_normal_vector[0] == 0. and injection_plane_normal_vector[1] == 0.,\
Exception('Rigid injection can only be done along z')
pywarpx.particles.rigid_injected_species.append(self.name)
self.species.rigid_advance = 1
self.species.zinject_plane = injection_plane_position
for interaction in self.interactions:
assert interaction[0] == 'ionization'
assert interaction[1] == 'ADK', 'WarpX only has ADK ionization model implemented'
self.species.do_field_ionization=1
self.species.physical_element=self.particle_type
self.species.ionization_product_species = interaction[2].name
self.species.ionization_initial_level = self.charge_state
self.species.charge = 'q_e'
picmistandard.PICMI_MultiSpecies.Species_class = Species
[docs]class MultiSpecies(picmistandard.PICMI_MultiSpecies):
def initialize_inputs(self, layout,
initialize_self_fields = False,
injection_plane_position = None,
injection_plane_normal_vector = None):
for species in self.species_instances_list:
species.initialize_inputs(layout,
initialize_self_fields,
injection_plane_position,
injection_plane_normal_vector)
[docs]class GaussianBunchDistribution(picmistandard.PICMI_GaussianBunchDistribution):
def initialize_inputs(self, species_number, layout, species, density_scale):
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 = constants.q_e
elif charge == '-q_e':
charge = -constants.q_e
species.q_tot = self.n_physical_particles*charge
if density_scale is not None:
species.q_tot *= density_scale
# --- 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
# --- Note that WarpX takes gamma*beta as input
if np.any(np.not_equal(self.velocity_divergence, 0.)):
species.momentum_distribution_type = "radial_expansion"
species.u_over_r = self.velocity_divergence[0]/constants.c
#species.u_over_y = self.velocity_divergence[1]/constants.c
#species.u_over_z = self.velocity_divergence[2]/constants.c
elif np.any(np.not_equal(self.rms_velocity, 0.)):
species.momentum_distribution_type = "gaussian"
species.ux_m = self.centroid_velocity[0]/constants.c
species.uy_m = self.centroid_velocity[1]/constants.c
species.uz_m = self.centroid_velocity[2]/constants.c
species.ux_th = self.rms_velocity[0]/constants.c
species.uy_th = self.rms_velocity[1]/constants.c
species.uz_th = self.rms_velocity[2]/constants.c
else:
species.momentum_distribution_type = "constant"
species.ux = self.centroid_velocity[0]/constants.c
species.uy = self.centroid_velocity[1]/constants.c
species.uz = self.centroid_velocity[2]/constants.c
[docs]class AnalyticDistribution(picmistandard.PICMI_AnalyticDistribution):
def init(self, kw):
self.mangle_dict = None
def initialize_inputs(self, species_number, layout, species, density_scale):
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]
if self.mangle_dict is None:
# Only do this once so that the same variables are used in this distribution
# is used multiple times
self.mangle_dict = pywarpx.my_constants.add_keywords(self.user_defined_kw)
expression = pywarpx.my_constants.mangle_expression(self.density_expression, self.mangle_dict)
species.profile = "parse_density_function"
if density_scale is None:
species.__setattr__('density_function(x,y,z)', expression)
else:
species.__setattr__('density_function(x,y,z)', "{}*({})".format(density_scale, expression))
# --- Note that WarpX takes gamma*beta as input
if np.any(np.not_equal(self.momentum_expressions, None)):
species.momentum_distribution_type = 'parse_momentum_function'
self.setup_parse_momentum_functions(species)
elif np.any(np.not_equal(self.rms_velocity, 0.)):
species.momentum_distribution_type = "gaussian"
species.ux_m = self.directed_velocity[0]/constants.c
species.uy_m = self.directed_velocity[1]/constants.c
species.uz_m = self.directed_velocity[2]/constants.c
species.ux_th = self.rms_velocity[0]/constants.c
species.uy_th = self.rms_velocity[1]/constants.c
species.uz_th = self.rms_velocity[2]/constants.c
else:
species.momentum_distribution_type = "constant"
species.ux = self.directed_velocity[0]/constants.c
species.uy = self.directed_velocity[1]/constants.c
species.uz = self.directed_velocity[2]/constants.c
if self.fill_in:
species.do_continuous_injection = 1
def setup_parse_momentum_functions(self, species):
for sdir, idir in zip(['x', 'y', 'z'], [0, 1, 2]):
if self.momentum_expressions[idir] is not None:
expression = pywarpx.my_constants.mangle_expression(self.momentum_expressions[idir], self.mangle_dict)
else:
expression = f'{self.directed_velocity[idir]}'
species.__setattr__(f'momentum_function_u{sdir}(x,y,z)', f'({expression})/{constants.c}')
[docs]class ParticleListDistribution(picmistandard.PICMI_ParticleListDistribution):
def init(self, kw):
pass
def initialize_inputs(self, species_number, layout, species, density_scale):
species.injection_style = "multipleparticles"
species.multiple_particles_pos_x = self.x
species.multiple_particles_pos_y = self.y
species.multiple_particles_pos_z = self.z
species.multiple_particles_vel_x = np.array(self.ux)/constants.c
species.multiple_particles_vel_y = np.array(self.uy)/constants.c
species.multiple_particles_vel_z = np.array(self.uz)/constants.c
species.multiple_particles_weight = self.weight
if density_scale is not None:
species.multiple_particles_weight = self.weight*density_scale
class ParticleDistributionPlanarInjector(picmistandard.PICMI_ParticleDistributionPlanarInjector):
pass
[docs]class GriddedLayout(picmistandard.PICMI_GriddedLayout):
pass
[docs]class PseudoRandomLayout(picmistandard.PICMI_PseudoRandomLayout):
def init(self, kw):
if self.seed is not None:
print('Warning: WarpX does not support specifying the random number seed in PseudoRandomLayout')
class BinomialSmoother(picmistandard.PICMI_BinomialSmoother):
def initialize_inputs(self, solver):
pywarpx.warpx.use_filter = 1
pywarpx.warpx.use_filter_compensation = bool(np.all(self.compensation))
if self.n_pass is None:
# If not specified, do at least one pass in each direction.
self.n_pass = 1
try:
# Check if n_pass is a vector
len(self.n_pass)
except TypeError:
# If not, make it a vector
self.n_pass = solver.grid.number_of_dimensions*[self.n_pass]
pywarpx.warpx.filter_npass_each_dir = self.n_pass
[docs]class CylindricalGrid(picmistandard.PICMI_CylindricalGrid):
"""
This assumes that WarpX was compiled with USE_RZ = TRUE
See `Input Parameters <https://warpx.readthedocs.io/en/latest/usage/parameters.html>`_ for more information.
Parameters
---------
warpx_max_grid_size: integer, default=32
Maximum block size in either direction
warpx_max_grid_size_x: integer, optional
Maximum block size in radial direction
warpx_max_grid_size_y: integer, optional
Maximum block size in longitudinal direction
warpx_blocking_factor: integer, optional
Blocking factor (which controls the block size)
warpx_blocking_factor_x: integer, optional
Blocking factor (which controls the block size) in the radial direction
warpx_blocking_factor_y: integer, optional
Blocking factor (which controls the block size) in the longitudinal direction
warpx_potential_lo_r: float, default=0.
Electrostatic potential on the lower radial boundary
warpx_potential_hi_r: float, default=0.
Electrostatic potential on the upper radial boundary
warpx_potential_lo_z: float, default=0.
Electrostatic potential on the lower longitudinal boundary
warpx_potential_hi_z: float, default=0.
Electrostatic potential on the upper longitudinal boundary
"""
def init(self, kw):
self.max_grid_size = kw.pop('warpx_max_grid_size', 32)
self.max_grid_size_x = kw.pop('warpx_max_grid_size_x', None)
self.max_grid_size_y = kw.pop('warpx_max_grid_size_y', None)
self.blocking_factor = kw.pop('warpx_blocking_factor', None)
self.blocking_factor_x = kw.pop('warpx_blocking_factor_x', None)
self.blocking_factor_y = kw.pop('warpx_blocking_factor_y', None)
self.potential_xmin = kw.pop('warpx_potential_lo_r', None)
self.potential_xmax = kw.pop('warpx_potential_hi_r', None)
self.potential_ymin = None
self.potential_ymax = None
self.potential_zmin = kw.pop('warpx_potential_lo_z', None)
self.potential_zmax = kw.pop('warpx_potential_hi_z', None)
# Geometry
# Set these as soon as the information is available
# (since these are needed to determine which shared object to load)
pywarpx.geometry.dims = 'RZ'
pywarpx.geometry.prob_lo = self.lower_bound # physical domain
pywarpx.geometry.prob_hi = self.upper_bound
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
pywarpx.amr.max_grid_size_x = self.max_grid_size_x
pywarpx.amr.max_grid_size_y = self.max_grid_size_y
pywarpx.amr.blocking_factor = self.blocking_factor
pywarpx.amr.blocking_factor_x = self.blocking_factor_x
pywarpx.amr.blocking_factor_y = self.blocking_factor_y
assert self.lower_bound[0] >= 0., Exception('Lower radial boundary must be >= 0.')
assert self.lower_boundary_conditions[0] != 'periodic' and self.upper_boundary_conditions[0] != 'periodic', Exception('Radial boundaries can not be periodic')
pywarpx.warpx.n_rz_azimuthal_modes = self.n_azimuthal_modes
# Boundary conditions
pywarpx.boundary.field_lo = [BC_map[bc] for bc in self.lower_boundary_conditions]
pywarpx.boundary.field_hi = [BC_map[bc] for bc in self.upper_boundary_conditions]
pywarpx.boundary.particle_lo = self.lower_boundary_conditions_particles
pywarpx.boundary.particle_hi = self.upper_boundary_conditions_particles
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 = 'r'
pywarpx.warpx.moving_window_v = self.moving_window_velocity[0]/constants.c # in units of the speed of light
if self.moving_window_velocity[1] != 0.:
pywarpx.warpx.moving_window_dir = 'z'
pywarpx.warpx.moving_window_v = self.moving_window_velocity[1]/constants.c # in units of the speed of light
if self.refined_regions:
assert len(self.refined_regions) == 1, Exception('WarpX only supports one refined region.')
assert self.refined_regions[0][0] == 1, Exception('The one refined region can only be level 1')
pywarpx.amr.max_level = 1
pywarpx.warpx.fine_tag_lo = self.refined_regions[0][1]
pywarpx.warpx.fine_tag_hi = self.refined_regions[0][2]
# The refinement_factor is ignored (assumed to be [2,2])
else:
pywarpx.amr.max_level = 0
[docs]class Cartesian1DGrid(picmistandard.PICMI_Cartesian1DGrid):
"""
See `Input Parameters <https://warpx.readthedocs.io/en/latest/usage/parameters.html>`_ for more information.
Parameters
---------
warpx_max_grid_size: integer, default=32
Maximum block size in either direction
warpx_max_grid_size_x: integer, optional
Maximum block size in longitudinal direction
warpx_blocking_factor: integer, optional
Blocking factor (which controls the block size)
warpx_blocking_factor_x: integer, optional
Blocking factor (which controls the block size) in the longitudinal direction
warpx_potential_lo_z: float, default=0.
Electrostatic potential on the lower longitudinal boundary
warpx_potential_hi_z: float, default=0.
Electrostatic potential on the upper longitudinal boundary
"""
def init(self, kw):
self.max_grid_size = kw.pop('warpx_max_grid_size', 32)
self.max_grid_size_x = kw.pop('warpx_max_grid_size_x', None)
self.blocking_factor = kw.pop('warpx_blocking_factor', None)
self.blocking_factor_x = kw.pop('warpx_blocking_factor_x', None)
self.potential_xmin = None
self.potential_xmax = None
self.potential_ymin = None
self.potential_ymax = None
self.potential_zmin = kw.pop('warpx_potential_lo_z', None)
self.potential_zmax = kw.pop('warpx_potential_hi_z', None)
# Geometry
# Set these as soon as the information is available
# (since these are needed to determine which shared object to load)
pywarpx.geometry.dims = '1'
pywarpx.geometry.prob_lo = self.lower_bound # physical domain
pywarpx.geometry.prob_hi = self.upper_bound
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
pywarpx.amr.max_grid_size_x = self.max_grid_size_x
pywarpx.amr.blocking_factor = self.blocking_factor
pywarpx.amr.blocking_factor_x = self.blocking_factor_x
# Boundary conditions
pywarpx.boundary.field_lo = [BC_map[bc] for bc in self.lower_boundary_conditions]
pywarpx.boundary.field_hi = [BC_map[bc] for bc in self.upper_boundary_conditions]
pywarpx.boundary.particle_lo = self.lower_boundary_conditions_particles
pywarpx.boundary.particle_hi = self.upper_boundary_conditions_particles
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 = 'z'
pywarpx.warpx.moving_window_v = self.moving_window_velocity[0]/constants.c # in units of the speed of light
if self.refined_regions:
assert len(self.refined_regions) == 1, Exception('WarpX only supports one refined region.')
assert self.refined_regions[0][0] == 1, Exception('The one refined region can only be level 1')
pywarpx.amr.max_level = 1
pywarpx.warpx.fine_tag_lo = self.refined_regions[0][1]
pywarpx.warpx.fine_tag_hi = self.refined_regions[0][2]
# The refinement_factor is ignored (assumed to be [2,2])
else:
pywarpx.amr.max_level = 0
[docs]class Cartesian2DGrid(picmistandard.PICMI_Cartesian2DGrid):
"""
See `Input Parameters <https://warpx.readthedocs.io/en/latest/usage/parameters.html>`_ for more information.
Parameters
---------
warpx_max_grid_size: integer, default=32
Maximum block size in either direction
warpx_max_grid_size_x: integer, optional
Maximum block size in x direction
warpx_max_grid_size_y: integer, optional
Maximum block size in z direction
warpx_blocking_factor: integer, optional
Blocking factor (which controls the block size)
warpx_blocking_factor_x: integer, optional
Blocking factor (which controls the block size) in the x direction
warpx_blocking_factor_y: integer, optional
Blocking factor (which controls the block size) in the z direction
warpx_potential_lo_x: float, default=0.
Electrostatic potential on the lower x boundary
warpx_potential_hi_x: float, default=0.
Electrostatic potential on the upper x boundary
warpx_potential_lo_z: float, default=0.
Electrostatic potential on the lower z boundary
warpx_potential_hi_z: float, default=0.
Electrostatic potential on the upper z boundary
"""
def init(self, kw):
self.max_grid_size = kw.pop('warpx_max_grid_size', 32)
self.max_grid_size_x = kw.pop('warpx_max_grid_size_x', None)
self.max_grid_size_y = kw.pop('warpx_max_grid_size_y', None)
self.blocking_factor = kw.pop('warpx_blocking_factor', None)
self.blocking_factor_x = kw.pop('warpx_blocking_factor_x', None)
self.blocking_factor_y = kw.pop('warpx_blocking_factor_y', None)
self.potential_xmin = kw.pop('warpx_potential_lo_x', None)
self.potential_xmax = kw.pop('warpx_potential_hi_x', None)
self.potential_ymin = None
self.potential_ymax = None
self.potential_zmin = kw.pop('warpx_potential_lo_z', None)
self.potential_zmax = kw.pop('warpx_potential_hi_z', None)
# Geometry
# Set these as soon as the information is available
# (since these are needed to determine which shared object to load)
pywarpx.geometry.dims = '2'
pywarpx.geometry.prob_lo = self.lower_bound # physical domain
pywarpx.geometry.prob_hi = self.upper_bound
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
pywarpx.amr.max_grid_size_x = self.max_grid_size_x
pywarpx.amr.max_grid_size_y = self.max_grid_size_y
pywarpx.amr.blocking_factor = self.blocking_factor
pywarpx.amr.blocking_factor_x = self.blocking_factor_x
pywarpx.amr.blocking_factor_y = self.blocking_factor_y
# Boundary conditions
pywarpx.boundary.field_lo = [BC_map[bc] for bc in self.lower_boundary_conditions]
pywarpx.boundary.field_hi = [BC_map[bc] for bc in self.upper_boundary_conditions]
pywarpx.boundary.particle_lo = self.lower_boundary_conditions_particles
pywarpx.boundary.particle_hi = self.upper_boundary_conditions_particles
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]/constants.c # in units of the speed of light
if self.moving_window_velocity[1] != 0.:
pywarpx.warpx.moving_window_dir = 'z'
pywarpx.warpx.moving_window_v = self.moving_window_velocity[1]/constants.c # in units of the speed of light
if self.refined_regions:
assert len(self.refined_regions) == 1, Exception('WarpX only supports one refined region.')
assert self.refined_regions[0][0] == 1, Exception('The one refined region can only be level 1')
pywarpx.amr.max_level = 1
pywarpx.warpx.fine_tag_lo = self.refined_regions[0][1]
pywarpx.warpx.fine_tag_hi = self.refined_regions[0][2]
# The refinement_factor is ignored (assumed to be [2,2])
else:
pywarpx.amr.max_level = 0
[docs]class Cartesian3DGrid(picmistandard.PICMI_Cartesian3DGrid):
"""
See `Input Parameters <https://warpx.readthedocs.io/en/latest/usage/parameters.html>`_ for more information.
Parameters
---------
warpx_max_grid_size: integer, default=32
Maximum block size in either direction
warpx_max_grid_size_x: integer, optional
Maximum block size in x direction
warpx_max_grid_size_y: integer, optional
Maximum block size in z direction
warpx_max_grid_size_z: integer, optional
Maximum block size in z direction
warpx_blocking_factor: integer, optional
Blocking factor (which controls the block size)
warpx_blocking_factor_x: integer, optional
Blocking factor (which controls the block size) in the x direction
warpx_blocking_factor_y: integer, optional
Blocking factor (which controls the block size) in the z direction
warpx_blocking_factor_z: integer, optional
Blocking factor (which controls the block size) in the z direction
warpx_potential_lo_x: float, default=0.
Electrostatic potential on the lower x boundary
warpx_potential_hi_x: float, default=0.
Electrostatic potential on the upper x boundary
warpx_potential_lo_y: float, default=0.
Electrostatic potential on the lower z boundary
warpx_potential_hi_y: float, default=0.
Electrostatic potential on the upper z boundary
warpx_potential_lo_z: float, default=0.
Electrostatic potential on the lower z boundary
warpx_potential_hi_z: float, default=0.
Electrostatic potential on the upper z boundary
"""
def init(self, kw):
self.max_grid_size = kw.pop('warpx_max_grid_size', 32)
self.max_grid_size_x = kw.pop('warpx_max_grid_size_x', None)
self.max_grid_size_y = kw.pop('warpx_max_grid_size_y', None)
self.max_grid_size_z = kw.pop('warpx_max_grid_size_z', None)
self.blocking_factor = kw.pop('warpx_blocking_factor', None)
self.blocking_factor_x = kw.pop('warpx_blocking_factor_x', None)
self.blocking_factor_y = kw.pop('warpx_blocking_factor_y', None)
self.blocking_factor_z = kw.pop('warpx_blocking_factor_z', None)
self.potential_xmin = kw.pop('warpx_potential_lo_x', None)
self.potential_xmax = kw.pop('warpx_potential_hi_x', None)
self.potential_ymin = kw.pop('warpx_potential_lo_y', None)
self.potential_ymax = kw.pop('warpx_potential_hi_y', None)
self.potential_zmin = kw.pop('warpx_potential_lo_z', None)
self.potential_zmax = kw.pop('warpx_potential_hi_z', None)
# Geometry
# Set these as soon as the information is available
# (since these are needed to determine which shared object to load)
pywarpx.geometry.dims = '3'
pywarpx.geometry.prob_lo = self.lower_bound # physical domain
pywarpx.geometry.prob_hi = self.upper_bound
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
pywarpx.amr.max_grid_size_x = self.max_grid_size_x
pywarpx.amr.max_grid_size_y = self.max_grid_size_y
pywarpx.amr.max_grid_size_z = self.max_grid_size_z
pywarpx.amr.blocking_factor = self.blocking_factor
pywarpx.amr.blocking_factor_x = self.blocking_factor_x
pywarpx.amr.blocking_factor_y = self.blocking_factor_y
pywarpx.amr.blocking_factor_z = self.blocking_factor_z
# Boundary conditions
pywarpx.boundary.field_lo = [BC_map[bc] for bc in self.lower_boundary_conditions]
pywarpx.boundary.field_hi = [BC_map[bc] for bc in self.upper_boundary_conditions]
pywarpx.boundary.particle_lo = self.lower_boundary_conditions_particles
pywarpx.boundary.particle_hi = self.upper_boundary_conditions_particles
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]/constants.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]/constants.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]/constants.c # in units of the speed of light
if self.refined_regions:
assert len(self.refined_regions) == 1, Exception('WarpX only supports one refined region.')
assert self.refined_regions[0][0] == 1, Exception('The one refined region can only be level 1')
pywarpx.amr.max_level = 1
pywarpx.warpx.fine_tag_lo = self.refined_regions[0][1]
pywarpx.warpx.fine_tag_hi = self.refined_regions[0][2]
# The refinement_factor is ignored (assumed to be [2,2,2])
else:
pywarpx.amr.max_level = 0
[docs]class ElectromagneticSolver(picmistandard.PICMI_ElectromagneticSolver):
"""
See `Input Parameters <https://warpx.readthedocs.io/en/latest/usage/parameters.html>`_ for more information.
Parameters
----------
warpx_pml_ncell: integer, optional
The depth of the PML, in number of cells
warpx_periodic_single_box_fft: bool, default=False
Whether to do the spectral solver FFTs assuming a single
simulation block
warpx_current_correction: bool, default=True
Whether to do the current correction for the spectral solver.
See documentation for exceptions to the default value.
warpx_psatd_update_with_rho: bool, optional
Whether to update with the actual rho for the spectral solver
warpx_psatd_do_time_averaging: bool, optional
Whether to do the time averaging for the spectral solver
warpx_do_pml_in_domain: bool, default=False
Whether to do the PML boundaries within the domain (versus
in the guard cells)
warpx_pml_has_particles: bool, default=False
Whether to allow particles in the PML region
warpx_do_pml_j_damping: bool, default=False
Whether to do damping of J in the PML
"""
def init(self, kw):
assert self.method is None or self.method in ['Yee', 'CKC', 'PSATD', 'ECT'], Exception("Only 'Yee', 'CKC', 'PSATD', and 'ECT' are supported")
self.pml_ncell = kw.pop('warpx_pml_ncell', None)
if self.method == 'PSATD':
self.psatd_periodic_single_box_fft = kw.pop('warpx_periodic_single_box_fft', None)
self.psatd_current_correction = kw.pop('warpx_current_correction', None)
self.psatd_update_with_rho = kw.pop('warpx_psatd_update_with_rho', None)
self.psatd_do_time_averaging = kw.pop('warpx_psatd_do_time_averaging', None)
self.do_pml_in_domain = kw.pop('warpx_do_pml_in_domain', None)
self.pml_has_particles = kw.pop('warpx_pml_has_particles', None)
self.do_pml_j_damping = kw.pop('warpx_do_pml_j_damping', None)
def initialize_inputs(self):
self.grid.initialize_inputs()
pywarpx.warpx.pml_ncell = self.pml_ncell
pywarpx.warpx.do_nodal = self.l_nodal
if self.method == 'PSATD':
pywarpx.psatd.periodic_single_box_fft = self.psatd_periodic_single_box_fft
pywarpx.psatd.current_correction = self.psatd_current_correction
pywarpx.psatd.update_with_rho = self.psatd_update_with_rho
pywarpx.psatd.do_time_averaging = self.psatd_do_time_averaging
if self.grid.guard_cells is not None:
pywarpx.psatd.nx_guard = self.grid.guard_cells[0]
if self.grid.number_of_dimensions == 3:
pywarpx.psatd.ny_guard = self.grid.guard_cells[1]
pywarpx.psatd.nz_guard = self.grid.guard_cells[-1]
if self.stencil_order is not None:
pywarpx.psatd.nox = self.stencil_order[0]
if self.grid.number_of_dimensions == 3:
pywarpx.psatd.noy = self.stencil_order[1]
pywarpx.psatd.noz = self.stencil_order[-1]
if self.galilean_velocity is not None:
if self.grid.number_of_dimensions == 2:
self.galilean_velocity = [self.galilean_velocity[0], 0., self.galilean_velocity[1]]
pywarpx.psatd.v_galilean = np.array(self.galilean_velocity)/constants.c
# --- Same method names are used, though mapped to lower case.
pywarpx.algo.maxwell_solver = self.method
if self.cfl is not None:
pywarpx.warpx.cfl = self.cfl
if self.source_smoother is not None:
self.source_smoother.initialize_inputs(self)
pywarpx.warpx.do_dive_cleaning = self.divE_cleaning
pywarpx.warpx.do_divb_cleaning = self.divB_cleaning
pywarpx.warpx.do_pml_dive_cleaning = self.pml_divE_cleaning
pywarpx.warpx.do_pml_divb_cleaning = self.pml_divB_cleaning
pywarpx.warpx.do_pml_in_domain = self.do_pml_in_domain
pywarpx.warpx.pml_has_particles = self.pml_has_particles
pywarpx.warpx.do_pml_j_damping = self.do_pml_j_damping
[docs]class ElectrostaticSolver(picmistandard.PICMI_ElectrostaticSolver):
"""
See `Input Parameters <https://warpx.readthedocs.io/en/latest/usage/parameters.html>`_ for more information.
Parameters
----------
warpx_relativistic: bool, default=False
Whether to use the relativistic solver or lab frame solver
warpx_absolute_tolerance: float, default=0.
Absolute tolerance on the lab fram solver
warpx_self_fields_verbosity: integer, default=2
Level of verbosity for the lab frame solver
"""
def init(self, kw):
self.relativistic = kw.pop('warpx_relativistic', False)
self.absolute_tolerance = kw.pop('warpx_absolute_tolerance', None)
self.self_fields_verbosity = kw.pop('warpx_self_fields_verbosity', None)
def initialize_inputs(self):
self.grid.initialize_inputs()
if self.relativistic:
pywarpx.warpx.do_electrostatic = 'relativistic'
else:
pywarpx.warpx.do_electrostatic = 'labframe'
pywarpx.warpx.self_fields_required_precision = self.required_precision
pywarpx.warpx.self_fields_absolute_tolerance = self.absolute_tolerance
pywarpx.warpx.self_fields_max_iters = self.maximum_iterations
pywarpx.warpx.self_fields_verbosity = self.self_fields_verbosity
pywarpx.boundary.potential_lo_x = self.grid.potential_xmin
pywarpx.boundary.potential_lo_y = self.grid.potential_ymin
pywarpx.boundary.potential_lo_z = self.grid.potential_zmin
pywarpx.boundary.potential_hi_x = self.grid.potential_xmax
pywarpx.boundary.potential_hi_y = self.grid.potential_ymax
pywarpx.boundary.potential_hi_z = self.grid.potential_zmax
[docs]class GaussianLaser(picmistandard.PICMI_GaussianLaser):
def initialize_inputs(self):
self.laser_number = len(pywarpx.lasers.names) + 1
if self.name is None:
self.name = 'laser{}'.format(self.laser_number)
self.laser = pywarpx.Lasers.newlaser(self.name)
self.laser.profile = "Gaussian"
self.laser.wavelength = self.wavelength # The wavelength of the laser (in meters)
self.laser.e_max = self.E0 # Maximum amplitude of the laser field (in V/m)
self.laser.polarization = self.polarization_direction # The main polarization vector
self.laser.profile_waist = self.waist # The waist of the laser (in meters)
self.laser.profile_duration = self.duration # The duration of the laser (in seconds)
self.laser.direction = self.propagation_direction
self.laser.zeta = self.zeta
self.laser.beta = self.beta
self.laser.phi2 = self.phi2
self.laser.phi0 = self.phi0
self.laser.do_continuous_injection = self.fill_in
[docs]class AnalyticLaser(picmistandard.PICMI_AnalyticLaser):
def init(self, kw):
self.mangle_dict = None
def initialize_inputs(self):
self.laser_number = len(pywarpx.lasers.names) + 1
if self.name is None:
self.name = 'laser{}'.format(self.laser_number)
self.laser = pywarpx.Lasers.newlaser(self.name)
self.laser.profile = "parse_field_function"
self.laser.wavelength = self.wavelength # The wavelength of the laser (in meters)
self.laser.e_max = self.Emax # Maximum amplitude of the laser field (in V/m)
self.laser.polarization = self.polarization_direction # The main polarization vector
self.laser.direction = self.propagation_direction
self.laser.do_continuous_injection = self.fill_in
if self.mangle_dict is None:
# Only do this once so that the same variables are used in this distribution
# is used multiple times
self.mangle_dict = pywarpx.my_constants.add_keywords(self.user_defined_kw)
expression = pywarpx.my_constants.mangle_expression(self.field_expression, self.mangle_dict)
self.laser.__setattr__('field_function(X,Y,t)', expression)
[docs]class LaserAntenna(picmistandard.PICMI_LaserAntenna):
def initialize_inputs(self, laser):
laser.laser.position = self.position # This point is on the laser plane
laser.laser.direction = self.normal_vector # The plane normal direction
if isinstance(laser, GaussianLaser):
laser.laser.profile_focal_distance = laser.focal_position[2] - self.position[2] # Focal distance from the antenna (in meters)
laser.laser.profile_t_peak = (self.position[2] - laser.centroid_position[2])/constants.c # The time at which the laser reaches its peak (in seconds)
[docs]class ConstantAppliedField(picmistandard.PICMI_ConstantAppliedField):
def initialize_inputs(self):
# Note that lower and upper_bound are not used by WarpX
if (self.Ex is not None or
self.Ey is not None or
self.Ez is not None):
pywarpx.particles.E_ext_particle_init_style = 'constant'
pywarpx.particles.E_external_particle = [self.Ex or 0., self.Ey or 0., self.Ez or 0.]
if (self.Bx is not None or
self.By is not None or
self.Bz is not None):
pywarpx.particles.B_ext_particle_init_style = 'constant'
pywarpx.particles.B_external_particle = [self.Bx or 0., self.By or 0., self.Bz or 0.]
[docs]class AnalyticAppliedField(picmistandard.PICMI_AnalyticAppliedField):
def init(self, kw):
self.mangle_dict = None
def initialize_inputs(self):
# Note that lower and upper_bound are not used by WarpX
if self.mangle_dict is None:
# Only do this once so that the same variables are used in this distribution
# is used multiple times
self.mangle_dict = pywarpx.my_constants.add_keywords(self.user_defined_kw)
if (self.Ex_expression is not None or
self.Ey_expression is not None or
self.Ez_expression is not None):
pywarpx.particles.E_ext_particle_init_style = 'parse_e_ext_particle_function'
for sdir, expression in zip(['x', 'y', 'z'], [self.Ex_expression, self.Ey_expression, self.Ez_expression]):
expression = pywarpx.my_constants.mangle_expression(expression, self.mangle_dict)
pywarpx.particles.__setattr__(f'E{sdir}_external_particle_function(x,y,z,t)', expression)
if (self.Bx_expression is not None or
self.By_expression is not None or
self.Bz_expression is not None):
pywarpx.particles.B_ext_particle_init_style = 'parse_b_ext_particle_function'
for sdir, expression in zip(['x', 'y', 'z'], [self.Bx_expression, self.By_expression, self.Bz_expression]):
expression = pywarpx.my_constants.mangle_expression(expression, self.mangle_dict)
pywarpx.particles.__setattr__(f'B{sdir}_external_particle_function(x,y,z,t)', expression)
[docs]class Mirror(picmistandard.PICMI_Mirror):
def initialize_inputs(self):
try:
pywarpx.warpx.num_mirrors
except AttributeError:
pywarpx.warpx.num_mirrors = 0
pywarpx.warpx.mirror_z = []
pywarpx.warpx.mirror_z_width = []
pywarpx.warpx.mirror_z_npoints = []
pywarpx.warpx.num_mirrors += 1
pywarpx.warpx.mirror_z.append(self.z_front_location)
pywarpx.warpx.mirror_z_width.append(self.depth)
pywarpx.warpx.mirror_z_npoints.append(self.number_of_cells)
[docs]class CoulombCollisions(picmistandard.base._ClassWithInit):
"""
Custom class to handle setup of binary Coulmb collisions in WarpX. If
collision initialization is added to picmistandard this can be changed to
inherit that functionality.
Parameters
----------
name: string
Name of instance (used in the inputs file)
species: list of species instances
The species involved in the collision. Must be of length 2.
CoulombLog: float, optional
Value of the Coulomb log to use in the collision cross section.
If not supplied, it is calculated from the local conditions.
ndt: integer, optional
The collisions will be applied every "ndt" steps. Must be 1 or larger.
"""
def __init__(self, name, species, CoulombLog=None, ndt=None, **kw):
self.name = name
self.species = species
self.CoulombLog = CoulombLog
self.ndt = ndt
self.handle_init(kw)
def initialize_inputs(self):
collision = pywarpx.Collisions.newcollision(self.name)
collision.type = 'pairwisecoulomb'
collision.species = [species.name for species in self.species]
collision.CoulombLog = self.CoulombLog
collision.ndt = self.ndt
[docs]class MCCCollisions(picmistandard.base._ClassWithInit):
"""
Custom class to handle setup of MCC collisions in WarpX. If collision
initialization is added to picmistandard this can be changed to inherit
that functionality.
Parameters
----------
name: string
Name of instance (used in the inputs file)
species: species instance
The species involved in the collision
background_density: float
The density of the background
background_temperature: float
The temperature of the background
scattering_processes: dictionary
The scattering process to use and any needed information
background_mass: float, optional
The mass of the background particle. If not supplied, the default depends
on the type of scattering process.
ndt: integer, optional
The collisions will be applied every "ndt" steps. Must be 1 or larger.
"""
def __init__(self, name, species, background_density,
background_temperature, scattering_processes,
background_mass=None, ndt=None, **kw):
self.name = name
self.species = species
self.background_density = background_density
self.background_temperature = background_temperature
self.background_mass = background_mass
self.scattering_processes = scattering_processes
self.ndt = ndt
self.handle_init(kw)
def initialize_inputs(self):
collision = pywarpx.Collisions.newcollision(self.name)
collision.type = 'background_mcc'
collision.species = self.species.name
collision.background_density = self.background_density
collision.background_temperature = self.background_temperature
collision.background_mass = self.background_mass
collision.ndt = self.ndt
collision.scattering_processes = self.scattering_processes.keys()
for process, kw in self.scattering_processes.items():
for key, val in kw.items():
if key == 'species':
val = val.name
collision.add_new_attr(process+'_'+key, val)
[docs]class EmbeddedBoundary(picmistandard.base._ClassWithInit):
"""
Custom class to handle set up of embedded boundaries specific to WarpX.
If embedded boundary initialization is added to picmistandard this can be
changed to inherit that functionality. The geometry can be specified either as
an implicit function or as an STL file (ASCII or binary). In the latter case the
geometry specified in the STL file can be scaled, translated and inverted.
Parameters
----------
implicit_function: string
Analytic expression describing the embedded boundary
stl_file: string
STL file path (string), file contains the embedded boundary geometry
stl_scale: float
Factor by which the STL geometry is scaled
stl_center: vector of floats
Vector by which the STL geometry is translated (in meters)
stl_reverse_normal: bool
If True inverts the orientation of the STL geometry
potential: string, default=0.
Analytic expression defining the potential. Can only be specified
when the solver is electrostatic.
Parameters used in the analytic expressions should be given as additional keyword arguments.
"""
def __init__(self, implicit_function=None, stl_file=None, stl_scale=None, stl_center=None, stl_reverse_normal=False,
potential=None, **kw):
assert stl_file is None or implicit_function is None, Exception('Only one between implicit_function and '
'stl_file can be specified')
self.implicit_function = implicit_function
self.stl_file = stl_file
if stl_file is None:
assert stl_scale is None, Exception('EB can only be scaled only when using an stl file')
assert stl_center is None, Exception('EB can only be translated only when using an stl file')
assert stl_reverse_normal is False, Exception('EB can only be reversed only when using an stl file')
self.stl_scale = stl_scale
self.stl_center = stl_center
self.stl_reverse_normal = stl_reverse_normal
self.potential = potential
# Handle keyword arguments used in expressions
self.user_defined_kw = {}
for k in list(kw.keys()):
if (implicit_function is not None and re.search(r'\b%s\b'%k, implicit_function) or
(potential is not None and re.search(r'\b%s\b'%k, potential))):
self.user_defined_kw[k] = kw[k]
del kw[k]
self.handle_init(kw)
def initialize_inputs(self, solver):
# Add the user defined keywords to my_constants
# The keywords are mangled if there is a conflicting variable already
# defined in my_constants with the same name but different value.
self.mangle_dict = pywarpx.my_constants.add_keywords(self.user_defined_kw)
if self.implicit_function is not None:
expression = pywarpx.my_constants.mangle_expression(self.implicit_function, self.mangle_dict)
pywarpx.warpx.eb_implicit_function = expression
if self.stl_file is not None:
pywarpx.eb2.geom_type = "stl"
pywarpx.eb2.stl_file = self.stl_file
pywarpx.eb2.stl_scale = self.stl_scale
pywarpx.eb2.stl_center = self.stl_center
pywarpx.eb2.stl_reverse_normal = self.stl_reverse_normal
if self.potential is not None:
assert isinstance(solver, ElectrostaticSolver), Exception('The potential is only supported with the ElectrostaticSolver')
expression = pywarpx.my_constants.mangle_expression(self.potential, self.mangle_dict)
pywarpx.warpx.__setattr__('eb_potential(x,y,z,t)', expression)
[docs]class PlasmaLens(picmistandard.base._ClassWithInit):
"""
Custom class to setup a plasma lens lattice.
The applied fields are dependent only on the transverse position.
Parameters
----------
period: float
Periodicity of the lattice (in lab frame, in meters)
starts: list of floats
The start of each lens relative to the periodic repeat
lengths: list of floats
The length of each lens
strengths_E=None: list of floats, default = 0.
The electric field strength of each lens
strengths_B=None: list of floats, default = 0.
The magnetic field strength of each lens
The field that is applied depends on the transverse position of the particle, (x,y)
- Ex = x*stengths_E
- Ey = y*stengths_E
- Bx = +y*stengths_B
- By = -x*stengths_B
"""
def __init__(self, period, starts, lengths, strengths_E=None, strengths_B=None, **kw):
self.period = period
self.starts = starts
self.lengths = lengths
self.strengths_E = strengths_E
self.strengths_B = strengths_B
assert (self.strengths_E is not None) or (self.strengths_B is not None),\
Exception('One of strengths_E or strengths_B must be supplied')
self.handle_init(kw)
def initialize_inputs(self):
pywarpx.particles.E_ext_particle_init_style = 'repeated_plasma_lens'
pywarpx.particles.B_ext_particle_init_style = 'repeated_plasma_lens'
pywarpx.particles.repeated_plasma_lens_period = self.period
pywarpx.particles.repeated_plasma_lens_starts = self.starts
pywarpx.particles.repeated_plasma_lens_lengths = self.lengths
pywarpx.particles.repeated_plasma_lens_strengths_E = self.strengths_E
pywarpx.particles.repeated_plasma_lens_strengths_B = self.strengths_B
[docs]class Simulation(picmistandard.PICMI_Simulation):
"""
See `Input Parameters <https://warpx.readthedocs.io/en/latest/usage/parameters.html>`_ for more information.
Parameters
----------
warpx_current_deposition_algo: {'direct', 'esirkepov', and 'vay'}, optional
Current deposition algorithm. The default depends on conditions.
warpx_charge_deposition_algo: {'standard'}, optional
Charge deposition algorithm.
warpx_field_gathering_algo: {'energy-conserving', 'momentum-conserving'}, optional
Field gathering algorithm. The default depends on conditions.
warpx_particle_pusher_algo: {'boris', 'vay', 'higuera'}, default='boris'
Particle pushing algorithm.
warpx_use_filter: bool, optional
Whether to use filtering. The default depends on the conditions.
warpx_serialize_initial_conditions: bool, default=False
Controls the random numbers used for initialization.
This parameter should only be used for testing and continuous integration.
warpx_do_dynamic_scheduling: bool, default=True
Whether to do dynamic scheduling with OpenMP
warpx_load_balance_intervals: string, default='0'
The intervals for doing load balancing
warpx_load_balance_efficiency_ratio_threshold: float, default=1.1
(See documentation)
warpx_load_balance_with_sfc: bool, default=0
(See documentation)
warpx_load_balance_knapsack_factor: float, default=1.24
(See documentation)
warpx_load_balance_costs_update: {'heuristic' or 'timers' or 'gpuclock'}, optional
(See documentation)
warpx_costs_heuristic_particles_wt: float, optional
(See documentation)
warpx_costs_heuristic_cells_wt: float, optional
(See documentation)
warpx_use_fdtd_nci_corr: bool, optional
Whether to use the NCI correction when using the FDTD solver
warpx_amr_check_input: bool, optional
Whether AMReX should perform checks on the input
(primarily related to the max grid size and blocking factors)
warpx_amr_restart: string, optional
The name of the restart to use
warpx_zmax_plasma_to_compute_max_step: float, optional
Sets the simulation run time based on the maximum z value
warpx_collisions: collision instance, optional
The collision instance specifying the particle collisions
warpx_embedded_boundary: embedded boundary instance, optional
warpx_break_signals: list of strings
Signals on which to break
warpx_checkpoint_signals: list of strings
Signals on which to write out a checkpoint
"""
# Set the C++ WarpX interface (see _libwarpx.LibWarpX) as an extension to
# Simulation objects. In the future, LibWarpX objects may actually be owned
# by Simulation objects to permit multiple WarpX runs simultaneously.
extension = pywarpx.libwarpx
def init(self, kw):
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_initial_conditions = kw.pop('warpx_serialize_initial_conditions', None)
self.do_dynamic_scheduling = kw.pop('warpx_do_dynamic_scheduling', None)
self.load_balance_intervals = kw.pop('warpx_load_balance_intervals', None)
self.load_balance_efficiency_ratio_threshold = kw.pop('warpx_load_balance_efficiency_ratio_threshold', None)
self.load_balance_with_sfc = kw.pop('warpx_load_balance_with_sfc', None)
self.load_balance_knapsack_factor = kw.pop('warpx_load_balance_knapsack_factor', None)
self.load_balance_costs_update = kw.pop('warpx_load_balance_costs_update', None)
self.costs_heuristic_particles_wt = kw.pop('warpx_costs_heuristic_particles_wt', None)
self.costs_heuristic_cells_wt = kw.pop('warpx_costs_heuristic_cells_wt', None)
self.use_fdtd_nci_corr = kw.pop('warpx_use_fdtd_nci_corr', None)
self.amr_check_input = kw.pop('warpx_amr_check_input', None)
self.amr_restart = kw.pop('warpx_amr_restart', None)
self.zmax_plasma_to_compute_max_step = kw.pop('warpx_zmax_plasma_to_compute_max_step', None)
self.collisions = kw.pop('warpx_collisions', None)
self.embedded_boundary = kw.pop('warpx_embedded_boundary', None)
self.break_signals = kw.pop('warpx_break_signals', None)
self.checkpoint_signals = kw.pop('warpx_checkpoint_signals', 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.time_step_size
if self.gamma_boost is not None:
pywarpx.warpx.gamma_boost = self.gamma_boost
pywarpx.warpx.boost_direction = 'z'
pywarpx.warpx.zmax_plasma_to_compute_max_step = self.zmax_plasma_to_compute_max_step
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.algo.load_balance_intervals = self.load_balance_intervals
pywarpx.algo.load_balance_efficiency_ratio_threshold = self.load_balance_efficiency_ratio_threshold
pywarpx.algo.load_balance_with_sfc = self.load_balance_with_sfc
pywarpx.algo.load_balance_knapsack_factor = self.load_balance_knapsack_factor
pywarpx.algo.load_balance_costs_update = self.load_balance_costs_update
pywarpx.algo.costs_heuristic_particles_wt = self.costs_heuristic_particles_wt
pywarpx.algo.costs_heuristic_cells_wt = self.costs_heuristic_cells_wt
pywarpx.warpx.use_filter = self.use_filter
pywarpx.warpx.serialize_initial_conditions = self.serialize_initial_conditions
pywarpx.warpx.do_dynamic_scheduling = self.do_dynamic_scheduling
pywarpx.particles.use_fdtd_nci_corr = self.use_fdtd_nci_corr
pywarpx.amr.check_input = self.amr_check_input
pywarpx.warpx.break_signals = self.break_signals
pywarpx.warpx.checkpoint_signals = self.checkpoint_signals
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 and (len(self.species) > 0 or len(self.lasers) > 0):
if isinstance(particle_shape, str):
interpolation_order = {'NGP':0, 'linear':1, 'quadratic':2, 'cubic':3}[particle_shape]
else:
interpolation_order = particle_shape
pywarpx.algo.particle_shape = interpolation_order
self.solver.initialize_inputs()
for i in range(len(self.species)):
self.species[i].initialize_inputs(self.layouts[i],
self.initialize_self_fields[i],
self.injection_plane_positions[i],
self.injection_plane_normal_vectors[i])
if self.collisions is not None:
pywarpx.collisions.collision_names = []
for collision in self.collisions:
pywarpx.collisions.collision_names.append(collision.name)
collision.initialize_inputs()
if self.embedded_boundary is not None:
self.embedded_boundary.initialize_inputs(self.solver)
for i in range(len(self.lasers)):
self.lasers[i].initialize_inputs()
self.laser_injection_methods[i].initialize_inputs(self.lasers[i])
for applied_field in self.applied_fields:
applied_field.initialize_inputs()
for diagnostic in self.diagnostics:
diagnostic.initialize_inputs()
if self.amr_restart:
pywarpx.amr.restart = self.amr_restart
def initialize_warpx(self, mpi_comm=None):
if self.warpx_initialized:
return
self.warpx_initialized = True
pywarpx.warpx.init(mpi_comm)
[docs] def step(self, nsteps=None, mpi_comm=None):
self.initialize_inputs()
self.initialize_warpx(mpi_comm)
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()
# ----------------------------
# Simulation frame diagnostics
# ----------------------------
class WarpXDiagnosticBase(object):
"""
Base class for all WarpX diagnostic containing functionality shared by
all WarpX diagnostic installations.
"""
def add_diagnostic(self):
# reduced diagnostics go in a different bucket than regular diagnostics
if isinstance(self, ReducedDiagnostic):
bucket = pywarpx.reduced_diagnostics
name_template = 'reduced_diag'
else:
bucket = pywarpx.diagnostics
name_template = 'diag'
name = getattr(self, 'name', None)
if name is None:
diagnostics_number = (len(bucket._diagnostics_dict) + 1)
self.name = f'{name_template}{diagnostics_number}'
try:
self.diagnostic = bucket._diagnostics_dict[self.name]
except KeyError:
self.diagnostic = pywarpx.Diagnostics.Diagnostic(
self.name, _species_dict={}
)
bucket._diagnostics_dict[self.name] = self.diagnostic
def set_write_dir(self):
if self.write_dir is not None or self.file_prefix is not None:
write_dir = (self.write_dir or 'diags')
file_prefix = (self.file_prefix or self.name)
self.diagnostic.file_prefix = os.path.join(write_dir, file_prefix)
[docs]class FieldDiagnostic(picmistandard.PICMI_FieldDiagnostic, WarpXDiagnosticBase):
"""
See `Input Parameters <https://warpx.readthedocs.io/en/latest/usage/parameters.html>`_ for more information.
Parameters
----------
warpx_plot_raw_fields: bool, optional
Flag whether to dump the raw fields
warpx_plot_raw_fields_guards: bool, optional
Flag whether the raw fields should include the guard cells
warpx_format: {plotfile, checkpoint, openpmd, ascent, sensei}, optional
Diagnostic file format
warpx_openpmd_backend: {bp, h5, json}, optional
Openpmd backend file format
warpx_file_prefix: string, optional
Prefix on the diagnostic file name
warpx_file_min_digits: integer, optional
Minimum number of digits for the time step number in the file name
warpx_dump_rz_modes: bool, optional
Flag whether to dump the data for all RZ modes
"""
def init(self, kw):
self.plot_raw_fields = kw.pop('warpx_plot_raw_fields', None)
self.plot_raw_fields_guards = kw.pop('warpx_plot_raw_fields_guards', None)
self.plot_finepatch = kw.pop('warpx_plot_finepatch', None)
self.plot_crsepatch = kw.pop('warpx_plot_crsepatch', None)
self.format = kw.pop('warpx_format', 'plotfile')
self.openpmd_backend = kw.pop('warpx_openpmd_backend', None)
self.file_prefix = kw.pop('warpx_file_prefix', None)
self.file_min_digits = kw.pop('warpx_file_min_digits', None)
self.dump_rz_modes = kw.pop('warpx_dump_rz_modes', None)
def initialize_inputs(self):
self.add_diagnostic()
self.diagnostic.diag_type = 'Full'
self.diagnostic.format = self.format
self.diagnostic.openpmd_backend = self.openpmd_backend
self.diagnostic.file_min_digits = self.file_min_digits
self.diagnostic.dump_rz_modes = self.dump_rz_modes
self.diagnostic.intervals = self.period
self.diagnostic.diag_lo = self.lower_bound
self.diagnostic.diag_hi = self.upper_bound
if self.number_of_cells is not None:
self.diagnostic.coarsening_ratio = (np.array(self.grid.number_of_cells)/np.array(self.number_of_cells)).astype(int)
# --- Use a set to ensure that fields don't get repeated.
fields_to_plot = set()
if self.data_list is not None:
for dataname in self.data_list:
if dataname == 'E':
fields_to_plot.add('Ex')
fields_to_plot.add('Ey')
fields_to_plot.add('Ez')
elif dataname == 'B':
fields_to_plot.add('Bx')
fields_to_plot.add('By')
fields_to_plot.add('Bz')
elif dataname == 'J':
fields_to_plot.add('jx')
fields_to_plot.add('jy')
fields_to_plot.add('jz')
elif dataname in ['Ex', 'Ey', 'Ez', 'Bx', 'By', 'Bz', 'rho', 'phi', 'F', 'proc_number', 'part_per_cell']:
fields_to_plot.add(dataname)
elif dataname in ['Jx', 'Jy', 'Jz']:
fields_to_plot.add(dataname.lower())
elif dataname.startswith('rho_'):
# Adds rho_species diagnostic
fields_to_plot.add(dataname)
elif dataname == 'dive':
fields_to_plot.add('divE')
elif dataname == 'divb':
fields_to_plot.add('divB')
elif dataname == 'raw_fields':
self.plot_raw_fields = 1
elif dataname == 'raw_fields_guards':
self.plot_raw_fields_guards = 1
elif dataname == 'finepatch':
self.plot_finepatch = 1
elif dataname == 'crsepatch':
self.plot_crsepatch = 1
elif dataname == 'none':
fields_to_plot = set(('none',))
# --- Convert the set to a sorted list so that the order
# --- is the same on all processors.
fields_to_plot = list(fields_to_plot)
fields_to_plot.sort()
self.diagnostic.fields_to_plot = fields_to_plot
self.diagnostic.plot_raw_fields = self.plot_raw_fields
self.diagnostic.plot_raw_fields_guards = self.plot_raw_fields_guards
self.diagnostic.plot_finepatch = self.plot_finepatch
self.diagnostic.plot_crsepatch = self.plot_crsepatch
self.set_write_dir()
ElectrostaticFieldDiagnostic = FieldDiagnostic
[docs]class Checkpoint(picmistandard.base._ClassWithInit, WarpXDiagnosticBase):
"""
Sets up checkpointing of the simulation, allowing for later restarts
See `Input Parameters <https://warpx.readthedocs.io/en/latest/usage/parameters.html>`_ for more information.
Parameters
----------
warpx_file_prefix: string
The prefix to the checkpoint directory names
warpx_file_min_digits: integer
Minimum number of digits for the time step number in the checkpoint
directory name.
"""
def __init__(self, period = 1, write_dir = None, name = None, **kw):
self.period = period
self.write_dir = write_dir
self.file_prefix = kw.pop('warpx_file_prefix', None)
self.file_min_digits = kw.pop('warpx_file_min_digits', None)
self.name = name
if self.name is None:
self.name = 'chkpoint'
self.handle_init(kw)
def initialize_inputs(self):
self.add_diagnostic()
self.diagnostic.intervals = self.period
self.diagnostic.diag_type = 'Full'
self.diagnostic.format = 'checkpoint'
self.diagnostic.file_min_digits = self.file_min_digits
self.set_write_dir()
[docs]class ParticleDiagnostic(picmistandard.PICMI_ParticleDiagnostic, WarpXDiagnosticBase):
"""
See `Input Parameters <https://warpx.readthedocs.io/en/latest/usage/parameters.html>`_ for more information.
Parameters
----------
warpx_format: {plotfile, checkpoint, openpmd, ascent, sensei}, optional
Diagnostic file format
warpx_openpmd_backend: {bp, h5, json}, optional
Openpmd backend file format
warpx_file_prefix: string, optional
Prefix on the diagnostic file name
warpx_file_min_digits: integer, optional
Minimum number of digits for the time step number in the file name
warpx_random_fraction: float, optional
Random fraction of particles to include in the diagnostic
warpx_uniform_stride: integer, optional
Stride to down select to the particles to include in the diagnostic
warpx_plot_filter_function: string, optional
Analytic expression to down select the particles to in the diagnostic
"""
def init(self, kw):
self.format = kw.pop('warpx_format', 'plotfile')
self.openpmd_backend = kw.pop('warpx_openpmd_backend', None)
self.file_prefix = kw.pop('warpx_file_prefix', None)
self.file_min_digits = kw.pop('warpx_file_min_digits', None)
self.random_fraction = kw.pop('warpx_random_fraction', None)
self.uniform_stride = kw.pop('warpx_uniform_stride', None)
self.plot_filter_function = kw.pop('warpx_plot_filter_function', None)
self.user_defined_kw = {}
if self.plot_filter_function is not None:
# This allows variables to be used in the plot_filter_function, but
# in order not to break other codes, the variables must begin with "warpx_"
for k in list(kw.keys()):
if k.startswith('warpx_') and re.search(r'\b%s\b'%k, self.plot_filter_function):
self.user_defined_kw[k] = kw[k]
del kw[k]
self.mangle_dict = None
def initialize_inputs(self):
self.add_diagnostic()
self.diagnostic.diag_type = 'Full'
self.diagnostic.format = self.format
self.diagnostic.openpmd_backend = self.openpmd_backend
self.diagnostic.file_min_digits = self.file_min_digits
self.diagnostic.intervals = self.period
self.set_write_dir()
# --- Use a set to ensure that fields don't get repeated.
variables = set()
if self.data_list is not None:
for dataname in self.data_list:
if dataname == 'position':
# --- The positions are alway written out anyway
pass
elif dataname == 'momentum':
variables.add('ux')
variables.add('uy')
variables.add('uz')
elif dataname == 'weighting':
variables.add('w')
elif dataname == 'fields':
variables.add('Ex')
variables.add('Ey')
variables.add('Ez')
variables.add('Bx')
variables.add('By')
variables.add('Bz')
elif dataname in ['ux', 'uy', 'uz', 'Ex', 'Ey', 'Ez', 'Bx', 'By', 'Bz']:
variables.add(dataname)
# --- Convert the set to a sorted list so that the order
# --- is the same on all processors.
variables = list(variables)
variables.sort()
if np.iterable(self.species):
species_list = self.species
else:
species_list = [self.species]
if self.mangle_dict is None:
# Only do this once so that the same variables are used in this distribution
# is used multiple times
self.mangle_dict = pywarpx.my_constants.add_keywords(self.user_defined_kw)
for specie in species_list:
diag = pywarpx.Bucket.Bucket(self.name + '.' + specie.name,
variables = variables,
random_fraction = self.random_fraction,
uniform_stride = self.uniform_stride)
expression = pywarpx.my_constants.mangle_expression(self.plot_filter_function, self.mangle_dict)
diag.__setattr__('plot_filter_function(t,x,y,z,ux,uy,uz)', expression)
self.diagnostic._species_dict[specie.name] = diag
# ----------------------------
# Lab frame diagnostics
# ----------------------------
[docs]class LabFrameFieldDiagnostic(picmistandard.PICMI_LabFrameFieldDiagnostic,
WarpXDiagnosticBase):
"""
See `Input Parameters <https://warpx.readthedocs.io/en/latest/usage/parameters.html>`_ for more information.
Parameters
----------
warpx_new_BTD: bool, optional
Use the new BTD diagnostics
warpx_format: string, optional
Passed to <diagnostic name>.format
warpx_openpmd_backend: string, optional
Passed to <diagnostic name>.openpmd_backend
warpx_file_prefix: string, optional
Passed to <diagnostic name>.file_prefix
warpx_file_min_digits: integer, optional
Passed to <diagnostic name>.file_min_digits
warpx_buffer_size: integer, optional
Passed to <diagnostic name>.buffer_size
warpx_lower_bound: vector of floats, optional
Passed to <diagnostic name>.lower_bound
warpx_upper_bound: vector of floats, optional
Passed to <diagnostic name>.upper_bound
"""
def init(self, kw):
self.use_new_BTD = kw.pop('warpx_new_BTD', False)
if self.use_new_BTD:
# The user is using the new BTD
self.format = kw.pop('warpx_format', None)
self.openpmd_backend = kw.pop('warpx_openpmd_backend', None)
self.file_prefix = kw.pop('warpx_file_prefix', None)
self.file_min_digits = kw.pop('warpx_file_min_digits', None)
self.buffer_size = kw.pop('warpx_buffer_size', None)
self.lower_bound = kw.pop('warpx_lower_bound', None)
self.upper_bound = kw.pop('warpx_upper_bound', None)
def initialize_inputs(self):
if self.use_new_BTD:
self.initialize_inputs_new()
else:
self.initialize_inputs_old()
def initialize_inputs_old(self):
pywarpx.warpx.check_consistency('num_snapshots_lab', self.num_snapshots, 'The number of snapshots must be the same in all lab frame diagnostics')
pywarpx.warpx.check_consistency('dt_snapshots_lab', self.dt_snapshots, 'The time between snapshots must be the same in all lab frame diagnostics')
pywarpx.warpx.check_consistency('lab_data_directory', self.write_dir, 'The write directory must be the same in all lab frame diagnostics')
pywarpx.warpx.do_back_transformed_diagnostics = 1
pywarpx.warpx.num_snapshots_lab = self.num_snapshots
pywarpx.warpx.dt_snapshots_lab = self.dt_snapshots
pywarpx.warpx.do_back_transformed_fields = 1
pywarpx.warpx.lab_data_directory = self.write_dir
def initialize_inputs_new(self):
self.add_diagnostic()
self.diagnostic.diag_type = 'BackTransformed'
self.diagnostic.format = self.format
self.diagnostic.openpmd_backend = self.openpmd_backend
self.diagnostic.file_min_digits = self.file_min_digits
self.diagnostic.diag_lo = self.lower_bound
self.diagnostic.diag_hi = self.upper_bound
self.diagnostic.do_back_transformed_fields = 1
self.diagnostic.num_snapshots_lab = self.num_snapshots
self.diagnostic.dt_snapshots_lab = self.dt_snapshots
self.diagnostic.buffer_size = self.buffer_size
# --- Use a set to ensure that fields don't get repeated.
fields_to_plot = set()
if self.data_list is not None:
for dataname in self.data_list:
if dataname == 'E':
fields_to_plot.add('Ex')
fields_to_plot.add('Ey')
fields_to_plot.add('Ez')
elif dataname == 'B':
fields_to_plot.add('Bx')
fields_to_plot.add('By')
fields_to_plot.add('Bz')
elif dataname == 'J':
fields_to_plot.add('jx')
fields_to_plot.add('jy')
fields_to_plot.add('jz')
elif dataname in ['Ex', 'Ey', 'Ez', 'Bx', 'By', 'Bz', 'rho']:
fields_to_plot.add(dataname)
elif dataname in ['Jx', 'Jy', 'Jz']:
fields_to_plot.add(dataname.lower())
elif dataname.startswith('rho_'):
# Adds rho_species diagnostic
fields_to_plot.add(dataname)
# --- Convert the set to a sorted list so that the order
# --- is the same on all processors.
fields_to_plot = list(fields_to_plot)
fields_to_plot.sort()
self.diagnostic.fields_to_plot = fields_to_plot
self.set_write_dir()
[docs]class LabFrameParticleDiagnostic(picmistandard.PICMI_LabFrameParticleDiagnostic):
def initialize_inputs(self):
pywarpx.warpx.check_consistency('num_snapshots_lab', self.num_snapshots, 'The number of snapshots must be the same in all lab frame diagnostics')
pywarpx.warpx.check_consistency('dt_snapshots_lab', self.dt_snapshots, 'The time between snapshots must be the same in all lab frame diagnostics')
pywarpx.warpx.check_consistency('lab_data_directory', self.write_dir, 'The write directory must be the same in all lab frame diagnostics')
pywarpx.warpx.do_back_transformed_diagnostics = 1
if isinstance(self.species, Species):
self.species.do_back_transformed_diagnostics = 1
else:
try:
for specie in self.species:
specie.do_back_transformed_diagnostics = 1
except TypeError:
pass
pywarpx.warpx.num_snapshots_lab = self.num_snapshots
pywarpx.warpx.dt_snapshots_lab = self.dt_snapshots
pywarpx.warpx.lab_data_directory = self.write_dir
class ReducedDiagnostic(picmistandard.base._ClassWithInit, WarpXDiagnosticBase):
"""
Sets up a reduced diagnostic in the simulation.
See `Input Parameters <https://warpx.readthedocs.io/en/latest/usage/parameters.html#reduced-diagnostics>`_
for more information.
Parameters
----------
diag_type: string
The type of reduced diagnostic. See the link above for all the different
types of reduced diagnostics available.
name: string
The name of this diagnostic which will also be the name of the data
file written to disk.
period: integer
The simulation step interval at which to output this diagnostic.
path: string
The file path in which the diagnostic file should be written.
extension: string
The file extension used for the diagnostic output.
separator: string
The separator between row values in the output file.
"""
def __init__(self, diag_type, name=None, period=1, path=None,
extension=None, separator=None, **kw):
self.name = name
self.type = diag_type
self.intervals = period
self.path = path
self.extension = extension
self.separator = separator
self._species = kw.pop('species', None)
# Now we need to handle all the specific inputs required for the
# different reduced diagnostic types.
# Note: only a limited number are presently supported.
# The simple diagnostics do not require any additional arguments
self._simple_reduced_diagnostics = [
'ParticleEnergy', 'ParticleMomentum', 'FieldEnergy',
'FieldMomentum', 'FieldMaximum', 'RhoMaximum', 'ParticleNumber',
'LoadBalanceCosts', 'LoadBalanceEfficiency',
]
# The species diagnostics require a species to be provided
self._species_reduced_diagnostics = [
'BeamRelevant', 'ParticleHistogram', 'ParticleExtrema',
]
if self.type in self._simple_reduced_diagnostics:
pass
elif self.type in self._species_reduced_diagnostics:
if self._species is None:
raise AttributeError(
f"{self.type} reduced diagnostic requires a species."
)
if self.type == 'ParticleHistogram':
raise NotImplementedError(
f"{self.type} reduced diagnostic is not yet supported "
"in pywarpx."
)
elif self.type == "FieldProbe":
kw = self._handle_field_probe(**kw)
else:
raise RuntimeError(
f"{self.type} reduced diagnostic is not yet supported "
"in pywarpx."
)
self.handle_init(kw)
def _handle_field_probe(self, **kw):
"""Utility function to grab required inputs for a field probe from kw"""
self.probe_geometry = kw.pop("probe_geometry")
self.x_probe = kw.pop("x_probe")
self.y_probe = kw.pop("y_probe")
self.z_probe = kw.pop("z_probe")
self.interp_order = kw.pop("interp_order", None)
self.integrate = kw.pop("integrate", None)
self.do_moving_window_FP = kw.pop("do_moving_window_FP", None)
if self.probe_geometry.lower() != 'point':
self.resolution = kw.pop("resolution")
if self.probe_geometry.lower() == 'line':
self.x1_probe = kw.pop("x1_probe")
self.y1_probe = kw.pop("y1_probe")
self.z1_probe = kw.pop("z1_probe")
if self.probe_geometry.lower() == 'plane':
self.detector_radius = kw.pop("detector_radius")
self.target_normal_x = kw.pop("target_normal_x")
self.target_normal_y = kw.pop("target_normal_y")
self.target_normal_z = kw.pop("target_normal_z")
self.target_up_x = kw.pop("target_up_x")
self.target_up_y = kw.pop("target_up_y")
self.target_up_z = kw.pop("target_up_z")
return kw
def initialize_inputs(self):
self.add_diagnostic()
for key in self.__dict__.keys():
if not key.startswith('_') and key not in ['name', 'diagnostic']:
self.diagnostic.__setattr__(key, self.__dict__[key])
if self._species is not None:
diag = pywarpx.Bucket.Bucket(self.name + '.' + self._species.name)
self.diagnostic._species_dict[self._species.name] = diag