Change parameters midway through simulation about fbpic HOT 4 OPEN

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

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')