def test_polarization_interpolation(self):
print("# ")
print("# Tests polarization with interpolation ")
print("# ")
#
# from arrays
#
x = numpy.linspace(-100, 100, 50)
y = numpy.linspace(-50, 50, 200)
X = numpy.outer(x, numpy.ones_like(y))
Y = numpy.outer(numpy.ones_like(x), y)
sigma = 10
Z = numpy.exp(-(X**2 + Y**2) / 2 / sigma**2) * 1j
Zp = 0.3333 * numpy.exp(-(X**2 + Y**2) / 2 / sigma**2) * 1j
wf2 = Wavefront2D.initialize_wavefront_from_arrays(x, y, Z, Zp)
ca1 = wf2.get_complex_amplitude(polarization=Polarization.SIGMA)
ca2 = wf2.get_interpolated_complex_amplitude(
X + 1e-3, Y + 1e-4, polarization=Polarization.SIGMA)
assert_array_almost_equal(ca1, ca2, 4)
ca1pi = wf2.get_complex_amplitude(polarization=Polarization.PI)
ca2pi = wf2.get_interpolated_complex_amplitude(
X + 1e-3, Y + 1e-4, polarization=Polarization.PI)
assert_array_almost_equal(ca1pi, ca2pi, 4)
def propagate_2D_fresnel_convolution(wavefront, propagation_distance,shift_half_pixel=1):
"""
2D Fresnel propagator using convolution via Fourier transform
:param wavefront:
:param propagation_distance: propagation distance
:param shift_half_pixel: set to 1 to shift half pixel (recommended using an even number of pixels) Set as default.
:return: a new 2D wavefront object with propagated wavefront
"""
from scipy.signal import fftconvolve
wavelength = wavefront.get_wavelength()
X = wavefront.get_mesh_x()
Y = wavefront.get_mesh_y()
if shift_half_pixel:
x = wavefront.get_coordinate_x()
y = wavefront.get_coordinate_y()
X += 0.5 * numpy.abs( x[0] - x[1] )
Y += 0.5 * numpy.abs( y[0] - y[1] )
kernel = numpy.exp(1j*2*numpy.pi/wavefront.get_wavelength() *
(X**2 + Y**2) / 2 / propagation_distance)
kernel *= numpy.exp(1j*2*numpy.pi/wavefront.get_wavelength() * propagation_distance)
kernel /= 1j * wavefront.get_wavelength() * propagation_distance
tmp = fftconvolve(wavefront.get_complex_amplitude(),kernel,mode='same')
wf_propagated = Wavefront2D.initialize_wavefront_from_arrays(wavefront.get_coordinate_x(),
wavefront.get_coordinate_y(),
tmp,
wavelength=wavelength)
return wf_propagated
def propagate_2D_fraunhofer(wavefront, propagation_distance=1.0,shift_half_pixel=1):
"""
2D Fraunhofer propagator using convolution via Fourier transform
:param wavefront:
:param propagation_distance: propagation distance. If set to zero, the abscissas
of the returned wavefront are in angle (rad)
:param shift_half_pixel: set to 1 to shift half pixel (recommended using an even number of pixels) Set as default.
:return: a new 2D wavefront object with propagated wavefront
"""
wavelength = wavefront.get_wavelength()
#
# check validity
#
x = wavefront.get_coordinate_x()
y = wavefront.get_coordinate_y()
half_max_aperture = 0.5 * numpy.array( (x[-1]-x[0],y[-1]-y[0])).max()
far_field_distance = half_max_aperture**2/wavelength
if propagation_distance < far_field_distance:
print("WARNING: Fraunhoffer diffraction valid for distances > > half_max_aperture^2/lambda = %f m (propagating at %4.1f)"%
(far_field_distance,propagation_distance))
#
#compute Fourier transform
#
F1 = numpy.fft.fft2(wavefront.get_complex_amplitude()) # Take the fourier transform of the image.
# Now shift the quadrants around so that low spatial frequencies are in
# the center of the 2D fourier transformed image.
F2 = numpy.fft.fftshift( F1 )
# frequency for axis 1
shape = wavefront.size()
delta = wavefront.delta()
pixelsize = delta[0] # p_x[1] - p_x[0]
npixels = shape[0]
freq_nyquist = 0.5/pixelsize
freq_n = numpy.linspace(-1.0,1.0,npixels)
freq_x = freq_n * freq_nyquist
freq_x *= wavelength
# frequency for axis 2
pixelsize = delta[1]
npixels = shape[1]
freq_nyquist = 0.5/pixelsize
freq_n = numpy.linspace(-1.0,1.0,npixels)
freq_y = freq_n * freq_nyquist
freq_y *= wavelength
if shift_half_pixel:
freq_x = freq_x - 0.5 * numpy.abs(freq_x[1] - freq_x[0])
freq_y = freq_y - 0.5 * numpy.abs(freq_y[1] - freq_y[0])
if propagation_distance != 1.0:
freq_x *= propagation_distance
freq_y *= propagation_distance
wf_propagated = Wavefront2D.initialize_wavefront_from_arrays(freq_x,freq_y,F2,wavelength=wavelength)
return wf_propagated
def propagate_2D_fresnel_convolution(wavefront,
propagation_distance,
shift_half_pixel=1):
"""
2D Fresnel propagator using convolution via Fourier transform
:param wavefront:
:param propagation_distance: propagation distance
:param shift_half_pixel: set to 1 to shift half pixel (recommended using an even number of pixels) Set as default.
:return: a new 2D wavefront object with propagated wavefront
"""
from scipy.signal import fftconvolve
wavelength = wavefront.get_wavelength()
X = wavefront.get_mesh_x()
Y = wavefront.get_mesh_y()
if shift_half_pixel:
x = wavefront.get_coordinate_x()
y = wavefront.get_coordinate_y()
X += 0.5 * numpy.abs(x[0] - x[1])
Y += 0.5 * numpy.abs(y[0] - y[1])
kernel = numpy.exp(1j * 2 * numpy.pi / wavefront.get_wavelength() *
(X**2 + Y**2) / 2 / propagation_distance)
kernel *= numpy.exp(1j * 2 * numpy.pi / wavefront.get_wavelength() *
propagation_distance)
kernel /= 1j * wavefront.get_wavelength() * propagation_distance
tmp = fftconvolve(wavefront.get_complex_amplitude(), kernel, mode='same')
# tmp = convolve(wavefront.get_complex_amplitude(), kernel, mode='full',method='auto')
# print(">>>>>>>>>>>> kernel: ",kernel.shape)
# print(">>>>>>>>>>>> wavefront: ", wavefront.get_complex_amplitude().shape)
# print(">>>>>>>>>>>> result: ", tmp.shape,numpy.arange(0, tmp.shape[1]).shape)
wf_propagated = Wavefront2D.initialize_wavefront_from_arrays(
wavefront.get_coordinate_x(),
wavefront.get_coordinate_y(),
tmp,
wavelength=wavelength)
# start = 0.75 * (tmp.shape)[0]
# start = int(start)
# end = start + (tmp.shape)[0]
#
# tmp2 = (numpy.copy(tmp))[start:end,start:end]
#
# print(">>>>>>>>>>> tmp2,start,end",tmp2.shape,start,end)
#
#
# wf_propagated = Wavefront2D.initialize_wavefront_from_arrays(numpy.arange(0,tmp2.shape[0]),
# numpy.arange(0, tmp2.shape[1]),
# tmp2,
# wavelength=wavelength)
return wf_propagated
def undulator(self, method, npixels_x, npixels_y, pixelsize_x, pixelsize_y,
wavelength, sigma, m, m_x, m_y, propagation_distance, radius,
show, image_source, apply_lens):
wavefront = 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)
wavefront.set_spherical_wave(radius=radius)
if image_source == False:
X = wavefront.get_mesh_x()
Y = wavefront.get_mesh_y()
wavefront.rescale_amplitude(
numpy.exp(-(X**2 + Y**2) / (2 * sigma**2))
) # applichiamo intensita' con profilo gaussiano
print("Input intensity: ", wavefront.get_intensity().sum())
if apply_lens == True:
wavefront.apply_ideal_lens(focal_length_x=propagation_distance,
focal_length_y=propagation_distance)
print("Intensity after lens: ", wavefront.get_intensity().sum())
if method == 'fft':
output_wf = propagate_2D_fresnel(
wavefront, propagation_distance=propagation_distance)
elif method == 'rescaling':
output_wf = propagator2d_fourier_rescaling(
wavefront, propagation_distance=propagation_distance, m=m)
elif method == 'rescaling_xy':
output_wf = propagator2d_fourier_rescaling_xy(
wavefront,
propagation_distance=propagation_distance,
m_x=m_x,
m_y=m_y)
plot_undulator(output_wf, method=method, show=show)
horizontal_profile_wf = line_image(wavefront.get_intensity(),
horizontal_or_vertical='H')
vertical_profile_wf = line_image(wavefront.get_intensity(),
horizontal_or_vertical='V')
horizontal_phase_profile_wf = line_image(wavefront.get_phase(),
horizontal_or_vertical='H')
if image_source == True:
return wavefront.get_coordinate_x(
), horizontal_profile_wf, horizontal_phase_profile_wf
def test_initializers(self,do_plot=do_plot):
print("# ")
print("# Tests for initializars (2D) ")
print("# ")
x = numpy.linspace(-100,100,50)
y = numpy.linspace(-50,50,200)
XY = numpy.meshgrid(x,y)
X = XY[0].T
Y = XY[1].T
sigma = 10
Z = numpy.exp(- (X**2 + Y**2)/2/sigma**2) * 1j
print("Shapes x,y,z: ",x.shape,y.shape,Z.shape)
wf0 = Wavefront2D.initialize_wavefront_from_steps(x[0],numpy.abs(x[1]-x[0]),y[0],numpy.abs(y[1]-y[0]),number_of_points=Z.shape)
wf0.set_complex_amplitude(Z)
wf1 = Wavefront2D.initialize_wavefront_from_range(x[0],x[-1],y[0],y[-1],number_of_points=Z.shape)
wf1.set_complex_amplitude(Z)
wf2 = Wavefront2D.initialize_wavefront_from_arrays(x,y,Z)
if do_plot:
from srxraylib.plot.gol import plot_image
plot_image(wf0.get_intensity(),wf0.get_coordinate_x(),wf0.get_coordinate_y(),
title="initialize_wavefront_from_steps",show=0)
plot_image(wf1.get_intensity(),wf1.get_coordinate_x(),wf1.get_coordinate_y(),
title="initialize_wavefront_from_range",show=0)
plot_image(wf2.get_intensity(),wf2.get_coordinate_x(),wf2.get_coordinate_y(),
title="initialize_wavefront_from_arrays",show=1)
numpy.testing.assert_almost_equal(numpy.abs(Z)**2,wf0.get_intensity(),11)
numpy.testing.assert_almost_equal(numpy.abs(Z)**2,wf1.get_intensity(),11)
numpy.testing.assert_almost_equal(numpy.abs(Z)**2,wf2.get_intensity(),11)
numpy.testing.assert_almost_equal(x,wf0.get_coordinate_x(),11)
numpy.testing.assert_almost_equal(x,wf1.get_coordinate_x(),11)
numpy.testing.assert_almost_equal(x,wf2.get_coordinate_x(),11)
numpy.testing.assert_almost_equal(y,wf0.get_coordinate_y(),11)
numpy.testing.assert_almost_equal(y,wf1.get_coordinate_y(),11)
numpy.testing.assert_almost_equal(y,wf2.get_coordinate_y(),11)
def test_plane_wave(self,do_plot=do_plot):
#
# plane wave
#
print("# ")
print("# Tests for a 2D plane wave ")
print("# ")
wavelength = 1.24e-10
wavefront_length_x = 400e-6
wavefront_length_y = wavefront_length_x
npixels_x = 1024
npixels_y = npixels_x
x = numpy.linspace(-0.5*wavefront_length_x,0.5*wavefront_length_x,npixels_x)
y = numpy.linspace(-0.5*wavefront_length_y,0.5*wavefront_length_y,npixels_y)
wavefront = Wavefront2D.initialize_wavefront_from_steps(
x_start=x[0],x_step=numpy.abs(x[1]-x[0]),
y_start=y[0],y_step=numpy.abs(y[1]-y[0]),
number_of_points=(npixels_x,npixels_y),wavelength=wavelength)
# possible modifications
wavefront.set_plane_wave_from_amplitude_and_phase(5.0,numpy.pi/2)
numpy.testing.assert_almost_equal(wavefront.get_intensity(),25,5)
wavefront.set_plane_wave_from_complex_amplitude(2.0+3j)
numpy.testing.assert_almost_equal(wavefront.get_intensity(),13,5)
phase_before = wavefront.get_phase()
wavefront.add_phase_shift(numpy.pi/2)
phase_after = wavefront.get_phase()
numpy.testing.assert_almost_equal(phase_before+numpy.pi/2,phase_after,5)
intensity_before = wavefront.get_intensity()
wavefront.rescale_amplitude(10.0)
intensity_after = wavefront.get_intensity()
numpy.testing.assert_almost_equal(intensity_before*100,intensity_after,5)
# interpolation
wavefront.set_plane_wave_from_complex_amplitude(2.0+3j)
test_value1 = wavefront.get_interpolated_complex_amplitude(0.01,1.3)
self.assertAlmostEqual( (2.0+3j).real, test_value1.real, 5)
self.assertAlmostEqual( (2.0+3j).imag, test_value1.imag, 5)
if do_plot:
from srxraylib.plot.gol import plot_image
plot_image(wavefront.get_intensity(),wavefront.get_coordinate_x(),wavefront.get_coordinate_y(),
title="Intensity (plane wave)",show=0)
plot_image(wavefront.get_phase(),wavefront.get_coordinate_x(),wavefront.get_coordinate_y(),
title="Phase (plane wave)",show=1)
def propagate_2D_fresnel(wavefront, propagation_distance, shift_half_pixel=1):
"""
2D Fresnel propagator using convolution via Fourier transform
:param wavefront:
:param propagation_distance: propagation distance
:param shift_half_pixel: set to 1 to shift half pixel (recommended using an even number of pixels) Set as default.
:return: a new 2D wavefront object with propagated wavefront
"""
wavelength = wavefront.get_wavelength()
#
# convolving with the Fresnel kernel via FFT multiplication
#
fft = numpy.fft.fft2(wavefront.get_complex_amplitude())
# frequency for axis 1
shape = wavefront.size()
delta = wavefront.delta()
pixelsize = delta[0] # p_x[1] - p_x[0]
npixels = shape[0]
freq_nyquist = 0.5 / pixelsize
freq_n = numpy.linspace(-1.0, 1.0, npixels)
freq_x = freq_n * freq_nyquist
# frequency for axis 2
pixelsize = delta[1]
npixels = shape[1]
freq_nyquist = 0.5 / pixelsize
freq_n = numpy.linspace(-1.0, 1.0, npixels)
freq_y = freq_n * freq_nyquist
if shift_half_pixel:
freq_x = freq_x - 0.5 * numpy.abs(freq_x[1] - freq_x[0])
freq_y = freq_y - 0.5 * numpy.abs(freq_y[1] - freq_y[0])
freq_xy = numpy.array(numpy.meshgrid(freq_y, freq_x))
fft *= numpy.exp(
(-1.0j) * numpy.pi * wavelength * propagation_distance *
numpy.fft.fftshift(freq_xy[0] * freq_xy[0] + freq_xy[1] * freq_xy[1]))
ifft = numpy.fft.ifft2(fft)
wf_propagated = Wavefront2D.initialize_wavefront_from_arrays(
wavefront.get_coordinate_x(),
wavefront.get_coordinate_y(),
ifft,
wavelength=wavelength)
return wf_propagated
def propagate_2D_fresnel(wavefront, propagation_distance,shift_half_pixel=1):
"""
2D Fresnel propagator using convolution via Fourier transform
:param wavefront:
:param propagation_distance: propagation distance
:param shift_half_pixel: set to 1 to shift half pixel (recommended using an even number of pixels) Set as default.
:return: a new 2D wavefront object with propagated wavefront
"""
wavelength = wavefront.get_wavelength()
#
# convolving with the Fresnel kernel via FFT multiplication
#
fft = numpy.fft.fft2(wavefront.get_complex_amplitude())
# frequency for axis 1
shape = wavefront.size()
delta = wavefront.delta()
pixelsize = delta[0] # p_x[1] - p_x[0]
npixels = shape[0]
freq_nyquist = 0.5/pixelsize
freq_n = numpy.linspace(-1.0,1.0,npixels)
freq_x = freq_n * freq_nyquist
# frequency for axis 2
pixelsize = delta[1]
npixels = shape[1]
freq_nyquist = 0.5/pixelsize
freq_n = numpy.linspace(-1.0,1.0,npixels)
freq_y = freq_n * freq_nyquist
if shift_half_pixel:
freq_x = freq_x - 0.5 * numpy.abs(freq_x[1] - freq_x[0])
freq_y = freq_y - 0.5 * numpy.abs(freq_y[1] - freq_y[0])
freq_xy = numpy.array(numpy.meshgrid(freq_y,freq_x))
fft *= numpy.exp((-1.0j) * numpy.pi * wavelength * propagation_distance *
numpy.fft.fftshift(freq_xy[0]*freq_xy[0] + freq_xy[1]*freq_xy[1]) )
ifft = numpy.fft.ifft2(fft)
wf_propagated = Wavefront2D.initialize_wavefront_from_arrays(wavefront.get_coordinate_x(),
wavefront.get_coordinate_y(),
ifft,
wavelength=wavelength)
return wf_propagated
def propagate_2D_fraunhofer(wavefront,
propagation_distance=1.0,
shift_half_pixel=0): #todo: modificato da giovanni
"""
2D Fraunhofer propagator using convolution via Fourier transform
:param wavefront:
:param propagation_distance: propagation distance. If set to zero, the abscissas
of the returned wavefront are in angle (rad)
:param shift_half_pixel: set to 1 to shift half pixel (recommended using an even number of pixels) Set as default.
:return: a new 2D wavefront object with propagated wavefront
"""
wavelength = wavefront.get_wavelength()
#
# check validity
#
x = wavefront.get_coordinate_x()
y = wavefront.get_coordinate_y()
half_max_aperture = 0.5 * numpy.array((x[-1] - x[0], y[-1] - y[0])).max()
far_field_distance = half_max_aperture**2 / wavelength
if propagation_distance < far_field_distance:
print(
"WARNING: Fraunhoffer diffraction valid for distances > > half_max_aperture^2/lambda = %f m (propagating at %4.1f)"
% (far_field_distance, propagation_distance))
#
#compute Fourier transform
#
# frequency for axis 1
shape = wavefront.size()
delta = wavefront.delta()
wavenumber = wavefront.get_wavenumber()
pixelsize = delta[0] # p_x[1] - p_x[0]
npixels = shape[0]
fft_scale = numpy.fft.fftfreq(npixels, d=pixelsize)
fft_scale = numpy.fft.fftshift(fft_scale)
x2 = fft_scale * propagation_distance * wavelength
# frequency for axis 2
pixelsize = delta[1]
npixels = shape[1]
fft_scale = numpy.fft.fftfreq(npixels, d=pixelsize)
fft_scale = numpy.fft.fftshift(fft_scale)
y2 = fft_scale * propagation_distance * wavelength
f_x, f_y = numpy.meshgrid(x2, y2, indexing='ij')
fsq = numpy.fft.fftshift(f_x**2 + f_y**2)
P1 = numpy.exp(1.0j * wavenumber * propagation_distance)
P2 = numpy.exp(1.0j * wavenumber / 2 / propagation_distance * fsq)
P3 = 1.0j * wavelength * propagation_distance
F1 = numpy.fft.fft2(wavefront.get_complex_amplitude()
) # Take the fourier transform of the image.
# Now shift the quadrants around so that low spatial frequencies are in
# the center of the 2D fourier transformed image.
F1 *= P1
F1 *= P2
F1 /= P3
F2 = numpy.fft.fftshift(F1)
if shift_half_pixel:
x2 = x2 - 0.5 * numpy.abs(x2[1] - x2[0])
y2 = y2 - 0.5 * numpy.abs(y2[1] - y2[0])
wf_propagated = Wavefront2D.initialize_wavefront_from_arrays(
x2, y2, F2, wavelength=wavelength)
return wf_propagated
def propagate_2D_integral(wavefront,
propagation_distance,
shuffle_interval=0,
calculate_grid_only=1):
"""
2D Fresnel-Kirchhoff propagator via simplified integral
NOTE: this propagator is experimental and much less performant than the ones using Fourier Optics
Therefore, it is not recommended to use.
:param wavefront:
:param propagation_distance: propagation distance
:param shuffle_interval: it is known that this method replicates the central diffraction spot
The distace of the replica is proportional to 1/pixelsize
To avoid that, it is possible to change a bit (randomly) the coordinates
of the wavefront. shuffle_interval controls this shift: 0=No shift. A typical
value can be 1e5.
The result shows a diffraction pattern without replica but with much noise.
:param calculate_grid_only: if set, it calculates only the horizontal and vertical profiles, but returns the
full image with the other pixels to zero. This is useful when calculating large arrays,
so it is set as the default.
:return: a new 2D wavefront object with propagated wavefront
"""
#
# Fresnel-Kirchhoff integral (neglecting inclination factor)
#
if calculate_grid_only == 0:
#
# calculation over the whole detector area
#
p_x = wavefront.get_coordinate_x()
p_y = wavefront.get_coordinate_y()
wavelength = wavefront.get_wavelength()
amplitude = wavefront.get_complex_amplitude()
det_x = p_x.copy()
det_y = p_y.copy()
p_X = wavefront.get_mesh_x()
p_Y = wavefront.get_mesh_y()
det_X = p_X
det_Y = p_Y
amplitude_propagated = numpy.zeros_like(amplitude, dtype='complex')
wavenumber = 2 * numpy.pi / wavelength
for i in range(det_x.size):
for j in range(det_y.size):
if shuffle_interval == 0:
rd_x = 0.0
rd_y = 0.0
else:
rd_x = (numpy.random.rand(p_x.size, p_y.size) -
0.5) * shuffle_interval
rd_y = (numpy.random.rand(p_x.size, p_y.size) -
0.5) * shuffle_interval
r = numpy.sqrt(
numpy.power(p_X + rd_x - det_X[i, j], 2) +
numpy.power(p_Y + rd_y - det_Y[i, j], 2) +
numpy.power(propagation_distance, 2))
amplitude_propagated[i, j] = (
amplitude / r * numpy.exp(1.j * wavenumber * r)).sum()
wavefront2 = Wavefront2D.initialize_wavefront_from_arrays(
det_x, det_y, amplitude_propagated)
else:
x = wavefront.get_coordinate_x()
y = wavefront.get_coordinate_y()
X = wavefront.get_mesh_x()
Y = wavefront.get_mesh_y()
wavenumber = 2 * numpy.pi / wavefront.get_wavelength()
amplitude = wavefront.get_complex_amplitude()
used_indices = wavefront.get_mask_grid(width_in_pixels=(1, 1),
number_of_lines=(1, 1))
indices_x = wavefront.get_mesh_indices_x()
indices_y = wavefront.get_mesh_indices_y()
indices_x_flatten = indices_x[numpy.where(used_indices == 1)].flatten()
indices_y_flatten = indices_y[numpy.where(used_indices == 1)].flatten()
X_flatten = X[numpy.where(used_indices == 1)].flatten()
Y_flatten = Y[numpy.where(used_indices == 1)].flatten()
complex_amplitude_propagated = amplitude * 0
print(
"propagate_2D_integral: Calculating %d points from a total of %d x %d = %d"
% (X_flatten.size, amplitude.shape[0], amplitude.shape[1],
amplitude.shape[0] * amplitude.shape[1]))
for i in range(X_flatten.size):
r = numpy.sqrt(
numpy.power(wavefront.get_mesh_x() - X_flatten[i], 2) +
numpy.power(wavefront.get_mesh_y() - Y_flatten[i], 2) +
numpy.power(propagation_distance, 2))
complex_amplitude_propagated[indices_x_flatten[i],
indices_y_flatten[i]] = (
amplitude / r * numpy.exp(
1.j * wavenumber * r)).sum()
wavefront2 = Wavefront2D.initialize_wavefront_from_arrays(
x, y, complex_amplitude_propagated, wavefront.get_wavelength())
return wavefront2
def propagate_2D_fresnel_srw(wavefront,
propagation_distance,
srw_autosetting=0):
"""
2D Fresnel propagator using convolution via Fourier transform
:param wavefront:
:param propagation_distance:
:param srw_autosetting:set to 1 for automatic SRW redimensionate wavefront
:return:
"""
#
# convolving with the Fresnel kernel via SRW package
#
from srxraylib.waveoptics.NumpyToSRW import numpyArrayToSRWArray, SRWWavefrontFromElectricField, SRWEFieldAsNumpy
import scipy.constants as codata
angstroms_to_eV = codata.h * codata.c / codata.e * 1e10
srw_wfr = SRWWavefrontFromElectricField(
wavefront.get_coordinate_x()[0],
wavefront.get_coordinate_x()[-1], wavefront.get_complex_amplitude(),
wavefront.get_coordinate_y()[0],
wavefront.get_coordinate_y()[-1],
numpy.zeros_like(wavefront.get_complex_amplitude()),
angstroms_to_eV / (wavefront.get_wavelength() * 1e10), 1.0, 1.0, 1e-3,
1.0, 1e-3)
#
# propagation
#
optDrift = srwlib.SRWLOptD(propagation_distance) #Drift space
#Wavefront Propagation Parameters:
#[0]: Auto-Resize (1) or not (0) Before propagation
#[1]: Auto-Resize (1) or not (0) After propagation
#[2]: Relative Precision for propagation with Auto-Resizing (1. is nominal)
#[3]: Allow (1) or not (0) for semi-analytical treatment of the quadratic (leading) phase terms at the propagation
#[4]: Do any Resizing on Fourier side, using FFT, (1) or not (0)
#[5]: Horizontal Range modification factor at Resizing (1. means no modification)
#[6]: Horizontal Resolution modification factor at Resizing
#[7]: Vertical Range modification factor at Resizing
#[8]: Vertical Resolution modification factor at Resizing
#[9]: Type of wavefront Shift before Resizing (not yet implemented)
#[10]: New Horizontal wavefront Center position after Shift (not yet implemented)
#[11]: New Vertical wavefront Center position after Shift (not yet implemented)
if srw_autosetting:
# 0 1 2 3 4 5 6 7 8 9 10 11
propagParDrift = [1, 1, 1., 0, 0, 1., 1., 1., 1., 0, 0, 0]
else:
# 0 1 2 3 4 5 6 7 8 9 10 11
propagParDrift = [0, 0, 1., 0, 0, 1., 1., 1., 1., 0, 0, 0]
optBL = srwlib.SRWLOptC(
[optDrift], [propagParDrift]
) #"Beamline" - Container of Optical Elements (together with the corresponding wavefront propagation instructions)
print(' Simulating Electric Field Wavefront Propagation by SRW ... ',
end='\n')
srwlib.srwl.PropagElecField(srw_wfr, optBL)
wavefront2 = Wavefront2D.initialize_wavefront_from_range(
srw_wfr.mesh.xStart,
srw_wfr.mesh.xFin,
srw_wfr.mesh.yStart,
srw_wfr.mesh.yFin,
number_of_points=(srw_wfr.mesh.nx, srw_wfr.mesh.ny),
wavelength=wavefront.get_wavelength())
amplitude2 = SRWEFieldAsNumpy(srw_wfr)
amplitude2 = amplitude2[0, :, :, 0]
wavefront2.set_complex_amplitude(amplitude2)
return wavefront2
def test_interpolator(self,do_plot=do_plot):
#
# interpolator
#
print("# ")
print("# Tests for 2D interpolator ")
print("# ")
x = numpy.linspace(-10,10,100)
y = numpy.linspace(-20,20,50)
xy = numpy.meshgrid(x,y)
X = xy[0].T
Y = xy[1].T
sigma = 3.0
Z = 3*numpy.exp(- (X**2+Y**2)/2/sigma**2) +4j
print("shape of Z",Z.shape)
wf = Wavefront2D.initialize_wavefront_from_steps(x[0],x[1]-x[0],y[0],y[1]-y[0],number_of_points=(100,50))
print("wf shape: ",wf.size())
wf.set_complex_amplitude( Z )
x1 = 3.2
y1 = -2.5
z1 = 3*numpy.exp(- (x1**2+y1**2)/2/sigma**2) + 4j
print("complex ampl at (%g,%g): %g+%gi (exact=%g+%gi)"%(x1,y1,
wf.get_interpolated_complex_amplitude(x1,y1).real,
wf.get_interpolated_complex_amplitude(x1,y1).imag,
z1.real,z1.imag))
self.assertAlmostEqual(wf.get_interpolated_complex_amplitude(x1,y1).real,z1.real,4)
self.assertAlmostEqual(wf.get_interpolated_complex_amplitude(x1,y1).imag,z1.imag,4)
#
print("intensity at (%g,%g): %g (exact=%g)"%(x1,y1,wf.get_interpolated_intensity(x1,y1),numpy.abs(z1)**2))
self.assertAlmostEqual(wf.get_interpolated_intensity(x1,y1),numpy.abs(z1)**2,4)
# interpolate on same grid
interpolated_complex_amplitude_on_same_grid = wf.get_interpolated_complex_amplitudes(X,Y)
print("Shape interpolated at same grid: ",interpolated_complex_amplitude_on_same_grid.shape)
numpy.testing.assert_array_almost_equal(wf.get_complex_amplitude(),interpolated_complex_amplitude_on_same_grid,4)
print("Total intensity original wavefront: %g, interpolated on the same grid: %g"%
(wf.get_intensity().sum(), (numpy.abs(interpolated_complex_amplitude_on_same_grid)**2).sum()))
if do_plot:
from srxraylib.plot.gol import plot_image
plot_image(wf.get_intensity(),wf.get_coordinate_x(),wf.get_coordinate_y(),title="Original",show=0)
plot_image(wf.get_interpolated_intensity(X,Y),wf.get_coordinate_x(),wf.get_coordinate_y(),
title="interpolated on same grid",show=1)
# rebin wavefront
# wf.set_plane_wave_from_complex_amplitude(3+4j)
wf_rebin = wf.rebin(2.0,5.0,0.5,0.8,keep_the_same_intensity=1,set_extrapolation_to_zero=1)
print("Shape before rebinning: ",wf.size())
print("Shape after rebinning: ",wf_rebin.size())
print("Total intensity original wavefront: %g, rebinned: %g"%
(wf.get_intensity().sum(), wf_rebin.get_intensity().sum() ))
if do_plot:
from srxraylib.plot.gol import plot_image
plot_image(wf.get_intensity(),wf.get_coordinate_x(),wf.get_coordinate_y(),title="BEFORE REBINNING",show=0)
plot_image(wf_rebin.get_intensity(),wf_rebin.get_coordinate_x(),wf_rebin.get_coordinate_y(),
title="REBINNED",show=1)
def propagate_2D_fresnel(self,
do_plot,
method='fft',
wavelength=1.24e-10,
aperture_type='square',
aperture_diameter=40e-6,
pixelsize_x=1e-6,
pixelsize_y=1e-6,
npixels_x=1024,
npixels_y=1024,
propagation_distance=30.0,
m=1,
m_x=1,
m_y=1,
show=0):
method_label = "fresnel (%s)" % method
print(
"\n# ")
print("# 2D near field fresnel (%s) diffraction from a %s aperture " %
(method_label, aperture_type))
print("# ")
# wf = Wavefront2D.initialize_wavefront_from_steps(x_start=-pixelsize_x*npixels_x/2,
# x_step=pixelsize_x,
# y_start=-pixelsize_y*npixels_y/2,
# y_step=pixelsize_y,
# wavelength=wavelength,
# number_of_points=(npixels_x,npixels_y))
wf = 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.set_plane_wave_from_complex_amplitude((1.0 + 0j))
if aperture_type == 'circle':
wf.apply_pinhole(aperture_diameter / 2)
elif aperture_type == 'square':
wf.apply_slit(-aperture_diameter / 2, aperture_diameter / 2,
-aperture_diameter / 2, aperture_diameter / 2)
else:
raise Exception("Not implemented! (accepted: circle, square)")
if method == 'fft':
wf1 = propagate_2D_fresnel(wf, propagation_distance)
elif method == 'rescaling':
wf1 = propagator2d_fourier_rescaling(wf, propagation_distance, m=m)
elif method == 'rescaling_xy':
wf1 = propagator2d_fourier_rescaling_xy(wf,
propagation_distance,
m_x=m_x,
m_y=m_y)
else:
raise Exception("Not implemented method: %s" % method)
if do_plot:
from srxraylib.plot.gol import plot_image
plot_image(wf.get_intensity(),
1e6 * wf.get_coordinate_x(),
1e6 * wf.get_coordinate_y(),
title="aperture intensity (%s), Diameter=%5.1f um" %
(aperture_type, 1e6 * aperture_diameter),
xtitle="X [um]",
ytitle="Y [um]",
show=0)
plot_image(
wf1.get_intensity(),
1e6 * wf1.get_coordinate_x() / propagation_distance,
1e6 * wf1.get_coordinate_y() / propagation_distance,
title=
"Diffracted intensity (%s) by a %s slit of aperture %3.1f um" %
(aperture_type, method_label, 1e6 * aperture_diameter),
xtitle="X [urad]",
ytitle="Y [urad]",
show=0)
# get the theoretical value
angle_x = wf.get_coordinate_x() / propagation_distance
intensity_theory = get_theoretical_diffraction_pattern(
angle_x,
aperture_type=aperture_type,
aperture_diameter=aperture_diameter,
wavelength=wavelength,
normalization=True)
intensity_calculated = wf1.get_intensity()[:, int(wf1.size()[1] / 2)]
intensity_calculated /= intensity_calculated.max()
if do_plot:
from srxraylib.plot.gol import plot
plot(
wf1.get_coordinate_x() * 1e6 / propagation_distance,
intensity_calculated,
angle_x * 1e6,
intensity_theory,
legend=[
"%s H profile" % method_label, "Theoretical (far field)"
],
legend_position=(0.60, 0.95),
title=
"%s diffraction of a %s slit of %3.1f um at wavelength of %3.1f A"
% (method_label, aperture_type, aperture_diameter * 1e6,
wavelength * 1e10),
xtitle="X (urad)",
ytitle="Intensity",
xrange=[-20, 20],
show=show)
if method == 'rescaling_xy' and do_plot:
angle_y = wf.get_coordinate_y() / propagation_distance
intensity_calculated_v = wf1.get_intensity()[int(wf1.size()[1] /
2), :]
intensity_calculated_v /= intensity_calculated_v.max()
from srxraylib.plot.gol import plot
plot(
wf1.get_coordinate_y() * 1e6 / propagation_distance,
intensity_calculated_v,
angle_y * 1e6,
intensity_theory,
legend=[
"%s V profile" % method_label, "Theoretical (far field)"
],
legend_position=(0.60, 0.95),
title=
"%s diffraction of a %s slit of %3.1f um at wavelength of %3.1f A"
% (method_label, aperture_type, aperture_diameter * 1e6,
wavelength * 1e10),
xtitle="Y (urad)",
ytitle="Intensity",
xrange=[-20, 20],
show=show)
return wf1.get_coordinate_x(
) / propagation_distance, intensity_calculated, angle_x, intensity_theory
def propagate_2D_fresnel_srw(wavefront, propagation_distance,
srw_autosetting=0):
"""
2D Fresnel propagator using convolution via Fourier transform
:param wavefront:
:param propagation_distance:
:param srw_autosetting:set to 1 for automatic SRW redimensionate wavefront
:return:
"""
#
# convolving with the Fresnel kernel via SRW package
#
from srxraylib.waveoptics.NumpyToSRW import numpyArrayToSRWArray, SRWWavefrontFromElectricField, SRWEFieldAsNumpy
import scipy.constants as codata
angstroms_to_eV = codata.h*codata.c/codata.e*1e10
srw_wfr = SRWWavefrontFromElectricField(
wavefront.get_coordinate_x()[0],wavefront.get_coordinate_x()[-1],wavefront.get_complex_amplitude(),
wavefront.get_coordinate_y()[0],wavefront.get_coordinate_y()[-1],numpy.zeros_like(wavefront.get_complex_amplitude()),
angstroms_to_eV/(wavefront.get_wavelength()*1e10), 1.0, 1.0, 1e-3, 1.0, 1e-3)
#
# propagation
#
optDrift = srwlib.SRWLOptD(propagation_distance) #Drift space
#Wavefront Propagation Parameters:
#[0]: Auto-Resize (1) or not (0) Before propagation
#[1]: Auto-Resize (1) or not (0) After propagation
#[2]: Relative Precision for propagation with Auto-Resizing (1. is nominal)
#[3]: Allow (1) or not (0) for semi-analytical treatment of the quadratic (leading) phase terms at the propagation
#[4]: Do any Resizing on Fourier side, using FFT, (1) or not (0)
#[5]: Horizontal Range modification factor at Resizing (1. means no modification)
#[6]: Horizontal Resolution modification factor at Resizing
#[7]: Vertical Range modification factor at Resizing
#[8]: Vertical Resolution modification factor at Resizing
#[9]: Type of wavefront Shift before Resizing (not yet implemented)
#[10]: New Horizontal wavefront Center position after Shift (not yet implemented)
#[11]: New Vertical wavefront Center position after Shift (not yet implemented)
if srw_autosetting:
# 0 1 2 3 4 5 6 7 8 9 10 11
propagParDrift = [1, 1, 1., 0, 0, 1., 1., 1., 1., 0, 0, 0]
else:
# 0 1 2 3 4 5 6 7 8 9 10 11
propagParDrift = [0, 0, 1., 0, 0, 1., 1., 1., 1., 0, 0, 0]
optBL = srwlib.SRWLOptC([optDrift], [propagParDrift]) #"Beamline" - Container of Optical Elements (together with the corresponding wavefront propagation instructions)
print(' Simulating Electric Field Wavefront Propagation by SRW ... ', end='\n')
srwlib.srwl.PropagElecField(srw_wfr, optBL)
wavefront2 = Wavefront2D.initialize_wavefront_from_range(srw_wfr.mesh.xStart, srw_wfr.mesh.xFin,
srw_wfr.mesh.yStart, srw_wfr.mesh.yFin,
number_of_points=(srw_wfr.mesh.nx,srw_wfr.mesh.ny),
wavelength=wavefront.get_wavelength())
amplitude2 = SRWEFieldAsNumpy(srw_wfr)
amplitude2 = amplitude2[0,:,:,0]
wavefront2.set_complex_amplitude(amplitude2)
return wavefront2
def test_propagate_2D_fraunhofer(self,do_plot=do_plot,aperture_type='square',aperture_diameter=40e-6,
pixelsize_x=1e-6,pixelsize_y=1e-6,npixels_x=1024,npixels_y=1024,wavelength=1.24e-10):
"""
:param do_plot: 0=No plot, 1=Do plot
:param aperture_type: 'circle' 'square' 'gaussian' (Gaussian sigma = aperture_diameter/2.35)
:param aperture_diameter:
:param pixelsize_x:
:param pixelsize_y:
:param npixels_x:
:param npixels_y:
:param wavelength:
:return:
"""
print("\n# ")
print("# far field 2D (fraunhofer) diffraction from a square aperture ")
print("# ")
method = "fraunhofer"
print("Fraunhoffer diffraction valid for distances > > a^2/lambda = %f m"%((aperture_diameter/2)**2/wavelength))
# wf = Wavefront2D.initialize_wavefront_from_steps(x_start=-pixelsize_x*npixels_x/2,
# x_step=pixelsize_x,
# y_start=-pixelsize_y*npixels_y/2,
# y_step=pixelsize_y,
# wavelength=wavelength,
# number_of_points=(npixels_x,npixels_y))
wf = 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.set_plane_wave_from_complex_amplitude((1.0+0j))
if aperture_type == 'circle':
wf.apply_pinhole(aperture_diameter/2)
elif aperture_type == 'square':
wf.apply_slit(-aperture_diameter/2, aperture_diameter/2,-aperture_diameter/2, aperture_diameter/2)
elif aperture_type == 'gaussian':
X = wf.get_mesh_x()
Y = wf.get_mesh_y()
window = numpy.exp(- (X*X + Y*Y)/2/(aperture_diameter/2.35)**2)
wf.rescale_amplitudes(window)
else:
raise Exception("Not implemented! (accepted: circle, square, gaussian)")
wf1 = propagate_2D_fraunhofer(wf, propagation_distance=1.0) # propagating at 1 m means the result is like in angles
if do_plot:
plot_image(wf.get_intensity(),1e6*wf.get_coordinate_x(),1e6*wf.get_coordinate_y(),
title="aperture intensity (%s), Diameter=%5.1f um"%
(aperture_type,1e6*aperture_diameter),xtitle="X [um]",ytitle="Y [um]",
show=0)
plot_image(wf1.get_intensity(),1e6*wf1.get_coordinate_x(),1e6*wf1.get_coordinate_y(),
title="2D Diffracted intensity (%s) by a %s slit of aperture %3.1f um"%
(aperture_type,method,1e6*aperture_diameter),
xtitle="X [urad]",ytitle="Y [urad]",
show=0)
angle_x = wf1.get_coordinate_x() # + 0.5*wf1.delta()[0] # shifted of half-pixel!!!
intensity_theory = get_theoretical_diffraction_pattern(angle_x,
aperture_type=aperture_type,aperture_diameter=aperture_diameter,
wavelength=wavelength,normalization=True)
intensity_calculated = wf1.get_intensity()[:,wf1.size()[1]/2]
intensity_calculated /= intensity_calculated.max()
if do_plot:
plot(wf1.get_coordinate_x()*1e6,intensity_calculated,
angle_x*1e6,intensity_theory,
legend=["Calculated (FT) H profile","Theoretical"],legend_position=(0.95, 0.95),
title="2D Fraunhofer Diffraction of a %s slit of %3.1f um at wavelength of %3.1f A"%
(aperture_type,aperture_diameter*1e6,wavelength*1e10),
xtitle="X (urad)", ytitle="Intensity",xrange=[-80,80])
numpy.testing.assert_almost_equal(intensity_calculated,intensity_theory,1)
def get_wavefront2d(self, i_mode, wavelength):
return Wavefront2D.initialize_wavefront_from_arrays(
self.x_coordinates(),
self.y_coordinates(),
self.mode(i_mode),
wavelength=wavelength)
pixelsize_y = 1e-7
npixels_x = 2024
npixels_y = 2024
propagation_distance = 1.0
show = 1
method_label = "fresnel (fft)"
print("\n# ")
print("# 2D near field fresnel (%s) diffraction from a a circular stop " %
(method_label))
print("# ")
wf = 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.set_plane_wave_from_complex_amplitude((1.0 + 0j))
wf.apply_pinhole(aperture_diameter / 2, negative=True)
wf1 = propagate_2D_fresnel(wf, propagation_distance)
plot_image(wf.get_intensity(),
1e6 * wf.get_coordinate_x(),
1e6 * wf.get_coordinate_y(),
title="intensity at screen/aperture plane, Diameter=%5.1f um" %
(1e6 * aperture_diameter),
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)
def propagate_2D(calculation_parameters, input_parameters):
shadow_oe = calculation_parameters.shadow_oe_end
if calculation_parameters.do_ff_z and calculation_parameters.do_ff_x:
global_phase_shift_profile = None
if shadow_oe._oe.F_MOVE == 1 and shadow_oe._oe.X_ROT != 0.0:
if input_parameters.ghy_calcType == 3 or input_parameters.ghy_calcType == 4:
global_phase_shift_profile = calculation_parameters.w_mirr_2D_values
elif input_parameters.ghy_calcType == 2:
global_phase_shift_profile = ScaledMatrix.initialize_from_range(numpy.zeros((3, 3)),
shadow_oe._oe.RWIDX2, shadow_oe._oe.RWIDX1,
shadow_oe._oe.RLEN2, shadow_oe._oe.RLEN1)
for x_index in range(global_phase_shift_profile.size_x()):
global_phase_shift_profile.z_values[x_index, :] += global_phase_shift_profile.get_y_values()*numpy.sin(numpy.radians(-shadow_oe._oe.X_ROT))
elif input_parameters.ghy_calcType == 3 or input_parameters.ghy_calcType == 4:
global_phase_shift_profile = calculation_parameters.w_mirr_2D_values
# only tangential slopes
if input_parameters.ghy_calcType == 3 or input_parameters.ghy_calcType == 4:
rms_slope = hy_findrmsslopefromheight(ScaledArray(np_array=global_phase_shift_profile.z_values[int(global_phase_shift_profile.size_x()/2), :],
scale=global_phase_shift_profile.get_y_values()))
print("Using RMS slope = " + str(rms_slope))
# ------------------------------------------
# far field calculation
# ------------------------------------------
focallength_ff = calculate_focal_length_ff_2D(calculation_parameters.ghy_x_min,
calculation_parameters.ghy_x_max,
calculation_parameters.ghy_z_min,
calculation_parameters.ghy_z_max,
input_parameters.ghy_npeak,
calculation_parameters.gwavelength)
if (input_parameters.ghy_calcType == 3 or input_parameters.ghy_calcType == 4) and rms_slope != 0:
focallength_ff = min(focallength_ff,(calculation_parameters.ghy_z_max-calculation_parameters.ghy_z_min) / 16 / rms_slope ) #xshi changed
elif input_parameters.ghy_calcType == 2 and not global_phase_shift_profile is None:
focallength_ff = min(focallength_ff, input_parameters.ghy_distance*4) #TODO: PATCH to be found with a formula
print("FF: calculated focal length: " + str(focallength_ff))
fftsize_x = int(calculate_fft_size(calculation_parameters.ghy_x_min,
calculation_parameters.ghy_x_max,
calculation_parameters.gwavelength,
focallength_ff,
input_parameters.ghy_fftnpts,
factor=20))
fftsize_z = int(calculate_fft_size(calculation_parameters.ghy_z_min,
calculation_parameters.ghy_z_max,
calculation_parameters.gwavelength,
focallength_ff,
input_parameters.ghy_fftnpts,
factor=20))
print("FF: creating plane wave begin, fftsize_x = " + str(fftsize_x) + ", fftsize_z = " + str(fftsize_z))
wavefront = Wavefront2D.initialize_wavefront_from_range(wavelength=calculation_parameters.gwavelength,
number_of_points=(fftsize_x, fftsize_z),
x_min=calculation_parameters.ghy_x_min,
x_max=calculation_parameters.ghy_x_max,
y_min=calculation_parameters.ghy_z_min,
y_max=calculation_parameters.ghy_z_max)
try:
for i in range(0, len(wavefront.electric_field_array.x_coord)):
for j in range(0, len(wavefront.electric_field_array.y_coord)):
interpolated = calculation_parameters.wIray_2d.interpolate_value(wavefront.electric_field_array.x_coord[i],
wavefront.electric_field_array.y_coord[j])
wavefront.electric_field_array.set_z_value(i, j, numpy.sqrt(0.0 if interpolated < 0 else interpolated))
except IndexError:
raise Exception("Unexpected Error during interpolation: try reduce Number of bins for I(Tangential) histogram")
wavefront.apply_ideal_lens(focallength_ff, focallength_ff)
shadow_oe = calculation_parameters.shadow_oe_end
if input_parameters.ghy_calcType == 3 or \
(input_parameters.ghy_calcType == 2 and not global_phase_shift_profile is None):
print("FF: calculating phase shift due to Height Error Profile")
phase_shifts = numpy.zeros(wavefront.size())
for index in range(0, phase_shifts.shape[0]):
np_array = numpy.zeros(global_phase_shift_profile.shape()[1])
for j in range(0, len(np_array)):
np_array[j] = global_phase_shift_profile.interpolate_value(wavefront.get_coordinate_x()[index], calculation_parameters.w_mirr_2D_values.get_y_value(j))
global_phase_shift_profile_z = ScaledArray.initialize_from_steps(np_array,
global_phase_shift_profile.y_coord[0],
global_phase_shift_profile.y_coord[1] - global_phase_shift_profile.y_coord[0])
phase_shifts[index, :] = get_mirror_phase_shift(wavefront.get_coordinate_y(),
calculation_parameters.gwavelength,
calculation_parameters.wangle_z,
calculation_parameters.wl_z,
global_phase_shift_profile_z)
wavefront.add_phase_shifts(phase_shifts)
elif input_parameters.ghy_calcType == 4:
print("FF: calculating phase shift due to Height Error Profile")
phase_shifts = numpy.zeros(wavefront.size())
for index in range(0, phase_shifts.shape[0]):
global_phase_shift_profile_z = ScaledArray.initialize_from_steps(global_phase_shift_profile.z_values[index, :],
global_phase_shift_profile.y_coord[0],
global_phase_shift_profile.y_coord[1] - global_phase_shift_profile.y_coord[0])
phase_shifts[index, :] = get_grating_phase_shift(wavefront.get_coordinate_y(),
calculation_parameters.gwavelength,
calculation_parameters.wangle_z,
calculation_parameters.wangle_ref_z,
calculation_parameters.wl_z,
global_phase_shift_profile_z)
wavefront.add_phase_shifts(phase_shifts)
elif input_parameters.ghy_calcType == 6:
for w_mirr_2D_values in calculation_parameters.w_mirr_2D_values:
phase_shift = get_crl_phase_shift(w_mirr_2D_values,
input_parameters,
calculation_parameters,
[wavefront.get_coordinate_x(), wavefront.get_coordinate_y()])
wavefront.add_phase_shift(phase_shift)
print("calculated plane wave: begin FF propagation (distance = " + str(focallength_ff) + ")")
propagated_wavefront = propagator2D.propagate_2D_fresnel(wavefront, focallength_ff)
print("dif_zp: begin calculation")
imagesize_x = min(abs(calculation_parameters.ghy_x_max), abs(calculation_parameters.ghy_x_min)) * 2
imagesize_x = min(imagesize_x,
input_parameters.ghy_npeak*2*0.88*calculation_parameters.gwavelength*focallength_ff/abs(calculation_parameters.ghy_x_max-calculation_parameters.ghy_x_min))
# TODO: this is a patch: to be rewritten
if shadow_oe._oe.F_MOVE==1 and not shadow_oe._oe.Y_ROT==0:
imagesize_x = max(imagesize_x, 8*(focallength_ff*numpy.tan(numpy.radians(numpy.abs(shadow_oe._oe.Y_ROT))) + numpy.abs(shadow_oe._oe.OFFX)))
delta_x = propagated_wavefront.delta()[0]
imagenpts_x = int(round(imagesize_x/delta_x/2) * 2 + 1)
imagesize_z = min(abs(calculation_parameters.ghy_z_max), abs(calculation_parameters.ghy_z_min)) * 2
imagesize_z = min(imagesize_z,
input_parameters.ghy_npeak*2*0.88*calculation_parameters.gwavelength*focallength_ff/abs(calculation_parameters.ghy_z_max-calculation_parameters.ghy_z_min))
# TODO: this is a patch: to be rewritten
if shadow_oe._oe.F_MOVE==1 and not shadow_oe._oe.X_ROT==0:
imagesize_z = max(imagesize_z, 8*(focallength_ff*numpy.tan(numpy.radians(numpy.abs(shadow_oe._oe.X_ROT))) + numpy.abs(shadow_oe._oe.OFFZ)))
delta_z = propagated_wavefront.delta()[1]
imagenpts_z = int(round(imagesize_z/delta_z/2) * 2 + 1)
dif_xpzp = ScaledMatrix.initialize_from_range(numpy.ones((imagenpts_x, imagenpts_z)),
min_scale_value_x = -(imagenpts_x - 1) / 2 * delta_x,
max_scale_value_x =(imagenpts_x - 1) / 2 * delta_x,
min_scale_value_y = -(imagenpts_z - 1) / 2 * delta_z,
max_scale_value_y =(imagenpts_z - 1) / 2 * delta_z)
for i in range(0, dif_xpzp.shape()[0]):
for j in range(0, dif_xpzp.shape()[1]):
dif_xpzp.set_z_value(i, j, numpy.absolute(propagated_wavefront.get_interpolated_complex_amplitude(
dif_xpzp.x_coord[i],
dif_xpzp.y_coord[j]))**2
)
dif_xpzp.set_scale_from_range(0,
-(imagenpts_x - 1) / 2 * delta_x / focallength_ff,
(imagenpts_x - 1) / 2 * delta_x / focallength_ff)
dif_xpzp.set_scale_from_range(1,
-(imagenpts_z - 1) / 2 * delta_z / focallength_ff,
(imagenpts_z - 1) / 2 * delta_z / focallength_ff)
calculation_parameters.dif_xpzp = dif_xpzp
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)
def propagate_2D_fresnel(self,do_plot=do_plot,method='fft',
wavelength=1.24e-10,aperture_type='square',aperture_diameter=40e-6,
pixelsize_x=1e-6,pixelsize_y=1e-6,npixels_x=1024,npixels_y=1024,
propagation_distance = 30.0,show=1):
method_label = "fresnel (%s)"%method
print("\n# ")
print("# 2D near field fresnel (%s) diffraction from a %s aperture "%(method_label,aperture_type))
print("# ")
# wf = Wavefront2D.initialize_wavefront_from_steps(x_start=-pixelsize_x*npixels_x/2,
# x_step=pixelsize_x,
# y_start=-pixelsize_y*npixels_y/2,
# y_step=pixelsize_y,
# wavelength=wavelength,
# number_of_points=(npixels_x,npixels_y))
wf = 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.set_plane_wave_from_complex_amplitude((1.0+0j))
if aperture_type == 'circle':
wf.apply_pinhole(aperture_diameter/2)
elif aperture_type == 'square':
wf.apply_slit(-aperture_diameter/2, aperture_diameter/2,-aperture_diameter/2, aperture_diameter/2)
elif aperture_type == 'gaussian':
X = wf.get_mesh_x()
Y = wf.get_mesh_y()
window = numpy.exp(- (X*X + Y*Y)/2/(aperture_diameter/2.35)**2)
wf.rescale_amplitudes(window)
else:
raise Exception("Not implemented! (accepted: circle, square, gaussian)")
if method == 'fft':
wf1 = propagate_2D_fresnel(wf, propagation_distance)
elif method == 'convolution':
wf1 = propagate_2D_fresnel_convolution(wf, propagation_distance)
elif method == 'srw':
wf1 = propagate_2D_fresnel_srw(wf, propagation_distance)
elif method == 'integral':
wf1 = propagate_2D_integral(wf, propagation_distance)
else:
raise Exception("Not implemented method: %s"%method)
if do_plot:
from srxraylib.plot.gol import plot_image
plot_image(wf.get_intensity(),1e6*wf.get_coordinate_x(),1e6*wf.get_coordinate_y(),
title="aperture intensity (%s), Diameter=%5.1f um"%
(aperture_type,1e6*aperture_diameter),xtitle="X [um]",ytitle="Y [um]",
show=0)
plot_image(wf1.get_intensity(),
1e6*wf1.get_coordinate_x()/propagation_distance,
1e6*wf1.get_coordinate_y()/propagation_distance,
title="Diffracted intensity (%s) by a %s slit of aperture %3.1f um"%
(aperture_type,method_label,1e6*aperture_diameter),
xtitle="X [urad]",ytitle="Y [urad]",
show=0)
# get the theoretical value
angle_x = wf1.get_coordinate_x() / propagation_distance
intensity_theory = get_theoretical_diffraction_pattern(angle_x,aperture_type=aperture_type,aperture_diameter=aperture_diameter,
wavelength=wavelength,normalization=True)
intensity_calculated = wf1.get_intensity()[:,wf1.size()[1]/2]
intensity_calculated /= intensity_calculated.max()
if do_plot:
from srxraylib.plot.gol import plot
plot(wf1.get_coordinate_x()*1e6/propagation_distance,intensity_calculated,
angle_x*1e6,intensity_theory,
legend=["%s H profile"%method_label,"Theoretical (far field)"],
legend_position=(0.95, 0.95),
title="%s diffraction of a %s slit of %3.1f um at wavelength of %3.1f A"%
(method_label,aperture_type,aperture_diameter*1e6,wavelength*1e10),
xtitle="X (urad)", ytitle="Intensity",xrange=[-20,20],
show=show)
return wf1.get_coordinate_x()/propagation_distance,intensity_calculated,angle_x,intensity_theory
def test_polarization(self, do_plot=0):
print("# ")
print("# Tests polarization ")
print("# ")
#
# from steps
#
x = numpy.linspace(-10, 10, 100)
y = numpy.linspace(-20, 20, 50)
# XY = numpy.meshgrid(y,x)
# sigma = 3.0
# Z = numpy.exp(- (XY[0]**2+XY[1]**2)/2/sigma**2)
xy = numpy.meshgrid(x, y)
X = xy[0].T
Y = xy[1].T
sigma = 3.0
Z = numpy.exp(-(X**2 + Y**2) / 2 / sigma**2)
Zp = 0.5 * numpy.exp(-(X**2 + Y**2) / 2 / sigma**2)
print("shape of Z", Z.shape)
wf = Wavefront2D.initialize_wavefront_from_steps(x[0],
x[1] - x[0],
y[0],
y[1] - y[0],
number_of_points=(100,
50))
print("wf shape: ", wf.size())
wf.set_complex_amplitude(Z, polarization=Polarization.SIGMA)
wf.set_complex_amplitude(Zp, polarization=Polarization.PI)
print("Total intensity: ",
wf.get_intensity(polarization=Polarization.TOTAL).sum())
print("Sigma intensity: ",
wf.get_intensity(polarization=Polarization.SIGMA).sum())
print("Pi intensity: ",
wf.get_intensity(polarization=Polarization.PI).sum())
self.assertAlmostEqual(wf.get_intensity(polarization=Polarization.TOTAL).sum(),
wf.get_intensity(polarization=Polarization.SIGMA).sum() +\
wf.get_intensity(polarization=Polarization.PI).sum())
#
# from range
#
x = numpy.linspace(-100, 100, 50)
y = numpy.linspace(-50, 50, 200)
X = numpy.outer(x, numpy.ones_like(y))
Y = numpy.outer(numpy.ones_like(x), y)
sigma = 10
Z = numpy.exp(-(X**2 + Y**2) / 2 / sigma**2) * 1j
Zp = 0.3333 * numpy.exp(-(X**2 + Y**2) / 2 / sigma**2) * 1j
print("Shapes x,y,z: ", x.shape, y.shape, Z.shape)
wf1 = Wavefront2D.initialize_wavefront_from_range(
x[0],
x[-1],
y[0],
y[-1],
number_of_points=Z.shape,
polarization=Polarization.SIGMA)
wf1.set_complex_amplitude(Z, polarization=Polarization.SIGMA)
wf1.set_complex_amplitude(Zp, polarization=Polarization.PI)
print("Total intensity: ",
wf1.get_intensity(polarization=Polarization.TOTAL).sum())
print("Sigma intensity: ",
wf1.get_intensity(polarization=Polarization.SIGMA).sum())
print("Pi intensity: ",
wf1.get_intensity(polarization=Polarization.PI).sum())
self.assertAlmostEqual(wf1.get_intensity(polarization=Polarization.TOTAL).sum(),
wf1.get_intensity(polarization=Polarization.SIGMA).sum() +\
wf1.get_intensity(polarization=Polarization.PI).sum())
#
# from arrays
#
wf2 = Wavefront2D.initialize_wavefront_from_arrays(x, y, Z, Zp)
self.assertAlmostEqual(wf2.get_intensity(polarization=Polarization.TOTAL).sum(),
wf2.get_intensity(polarization=Polarization.SIGMA).sum() +\
wf2.get_intensity(polarization=Polarization.PI).sum())
def propagation_with_lens(self,do_plot=do_plot,method='fft',
wavelength=1.24e-10,
pixelsize_x=1e-6,npixels_x=2000,pixelsize_y=1e-6,npixels_y=2000,
propagation_distance=30.0,defocus_factor=1.0,propagation_steps=1,show=1):
method_label = "fresnel (%s)"%method
print("\n# ")
print("# near field fresnel (%s) diffraction and focusing "%(method_label))
print("# ")
# \ | /
# * | | | *
# / | \
# <------- d ---------------><--------- d ------->
# d is propagation_distance
# wf = Wavefront2D.initialize_wavefront_from_steps(x_start=-pixelsize_x*npixels_x/2,
# x_step=pixelsize_x,
# y_start=-pixelsize_y*npixels_y/2,
# y_step=pixelsize_y,
# wavelength=wavelength,
# number_of_points=(npixels_x,npixels_y))
wf = 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)
spherical_or_plane_and_lens = 0
if spherical_or_plane_and_lens == 0:
# set spherical wave at the lens entrance (radius=distance)
wf.set_spherical_wave(complex_amplitude=1.0,radius=-propagation_distance)
else:
# apply lens that will focus at propagation_distance downstream the lens.
# Note that the vertical is a bit defocused
wf.set_plane_wave_from_complex_amplitude(1.0+0j)
focal_length = propagation_distance # / 2
wf.apply_ideal_lens(focal_length,focal_length)
print("Incident intensity: ",wf.get_intensity().sum())
# propagation downstream the lens to image plane
for i in range(propagation_steps):
if propagation_steps > 1:
print(">>> Propagating step %d of %d; propagation_distance=%g m"%(i+1,propagation_steps,
propagation_distance*defocus_factor/propagation_steps))
if method == 'fft':
wf = propagate_2D_fresnel(wf, propagation_distance*defocus_factor/propagation_steps)
elif method == 'convolution':
wf = propagate_2D_fresnel_convolution(wf, propagation_distance*defocus_factor/propagation_steps)
elif method == 'srw':
wf = propagate_2D_fresnel_srw(wf, propagation_distance*defocus_factor/propagation_steps)
elif method == 'fraunhofer':
wf = propagate_2D_fraunhofer(wf, propagation_distance*defocus_factor/propagation_steps)
else:
raise Exception("Not implemented method: %s"%method)
horizontal_profile = wf.get_intensity()[:,wf.size()[1]/2]
horizontal_profile /= horizontal_profile.max()
print("FWHM of the horizontal profile: %g um"%(1e6*line_fwhm(horizontal_profile)*wf.delta()[0]))
vertical_profile = wf.get_intensity()[wf.size()[0]/2,:]
vertical_profile /= vertical_profile.max()
print("FWHM of the vertical profile: %g um"%(1e6*line_fwhm(vertical_profile)*wf.delta()[1]))
if do_plot:
from srxraylib.plot.gol import plot,plot_image
plot_image(wf.get_intensity(),wf.get_coordinate_x(),wf.get_coordinate_y(),title='intensity (%s)'%method,show=0)
# plot_image(wf.get_amplitude(),wf.get_coordinate_x(),wf.get_coordinate_y(),title='amplitude (%s)'%method,show=0)
plot_image(wf.get_phase(),wf.get_coordinate_x(),wf.get_coordinate_y(),title='phase (%s)'%method,show=0)
plot(wf.get_coordinate_x(),horizontal_profile,
wf.get_coordinate_y(),vertical_profile,
legend=['Horizontal profile','Vertical profile'],title="%s"%method,show=show)
print("Output intensity: ",wf.get_intensity().sum())
return wf.get_coordinate_x(),horizontal_profile
def test_plane_wave(self, do_plot=do_plot):
#
# plane wave
#
print("# ")
print(
"# Tests for a 2D plane wave "
)
print("# ")
wavelength = 1.24e-10
wavefront_length_x = 400e-6
wavefront_length_y = wavefront_length_x
npixels_x = 1024
npixels_y = npixels_x
x = numpy.linspace(-0.5 * wavefront_length_x, 0.5 * wavefront_length_x,
npixels_x)
y = numpy.linspace(-0.5 * wavefront_length_y, 0.5 * wavefront_length_y,
npixels_y)
wavefront = Wavefront2D.initialize_wavefront_from_steps(
x_start=x[0],
x_step=numpy.abs(x[1] - x[0]),
y_start=y[0],
y_step=numpy.abs(y[1] - y[0]),
number_of_points=(npixels_x, npixels_y),
wavelength=wavelength)
# possible modifications
wavefront.set_plane_wave_from_amplitude_and_phase(5.0, numpy.pi / 2)
numpy.testing.assert_almost_equal(wavefront.get_intensity(), 25, 5)
wavefront.set_plane_wave_from_complex_amplitude(2.0 + 3j)
numpy.testing.assert_almost_equal(wavefront.get_intensity(), 13, 5)
phase_before = wavefront.get_phase()
wavefront.add_phase_shift(numpy.pi / 2)
phase_after = wavefront.get_phase()
numpy.testing.assert_almost_equal(phase_before + numpy.pi / 2,
phase_after, 5)
intensity_before = wavefront.get_intensity()
wavefront.rescale_amplitude(10.0)
intensity_after = wavefront.get_intensity()
numpy.testing.assert_almost_equal(intensity_before * 100,
intensity_after, 5)
# interpolation
wavefront.set_plane_wave_from_complex_amplitude(2.0 + 3j)
test_value1 = wavefront.get_interpolated_complex_amplitude(0.01, 1.3)
self.assertAlmostEqual((2.0 + 3j).real, test_value1.real, 5)
self.assertAlmostEqual((2.0 + 3j).imag, test_value1.imag, 5)
if do_plot:
from srxraylib.plot.gol import plot_image
plot_image(wavefront.get_intensity(),
wavefront.get_coordinate_x(),
wavefront.get_coordinate_y(),
title="Intensity (plane wave)",
show=0)
plot_image(wavefront.get_phase(),
wavefront.get_coordinate_x(),
wavefront.get_coordinate_y(),
title="Phase (plane wave)",
show=1)
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 test_interpolator(self, do_plot=do_plot):
#
# interpolator
#
print("# ")
print("# Tests for 2D interpolator ")
print("# ")
x = numpy.linspace(-10, 10, 100)
y = numpy.linspace(-20, 20, 50)
xy = numpy.meshgrid(x, y)
X = xy[0].T
Y = xy[1].T
sigma = 3.0
Z = 3 * numpy.exp(-(X**2 + Y**2) / 2 / sigma**2) + 4j
print("shape of Z", Z.shape)
wf = Wavefront2D.initialize_wavefront_from_steps(x[0],
x[1] - x[0],
y[0],
y[1] - y[0],
number_of_points=(100,
50))
print("wf shape: ", wf.size())
wf.set_complex_amplitude(Z)
x1 = 3.2
y1 = -2.5
z1 = 3 * numpy.exp(-(x1**2 + y1**2) / 2 / sigma**2) + 4j
print("complex ampl at (%g,%g): %g+%gi (exact=%g+%gi)" %
(x1, y1, wf.get_interpolated_complex_amplitude(
x1, y1).real, wf.get_interpolated_complex_amplitude(
x1, y1).imag, z1.real, z1.imag))
self.assertAlmostEqual(
wf.get_interpolated_complex_amplitude(x1, y1).real, z1.real, 4)
self.assertAlmostEqual(
wf.get_interpolated_complex_amplitude(x1, y1).imag, z1.imag, 4)
#
print(
"intensity at (%g,%g): %g (exact=%g)" %
(x1, y1, wf.get_interpolated_intensity(x1, y1), numpy.abs(z1)**2))
self.assertAlmostEqual(wf.get_interpolated_intensity(x1, y1),
numpy.abs(z1)**2, 4)
# interpolate on same grid
interpolated_complex_amplitude_on_same_grid = wf.get_interpolated_complex_amplitudes(
X, Y)
print("Shape interpolated at same grid: ",
interpolated_complex_amplitude_on_same_grid.shape)
numpy.testing.assert_array_almost_equal(
wf.get_complex_amplitude(),
interpolated_complex_amplitude_on_same_grid, 4)
print(
"Total intensity original wavefront: %g, interpolated on the same grid: %g"
% (wf.get_intensity().sum(),
(numpy.abs(interpolated_complex_amplitude_on_same_grid)**
2).sum()))
if do_plot:
from srxraylib.plot.gol import plot_image
plot_image(wf.get_intensity(),
wf.get_coordinate_x(),
wf.get_coordinate_y(),
title="Original",
show=0)
plot_image(wf.get_interpolated_intensity(X, Y),
wf.get_coordinate_x(),
wf.get_coordinate_y(),
title="interpolated on same grid",
show=1)
# rebin wavefront
# wf.set_plane_wave_from_complex_amplitude(3+4j)
wf_rebin = wf.rebin(2.0,
5.0,
0.5,
0.8,
keep_the_same_intensity=1,
set_extrapolation_to_zero=1)
print("Shape before rebinning: ", wf.size())
print("Shape after rebinning: ", wf_rebin.size())
print("Total intensity original wavefront: %g, rebinned: %g" %
(wf.get_intensity().sum(), wf_rebin.get_intensity().sum()))
if do_plot:
from srxraylib.plot.gol import plot_image
plot_image(wf.get_intensity(),
wf.get_coordinate_x(),
wf.get_coordinate_y(),
title="BEFORE REBINNING",
show=0)
plot_image(wf_rebin.get_intensity(),
wf_rebin.get_coordinate_x(),
wf_rebin.get_coordinate_y(),
title="REBINNED",
show=1)
def test_spherical_wave(self, do_plot=do_plot):
#
# plane wave
#
print("# ")
print("# Tests for a 2D spherical wave ")
print("# ")
wavelength = 1.24e-10
wavefront_length_x = 400e-6
wavefront_length_y = wavefront_length_x
npixels_x = 1024
npixels_y = npixels_x
x = numpy.linspace(-0.5 * wavefront_length_x, 0.5 * wavefront_length_x,
npixels_x)
y = numpy.linspace(-0.5 * wavefront_length_y, 0.5 * wavefront_length_y,
npixels_y)
wf1 = Wavefront2D.initialize_wavefront_from_steps(
x_start=x[0],
x_step=numpy.abs(x[1] - x[0]),
y_start=y[0],
y_step=numpy.abs(y[1] - y[0]),
number_of_points=(npixels_x, npixels_y),
wavelength=wavelength)
numpy.testing.assert_almost_equal(x, wf1.get_coordinate_x(), 9)
numpy.testing.assert_almost_equal(y, wf1.get_coordinate_y(), 9)
wf2 = Wavefront2D.initialize_wavefront_from_steps(
x_start=x[0],
x_step=numpy.abs(x[1] - x[0]),
y_start=y[0],
y_step=numpy.abs(y[1] - y[0]),
number_of_points=(npixels_x, npixels_y),
wavelength=wavelength)
numpy.testing.assert_almost_equal(x, wf2.get_coordinate_x(), 9)
numpy.testing.assert_almost_equal(y, wf2.get_coordinate_y(), 9)
# an spherical wavefront is obtained 1) by creation, 2) focusing a planewave
wf1.set_spherical_wave(-5.0, 3 + 0j)
wf1.apply_slit(-50e-6, 10e-6, -20e-6, 40e-6)
wf2.set_plane_wave_from_complex_amplitude(3 + 0j)
wf2.apply_ideal_lens(5.0, 5.0)
wf2.apply_slit(-50e-6, 10e-6, -20e-6, 40e-6)
if do_plot:
from srxraylib.plot.gol import plot_image
plot_image(wf1.get_phase(),
wf2.get_coordinate_x(),
wf2.get_coordinate_y(),
title="Phase of spherical wavefront",
show=0)
plot_image(wf2.get_phase(),
wf2.get_coordinate_x(),
wf2.get_coordinate_y(),
title="Phase of focused plane wavefront",
show=0)
plot_image(
wf1.get_phase(from_minimum_intensity=0.1),
wf2.get_coordinate_x(),
wf2.get_coordinate_y(),
title="Phase of spherical wavefront (for intensity > 0.1)",
show=0)
plot_image(
wf2.get_phase(from_minimum_intensity=0.1),
wf2.get_coordinate_x(),
wf2.get_coordinate_y(),
title="Phase of focused plane wavefront (for intensity > 0.1)",
show=1)
numpy.testing.assert_almost_equal(wf1.get_phase(), wf2.get_phase(), 5)
def propagator2d_fourier_rescaling_xy(wf,
propagation_distance,
shift_half_pixel=1,
m_x=1,
m_y=1):
wavefront = wf.duplicate()
wavenumber = wavefront.get_wavenumber()
wavelength = wavefront.get_wavelength()
# frequency for axis 1
shape = wavefront.size()
delta = wavefront.delta()
pixelsize = delta[0] # p_x[1] - p_x[0]
npixels = shape[0]
freq_nyquist = 0.5 / pixelsize
freq_n = numpy.linspace(-1.0, 1.0, npixels)
freq_x = freq_n * freq_nyquist
# frequency for axis 2
pixelsize = delta[1]
npixels = shape[1]
freq_nyquist = 0.5 / pixelsize
freq_n = numpy.linspace(-1.0, 1.0, npixels)
freq_y = freq_n * freq_nyquist
if shift_half_pixel:
freq_x = freq_x - 0.5 * numpy.abs(freq_x[1] - freq_x[0])
freq_y = freq_y - 0.5 * numpy.abs(freq_y[1] - freq_y[0])
f_x, f_y = numpy.meshgrid(freq_x, freq_y, indexing='ij')
fsq = numpy.fft.fftshift(f_x**2 / m_x + f_y**2 / m_y)
x = wavefront.get_mesh_x()
y = wavefront.get_mesh_y()
x_rescaling = wavefront.get_mesh_x() * m_x
y_rescaling = wavefront.get_mesh_y() * m_y
r1sq = x**2 * (1 - m_x) + y**2 * (1 - m_y)
r2sq = x_rescaling**2 * (m_x - 1 / m_x) + y_rescaling**2 * (m_y - 1 / m_y)
Q1 = wavenumber / 2 / propagation_distance * r1sq
Q2 = numpy.exp(-1.0j * numpy.pi * wavelength * propagation_distance * fsq)
Q3 = numpy.exp(1.0j * wavenumber / 2 / propagation_distance * r2sq)
wavefront.add_phase_shift(Q1)
fft = numpy.fft.fft2(wavefront.get_complex_amplitude())
ifft = numpy.fft.ifft2(fft * Q2) * Q3 / numpy.sqrt(m_x * m_y)
wf_propagated = Wavefront2D.initialize_wavefront_from_arrays(
wavefront.get_coordinate_x() *
m_x, ####################NON SONO SICURO CHE SIA GIUSTO QUELLO CHE RETURN
wavefront.get_coordinate_y() * m_y,
ifft,
wavelength=wavelength)
return wf_propagated
def propagation_with_lens(self,
do_plot=do_plot,
method='fft',
wavelength=1.24e-10,
pixelsize_x=1e-6,
npixels_x=2000,
pixelsize_y=1e-6,
npixels_y=2000,
propagation_distance=30.0,
defocus_factor=1.0,
propagation_steps=1,
show=1):
method_label = "fresnel (%s)" % method
print(
"\n# ")
print("# near field fresnel (%s) diffraction and focusing " %
(method_label))
print("# ")
# \ | /
# * | | | *
# / | \
# <------- d ---------------><--------- d ------->
# d is propagation_distance
# wf = Wavefront2D.initialize_wavefront_from_steps(x_start=-pixelsize_x*npixels_x/2,
# x_step=pixelsize_x,
# y_start=-pixelsize_y*npixels_y/2,
# y_step=pixelsize_y,
# wavelength=wavelength,
# number_of_points=(npixels_x,npixels_y))
wf = 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)
spherical_or_plane_and_lens = 0
if spherical_or_plane_and_lens == 0:
# set spherical wave at the lens entrance (radius=distance)
wf.set_spherical_wave(complex_amplitude=1.0,
radius=-propagation_distance)
else:
# apply lens that will focus at propagation_distance downstream the lens.
# Note that the vertical is a bit defocused
wf.set_plane_wave_from_complex_amplitude(1.0 + 0j)
focal_length = propagation_distance # / 2
wf.apply_ideal_lens(focal_length, focal_length)
print("Incident intensity: ", wf.get_intensity().sum())
# propagation downstream the lens to image plane
for i in range(propagation_steps):
if propagation_steps > 1:
print(
">>> Propagating step %d of %d; propagation_distance=%g m"
% (i + 1, propagation_steps, propagation_distance *
defocus_factor / propagation_steps))
if method == 'fft':
wf = propagate_2D_fresnel(
wf,
propagation_distance * defocus_factor / propagation_steps)
elif method == 'convolution':
wf = propagate_2D_fresnel_convolution(
wf,
propagation_distance * defocus_factor / propagation_steps)
elif method == 'srw':
wf = propagate_2D_fresnel_srw(
wf,
propagation_distance * defocus_factor / propagation_steps)
elif method == 'fraunhofer':
wf = propagate_2D_fraunhofer(
wf,
propagation_distance * defocus_factor / propagation_steps)
else:
raise Exception("Not implemented method: %s" % method)
horizontal_profile = wf.get_intensity()[:, int(wf.size()[1] / 2)]
horizontal_profile /= horizontal_profile.max()
print("FWHM of the horizontal profile: %g um" %
(1e6 * line_fwhm(horizontal_profile) * wf.delta()[0]))
vertical_profile = wf.get_intensity()[int(wf.size()[0] / 2), :]
vertical_profile /= vertical_profile.max()
print("FWHM of the vertical profile: %g um" %
(1e6 * line_fwhm(vertical_profile) * wf.delta()[1]))
if do_plot:
from srxraylib.plot.gol import plot, plot_image
plot_image(wf.get_intensity(),
wf.get_coordinate_x(),
wf.get_coordinate_y(),
title='intensity (%s)' % method,
show=0)
# plot_image(wf.get_amplitude(),wf.get_coordinate_x(),wf.get_coordinate_y(),title='amplitude (%s)'%method,show=0)
plot_image(wf.get_phase(),
wf.get_coordinate_x(),
wf.get_coordinate_y(),
title='phase (%s)' % method,
show=0)
plot(wf.get_coordinate_x(),
horizontal_profile,
wf.get_coordinate_y(),
vertical_profile,
legend=['Horizontal profile', 'Vertical profile'],
title="%s" % method,
show=show)
print("Output intensity: ", wf.get_intensity().sum())
return wf.get_coordinate_x(), horizontal_profile
def test_spherical_wave(self,do_plot=do_plot):
#
# plane wave
#
print("# ")
print("# Tests for a 2D spherical wave ")
print("# ")
wavelength = 1.24e-10
wavefront_length_x = 400e-6
wavefront_length_y = wavefront_length_x
npixels_x = 1024
npixels_y = npixels_x
x = numpy.linspace(-0.5*wavefront_length_x,0.5*wavefront_length_x,npixels_x)
y = numpy.linspace(-0.5*wavefront_length_y,0.5*wavefront_length_y,npixels_y)
wf1 = Wavefront2D.initialize_wavefront_from_steps(
x_start=x[0],x_step=numpy.abs(x[1]-x[0]),
y_start=y[0],y_step=numpy.abs(y[1]-y[0]),
number_of_points=(npixels_x,npixels_y),wavelength=wavelength)
numpy.testing.assert_almost_equal(x,wf1.get_coordinate_x(),9)
numpy.testing.assert_almost_equal(y,wf1.get_coordinate_y(),9)
wf2 = Wavefront2D.initialize_wavefront_from_steps(
x_start=x[0],x_step=numpy.abs(x[1]-x[0]),
y_start=y[0],y_step=numpy.abs(y[1]-y[0]),
number_of_points=(npixels_x,npixels_y),wavelength=wavelength)
numpy.testing.assert_almost_equal(x,wf2.get_coordinate_x(),9)
numpy.testing.assert_almost_equal(y,wf2.get_coordinate_y(),9)
# an spherical wavefront is obtained 1) by creation, 2) focusing a planewave
wf1.set_spherical_wave(-5.0, 3+0j)
wf1.apply_slit(-50e-6,10e-6,-20e-6,40e-6)
wf2.set_plane_wave_from_complex_amplitude(3+0j)
wf2.apply_ideal_lens(5.0,5.0)
wf2.apply_slit(-50e-6,10e-6,-20e-6,40e-6)
if do_plot:
from srxraylib.plot.gol import plot_image
plot_image(wf1.get_phase(),wf2.get_coordinate_x(),wf2.get_coordinate_y(),
title="Phase of spherical wavefront",show=0)
plot_image(wf2.get_phase(),wf2.get_coordinate_x(),wf2.get_coordinate_y(),
title="Phase of focused plane wavefront",show=0)
plot_image(wf1.get_phase(from_minimum_intensity=0.1),wf2.get_coordinate_x(),wf2.get_coordinate_y(),
title="Phase of spherical wavefront (for intensity > 0.1)",show=0)
plot_image(wf2.get_phase(from_minimum_intensity=0.1),wf2.get_coordinate_x(),wf2.get_coordinate_y(),
title="Phase of focused plane wavefront (for intensity > 0.1)",show=1)
numpy.testing.assert_almost_equal(wf1.get_phase(),wf2.get_phase(),5)
def propagate_2D_integral(wavefront, propagation_distance,
shuffle_interval=0,calculate_grid_only=1):
"""
2D Fresnel-Kirchhoff propagator via simplified integral
NOTE: this propagator is experimental and much less performant than the ones using Fourier Optics
Therefore, it is not recommended to use.
:param wavefront:
:param propagation_distance: propagation distance
:param shuffle_interval: it is known that this method replicates the central diffraction spot
The distace of the replica is proportional to 1/pixelsize
To avoid that, it is possible to change a bit (randomly) the coordinates
of the wavefront. shuffle_interval controls this shift: 0=No shift. A typical
value can be 1e5.
The result shows a diffraction pattern without replica but with much noise.
:param calculate_grid_only: if set, it calculates only the horizontal and vertical profiles, but returns the
full image with the other pixels to zero. This is useful when calculating large arrays,
so it is set as the default.
:return: a new 2D wavefront object with propagated wavefront
"""
#
# Fresnel-Kirchhoff integral (neglecting inclination factor)
#
if calculate_grid_only == 0:
#
# calculation over the whole detector area
#
p_x = wavefront.get_coordinate_x()
p_y = wavefront.get_coordinate_y()
wavelength = wavefront.get_wavelength()
amplitude = wavefront.get_complex_amplitude()
det_x = p_x.copy()
det_y = p_y.copy()
p_X = wavefront.get_mesh_x()
p_Y = wavefront.get_mesh_y()
det_X = p_X
det_Y = p_Y
amplitude_propagated = numpy.zeros_like(amplitude,dtype='complex')
wavenumber = 2 * numpy.pi / wavelength
for i in range(det_x.size):
for j in range(det_y.size):
if shuffle_interval == 0:
rd_x = 0.0
rd_y = 0.0
else:
rd_x = (numpy.random.rand(p_x.size,p_y.size)-0.5)*shuffle_interval
rd_y = (numpy.random.rand(p_x.size,p_y.size)-0.5)*shuffle_interval
r = numpy.sqrt( numpy.power(p_X + rd_x - det_X[i,j],2) +
numpy.power(p_Y + rd_y - det_Y[i,j],2) +
numpy.power(propagation_distance,2) )
amplitude_propagated[i,j] = (amplitude / r * numpy.exp(1.j * wavenumber * r)).sum()
wavefront2 = Wavefront2D.initialize_wavefront_from_arrays(det_x,det_y,amplitude_propagated)
else:
x = wavefront.get_coordinate_x()
y = wavefront.get_coordinate_y()
X = wavefront.get_mesh_x()
Y = wavefront.get_mesh_y()
wavenumber = 2 * numpy.pi / wavefront.get_wavelength()
amplitude = wavefront.get_complex_amplitude()
used_indices = wavefront.get_mask_grid(width_in_pixels=(1,1),number_of_lines=(1,1))
indices_x = wavefront.get_mesh_indices_x()
indices_y = wavefront.get_mesh_indices_y()
indices_x_flatten = indices_x[numpy.where(used_indices == 1)].flatten()
indices_y_flatten = indices_y[numpy.where(used_indices == 1)].flatten()
X_flatten = X[numpy.where(used_indices == 1)].flatten()
Y_flatten = Y[numpy.where(used_indices == 1)].flatten()
complex_amplitude_propagated = amplitude*0
print("propagate_2D_integral: Calculating %d points from a total of %d x %d = %d"%(
X_flatten.size,amplitude.shape[0],amplitude.shape[1],amplitude.shape[0]*amplitude.shape[1]))
for i in range(X_flatten.size):
r = numpy.sqrt( numpy.power(wavefront.get_mesh_x() - X_flatten[i],2) +
numpy.power(wavefront.get_mesh_y() - Y_flatten[i],2) +
numpy.power(propagation_distance,2) )
complex_amplitude_propagated[indices_x_flatten[i],indices_y_flatten[i]] = (amplitude / r * numpy.exp(1.j * wavenumber * r)).sum()
wavefront2 = Wavefront2D.initialize_wavefront_from_arrays(x,y,complex_amplitude_propagated,wavefront.get_wavelength())
return wavefront2
def test_propagate_2D_fraunhofer(self,
do_plot=do_plot,
aperture_type='square',
aperture_diameter=40e-6,
pixelsize_x=1e-6,
pixelsize_y=1e-6,
npixels_x=1024,
npixels_y=1024,
wavelength=1.24e-10):
"""
:param do_plot: 0=No plot, 1=Do plot
:param aperture_type: 'circle' 'square' 'gaussian' (Gaussian sigma = aperture_diameter/2.35)
:param aperture_diameter:
:param pixelsize_x:
:param pixelsize_y:
:param npixels_x:
:param npixels_y:
:param wavelength:
:return:
"""
print(
"\n# ")
print(
"# far field 2D (fraunhofer) diffraction from a square aperture ")
print("# ")
method = "fraunhofer"
print(
"Fraunhoffer diffraction valid for distances > > a^2/lambda = %f m"
% ((aperture_diameter / 2)**2 / wavelength))
# wf = Wavefront2D.initialize_wavefront_from_steps(x_start=-pixelsize_x*npixels_x/2,
# x_step=pixelsize_x,
# y_start=-pixelsize_y*npixels_y/2,
# y_step=pixelsize_y,
# wavelength=wavelength,
# number_of_points=(npixels_x,npixels_y))
wf = 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.set_plane_wave_from_complex_amplitude((1.0 + 0j))
if aperture_type == 'circle':
wf.apply_pinhole(aperture_diameter / 2)
elif aperture_type == 'square':
wf.apply_slit(-aperture_diameter / 2, aperture_diameter / 2,
-aperture_diameter / 2, aperture_diameter / 2)
elif aperture_type == 'gaussian':
X = wf.get_mesh_x()
Y = wf.get_mesh_y()
window = numpy.exp(-(X * X + Y * Y) / 2 /
(aperture_diameter / 2.35)**2)
wf.rescale_amplitudes(window)
else:
raise Exception(
"Not implemented! (accepted: circle, square, gaussian)")
wf1 = propagate_2D_fraunhofer(
wf, propagation_distance=1.0
) # propagating at 1 m means the result is like in angles
if do_plot:
plot_image(wf.get_intensity(),
1e6 * wf.get_coordinate_x(),
1e6 * wf.get_coordinate_y(),
title="aperture intensity (%s), Diameter=%5.1f um" %
(aperture_type, 1e6 * aperture_diameter),
xtitle="X [um]",
ytitle="Y [um]",
show=0)
plot_image(
wf1.get_intensity(),
1e6 * wf1.get_coordinate_x(),
1e6 * wf1.get_coordinate_y(),
title=
"2D Diffracted intensity (%s) by a %s slit of aperture %3.1f um"
% (aperture_type, method, 1e6 * aperture_diameter),
xtitle="X [urad]",
ytitle="Y [urad]",
show=0)
angle_x = wf1.get_coordinate_x(
) # + 0.5*wf1.delta()[0] # shifted of half-pixel!!!
intensity_theory = get_theoretical_diffraction_pattern(
angle_x,
aperture_type=aperture_type,
aperture_diameter=aperture_diameter,
wavelength=wavelength,
normalization=True)
intensity_calculated = wf1.get_intensity()[:, int(wf1.size()[1] / 2)]
intensity_calculated /= intensity_calculated.max()
if do_plot:
plot(
wf1.get_coordinate_x() * 1e6,
intensity_calculated,
angle_x * 1e6,
intensity_theory,
legend=["Calculated (FT) H profile", "Theoretical"],
legend_position=(0.95, 0.95),
title=
"2D Fraunhofer Diffraction of a %s slit of %3.1f um at wavelength of %3.1f A"
% (aperture_type, aperture_diameter * 1e6, wavelength * 1e10),
xtitle="X (urad)",
ytitle="Intensity",
xrange=[-80, 80])
numpy.testing.assert_almost_equal(intensity_calculated,
intensity_theory, 1)
def get_wavefront2d(self,i_mode,wavelength):
return Wavefront2D.initialize_wavefront_from_arrays(self.x_coordinates(),self.y_coordinates(),
self.mode(i_mode),wavelength=wavelength)
def test_initializers(self, do_plot=do_plot):
print("# ")
print("# Tests for initializars (2D) ")
print("# ")
x = numpy.linspace(-100, 100, 50)
y = numpy.linspace(-50, 50, 200)
XY = numpy.meshgrid(x, y)
X = XY[0].T
Y = XY[1].T
sigma = 10
Z = numpy.exp(-(X**2 + Y**2) / 2 / sigma**2) * 1j
print("Shapes x,y,z: ", x.shape, y.shape, Z.shape)
wf0 = Wavefront2D.initialize_wavefront_from_steps(
x[0],
numpy.abs(x[1] - x[0]),
y[0],
numpy.abs(y[1] - y[0]),
number_of_points=Z.shape)
wf0.set_complex_amplitude(Z)
wf1 = Wavefront2D.initialize_wavefront_from_range(
x[0], x[-1], y[0], y[-1], number_of_points=Z.shape)
wf1.set_complex_amplitude(Z)
wf2 = Wavefront2D.initialize_wavefront_from_arrays(x, y, Z)
if do_plot:
from srxraylib.plot.gol import plot_image
plot_image(wf0.get_intensity(),
wf0.get_coordinate_x(),
wf0.get_coordinate_y(),
title="initialize_wavefront_from_steps",
show=0)
plot_image(wf1.get_intensity(),
wf1.get_coordinate_x(),
wf1.get_coordinate_y(),
title="initialize_wavefront_from_range",
show=0)
plot_image(wf2.get_intensity(),
wf2.get_coordinate_x(),
wf2.get_coordinate_y(),
title="initialize_wavefront_from_arrays",
show=1)
numpy.testing.assert_almost_equal(
numpy.abs(Z)**2, wf0.get_intensity(), 11)
numpy.testing.assert_almost_equal(
numpy.abs(Z)**2, wf1.get_intensity(), 11)
numpy.testing.assert_almost_equal(
numpy.abs(Z)**2, wf2.get_intensity(), 11)
numpy.testing.assert_almost_equal(x, wf0.get_coordinate_x(), 11)
numpy.testing.assert_almost_equal(x, wf1.get_coordinate_x(), 11)
numpy.testing.assert_almost_equal(x, wf2.get_coordinate_x(), 11)
numpy.testing.assert_almost_equal(y, wf0.get_coordinate_y(), 11)
numpy.testing.assert_almost_equal(y, wf1.get_coordinate_y(), 11)
numpy.testing.assert_almost_equal(y, wf2.get_coordinate_y(), 11)
