def test_hisotry():
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())
# b.propagate(wf)
# srwl.ResizeElecField(srw_wf, 'c', [0, 0.25, 1, 0.25, 1])
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_history.h5")
wf.store_hdf5(wf_hdf5_out_file_path)
wf.custom_fields["/history/params/beamline/printout"] = str(b)
wf.store_hdf5(wf_hdf5_out_file_path)
wf_out = wpg.Wavefront()
wf_out.load_hdf5(wf_hdf5_out_file_path)
wf_out.custom_fields["/history/params/beamline/printout"] = str(b)
wf_out.store_hdf5(wf_hdf5_out_file_path)
return wf
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 __init__(
self,
input_path=None,
):
""" Constructor for the XFELPhotonAnalysis class.
:param input_path: Name of file or directory that contains data to analyse.
:type input_path: str
"""
print "\n Start initialization."
# Initialize base class. This takes care of parameter checking.
super(XFELPhotonAnalysis, self).__init__(input_path)
# Get wavefront file name.
wavefront = wpg.Wavefront()
print "\n Loading wavefront from %s." % (self.input_path)
wavefront.load_hdf5(self.input_path)
print " ... done."
# Init intensity.
self.__intensity = None
# Init wavefront, triggers assignment of intensity.
self.wavefront = wavefront
def __init__(self, input_path=None):
"""
:param input_path: Name of the s2e wavefront .h5 file
:type input_path: str
"""
self.input_path = input_path
wavefront = wpg.Wavefront()
wavefront.load_hdf5(self.input_path)
self.wavefront = wavefront
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 main():
d2waist = 270.
# 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 = 600
ny = nx
nz = 5
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())
b.propagate(wf)
# srwl.ResizeElecField(srw_wf, 'c', [0, 0.25, 1, 0.25, 1])
if not os.path.exists('tests_data'):
os.mkdir('tests_data')
wf_hdf5_out_file_path = os.path.join('tests_data', '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 testLoadOPMDWavefront(self):
""" Test if loading a wavefront from openpmd-hdf into a WPG structure works."""
# Get sample file.
h5_input = generateTestFilePath('prop_out/prop_out_0000011.h5')
# Convert.
convertToOPMD(h5_input)
# New file name.
opmd_h5_file = h5_input.replace(".h5", ".opmd.h5")
self.__files_to_remove.append(opmd_h5_file)
# Reconstruct the series.
series = opmd.Series(opmd_h5_file, opmd.Access_Type.read_only)
wavefront = wpg.Wavefront()
def plotIntensityMap(self, qspace=False, logscale=False):
""" Plot the integrated intensity as function of x,y or qx, qy on a colormap.
:param qspace: Whether to plot the reciprocal space intensity map (default False).
:type qspace: bool
:param logscale: Whether to plot the intensity on a logarithmic scale (z-axis) (default False).
:type logscale: bool
"""
print("\n Plotting intensity map.")
# Setup new figure.
plt.figure()
wf = self.wavefront
wf_intensity = self.intensity
# Switch to q-space if requested.
if qspace:
print("\n Switching to reciprocal space.")
srwl_wf_a = copy.deepcopy(self.wavefront._srwl_wf)
wpg.srwlib.srwl.SetRepresElecField(srwl_wf_a, 'a')
wf = wpg.Wavefront(srwl_wf_a)
print(" ... done.")
wf_intensity = wf.get_intensity()
nans = mask_nans(wf_intensity)
# Get average and time slicing.
wf_intensity = wf_intensity.sum(axis=-1)
nslices = wf.params.Mesh.nSlices
if (nslices>1):
dt = (wf.params.Mesh.sliceMax-wf.params.Mesh.sliceMin)/(nslices-1)
t0 = dt*nslices/2 + wf.params.Mesh.sliceMin
else:
t0 = (wf.params.Mesh.sliceMax+wf.params.Mesh.sliceMin)/2
# Setup a figure.
figure = plt.figure(figsize=(10, 10), dpi=100)
plt.axis('tight')
# Profile plot.
profile = plt.subplot2grid((3, 3), (1, 0), colspan=2, rowspan=2)
# Get limits.
xmin, xmax, ymax, ymin = wf.get_limits()
mn, mx = wf_intensity.min(), wf_intensity.max()
# Plot profile as 2D colorcoded map.
if logscale:
profile.imshow(wf_intensity, norm=mpl.colors.LogNorm(vmin=mn, vmax=mx), extent=[xmin*1.e6, xmax*1.e6, ymax*1.e6, ymin*1.e6], cmap="viridis")
else:
profile.imshow(wf_intensity, norm=mpl.colors.Normalize(vmin=mn, vmax=mx), extent=[xmin*1.e6, xmax*1.e6, ymax*1.e6, ymin*1.e6], cmap="viridis")
# Get x and y ranges.
x = numpy.linspace(xmin*1.e6, xmax*1.e6, wf_intensity.shape[1])
y = numpy.linspace(ymin*1.e6, ymax*1.e6, wf_intensity.shape[0])
# Labels.
if qspace:
profile.set_xlabel(r'$q_{x}\ \mathrm{(\mu rad)}$')
profile.set_ylabel(r'$q_{y}\ \mathrm{(\mu rad)}$')
else:
profile.set_xlabel(r'$x\ \mathrm{(\mu m)}$')
profile.set_ylabel(r'$y\ \mathrm{(\mu m)}$')
# x-projection plots above main plot.
x_projection = plt.subplot2grid((3, 3), (0, 0), sharex=profile, colspan=2)
x_projection.plot(x, wf_intensity.sum(axis=0), label='x projection')
# Set range according to input.
profile.set_xlim([xmin*1.e6, xmax*1.e6])
# y-projection plot right of main plot.
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, they disturb here.
plt.minorticks_off()
# Set range according to input.
profile.set_ylim([ymin*1.e6, ymax*1.e6])
# Cleanup.
if qspace:
del wf
plt.title("Real part")
plt.subplot(132)
plt.imshow(wf["imag"])
plt.colorbar()
plt.title("Imaginary part")
plt.subplot(133)
plt.imshow(np.sqrt(wf["real"] ** 2 + wf["imag"] ** 2))
plt.colorbar()
plt.title("Intensity")
## Writing EDF
# In[ ]:
wf = wpg.Wavefront()
wf.load_hdf5("../wpg/tests/tests_data/my_gauss.h5")
# In[ ]:
wf.params.Mesh.nSlices
# In[ ]:
real_part = wf.get_real_part(polarization="vertical").sum(axis=-1)
plt.imshow(real_part)
### What do with polarisation?
