def __init__(self, diag_steps, diag_step_offset=0,
extended_interval_level=None,
manual_timesteps=None,
):
"""Initialize timing. Will often be called within a child object's
init.
Arguments:
diag_steps (int): Run the diagnostic with this period.
diag_step_offset (int): Run when ``simulation_step % diag_steps =
diag_step_offset``. Default 0.
extended_interval_level (int): Enable an exponentially-increasing
interval as the run progresses. If None (default), calculate
every time period. Otherwise, the interval increases on an
exponential scale. 0 is most aggressive at skipping diagnostic
timesteps; 1, 2, etc. are progressively less aggressive but
still exponential. Specifically, 0 executes once during each
power-of-2 set of diagnostic steps; 1 twice; 2 four times, etc.
manual_timesteps (list of ints): A list of timesteps to execute on
within each diag_steps. If this list contains steps for the
entire run, set diag_steps to a very large number. Default
None.
"""
self.diag_steps = int(round(diag_steps))
self.diag_step_offset = int(round(diag_step_offset))
self.extended_interval_level = extended_interval_level
if self.extended_interval_level is not None:
self.extended_interval_level = int(round(
self.extended_interval_level))
self.manual_timesteps = manual_timesteps
if self.manual_timesteps is not None:
self.manual_timesteps = [int(round(x))
for x in self.manual_timesteps]
if (
(self.manual_timesteps is not None)
and (self.extended_interval_level is not None)
):
raise ValueError("manual_timesteps and extended_interval_level "
"cannot be used together; one must be None.")
if mwxrun.initialized:
raise RuntimeError("You must install all diagnostics before "
"initializing warpx")
# handle creating the folder to save diagnostics
callbacks.installafterinit(self._create_diag_folder)
def init(self, kw):
"""PICMI relies on C++ WarpX to translate charge & mass strings to
floats. To get around that we have our own variables sq/sm (species
charge/mass) that are always floats.
Automatically adds the species to the simulation
"""
super(Species, self).init(kw)
if isinstance(self.charge, str):
if self.charge == 'q_e':
self.sq = picmi.constants.q_e
elif self.charge == '-q_e':
self.sq = -picmi.constants.q_e
else:
raise ValueError("Unrecognized charge {}".format(self.charge))
else:
self.sq = self.charge
if isinstance(self.mass, str):
if self.mass == 'm_e':
self.sm = picmi.constants.m_e
elif self.mass == 'm_p':
self.sm = picmi.constants.m_p
else:
raise ValueError("Unrecognized mass {}".format(self.mass))
else:
self.sm = self.mass
mwxrun.simulation.add_species(self,
layout=picmi.GriddedLayout(
n_macroparticle_per_cell=[0, 0],
grid=mwxrun.grid))
self.pids_initialized = False
# Only keys are used; ensures pids are both unique and ordered.
self.waiting_extra_pids = {}
# add a callback to initialize the extra PIDs after sim init
callbacks.installafterinit(self.init_pid_dict)
def __init__(self, diag_steps, preset_string='default',
custom_string=None, install=True, **kwargs):
"""Generate and install function to write out step #.
Arguments:
diag_steps (int): Number of steps between each output
simulation (mespecies.Simulation): Main simulation object
preset_string (str): Defaults to choose between:
- ``default`` - just the step number and total particle num
- ``perfdebug`` - like particledebug, plus interval wall time,
step rate, and particle-step rate
- ``memdebug`` - print out verbose memory usage information
custom_string (str): Overrides preset_string if not None. The full
string to output, with:
- ``{step}`` formatting symbol for where the step number should
go
- ``{wall_time}`` run time of the last diag_steps steps
- ``{step_rate}`` diag_steps / wall_time
- ``{particle_step_rate}`` nplive * diag_steps / wall_time
- ``{nplive}`` for number of live particles (global).
- ``{npperspecies}`` for number of particles per species
(global).
- ``{iproc}`` for the current processor number
- ``{system_memory}`` for verbose information on system memory
usage.
- ``{memory_usage}`` for memory usage of the current process
only.
install (bool): If False, don't actually install this into WarpX.
Use if you want to call manually for debugging.
kwargs: See :class:`mewarpx.mewarpx.diags_store.diag_base.WarpXDiagnostic`
for more timing options.
"""
self.defaults_dict = {
'default': "Step #{step:6d}; {nplive:8d} particles",
'perfdebug': ("Step #{step:6d}; {nplive:8d} particles "
"{npperspecies} "
"{wall_time:6.1f} s wall time; "
"{step_rate:4.2f} steps/s; "
"{particle_step_rate:4.2f} particle*steps/s in the last {diag_steps} steps; "
"{particle_step_rate_total:4.2f} particle*steps/s overall"),
'memdebug': ("{system_memory}\n"
"Proc {iproc} usage:\n{memory_usage}"),
}
if custom_string is not None:
self.diag_string = custom_string
else:
if preset_string not in self.defaults_dict:
logger.warning(("Preset {} not found for set_step_diag, "
"using default").format(preset_string))
preset_string = 'default'
self.diag_string = self.defaults_dict[preset_string]
super(TextDiag, self).__init__(diag_steps=diag_steps, **kwargs)
callbacks.installafterinit(self.init_timers_and_counters)
if install:
callbacks.installafterstep(self.text_diag)
def __init__(self, diag_steps, runinfo, overwrite=True,
history_maxlen=5000, sig_figs=6,
printed_qtys=None,
check_charge_conservation=True,
print_per_diagnostic=True, print_total=False,
plot=True, save_csv=False, profile_decorator=None,
**kwargs):
"""Generate and install function to write out fluxes.
Arguments:
diag_steps (int): Number of steps between each output
runinfo (:class:`mewarpx.runinfo.RunInfo`): RunInfo object is used
to get the species, injectors, surfaces, and system area.
overwrite (bool): If True the dill pickled save file will overwrite
the previous diagnostic period's saved file.
history_maxlen (int): Maximum length of full history to keep. If
this is exceeded, history is resampled to a 2x lower frequency.
Default 5000.
sig_figs (int): Number of significant figures in text output.
Default 6.
printed_qtys (dict): Override individual values of
default_printed_qtys; same input format but keys can be omitted
to use defaults.
check_charge_conservation (bool): Whether to check if charge is
conserved in simulation.
print_per_diagnostic (bool): Whether to print current results for
the latest diagnostic period.
print_total (bool): Whether to print total history of fluxes after a
diagnostic period.
plot (bool): Whether to save a plot of fluxes after each diagnostic
period.
save_csv (bool): Whether to save csv files of scraped / injected
particles.
profile_decorator (decorator): A decorator used to profile the
timeseries update methods and related functions.
kwargs: See :class:`mewarpx.diags_store.diag_base.WarpXDiagnostic`
for more timing options.
"""
# Save input variables
self.runinfo = runinfo
self.history_maxlen = history_maxlen
self.history_dt = mwxrun.get_dt()
self.check_charge_conservation = check_charge_conservation
self.print_per_diagnostic = print_per_diagnostic
self.print_total = print_total
self.plot = plot
if save_csv:
csv_write_dir = os.path.join(self.DIAG_DIR, self.FLUX_DIAG_DIR)
else:
csv_write_dir = None
if profile_decorator is not None:
self.update_ts_dict = profile_decorator(self.update_ts_dict)
self.update_fullhist_dict = (
profile_decorator(self.update_fullhist_dict)
)
self.print_fluxes = profile_decorator(self.print_fluxes)
self.plot_fluxes = profile_decorator(self.plot_fluxes)
self.injector_dict = runinfo.injector_dict
self.surface_dict = runinfo.surface_dict
self.diags_dict = collections.OrderedDict()
for key, val in self.injector_dict.items():
self.injector_dict[key] = mwxutil.return_iterable(val)
self.diags_dict[('inject', key)] = [
InjectorFluxDiag(diag_steps=diag_steps,
injector=injector,
write_dir=csv_write_dir,
**kwargs)
for injector in self.injector_dict[key]
]
for key, val in self.surface_dict.items():
self.surface_dict[key] = mwxutil.return_iterable(val)
self.diags_dict[('scrape', key)] = [
SurfaceFluxDiag(diag_steps=diag_steps,
surface=surface,
write_dir=csv_write_dir,
**kwargs)
for surface in self.surface_dict[key]
]
# Initialize other variables
self.last_run_step = 0
super(FluxDiagnostic, self).__init__(
diag_steps=diag_steps,
runinfo=runinfo,
overwrite=overwrite,
sig_figs=sig_figs,
printed_qtys=printed_qtys,
**kwargs
)
self.check_scraping()
# if this run is from a restart, try to load flux data up to the
# current point
# mwxrun.restart isn't set until mwxrun.init_run is called, so the
# lambda function is needed here
callbacks.installafterinit(
lambda _=None:
self._load_checkpoint_flux() if mwxrun.restart else None
)
callbacks.installafterstep(self._flux_ana)
def __init__(self,
injector_dict,
surface_dict,
electrode_params,
user_var_dict=None,
local_vars=None,
**kwargs):
"""**Initialization**: Set basic run parameters
Call just before creating final diagnostics. This uses built-in WARP
variables, in addition to the passed quantities, to save relevant
parameters for diagnostic calculations and postprocessing.
Arguments:
injector_dict (dict): Dictionary where keys are strings for groups
of surfaces or interaction physics (see component_labels keys
for main possibilities) and values are injectors or lists of
injectors. Used by various diagnostics, so should be complete.
``cathode`` is a required key.
surface_dict (dict): Dictionary where keys are groups of surfaces
which scrape particles and values are surfaces or lists of
surfaces. Used by various diagnostics, so should be complete.
``cathode`` is a required key.
user_var_dict (dict): A dictionary of all parameters set by user.
This will depend on the given system being investigated.
local_vars (dict): Instead of user_var_dict, pass locals() from the
executing script. Then all UPPER_CASE variables will be saved to
user_var_dict internally.
electrode_params (dict): A dictionary of electrode parameters:
* `CATHODE_A` (Amp/m^2/K^2)
* `CATHODE_TEMP` (K)
* `ANODE_TEMP` (K) (if back-emission used)
* `ANODE_A` (Amp/m^2/K^2) (if back-emission used)
* Additional elements if desired
run_param_dict (dict): Optional, a dictionary of run parameters such
as diag_timesteps and total_timesteps.
run_file (str): Path to the main simulation file, in order to store
it in run_param_dict for easy access later.
area (float): Optional, cross-sectional area in m^2 to use for flux
calculations from this simulation. If not supplied, will be
calculated automatically from the geometry.
"""
# Parameter dictionaries and injectors
if local_vars is not None:
if user_var_dict is not None:
raise ValueError(
"Specify local_vars or user_var_dict, not both.")
self.user_var_dict = self._process_local_vars(local_vars)
else:
self.user_var_dict = user_var_dict
self.run_param_dict = kwargs.pop("run_param_dict", None)
self.run_file = kwargs.pop("run_file", None)
self.electrode_params = electrode_params
self.init_species(mwxrun.simulation.species)
self.init_injectors_surfaces(injector_dict, surface_dict)
# Store cross-sectional area for flux calculations
self.area = kwargs.pop('area', None)
self.geom = mwxrun.geom_str
self.pos_lims = [
mwxrun.xmin, mwxrun.xmax, mwxrun.ymin, mwxrun.ymax, mwxrun.zmin,
mwxrun.zmax
]
self.nxyz = [mwxrun.nx, mwxrun.ny, mwxrun.nz]
self.dt = mwxrun.get_dt()
self.periodic = (
#TODO: Should be implemented when needed
#warp.w3d.boundxy is warp.periodic
# The rwall in RZ simulations can be set even if boundxy is
# periodic. It defaults to 1e36 so this should be fine.
#and warp.top.prwall > 1e10
)
if self.area is None:
self.area = mwxrun.get_domain_area()
# All kwargs should have been popped, so raise an error if that isn't
# the case.
if len(kwargs) > 0:
raise ValueError(f'Unrecognized kwarg(s) {list(kwargs.keys())}')
# Shapes/injectors just adds all the lists of shapes together into one
# list
self.shapes = reduce((lambda x, y: x + y),
list(self.surface_dict.values()))
self.injectors = reduce((lambda x, y: x + y),
list(self.injector_dict.values()))
# _reserved_names ensures results dictionary will not have duplicated
# entries that overwrite each other.
self._reserved_names = (
["cathode_all", "anode_all", "aside", "aside_nodiel"] +
[x + "_all" for x in self.component_list])
callbacks.installafterinit(self._runinfo_after_init)
def init_grid(self,
lower_bound,
upper_bound,
number_of_cells,
use_rz=False,
**kwargs):
"""Function to set up the simulation grid.
Arguments:
lower_bound (list): Minimum coordinates for all direction; length
of list has to equal number of dimensions.
upper_bound (list): Maximum coordinates for all direction; length
of list has to equal number of dimensions.
number_of_cells (list): Number of grid cells in each direction.
use_rz (bool): If True, cylindrical coordinates will be used.
kwargs (dict): Dictionary containing extra arguments. These can
include:
- ``bc_fields_x/y/z/r_min/max``: field boundary condition
settings (can be 'periodic', 'dirichlet' or 'none').
- ``bc_particles_x/y/z/r_min/max``: particle boundary
condition settings (can be 'periodic', 'absorbing' or
'reflecting').
- ``min_tiles``: the minimum number of tiles. See function
``_set_max_grid_size()`` below for details.
"""
self.dim = len(lower_bound)
# sanity check inputs
assert len(upper_bound) == self.dim
assert len(number_of_cells) == self.dim
if use_rz and self.dim != 2:
raise RuntimeError(
f"Cannot use cylindrical coordinates in {self.dim} dimensions."
)
self._set_geom_str(use_rz)
boundary_conditions = self._get_default_boundary_conditions()
# update boundary conditions if provided
for key in [key for key in kwargs.keys() if key.startswith('bc')]:
kind, dim, side = key.split('_')[-3:]
boundary_conditions[kind][side][self.coord_map[dim]] = kwargs[key]
if self.geom_str == 'RZ':
grid = picmi.CylindricalGrid
elif self.geom_str == 'Z':
grid = picmi.Cartesian1DGrid
elif self.geom_str == 'XZ':
grid = picmi.Cartesian2DGrid
elif self.geom_str == 'XYZ':
grid = picmi.Cartesian3DGrid
self.grid = grid(
number_of_cells=number_of_cells,
lower_bound=lower_bound,
upper_bound=upper_bound,
lower_boundary_conditions=boundary_conditions['fields']['min'],
upper_boundary_conditions=boundary_conditions['fields']['max'],
lower_boundary_conditions_particles=boundary_conditions[
'particles']['min'],
upper_boundary_conditions_particles=boundary_conditions[
'particles']['max'],
)
self._set_grid_params()
self._set_max_grid_size(kwargs.get('min_tiles', None))
self._print_grid_params()
# there are a number of initialization tasks that have to happen
# immediately after init so this function is installed as the
# first afterinit callback
if len(callbacks._afterinit) != 0:
raise RuntimeError(
"mwxrun._after_init must be the first afterinit callback to "
"execute, but one was already installed. Please ensure "
"init_grid is called before other items in the run script.")
callbacks.installafterinit(self._after_init)