annotation of 1D FEL code
1.function sase in sase1d.py
import numpy as np
import scipy
import load_bucket
import matplotlib.pyplot as plt
def sase(inp_struct):
'''
SASE 1D FEL run function
Input:
Nruns # Number of runs
npart # n-macro-particles per bucket
s_steps # n-sample points along bunch length
z_steps # n-sample points along undulator
energy # electron energy [MeV]
eSpread # relative rms energy spread [ ]
emitN # normalized transverse emittance [mm-mrad]
currentMax # peak current [Ampere]
beta # mean beta [meter]
unduPeriod # undulator period [meter]
unduK # undulator parameter, K [ ]
unduL # length of undulator [meter]
radWavelength # seed wavelength? [meter], used only in single-freuqency runs
dEdz # rate of relative energy gain or taper [keV/m], optimal~130
iopt # 5=SASE, 4=seeded
P0 # small seed input power [W]
constseed # whether we want to use constant random seed for reproducibility, 1 Yes, 0 No
Output:
z # longitudinal steps along undulator
power_z # power profile along undulator
s # longitudinal steps along beam
power_s # power profile along beam
rho # FEL Pierce parameter
detune # deviation from the central energy
field # final output field along beam
field_s # output field along beam for different z position
gainLength # 1D FEL gain Length
resWavelength # resonant wavelength
thet_out # output phase
gam_out # output energy in unit of mc2
bunching # bunching factor
'''
#export variables
Nruns=inp_struct['Nruns']
npart=inp_struct['npart']
s_steps=inp_struct['s_steps']
z_steps=inp_struct['z_steps']
energy=inp_struct['energy']
eSpread=inp_struct['eSpread']
emitN=inp_struct['emitN']
currentMax=inp_struct['currentMax']
beta=inp_struct['beta']
unduPeriod=inp_struct['unduPeriod']
unduK=inp_struct['unduK']
unduL=inp_struct['unduL']
radWavelength=inp_struct['radWavelength']
dEdz=inp_struct['dEdz']
iopt=inp_struct['iopt']
P0=inp_struct['P0']
constseed=inp_struct['constseed']
# whether to use constant random seed for reproducibility
if constseed==1:
np.random.seed(22)
# Some constant values
alfvenCurrent = 17045.0 # Alfven current ~ 17 kA
mc2 = 510.99906E-3 # Electron rest mass in MeV
c = 2.99792458E8 # light speed in meter
e = 1.60217733E-19 # electron charge in Coulomb
#calculating intermediate parameters
unduJJ = scipy.special.jv(0,unduK**2/(4+2*unduK**2))\
-scipy.special.jv(1,unduK**2/(4+2*unduK**2)) # undulator JJ
gamma0 = energy/mc2 # central energy of the beam in unit of mc2
sigmaX2 = emitN*beta/gamma0 # rms transverse size, divergence of the electron beam
rho = (0.5/gamma0)*((currentMax/alfvenCurrent)\
*(unduPeriod*unduK*unduJJ/(2*np.pi))**2\
/(2*sigmaX2))**(1/3) # FEL Pierce parameter
resWavelength = unduPeriod*(1+unduK**2/2.0)\
/(2*gamma0**2) # resonant wavelength
rhoPbeam = rho*energy*currentMax/1000.0 # rho times beam power [GW]
coopLength = resWavelength/(4*np.pi*rho) # cooperation length
gainLength = unduPeriod/(4*np.pi*rho) # rough gain length
#cs0 = bunchLength/coopLength # bunch length in units of cooperation length
z0 = unduL/gainLength # wiggler length in units of gain length
delt = z0/z_steps # integration step in z0 ~ 0.1 gain length
dels = delt # integration step in s0 must be same as in z0
a02 = P0*1E-9/rhoPbeam # scaled input power
gbar = (resWavelength-radWavelength)\
/(radWavelength*rho) # scaled detune parameter
delg = eSpread/rho # Gaussian energy spread in units of rho
Ns = currentMax*unduL/unduPeriod/z_steps\
*resWavelength/c/e # N electrons per s-slice [ ]
Eloss = -dEdz*1E-3/energy/rho*gainLength # convert dEdz to alpha parameter
s = np.arange(1,s_steps+1)*dels*coopLength*1.0e6 # longitundinal steps along beam in meter
z = np.arange(1,z_steps+1)*delt*gainLength # longitundinal steps along undulator in meter
bunchLength=s[-1]*1e-6/2.5 # beam length in meter
bunch_steps=np.round(bunchLength/delt/coopLength) # rms (Gaussian) or half width (flattop) bunch length in s_step
shape = np.zeros((1,s_steps)) # initialization of the beam current shape
shape= 0.5*(np.tanh(10*(np.arange(1,s_steps+1)\
-s_steps/2+bunch_steps)/bunch_steps)\
-np.tanh(10*(np.arange(1,s_steps+1)\
-s_steps/2-bunch_steps)/bunch_steps)) # filling the shape of current and plot it
plt.plot(shape)
# initialization of variables during the 1D FEL process
ar=np.zeros((s_steps+1,z_steps+1))
ai=np.zeros((s_steps+1,z_steps+1))
gam=np.zeros((npart,z_steps+1))
thethalf=np.zeros((npart,z_steps+1))
thet_out=np.zeros((s_steps,1))
bunching=np.zeros((s_steps,z_steps),dtype=complex)
# sase mode is chosen, go over all slices of the bunch starting from the tail k=1
if iopt==5:
for k in range(s_steps):
ar[k,0] = np.sqrt(a02) # input seed signal
ai[k,0] = 0.0
[thet0,gam0] = load_bucket.load_bucket(npart,gbar,delg,iopt,Ns) # load each bucket
gam[:,0] = gam0.T # gamma at j=1
thethalf[:,0] = thet0.T-gam[:,0]*delt/2 # half back
thet_out[k,0]=np.mean(thet0.T)
for j in range(z_steps): # evolve e and gamma in s and t by leap-frog
thet = thethalf[:,j]+gam[:,j]*delt/2
sumsin = np.sum(np.sin(thet))
sumcos = np.sum(np.cos(thet))
sinavg = shape[k]*sumsin/npart
cosavg = shape[k]*sumcos/npart
arhalf = ar[k,j]+cosavg*dels/2
aihalf = ai[k,j]-sinavg*dels/2
thethalf[:,j+1] = thethalf[:,j]+gam[:,j]*delt
gam[:,j+1] = gam[:,j]-2*arhalf*np.cos(thethalf[:,j+1])*delt\
+2*aihalf*np.sin(thethalf[:,j+1])*delt-Eloss*delt #Eloss*delt to simulate the taper
sumsin = np.sum(np.sin(thethalf[:,j+1]))
sumcos = np.sum(np.cos(thethalf[:,j+1]))
sinavg = shape[k]*sumsin/npart
cosavg = shape[k]*sumcos/npart
ar[k+1,j+1] = ar[k,j]+cosavg*dels # apply slippage condition
ai[k+1,j+1] = ai[k,j]-sinavg*dels
bunching[k,j]=np.mean(np.real(np.exp(-1j*thet)))\
+np.mean(np.imag(np.exp(-1j*thet)))*1j #bunching factor calculation
#converting a and gam to field, power and gamavg
power_s=np.zeros((z_steps,s_steps))
power_z=np.zeros(z_steps)
gamavg=np.zeros(z_steps)
for j in range(z_steps):
for k in range(s_steps):
power_s[j,k] = (ar[k+1,j]**2+ai[k+1,j]**2)*rhoPbeam
power_z[j] = np.sum(ar[:,j]**2+ai[:,j]**2)/s_steps*rhoPbeam
gamavg[j] = np.sum(gam[:,j+1])/npart # average electron energy at every z position
thet_out=0 # don't output phase space
gam_out=0
detune = 2*np.pi/(dels*s_steps)*np.arange(-s_steps/2,s_steps/2+1)
field = (ar[:,z_steps]+ai[:,z_steps]*1j)*np.sqrt(rhoPbeam)
field_s = (ar[:,:]+ai[:,:]*1j)*np.sqrt(rhoPbeam)
return z,power_z,s,power_s,rho,detune,field,field_s,gainLength,resWavelength,thet_out,gam_out,bunching
1.1 Some constant values
alfvenCurrent = 17045.0 Alfven current ~ 17 kA
mc2 = 510.99906E-3 Electron rest mass in MeV
c = 2.99792458E8 light speed in meter
e = 1.60217733E-19 electron charge in Coulomb
1.2 calculating intermediate parameters
code:
unduJJ = scipy.special.jv(0,unduK**2/(4+2*unduK**2))\
-scipy.special.jv(1,unduK**2/(4+2*unduK**2))
equation: $[JJ]=J_0(\frac{K^2}{4+2K^2})-J_1(\frac{K^2}{4+2K^2})$
code:
sigmaX2 = emitN*beta/gamma0
equation:
$\sigma_x^2=\epsilon\beta$ scaled by central energy
code:
rho= (0.5/gamma0)*((currentMax/alfvenCurrent)\
*(unduPeriod*unduK*unduJJ/(2*np.pi))**2\
/(2*sigmaX2))**(1/3)
equation:
$\rho=[\frac{1}{8\pi} \frac{I}{I_A} (\frac{K[JJ]}{1+\frac{K^2}{2}})^2 \frac{\gamma \lambda_1^{2}}{2\pi\sigma_x^2}]^{\frac{1}{3}}$
# FEL Pierce parameter
resWavelength = unduPeriod*(1+unduK**2/2.0)\
/(2*gamma0**2) # resonant wavelength
equation: resonance condition
$\lambda_s=\frac{\lambda_u}{2\gamma^2}(1+\frac{K^2}{2})$
Some defination and scaling :
rhoPbeam = rho*energy*currentMax/1000.0 # rho times beam power [GW]
coopLength = resWavelength/(4*np.pi*rho) # cooperation length
gainLength = unduPeriod/(4*np.pi*rho) # rough gain length
#cs0 = bunchLength/coopLength # bunch length in units of cooperation length
z0 = unduL/gainLength # wiggler length in units of gain length
delt = z0/z_steps # integration step in z0 ~ 0.1 gain length
dels = delt # integration step in s0 must be same as in z0
$P=a^2 \rho P_{beam}$
Some definations
Eloss is an effective term for the taper optimization
a02 = P0*1E-9/rhoPbeam # scaled input power
gbar = (resWavelength-radWavelength)\
/(radWavelength*rho) # scaled detune parameter
delg = eSpread/rho # Gaussian energy spread in units of rho
Ns = currentMax*unduL/unduPeriod/z_steps\
*resWavelength/c/e # N electrons per s-slice [ ]
Eloss = -dEdz*1E-3/energy/rho*gainLength # convert dEdz to alpha parameter
s = np.arange(1,s_steps+1)*dels*coopLength*1.0e6 # longitundinal steps along beam in meter
z = np.arange(1,z_steps+1)*delt*gainLength # longitundinal steps along undulator in meter
Use a self-designed shape function to simulate the current profile
bunchLength=s[-1]*1e-6/2.5 # beam length in meter
bunch_steps=np.round(bunchLength/delt/coopLength) # rms (Gaussian) or half width (flattop) bunch length in s_step
shape = np.zeros((1,s_steps)) # initialization of the beam current shape
shape= 0.5*(np.tanh(10*(np.arange(1,s_steps+1)\
-s_steps/2+bunch_steps)/bunch_steps)\
-np.tanh(10*(np.arange(1,s_steps+1)\
-s_steps/2-bunch_steps)/bunch_steps)) # filling the shape of current and plot it
plt.plot(shape)
Initialization of some variables
# initialization of variables during the 1D FEL process
ar=np.zeros((s_steps+1,z_steps+1))
ai=np.zeros((s_steps+1,z_steps+1))
gam=np.zeros((npart,z_steps+1))
thethalf=np.zeros((npart,z_steps+1))
thet_out=np.zeros((s_steps,1))
bunching=np.zeros((s_steps,z_steps),dtype=complex)
1.3 FEL 1D process
Load bucket for each slice, calculate the initialized particle position and energy of the particles
equation:
$\frac{d\theta_j}{d\hat{z}}=\hat{\eta_j}$
code:
#sase mode is chosen, go over all slices of the bunch starting from the tail k=1
if iopt==5:
for k in range(s_steps):
ar[k,0] = np.sqrt(a02) # input seed signal
ai[k,0] = 0.0
[thet0,gam0] = load_bucket.load_bucket(npart,gbar,delg,iopt,Ns) # load each bucket
gam[:,0] = gam0.T # gamma at j=1
thethalf[:,0] = thet0.T-gam[:,0]*delt/2 # half back
thet_out[k,0]=np.mean(thet0.T)
equation:
$\frac{da}{d\hat{z}}=-<e^{-i\theta_{j}}>_{\Delta}$
code:
for j in range(z_steps): # evolve e and gamma in s and t by leap-frog
thet = thethalf[:,j]+gam[:,j]*delt/2
sumsin = np.sum(np.sin(thet))
sumcos = np.sum(np.cos(thet))
sinavg = shape[k]*sumsin/npart
cosavg = shape[k]*sumcos/npart
arhalf = ar[k,j]+cosavg*dels/2
aihalf = ai[k,j]-sinavg*dels/2
equation:
$\frac{d\theta_j}{d\hat{z}}=\hat{\eta_j}$
$\frac{d\hat{\eta_j}}{dz}=ae^{i\theta_{j}}+a^{*}e^{i\theta_{-j}}$
code:
thethalf[:,j+1] = thethalf[:,j]+gam[:,j]*delt
gam[:,j+1] = gam[:,j]-2*arhalf*np.cos(thethalf[:,j+1])*delt\
+2*aihalf*np.sin(thethalf[:,j+1])*delt-Eloss*delt
#Eloss*delt to simulate the taper
equation (same thing):
$\frac{da}{d\hat{z}}=-<e^{-i\theta_{j}}>_{\Delta}$
$\frac{d\theta_j}{d\hat{z}}=\hat{\eta_j}$
$\frac{d\hat{\eta_j}}{dz}=ae^{i\theta_{j}}+a^{*}e^{i\theta_{-j}}$
code:
sumsin = np.sum(np.sin(thethalf[:,j+1]))
sumcos = np.sum(np.cos(thethalf[:,j+1]))
sinavg = shape[k]*sumsin/npart
cosavg = shape[k]*sumcos/npart
ar[k+1,j+1] = ar[k,j]+cosavg*dels # apply slippage condition
ai[k+1,j+1] = ai[k,j]-sinavg*dels
bunching[k,j]=np.mean(np.real(np.exp(-1j*thet)))\
+np.mean(np.imag(np.exp(-1j*thet)))*1j #bunching factor calculation
1.4 converting a and gam to field, power and gamavg
Initialization of power_s (power along beam at different longitudinal undulator positions), power_z (averaged power at different longitundinal undulator positions), gamavg (averaged energy gamma at different longitundinal undulator positions)
power_s=np.zeros((z_steps,s_steps))
power_z=np.zeros(z_steps)
gamavg=np.zeros(z_steps)
$P=a^2 \rho P_{beam}$
detune is the frequency array of the beam
$field=\sqrt{P}$
for j in range(z_steps):
for k in range(s_steps):
power_s[j,k] = (ar[k+1,j]**2+ai[k+1,j]**2)*rhoPbeam
power_z[j] = np.sum(ar[:,j]**2+ai[:,j]**2)/s_steps*rhoPbeam
gamavg[j] = np.sum(gam[:,j+1])/npart # average electron energy at every z position
thet_out=0 # don't output phase space
gam_out=0
detune = 2*np.pi/(dels*s_steps)*np.arange(-s_steps/2,s_steps/2+1)
field = (ar[:,z_steps]+ai[:,z_steps]*1j)*np.sqrt(rhoPbeam)
field_s = (ar[:,:]+ai[:,:]*1j)*np.sqrt(rhoPbeam)
2.Load bucket function
import numpy as np
def load_bucket(n,gbar,delg,iopt,Ns):
'''
random initialization of the beam load_bucket
inputs:
n # n-macro-particles per bucket
gbar # scaled detune parameter
delg # Gaussian energy spread in units of rho
iopt # 5=SASE, 4=seeded
Ns # N electrons per s-slice ??
outputs:
thet # bucket macro particles position
gam # bucket macro particles energy
'''
nmax = 10000;
if n>nmax:
raise ValueError('increase nmax, subr load')
gam=np.zeros(n)
thet=np.zeros(n)
if iopt==4:
M=128 # number of particles in each beamlet
nb= int(np.round(n/M)) # number of beamlet via Fawley between 64 to 256 (x16=1024 to 4096)
if M*nb!=n:
raise ValueError('n must be a multiple of 4')
for i in range(nb):
gamma=delg*np.random.randn(1)+gbar
for j in range(M):
gam[i*M+j]=gamma
thet[i*M+j]=2*np.pi*j/M
elif iopt==5:
M=4 # number of particles in each beamlet
nb= int(np.round(n/M) ) #number of beamlet via Fawley between 64 to 256 (x16=1024 to 4096)
if M*nb!=n:
raise ValueError('n must be a multiple of 4')
effnoise = np.sqrt(3*M/(Ns/nb)) # Penman algorithm for Ns/nb >> M
for i in range(nb):
gamma=delg*np.random.randn(1)+gbar
for j in range(M):
gam[i*M+j]=gamma
thet[i*M+j]=2*np.pi*(j+1)/M+2*np.random.rand(1)*effnoise
return thet,gam
2.1 Set a limitation of the maximum number of macro-particles
And initialization of energy and position of every macro-particle
nmax = 10000;
if n>nmax:
raise ValueError('increase nmax, subr load')
gam=np.zeros(n)
thet=np.zeros(n)
2.2 Seeded mode
if iopt==4:
M=128 # number of particles in each beamlet
nb= int(np.round(n/M)) # number of beamlet via Fawley between 64 to 256 (x16=1024 to 4096)
if M*nb!=n:
raise ValueError('n must be a multiple of 4')
For every beamlet, use a random gaussian function scaled by the energy spread to sample each beamlet energy
For every macro-particle, the energy in each beamlet is set to be the same, the positions are set to be uniformly distributed along the bucket
for i in range(nb):
gamma=delg*np.random.randn(1)+gbar
for j in range(M):
gam[i*M+j]=gamma
thet[i*M+j]=2*np.pi*j/M
2.3 SASE mode
elif iopt==5:
M=4 # number of particles in each beamlet
nb= int(np.round(n/M) ) #number of beamlet via Fawley between 64 to 256 (x16=1024 to 4096)
if M*nb!=n:
raise ValueError('n must be a multiple of 4')
calculate an effnoise term that is propotional to the square root of the number of $\lambda_s$ in each bucket
For every beamlet, use a random gaussian function scaled by the energy spread to sample each beamlet energy
For every macro-particle, the energy in each beamlet is set to be the same, the positions are set to be uniformly distributed along the bucket with an random offset that is propotional to the effnoise
effnoise = np.sqrt(3*M/(Ns/nb)) # Penman algorithm for Ns/nb >> M
for i in range(nb):
gamma=delg*np.random.randn(1)+gbar
for j in range(M):
gam[i*M+j]=gamma
thet[i*M+j]=2*np.pi*(j+1)/M+2*np.random.rand(1)*effnoise