def test_create_checkpoints_with_fluxdiag():
testing_util.initialize_testingdir("test_create_checkpoints_with_fluxdiag")
# use a fixed random seed
np.random.seed(47239475)
run = get_run()
run.init_injectors()
checkpoint = CheckPointDiagnostic(DIAG_STEPS,
CHECKPOINT_NAME,
clear_old_checkpoints=True)
run.init_runinfo()
run.init_fluxdiag()
mwxrun.init_run(restart=False)
checkpoint.flux_diag = run.fluxdiag
# Run the main WARP loop
mwxrun.simulation.step(MAX_STEPS)
checkpoint_names = [
f"{CHECKPOINT_NAME}{i:05}"
for i in range(DIAG_STEPS, MAX_STEPS + 1, DIAG_STEPS)
]
for name in checkpoint_names:
print(f"Looking for checkpoint file 'diags/{name}'...")
assert os.path.isdir(os.path.join("diags", name))
assert os.path.isfile("diags/checkpoint00004/fluxdata.ckpt")
def test_restart_from_checkpoint(caplog, force, files_exist):
caplog.set_level(logging.WARNING)
testing_util.initialize_testingdir(
f"test_restart_from_checkpoint_{force}_{files_exist}")
run = get_run()
if force:
restart = True
else:
restart = None
if not files_exist:
prefix = "nonexistent_prefix"
else:
prefix = "checkpoint"
try:
mwxrun.init_run(restart=restart,
checkpoint_dir=os.path.join(testing_util.test_dir,
"checkpoint"),
checkpoint_prefix=prefix)
except RuntimeError as e:
if ("There were no checkpoint directories starting with "
"nonexistent_prefix!" in str(e)):
# There should only be an exception if this restart was forced
assert force
# if the files didn't exist then we didn't restart,
# so there's no need to verify the correct number of steps passed
if files_exist:
new_max_steps = mwxrun.simulation.max_steps
start_step = mwxrun.get_it()
if not force and not files_exist:
log_lines = [r.msg for r in caplog.records]
assert any([
f"There were no checkpoint directories starting with {prefix}!"
in l for l in log_lines
])
mwxrun.simulation.step()
end_step = mwxrun.get_it()
assert end_step - start_step == new_max_steps
def test_checkpoints_fluxdiag():
testing_util.initialize_testingdir("test_checkpoints_fluxdiag")
# use a fixed random seed
np.random.seed(47239475)
run = get_run()
run.init_injectors()
run.init_runinfo()
run.init_fluxdiag()
mwxrun.init_run(
restart=True,
checkpoint_dir=os.path.join(testing_util.test_dir, "checkpoint"),
)
# Run the main WarpX loop
mwxrun.simulation.step()
# load flux diagnostic from a completed run for comparison
basedir = os.path.join(testing_util.test_dir, "checkpoint")
original_flux_file = os.path.join(
testing_util.test_dir, "checkpoint/fluxes/fluxdata_0000000008.dpkl")
original_flux = FluxDiagFromFile(basedir, original_flux_file)
# compare injected flux from the restarted history with the original history
flux_key = ('inject', 'cathode', 'electrons')
for key in ['J', 'P', 'n']:
# if key == 'J':
# continue
old = original_flux.fullhist_dict[flux_key].get_timeseries_by_key(key)
new = run.fluxdiag.fullhist_dict[flux_key].get_timeseries_by_key(key)
print("~~~~~~~~~~~~~~~~~~~~~~~~~~~")
print(f"The key is {key}")
print(f"old: \n {old}")
print(f"new: \n {new}")
assert np.allclose(old, new)
def test_extra_steps_after_restart():
testing_util.initialize_testingdir("test_extra_steps_after_restart")
# use a fixed random seed
np.random.seed(47239475)
run = get_run()
additional_steps = 8
# restart from checkpoint created by test_create_checkpoints
mwxrun.init_run(restart=True,
checkpoint_dir=os.path.join(testing_util.test_dir,
"checkpoint"),
additional_steps=additional_steps)
start_step = mwxrun.get_it()
mwxrun.simulation.step()
end_step = mwxrun.get_it()
assert start_step + additional_steps == end_step
restart_net_charge_density = np.load(
os.path.join(run.field_diag.write_dir,
"Net_charge_density_0000000008.npy"))
# compare against data from test_create_checkpoints
original_net_charge_density = np.load(
os.path.join(testing_util.test_dir, "checkpoint",
"Net_charge_density_0000000008.npy"))
assert np.allclose(restart_net_charge_density,
original_net_charge_density,
rtol=0.1)
def setup_run(self, n=0, steps=None):
"""Setup run for the specific case (n) desired."""
self.steps = steps
# Case specific input parameters
self.ANODE_VOLTAGE = f"{self.VOLTAGE[n]}*sin(2*pi*{self.FREQ:.5e}*t)"
self.N_INERT = self.N_INERT[n]
self.PLASMA_DENSITY = self.PLASMA_DENSITY[n]
self.SEED_NPPC = self.SEED_NPPC[n]
self.NZ = self.NZ[n]
self.DT = self.DT[n]
self.MAX_STEPS = int(self.TOTAL_TIME[n] / self.DT)
self.DIAG_STEPS = int(self.DIAG_INTERVAL / self.DT)
if self.steps is not None:
self.MAX_STEPS = self.steps
self.DIAG_STEPS = self.MAX_STEPS // 5
#######################################################################
# Set geometry, boundary conditions and timestep #
#######################################################################
mwxrun.init_grid(
lower_bound=[-self.D_CA / self.NZ * self.NX / 2.0, 0],
upper_bound=[self.D_CA / self.NZ * self.NX / 2.0, self.D_CA],
number_of_cells=[self.NX, self.NZ],
min_tiles=16)
mwxrun.grid.potential_zmax = self.ANODE_VOLTAGE
mwxrun.init_timestep(DT=self.DT)
mwxrun.simulation.max_steps = self.MAX_STEPS
mwxrun.simulation.load_balance_intervals = self.MAX_STEPS // 5000
#######################################################################
# Field solver #
#######################################################################
# self.solver = picmi.ElectrostaticSolver(
# grid=mwxrun.grid, method='Multigrid', required_precision=1e-12
# )
self.solver = poisson_solvers.PoissonSolverPseudo1D(grid=mwxrun.grid)
mwxrun.simulation.solver = self.solver
#######################################################################
# Particle types setup #
#######################################################################
self.electrons = mespecies.Species(particle_type='electron',
name='electrons')
self.ions = mespecies.Species(particle_type='He',
name='he_ions',
charge='q_e')
#######################################################################
# Collision initialization #
#######################################################################
self.mcc = mcc_wrapper.MCC(self.electrons,
self.ions,
T_INERT=self.T_INERT,
N_INERT=self.N_INERT,
exclude_collisions=['charge_exchange'])
#######################################################################
# Neutral plasma injection #
#######################################################################
self.vol_emitter = emission.UniformDistributionVolumeEmitter(
T=self.T_ELEC)
self.plasma_injector = emission.PlasmaInjector(
emitter=self.vol_emitter,
species1=self.electrons,
species2=self.ions,
npart=2 * self.SEED_NPPC * self.NX * self.NZ,
T_2=self.T_INERT,
plasma_density=self.PLASMA_DENSITY)
#######################################################################
# Add diagnostics #
#######################################################################
self.text_diag = diag_base.TextDiag(diag_steps=self.DIAG_STEPS,
preset_string='perfdebug')
self.field_diag = FieldDiagnostic(diag_steps=self.DIAG_STEPS,
barrier_slices=[0],
save_pdf=False,
style='roelof',
min_dim=2.0)
# array to save ion density
self.rho_array = np.zeros(self.NZ + 1)
#######################################################################
# Initialize run and print diagnostic info #
#######################################################################
mwxrun.init_run()
def test_utils_pulsing_sim():
name = "pulsing_utility_test"
# Include a random run number to allow parallel runs to not collide. Using
# python randint prevents collisions due to numpy rseed below
testing_util.initialize_testingdir(name)
# Initialize each run with consistent, randomly-chosen, rseed. Use a random
# seed instead for initial dataframe generation.
# np.random.seed()
np.random.seed(9216024)
nx = 8
nz = 128
D_CA = 1e-4 # m
xmin = 0.0
zmin = 0.0
xmax = D_CA / nz * nx
zmax = D_CA
max_steps = 50
DT = 1e-10
#####################################
# grid, solver and timesteps
#####################################
mwxrun.init_grid([xmin, zmin], [xmax, zmax], [nx, nz])
solver = PoissonSolverPseudo1D(grid=mwxrun.grid)
mwxrun.simulation.solver = solver
mwxrun.init_timestep(DT=DT)
mwxrun.simulation.max_steps = max_steps
pulse_expr = pulsing.linear_pulse_function(
V_off=-1.0, V_on=7.5, pulse_period=20e-9,
pulse_length=2e-9, t_rise=1e-9, t_fall=1e-9,
wait_time=0.5e-9, plot=True
)
anode = assemblies.Anode(D_CA, pulse_expr, 100, 3.0)
electrons = mespecies.Species(
particle_type='electron',
name='electrons',
)
# mwxrun.simulation.write_input_file()
# exit()
#################################
# simulation setup
################################
mwxrun.init_run()
# check that plot was saved successfully
assert os.path.isfile('diags/circuit/pulse_schematic.png')
################################
# Simulation run
################################
times = np.arange(max_steps)*DT
Vs = np.zeros(max_steps)
# Run the main loop
for ii in range(max_steps):
Vs[ii] = anode.getvoltage()
mwxrun.simulation.step(1)
print(repr(Vs))
ref_Vs = np.array(
[-1. , -1. , -1. , -1. , -1. , -1. , -1. , -0.15, 0.7 ,
1.55, 2.4 , 3.25, 4.1 , 4.95, 5.8 , 6.65, 7.5 , 7.5 ,
7.5 , 7.5 , 7.5 , 7.5 , 7.5 , 7.5 , 7.5 , 7.5 , 7.5 ,
7.5 , 7.5 , 7.5 , 7.5 , 7.5 , 7.5 , 7.5 , 7.5 , 7.5 ,
7.5 , 6.65, 5.8 , 4.95, 4.1 , 3.25, 2.4 , 1.55, 0.7 ,
-0.15, -1. , -1. , -1. , -1. ]
)
assert np.allclose(Vs, ref_Vs, atol=0)
def test_coulomb_scattering():
name = "coulomb_scattering_langevin"
# Include a random run number to allow parallel runs to not collide. Using
# python randint prevents collisions due to numpy rseed below
testing_util.initialize_testingdir(name)
# Initialize each run with consistent, randomly-chosen, rseed. Use a random
# seed instead for initial dataframe generation.
# np.random.seed()
np.random.seed(92160000)
class BeamEmitter(emission.UniformDistributionVolumeEmitter):
"""Child class of UniformDistributionVolumeEmitter used to overwrite the
get_velocities function when sampling velocities for the electrons since
we want to inject them with a delta function."""
def __init__(self, T, v_beam):
self.v_beam = v_beam
super(BeamEmitter, self).__init__(T=T)
def _get_xv_coords(self, npart, m, rseed):
"""Get velocities and call specialized function for position."""
x_coords = self._get_x_coords(npart)
if m == 1:
# ion velocities can be Maxwellian
v_coords = mwxutil.get_velocities(
npart, self.T, m=m, emission_type='random', rseed=None)
else:
# electron velocities should be a delta function, we need to
# have small non-zero values for x and y otherwise the Coulomb
# scattering algorithm fails due to divide by zero errors
v_coords = np.zeros((3, npart))
v_coords[0] = np.random.normal(0.0, 1.0, npart)
v_coords[1] = np.random.normal(0.0, 1.0, npart)
v_coords[2] = self.v_beam
return (
x_coords[:, 0], x_coords[:, 1], x_coords[:, 2],
v_coords[0], v_coords[1], v_coords[2]
)
class VarianceTracker(diags.WarpXDiagnostic):
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 _record_variance(self):
if not self.check_timestep():
return
ux = np.concatenate(mwxrun.sim_ext.get_particle_ux('electrons'))
self.results_array[self.i, 0] = mwxrun.get_t()
self.results_array[self.i, 1] = np.std(ux)
self.i += 1
#######################################################################
# Simulation parameters #
#######################################################################
ZMAX = 250e-6
NZ = 256
NX = 8
DT = 1e-13
MAX_STEPS = 150
DIAG_STEPS = 15
SEED_DENSITY = 5e17
SEED_COUNT = 20000
LOGLAMBDA = 7.5
VBEAM = 1e6
# Derived quantities
PERIOD = ZMAX / NZ * NX
#######################################################################
# Set geometry, boundary conditions and timestep #
#######################################################################
mwxrun.init_grid(
lower_bound=[-PERIOD/2.0, 0.], upper_bound=[PERIOD/2.0, ZMAX],
number_of_cells=[NX, NZ],
bc_fields_z_min='periodic',
bc_fields_z_max='periodic',
bc_particles_z_min='periodic',
bc_particles_z_max='periodic'
)
mwxrun.init_timestep(DT=DT)
mwxrun.simulation.max_steps = MAX_STEPS
#######################################################################
# Dummy Poisson solver to effectively turn off the field solve #
#######################################################################
solver = DummyPoissonSolver(grid=mwxrun.grid)
mwxrun.simulation.solver = solver
#######################################################################
# Particle types setup #
#######################################################################
electrons = mespecies.Species(
particle_type='electron', name='electrons'
)
# artificially increase ion mass to match Lorentz gas better
ions = mespecies.Species(
particle_type='Xe', name='xe_ions', charge='q_e', mass=1.0
)
#######################################################################
# Seed simulation with quasi-neutral plasma #
#######################################################################
volume_emitter = BeamEmitter(1.0, VBEAM)
emission.PlasmaInjector(
volume_emitter, electrons, ions, npart=SEED_COUNT*2,
plasma_density=SEED_DENSITY
)
#######################################################################
# Coulomb scattering #
#######################################################################
langevin = coulomb_scattering.LangevinElectronIonScattering(
electron_species=electrons, ion_species=ions,
log_lambda=LOGLAMBDA, subcycling_steps=1
)
#######################################################################
# Diagnostics #
#######################################################################
diags.TextDiag(MAX_STEPS // 5, preset_string='perfdebug')
variance_tracker = VarianceTracker(MAX_STEPS, DIAG_STEPS)
#######################################################################
# Run simulation #
#######################################################################
mwxrun.init_run()
mwxrun.simulation.step()
#######################################################################
# Compare results to theory #
#######################################################################
D = (
SEED_DENSITY * mwxconstants.e**4 * LOGLAMBDA
/ (4 * np.pi * mwxconstants.epsilon_0**2 * mwxconstants.m_e**2 * VBEAM)
)
times = variance_tracker.results_array[:, 0]
calculated_values = variance_tracker.results_array[:, 1] / VBEAM
expected_values = np.sqrt(D*times) / VBEAM
print("Calculated ratio: ", calculated_values )
print("Expected ratio: ", expected_values)
assert np.allclose(calculated_values, expected_values, rtol=0.01)
'''
def init_warpx(self):
logger.info("### Init Simulation Run ###")
mwxrun.init_run()
def setup_run(self):
self.rmax = 1e-3
self.rmin = self.rmax - self.D_CA - self.USE_EB * 10e-6
self.zmax = (self.rmax - self.rmin) / self.NR * self.NZ / 2.0
self.zmin = -self.zmax
self.V_ANODE = self.V_CATHODE + self.V_ANODE_CATHODE
if self.DIAG_STEPS is None:
self.DIAG_STEPS = ((self.TOTAL_TIMESTEPS //
5) if self.TOTAL_TIMESTEPS > 10 else
self.TOTAL_TIMESTEPS)
#######################################################################
# Set geometry, boundary conditions and timestep #
#######################################################################
mwxrun.init_grid(
lower_bound=[self.rmin, self.zmin],
upper_bound=[self.rmax, self.zmax],
number_of_cells=[self.NR, self.NZ], # min_tiles=1,
use_rz=True,
bc_fields_z_min='periodic',
bc_fields_z_max='periodic',
bc_particles_z_min='periodic',
bc_particles_z_max='periodic',
bc_fields_r_min='dirichlet',
bc_fields_r_max='dirichlet',
)
mwxrun.init_timestep(DT=self.DT)
mwxrun.simulation.max_steps = self.TOTAL_TIMESTEPS
mwxrun.simulation.load_balance_intervals = self.TOTAL_TIMESTEPS // 5000
#######################################################################
# Field solver #
#######################################################################
self.solver = picmi.ElectrostaticSolver(grid=mwxrun.grid,
method='Multigrid',
required_precision=1e-6,
warpx_self_fields_verbosity=0,
maximum_iterations=1000)
mwxrun.simulation.solver = self.solver
#######################################################################
# Conductors setup and installation #
#######################################################################
self.cathode = assemblies.CylinderZ(V=self.V_CATHODE,
T=self.CATHODE_TEMP,
WF=self.CATHODE_PHI,
name='cathode',
r_outer=self.rmax + 1e-6,
r_inner=self.rmax)
self.anode = assemblies.CylinderZ(
V=self.V_ANODE,
T=self.ANODE_TEMP,
WF=self.ANODE_PHI,
name='anode',
r_outer=(self.rmin if not self.USE_EB else self.rmax - self.D_CA))
#######################################################################
# Particle types setup #
#######################################################################
self.electrons = mespecies.Species(particle_type='electron',
name='electrons')
#######################################################################
# Scraper setup #
#######################################################################
for species in [self.electrons]:
species.save_particles_at_xhi = 1
if self.USE_EB:
species.save_particles_at_eb = 1
else:
species.save_particles_at_xlo = 1
#######################################################################
# Thermionic emission #
#######################################################################
self.cathode_emitter = emission.ZCylinderEmitter(
conductor=self.cathode, rdir=-1)
self.cathode_injector = emission.ThermionicInjector(
emitter=self.cathode_emitter,
species=self.electrons,
A=self.CATHODE_A,
npart_per_cellstep=self.NPPC,
use_Schottky=self.USE_SCHOTTKY)
#######################################################################
# Clean up and output RunInfo #
#######################################################################
runvars = CylinderVacuumTEC.__dict__.copy()
runvars.update(self.__dict__)
self.runinfo = runinfo.RunInfo(
local_vars=runvars,
electrode_params={
"CATHODE_A": self.CATHODE_A,
"CATHODE_TEMP": self.CATHODE_TEMP,
"CATHODE_PHI": self.CATHODE_PHI,
"ANODE_PHI": self.ANODE_PHI
},
surface_dict={
'cathode': self.cathode,
'anode': self.anode
},
injector_dict={'cathode': self.cathode_injector},
)
# overwrite runinfo area to properly normalize currents
self.runinfo.area = self.cathode_emitter.area
#######################################################################
# Add diagnostics #
#######################################################################
self.text_diag = diags.TextDiag(diag_steps=self.DIAG_STEPS,
preset_string='perfdebug')
# self.field_diag = diags.FieldDiagnostic(
# diag_steps=self.DIAG_STEPS, style='roelof', save_pdf=False
# )
self.fluxdiag = diags.FluxDiagnostic(diag_steps=self.DIAG_STEPS,
runinfo=self.runinfo,
check_charge_conservation=False,
overwrite=True)
#######################################################################
# Initialize run and print diagnostic info #
#######################################################################
mwxrun.init_run()