def back_propagate(params):
'''
Propagate pulse from file params[0] at the distance params[1] and save result to HDF5 file.
If output files exists - skip calculations.
'''
(input_path, distance, propagation_parameters) = params
input_dir, input_file_name = os.path.split(input_path)
out_file_name = '{}_{:0.4f}.h5'.format(
'.'.join(input_file_name.split('.')[:-1]), distance)
out_path = os.path.join(input_dir, out_file_name)
if os.path.exists(out_path):
return
wf_L1 = Wavefront()
wf_L1.load_hdf5(input_path)
drift1 = optical_elements.Drift(distance)
srwl_bl1 = SRWLOptC([drift1, ], [propagation_parameters, ])
bl1 = Beamline(srwl_bl1)
wpg.srwlib.srwl.SetRepresElecField(wf_L1._srwl_wf, 'f')
bl1.propagate(wf_L1)
wpg.srwlib.srwl.SetRepresElecField(wf_L1._srwl_wf, 't')
fit_gaussian_pulse(wf_L1)
wf_L1.store_hdf5(out_path)
del wf_L1
gc.collect()
return out_path
def setupTestWavefront():
""" Utility to setup a Gaussian wavefront. Geometry corresponds to SPB-SFX Day1 configuration. """
### Geometry ###
src_to_hom1 = 257.8 # Distance source to HOM 1 [m]
src_to_hom2 = 267.8 # Distance source to HOM 2 [m]
src_to_crl = 887.8 # Distance source to CRL [m]
src_to_exp = 920.42 # Distance source to experiment [m]
# Central photon energy.
ekev = 8.4 # Energy [keV]
# Pulse parameters.
qnC = 0.5 # e-bunch charge, [nC]
pulse_duration = 9.e-15 # [s]
pulseEnergy = 1.5e-3 # total pulse energy, J
# Coherence time
coh_time = 0.24e-15 # [s]
# Distance to first HOM.
z1 = src_to_hom1
# Angular distribution
theta_fwhm = 2.124e-6 # Beam divergence # From Patrick's raytrace.
wlambda = 12.4 * 1e-10 / ekev # wavelength [AKM]
w0 = wlambda / (numpy.pi * theta_fwhm) # beam waist
zR = (numpy.pi * w0**2) / wlambda #Rayleigh range
fwhm_at_zR = theta_fwhm * zR #FWHM at Rayleigh range
sigmaAmp = w0 / (2 * numpy.sqrt(numpy.log(2))) #sigma of amplitude
# Number of points in each x and y dimension.
np = 100
bSaved = False
dx = 10.e-6
range_xy = dx * (np - 1)
#print ('range_xy = ', range_xy)
nslices = 10
srwl_wf = build_gauss_wavefront(np,
np,
nslices,
ekev,
-range_xy / 2,
range_xy / 2,
-range_xy / 2,
range_xy / 2,
coh_time / numpy.sqrt(2),
sigmaAmp,
sigmaAmp,
src_to_hom1,
pulseEn=pulseEnergy,
pulseRange=8.)
wf = Wavefront(srwl_wf)
wf.store_hdf5("source.h5")
def load_wavefront(nslice_t, dirname_prop):
wf_holder = []
for i in range(nslice_t):
fname = dirname_prop + 'wavefront_focused_slice_' + str(i) + '.h5'
mwf_temp = Wavefront()
mwf_temp.load_hdf5(fname)
wf_holder.append(mwf_temp)
return wf_holder
def show_diagnostics(FELsource_out_number):
# read FELsource_out_.h5
if not FELsource_out_number == 'FELsource_out_0000001.h5':
#FELsource_out_file = "FELsource_out_{}.h5".format(FELsource_out_number.zfill(7))
FELsource_out_file = "{}.h5".format(FELsource_out_number.zfill(7))
else:
FELsource_out_file = FELsource_out_number
if not os.path.exists(FELsource_out_file):
print 'Input file {} not found.'.format(FELsource_out_file)
return
wf = Wavefront()
wf.load_hdf5(FELsource_out_file)
# show two figures window 1: image of I(x,y) integral intensity, with real
# x and y axis and title with file name
J2eV = 6.24150934e18;
mesh = wf.params.Mesh
tmin = mesh.sliceMin;
tmax = mesh.sliceMax;
dt = (tmax - tmin) / (mesh.nSlices - 1);
dx = (mesh.xMax - mesh.xMin) / (mesh.nx - 1);
dy = (mesh.yMax - mesh.yMin) / (mesh.ny - 1);
wf_intensity = wf.get_intensity(polarization='horizontal');
total_intensity = wf_intensity.sum(axis=-1);
data = total_intensity * dt
plt.figure()
plt.imshow(data*dx*dy*1e6*J2eV/wf.params.photonEnergy,extent=[mesh.xMin*1e6,mesh.xMax*1e6,mesh.yMin*1e6,mesh.yMax * 1e6])
title = 'Number of photons per %.2f x %.2f $\mu m ^2$ pixel' % (dx*1e6, dx*1e6)
plt.title(title)
plt.colorbar(); plt.xlabel('[$\mu m$]');
# window 2: plot of 2 curves:
#(1) history/parent/temporal_struct - FAST post-processing
temporal_struct = wf.custom_fields['history']['parent']['misc']['temporal_struct']
t0 = (temporal_struct[:, 0].max() + temporal_struct[:, 0].min()) / 2
plt.figure()
plt.plot(temporal_struct[:, 0] - t0, temporal_struct[:, 1] * 1e-9, 'b',label = 'output FAST-pp')
plt.hold(True)
#(2) integral intensity I(t) calculated for wavefront written in h5
t = np.linspace(tmin, tmax, wf.params.Mesh.nSlices)
pulse_energy = wf.get_intensity().sum(axis=0).sum(axis=0) #check it
plt.plot(t * 1e15, pulse_energy*dx*dy*1e6*1e-9,'ro', label = 'wavefront data')
title = 'FEL pulse energy %.2f %s ' % (pulse_energy.sum(axis=0) * dx * dy * 1e6 * dt * 1e3, 'mJ')
plt.title(title)
plt.xlabel('time [fs]');
plt.ylabel('Instantaneous power [GW]');
plt.legend()
plt.grid(True)
plt.show()
def propagate_run(ifname, ofname, optBL, bSaved=False):
"""
Propagate wavefront through a beamline and save the result (optionally).
:param ifname: input hdf5 file name with wavefront to be propagated
:param ofname: output hdf5 file name
:param optBL: beamline
:param bSaved: if True, save propagated wavefront in h5 file
:return: propagated wavefront
"""
print_beamline(optBL)
startTime = time.time()
print('*****reading wavefront from h5 file...')
w2 = Wavefront()
w2.load_hdf5(ifname + '.h5')
wfr = w2._srwl_wf
print('*****propagating wavefront (with resizing)...')
srwl.PropagElecField(wfr, optBL)
mwf = Wavefront(wfr)
print('[nx, ny, xmin, xmax, ymin, ymax]', get_mesh(mwf))
if bSaved:
print('save hdf5:', ofname + '.h5')
mwf.store_hdf5(ofname + '.h5')
print('done')
print('propagation lasted:',
round((time.time() - startTime) / 6.) / 10., 'min')
return wfr
def _readH5(self):
""" """
""" Private method for reading the hdf5 input and extracting the parameters and data relevant to initialize the object. """
# Check input.
try:
self.__h5 = h5py.File(self.input_path, 'r')
except:
raise IOError(
'The input_path argument (%s) is not a path to a valid hdf5 file.'
% (self.input_path))
# Construct wpg wavefront based on input data.
self.__wavefront = Wavefront()
self.__wavefront.load_hdf5(self.input_path)
def do_wpg_calculation(self):
theta_fwhm = self.calculate_theta_fwhm_cdr(self.ekev, self.qnC)
k = 2*numpy.sqrt(2*numpy.log(2))
sigX = 12.4e-10*k/(self.ekev*4*numpy.pi*theta_fwhm)
print('sigX, waist_fwhm [um], far field theta_fwhms [urad]: {}, {},{}'.format(
sigX*1e6, sigX*k*1e6, theta_fwhm*1e6)
)
#define limits
range_xy = theta_fwhm/k*self.distance*7. # sigma*7 beam size
self.range_xy = range_xy
npoints=180
wfr0 = build_gauss_wavefront_xy(npoints,
npoints,
self.ekev,
-range_xy/2,
range_xy/2,
-range_xy/2,
range_xy/2,
sigX,
sigX,
self.distance,
pulseEn=self.pulseEnergy,
pulseTau=self.coh_time/numpy.sqrt(2),
repRate=1/(numpy.sqrt(2)*self.pulse_duration))
return Wavefront(wfr0)
def forward_propagate(root_dir, distance, propagation_parameters):
"""
Forward_propagate_wavefront
the result will saved in root_dir\distance\distance.h5 file
:param root_dir: directory, where '0.h' file located
:param distance: distance to forward propagate initial wvefront
:param propagation_parameters: SRW propagation parameters
"""
out_dir = os.path.join(root_dir, '{:0.4f}'.format(distance))
mkdir_p(out_dir)
out_file_name = '{:0.4f}.h5'.format(distance)
out_path = os.path.join(out_dir, out_file_name)
if os.path.exists(out_path):
print 'File exists: {}. Skiping.'.format(out_path)
return out_path
ppDrift0 = propagation_parameters
drift0 = optical_elements.Drift(distance)
srwl_bl0 = SRWLOptC([drift0, ], [ppDrift0, ])
bl0 = Beamline(srwl_bl0)
# forward propagate to L0 meters
wf_L0 = Wavefront()
wf_L0.load_hdf5(os.path.join(root_dir, '0.h5'))
tmin = wf_L0.params.Mesh.sliceMin
tmax = wf_L0.params.Mesh.sliceMax
wf_L0.params.Mesh.sliceMin = -(tmax-tmin)/2
wf_L0.params.Mesh.sliceMax = (tmax-tmin)/2
# wpg.srwlib.srwl.ResizeElecField(wf_L0._srwl_wf, 't',[0,3.,1.])
wpg.srwlib.srwl.SetRepresElecField(wf_L0._srwl_wf, 'f')
bl0.propagate(wf_L0)
wpg.srwlib.srwl.SetRepresElecField(wf_L0._srwl_wf, 't')
fit_gaussian_pulse(wf_L0)
wf_L0.store_hdf5(out_path)
print 'Save file : {}'.format(out_path)
del wf_L0
return out_path
def read_genesis_file(gen_fname, lmb_und, pol="lin"):
speed_of_light = const.codata.value("speed of light in vacuum")
h_eV_s = const.codata.value("Planck constant in eV s")
wf = Wavefront()
### Open the hdf5 output GENESIS file
with h5py.File(gen_fname, "r") as hf:
grid_size = hf["gridsize"].value
slice_count = hf["slicecount"].value
print(slice_count)
slice_spacing = hf["slicespacing"].value
wavelength = hf["wavelength"].value
data1 = open_slice_data(hf, 1)
npoints = int(np.sqrt(len(data1) / 2))
range_xy = (npoints - 1) * grid_size
### Fill in the fields of the wavefront object
wf.params.wEFieldUnit = "sqrt(W/mm^2)"
wf.params.photonEnergy = h_eV_s * speed_of_light / wavelength
print("Photon energy: ", wf.params.photonEnergy, "eV")
wf.params.wDomain = "time"
wf.params.Mesh.nSlices = slice_count
wf.params.Mesh.nx = npoints
wf.params.Mesh.ny = npoints
pulse_length = (slice_count - 1) * slice_spacing / (speed_of_light)
print("Pulse length: ", pulse_length)
wf.params.Mesh.sliceMin = -pulse_length / 2.0
wf.params.Mesh.sliceMax = pulse_length / 2.0
wf.params.Mesh.xMin = -range_xy / 2.0
wf.params.Mesh.xMax = range_xy / 2.0
wf.params.Mesh.yMin = -range_xy / 2.0
wf.params.Mesh.yMax = range_xy / 2.0
wf.params.nval = 2
### Extract the field data from the h5 file and fill in the data of the Electric field, by calling the
### vector_grid_conversion function
if pol == "lin":
wf.data.arrEhor = vector_grid_conversion(hf, npoints, slice_count,
range_xy, wavelength,
lmb_und)
wf.data.arrEver = np.zeros_like(wf.data.arrEhor)
else:
matrix_E = vector_grid_conversion(hf, npoints, slice_count,
range_xy, wavelength, lmb_und)
Re_Etotal = deepcopy(matrix_E[:, :, :, 0])
Im_Etotal = deepcopy(matrix_E[:, :, :, 1])
wf.data.arrEhor = deepcopy(matrix_E)
wf.data.arrEver = deepcopy(wf.data.arrEhor)
wf.data.arrEhor = np.multiply(np.sqrt(0.5), wf.data.arrEhor)
wf.data.arrEver[:, :, :, 0] = deepcopy(
np.multiply(-1.0, wf.data.arrEhor[:, :, :, 1]))
wf.data.arrEver[:, :, :, 1] = deepcopy(wf.data.arrEhor[:, :, :, 0])
return wf
def propagate(in_fname, out_fname, get_beamline):
"""
Propagate wavefront
:param in_file: input wavefront file
:param out_file: output file
:param get_beamline: function to build beamline
"""
print("#" * 80)
print("Setup initial wavefront.")
wf = Wavefront()
# Load wavefront data.
print("Load " + in_fname)
wf.load_hdf5(in_fname)
# Get beamline.
bl0 = get_beamline()
# Switch to frequency domain.
wpg.srwlib.srwl.SetRepresElecField(wf._srwl_wf, "f")
# Save spectrum for later reference.
sz0 = get_intensity_on_axis(wf)
wf.custom_fields["/misc/spectrum0"] = sz0
# Propagate.
bl0.propagate(wf)
# Save spectrum after propagation for later reference.
sz1 = get_intensity_on_axis(wf)
wf.custom_fields["/misc/spectrum1"] = sz1
# Switch back to time domain.
wpg.srwlib.srwl.SetRepresElecField(wf._srwl_wf, "t")
# Resizing: decreasing Range of Horizontal and Vertical Position:
wpg.srwlib.srwl.ResizeElecField(wf._srwl_wf, "c", [0, 0.5, 1, 0.5, 1])
add_custom_data(wf, bl0)
print("Saving propagated wavefront to " + out_fname)
mkdir_p(os.path.dirname(out_fname))
wf.store_hdf5(out_fname)
print("Saving history.")
add_history(out_fname, in_fname)
print("ALL DONE.")
print("#" * 80)
def _readH5(self):
""" """
""" Private method for reading the hdf5 input and extracting the parameters and data relevant to initialize the object. """
# Check input.
try:
self.__h5 = h5py.File( self.input_path, 'r' )
except:
raise IOError( 'The input_path argument (%s) is not a path to a valid hdf5 file.' % (self.input_path) )
# Construct wpg wavefront based on input data.
self.__wavefront = Wavefront()
self.__wavefront.load_hdf5(self.input_path)
def backengine(self):
# The rms of the amplitude distribution (Gaussian)
theta = self.parameters.divergence.m_as(radian)
E = self.parameters.photon_energy
coherence_time = 2. * math.pi * hbar / self.parameters.photon_energy_relative_bandwidth / E.m_as(
joule)
beam_waist = 2. * hbar * c / theta / E.m_as(joule)
wavelength = 1239.8e-9 / E.m_as(electronvolt)
rayleigh_length = math.pi * beam_waist**2 / wavelength
print(rayleigh_length)
beam_diameter_fwhm = self.parameters.beam_diameter_fwhm.m_as(meter)
beam_waist_radius = beam_diameter_fwhm / math.sqrt(2. * math.log(2.))
# x-y range at beam waist.
range_xy = 30.0 * beam_waist_radius
# Set number of sampling points in x and y and number of temporal slices.
np = self.parameters.number_of_transverse_grid_points
nslices = self.parameters.number_of_time_slices
# Distance from source position.
z = self.parameters.z.m_as(meter)
# Build wavefront
srwl_wf = build_gauss_wavefront(
np,
np,
nslices,
E.m_as(electronvolt) / 1.0e3,
-range_xy / 2,
range_xy / 2,
-range_xy / 2,
range_xy / 2,
coherence_time / math.sqrt(2),
beam_waist_radius / 2,
beam_waist_radius /
2, # Scaled such that fwhm comes out as demanded by parameters.
d2waist=z,
pulseEn=self.parameters.pulse_energy.m_as(joule),
pulseRange=8.)
# Correct radius of curvature.
Rx = Ry = z * math.sqrt(1. + (rayleigh_length / z)**2)
# Store on class.
srwl_wf.Rx = Rx
srwl_wf.Ry = Ry
self.__wavefront = Wavefront(srwl_wf)
def propagate_wavefront(wavefront, beamline, output_file = None):
"""
Propagate wavefront and store it in output file.
:param wavefront: Wavefront object or path to HDF5 file
:param beamline: SRWLOptC container of beamline
:param output_file: if parameter present - store propagaed wavefront to file
:return: propagated wavefront object:
"""
if not isinstance(beamline, Beamline):
bl = Beamline(beamline)
else:
bl = beamline
if isinstance(wavefront, Wavefront):
wfr = Wavefront(srwl_wavefront=wavefront._srw_wf)
else:
print '*****reading wavefront from h5 file...'
wfr = Wavefront()
wfr.load_hdf5(wavefront)
print '*****propagating wavefront (with resizing)...'
bl.propagate(wfr)
print '[nx, ny, xmin, xmax, ymin, ymax]', get_mesh(wfr)
if not output_file is None:
print 'save hdf5:', output_file
wfr.store_hdf5(output_file)
print 'done'
return wfr
def back_propagate(params):
"""
Propagate pulse from file params[0] at the distance params[1] and save result to HDF5 file.
If output files exists - skip calculations.
"""
(input_path, distance, propagation_parameters) = params
input_dir, input_file_name = os.path.split(input_path)
out_file_name = "{}_{:0.4f}.h5".format(
".".join(input_file_name.split(".")[:-1]), distance
)
out_path = os.path.join(input_dir, out_file_name)
if os.path.exists(out_path):
return
wf_L1 = Wavefront()
wf_L1.load_hdf5(input_path)
drift1 = optical_elements.Drift(distance)
srwl_bl1 = SRWLOptC([drift1,], [propagation_parameters,])
bl1 = Beamline(srwl_bl1)
wpg.srwlib.srwl.SetRepresElecField(wf_L1._srwl_wf, "f")
bl1.propagate(wf_L1)
wpg.srwlib.srwl.SetRepresElecField(wf_L1._srwl_wf, "t")
fit_gaussian_pulse(wf_L1)
wf_L1.store_hdf5(out_path)
del wf_L1
gc.collect()
return out_path
def propagate_run(ifname, ofname, optBL, bSaved=False):
"""
Propagate wavefront through a beamline and save the result (optionally).
:param ifname: input hdf5 file name with wavefront to be propagated
:param ofname: output hdf5 file name
:param optBL: beamline
:param bSaved: if True, save propagated wavefront in h5 file
:return: propagated wavefront
"""
print_beamline(optBL)
startTime = time.time()
print '*****reading wavefront from h5 file...'
w2 = Wavefront()
w2.load_hdf5(ifname + '.h5')
wfr = w2._srwl_wf
print '*****propagating wavefront (with resizing)...'
srwl.PropagElecField(wfr, optBL)
mwf = Wavefront(wfr)
print '[nx, ny, xmin, xmax, ymin, ymax]', get_mesh(mwf)
if bSaved:
print 'save hdf5:', ofname + '.h5'
mwf.store_hdf5(ofname + '.h5')
print 'done'
print 'propagation lasted:', round((time.time() - startTime) / 6.) / 10., 'min'
return wfr
def main(wpg_out):
# Backup
backup = wpg_out + ".backup"
shutil.copy2(wpg_out, backup)
print("Processing %s." % (wpg_out))
# Load wavefront.
wavefront = Wavefront()
wavefront.load_hdf5(wpg_out)
# Get width from intensity profile.
fwhm = wpg_uti_wf.calculate_fwhm(wavefront)
# Carefully insert new datasets.
try:
with h5py.File(wpg_out, 'a') as h5_handle:
if not 'misc' in h5_handle.keys():
misc = h5_handle.create_group("misc")
else:
misc = h5_handle['misc']
if not "xFWHM" in misc.keys():
misc.create_dataset("xFWHM", data=fwhm["fwhm_x"])
if not "yFWHM" in misc.keys():
misc.create_dataset("yFWHM", data=fwhm["fwhm_y"])
# If anything went wrong, restore the original file and raise.
except:
shutil.move(backup, wpg_out)
raise
# We only reach this point if everything went well, so can safely remove the backup."
os.remove(backup)
def read_genesis_file(gen_fname):
speed_of_light = const.codata.value("speed of light in vacuum")
h_eV_s = const.codata.value("Planck constant in eV s")
lmb_und = 2.75e-2
wf = Wavefront()
### Open the hdf5 output GENESIS file
with h5py.File(gen_fname, "r") as hf:
grid_size = hf["gridsize"].value
slice_count = hf["slicecount"].value
print(slice_count)
slice_spacing = hf["slicespacing"].value
wavelength = hf["wavelength"].value
data1 = open_slice_data(hf, 1)
npoints = int(np.sqrt(len(data1) / 2))
### Definition of the Electric field field arrays where the electric field from the GENESIS output file
### will be copied
wf.data.arrEhor = np.zeros(shape=(npoints, npoints, slice_count, 2))
wf.data.arrEver = np.zeros(shape=(npoints, npoints, slice_count, 2))
### Fill in the fields of the wavefront object
wf.params.wEFieldUnit = "sqrt(W/mm^2)"
wf.params.photonEnergy = h_eV_s * speed_of_light / wavelength
print("Photon energy: ", wf.params.photonEnergy, "eV")
wf.params.wDomain = "time"
wf.params.Mesh.nSlices = slice_count
wf.params.Mesh.nx = npoints
wf.params.Mesh.ny = npoints
pulse_length = (slice_count - 1) * slice_spacing / (speed_of_light)
print("Pulse length: ", pulse_length)
wf.params.Mesh.sliceMin = -pulse_length / 2.0
wf.params.Mesh.sliceMax = pulse_length / 2.0
range_xy = (npoints - 1) * grid_size
wf.params.Mesh.xMin = -range_xy / 2.0
wf.params.Mesh.xMax = range_xy / 2.0
wf.params.Mesh.yMin = -range_xy / 2.0
wf.params.Mesh.yMax = range_xy / 2.0
### Extract the field data from the h5 file and fill in the data of the Electric field, by calling the
### vector_grid_conversion function
wf.data.arrEhor = vector_grid_conversion(hf, npoints, slice_count,
range_xy, wavelength, lmb_und)
return wf
def backengine(self):
# The rms of the amplitude distribution (Gaussian)
theta = self.parameters.divergence.m_as(radian)
E = self.parameters.photon_energy
coherence_time = 2. * math.pi * hbar / self.parameters.photon_energy_relative_bandwidth / E.m_as(
joule)
beam_waist = 2. * hbar * c / theta / E.m_as(joule)
beam_diameter_fwhm = self.parameters.beam_diameter_fwhm.m_as(meter)
beam_waist_radius = beam_diameter_fwhm / math.sqrt(2. * math.log(2.))
print("beam waist radius from divergence angle = {0:4.3e}".format(
beam_waist))
print(
"beam waist radius from fwhm = {0:4.3e}".format(beam_waist_radius))
# Rule of thumb: 7 times w0
# x-y range at beam waist.
range_xy = 30.0 * beam_waist_radius
# Set number of sampling points in x and y and number of temporal slices.
np = self.parameters.number_of_transverse_grid_points
nslices = self.parameters.number_of_time_slices
# Build wavefront
srwl_wf = build_gauss_wavefront(
np,
np,
nslices,
E.m_as(electronvolt) / 1.0e3,
-range_xy / 2,
range_xy / 2,
-range_xy / 2,
range_xy / 2,
coherence_time / math.sqrt(2),
# beam_diameter_fwhm, beam_diameter_fwhm,
beam_waist_radius / 2,
beam_waist_radius /
2, # Scaled such that fwhm comes out as demanded by parameters.
0.0,
pulseEn=self.parameters.pulse_energy.m_as(joule),
pulseRange=8.)
# Store on class.
self.__wavefront = Wavefront(srwl_wf)
def show_diagnostics(FELsource_out_number):
FELsource_out_file = FELsource_out_number
if not os.path.exists(FELsource_out_file):
print('Input file {} not found.'.format(FELsource_out_file))
return
wf = Wavefront()
wf.load_hdf5(FELsource_out_file)
plot_t_wf(wf)
look_at_q_space(wf)
# show two figures window 1: image of I(x,y) integral intensity, with real
# x and y axis and title with file name
J2eV = 6.24150934e18;
mesh = wf.params.Mesh
tmin = mesh.sliceMin;
tmax = mesh.sliceMax;
dt = (tmax - tmin) / (mesh.nSlices - 1);
dx = (mesh.xMax - mesh.xMin) / (mesh.nx - 1);
dy = (mesh.yMax - mesh.yMin) / (mesh.ny - 1);
wf_intensity = wf.get_intensity(polarization='horizontal');
total_intensity = wf_intensity.sum(axis=-1);
data = total_intensity * dt
plt.figure()
plt.imshow(data*dx*dy*1e6*J2eV/wf.params.photonEnergy,extent=[mesh.xMin*1e6,mesh.xMax*1e6,mesh.yMin*1e6,mesh.yMax * 1e6], cmap="YlGnBu_r")
title = 'Number of photons per %.2f x %.2f $\mu m ^2$ pixel' % (dx*1e6, dx*1e6)
plt.title(title)
plt.colorbar(); plt.xlabel('[$\mu m$]');
# window 2: plot of 2 curves:
#(1) history/parent/temporal_struct - FAST post-processing
temporal_struct = wf.custom_fields['history']['parent']['misc']['temporal_struct']
t0 = (temporal_struct[:, 0].max() + temporal_struct[:, 0].min()) / 2
plt.figure()
plt.plot(temporal_struct[:, 0] - t0, temporal_struct[:, 1] * 1e-9, 'b',label = 'output FAST-pp')
plt.hold(True)
#(2) integral intensity I(t) calculated for wavefront written in h5
t = np.linspace(tmin, tmax, wf.params.Mesh.nSlices)
pulse_energy = wf.get_intensity().sum(axis=0).sum(axis=0) #check it
plt.plot(t * 1e15, pulse_energy*dx*dy*1e6*1e-9,'ro', label = 'wavefront data')
title = 'FEL pulse energy %.2f %s ' % (pulse_energy.sum(axis=0) * dx * dy * 1e6 * dt * 1e3, 'mJ')
plt.title(title)
plt.xlabel('time [fs]');
plt.ylabel('Instantaneous power [GW]');
plt.legend()
plt.grid(True)
plt.show()
def setup_gauss_wavefront(sanity=True):
""" Gaussian wavefront builder """
if sanity:
np = 700
nslices = 100
else:
np = 10
nslices = 10
photon_energy = 1.6e3
theta_fwhm = 6.0e-6 ### ADJUST ME
wlambda = 1.24 * 1e-6 / photon_energy # wavelength [AKM]
w0 = wlambda / (numpy.pi * theta_fwhm) # beam waist
zR = (numpy.pi * w0 ** 2) / wlambda # Rayleigh range
sigmaAmp = w0 / (2 * numpy.sqrt(numpy.log(2))) # sigma of amplitude
src_to_aperture_distance = 170.0
pulse_energy = 4e-4 # [J]
# Coherence time
pulse_duration = 30.0e-15 # [s]
coh_time = pulse_duration / 10.0 # estimate, [s]
# expected beam radius at HOM1 position to get the range of the wavefront
range_xy = w0 * numpy.sqrt(1 + (src_to_aperture_distance / zR) ** 2) * 5.5
srwl_wf = build_gauss_wavefront(
np,
np,
nslices,
photon_energy / 1e3,
-range_xy / 2,
range_xy / 2,
-range_xy / 2,
range_xy / 2,
coh_time / numpy.sqrt(2),
sigmaAmp,
sigmaAmp,
src_to_aperture_distance,
pulseEn=pulse_energy,
pulseRange=8.0,
)
return Wavefront(srwl_wf)
def forward_propagate(root_dir, distance, propagation_parameters):
"""
Forward_propagate_wavefront
the result will saved in root_dir\distance\distance.h5 file
:param root_dir: directory, where '0.h' file located
:param distance: distance to forward propagate initial wvefront
:param propagation_parameters: SRW propagation parameters
"""
out_dir = os.path.join(root_dir, '{:0.4f}'.format(distance))
mkdir_p(out_dir)
out_file_name = '{:0.4f}.h5'.format(distance)
out_path = os.path.join(out_dir, out_file_name)
if os.path.exists(out_path):
print('File exists: {}. Skiping.'.format(out_path))
return out_path
ppDrift0 = propagation_parameters
drift0 = optical_elements.Drift(distance)
srwl_bl0 = SRWLOptC([
drift0,
], [
ppDrift0,
])
bl0 = Beamline(srwl_bl0)
# forward propagate to L0 meters
wf_L0 = Wavefront()
wf_L0.load_hdf5(os.path.join(root_dir, '0.h5'))
tmin = wf_L0.params.Mesh.sliceMin
tmax = wf_L0.params.Mesh.sliceMax
wf_L0.params.Mesh.sliceMin = -(tmax - tmin) / 2
wf_L0.params.Mesh.sliceMax = (tmax - tmin) / 2
# wpg.srwlib.srwl.ResizeElecField(wf_L0._srwl_wf, 't',[0,3.,1.])
wpg.srwlib.srwl.SetRepresElecField(wf_L0._srwl_wf, 'f')
bl0.propagate(wf_L0)
wpg.srwlib.srwl.SetRepresElecField(wf_L0._srwl_wf, 't')
fit_gaussian_pulse(wf_L0)
wf_L0.store_hdf5(out_path)
print('Save file : {}'.format(out_path))
del wf_L0
return out_path
def plot_intensity(wavefront):
"""IN 2D
PLOT THE INTENSITY OF A WAVFRONT"""
print("******plotting wavefront intensity")
# calculate intensity
wavefront_intensity = Wavefront.get_intensity(wavefront)
wavefront_intensity = wavefront_intensity.sum(-1)
plt.figure()
plt.set_cmap("jet")
plt.imshow(wavefront_intensity)
# remove ticks
plt.gca().set_xticks([])
plt.gca().set_yticks([])
# color bar
cbar = plt.colorbar()
cbar.ax.set_ylabel("Intensity")
def add_wf_attributes(fname0):
# use srwlib glossary to add attributes to wavefront datasets
in_fname = fname0+'.h5'
bare_fname = fname0+'_bare.h5'
if doPrint: print('Loading wavefront data from the file: '+in_fname)
wf_struct=Wavefront()
wf_struct.load_hdf5(in_fname)
wfr = wf_struct._srwl_wf
wf_struct = Wavefront(wfr)
if doPrint: print('Saving the wavefront data with attributes:'+bare_fname)
wf_struct.store_hdf5(bare_fname)
if doPrint: print('Replacing data with attributes from '+bare_fname)
with h5py.File(bare_fname) as h2:
with h5py.File(in_fname) as h1:
try:
del h1['params'] # delete group
except KeyError:
pass
h2.copy('params',h1) #copy h2['params'] to h1
def show_diagnostics(prop_out_number):
# read prop_out_.h5
if not prop_out_number == 'prop_out_1.h5':
#prop_out_file = "prop_out_{}.h5".format(prop_out_number.zfill(7))
prop_out_file = "{}.h5".format(prop_out_number.zfill(7))
else:
prop_out_file = prop_out_number
if not os.path.exists(prop_out_file):
print 'Input file {} not found.'.format(prop_out_file)
return
wf = Wavefront()
wf.load_hdf5(prop_out_file)
# show two figures window 1: image of I(x,y) integral intensity, with real
# x and y axis and title with file name
J2eV = 6.24150934e18
mesh = wf.params.Mesh
tmin = mesh.sliceMin
tmax = mesh.sliceMax
dt = (tmax - tmin) / (mesh.nSlices - 1)
dx = (mesh.xMax - mesh.xMin) / (mesh.nx - 1)
dy = (mesh.yMax - mesh.yMin) / (mesh.ny - 1)
wf_intensity = wf.get_intensity(polarization='horizontal')
total_intensity = wf_intensity.sum(axis=-1)
data = total_intensity * dt
plt.figure()
plt.imshow(data*dx*dy*1e6*J2eV/wf.params.photonEnergy,extent=[mesh.xMin*1e6,mesh.xMax*1e6,mesh.yMin*1e6,mesh.yMax * 1e6])
title = 'Number of photons per %.2f x %.2f $nm ^2$ pixel' % (dx*1e9, dx*1e9)
plt.title(title)
plt.colorbar()
plt.xlabel(r'[$\mu$m]')
# window 2: plot of 2 curves:
#(1) history/parent/parent/temporal_struct - before propagating
temporal_struct = wf.custom_fields['history']['parent']['parent']['misc']['temporal_struct']
t0 = (temporal_struct[:, 0].max() + temporal_struct[:, 0].min()) / 2
plt.figure()
plt.plot(temporal_struct[:, 0] - t0, temporal_struct[:, 1] * 1e-9, 'b',label = 'original')
plt.hold(True)
#(2) integral intensity I(t) after propagating
t = np.linspace(tmin, tmax, wf.params.Mesh.nSlices)
pulse_energy = wf.get_intensity().sum(axis=0).sum(axis=0) #check it
plt.plot(t * 1e15, pulse_energy*dx*dy*1e6*1e-9,'r', label = 'propag')
title = 'The propagated pulse energy %.2f %s ' % (pulse_energy.sum(axis=0) * dx * dy * 1e6 * dt * 1e3, 'mJ')
plt.title(title)
plt.xlabel('time [fs]')
plt.ylabel('Instantaneous power [GW]')
plt.legend()
plt.grid(True)
sz0 = wf.custom_fields['misc']['spectrum0']
sz1 = wf.custom_fields['misc']['spectrum1']
plt.figure()
plt.plot(sz0[:,0],sz0[:,1], label='before propagating')
plt.hold(True)
plt.plot(sz1[:,0],sz1[:,1],label='after propagating')
plt.grid(True)
plt.title('Spectrum (x=y=0)')
plt.xlabel('[eV]')
plt.ylabel('[arb. unit]')
plt.legend()
plt.show()
def propagate(in_fname, out_fname):
"""
Propagate wavefront
:param in_file: input wavefront file
:param out_file: output file
"""
print('Start propagating:' + in_fname)
wf=Wavefront()
wf.load_hdf5(in_fname)
distance0 = 300.
distance1 = 630.
distance = distance0 + distance1
f_hfm = 3.0 # nominal focal length for HFM KB
f_vfm = 1.9 # nominal focal length for VFM KB
distance_hfm_vfm = f_hfm - f_vfm
distance_foc = 1. /(1./f_vfm + 1. / (distance + distance_hfm_vfm))
theta_om = 3.5e-3 # offset mirrors incidence angle
theta_kb = 3.5e-3 # KB mirrors incidence angle
drift0 = wpg.optical_elements.Drift(distance0)
drift1 = wpg.optical_elements.Drift(distance1)
drift_in_kb = wpg.optical_elements.Drift(distance_hfm_vfm)
drift_to_foc = wpg.optical_elements.Drift(distance_foc)
om_mirror_length = 0.8; om_clear_ap = om_mirror_length*theta_om
kb_mirror_length = 0.9; kb_clear_ap = kb_mirror_length*theta_kb
ap0 = wpg.optical_elements.Aperture('r','a', 120.e-6, 120.e-6)
ap1 = wpg.optical_elements.Aperture('r','a', om_clear_ap, 2*om_clear_ap)
ap_kb = wpg.optical_elements.Aperture('r','a', kb_clear_ap, kb_clear_ap)
hfm = wpg.optical_elements.Mirror_elliptical(
orient='x',p=distance, q=(distance_hfm_vfm+distance_foc),
thetaE=theta_kb, theta0=theta_kb, length=0.9)
vfm = wpg.optical_elements.Mirror_elliptical(
orient='y',p=(distance+distance_hfm_vfm), q=distance_foc,
thetaE=theta_kb, theta0=theta_kb, length=0.9)
wf_dist_om = wpg.optical_elements.WF_dist(1500, 100, om_clear_ap, 2*om_clear_ap)
defineOPD(wf_dist_om, os.path.join(mirror_data_dir,'mirror2.dat'), 2, '\t', 'x',
theta_kb, scale=2)
if isIpynb:
meshT = wf_dist_om.mesh
opdTmp=np.array(wf_dist_om.arTr)[1::2].reshape(meshT.ny,meshT.nx)
figure(); pylab.imshow(opdTmp,extent=[meshT.xStart,meshT.xFin,meshT.yStart,meshT.yFin])
pylab.title('OPD [m]');pylab.xlabel('x (m)'); pylab.ylabel('y (m)')
wf_dist_hfm = wpg.optical_elements.WF_dist(1500, 100, kb_clear_ap, kb_clear_ap)
defineOPD(wf_dist_hfm, os.path.join(mirror_data_dir,'mirror1.dat'), 2, '\t', 'x',
theta_kb, scale=2, stretching=kb_mirror_length/0.8)
if isIpynb:
meshT = wf_dist_hfm.mesh
opdTmp=np.array(wf_dist_hfm.arTr)[1::2].reshape(meshT.ny,meshT.nx)
figure(); pylab.imshow(opdTmp,extent=[meshT.xStart,meshT.xFin,meshT.yStart,meshT.yFin])
pylab.title('OPD [m]');pylab.xlabel('x (m)'); pylab.ylabel('y (m)')
wf_dist_vfm = wpg.optical_elements.WF_dist(1100, 1500, kb_clear_ap, kb_clear_ap)
defineOPD(wf_dist_vfm, os.path.join(mirror_data_dir,'mirror2.dat'), 2, ' ', 'y',
theta_kb, scale=2, stretching=kb_mirror_length/0.8)
if isIpynb:
meshT = wf_dist_vfm.mesh
opdTmp=np.array(wf_dist_vfm.arTr)[1::2].reshape(meshT.ny,meshT.nx)
figure(); pylab.imshow(opdTmp,extent=[meshT.xStart,meshT.xFin,meshT.yStart,meshT.yFin])
pylab.title('OPD [m]');pylab.xlabel('x (m)'); pylab.ylabel('y (m)')
bl0 = wpg.Beamline()
bl0.append(ap0, Use_PP(semi_analytical_treatment=0, zoom=14.4, sampling=1/1.6))
bl0.append(drift0,Use_PP(semi_analytical_treatment=0))
bl0.append(ap1, Use_PP(zoom=0.8)) #bl0.append(ap1, Use_PP(zoom=1.6, sampling=1/1.5))
bl0.append(wf_dist_om, Use_PP())
bl0.append(drift1, Use_PP(semi_analytical_treatment=1))
bl0.append(ap_kb, Use_PP(zoom = 6.4, sampling = 1/16.))#bl0.append(ap_kb, Use_PP(zoom=5.4, sampling=1/6.4))
bl0.append(hfm, Use_PP())
bl0.append(wf_dist_hfm, Use_PP())
bl0.append(drift_in_kb, Use_PP(semi_analytical_treatment=1))
bl0.append(vfm, Use_PP())
bl0.append(wf_dist_vfm, Use_PP())
bl0.append(drift_to_foc, Use_PP(semi_analytical_treatment=1))
if isIpynb:
print bl0
wpg.srwlib.srwl.SetRepresElecField(wf._srwl_wf, 'f')
sz0 = get_intensity_on_axis(wf);
wf.custom_fields['/misc/spectrum0'] = sz0
bl0.propagate(wf)
sz1 = get_intensity_on_axis(wf);
wf.custom_fields['/misc/spectrum1'] = sz1
wpg.srwlib.srwl.SetRepresElecField(wf._srwl_wf, 't')
#Resizing: decreasing Range of Horizontal and Vertical Position:
wpg.srwlib.srwl.ResizeElecField(wf._srwl_wf, 'c', [0, 0.25, 1, 0.25, 1]);
fwhm = calculate_fwhm(wf)
wf.custom_fields['/misc/xFWHM'] = fwhm['fwhm_x']
wf.custom_fields['/misc/yFWHM'] = fwhm['fwhm_y']
wf.custom_fields['/params/beamline/printout'] = str(bl0)
wf.custom_fields['/info/contact'] = [
'Name: Liubov Samoylova', 'Email: [email protected]',
'Name: Alexey Buzmakov', 'Email: [email protected]']
wf.custom_fields['/info/data_description'] = 'This dataset contains infromation about wavefront propagated through beamline (WPG and SRW frameworks).'
wf.custom_fields['/info/method_description'] = """WPG, WaveProperGator (http://github.com/samoylv/WPG)is an interactive simulation framework for coherent X-ray wavefront propagation.\nSRW, Synchrotron Radiation Workshop (http://github.com/ochubar/SRW), is a physical optics computer code for simulation of the radiation wavefront propagation through optical systems of beamlines as well as detailed characteristics of Synchrotron Radiation (SR) generated by relativistic electrons in magnetic fields of arbitrary configuration."""
wf.custom_fields['/info/package_version'] = '2014.1'
print('Saving the wavefront data after propagation:' + out_fname)
mkdir_p(os.path.dirname(out_fname))
wf.store_hdf5(out_fname)
add_history(out_fname, in_fname)
print('...done')
def read_genesis_file(genesis_out, genesis_dfl):
"""
Based on WPG/wpg/converters/genesis.py
"""
# Needed constants.
speed_of_light = const.codata.value("speed of light in vacuum")
h_eV_s = const.codata.value("Planck constant in eV s")
# Hard-code undulator period, should be read from genesis out object.
lmb_und = 2.75e-2
# Initialize empty wavefront.
wf = Wavefront()
# Get the radiation field from genesis output.
genesis_out = read_out_file(genesis_out)
genesis_radiation_field = read_dfl_file_out(out=genesis_out,
filePath=genesis_dfl,
debug=0)
# Extract geometry.
slice_count = genesis_radiation_field.Nz()
slice_spacing = genesis_radiation_field.dz
wavelength = genesis_radiation_field.xlamds
npoints = genesis_radiation_field.Nx()
# Definition of the Electric field arrays where the electric field from the GENESIS output file
# will be copied
# Setup E-field.
wf.data.arrEhor = numpy.zeros(shape=(npoints, npoints, slice_count, 2))
wf.data.arrEver = numpy.zeros(shape=(npoints, npoints, slice_count, 2))
# Fill in the fields of the wavefront object
wf.params.wEFieldUnit = "sqrt(W/mm^2)"
wf.params.photonEnergy = h_eV_s * speed_of_light / wavelength
wf.params.wDomain = "time"
wf.params.Mesh.nSlices = slice_count
wf.params.Mesh.nx = npoints
wf.params.Mesh.ny = npoints
pulse_length = (slice_count - 1) * slice_spacing / (speed_of_light)
wf.params.Mesh.sliceMin = -pulse_length / 2.0
wf.params.Mesh.sliceMax = pulse_length / 2.0
range_xy = genesis_radiation_field.Lx()
wf.params.Mesh.xMin = -range_xy / 2.0
wf.params.Mesh.xMax = range_xy / 2.0
wf.params.Mesh.yMin = -range_xy / 2.0
wf.params.Mesh.yMax = range_xy / 2.0
print("Photon energy: ", wf.params.photonEnergy, "eV")
print("Pulse length: %4.3e" % (pulse_length))
# Extract the field data from the h5 file and fill in the data of the Electric field, by calling the
# vector_grid_conversion function
wf.data.arrEhor = vector_grid_conversion(genesis_radiation_field.fld,
npoints, slice_count, range_xy,
wavelength, lmb_und)
return wf
def propagate(in_fname, out_fname, get_beamline):
"""
Propagate wavefront
:param in_file: input wavefront file
:param out_file: output file
:param get_beamline: function to build beamline
"""
print('Start propagating:' + in_fname)
wf = Wavefront()
wf.load_hdf5(in_fname)
bl0 = get_beamline()
if isIpynb:
print bl0
wpg.srwlib.srwl.SetRepresElecField(wf._srwl_wf, 'f')
sz0 = get_intensity_on_axis(wf)
wf.custom_fields['/misc/spectrum0'] = sz0
bl0.propagate(wf)
sz1 = get_intensity_on_axis(wf)
wf.custom_fields['/misc/spectrum1'] = sz1
wpg.srwlib.srwl.SetRepresElecField(wf._srwl_wf, 't')
#Resizing: decreasing Range of Horizontal and Vertical Position:
wpg.srwlib.srwl.ResizeElecField(wf._srwl_wf, 'c', [0, 0.5, 1, 0.5, 1])
fwhm = calculate_fwhm(wf)
wf.custom_fields['/misc/xFWHM'] = fwhm['fwhm_x']
wf.custom_fields['/misc/yFWHM'] = fwhm['fwhm_y']
wf.custom_fields['/params/beamline/printout'] = str(bl0)
wf.custom_fields['/info/contact'] = [
'Name: Liubov Samoylova', 'Email: [email protected]',
'Name: Alexey Buzmakov', 'Email: [email protected]'
]
wf.custom_fields[
'/info/data_description'] = 'This dataset contains infromation about wavefront propagated through beamline (WPG and SRW frameworks).'
wf.custom_fields[
'/info/method_description'] = """WPG, WaveProperGator (http://github.com/samoylv/WPG)is an interactive simulation framework for coherent X-ray wavefront propagation.\nSRW, Synchrotron Radiation Workshop (http://github.com/ochubar/SRW), is a physical optics computer code for simulation of the radiation wavefront propagation through optical systems of beamlines as well as detailed characteristics of Synchrotron Radiation (SR) generated by relativistic electrons in magnetic fields of arbitrary configuration."""
wf.custom_fields['/info/package_version'] = '2014.1'
print('Saving the wavefront data after propagation:' + out_fname)
mkdir_p(os.path.dirname(out_fname))
wf.store_hdf5(out_fname)
add_history(out_fname, in_fname)
def propagate(in_fname, out_fname, get_beamline):
"""
Propagate wavefront
:param in_file: input wavefront file
:param out_file: output file
:param get_beamline: function to build beamline
"""
print('Start propagating:' + in_fname)
wf=Wavefront()
wf.load_hdf5(in_fname)
bl0 = get_beamline()
if isIpynb:
print bl0
wpg.srwlib.srwl.SetRepresElecField(wf._srwl_wf, 'f')
sz0 = get_intensity_on_axis(wf);
wf.custom_fields['/misc/spectrum0'] = sz0
bl0.propagate(wf)
sz1 = get_intensity_on_axis(wf);
wf.custom_fields['/misc/spectrum1'] = sz1
wpg.srwlib.srwl.SetRepresElecField(wf._srwl_wf, 't')
#Resizing: decreasing Range of Horizontal and Vertical Position:
wpg.srwlib.srwl.ResizeElecField(wf._srwl_wf, 'c', [0, 0.5, 1, 0.5, 1]);
fwhm = calculate_fwhm(wf)
wf.custom_fields['/misc/xFWHM'] = fwhm['fwhm_x']
wf.custom_fields['/misc/yFWHM'] = fwhm['fwhm_y']
wf.custom_fields['/params/beamline/printout'] = str(bl0)
wf.custom_fields['/info/contact'] = [
'Name: Liubov Samoylova', 'Email: [email protected]',
'Name: Alexey Buzmakov', 'Email: [email protected]']
wf.custom_fields['/info/data_description'] = 'This dataset contains infromation about wavefront propagated through beamline (WPG and SRW frameworks).'
wf.custom_fields['/info/method_description'] = """WPG, WaveProperGator (http://github.com/samoylv/WPG)is an interactive simulation framework for coherent X-ray wavefront propagation.\nSRW, Synchrotron Radiation Workshop (http://github.com/ochubar/SRW), is a physical optics computer code for simulation of the radiation wavefront propagation through optical systems of beamlines as well as detailed characteristics of Synchrotron Radiation (SR) generated by relativistic electrons in magnetic fields of arbitrary configuration."""
wf.custom_fields['/info/package_version'] = '2014.1'
print('Saving the wavefront data after propagation:' + out_fname)
mkdir_p(os.path.dirname(out_fname))
wf.store_hdf5(out_fname)
add_history(out_fname, in_fname)
def testGaussianVsAnalytic(self, debug=False):
""" Check that propagation of a Gaussian pulse (in t,x,y) through vacuum gives the correct result, compare
to analytic solution. """
# Central photon energy.
ekev = 8.4 # Energy [keV]
# Pulse parameters.
qnC = 0.5 # e-bunch charge, [nC]
pulse_duration = 9.0e-15 # [s]
pulseEnergy = 1.5e-3 # total pulse energy, J
# Coherence time
coh_time = 0.25e-15 # [s]
# Distance in free space.
z0 = 10. # (m), position where to build the wavefront.
z1 = 20. # (m), distance to travel in free space.
z2 = z0 + z1 # distance where to build the reference wavefront.
# Beam divergence.
theta_fwhm = 2.5e-6 # rad
wlambda = 12.4*1e-10/ekev # wavelength, m
w0 = wlambda/(numpy.pi*theta_fwhm) # beam waist, m
zR = (math.pi*w0**2)/wlambda #Rayleigh range, m
fwhm_at_zR = theta_fwhm*zR #FWHM at Rayleigh range, m
sigmaAmp = w0/(2.0*math.sqrt(math.log(2.0))) #sigma of amplitude, m
if debug:
print (" *** Pulse properties ***")
print (" lambda = %4.3e m" % (wlambda) )
print (" w0 = %4.3e m" % (w0) )
print (" zR = %4.3e m" % (zR) )
print (" fwhm at zR = %4.3e m" % (fwhm_at_zR) )
print (" sigma = %4.3e m" % (sigmaAmp) )
# expected beam radius after free space drift.
expected_beam_radius = w0*math.sqrt(1.0+(z0/zR)**2)
# Number of points in each x and y dimension.
np=600
# Sampling window = 6 sigma of initial beam.
range_xy = 6.*expected_beam_radius
dx = range_xy / (np-1)
nslices = 20
#if debug:
#print (" Expected beam waist at z=%4.3f m : %4.3e m." % (z0, expected_beam_radius) )
#print ("Setting up mesh of %d points per dimension on a %4.3e x %4.3e m^2 grid with grid spacing %4.3e m." % (np, range_xy, range_xy, dx) )
# Construct srw wavefront.
srwl_wf = build_gauss_wavefront(np, np, nslices, ekev, -range_xy/2., range_xy/2.,
-range_xy/2., range_xy/2., coh_time/math.sqrt(2.),
sigmaAmp, sigmaAmp, z0,
pulseEn=pulseEnergy, pulseRange=8.)
# Convert to wpg.
wf = Wavefront(srwl_wf)
# Construct reference srw wavefront.
reference_srwl_wf = build_gauss_wavefront(np, np, nslices, ekev, -1.5*range_xy/2., 1.5*range_xy/2.,
-1.5*range_xy/2., 1.5*range_xy/2., coh_time/math.sqrt(2.),
sigmaAmp, sigmaAmp, z2,
pulseEn=pulseEnergy, pulseRange=8.)
reference_wf = Wavefront(reference_srwl_wf)
if debug:
print('*** z=%4.3e m ***' % (z0))
fwhm = calculate_fwhm(wf)
print('wf:\nfwhm_x = %4.3e\nfwhm_y = %4.3e' % (fwhm['fwhm_x'], fwhm['fwhm_y']) )
plot_t_wf(wf)
#look_at_q_space(wf)
# Construct the beamline.
beamline = Beamline()
# Add free space drift.
drift = Drift(z1)
beamline.append( drift, Use_PP(semi_analytical_treatment=0, zoom=2.0, sampling=0.5))
# Propagate
srwl.SetRepresElecField(wf._srwl_wf, 'f')
beamline.propagate(wf)
srwl.SetRepresElecField(wf._srwl_wf, 't')
fwhm = calculate_fwhm(wf)
reference_fwhm = calculate_fwhm(reference_wf)
if debug:
print('*** z=%4.3e m ***' % (z0+z1))
print('wf :\nfwhm_x = %4.3e\nfwhm_y = %4.3e' % (fwhm['fwhm_x'], fwhm['fwhm_y']) )
plot_t_wf(wf)
print('ref:\nfwhm_x = %4.3e\nfwhm_y = %4.3e' % (reference_fwhm['fwhm_x'], reference_fwhm['fwhm_y']) )
plot_t_wf(reference_wf)
#look_at_q_space(wf)
# Calculate difference
reference_norm = numpy.linalg.norm(numpy.array([reference_fwhm['fwhm_x'], reference_fwhm['fwhm_y']]))
difference_norm = numpy.linalg.norm(numpy.array([fwhm['fwhm_x'], fwhm['fwhm_y']]) - numpy.array([reference_fwhm['fwhm_x'], reference_fwhm['fwhm_y']]))
if debug:
print ("|ref_fwhm_xy| = %4.3e" % (reference_norm) )
print ("|ref_fwhm_xy - fhwm_xy| = %4.3e" % (difference_norm) )
self.assertLess(difference_norm / reference_norm, 0.01)
def propToBeamParameters(prop_output_path):
""" Utility to setup a PhotonBeamParameters instance from propagation output. """
# Check prop out exists.
if not os.path.isfile(prop_output_path):
raise IOError("File not found: %s." % (prop_output_path))
# Construct the wavefront.
wavefront = Wavefront()
wavefront.load_hdf5(prop_output_path)
pulse_energy = wpg_uti_wf.calc_pulse_energy(wavefront)
mesh = wavefront.params.Mesh
dx = (mesh.xMax - mesh.xMin) / (mesh.nx - 1)
dy = (mesh.yMax - mesh.yMin) / (mesh.ny - 1)
int0 = wavefront.get_intensity().sum(axis=0).sum(axis=0) # I(slice_num)
total_intensity = int0 * dx * dy # [J]
times = numpy.linspace(mesh.sliceMin, mesh.sliceMax, mesh.nSlices)
t = times[:-1]
dt = times[1:] - times[:-1]
It = total_intensity[:-1]
m0 = numpy.sum(It * dt)
m1 = numpy.sum(It * t * dt) / m0
m2 = numpy.sum(It * t**2 * dt) / m0
rms = math.sqrt(m2 - m1**2)
spike_fwhm_J = constants.hbar / rms
spike_fwhm_eV = spike_fwhm_J / constants.e
# Switch to energy domain
srwl.SetRepresElecField(wavefront._srwl_wf, 'f')
mesh = wavefront.params.Mesh
spectrum = wavefront.get_intensity().sum(axis=0).sum(
axis=0) # I(slice_num)
energies = numpy.linspace(mesh.sliceMin, mesh.sliceMax, mesh.nSlices)
w = energies[:-1]
dw = energies[1:] - energies[:-1]
Iw = spectrum[:-1]
m0 = numpy.sum(Iw * dw)
m1 = numpy.sum(Iw * w * dw) / m0
m2 = numpy.sum(Iw * w**2 * dw) / m0
rms = math.sqrt(m2 - m1**2)
photon_energy = m1
#spec_fwhm_eV = rms
# Extract beam diameter fwhm
xy_fwhm = wpg_uti_wf.calculate_fwhm(wavefront)
# Extract divergence
# Switch to reciprocal space
srwl.SetRepresElecField(wavefront._srwl_wf, 'a')
qxqy_fwhm = wpg_uti_wf.calculate_fwhm(wavefront)
del wavefront
beam_parameters = PhotonBeamParameters(
photon_energy=photon_energy * electronvolt,
photon_energy_relative_bandwidth=spike_fwhm_eV / photon_energy,
pulse_energy=pulse_energy * joule,
divergence=max([qxqy_fwhm['fwhm_x'], qxqy_fwhm['fwhm_y']]) / 2. *
radian,
beam_diameter_fwhm=max([xy_fwhm['fwhm_x'], xy_fwhm['fwhm_y']]) * meter,
photon_energy_spectrum_type="SASE",
)
return beam_parameters
def testGaussianReference(self, debug=False):
""" Check that propagation of a Gaussian pulse (in t,x,y) through vacuum reproduces reference data. """
# Central photon energy.
ekev = 8.4 # Energy [keV]
# Pulse parameters.
qnC = 0.5 # e-bunch charge, [nC]
pulse_duration = 9.0e-15 # [s]
pulseEnergy = 1.5e-3 # total pulse energy, J
# Coherence time
coh_time = 0.25e-15 # [s]
# Distance in free space.
z0 = 10. # (m), position where to build the wavefront.
z1 = 10. # (m), distance to travel in free space.
# Beam divergence.
theta_fwhm = 2.5e-6 # rad
wlambda = 12.4*1e-10/ekev # wavelength, m
w0 = wlambda/(numpy.pi*theta_fwhm) # beam waist, m
zR = (math.pi*w0**2)/wlambda #Rayleigh range, m
fwhm_at_zR = theta_fwhm*zR #FWHM at Rayleigh range, m
sigmaAmp = w0/(2.0*math.sqrt(math.log(2.0))) #sigma of amplitude, m
if debug:
print (" *** Pulse properties ***")
print (" lambda = %4.3e m" % (wlambda) )
print (" w0 = %4.3e m" % (w0) )
print (" zR = %4.3e m" % (zR) )
print (" fwhm at zR = %4.3e m" % (fwhm_at_zR) )
print (" sigma = %4.3e m" % (sigmaAmp) )
# expected beam radius after free space drift.
expected_beam_radius = w0*math.sqrt(1.0+(z0/zR)**2)
# Number of points in each x and y dimension.
np=400
# Sampling window = 6 sigma of initial beam.
range_xy = 6.*expected_beam_radius
dx = range_xy / (np-1)
nslices = 20
if debug:
print (" Expected beam waist at z=%4.3f m : %4.3e m." % (z0, expected_beam_radius) )
print ("Setting up mesh of %d points per dimension on a %4.3e x %4.3e m^2 grid with grid spacing %4.3e m." % (np, range_xy, range_xy, dx) )
# Construct srw wavefront.
srwl_wf = build_gauss_wavefront(np, np, nslices, ekev, -range_xy/2., range_xy/2.,
-range_xy/2., range_xy/2., coh_time/math.sqrt(2.),
sigmaAmp, sigmaAmp, z0,
pulseEn=pulseEnergy, pulseRange=8.)
# Convert to wpg.
wf = Wavefront(srwl_wf)
if debug:
print('*** z=%4.3e m ***' % (z0))
fwhm = calculate_fwhm(wf)
print('fwhm_x = %4.3e\nfwhm_y = %4.3e' % (fwhm['fwhm_x'], fwhm['fwhm_y']) )
plot_t_wf(wf)
look_at_q_space(wf)
# Construct the beamline.
beamline = Beamline()
# Add free space drift.
drift = Drift(z1)
beamline.append( drift, Use_PP(semi_analytical_treatment=1))
# Propagate
srwl.SetRepresElecField(wf._srwl_wf, 'f') # <---- switch to frequency domain
beamline.propagate(wf)
srwl.SetRepresElecField(wf._srwl_wf, 't')
if debug:
print('*** z=%4.3e m ***' % (z0+z1))
fwhm = calculate_fwhm(wf)
print('fwhm_x = %4.3e\nfwhm_y = %4.3e' % (fwhm['fwhm_x'], fwhm['fwhm_y']) )
plot_t_wf(wf)
look_at_q_space(wf)
# Get propagated wavefront data.
wf_intensity = wf.get_intensity()
# Project on t axis.
wf_onaxis = wf_intensity.sum(axis=(0,1))
# Get hash of the data.
wf_hash = hash( wf_intensity.tostring() )
# Load reference hash.
with open(TestUtilities.generateTestFilePath("reference_wf_gauss_10m.hash.txt"), 'r') as hashfile:
ref_hash = hashfile.readline()
hashfile.close()
ref_onaxis = numpy.loadtxt(TestUtilities.generateTestFilePath("reference_wf_gauss_onaxis_10m.txt"))
# Weak test.
for x,y in zip(wf_onaxis, ref_onaxis):
self.assertAlmostEqual( x, y, 14 )
# Strong test.
self.assertEqual( str(wf_hash), ref_hash)
def testGaussianReference(self, debug=False):
""" Check that propagation of a Gaussian pulse (in t,x,y) through vacuum reproduces reference data. """
# Central photon energy.
ekev = 8.4 # Energy [keV]
# Pulse parameters.
qnC = 0.5 # e-bunch charge, [nC]
pulse_duration = 9.0e-15 # [s]
pulseEnergy = 1.5e-3 # total pulse energy, J
# Coherence time
coh_time = 0.25e-15 # [s]
# Distance in free space.
z0 = 10. # (m), position where to build the wavefront.
z1 = 10. # (m), distance to travel in free space.
# Beam divergence.
theta_fwhm = 2.5e-6 # rad
wlambda = 12.4*1e-10/ekev # wavelength, m
w0 = wlambda/(numpy.pi*theta_fwhm) # beam waist, m
zR = (math.pi*w0**2)/wlambda #Rayleigh range, m
fwhm_at_zR = theta_fwhm*zR #FWHM at Rayleigh range, m
sigmaAmp = w0/(2.0*math.sqrt(math.log(2.0))) #sigma of amplitude, m
if debug:
print (" *** Pulse properties ***")
print (" lambda = %4.3e m" % (wlambda) )
print (" w0 = %4.3e m" % (w0) )
print (" zR = %4.3e m" % (zR) )
print (" fwhm at zR = %4.3e m" % (fwhm_at_zR) )
print (" sigma = %4.3e m" % (sigmaAmp) )
# expected beam radius after free space drift.
expected_beam_radius = w0*math.sqrt(1.0+(z0/zR)**2)
# Number of points in each x and y dimension.
np=400
# Sampling window = 6 sigma of initial beam.
range_xy = 6.*expected_beam_radius
dx = range_xy / (np-1)
nslices = 20
if debug:
print (" Expected beam waist at z=%4.3f m : %4.3e m." % (z0, expected_beam_radius) )
print ("Setting up mesh of %d points per dimension on a %4.3e x %4.3e m^2 grid with grid spacing %4.3e m." % (np, range_xy, range_xy, dx) )
# Construct srw wavefront.
srwl_wf = build_gauss_wavefront(np, np, nslices, ekev, -range_xy/2., range_xy/2.,
-range_xy/2., range_xy/2., coh_time/math.sqrt(2.),
sigmaAmp, sigmaAmp, z0,
pulseEn=pulseEnergy, pulseRange=8.)
# Convert to wpg.
wf = Wavefront(srwl_wf)
if debug:
print('*** z=%4.3e m ***' % (z0))
fwhm = calculate_fwhm(wf)
print('fwhm_x = %4.3e\nfwhm_y = %4.3e' % (fwhm['fwhm_x'], fwhm['fwhm_y']) )
plot_t_wf(wf)
look_at_q_space(wf)
# Construct the beamline.
beamline = Beamline()
# Add free space drift.
drift = Drift(z1)
beamline.append( drift, Use_PP(semi_analytical_treatment=1))
# Propagate
srwl.SetRepresElecField(wf._srwl_wf, 'f') # <---- switch to frequency domain
beamline.propagate(wf)
srwl.SetRepresElecField(wf._srwl_wf, 't')
if debug:
print('*** z=%4.3e m ***' % (z0+z1))
fwhm = calculate_fwhm(wf)
print('fwhm_x = %4.3e\nfwhm_y = %4.3e' % (fwhm['fwhm_x'], fwhm['fwhm_y']) )
plot_t_wf(wf)
look_at_q_space(wf)
# Get propagated wavefront data.
wf_intensity = wf.get_intensity()
# Project on t axis.
wf_onaxis = wf_intensity.sum(axis=(0,1))
# Get hash of the data.
wf_hash = hash( wf_intensity.tostring() )
# Load reference hash.
with open(TestUtilities.generateTestFilePath("reference_wf_gauss_10m.hash.txt"), 'r') as hashfile:
ref_hash = hashfile.readline()
hashfile.close()
ref_onaxis = numpy.loadtxt(TestUtilities.generateTestFilePath("reference_wf_gauss_onaxis_10m.txt"))
# Weak test.
for x,y in zip(wf_onaxis, ref_onaxis):
self.assertAlmostEqual( x, y, 14 )
# Strong test.
self.assertEqual( str(wf_hash), ref_hash)
class WavePropagator(AbstractPhotonPropagator):
"""
Class representing a photon propagator that uses wave optics.
"""
def __init__(self, parameters=None, input_path=None, output_path=None):
"""
:param parameters: Parameters steering the propagation of photons.
:type parameters: WavePropagatorParameters
:param input_path: Location of input data for the photon propagation.
:type input_path: str
:param output_path: Location of output data for the photon propagation.
:type output_path: str
"""
# DCheck (and set) parameters.
parameters = checkAndSetInstance(WavePropagatorParameters, parameters, WavePropagatorParameters() )
# Initialize base class.
super(WavePropagator, self).__init__(parameters,input_path,output_path)
def backengine(self):
""" This method drives the backengine code, in this case the WPG interface to SRW.
:return: 0 if WPG run was successful, 1 if not.
"""
# Switch to frequency representation.
srwl.SetRepresElecField(self.__wavefront._srwl_wf, 'f') # <---- switch to frequency domain
# Propagate through beamline.
self.parameters.beamline.propagate(self.__wavefront)
# Switch back to time representation.
srwl.SetRepresElecField(self.__wavefront._srwl_wf, 't')
return 0
@property
def data(self):
""" Query for the field data.
:return: The WPG wavefront data.
"""
return self.__data
def _readH5(self):
""" """
""" Private method for reading the hdf5 input and extracting the parameters and data relevant to initialize the object. """
# Check input.
try:
self.__h5 = h5py.File( self.input_path, 'r' )
except:
raise IOError( 'The input_path argument (%s) is not a path to a valid hdf5 file.' % (self.input_path) )
# Construct wpg wavefront based on input data.
self.__wavefront = Wavefront()
self.__wavefront.load_hdf5(self.input_path)
def saveH5(self):
"""
Method to save the object to a file.
:param output_path: The file where to save the wavefront data.
:type output_path: str, default 'prop_out.h5'
"""
# Write data to hdf file using wpg interface function.
self.__wavefront.store_hdf5(self.output_path)
# Write openPMD file if requested.
if self.parameters.use_opmd:
wpg_to_opmd.convertToOPMD( self.output_path )
class WavePropagator(AbstractPhotonPropagator):
"""
Class representing a photon propagator using wave optics through WPG.
"""
def __init__(self, parameters=None, input_path=None, output_path=None):
"""
Constructor for the xfel photon propagator.
@param parameters : Parameters steering the propagation of photons.
<br/><b>type</b> : dict
@param input_path : Location of input data for the photon propagation.
<br/><b>type</b> : string
@param output_path : Location of output data for the photon propagation.
<br/><b>type</b> : string
"""
# Check if beamline was given.
if isinstance(parameters, Beamline):
parameters = {'beamline' : parameters}
# Raise if no beamline in parameters.
if parameters is None or not 'beamline' in parameters.keys():
raise RuntimeError( 'The parameters argument must be an instance of wpg.Beamline or a dict containing the key "beamline" and an instance of wpg.Beamline as the corresponding value.')
# Initialize base class.
super(WavePropagator, self).__init__(parameters,input_path,output_path)
# Take reference to beamline.
self.__beamline = parameters['beamline']
def backengine(self):
""" This method drives the backengine code, in this case the WPG interface to SRW."""
# Switch to frequency representation.
srwl.SetRepresElecField(self.__wavefront._srwl_wf, 'f') # <---- switch to frequency domain
# Propagate through beamline.
self.__beamline.propagate(self.__wavefront)
# Switch back to time representation.
srwl.SetRepresElecField(self.__wavefront._srwl_wf, 't')
return 0
@property
def data(self):
""" Query for the field data. """
return self.__data
def _readH5(self):
""" """
""" Private method for reading the hdf5 input and extracting the parameters and data relevant to initialize the object. """
# Check input.
try:
self.__h5 = h5py.File( self.input_path, 'r' )
except:
raise IOError( 'The input_path argument (%s) is not a path to a valid hdf5 file.' % (self.input_path) )
# Construct wpg wavefront based on input data.
self.__wavefront = Wavefront()
self.__wavefront.load_hdf5(self.input_path)
def saveH5(self):
""" """
"""
Private method to save the object to a file.
@param output_path : The file where to save the object's data.
<br/><b>type</b> : string
<br/><b>default</b> : None
"""
# Write data to hdf file using wpg interface function.
self.__wavefront.store_hdf5(self.output_path)
class WavePropagator(AbstractPhotonPropagator):
"""
Class representing a photon propagator using wave optics through WPG.
"""
def __init__(self, parameters=None, input_path=None, output_path=None):
"""
Constructor for the xfel photon propagator.
@param parameters : Parameters steering the propagation of photons.
<br/><b>type</b> : dict
@param input_path : Location of input data for the photon propagation.
<br/><b>type</b> : string
@param output_path : Location of output data for the photon propagation.
<br/><b>type</b> : string
"""
# Check if beamline was given.
if isinstance(parameters, Beamline):
parameters = {'beamline': parameters}
# Raise if no beamline in parameters.
if parameters is None or not 'beamline' in parameters.keys():
raise RuntimeError(
'The parameters argument must be an instance of wpg.Beamline or a dict containing the key "beamline" and an instance of wpg.Beamline as the corresponding value.'
)
# Initialize base class.
super(WavePropagator, self).__init__(parameters, input_path,
output_path)
# Take reference to beamline.
self.__beamline = parameters['beamline']
def backengine(self):
""" This method drives the backengine code, in this case the WPG interface to SRW."""
# Switch to frequency representation.
srwl.SetRepresElecField(self.__wavefront._srwl_wf,
'f') # <---- switch to frequency domain
# Propagate through beamline.
self.__beamline.propagate(self.__wavefront)
# Switch back to time representation.
srwl.SetRepresElecField(self.__wavefront._srwl_wf, 't')
return 0
@property
def data(self):
""" Query for the field data. """
return self.__data
def _readH5(self):
""" """
""" Private method for reading the hdf5 input and extracting the parameters and data relevant to initialize the object. """
# Check input.
try:
self.__h5 = h5py.File(self.input_path, 'r')
except:
raise IOError(
'The input_path argument (%s) is not a path to a valid hdf5 file.'
% (self.input_path))
# Construct wpg wavefront based on input data.
self.__wavefront = Wavefront()
self.__wavefront.load_hdf5(self.input_path)
def saveH5(self):
""" """
"""
Private method to save the object to a file.
@param output_path : The file where to save the object's data.
<br/><b>type</b> : string
<br/><b>default</b> : None
"""
# Write data to hdf file using wpg interface function.
self.__wavefront.store_hdf5(self.output_path)
def constructWavefield(
ekeV,
qnC,
z1,
range_xy,
strOutputDataFolder=None,
dimension=2,
beta=0.7, # 'tunable' parameter — sets global degree of coherence
threshold=0.8, # ignore modes with eigenvalue < this value
npoints=512,
nslices=10,
display=True,
):
"""
Define Transverse Optical Modes
"""
k = 2 * np.sqrt(2 * np.log(2))
wlambda = 12.4 * 1e-10 / ekeV # wavelength
theta_fwhm = calculate_theta_fwhm_cdr(ekeV, qnC)
sigX = 12.4e-10 * k / (ekeV * 4 * np.pi * theta_fwhm)
# coherence width:
widthCoherence = beta * sigX
# get first n eigenvalues for transverse modes
n = 99
e0 = eigenvaluePartCoh(sigX, widthCoherence, range(0, n))
# keep only modes for which eigenvalue > threshold
e = e0[e0 > threshold]
if display:
plotEigenValues(e0, threshold)
# generate mode indices
modes = modes2D(9, N=len(e))
dimension = 2 # value should be 3 for 3D wavefront, 2 for 2D wavefront
wf = []
for mx, my in modes:
# define unique filename for storing results
ip = np.floor(ekeV)
frac = np.floor((ekeV - ip) * 1e3)
# build initial gaussian wavefront
if dimension == 2:
wfr0 = build_gauss_wavefront_xy(
nx=npoints,
ny=npoints,
ekev=ekeV,
xMin=-range_xy / 2,
xMax=range_xy / 2,
yMin=-range_xy / 2,
yMax=range_xy / 2,
sigX=sigX,
sigY=sigX,
d2waist=z1,
_mx=mx,
_my=my,
)
else:
# build initial 3d gaussian beam
tau = 1
# not sure if this parameter is even used - check meaning.
wfr0 = build_gauss_wavefront(
nx=npoints,
ny=npoints,
nz=nslices,
ekev=ekev,
xMin=-range_xy / 2,
xMax=range_xy / 2,
yMin=-range_xy / 2,
yMax=range_xy / 2,
tau=tau,
sigX=sigX,
sigY=sigX,
d2waist=z1,
_mx=mx,
_my=my,
)
if display == True:
print(
"dy {:.1f} um".format(
(mwf.params.Mesh.yMax - mwf.params.Mesh.yMin)
* 1e6
/ (mwf.params.Mesh.ny - 1.0)
)
)
print(
"dx {:.1f} um".format(
(mwf.params.Mesh.xMax - mwf.params.Mesh.xMin)
* 1e6
/ (mwf.params.Mesh.nx - 1.0)
)
)
plot_t_wf(mwf)
look_at_q_space(mwf)
# init WPG Wavefront helper class
mwf = Wavefront(wfr0)
# store wavefront to HDF5 file
if strOutputDataFolder:
fname0 = (
"g"
+ str(int(ip))
+ "_"
+ str(int(frac))
+ "kev"
+ "_tm"
+ str(mx)
+ str(my)
)
ifname = os.path.join(strOutputDataFolder, fname0 + ".h5")
mwf.store_hdf5(ifname)
print("Saved wavefront to HDF5 file: {}".format(ifname))
else:
ifname = None
wf.append([mx, my, ifname, mwf])
# plotWavefront(mwf, 'at '+str(z1)+' m')
# look_at_q_space(mwf)
fwhm_x = calculate_fwhm_x(mwf)
print(
"FWHMx [mm], theta_fwhm [urad]: {}, {}".format(
fwhm_x * 1e3, fwhm_x / z1 * 1e6
)
)
# show_slices_hsv(mwf, slice_numbers=None, pretitle='SLICETYSLICE')
return wf, modes, e
def backengine(self):
# check for WPG first
if not WPG_AVAILABLE:
raise ModuleNotFoundError(
'Cannot find the "WPG" module, which is required to run '
"GaussianSourceCalculator.backengine(). Is it included in PYTHONPATH?"
)
# The rms of the amplitude distribution (Gaussian)
theta = self.parameters["divergence"].value_no_conversion.to(
"radian").magnitude
E_joule = (self.parameters["photon_energy"].value_no_conversion.to(
"joule").magnitude)
E_eV = self.parameters["photon_energy"].value_no_conversion.to(
"eV").magnitude
relative_bandwidth = self.parameters[
"photon_energy_relative_bandwidth"].value
coherence_time = 2.0 * np.pi * hbar / relative_bandwidth / E_joule
pulse_energy = self.parameters["pulse_energy"].value_no_conversion
beam_waist = 2.0 * hbar * c / theta / E_joule
wavelength = 1239.8e-9 / E_eV
rayleigh_length = np.pi * beam_waist**2 / wavelength
logger.info(f"rayleigh_length = {rayleigh_length}")
beam_diameter_fwhm = (self.parameters["beam_diameter_fwhm"].
value_no_conversion.to("meter").magnitude)
beam_waist_radius = beam_diameter_fwhm / np.sqrt(2.0 * np.log(2.0))
# x-y range at beam waist.
range_xy = 30.0 * beam_waist_radius
# Set number of sampling points in x and y and number of temporal slices.
npoints = self.parameters["number_of_transverse_grid_points"].value
nslices = self.parameters["number_of_time_slices"].value
# Distance from source position.
z = self.parameters["z"].value_no_conversion.to("meter").magnitude
# Build wavefront
srwl_wf = build_gauss_wavefront(
npoints,
npoints,
nslices,
E_eV / 1.0e3,
-range_xy / 2,
range_xy / 2,
-range_xy / 2,
range_xy / 2,
coherence_time / np.sqrt(2),
beam_waist_radius / 2,
beam_waist_radius /
2, # Scaled such that fwhm comes out as demanded by parameters.
d2waist=z,
pulseEn=pulse_energy.to("joule").magnitude,
pulseRange=8.0,
)
# Correct radius of curvature.
Rx = Ry = z * np.sqrt(1.0 + (rayleigh_length / z)**2)
# Store on class.
srwl_wf.Rx = Rx
srwl_wf.Ry = Ry
key = self.output_keys[0]
filename = self.output_file_paths[0]
output_data = self.output[key]
wavefront = Wavefront(srwl_wf)
wavefront.store_hdf5(filename)
output_data.set_file(filename, WPGFormat)
return self.output
def stepwise(in_fname, get_beamline):
"""
Propagate wavefront stepwise, dumping the wavefront at every step.
:param in_file: input wavefront file
:param get_beamline: function to build beamline
"""
print("#" * 80)
print("Setup initial wavefront.")
wf = Wavefront()
# Load wavefront data.
print("Load " + in_fname)
wf.load_hdf5(in_fname)
# Get beamline.
bl0 = get_beamline()
beamline = bl0.propagation_options
if len(beamline) > 1:
raise RuntimeError("Beamline configuration not supported.")
beamline = beamline[0]
elements = beamline["optical_elements"]
options = beamline["propagation_parameters"]
if len(elements) != len(options):
raise RuntimeError("Beamline configuration not supported.")
i = 0
for element, option in zip(elements, options):
print("\n")
print("#" * 80)
print("Propagation step %d." % (i))
print("Setting up incremental beamline.")
beamline_step = Beamline()
beamline_step.append(element, option) ### <== CHECKME
# Switch to frequency domain.
wpg.srwlib.srwl.SetRepresElecField(wf._srwl_wf, "f")
# Save spectrum for later reference.
sz0 = get_intensity_on_axis(wf)
wf.custom_fields["/misc/spectrum0"] = sz0
# Propagate.
beamline_step.propagate(wf)
# Save spectrum after propagation for later reference.
sz1 = get_intensity_on_axis(wf)
wf.custom_fields["/misc/spectrum1"] = sz1
# Switch back to time domain.
wpg.srwlib.srwl.SetRepresElecField(wf._srwl_wf, "t")
incremental_filename = "%04d.h5" % (i)
print("Saving propagated wavefront to " + incremental_filename)
mkdir_p(os.path.dirname(incremental_filename))
wf.store_hdf5(incremental_filename)
print("Done with propagation step %d." % (i))
print("#" * 80)
# Increment running index.
i += 1
