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 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
def plot_lineout_from_wf(wf,
color,
label=None,
fov=1e30,
ori='V',
if_log=1,
if_norm=0):
srwl.SetRepresElecField(wf._srwl_wf, 't')
mesh = wf.params.Mesh
[xmin, xmax, ymin, ymax] = wf.get_limits()
img = wf.get_intensity().sum(axis=-1)
if ori == 'V':
# vertical plane
axis = np.linspace(ymin, ymax, mesh.ny) * 1e6
lineout = img[:, int(mesh.nx / 2)]
else:
# horizontal plane
axis = np.linspace(xmin, xmax, mesh.nx) * 1e6
lineout = img[int(mesh.ny / 2), :]
plot_lineout(axis,
lineout,
color,
label=label,
fov=fov,
ori=ori,
if_log=if_log,
if_norm=if_norm)
def get_field(wf):
srwl.SetRepresElecField(wf._srwl_wf, 't')
mesh = wf.params.Mesh
E_real = wf.get_real_part()
E_img = wf.get_imag_part()
axis_x = np.linspace(mesh.xMin, mesh.xMax, mesh.nx)
axis_y = np.linspace(mesh.yMin, mesh.yMax, mesh.ny)
return axis_x, axis_y, E_real, E_img
def get_axis(wfr, axis='x'):
if axis == 'x':
axis = np.linspace(wfr.params.Mesh.xMin, wfr.params.Mesh.xMax,
wfr.params.Mesh.nx)
elif axis == 'y':
axis = np.linspace(wfr.params.Mesh.yMin, wfr.params.Mesh.yMax,
wfr.params.Mesh.ny)
elif axis == 't':
if wfr.params.wDomain != 'time':
srwl.SetRepresElecField(wfr._srwl_wf, 't')
axis = np.linspace(wfr.params.Mesh.sliceMin, wfr.params.Mesh.sliceMax,
wfr.params.Mesh.nSlices)
if wfr.params.wDomain != 'frequency':
srwl.SetRepresElecField(wfr._srwl_wf, 't')
return axis
def plot_spatial_from_wf(wf):
srwl.SetRepresElecField(wf._srwl_wf, 't')
[xmin, xmax, ymin, ymax] = wf.get_limits()
img = wf.get_intensity().sum(axis=-1)
plt.figure()
plt.imshow(img,
cmap='jet',
extent=[xmin * 1e6, xmax * 1e6, ymin * 1e6, ymax * 1e6])
plt.colorbar()
plt.xlabel(r'x ($\mu$m)', fontsize=18)
plt.ylabel(r'y ($\mu$m)', fontsize=18)
def get_tilt(wf, ori='V'):
srwl.SetRepresElecField(wf._srwl_wf, 't')
[xmin, xmax, ymin, ymax] = wf.get_limits()
mesh = wf.params.Mesh
if ori == 'V':
axis = np.linspace(ymin, ymax, mesh.ny)
tilt = wf.get_intensity()[:, int(mesh.nx / 2), :]
else:
axis = np.linspace(xmin, xmax, mesh.nx)
tilt = wf.get_intensity()[int(mesh.ny / 2), :, :]
return axis, tilt
def _readH5(self):
"""
Private method for reading the hdf5 input and extracting the parameters and data relevant to initialize the object. """
# Import wpg only here if needed.
import wpg
from wpg.srwlib import srwl
# Construct the wave.
wavefront = wpg.Wavefront()
wavefront.load_hdf5(self.input_path)
### Switch to frequency domain and get spectrum
srwl.SetRepresElecField(wavefront._srwl_wf, 'f')
# Get mesh
mesh = wavefront.params.Mesh
dx = (mesh.xMax - mesh.xMin) / (mesh.nx - 1)
dy = (mesh.yMax - mesh.yMin) / (mesh.ny - 1)
# Get intensity and sum over x,y coordinates.
int0 = wavefront.get_intensity().sum(axis=(0, 1))
int0 = int0 * (dx * dy * 1.e6) # scale to proper unit sqrt(W/mm^2)
dSlice = 0
if mesh.nSlices > 1:
dSlice = (mesh.sliceMax - mesh.sliceMin) / (mesh.nSlices - 1)
energies = numpy.arange(mesh.nSlices) * dSlice + mesh.sliceMin
radiated_energy_per_photon_energy = int0
# Get mean of spectrum
spectrum_mean = numpy.sum(
energies * radiated_energy_per_photon_energy) / numpy.sum(
radiated_energy_per_photon_energy)
self._input_data = {}
photon_energy = wavefront.params.photonEnergy
if self.parameters.photon_energy != photon_energy:
print(
"WARNING: Parameter 'photon_energy' (%4.3e eV) not equal to source photon energy (%4.3e eV). Will proceed with source photon energy."
% (self.parameters.photon_energy, photon_energy))
if abs(spectrum_mean - photon_energy) > 1.0:
print(
"WARNING: Given photon energy (%4.3e eV) deviates from spectral mean (%4.3e) by > 1 eV. Will proceed with spectral mean."
% (photon_energy, spectrum_mean))
photon_energy = spectrum_mean
# Finally store the photon energy on the class.
self.parameters.photon_energy = photon_energy
source_data = numpy.vstack(
(energies - photon_energy,
radiated_energy_per_photon_energy)).transpose()
self._input_data['source_spectrum'] = source_data
def get_pulse_energy(wfr, mode='integrated'):
"""
calculate pulse energy (in J), number of photons and fluence in the time-domain
:param wfr: wpg.Wavefront structure
:param mode: treat integrated or pulse wavefront intensity
:return E: pulse energy statistics [(pulse energy, nphotons, fluence/mm^2), slices]
"""
J2eV = 6.24150934e18
if wfr.params.wDomain != "time":
pass
else:
srwl.SetRepresElecField(wfr._srwl_wf, 't')
x = get_axis(wfr, 'x')
y = get_axis(wfr, 'y')
t = get_axis(wfr, 't')
dx = x[1] - x[0]
dy = y[1] - y[0]
dt = t[1] - t[0]
ax = (np.max(x) * 1e-03) - (np.min(x) * 1e-03)
ay = (np.max(y) * 1e-03) - (np.min(y) * 1e-03)
if mode == 'integrated':
E = np.zeros([1, 3])
E[0, 0] = wfr.get_intensity().sum()
E[0, 0] *= dx * dy * 1e6 * dt
E[0, 1] = E[0, 0] * J2eV / wfr.params.photonEnergy
E[0, 2] = E[0, 1] / (ax * ay)
elif mode == 'pulse':
ii = wfr.get_intensity()
E = np.zeros([ii.shape[-1], 3])
for slc in range(ii.shape[-1]):
E[slc, 0] = wfr.get_intensity()[:, :, slc].sum()
E[slc, 0] *= dx * dy * 1e6 * dt
E[slc, 1] = E[slc, 0] * J2eV / wfr.params.photonEnergy
E[slc, 2] = E[slc, 1] / (ax * ay)
return E
def get_axial_power_density(wfr, spectrum=False):
if spectrum:
srwl.SetRepresElecField(wfr._srwl_wf, 'f')
intensity = wfr.get_intensity()
elif spectrum == False:
srwl.SetRepresElecField(wfr._srwl_wf, 't')
intensity = wfr.get_intensity()
mesh = wfr.params.Mesh
dx = (mesh.xMax - mesh.xMin) / (mesh.nx - 1)
dy = (mesh.yMax - mesh.yMin) / (mesh.ny - 1)
# Get center pixel numbers.
center_nx = int(mesh.nx / 2)
center_ny = int(mesh.ny / 2)
# Get time slices of intensity.
intensity = wfr.get_intensity()
# Get on-axis intensity.
int0_00 = intensity[center_ny, center_nx, :]
int0 = intensity.sum(axis=(0, 1)) * (dx * dy * 1.e6
) # amplitude units sqrt(W/mm^2)
int0max = int0.max()
# Get meaningful slices.
aw = [a[0] for a in numpy.argwhere(int0 > int0max * 0.01)]
if aw == []:
raise RuntimeError("No significant intensities found.")
dSlice = (mesh.sliceMax - mesh.sliceMin) / (mesh.nSlices - 1)
xs = numpy.arange(mesh.nSlices) * dSlice + mesh.sliceMin
xs_mf = numpy.arange(min(aw), max(aw)) * dSlice + mesh.sliceMin
return xs, int0
def get_temporal(wf):
# change the wavefront into time domain, where each slice in z represent a time
srwl.SetRepresElecField(wf._srwl_wf, 't')
mesh = wf.params.Mesh
dx = (mesh.xMax - mesh.xMin) / (mesh.nx - 1) # spatial sampling resolution
dy = (mesh.yMax - mesh.yMin) / (mesh.ny - 1)
int0 = wf.get_intensity().sum(axis=0).sum(axis=0) # intensity per slice
#int0_00 = wf.get_intensity()[int(mesh.ny/2), int(mesh.nx/2),:] # central intensity
dSlice = (mesh.sliceMax - mesh.sliceMin)/(mesh.nSlices-1) # time sampling resolution
axis_t = np.arange(mesh.nSlices) * dSlice + mesh.sliceMin # time axis
# get meaningful slices (>1% maximum intensity)
int0max = max(int0) # maximum intensity slice
aw = [a[0] for a in np.argwhere(int0>int0max*0.01)] # meaningful indices
aw = np.asarray(aw)
#int0_mean = int0[min(aw):max(aw)] # meaningful intensity slices
#axis_t = np.arange(min(aw),max(aw)) * dSlice + mesh.sliceMin # meaningful time axis
return aw, axis_t, int0
def look_at_q_space(wf, output_file=None, save='', range_x=None, range_y=None):
"""
change wavefront representation from R- to Q-space and store it in output file.
:params wf: Wavefront object in R-space representation
:params output_file: if parameter present - store propagaed wavefront to file
:params save: Whether to save the figure on disk
:type: string for filename. Empty string '' means don't save.
:default: '', do not save the figure.
:params range_x: x-axis range.
:type: float
:default: None, take entire x range.
:params range_y: y-ayis range.
:type: float
:default: None, take entire y range.
:return: propagated wavefront object:
"""
wfr = Wavefront(srwl_wavefront=wf._srwl_wf)
if not wf.params.wSpace == 'R-space':
print('space should be in R-space, but not ' + wf.params.wSpace)
return
srwl_wf = wfr._srwl_wf
srwl_wf_a = copy.deepcopy(srwl_wf)
srwl.SetRepresElecField(srwl_wf_a, 'a')
wf_a = Wavefront(srwl_wf_a)
if output_file is not None:
print('store wavefront to HDF5 file: ' + output_file+'...')
wf_a.store_hdf5(output_file)
print('done')
print(calculate_fwhm(wf_a))
plot_t_wf_a(wf_a, save=save, range_x=range_x, range_y=range_y)
return
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 plot_axial_power_density(wfr, spectrum=False, outdir=None):
""" Method to plot the on-axis power density.
:param spectrum: Whether to plot the power density in energy domain (True) or time domain (False, default).
:type spectrum: bool
"""
""" Adapted from github:Samoylv/WPG/wpg/wpg_uti_wf.integral_intensity() """
#print("\n Plotting on-axis power density.")
# Setup new figure.
plt.figure()
# Switch to frequency (energy) domain if requested.
if spectrum:
srwl.SetRepresElecField(wfr._srwl_wf, 'f')
intensity = wfr.get_intensity()
# Get dimensions.
mesh = wfr.params.Mesh
dx = (mesh.xMax - mesh.xMin) / (mesh.nx - 1)
dy = (mesh.yMax - mesh.yMin) / (mesh.ny - 1)
# Get center pixel numbers.
center_nx = int(mesh.nx / 2)
center_ny = int(mesh.ny / 2)
# Get time slices of intensity.
intensity = wfr.get_intensity()
# Get on-axis intensity.
int0_00 = intensity[center_ny, center_nx, :]
int0 = intensity.sum(axis=(0, 1)) * (dx * dy * 1.e6
) # amplitude units sqrt(W/mm^2)
int0max = int0.max()
# Get meaningful slices.
aw = [a[0] for a in numpy.argwhere(int0 > int0max * 0.01)]
if aw == []:
raise RuntimeError("No significant intensities found.")
dSlice = (mesh.sliceMax - mesh.sliceMin) / (mesh.nSlices - 1)
xs = numpy.arange(mesh.nSlices) * dSlice + mesh.sliceMin
xs_mf = numpy.arange(min(aw), max(aw)) * dSlice + mesh.sliceMin
# Switch back to time domain.
srwl.SetRepresElecField(wfr._srwl_wf, 't')
# Plot.
if (wfr.params.wDomain == 'time'):
plt.plot(xs * 1e15, int0_00)
plt.plot(xs_mf * 1e15, int0_00[min(aw):max(aw)], 'ro')
plt.title('On-Axis Power Density')
plt.xlabel('time (fs)')
plt.ylabel(r'Power density (W/mm${}^{2}$)')
else: #frequency domain
plt.plot(xs, int0_00)
plt.plot(xs_mf, int0_00[min(aw):max(aw)], 'ro')
plt.title('On-Axis Spectral Fluence')
plt.xlabel('photon energy (eV)')
plt.ylabel(r'fluence (J/eV/mm${}^{2}$)')
# Switch back to time domain.
srwl.SetRepresElecField(wfr._srwl_wf, 't')
intensity = wfr.get_intensity()
if outdir is not None:
if spectrum == True:
mode = 'spectrum'
else:
mode = 'time'
plt.savefig(outdir + "/OnAxisPowerDensity_{}.png".format(mode))
def plot_intensity_qmap(wf,
output_file=None,
save="",
range_x=None,
range_y=None,
im_aspect="equal"):
"""
change wavefront representation from R- to Q-space and plot it the resulting wavefront.
:param wf: Wavefront object in R-space representation
:param output_file: if parameter present - store wavefront in Q-space to the file
:param save: string for filename. Empty string '' means don't save.
:param range_x: x-axis range, _float_. If None, take entire x range.
:param range_y: y-ayis range, float. If None, take entire y range.
:return: propagated wavefront object:
"""
wfr = Wavefront(srwl_wavefront=wf._srwl_wf)
if not wf.params.wSpace == "R-space":
#print("space should be in R-space, but not " + wf.params.wSpace)
return
srwl_wf = wfr._srwl_wf
srwl_wf_a = copy.deepcopy(srwl_wf)
srwl.SetRepresElecField(srwl_wf_a, "a")
wf_a = Wavefront(srwl_wf_a)
if output_file is not None:
#print("store wavefront to HDF5 file: " + output_file + "...")
wf_a.store_hdf5(output_file)
#print("done")
#print(calculate_fwhm(wf_a))
# plot_t_wf_a(wf_a, save=save, range_x=range_x, range_y=range_y)
import matplotlib.pyplot as plt
wf_intensity = wf_a.get_intensity().sum(axis=-1)
plt.figure(figsize=(10, 10), dpi=100)
plt.axis("tight")
profile = plt.subplot2grid((3, 3), (1, 0), colspan=2, rowspan=2)
xmin, xmax, ymax, ymin = wf_a.get_limits()
profile.imshow(
wf_intensity,
extent=[xmin * 1.0e6, xmax * 1.0e6, ymax * 1.0e6, ymin * 1.0e6],
cmap="YlGnBu_r",
)
profile.set_aspect(im_aspect)
# [LS:2016-03-17]
# change shape dimension, otherwise, in case nx!=ny ,
# 'x, y should have the same dimension' error from py plot
# x = numpy.linspace(xmin*1.e6,xmax*1.e6,wf_intensity.shape[0])
# y = numpy.linspace(ymin*1.e6,ymax*1.e6,wf_intensity.shape[1])
x = numpy.linspace(xmin * 1.0e6, xmax * 1.0e6, wf_intensity.shape[1])
y = numpy.linspace(ymin * 1.0e6, ymax * 1.0e6, wf_intensity.shape[0])
profile.set_xlabel(r"$\mu$rad", fontsize=12)
profile.set_ylabel(r"$\mu$rad", fontsize=12)
x_projection = plt.subplot2grid((3, 3), (0, 0), sharex=profile, colspan=2)
#print(x.shape, wf_intensity.sum(axis=0).shape)
x_projection.plot(x, wf_intensity.sum(axis=0), label="x projection")
if range_x is None:
profile.set_xlim([xmin * 1.0e6, xmax * 1.0e6])
else:
profile.set_xlim([-range_x / 2.0, range_x / 2.0])
y_projection = plt.subplot2grid((3, 3), (1, 2), rowspan=2, sharey=profile)
y_projection.plot(wf_intensity.sum(axis=1), y, label="y projection")
# Hide minor tick labels.
plt.minorticks_off()
if range_y is None:
profile.set_ylim([ymin * 1.0e6, ymax * 1.0e6])
else:
profile.set_ylim([-range_y / 2.0, range_y / 2.0])
if save != "":
plt.savefig(save)
else:
plt.show()
return
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 plot_total_power(wfr, spectrum=False, outdir=None):
""" Method to plot the total power.
:param spectrum: Whether to plot the power density in energy domain (True) or time domain (False, default).
:type spectrum: bool
"""
""" Adapted from github:Samoylv/WPG/wpg/wpg_uti_wf.integral_intensity() """
#print("\n Plotting total power.")
# Setup new figure.
plt.figure()
# Switch to frequency (energy) domain if requested.
if spectrum:
#print("\n Switching to frequency domain.")
srwl.SetRepresElecField(wfr._srwl_wf, 'f')
intensity = wfr.get_intensity()
else:
intensity = wfr.get_intensity()
# Get dimensions.
mesh = wfr.params.Mesh
dx = (mesh.xMax - mesh.xMin) / (mesh.nx - 1)
dy = (mesh.yMax - mesh.yMin) / (mesh.ny - 1)
# Get intensity by integrating over transverse dimensions.
int0 = intensity.sum(axis=(0, 1))
# Scale to get unit W/mm^2
int0 = int0 * (dx * dy * 1.e6) # amplitude units sqrt(W/mm^2)
int0max = int0.max()
# Get center pixel numbers.
center_nx = int(mesh.nx / 2)
center_ny = int(mesh.ny / 2)
# Get meaningful slices.
aw = [a[0] for a in np.argwhere(int0 > int0max * 0.01)]
int0_mean = int0[min(aw):max(aw)] # meaningful range of pulse
dSlice = (mesh.sliceMax - mesh.sliceMin) / (mesh.nSlices - 1)
xs = np.arange(mesh.nSlices) * dSlice + mesh.sliceMin
xs_mf = np.arange(min(aw), max(aw)) * dSlice + mesh.sliceMin
if (wfr.params.wDomain == 'time'):
plt.plot(xs * 1e15, int0) # time axis converted to fs.
plt.plot(xs_mf * 1e15, int0_mean, 'ro')
plt.title('Power')
plt.xlabel('time (fs)')
plt.ylabel('Power (W)')
dt = (mesh.sliceMax - mesh.sliceMin) / (mesh.nSlices - 1)
#print(('Pulse energy {:1.2g} J'.format(int0_mean.sum()*dt)))
else: #frequency domain
plt.plot(xs, int0)
plt.plot(xs_mf, int0_mean, 'ro')
plt.title('Spectral Energy')
plt.xlabel('eV')
plt.ylabel('J/eV')
# Switch back to time domain.
srwl.SetRepresElecField(wfr._srwl_wf, 't')
wfr.intensity = wfr.get_intensity()
if outdir is not None:
if spectrum == True:
mode = 'spectrum'
else:
mode = 'time'
plt.savefig(outdir + "/TotalPower{}.png".format(mode))
