Change parameters midway through simulation

[SamuelMcloughlin]

Hello,

Is it possible to run a portion of a simulation, pause the simulation, then change some values i.e. the plasma density, and then continue the simulation from where it was paused with the changes implemented?

Thanks a million!

[RemiLehe]

Hi,

Yes, it is possible in principle, but whether it will work properly depends a lot on what you are modeling, and on the details of the Python code that implements the change of value.

One way to do this is by calling step multiple times.

sim.step(3000)
<Python code that modifies some variables>
sim.step(2000)

In this case, make sure that the Python code is modifying a variable that is actually used by the simulation, and not some intermediate variables.

Another way to do this would be to use checkpoints and restarts (and thus modify the value of e.g. the density in the script for the restarted simulation)

Finally, note that, if you are running a simulation with a moving window, you can effectively change the value of the density during the simulation by providing a function that varies on z.

[SamuelMcloughlin]

Hello,

I had tried that first method before, but found that only the information from the first step was being saved.

[RemiLehe]

OK, thanks for letting me know.
As I was mentioning, whether this will work depends largely on the details of the Python code that implements the change of value. Would you be willing to share your Python script? Maybe we can find a way to make it work.

[SamuelMcloughlin]

Of course, here is the script

import numpy as np                                         # Array creation & manipulation

from scipy.constants import c, e, m_e, m_p, pi, epsilon_0  # Physical constants
import os
import shutil
from scipy.signal import find_peaks


from openpmd_viewer.addons import LpaDiagnostics          # OpenPMD Viewer (add-on for plasma acceleration)
import matplotlib.pyplot as plt                           # Plotting

from fbpic.main import Simulation
from fbpic.lpa_utils.laser import add_laser_pulse
from fbpic.lpa_utils.laser.laser_profiles import GaussianLaser
from fbpic.lpa_utils.boosted_frame import BoostConverter
from fbpic.openpmd_diag import BackTransformedFieldDiagnostic, BackTransformedParticleDiagnostic, ParticleChargeDensityDiagnostic
# Note: You can ignore the MPI error since the simulation is running on a single GPU only.

#SIMULATION SETUP

# Define the Lorentz factor of the Boosted frame
gamma_boost = 10.
# Create a boosted frame converter
boost = BoostConverter(gamma_boost)

# Simulation box
Nz = 3600         # Number of gridpoints along z
zmax = 20.e-6    # Right boundary of the box along z (meters)
zmin = -40.e-6   # Left boundary of the box along z (meters)
Nr = 50          # Number of gridpoints along r
rmax = 30.e-6    # Length of the box along r (meters)
Nm = 2           # Number of azimuthal modes used

# Simulation timestep
dt = min( rmax/(2*boost.gamma0*Nr)/c, (zmax-zmin)/Nz/c )  # Timestep (seconds)

# Simulation parameters
sim_params = {
    'Nz': Nz, 'zmin': zmin, 'zmax': zmax,            # Box parameters
    'Nr': Nr, 'rmax': rmax,                          # Box parameters
    'Nm': Nm, 'dt': dt,                              # Box parameters
    'boundaries': {'z':'open', 'r':'open'},          # Boundary conditions
    'particle_shape': 'cubic',                       # Marcoparticle shape factor
    'gamma_boost': boost.gamma0,                     # Lorentz factor of the Boosted frame
    'v_comoving': -c*np.sqrt(1.-1./boost.gamma0**2), # Velocity of the Galilean frame (NCI suppression)
    'use_cuda': True,                                # Simulate on CUDA-enabled GPU
}





#SIMULATION

# Initialize the simulation
sim = Simulation(**sim_params)

#Moving Window
# Moving window speed
v_window = c
# The interaction length of the simulation, in the lab frame (meters)
L_interact = 50.e-6
# Interaction time, in the lab frame (seconds)
T_interact = (L_interact + (zmax-zmin)) / v_window 
# Add the moving window to the simulation
sim.set_moving_window( v = boost.velocity([v_window])[0] )


