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))