From d01b6cb982d7fe755d9315bbc1fb83e36978fbce Mon Sep 17 00:00:00 2001
From: Nadezhda Khachatrian <nadezhda.khachatrian@studium.uni-hamburg.de>
Date: Thu, 26 Oct 2023 10:41:59 +0200
Subject: [PATCH] examples of the lasy beam propagation including resampling of
 the beam while focusing it; focused beam is also propagated in fbpic, both
 beams are compared

---
 lasy/examples/FBPICbeam.py   | 103 +++++++++++++++++++++++
 lasy/examples/LASYbeam.py    | 159 +++++++++++++++++++++++++++++++++++
 lasy/examples/LASYvsFBPIC.py | 106 +++++++++++++++++++++++
 3 files changed, 368 insertions(+)
 create mode 100644 lasy/examples/FBPICbeam.py
 create mode 100644 lasy/examples/LASYbeam.py
 create mode 100644 lasy/examples/LASYvsFBPIC.py

diff --git a/lasy/examples/FBPICbeam.py b/lasy/examples/FBPICbeam.py
new file mode 100644
index 00000000..cd35c72b
--- /dev/null
+++ b/lasy/examples/FBPICbeam.py
@@ -0,0 +1,103 @@
+import numpy as np
+from scipy.constants import c, e, m_e, m_p, k
+# Import the relevant structures from fbpic
+from fbpic.main import Simulation
+from fbpic.lpa_utils.laser import add_laser, add_laser_pulse
+from fbpic.lpa_utils.laser.laser_profiles import FlattenedGaussianLaser, GaussianLaser
+from fbpic.openpmd_diag import FieldDiagnostic, ParticleDiagnostic, \
+     set_periodic_checkpoint, restart_from_checkpoint, BoostedFieldDiagnostic, BoostedParticleDiagnostic
+from fbpic.lpa_utils.boosted_frame import BoostConverter
+from fbpic.lpa_utils.laser import add_laser_pulse, FromLasyFileLaser
+
+import os
+from mpi4py.MPI import COMM_WORLD as comm
+
+
+# Whether to use the GPU
+use_cuda = True
+
+# The simulation box
+Nz = 6000 #4000 #3200        # Number of gridpoints along z
+zmax = 0.e-6     # Right end of the simulation box (meters)
+#equal to the size of LASY box for the laser to be centered 
+zmin = -2*90.e-15*c   # Left end of the simulation box (meters)
+Nr = 500 #405 #270         # Number of gridpoints along r
+rmax = 200.e-6    # Length of the box along r (meters)
+Nm = 2 # Number of modes used
+# Boost factor
+gamma_boost = 5.
+# Maximum simulation length
+Lmax = 5.e-3 #LASY propagation distance of a pulse 
+# The simulation timestep
+dt =  min( rmax/(2*gamma_boost*Nr), (zmax-zmin)/Nz/c )  # Timestep (seconds)
+
+n_order = 16
+
+boost = BoostConverter(gamma_boost)
+
+
+
+# The moving window
+v_window = c
+
+# Velocity of the Galilean frame (for suppression of the NCI)
+v_comoving = -np.sqrt(gamma_boost**2-1.)/gamma_boost * c
+
+# The diagnostics
+diag_period = 500 # Period of the diagnostics in number of timesteps
+
+# Whether to write the fields in the lab frame
+Ntot_snapshot_lab = 40 
+dt_snapshot_lab = (Lmax) / v_window / (Ntot_snapshot_lab - 1)
+
+track_bunch = False  # Whether to tag and track the particles of the bunch
+
+# The interaction length of the simulation (meters)
+L_interact = Lmax # the plasma length
+# Interaction time (seconds) (to calculate number of PIC iterations)
+T_interact = boost.interaction_time( L_interact, (zmax-zmin), v_window )
+# (i.e. the time it takes for the moving window to slide across the plasma)
+
+#take the first file
+PathToTheLasyFile = './DiagLASYforFBPIC/laser00000001_00001.h5' 
+laser_profile = FromLasyFileLaser(PathToTheLasyFile) 
+
+
+# ---------------------------
+# Carrying out the simulation
+# ---------------------------
+
+if __name__ == '__main__':
+
+    # Initialize the simulation object
+    sim = Simulation( Nz, zmax, Nr, rmax, Nm, dt,
+        zmin=zmin, boundaries={'z':'open', 'r':'open'}, initialize_ions=False,
+        n_order=n_order, use_cuda=use_cuda, v_comoving=v_comoving,
+        gamma_boost=gamma_boost, verbose_level=2, particle_shape='cubic',
+        use_galilean=True)
+    # By default the simulation initializes an electron species (sim.ptcl[0])
+    # Because we did not pass the arguments `n`, `p_nz`, `p_nr`, `p_nz`,
+    # this electron species does not contain any macroparticles.
+    # It is okay to just remove it from the list of species.
+    sim.ptcl = []
+
+
+    add_laser_pulse( sim, laser_profile, gamma_boost=gamma_boost,
+                method='antenna', z0_antenna=0.e-6, v_antenna=0.)
+    # Convert parameter to boosted frame
+    v_window, = boost.velocity( [ v_window ] )
+    # Configure the moving window
+    sim.set_moving_window( v=v_window )
+
+    # Add a diagnostics
+    write_dir = 'DiagsFBPIC'
+    sim.diags = [
+                 BoostedFieldDiagnostic( zmin, zmax, c,
+                    dt_snapshot_lab, Ntot_snapshot_lab, gamma_boost,
+                    period=diag_period, fldobject=sim.fld, comm=sim.comm,
+                    fieldtypes=["E", "B", "rho"], write_dir=write_dir)
+                    ]
+
+    N_step = int(T_interact/sim.dt)
+    ### Run the simulation
+    sim.step( N_step )
\ No newline at end of file
diff --git a/lasy/examples/LASYbeam.py b/lasy/examples/LASYbeam.py
new file mode 100644
index 00000000..6727eddb
--- /dev/null
+++ b/lasy/examples/LASYbeam.py
@@ -0,0 +1,159 @@
+# Lasy Imports
+from lasy.profiles.gaussian_profile import GaussianProfile
+from lasy.laser import Laser, Grid
+from lasy.utils.zernike import zernike
+
+# OpenPMD Viewer (add-on for plasma acceleration)
+from openpmd_viewer.addons.pic import LpaDiagnostics
+
+# Standard Imports
+from scipy.constants import c
+import matplotlib.pyplot as plt
+import numpy as np
+import copy
+import os, shutil
+
+
+def add_parabola(laser, f0):
+    dt = laser.grid.dx[-1]
+    omega0 = laser.profile.omega0
+    Nt = laser.grid.field.shape[-1]
+    omega = 2 * np.pi * np.fft.fftfreq(Nt, dt) + omega0
+    
+    #on-axis parabola
+    
+    r = laser.grid.axes[0]
+    #transform the field from temporal to frequency domain
+    field_fft = np.fft.ifft(laser.grid.field, axis=-1, norm="backward") 
+     
+    #multiply spectral image by a parabolic mirror phase
+    field_fft *= mirror_parabolic(f0, omega/c, r).T   
+    
+    #rewrite laser field
+    laser.grid.field[:, :, :] = np.fft.fft(field_fft, axis=-1, norm="backward"
+        )
+
+def mirror_parabolic(f0, kz, r):
+    """
+    Generate array with spectra-radial phase
+    representing the on-axis Parabolic Mirror
+    
+    f0: focal length of a parabolic mirror
+    kz = omega/c
+    r: radial axis 
+    """
+   
+    s_ax = r**2/4/f0
+
+    val = np.exp(   -2j * s_ax[None,:] * \
+                 ( kz * np.ones((*kz.shape, *r.shape)).T ).T)
+    return val
+
+
+#laser parameters
+
+#laser wavelength (m)
+wavelength = 800*1.e-9
+#polarization vector
+pol = (1, 0)
+#laser energy (J)
+laser_energy = 1
+#laser waist (m)
+w0 = 20.e-3
+#pulse duration (s)
+tau = 30e-15
+#maximum intensity time 
+t_peak = 0
+
+pulseProfile = GaussianProfile(wavelength, pol, laser_energy, w0, tau, t_peak, z_foc=0.e-3)
+
+# Create initial grid for the pulse
+Nr = 3000
+R_max = 60.e-3
+
+dim = "rt"
+lo = (0e-6, -90e-15) 
+hi = (R_max, 90e-15)  
+
+npoints = (Nr, 500)
+
+laser = Laser(dim=dim, lo=lo, hi=hi, npoints=npoints, profile=pulseProfile)
+
+#add sperical abberations 
+R = np.linspace(0, R_max, Nr)
+#(n, m) = (2, 0)
+jthree = 4
+#(n, m) = (5, 0)
+jfive = 12
+
+pupilRadius = 4 * w0
+pupilCoord = (0, 0, pupilRadius)
+
+third = zernike(0, R, pupilCoord, jthree)
+fifth = zernike(0, R, pupilCoord, jfive)
+
+coeff1 = -0.1 #defocus
+coeff2 = 0.9 #primary spherical
+
+phase = coeff1*third + coeff2*fifth
+
+phase2D = np.repeat(phase[:, np.newaxis], npoints[1], axis=1)
+laser.grid.field = laser.grid.field * np.exp(1j * phase2D)
+
+#propagate to OAP
+laser.propagate(40)
+
+#add parabolic mirror
+add_parabola(laser, 2000.e-3)
+
+#define a new grid for a focused beam 
+new_grid = Grid(dim, lo, (200.e-6, 90e-15), npoints, n_azimuthal_modes=laser.grid.n_azimuthal_modes)    
+laser.propagate((1995.e-3), grid=new_grid, nr_boundary = 128, backend = 'NP') #resample the grid
+
+laser_fbpic = copy.deepcopy(laser)
+
+#propagate laser by a box size to align LASY and FBPIC
+#laser pulse in FBPIC is emitted first and then propagated 
+#in FBPIC laser is emitted by antenna and consequently it is fully iniallized after a full box size
+#initiallised laser should be at the same position as in LASY after propagation over one box
+!rm -rf './DiagLASYforFBPIC'
+path = './DiagLASYforFBPIC'
+os.mkdir(path)
+os.chdir(path)
+
+laser_fbpic.propagate(2*90e-15*c)
+file_format ='h5'
+filename = "laser%08d"%1
+laser_fbpic.write_to_file(0, 1995.e-3 , filename, file_format)
+
+os.chdir('../')
+
+
+#propagating in a small grid to a focal point, step after step
+propDist = 1995.e-3 
+
+N = 40
+prop_dz = (5.e-3)/(N-1)
+
+!rm -rf './DiagsLASY'
+path = './DiagsLASY'
+os.mkdir(path)
+os.chdir(path)
+
+
+file_format ='h5'
+filename = "laser%08d"%0
+laser.write_to_file(0, propDist, filename, file_format) #writing propagation from 0 to 1995e-3 
+
+#write a file for separate iteration
+for j in range(N-1):
+    laser.propagate(prop_dz)
+    propDist += prop_dz
+    
+    i = j+1
+    file_format ='h5'
+    filename = "laser%08d"%i
+    currentit = i
+    laser.write_to_file(currentit, propDist, filename, file_format)
+    
+os.chdir('../')
\ No newline at end of file
diff --git a/lasy/examples/LASYvsFBPIC.py b/lasy/examples/LASYvsFBPIC.py
new file mode 100644
index 00000000..26f4dc5c
--- /dev/null
+++ b/lasy/examples/LASYvsFBPIC.py
@@ -0,0 +1,106 @@
+#Analyse LASYvsFBPIC
+#compares propagation of two laser beams 
+
+# OpenPmd (for reading information from file)
+from openpmd_viewer.addons.pic import LpaDiagnostics
+#from openpmd_viewer.addons import LpaDiagnostics          # OpenPMD Viewer (add-on for plasma acceleration)
+
+# Standard Imports
+from scipy.constants import c
+import matplotlib.pyplot as plt
+import numpy as np
+import copy
+import mpl_style
+mpl_style.load_preset()
+
+
+#FBPIC
+#data corresponding to FBPIC laser 
+path_fbpic = './DiagsFBPICtwo/hdf5'
+data_fbpic = LpaDiagnostics(path_fbpic, check_all_files=True, backend="h5py")
+
+
+#all of the values correspond to the intensity 
+num_iteration = 40
+
+laser_fbpic, info_fbpic = data_fbpic.get_laser_envelope(iteration = 1 ,pol = 'x',slice_across='z', plot=False)
+
+#create an array containing propagation of the laser waist (intensity)
+array2d_fbpic = np.zeros((laser_fbpic.shape[0], num_iteration))
+
+for iteration in range(num_iteration):
+    laser_fbpic, info_fbpic = data_fbpic.get_laser_envelope(iteration = iteration ,pol = 'x', slice_across='z',       plot=False)
+    array2d_fbpic[:, iteration] = np.abs(laser_fbpic)**2
+
+#plt.imshow(array2d_fbpic[:,:], extent=[1995, 2000, 0, 200], aspect='auto')
+#plt.xlabel('$z$ (mm)')
+#plt.ylabel('$r$ (µm)')
+#plt.colorbar()
+#plt.savefig('FBPIC.pdf')
+
+
+
+#LASY
+#data corresponding to the LASY laser
+path = './DiagLASY'
+data_lasy = LpaDiagnostics(path, check_all_files=True, backend="h5py")
+
+#lasy field has a shape of (6000, 40), to compare it with fbpic we need to reduce it to (1000, 40), manually pick #out appropriate entries for lasy (to avoid duplication of entries)
+def rewrite_lasy(laser_field, info_lasy, info_fbpic):
+    array_half_size = info_lasy.r.size//2 #size of the lasy array
+    divide_factor = info_lasy.r.size // info_fbpic.r.size #match sizes of lasy and fbpic
+
+    r_second_half = info_lasy.r[array_half_size:][::divide_factor] 
+    r_first_half = info_lasy.r[::-1][array_half_size:][::divide_factor][::-1]
+    r_new = np.concatenate([r_first_half,r_second_half])
+
+    field_second_half = laser_field[array_half_size:][::divide_factor]
+    field_first_half = laser_field[::-1][array_half_size:][::divide_factor][::-1]
+    field_new = np.concatenate([field_first_half,field_second_half])
+    return(field_new)
+
+
+#create a lasy field array of appropriate size
+lasy_iterations = []
+for iteration in range(0, num_iteration):
+    laser_fbpic, info_fbpic = data_fbpic.get_laser_envelope(iteration = iteration ,pol = 'x', slice_across='z',       plot=False)
+    laser_field,laser_field_info = data_lasy.get_field(iteration = iteration ,field='laserEnvelope',                   slice_across='t', plot=False)
+    #write a separate laser field for each iteration
+    laser_lasy = rewrite_lasy(laser_field, laser_field_info, info_fbpic)
+    lasy_iterations.append(laser_lasy)
+    
+#propagation of the LASY laser waist 
+array2d_lasy = np.zeros((lasy_iterations[0].shape[0], num_iteration))
+
+for iteration in range(num_iteration):
+    array2d_lasy[:, iteration] = np.abs(np.array(lasy_iterations)[iteration, :])**2    
+
+#plt.imshow(array2d_lasy[:,:], extent=[1995, 2000, 0, 200])
+#plt.xlabel('$z$ (mm)')
+#plt.ylabel('$r$ (µm)')
+#plt.colorbar()
+#plt.savefig('LASY.pdf')
+
+
+#LASY vs FBPIC
+#0 iteration is irrelevant since in fbpic the laser is getting emmited by antenna 
+plt.imshow(array2d_lasy[:,1:] - array2d_fbpic[:,1:], extent=[1995.000128, 2000, 0, 200])
+plt.colorbar()
+plt.xlabel('$z$ (mm)')
+plt.ylabel('$r$ (µm)')
+plt.title("LASY minus FBPIC")
+plt.savefig('LASYvsFBPIC.pdf')
+
+#slice comparison
+laser_fbpic, info_fbpic = data_fbpic.get_laser_envelope(iteration = 25 ,pol = 'x',slice_across='z', plot=False)
+plt.plot(np.abs(laser_fbpic)**2, 'b', label='fbpic')
+plt.plot(np.abs(np.array(lasy_iterations)[25,:])**2, 'y', label='lasy')
+plt.legend(loc='best')
+
+#max intensities at each iteration
+plt.plot(array2d_lasy[:, 1:].max(axis=0), 'y', label='lasy')
+plt.plot(array2d_fbpic[:,1:].max(axis=0), 'b', label='fbpic')
+plt.xlabel('$z$ (mm)')
+plt.ylabel('intensity')
+plt.legend(loc='best')
+#0.000128 mm to 5mm
\ No newline at end of file