def run_sim(self):
self.sim.step(self.max_steps - self.diag_steps)
callbacks.installafterstep(self._get_rho_ions)
self.sim.step(self.diag_steps)
if self.sim.extension.getMyProc() == 0:
np.save('avg_ion_density.npy', self.ion_density_array)
def __init__(self, total_steps, diag_steps):
"""Diagnostic to record the electron velocity variance at every
diagnostic step."""
self.i = 0
self.results_array = np.zeros((total_steps//diag_steps, 2))
super(VarianceTracker, self).__init__(diag_steps)
callbacks.installafterstep(self._record_variance)
def __init__(self, diag_steps, name=None, species=None,
data_list=None, post_processing=False, plot_data_list=None,
plot_species=None, **kwargs):
"""Initializes the picmi.ParticleDiagnostic and adds the diagnostic to
the simulation
Arguments:
diag_steps (int): Number of steps between each diagnostic output
name (str): name of the diag output folder, defaults to
``particle_diag``
species (:class:`mewarpx.mespecies.Species`): species in the
simulation, if None then uses all particles in the simulation
data_list (list str): list of attributes to be outputted by the
diagnostic, default uses ``["position", "momentum", "weighting"]``
post_process (bool): generate plots for each diagnostic data
directory produced
plot_data_list (list str): list of data to be plotted for each
diagnostic step ("particle_position_x", "particle_position_y",
"particle_position_z", "particle_momentum_x",
"particle_momentum_y", "particle_momentum_z")
plot_species (list str): list of species names to be plotted,
defaults to all species if not specified. Name variable in
:class:`mewarpx.mespecies.Species` must be set for each species
in the simulation
"""
self.name = name
self.species = species
self.data_list = data_list
self.write_dir = os.path.join(self.DIAG_DIR, self.PARTICLE_DIAG_DIR)
self.post_processing = post_processing
self.plot_data_list = plot_data_list
self.plot_species = plot_species
if self.name is None:
self.name = 'particle_diag'
if self.species is None:
self.species = mwxrun.simulation.species
if self.data_list is None:
self.data_list = ['position', 'momentum', 'weighting']
if self.plot_species is None:
self.plot_species = []
for specimen in self.species:
self.plot_species.append(specimen.name)
super(ParticleDiagnostic, self).__init__(
diag_steps=diag_steps, **kwargs)
self.add_particle_diag()
if self.post_processing:
if self.plot_data_list is None:
warnings.warn('Warning: post_process is True,'
'but plot_data_list is None!')
else:
callbacks.installafterstep(self.check_for_end_of_sim)
def run_sim(self):
self.sim.step(self.max_steps - self.diag_steps)
self.rho_wrapper = fields.RhoFPWrapper(0, False)
callbacks.installafterstep(self._get_rho_ions)
self.sim.step(self.diag_steps)
if self.sim.extension.getMyProc() == 0:
np.save(f'ion_density_case_{self.n+1}.npy', self.ion_density_array)
def run_sim(self):
if self.steps is not None:
mwxrun.simulation.step(self.steps)
self.text_diag.print_performance_summary()
else:
mwxrun.simulation.step(self.MAX_STEPS - self.DIAG_STEPS)
callbacks.installafterstep(self._get_rho_ions)
mwxrun.simulation.step(self.DIAG_STEPS)
if mwxrun.me == 0:
np.save('avg_rho_data.npy', self.rho_array)
def __init__(self,
electron_species,
ion_species,
log_lambda=None,
subcycling_steps=1):
"""
Arguments:
electron_species (:class:`mewarpx.mespecies.Species`): Electron
species that will be scattered off the ion species.
ion_species (:class:`mewarpx.mespecies.Species`): Ion species from
which electrons will be scattered.
log_lambda (float): If specified, a fixed value for the Coulomb
logarithm. If not specified it will be calculated according to
the NRL formulary.
subcycling_steps (int): Number of steps between updating the grid
quantities. Default 1.
"""
self.collider = electron_species
self.field = ion_species
self.log_lambda = log_lambda
self.subcycling_steps = subcycling_steps
self.nu_coef = (
self.collider.sq**2 * self.field.sq**2 /
(4.0 * np.pi * constants.epsilon_0**2 * self.collider.sm**2))
if mwxrun.geom_str not in ['Z', 'XZ']:
raise NotImplementedError(
"Currently LangevinElectronIonScattering is only implemented "
"for Z and XZ geometries.")
print_str = ("Initialized electron-ion Coulomb scattering for species "
f"{self.collider.name} and {self.field.name}.")
if self.log_lambda is None:
print_str += 'The Coulomb logarithm will be calculated.'
else:
print_str += ('The Coulomb logarithm is fixed at %.2f.' %
self.log_lambda)
logger.info(print_str)
callbacks.installafterstep(self.run_scattering_method)
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,
diag_steps,
process_phi=True,
process_E=False,
process_rho=True,
species_list=None,
plot=True,
barrier_slices=None,
max_dim=16.0,
min_dim=0.0,
dpi=300,
install_field_diagnostic=False,
post_processing=False,
**kwargs):
"""
This class handles diagnostics for field quantities (output and
plotting) typically of interest in Modern Electron simulations.
Optionally a picmi FieldDiagnostic can also be installed.
Arguments:
diag_steps (int): Run the diagnostic with this period.
Also plot on this period if enabled.
process_phi (bool): If True, output phi. Default True.
process_E (bool): If True, output E. Default False.
process_rho (bool): If True, output rho. Default True.
species_list (list): Optional list of picmi.Species objects for
which output density. If not specified all species will be
tabulated.
plot (bool): If True, also generate plots. Default True.
barrier_slices (list): If provided, also plot potential energy
slices as a function of z for each value of x/r in the list.
Units are in meters.
max_dim (float): Maximum figure dimension in inches for field plots.
min_dim (float): Minimum figure dimension in inches for field plots.
dpi (int): Resolution to use for saved plots
install_field_diagnostic (bool): If true, install a picmi
FieldDiagnostic. All parameters for the diagnostic should be
passed as keyword arguments.
post_processing (bool): Whether or not to plot data after simulation
ends from any yt files generated during the run.
kwargs: For a list of valid keyword arguments see
diag_base.WarpXDiagnostic
"""
self.diag_steps = diag_steps
self.write_dir = os.path.join(self.DIAG_DIR, self.FIELD_DIAG_DIR)
# field quantities to process
self.process_phi = process_phi
self.process_E = process_E
self.process_rho = process_rho
self.species_list = species_list
if self.species_list is None:
self.species_list = mwxrun.simulation.species
self.barrier_slices = barrier_slices
self.plot = plot
self.max_dim = max_dim
self.min_dim = min_dim
self.dpi = dpi
self.a_ax = 'z'
self.o_ax = 'x'
super(FieldDiagnostic,
self).__init__(diag_steps, kwargs.pop('diag_step_offset', 0),
kwargs.pop('extended_interval_level', None),
kwargs.pop('manual_timesteps', None))
self.install_field_diagnostic = install_field_diagnostic
self.post_processing = post_processing
self.kwargs = kwargs
callbacks.installafterstep(self.fields_diag)
if self.install_field_diagnostic:
self.add_field_diag()