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