#LASER
# Laser parameters
laser_params = {
    'a0': 2,          # Laser amplitude 
    'waist': 10.e-6,    # Laser waist (meters)
    'tau': 12.e-15,     # Laser duration (seconds)
    'z0': -8.e-6,       # Laser centroid (meters)
    'zf': 250.e-6,      # Focal position (meters)
    'lambda0': 800.e-9  # Laser wavelength (meters)
}

# Create the laser profile
laser_profile = GaussianLaser(**laser_params)

# Add the laser to the simulation
add_laser_pulse( sim, laser_profile, gamma_boost=boost.gamma0, method='antenna', z0_antenna=0 )




#GAS AND PLASMA DENSITY

# Gas density
n0 = 6.e24 # Gas atoms per unit volume [1/m^3]

# Gas profile
ramp_up = 100.e-6    # Length of density up-ramp (rising half-cosine)
plateau = 1900.e-6   # Length of density plateau
ramp_down = 100.e-6  # Length of density down-ramp (falling half-cosine)

# Gas density (Hydrogen gas is assumed)
def hydrogen_density(z, r):
    # Conditional expression describing the profile
    n = np.ones_like(z)
    n=np.where(z<ramp_up, n0*0.5*(1-np.cos(pi*(z)/ramp_up)), n)
    n=np.where(z>=ramp_up, n0 , n)
    return(n)

# Dopant gas density fraction
f = 1e-10

# Dopant gas density profile (Nitrogen gas is assumed)
def nitrogen_density(z, r):
    # Conditional expression describing the profile
    n = np.ones_like(z)
    n= np.where(z<ramp_up, f*n0*0.5*(1-np.cos(pi*(z)/ramp_up)) , n)
    n= np.where(z<ramp_up+ramp_down, f*n0*0.5*(1+np.cos(pi*(z-(ramp_up))/ramp_down)), n )
    return n

# Plasma electron density (assuming full gas ionization)
def electron_density(z, r):
    # Inferred from the preionized Hydrogen & Nitrogen (only up to 5th level) gas profile
    n = hydrogen_density(z, r) + 5 * nitrogen_density(z, r) # 1 e- per H-ion & 5 e- per N-ion
    return n

#PARTICLES

# Macroparticle parameters
part_params = {
    'p_zmin': 0.e-6,   # Start of particles (meters)
    'p_zmax': 2100.e-6, # End of particles (meters)
    'p_rmax': 30.e-6,  # Radial limit of particles (meters)
    'p_nz': 1,         # Number of particles per cell along z
    'p_nr': 2,         # Number of particles per cell along r
    'p_nt': 6          # Number of particles per cell along theta
} 

# Clear all species in the simulation
sim.ptcl = []

# Add the nitrogen ions (assuming pre-ionization up to the 5th level)
nitrogen_ions = sim.add_new_species( q=5*e, m=14.*m_p, n=1, dens_func=nitrogen_density, p_zmin=0e-6, p_zmax=2100e-6, p_rmax=30e-6, p_nz=1, p_nr=2, p_nt=6, boost_positions_in_dens_func=True )

# Add the plasma electrons
electrons = sim.add_new_species( q=-e, m=m_e, n=1, dens_func=electron_density,p_zmin=0e-6, p_zmax=2100e-6, p_rmax=30e-6, p_nz=1, p_nr=2, p_nt=6, boost_positions_in_dens_func=True )

# Add the hydrogen ions
hydrogen_ions = sim.add_new_species( q=e, m=m_p, n=1, dens_func=hydrogen_density, p_zmin=0e-6, p_zmax=2100e-6, p_rmax=30e-6, p_nz=1, p_nr=2, p_nt=6, boost_positions_in_dens_func=True )



# Add a second electron species for storing the ionized electrons separately
ionized_electrons = sim.add_new_species( q=-e, m=m_e )

