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?