def __init__(self, diag_steps,
name=init_restart_util.DEFAULT_CHECKPOINT_NAME,
clear_old_checkpoints=True, **kwargs):
"""
This class is a wrapper for creating checkpoints from which
to restart a simulation from. Adding flux diagnostic data to
checkpoints are supported, but since this class has to be initialized
before the simulation and flux diagnostics are initialized after
the simulation, the user is responsible for adding the ``flux_diag``
attribute to this object in a simulation input file.
Arguments:
diag_steps (int): Run the diagnostic with this period.
Also plot on this period if enabled.
name (str): The name of the diagnostic to be passed into the
picmi checkpoint diagnostic.
clear_old_checkpoints (bool): If True old checkpoints will be
deleted after new ones are created.
kwargs: For a list of valid keyword arguments see
:class:`mewarpx.diags_store.diag_base.WarpXDiagnostic`
"""
self.checkpoint_steps = diag_steps
self.name = name
self.clear_old_checkpoints = clear_old_checkpoints
self.flux_diag = None
if kwargs.pop('diag_step_offset', 1) != 1:
raise ValueError(
"diag_step_offset in CheckPointDiagnostic must be 1.")
super(CheckPointDiagnostic, self).__init__(
diag_steps=diag_steps, diag_step_offset=1, **kwargs)
self.write_dir = self.DIAG_DIR
self.add_checkpoint()
# note that WarpX diagnostics (including checkpoints) are output
# after the "afterstep" functions, therefore we install this function
# before steps and add 1 to the diag_steps above
callbacks.installbeforestep(self.checkpoint_manager)
def __init__(self, diag_steps, surface, write_dir, **kwargs):
"""Initialize surface-specific features.
Arguments:
diag_steps (int): Number of steps between each output
surface (mewarpx.Assembly object): The assembly object in which
particles are scraped.
write_dir (string): Directory to write CSV files to.
kwargs: See :class:`mewarpx.diags_store.diag_base.WarpXDiagnostic`
for more timing options.
"""
self.surface = surface
# Register with the surface, make sure we don't have two diagnostics
# attached, and initialize with fields.
if self.surface.scraper_diag is not None:
raise RuntimeError("Cannot attach two SurfaceFluxDiag objects "
"to the same surface.")
self.surface.scraper_diag = self
self.surface.init_scrapedparticles(self.surface.fields)
save_name = self.surface.name + '_scraped.csv'
super(SurfaceFluxDiag, self).__init__(
diag_steps=diag_steps,
write_dir=write_dir,
save_name=save_name,
**kwargs
)
# install callback that will have the assembly object check the
# particle buffer and record the scraped particles data
# TODO change this to only process the particle buffer on a diagnostic
# step instead of after every step
callbacks.installbeforestep(
mwxrun.sim_ext.libwarpx_so.warpx_clearParticleBoundaryBuffer)
callbacks.installbeforeEsolve(self.surface.record_scrapedparticles)
nps = 10
x = np.random.rand(nps) * 0.03
y = np.zeros(nps)
z = np.random.random(nps) * 0.03
ux = np.random.normal(loc=0, scale=1e3, size=nps)
uy = np.random.normal(loc=0, scale=1e3, size=nps)
uz = np.random.normal(loc=0, scale=1e3, size=nps)
w = np.ones(nps) * 2.0
newPid = 5.0
_libwarpx.add_particles(
species_name='electrons', x=x, y=y, z=z, ux=ux, uy=uy, uz=uz,
w=w, newPid=newPid, unique_particles=(not color)
)
callbacks.installbeforestep(add_particles)
##########################
# simulation run
##########################
sim.step(max_steps - 1, mpi_comm=new_comm)
##########################
# check that the new PIDs are properly set
##########################
if color == 0:
assert (_libwarpx.get_particle_count('electrons') == 90)
else:
assert (_libwarpx.get_particle_count('electrons') == 90)