# Make Nitrogen atoms ionizable (assuming pre-ionization up to the 5th level)
nitrogen_ions.make_ionizable( 'N', target_species=ionized_electrons, level_start=5 )



#DIAGNOSTICS


# Number of discrete diagnostic snapshots, back-transformed from the boosted simulation frame to the lab frame
N_diag = 100+1
# Time interval between diagnostic snapshots *in the lab frame* 
# (first at t=0, the last at t=T_interact)
dt_diag_period = T_interact / (N_diag - 1) 

# Field diagnostics (back-transformed to lab frame)
field_diag = BackTransformedFieldDiagnostic( zmin, zmax, v_window, dt_diag_period, N_diag, boost.gamma0,
                                             period=50, fldobject=sim.fld, comm=sim.comm, 
                                             fieldtypes=['E','rho'] )
# Add the diagnostic to the simulation
sim.diags.append( field_diag )

# Particle diagnostics (back-transformed to lab frame)
particle_diag = BackTransformedParticleDiagnostic( zmin, zmax, v_window, dt_diag_period, N_diag, boost.gamma0,
                                                   period=50, fldobject=sim.fld, comm=sim.comm,
                                                   species={'bunch': ionized_electrons, 'electrons': electrons} )
            

# Add the diagnostic to the simulation
sim.diags.append( particle_diag )


#PIC Loop


# Number of PIC loop iterations (Interaction time in the Boosted frame divided by the timestep)
N_step = int( boost.interaction_time( L_interact, (zmax-zmin), v_window ) / sim.dt )


L_start=200e-6
N_step1 = int( boost.interaction_time( L_start, (zmax-zmin), v_window ) / sim.dt )


sim.step(N_step1)
for i in range(30):
    data = LpaDiagnostics('./lab_diags/hdf5',check_all_files=True)

    Ez_1D, info = data.get_field('E', 'z', iteration=i+1, slice_across='r')
    troughs, _= find_peaks(-Ez_1D, height=-0.5*Ez_1D.min(), distance=150)
    z_min=info.z[troughs[-1]]

    z, uz = data.get_particle(var_list=['z','uz'], species='bunch', iteration=i+1)
    topuz=np.argpartition(uz,-5)[-5:]
    bunch_uz_z=np.mean(z[topuz])

    delta_uz_z=bunch_uz_z-z_min
 
    density_factor=1.125

    def dens2_func(z,r):
        return density_factor*hydrogen_density(z,r)+5*nitrogen_density(z,r)

    sim.ptcl[1].injector.dens_func = dens2_func
    sim.step(N_step)


data = LpaDiagnostics('./lab_diags/hdf5',check_all_files=True)

z_prop = data.t*c

# Get the mean gamma along the propagation
mean_gamma, std_gamma = data.iterate( data.get_mean_gamma, species='bunch' )

laser_energy_list=[]
for i in range(len(z_prop)):
    Er,Er_info = data.get_field(field='E',coord='r', iteration=i,m='all')
    r = Er_info.r
    z = Er_info.z
    dr = np.mean(np.diff(r))
    dz = np.mean(np.diff(z))
    laser_energy = epsilon_0*np.sum(np.abs(r).reshape(-1,1)*Er**2)*dr*dz
    laser_energy_list.append(laser_energy)

# Plot the mean gamma and energy of the laser
plt.figure(figsize=(5,4), dpi=150)
lns1=plt.plot(z_prop*1.e6, mean_gamma, color='blue', label='Lorentz Factor')
plt.xlabel('$z_{prop}$ ($\mu$m)')
plt.ylabel('$\gamma$')
plt.twinx()
lns2=plt.plot(z_prop*1e6, laser_energy_list, color='cyan', label='Laser Energy')
plt.xlabel('$z_{prop}$ ($\mu$m)')
plt.ylabel('$E$ (J)')
# To Create Legend for Both
lns = lns1+lns2
labs = [l.get_label() for l in lns]
plt.legend(lns, labs)
plt.savefig(f'Mean Gamma and Laser Energy.png')