def test_propagate_2D_fresnel_all_circle(self):
xcalc_fft, ycalc_fft, xtheory, ytheory = self.propagate_2D_fresnel(do_plot=0,method='fft',aperture_type='square',
aperture_diameter=40e-6,
pixelsize_x=1e-6,pixelsize_y=1e-6,npixels_x=1024,npixels_y=1024,
propagation_distance=5.0,wavelength=1.24e-10)
if SRWLIB_AVAILABLE:
xcalc_srw, ycalc_srw, xtheory, ytheory = self.propagate_2D_fresnel(do_plot=0,method='srw',aperture_type='square',
aperture_diameter=40e-6,
pixelsize_x=1e-6,pixelsize_y=1e-6,npixels_x=1024,npixels_y=1024,
propagation_distance=5.0,wavelength=1.24e-10)
xcalc_srw, ycalc_convolution, xtheory, ytheory = self.propagate_2D_fresnel(do_plot=0,method='convolution',aperture_type='square',
aperture_diameter=40e-6,
pixelsize_x=1e-6,pixelsize_y=1e-6,npixels_x=1024,npixels_y=1024,
propagation_distance=5.0,wavelength=1.24e-10)
if do_plot:
x = xcalc_srw
y = numpy.vstack((ycalc_fft,ycalc_srw,ycalc_convolution))
plot_table(1e6*x,y,legend=["fft","srw","convolution"],ytitle="Intensity",xtitle="x coodinate [um]",
title="Comparison circular aperture - near field")
numpy.testing.assert_almost_equal(ycalc_fft,ycalc_srw,1)
numpy.testing.assert_almost_equal(ycalc_convolution,ycalc_srw,1)
def test_propagate_2D_fresnel_all_circle(self):
xcalc_fft, ycalc_fft, xtheory, ytheory = self.propagate_2D_fresnel(do_plot=0,method='fft',aperture_type='square',
aperture_diameter=40e-6,
pixelsize_x=1e-6,pixelsize_y=1e-6,npixels_x=1024,npixels_y=1024,
propagation_distance=5.0,wavelength=1.24e-10)
if SRWLIB_AVAILABLE:
xcalc_srw, ycalc_srw, xtheory, ytheory = self.propagate_2D_fresnel(do_plot=0,method='srw',aperture_type='square',
aperture_diameter=40e-6,
pixelsize_x=1e-6,pixelsize_y=1e-6,npixels_x=1024,npixels_y=1024,
propagation_distance=5.0,wavelength=1.24e-10)
xcalc_srw, ycalc_convolution, xtheory, ytheory = self.propagate_2D_fresnel(do_plot=0,method='convolution',aperture_type='square',
aperture_diameter=40e-6,
pixelsize_x=1e-6,pixelsize_y=1e-6,npixels_x=1024,npixels_y=1024,
propagation_distance=5.0,wavelength=1.24e-10)
if do_plot:
x = xcalc_srw
y = numpy.vstack((ycalc_fft,ycalc_srw,ycalc_convolution))
plot_table(1e6*x,y,legend=["fft","srw","convolution"],ytitle="Intensity",xtitle="x coodinate [um]",
title="Comparison circular aperture - near field")
numpy.testing.assert_almost_equal(ycalc_fft,ycalc_srw,1)
numpy.testing.assert_almost_equal(ycalc_convolution,ycalc_srw,1)
def test_lens(self):
lens_diameter = 0.002
npixels_x = 2048
pixelsize_x = lens_diameter / npixels_x
print("pixelsize: ",pixelsize_x)
pixelsize_y = pixelsize_x
npixels_y = npixels_x
wavelength = 1.24e-10
propagation_distance = 30.0
defocus_factor = 1.0 # 1.0 is at focus
propagation_steps = 1
x_fft, y_fft = self.propagation_with_lens(do_plot=0,method='fft',
propagation_steps=propagation_steps,
wavelength=wavelength,
pixelsize_x=pixelsize_x,npixels_x=npixels_x,pixelsize_y=pixelsize_y,npixels_y=npixels_y,
propagation_distance = propagation_distance, defocus_factor=defocus_factor)
if SRWLIB_AVAILABLE:
x_srw, y_srw = self.propagation_with_lens(do_plot=0,method='srw',
propagation_steps=propagation_steps,
wavelength=wavelength,
pixelsize_x=pixelsize_x,npixels_x=npixels_x,pixelsize_y=pixelsize_y,npixels_y=npixels_y,
propagation_distance = propagation_distance, defocus_factor=defocus_factor)
x_convolution, y_convolution = self.propagation_with_lens(do_plot=0,method='convolution',
propagation_steps=propagation_steps,
wavelength=wavelength,
pixelsize_x=pixelsize_x,npixels_x=npixels_x,pixelsize_y=pixelsize_y,npixels_y=npixels_y,
propagation_distance = propagation_distance, defocus_factor=defocus_factor)
if do_plot:
if SRWLIB_AVAILABLE:
x = x_fft
y = numpy.vstack((y_fft,y_srw,y_convolution))
plot_table(1e6*x,y,legend=["fft","srw","convolution"],ytitle="Intensity",xtitle="x coordinate [um]",
title="Comparison 1:1 focusing")
else:
x = x_fft
y = numpy.vstack((y_fft,y_convolution))
plot_table(1e6*x,y,legend=["fft","convolution"],ytitle="Intensity",xtitle="x coordinate [um]",
title="Comparison 1:1 focusing")
numpy.testing.assert_almost_equal(y_fft,y_convolution,1)
if SRWLIB_AVAILABLE:
numpy.testing.assert_almost_equal(y_fft,y_srw,1)
numpy.testing.assert_almost_equal(y_convolution,y_srw,1)
def test_lens(self):
lens_diameter = 0.002
npixels_x = 2048
pixelsize_x = lens_diameter / npixels_x
print("pixelsize: ",pixelsize_x)
pixelsize_y = pixelsize_x
npixels_y = npixels_x
wavelength = 1.24e-10
propagation_distance = 30.0
defocus_factor = 1.0 # 1.0 is at focus
propagation_steps = 1
x_fft, y_fft = self.propagation_with_lens(do_plot=0,method='fft',
propagation_steps=propagation_steps,
wavelength=wavelength,
pixelsize_x=pixelsize_x,npixels_x=npixels_x,pixelsize_y=pixelsize_y,npixels_y=npixels_y,
propagation_distance = propagation_distance, defocus_factor=defocus_factor)
if SRWLIB_AVAILABLE:
x_srw, y_srw = self.propagation_with_lens(do_plot=0,method='srw',
propagation_steps=propagation_steps,
wavelength=wavelength,
pixelsize_x=pixelsize_x,npixels_x=npixels_x,pixelsize_y=pixelsize_y,npixels_y=npixels_y,
propagation_distance = propagation_distance, defocus_factor=defocus_factor)
x_convolution, y_convolution = self.propagation_with_lens(do_plot=0,method='convolution',
propagation_steps=propagation_steps,
wavelength=wavelength,
pixelsize_x=pixelsize_x,npixels_x=npixels_x,pixelsize_y=pixelsize_y,npixels_y=npixels_y,
propagation_distance = propagation_distance, defocus_factor=defocus_factor)
if do_plot:
if SRWLIB_AVAILABLE:
x = x_fft
y = numpy.vstack((y_fft,y_srw,y_convolution))
plot_table(1e6*x,y,legend=["fft","srw","convolution"],ytitle="Intensity",xtitle="x coordinate [um]",
title="Comparison 1:1 focusing")
else:
x = x_fft
y = numpy.vstack((y_fft,y_convolution))
plot_table(1e6*x,y,legend=["fft","convolution"],ytitle="Intensity",xtitle="x coordinate [um]",
title="Comparison 1:1 focusing")
numpy.testing.assert_almost_equal(y_fft,y_convolution,1)
if SRWLIB_AVAILABLE:
numpy.testing.assert_almost_equal(y_fft,y_srw,1)
numpy.testing.assert_almost_equal(y_convolution,y_srw,1)
if method == "kim":
coeff = 1.0 / 4 / numpy.pi
else:
coeff = 1.0 / 2 / numpy.pi
CF = (lambdan * coeff)**2 / numpy.sqrt(Sx2 * Sz2 * Sxp2 * Szp2)
Wx = f2dot35 * Sx
Wxp = f2dot35 * Sxp
Wz = f2dot35 * Sz
Wzp = f2dot35 * Szp
print(
">>> Energy: %8.1f eV, FWHM H: %5.2f urad, FWHM V: %5.2f urad, CF: %g"
% (photon_energy, Wxp, Wzp, CF))
out[iphoton_energy, 0] = photon_energy
out[iphoton_energy, 1 + imethod] = CF
# plot(out[:,0],out[:,1],xlog=1,ylog=1)
p = plot_table(out[:, 0].T,
out[:, 1:4].T,
xlog=1,
ylog=1,
legend=methods_labels,
xtitle="photon enetgy [eV]",
ytitle="Coherent Fraction",
show=0)
plt.savefig("CF.eps")
p.show()
def main(beamline,pixels=100):
npixels_x = pixels
npixels_y = npixels_x
pixelsize_x = beamline['gapH'] / npixels_x
pixelsize_y = beamline['gapV'] / npixels_y
print("pixelsize X=%f,Y=%f: "%(pixelsize_x,pixelsize_y))
propagation_distance = beamline['distance']
#
# initialize wavefronts of dimension equal to the lens
#
wf_fft = Wavefront2D.initialize_wavefront_from_range(x_min=-pixelsize_x*npixels_x/2,x_max=pixelsize_x*npixels_x/2,
y_min=-pixelsize_y*npixels_y/2,y_max=pixelsize_y*npixels_y/2,
number_of_points=(npixels_x,npixels_y),wavelength=1e-10)
gamma = beamline['ElectronEnergy'] / (codata_mee * 1e-3)
print ("Gamma: %f \n"%(gamma))
# photon_wavelength = (1 + beamline['Kv']**2 / 2.0) / 2 / gamma**2 * beamline["PeriodID"] / beamline['harmonicID']
photon_wavelength = m2ev / beamline["photonEnergy"]
print ("Photon wavelength [A]: %g \n"%(1e10*photon_wavelength))
print ("Photon energy [eV]: %g \n"%(beamline["photonEnergy"]))
myBeam = ElectronBeam(Electron_energy=beamline['ElectronEnergy'], I_current=beamline['ElectronCurrent'])
myUndulator = MagneticStructureUndulatorPlane(K=beamline['Kv'], period_length=beamline['PeriodID'],
length=beamline['PeriodID']*beamline['NPeriods'])
XX = wf_fft.get_mesh_x()
YY = wf_fft.get_mesh_y()
X = wf_fft.get_coordinate_x()
Y = wf_fft.get_coordinate_y()
source = SourceUndulatorPlane(undulator=myUndulator,
electron_beam=myBeam, magnetic_field=None)
omega = beamline["photonEnergy"] * codata.e / codata.hbar
Nb_pts_trajectory = int(source.choose_nb_pts_trajectory(0.01,photon_frequency=omega))
print("Number of trajectory points: ",Nb_pts_trajectory)
traj_fact = TrajectoryFactory(Nb_pts=Nb_pts_trajectory,method=TRAJECTORY_METHOD_ODE,
initial_condition=None)
print("Number of trajectory points: ",traj_fact.Nb_pts)
if (traj_fact.initial_condition == None):
traj_fact.initial_condition = source.choose_initial_contidion_automatic()
print("Number of trajectory points: ",traj_fact.Nb_pts,traj_fact.initial_condition)
#print('step 2')
rad_fact = RadiationFactory(method=RADIATION_METHOD_NEAR_FIELD, photon_frequency=omega)
#print('step 3')
trajectory = traj_fact.create_from_source(source=source)
#print('step 4')
radiation = rad_fact.create_for_one_relativistic_electron(trajectory=trajectory, source=source,
XY_are_list=False,distance=beamline['distance'], X=X, Y=Y)
efield = rad_fact.calculate_electrical_field(trajectory=trajectory,source=source,
distance=beamline['distance'],X_array=XX,Y_array=YY)
tmp = efield.electrical_field()[:,:,0]
wf_fft.set_photon_energy(beamline["photonEnergy"])
wf_fft.set_complex_amplitude( tmp )
# plot
plot_image(wf_fft.get_intensity(),1e6*wf_fft.get_coordinate_x(),1e6*wf_fft.get_coordinate_y(),
xtitle="X um",ytitle="Y um",title="UND source at lens plane",show=1)
# apply lens
focal_length = propagation_distance / 2
wf_fft.apply_ideal_lens(focal_length,focal_length)
plot_image(wf_fft.get_phase(),1e6*wf_fft.get_coordinate_x(),1e6*wf_fft.get_coordinate_y(),
title="Phase just after the lens",xtitle="X um",ytitle="Y um",show=1)
wf_fft, x_fft, y_fft = propagation_to_image(wf_fft,do_plot=0,method='fft',
propagation_steps=1,
propagation_distance = propagation_distance, defocus_factor=1.0)
plot_image(wf_fft.get_intensity(),1e6*wf_fft.get_coordinate_x(),1e6*wf_fft.get_coordinate_y(),
title="Intensity at image plane",xtitle="X um",ytitle="Y um",show=1)
if do_plot:
plot_table(1e6*x_fft,y_fft,ytitle="Intensity",xtitle="x coordinate [um]",
title="Comparison 1:1 focusing ")
def main(mode_wavefront_before_lens):
lens_diameter = 0.002 # 0.001 # 0.002
if mode_wavefront_before_lens == 'Undulator with lens':
npixels_x = 512
else:
npixels_x = int(2048*1.5)
pixelsize_x = lens_diameter / npixels_x
print("pixelsize: ",pixelsize_x)
pixelsize_y = pixelsize_x
npixels_y = npixels_x
wavelength = 1.24e-10
propagation_distance = 30.0
defocus_factor = 1.0 # 1.0 is at focus
propagation_steps = 1
# for Gaussian source
sigma_x = lens_diameter / 400 # 5e-6
sigma_y = sigma_x # 5e-6
# for Hermite-Gauss, the H and V mode index (start from 0)
hm = 3
hn = 1
#
# initialize wavefronts of dimension equal to the lens
#
wf_fft = Wavefront2D.initialize_wavefront_from_range(x_min=-pixelsize_x*npixels_x/2,x_max=pixelsize_x*npixels_x/2,
y_min=-pixelsize_y*npixels_y/2,y_max=pixelsize_y*npixels_y/2,
number_of_points=(npixels_x,npixels_y),wavelength=wavelength)
wf_convolution = Wavefront2D.initialize_wavefront_from_range(x_min=-pixelsize_x*npixels_x/2,x_max=pixelsize_x*npixels_x/2,
y_min=-pixelsize_y*npixels_y/2,y_max=pixelsize_y*npixels_y/2,
number_of_points=(npixels_x,npixels_y),wavelength=wavelength)
if SRWLIB_AVAILABLE:
wf_srw = Wavefront2D.initialize_wavefront_from_range(x_min=-pixelsize_x*npixels_x/2,x_max=pixelsize_x*npixels_x/2,
y_min=-pixelsize_y*npixels_y/2,y_max=pixelsize_y*npixels_y/2,
number_of_points=(npixels_x,npixels_y),wavelength=wavelength)
#
# calculate/define wavefront at zero distance downstream the lens
#
if mode_wavefront_before_lens == 'convergent spherical':
# no need to propagate nor define lens
wf_fft.set_spherical_wave(complex_amplitude=1.0,radius=-propagation_distance)
wf_convolution.set_spherical_wave(complex_amplitude=1.0,radius=-propagation_distance)
if SRWLIB_AVAILABLE: wf_srw.set_spherical_wave(complex_amplitude=1.0,radius=-propagation_distance)
elif mode_wavefront_before_lens == 'divergent spherical with lens':
# define wavefront at zero distance upstream the lens and apply lens
focal_length = propagation_distance / 2.
wf_fft.set_spherical_wave(complex_amplitude=1.0,radius=propagation_distance)
wf_fft.apply_ideal_lens(focal_length,focal_length)
wf_convolution.set_spherical_wave(complex_amplitude=1.0,radius=propagation_distance)
wf_convolution.apply_ideal_lens(focal_length,focal_length)
if SRWLIB_AVAILABLE:
wf_srw.set_spherical_wave(complex_amplitude=1.0,radius=propagation_distance)
wf_srw.apply_ideal_lens(focal_length,focal_length)
elif mode_wavefront_before_lens == 'plane with lens':
# define wavefront at zero distance upstream the lens and apply lens
focal_length = propagation_distance
wf_fft.set_plane_wave_from_complex_amplitude(1.0+0j)
wf_fft.apply_ideal_lens(focal_length,focal_length)
wf_convolution.set_plane_wave_from_complex_amplitude(1.0+0j)
wf_convolution.apply_ideal_lens(focal_length,focal_length)
if SRWLIB_AVAILABLE:
wf_srw.set_plane_wave_from_complex_amplitude(1.0+0j)
wf_srw.apply_ideal_lens(focal_length,focal_length)
elif mode_wavefront_before_lens == 'Gaussian with lens':
# define wavefront at source point, propagate to the lens and apply lens
X = wf_fft.get_mesh_x()
Y = wf_fft.get_mesh_y()
intensity = numpy.exp( - X**2/(2*sigma_x**2)) * numpy.exp( - Y**2/(2*sigma_y**2))
wf_fft.set_complex_amplitude( numpy.sqrt(intensity) )
wf_convolution.set_complex_amplitude( numpy.sqrt(intensity) )
if SRWLIB_AVAILABLE: wf_srw.set_complex_amplitude( numpy.sqrt(intensity) )
# plot
plot_image(wf_fft.get_intensity(),1e6*wf_fft.get_coordinate_x(),1e6*wf_fft.get_coordinate_y(),
xtitle="X um (%d pixels)"%(wf_fft.size()[0]),ytitle="Y um (%d pixels)"%(wf_fft.size()[1]),title="Gaussian source",show=1)
wf_fft, tmp1, tmp2 = propagation_to_image(wf_fft,method='fft',propagation_distance=propagation_distance,
do_plot=0,plot_title="Before lens")
wf_convolution, tmp1, tmp2 = propagation_to_image(wf_convolution,method='convolution',propagation_distance=propagation_distance,
do_plot=0,plot_title="Before lens")
if SRWLIB_AVAILABLE:
wf_srw, tmp1, tmp2 = propagation_to_image(wf_srw,method='srw',propagation_distance=propagation_distance,
do_plot=0,plot_title="Before lens")
plot_image(wf_fft.get_intensity(),1e6*wf_fft.get_coordinate_x(),1e6*wf_fft.get_coordinate_y(),
xtitle="X um (%d pixels)"%(wf_fft.size()[0]),ytitle="Y um",title="Before lens fft",show=1)
plot_image(wf_convolution.get_intensity(),1e6*wf_fft.get_coordinate_x(),1e6*wf_fft.get_coordinate_y(),
xtitle="X um (%d pixels)"%(wf_convolution.size()[0]),ytitle="Y um",title="Before lens convolution",show=1)
focal_length = propagation_distance / 2
wf_fft.apply_ideal_lens(focal_length,focal_length)
wf_convolution.apply_ideal_lens(focal_length,focal_length)
if SRWLIB_AVAILABLE: wf_srw.apply_ideal_lens(focal_length,focal_length)
elif mode_wavefront_before_lens == 'Hermite with lens':
# define wavefront at source point, propagate to the lens and apply lens
X = wf_fft.get_mesh_x()
Y = wf_fft.get_mesh_y()
efield = (hermite(hm)(numpy.sqrt(2)*X/sigma_x)*numpy.exp(-X**2/sigma_x**2))**2 \
* (hermite(hn)(numpy.sqrt(2)*Y/sigma_y)*numpy.exp(-Y**2/sigma_y**2))**2
wf_fft.set_complex_amplitude( efield )
wf_convolution.set_complex_amplitude( efield )
if SRWLIB_AVAILABLE: wf_srw.set_complex_amplitude( efield )
# plot
plot_image(wf_fft.get_intensity(),1e6*wf_fft.get_coordinate_x(),1e6*wf_fft.get_coordinate_y(),
xtitle="X um (%d pixels)"%(wf_fft.size()[0]),ytitle="Y um (%d pixels)"%(wf_fft.size()[0]),title="Hermite-Gauss source",show=1)
wf_fft, tmp1, tmp2 = propagation_to_image(wf_fft,method='fft',propagation_distance=propagation_distance,
do_plot=0,plot_title="Before lens")
wf_convolution, tmp1, tmp2 = propagation_to_image(wf_convolution,method='convolution',propagation_distance=propagation_distance,
do_plot=0,plot_title="Before lens")
if SRWLIB_AVAILABLE:
wf_srw, tmp1, tmp2 = propagation_to_image(wf_srw,method='srw',propagation_distance=propagation_distance,
do_plot=0,plot_title="Before lens")
plot_image(wf_fft.get_intensity(),1e6*wf_fft.get_coordinate_x(),1e6*wf_fft.get_coordinate_y(),
xtitle="X um (%d pixels)"%(wf_fft.size()[0]),ytitle="Y um (%d pixels)"%(wf_fft.size()[0]),title="Before lens fft",show=1)
plot_image(wf_convolution.get_intensity(),1e6*wf_fft.get_coordinate_x(),1e6*wf_fft.get_coordinate_y(),
xtitle="X um (%d pixels)"%(wf_fft.size()[0]),ytitle="Y um (%d pixels)"%(wf_fft.size()[0]),title="Before lens convolution",show=1)
focal_length = propagation_distance / 2
wf_fft.apply_ideal_lens(focal_length,focal_length)
wf_convolution.apply_ideal_lens(focal_length,focal_length)
if SRWLIB_AVAILABLE: wf_srw.apply_ideal_lens(focal_length,focal_length)
elif mode_wavefront_before_lens == 'Undulator with lens':
beamline = {}
beamline['ElectronCurrent'] = 0.2
beamline['ElectronEnergy'] = 6.0
beamline['Kv'] = 1.68 # 1.87
beamline['NPeriods'] = 111 # 14
beamline['PeriodID'] = 0.018 # 0.035
beamline['distance'] = propagation_distance
gamma = beamline['ElectronEnergy'] / (codata_mee * 1e-3)
print ("Gamma: %f \n"%(gamma))
resonance_wavelength = (1 + beamline['Kv']**2 / 2.0) / 2 / gamma**2 * beamline["PeriodID"]
resonance_energy = m2ev / resonance_wavelength
print ("Resonance wavelength [A]: %g \n"%(1e10*resonance_wavelength))
print ("Resonance energy [eV]: %g \n"%(resonance_energy))
# red shift 100 eV
resonance_energy = resonance_energy - 100
myBeam = ElectronBeam(Electron_energy=beamline['ElectronEnergy'], I_current=beamline['ElectronCurrent'])
myUndulator = MagneticStructureUndulatorPlane(K=beamline['Kv'], period_length=beamline['PeriodID'],
length=beamline['PeriodID']*beamline['NPeriods'])
XX = wf_fft.get_mesh_x()
YY = wf_fft.get_mesh_y()
X = wf_fft.get_coordinate_x()
Y = wf_fft.get_coordinate_y()
source = SourceUndulatorPlane(undulator=myUndulator,
electron_beam=myBeam, magnetic_field=None)
omega = resonance_energy * codata.e / codata.hbar
Nb_pts_trajectory = int(source.choose_nb_pts_trajectory(0.01,photon_frequency=omega))
print("Number of trajectory points: ",Nb_pts_trajectory)
traj_fact = TrajectoryFactory(Nb_pts=Nb_pts_trajectory,method=TRAJECTORY_METHOD_ODE,
initial_condition=None)
print("Number of trajectory points: ",traj_fact.Nb_pts)
if (traj_fact.initial_condition == None):
traj_fact.initial_condition = source.choose_initial_contidion_automatic()
print("Number of trajectory points: ",traj_fact.Nb_pts,traj_fact.initial_condition)
#print('step 2')
rad_fact = RadiationFactory(method=RADIATION_METHOD_NEAR_FIELD, photon_frequency=omega)
#print('step 3')
trajectory = traj_fact.create_from_source(source=source)
#print('step 4')
radiation = rad_fact.create_for_one_relativistic_electron(trajectory=trajectory, source=source,
XY_are_list=False,distance=beamline['distance'], X=X, Y=Y)
efield = rad_fact.calculate_electrical_field(trajectory=trajectory,source=source,
distance=beamline['distance'],X_array=XX,Y_array=YY)
tmp = efield.electrical_field()[:,:,0]
wf_fft.set_photon_energy(resonance_energy)
wf_convolution.set_photon_energy(resonance_energy)
if SRWLIB_AVAILABLE: wf_srw.set_photon_energy(resonance_energy)
wf_fft.set_complex_amplitude( tmp )
wf_convolution.set_complex_amplitude( numpy.sqrt(tmp) )
if SRWLIB_AVAILABLE: wf_srw.set_complex_amplitude( numpy.sqrt(tmp) )
# plot
plot_image(wf_fft.get_intensity(),1e6*wf_fft.get_coordinate_x(),1e6*wf_fft.get_coordinate_y(),
xtitle="X um (%d pixels)"%(wf_fft.size()[0]),ytitle="Y um (%d pixels)"%(wf_fft.size()[0]),title="UND source at lens plane",show=1)
# apply lens
focal_length = propagation_distance / 2
wf_fft.apply_ideal_lens(focal_length,focal_length)
wf_convolution.apply_ideal_lens(focal_length,focal_length)
if SRWLIB_AVAILABLE: wf_srw.apply_ideal_lens(focal_length,focal_length)
else:
raise Exception("Unknown mode")
plot_image(wf_fft.get_phase(),1e6*wf_fft.get_coordinate_x(),1e6*wf_fft.get_coordinate_y(),
title="Phase just after the lens",xtitle="X um (%d pixels)"%(wf_fft.size()[0]),ytitle="Y um (%d pixels)"%(wf_fft.size()[0]),show=1)
wf_fft, x_fft, y_fft = propagation_to_image(wf_fft,do_plot=0,method='fft',
propagation_steps=propagation_steps,
propagation_distance = propagation_distance, defocus_factor=defocus_factor)
wf_convolution, x_convolution, y_convolution = propagation_to_image(wf_convolution,do_plot=0,method='convolution',
propagation_steps=propagation_steps,
propagation_distance = propagation_distance, defocus_factor=defocus_factor)
if SRWLIB_AVAILABLE:
wf_srw, x_srw, y_srw = propagation_to_image(wf_srw,do_plot=0,method='srw',
propagation_steps=propagation_steps,
propagation_distance = propagation_distance, defocus_factor=defocus_factor)
plot_image(wf_fft.get_intensity(),1e6*wf_fft.get_coordinate_x(),1e6*wf_fft.get_coordinate_y(),
title="Intensity at image plane",xtitle="X um (%d pixels)"%(wf_fft.size()[0]),ytitle="Y um (%d pixels)"%(wf_fft.size()[0]),show=1)
if do_plot:
if SRWLIB_AVAILABLE:
x = x_fft
y = numpy.vstack((y_fft,y_srw,y_convolution))
plot_table(1e6*x,y,legend=["fft","srw","convolution"],ytitle="Intensity",xtitle="x coordinate [um]",
title="Comparison 1:1 focusing "+mode_wavefront_before_lens)
else:
x = x_fft
y = numpy.vstack((y_fft,y_convolution))
plot_table(1e6*x,y,legend=["fft","convolution"],ytitle="Intensity",xtitle="x coordinate [um]",
title="Comparison 1:1 focusing "+mode_wavefront_before_lens)
input_array = numpy.loadtxt("aps_axo_influence_functions2019.dat")
abscissas = input_array[:, 0].copy()
print("abscisas: ", abscissas)
# prepare input format for orthonormalize_a
col19 = input_array[:, 0].copy() * 0 + 1
col20 = numpy.linspace(-1, 1, input_array.shape[0])
a = []
a.append({'a': col19, 'total_squared': 0})
a.append({'a': col20, 'total_squared': 0})
for i in [9, 10, 8, 11, 7, 12, 6, 13, 5, 14, 4, 15, 3, 16, 2, 17, 1, 18]:
a.append({'a': input_array[:, i], 'total_squared': 0})
plot_table(abscissas, input_array[:, 1:].T, title="influence functions")
# prepare a Gaussian (data to fit)
sigma = (abscissas[-1] - abscissas[0])
u = 15 * numpy.exp(-abscissas**2 / 2 / sigma)
mask = None # u
# compute the basis
b, matrix = orthonormalize_a(a, mask=mask)
print(">>>>> matrix", matrix.shape)
# plot basis
b_array = numpy.zeros((input_array.shape[0], 20))
def fit_correction(filein, fileout="", calculate=False):
from axo import orthonormalize_a, linear_2dgsfit1, linear_basis
from srxraylib.plot.gol import plot, plot_table
# loads file with data to fit
input_array = numpy.loadtxt("aps_axo_influence_functions2019.dat")
abscissas = input_array[:, 0].copy()
print("abscisas: ", abscissas)
tmp = numpy.loadtxt(filein)
print(">>>>>>>>>>>>>>>>>>", tmp.shape)
plot(tmp[:, 0], tmp[:, 1], title="data to fit")
u = numpy.interp(abscissas, 1000 * tmp[:, 0], tmp[:, 1])
plot(abscissas, u, title="Result of fit")
sigma = (abscissas[-1] - abscissas[0])
g = 15 * numpy.exp(- abscissas ** 2 / 2 / sigma)
mask = None # g
if calculate:
# prepare input format for orthonormalize_a
col19 = input_array[:, 0].copy() * 0 + 1
col20 = numpy.linspace(-1,1,input_array.shape[0])
a = []
a.append({'a': col19, 'total_squared': 0})
a.append({'a': col20, 'total_squared': 0})
for i in [9, 10, 8, 11, 7, 12, 6, 13, 5, 14, 4, 15, 3, 16, 2, 17, 1, 18]:
a.append({'a': input_array[:, i], 'total_squared':0})
plot_table(abscissas, input_array[:, 1:].T, title="influence functions")
# compute the basis
b, matrix = orthonormalize_a(a, mask=mask)
# plot basis
b_array = numpy.zeros((input_array.shape[0],20))
for i in range(20):
b_array[:,i] = b[i]["a"]
plot_table(abscissas, b_array.T, title="basis functions")
numpy.savetxt("aps_axo_orthonormal_functions2019.dat",b_array)
print("File written to disk aps_axo_orthonormal_functions2019.dat")
else:
b_array = numpy.loadtxt("aps_axo_orthonormal_functions2019.dat")
print(">>>>>>>>>>>>>>>>>>>>>>>>>>>>>>",b_array.shape)
b = []
for i in range(b_array.shape[1]):
b.append({'a': b_array[:, i], 'total_squared':(b_array[:, i]**2).sum()})
# perform the fit
v = linear_2dgsfit1(u, b, mask=mask)
print("coefficients: ",v)
# evaluate the fitted data form coefficients and basis
y = linear_basis(v, b)
plot(abscissas,u,abscissas,y,legend=["Data","Fit"])
if fileout != "":
f = open(fileout,'w')
for i in range(abscissas.size):
f.write("%g %g \n"%(1e-3*abscissas[i],y[i]))
f.close()
print("File %s written to disk"%fileout)
return v
def main(mode_wavefront_before_lens):
lens_diameter = 0.002 # 0.001 # 0.002
if mode_wavefront_before_lens == 'Undulator with lens':
npixels_x = 512
else:
npixels_x = 2048*1.5
pixelsize_x = lens_diameter / npixels_x
print("pixelsize: ",pixelsize_x)
pixelsize_y = pixelsize_x
npixels_y = npixels_x
wavelength = 1.24e-10
propagation_distance = 30.0
defocus_factor = 1.0 # 1.0 is at focus
propagation_steps = 1
# for Gaussian source
sigma_x = lens_diameter / 400 # 5e-6
sigma_y = sigma_x # 5e-6
# for Hermite-Gauss, the H and V mode index (start from 0)
hm = 3
hn = 1
#
# initialize wavefronts of dimension equal to the lens
#
wf_fft = Wavefront2D.initialize_wavefront_from_range(x_min=-pixelsize_x*npixels_x/2,x_max=pixelsize_x*npixels_x/2,
y_min=-pixelsize_y*npixels_y/2,y_max=pixelsize_y*npixels_y/2,
number_of_points=(npixels_x,npixels_y),wavelength=wavelength)
wf_convolution = Wavefront2D.initialize_wavefront_from_range(x_min=-pixelsize_x*npixels_x/2,x_max=pixelsize_x*npixels_x/2,
y_min=-pixelsize_y*npixels_y/2,y_max=pixelsize_y*npixels_y/2,
number_of_points=(npixels_x,npixels_y),wavelength=wavelength)
if SRWLIB_AVAILABLE:
wf_srw = Wavefront2D.initialize_wavefront_from_range(x_min=-pixelsize_x*npixels_x/2,x_max=pixelsize_x*npixels_x/2,
y_min=-pixelsize_y*npixels_y/2,y_max=pixelsize_y*npixels_y/2,
number_of_points=(npixels_x,npixels_y),wavelength=wavelength)
#
# calculate/define wavefront at zero distance downstream the lens
#
if mode_wavefront_before_lens == 'convergent spherical':
# no need to propagate nor define lens
wf_fft.set_spherical_wave(complex_amplitude=1.0,radius=-propagation_distance)
wf_convolution.set_spherical_wave(complex_amplitude=1.0,radius=-propagation_distance)
if SRWLIB_AVAILABLE: wf_srw.set_spherical_wave(complex_amplitude=1.0,radius=-propagation_distance)
elif mode_wavefront_before_lens == 'divergent spherical with lens':
# define wavefront at zero distance upstream the lens and apply lens
focal_length = propagation_distance / 2.
wf_fft.set_spherical_wave(complex_amplitude=1.0,radius=propagation_distance)
wf_fft.apply_ideal_lens(focal_length,focal_length)
wf_convolution.set_spherical_wave(complex_amplitude=1.0,radius=propagation_distance)
wf_convolution.apply_ideal_lens(focal_length,focal_length)
if SRWLIB_AVAILABLE:
wf_srw.set_spherical_wave(complex_amplitude=1.0,radius=propagation_distance)
wf_srw.apply_ideal_lens(focal_length,focal_length)
elif mode_wavefront_before_lens == 'plane with lens':
# define wavefront at zero distance upstream the lens and apply lens
focal_length = propagation_distance
wf_fft.set_plane_wave_from_complex_amplitude(1.0+0j)
wf_fft.apply_ideal_lens(focal_length,focal_length)
wf_convolution.set_plane_wave_from_complex_amplitude(1.0+0j)
wf_convolution.apply_ideal_lens(focal_length,focal_length)
if SRWLIB_AVAILABLE:
wf_srw.set_plane_wave_from_complex_amplitude(1.0+0j)
wf_srw.apply_ideal_lens(focal_length,focal_length)
elif mode_wavefront_before_lens == 'Gaussian with lens':
# define wavefront at source point, propagate to the lens and apply lens
X = wf_fft.get_mesh_x()
Y = wf_fft.get_mesh_y()
intensity = numpy.exp( - X**2/(2*sigma_x**2)) * numpy.exp( - Y**2/(2*sigma_y**2))
wf_fft.set_complex_amplitude( numpy.sqrt(intensity) )
wf_convolution.set_complex_amplitude( numpy.sqrt(intensity) )
if SRWLIB_AVAILABLE: wf_srw.set_complex_amplitude( numpy.sqrt(intensity) )
# plot
plot_image(wf_fft.get_intensity(),1e6*wf_fft.get_coordinate_x(),1e6*wf_fft.get_coordinate_y(),
xtitle="X um",ytitle="Y um",title="Gaussian source",show=1)
wf_fft, tmp1, tmp2 = propagation_to_image(wf_fft,method='fft',propagation_distance=propagation_distance,
do_plot=0,plot_title="Before lens")
wf_convolution, tmp1, tmp2 = propagation_to_image(wf_convolution,method='convolution',propagation_distance=propagation_distance,
do_plot=0,plot_title="Before lens")
if SRWLIB_AVAILABLE:
wf_srw, tmp1, tmp2 = propagation_to_image(wf_srw,method='srw',propagation_distance=propagation_distance,
do_plot=0,plot_title="Before lens")
plot_image(wf_fft.get_intensity(),1e6*wf_fft.get_coordinate_x(),1e6*wf_fft.get_coordinate_y(),
xtitle="X um",ytitle="Y um",title="Before lens fft",show=1)
plot_image(wf_convolution.get_intensity(),1e6*wf_fft.get_coordinate_x(),1e6*wf_fft.get_coordinate_y(),
xtitle="X um",ytitle="Y um",title="Before lens convolution",show=1)
focal_length = propagation_distance / 2
wf_fft.apply_ideal_lens(focal_length,focal_length)
wf_convolution.apply_ideal_lens(focal_length,focal_length)
if SRWLIB_AVAILABLE: wf_srw.apply_ideal_lens(focal_length,focal_length)
elif mode_wavefront_before_lens == 'Hermite with lens':
# define wavefront at source point, propagate to the lens and apply lens
X = wf_fft.get_mesh_x()
Y = wf_fft.get_mesh_y()
efield = (hermite(hm)(numpy.sqrt(2)*X/sigma_x)*numpy.exp(-X**2/sigma_x**2))**2 \
* (hermite(hn)(numpy.sqrt(2)*Y/sigma_y)*numpy.exp(-Y**2/sigma_y**2))**2
wf_fft.set_complex_amplitude( efield )
wf_convolution.set_complex_amplitude( efield )
if SRWLIB_AVAILABLE: wf_srw.set_complex_amplitude( efield )
# plot
plot_image(wf_fft.get_intensity(),1e6*wf_fft.get_coordinate_x(),1e6*wf_fft.get_coordinate_y(),
xtitle="X um",ytitle="Y um",title="Hermite-Gauss source",show=1)
wf_fft, tmp1, tmp2 = propagation_to_image(wf_fft,method='fft',propagation_distance=propagation_distance,
do_plot=0,plot_title="Before lens")
wf_convolution, tmp1, tmp2 = propagation_to_image(wf_convolution,method='convolution',propagation_distance=propagation_distance,
do_plot=0,plot_title="Before lens")
if SRWLIB_AVAILABLE:
wf_srw, tmp1, tmp2 = propagation_to_image(wf_srw,method='srw',propagation_distance=propagation_distance,
do_plot=0,plot_title="Before lens")
plot_image(wf_fft.get_intensity(),1e6*wf_fft.get_coordinate_x(),1e6*wf_fft.get_coordinate_y(),
xtitle="X um",ytitle="Y um",title="Before lens fft",show=1)
plot_image(wf_convolution.get_intensity(),1e6*wf_fft.get_coordinate_x(),1e6*wf_fft.get_coordinate_y(),
xtitle="X um",ytitle="Y um",title="Before lens convolution",show=1)
focal_length = propagation_distance / 2
wf_fft.apply_ideal_lens(focal_length,focal_length)
wf_convolution.apply_ideal_lens(focal_length,focal_length)
if SRWLIB_AVAILABLE: wf_srw.apply_ideal_lens(focal_length,focal_length)
elif mode_wavefront_before_lens == 'Undulator with lens':
beamline = {}
# beamline['name'] = "ESRF_NEW_OB"
# beamline['ElectronBeamDivergenceH'] = 5.2e-6 # these values are not used (zero emittance)
# beamline['ElectronBeamDivergenceV'] = 1.4e-6 # these values are not used (zero emittance)
# beamline['ElectronBeamSizeH'] = 27.2e-6 # these values are not used (zero emittance)
# beamline['ElectronBeamSizeV'] = 3.4e-6 # these values are not used (zero emittance)
# beamline['ElectronEnergySpread'] = 0.001 # these values are not used (zero emittance)
beamline['ElectronCurrent'] = 0.2
beamline['ElectronEnergy'] = 6.0
beamline['Kv'] = 1.68 # 1.87
beamline['NPeriods'] = 111 # 14
beamline['PeriodID'] = 0.018 # 0.035
beamline['distance'] = propagation_distance
# beamline['gapH'] = pixelsize_x*npixels_x
# beamline['gapV'] = pixelsize_x*npixels_x
gamma = beamline['ElectronEnergy'] / (codata_mee * 1e-3)
print ("Gamma: %f \n"%(gamma))
resonance_wavelength = (1 + beamline['Kv']**2 / 2.0) / 2 / gamma**2 * beamline["PeriodID"]
resonance_energy = m2ev / resonance_wavelength
print ("Resonance wavelength [A]: %g \n"%(1e10*resonance_wavelength))
print ("Resonance energy [eV]: %g \n"%(resonance_energy))
# red shift 100 eV
resonance_energy = resonance_energy - 100
myBeam = ElectronBeam(Electron_energy=beamline['ElectronEnergy'], I_current=beamline['ElectronCurrent'])
myUndulator = MagneticStructureUndulatorPlane(K=beamline['Kv'], period_length=beamline['PeriodID'],
length=beamline['PeriodID']*beamline['NPeriods'])
XX = wf_fft.get_mesh_x()
YY = wf_fft.get_mesh_y()
X = wf_fft.get_coordinate_x()
Y = wf_fft.get_coordinate_y()
source = SourceUndulatorPlane(undulator=myUndulator,
electron_beam=myBeam, magnetic_field=None)
omega = resonance_energy * codata.e / codata.hbar
Nb_pts_trajectory = int(source.choose_nb_pts_trajectory(0.01,photon_frequency=omega))
print("Number of trajectory points: ",Nb_pts_trajectory)
traj_fact = TrajectoryFactory(Nb_pts=Nb_pts_trajectory,method=TRAJECTORY_METHOD_ODE,
initial_condition=None)
print("Number of trajectory points: ",traj_fact.Nb_pts)
if (traj_fact.initial_condition == None):
traj_fact.initial_condition = source.choose_initial_contidion_automatic()
print("Number of trajectory points: ",traj_fact.Nb_pts,traj_fact.initial_condition)
#print('step 2')
rad_fact = RadiationFactory(method=RADIATION_METHOD_NEAR_FIELD, photon_frequency=omega)
#print('step 3')
trajectory = traj_fact.create_from_source(source=source)
#print('step 4')
radiation = rad_fact.create_for_one_relativistic_electron(trajectory=trajectory, source=source,
XY_are_list=False,distance=beamline['distance'], X=X, Y=Y)
efield = rad_fact.calculate_electrical_field(trajectory=trajectory,source=source,
distance=beamline['distance'],X_array=XX,Y_array=YY)
tmp = efield.electrical_field()[:,:,0]
wf_fft.set_photon_energy(resonance_energy)
wf_convolution.set_photon_energy(resonance_energy)
if SRWLIB_AVAILABLE: wf_srw.set_photon_energy(resonance_energy)
wf_fft.set_complex_amplitude( tmp )
wf_convolution.set_complex_amplitude( numpy.sqrt(tmp) )
if SRWLIB_AVAILABLE: wf_srw.set_complex_amplitude( numpy.sqrt(tmp) )
# plot
plot_image(wf_fft.get_intensity(),1e6*wf_fft.get_coordinate_x(),1e6*wf_fft.get_coordinate_y(),
xtitle="X um",ytitle="Y um",title="UND source at lens plane",show=1)
# apply lens
focal_length = propagation_distance / 2
wf_fft.apply_ideal_lens(focal_length,focal_length)
wf_convolution.apply_ideal_lens(focal_length,focal_length)
if SRWLIB_AVAILABLE: wf_srw.apply_ideal_lens(focal_length,focal_length)
else:
raise Exception("Unknown mode")
plot_image(wf_fft.get_phase(),1e6*wf_fft.get_coordinate_x(),1e6*wf_fft.get_coordinate_y(),
title="Phase just after the lens",xtitle="X um",ytitle="Y um",show=1)
wf_fft, x_fft, y_fft = propagation_to_image(wf_fft,do_plot=0,method='fft',
propagation_steps=propagation_steps,
propagation_distance = propagation_distance, defocus_factor=defocus_factor)
wf_convolution, x_convolution, y_convolution = propagation_to_image(wf_convolution,do_plot=0,method='convolution',
propagation_steps=propagation_steps,
propagation_distance = propagation_distance, defocus_factor=defocus_factor)
if SRWLIB_AVAILABLE:
wf_srw, x_srw, y_srw = propagation_to_image(wf_srw,do_plot=0,method='srw',
propagation_steps=propagation_steps,
propagation_distance = propagation_distance, defocus_factor=defocus_factor)
plot_image(wf_fft.get_intensity(),1e6*wf_fft.get_coordinate_x(),1e6*wf_fft.get_coordinate_y(),
title="Intensity at image plane",xtitle="X um",ytitle="Y um",show=1)
if do_plot:
if SRWLIB_AVAILABLE:
x = x_fft
y = numpy.vstack((y_fft,y_srw,y_convolution))
plot_table(1e6*x,y,legend=["fft","srw","convolution"],ytitle="Intensity",xtitle="x coordinate [um]",
title="Comparison 1:1 focusing "+mode_wavefront_before_lens)
else:
x = x_fft
y = numpy.vstack((y_fft,y_convolution))
plot_table(1e6*x,y,legend=["fft","convolution"],ytitle="Intensity",xtitle="x coordinate [um]",
title="Comparison 1:1 focusing "+mode_wavefront_before_lens)
a = []
col19 = input_array[:, 0].copy() * 0 + 1
col20 = numpy.linspace(-1, 1, input_array.shape[0])
a.append({'a': col19, 'total_squared': 0})
a.append({'a': col20, 'total_squared': 0})
for i in [9, 10, 8, 11, 7, 12, 6, 13, 5, 14, 4, 15, 3, 16, 2, 17, 1, 18]:
a.append({'a': input_array[:, i], 'total_squared': 0})
a_array = numpy.zeros((input_array.shape[0], 20))
for i in range(20):
a_array[:, i] = a[i]["a"]
print(a_array.shape, abscissas.shape)
plot_table(abscissas, a_array[:, :].T, title="input data")
#
b, matrix = orthonormalize_a(a, mask=None)
#
# print(matrix.shape)
#
b_array = numpy.zeros((a_array.shape[0], 20))
for i in range(20):
b_array[:, i] = b[i]["a"]
# #
# plot(a_array[:, 0], b_array[:, 18])
print(b)
legend = []
for i in range(20):