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 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 test_simple_gauusina_propagation():
# TODO: fix propagation for coerrect results
import sys
sys.path.insert(0, "..")
import os
import wpg
from wpg.generators import build_gauss_wavefront
from wpg.beamline import Beamline
from wpg.optical_elements import Drift, Use_PP
from wpg.srwlib import srwl
import numpy as np
d2waist = 270.0
# beam parameters:
qnC = 0.1 # [nC] e-bunch charge
thetaOM = 3.6e-3
ekev = 5.0
# calculate angular divergence:
theta_fwhm = (17.2 - 6.4 * np.sqrt(qnC)) * 1e-6 / ekev**0.85
theta_rms = theta_fwhm / 2.35
sigX = 12.4e-10 / (ekev * 4 * np.pi * theta_rms)
# define limits
xmax = theta_rms * d2waist * 3.5
xmin = -xmax
ymin = xmin
ymax = xmax
nx = 300
ny = nx
nz = 3
tau = 0.12e-15
srw_wf = build_gauss_wavefront(nx, ny, nz, ekev, xmin, xmax, ymin, ymax,
tau, sigX, sigX, d2waist)
wf = wpg.Wavefront(srw_wf)
b = Beamline()
b.append(Drift(5), Use_PP())
srwl.SetRepresElecField(wf._srwl_wf, "f")
b.propagate(wf)
srwl.SetRepresElecField(wf._srwl_wf, "c")
srwl.ResizeElecField(srw_wf, "c", [0, 0.25, 1, 0.25, 1])
ti = wf.get_intensity()
out_folder = os.path.join(os.path.dirname(__file__), "tests_data")
if not os.path.exists(out_folder):
os.mkdir(out_folder)
wf_hdf5_out_file_path = os.path.join(out_folder, "my_gauss.h5")
wf.store_hdf5(wf_hdf5_out_file_path)
wf_out = wpg.Wavefront()
wf_out.load_hdf5(wf_hdf5_out_file_path)
return 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 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 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 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 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
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 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)
