def propagate_single_mode(af, i, beamline):
from wofry.propagator.propagator import PropagationManager, PropagationElements, PropagationParameters
from syned.beamline.beamline_element import BeamlineElement
from syned.beamline.element_coordinates import ElementCoordinates
from wofry.propagator.propagators2D.fresnel_zoom_xy import FresnelZoomXY2D
from wofry.propagator.wavefront2D.generic_wavefront import GenericWavefront2D
from wofry.beamline.optical_elements.ideal_elements.screen import WOScreen
mi = af.coherentMode(i)
evi = af.eigenvalue(i)
print("propagating mode index", i, evi, mi.shape)
input_wavefront = GenericWavefront2D.initialize_wavefront_from_arrays(
x_array=af.xCoordinates(),
y_array=af.yCoordinates(),
z_array=mi * numpy.sqrt(evi),
)
i0 = input_wavefront.get_integrated_intensity()
input_wavefront.set_photon_energy(17226.0)
propagator = PropagationManager.Instance()
try:
propagator.add_propagator(FresnelZoomXY2D())
except:
pass
wfp = beamline.propagate(input_wavefront, propagator)
i1 = wfp[-1].get_integrated_intensity()
return wfp, i1 / i0
def propagate_numpy_wavefront(cls, filename_in, filename_out, beamline, mypropagator, return_wavefront_list=True):
file_content = numpy.load(filename_in)
e_field = file_content["e_field"]
coordinates = file_content["coordinates"]
energies = file_content["energies"]
x = numpy.linspace(coordinates[0], coordinates[1], e_field.shape[1])
y = numpy.linspace(coordinates[2], coordinates[3], e_field.shape[2])
wofry_wf_in = GenericWavefront2D.initialize_wavefront_from_arrays(x, y, e_field[0, :, :, 0].copy())
wofry_wf_in.set_photon_energy(energies[0])
# wofry_wf_out_list = cls.propagate_classmethod(wofry_wf_in,beamline,mypropagator)
wofry_wf_out = beamline.propagate(wofry_wf_in, mypropagator, return_wavefront_list=return_wavefront_list)
if return_wavefront_list:
wofry_wf_out_list = wofry_wf_out
wofry_wf_out = wofry_wf_out_list[-1]
e_field[0, :, :, 0] = wofry_wf_out.get_complex_amplitude()
coordinates[0] = wofry_wf_out.get_coordinate_x()[0]
coordinates[1] = wofry_wf_out.get_coordinate_x()[-1]
coordinates[2] = wofry_wf_out.get_coordinate_y()[0]
coordinates[3] = wofry_wf_out.get_coordinate_y()[-1]
numpy.savez(filename_out,
e_field=e_field,
coordinates=coordinates,
energies=energies)
if return_wavefront_list:
return wofry_wf_out_list
else:
return wofry_wf_out
def load_h5_file(filename, filepath):
try:
f = h5py.File(filename, 'r')
mesh_X = f[filepath + "/wfr_mesh_X"].value
mesh_Y = f[filepath + "/wfr_mesh_Y"].value
complex_amplitude_s = f[filepath + "/wfr_complex_amplitude_s"].value.T
wfr = GenericWavefront2D.initialize_wavefront_from_arrays(
x_array=numpy.linspace(mesh_X[0], mesh_X[1], int(mesh_X[2])),
y_array=numpy.linspace(mesh_Y[0], mesh_Y[1], int(mesh_Y[2])),
z_array=complex_amplitude_s)
wfr.set_photon_energy(f[filepath + "/wfr_photon_energy"].value)
f.close()
return wfr
except:
raise Exception("Failed to load 2D wavefront to h5 file: " + filename)
def propagate_numpy_wavefront(cls,
filename_in,
filename_out,
beamline,
mypropagator,
return_wavefront_list=True):
file_content = numpy.load(filename_in)
e_field = file_content["e_field"]
coordinates = file_content["coordinates"]
energies = file_content["energies"]
x = numpy.linspace(coordinates[0], coordinates[1], e_field.shape[1])
y = numpy.linspace(coordinates[2], coordinates[3], e_field.shape[2])
wofry_wf_in = GenericWavefront2D.initialize_wavefront_from_arrays(
x, y, e_field[0, :, :, 0].copy())
wofry_wf_in.set_photon_energy(energies[0])
# wofry_wf_out_list = cls.propagate_classmethod(wofry_wf_in,beamline,mypropagator)
wofry_wf_out = beamline.propagate(
wofry_wf_in,
mypropagator,
return_wavefront_list=return_wavefront_list)
if return_wavefront_list:
wofry_wf_out_list = wofry_wf_out
wofry_wf_out = wofry_wf_out_list[-1]
e_field[0, :, :, 0] = wofry_wf_out.get_complex_amplitude()
coordinates[0] = wofry_wf_out.get_coordinate_x()[0]
coordinates[1] = wofry_wf_out.get_coordinate_x()[-1]
coordinates[2] = wofry_wf_out.get_coordinate_y()[0]
coordinates[3] = wofry_wf_out.get_coordinate_y()[-1]
numpy.savez(filename_out,
e_field=e_field,
coordinates=coordinates,
energies=energies)
if return_wavefront_list:
return wofry_wf_out_list
else:
return wofry_wf_out
def fresnel(wavefront, propagation_distance, shift_half_pixel=False):
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]) )
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]))
wf_propagated = GenericWavefront2D.initialize_wavefront_from_arrays(
x_array=wavefront.get_coordinate_x(),
y_array=wavefront.get_coordinate_y(),
z_array=numpy.fft.ifft2(fft),
wavelength=wavelength)
return wf_propagated
def send_mode(self):
wf = GenericWavefront2D.initialize_wavefront_from_arrays(
self.af.x_coordinates(), self.af.y_coordinates(),
self.af.mode(self.MODE_INDEX))
wf.set_photon_energy(self.af.photon_energy())
ampl = wf.get_complex_amplitude()
if self.TYPE_PRESENTATION == 5:
eigen = self.af.eigenvalues_old()
wf.set_complex_amplitude(ampl * numpy.sqrt(eigen[self.MODE_INDEX]))
else:
wf.set_complex_amplitude(ampl)
beamline = WOBeamline(light_source=self.get_light_source())
print(">>> sending mode: ", int(self.MODE_INDEX))
self.send("WofryData", WofryData(wavefront=wf, beamline=beamline))
# script
self.wofry_python_script.set_code(beamline.to_python_code())
def propagate_wavefront(cls,
wavefront,
propagation_distance,
shift_half_pixel=False):
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
wavefront_out = GenericWavefront2D.initialize_wavefront_from_arrays(
x_array=wavefront.get_coordinate_x(),
y_array=wavefront.get_coordinate_y(),
z_array=fftconvolve(wavefront.get_complex_amplitude(),
kernel,
mode='same'),
wavelength=wavelength)
# added [email protected] 2018-03-23 to conserve energy - TODO: review method!
wavefront_out.rescale_amplitude(
numpy.sqrt(wavefront.get_intensity().sum() /
wavefront_out.get_intensity().sum()))
return wavefront_out
def propagate_wavefront(cls,
wavefront1,
propagation_distance,
magnification_x=1.0,
magnification_y=1.0,
shift_half_pixel=False):
wavefront = wavefront1.duplicate()
wavelength = wavefront.get_wavelength()
wavenumber = wavefront.get_wavenumber()
shape = wavefront.size()
delta = wavefront.delta()
# frequency for axis 1
pixelsize = delta[0]
npixels = shape[0]
freq_nyquist = 0.5 / pixelsize
freq_n = numpy.linspace(-1.0, 1.0, npixels)
freq_x0 = freq_n * freq_nyquist
freq_x1 = numpy.fft.fftfreq(npixels, pixelsize)
freq_x1 = numpy.fft.ifftshift(freq_x1)
# 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_y0 = freq_n * freq_nyquist
freq_y1 = numpy.fft.fftfreq(npixels, pixelsize)
freq_y1 = numpy.fft.ifftshift(freq_y1)
print(freq_x0[0:10], freq_x1[0:10])
# It happens that with method=0 (old) the propagation of a centro-symmetric beam
# is not longer center but shifted.
# This is due to "shifted" storage of the frequencies that is dependent on the
# even or odd number of pixels.
# It seems that with the new method (method=1) the beam is center.
# Note: The new method ignores the shift_half_pixel keyword
# See also doscussion with V. Favre-Nicolin email to srio on 2018-05-02
method = 1 # 0-old, 1-new
if method == 0:
freq_x = freq_x0
freq_y = freq_y0
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])
else:
freq_x = freq_x1
freq_y = freq_y1
f_x, f_y = numpy.meshgrid(freq_x, freq_y, indexing='ij')
fsq = numpy.fft.fftshift(f_x**2 / magnification_x +
f_y**2 / magnification_y)
x = wavefront.get_mesh_x()
y = wavefront.get_mesh_y()
x_rescaling = wavefront.get_mesh_x() * magnification_x
y_rescaling = wavefront.get_mesh_y() * magnification_y
r1sq = x**2 * (1 - magnification_x) + y**2 * (1 - magnification_y)
r2sq = x_rescaling**2 * (
(magnification_x - 1) / magnification_x) + y_rescaling**2 * (
(magnification_y - 1) / magnification_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(
magnification_x * magnification_y)
wf_propagated = GenericWavefront2D.initialize_wavefront_from_arrays(
x_array=wavefront.get_coordinate_x() * magnification_x,
y_array=wavefront.get_coordinate_y() * magnification_y,
z_array=ifft,
wavelength=wavelength)
return wf_propagated
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 = GenericWavefront2D.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 = GenericWavefront2D.initialize_wavefront_from_range(
x[0], x[-1], y[0], y[-1], number_of_points=Z.shape)
wf1.set_complex_amplitude(Z)
wf2 = GenericWavefront2D.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 propagate_wavefront(cls,
wavefront1,
propagation_distance,
magnification_x=1.0,
magnification_y=1.0,
shift_half_pixel=False):
wavefront = wavefront1.duplicate()
wavelength = wavefront.get_wavelength()
wavenumber = wavefront.get_wavenumber()
shape = wavefront.size()
delta = wavefront.delta()
pixelsize = delta[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 / magnification_x +
f_y**2 / magnification_y)
x = wavefront.get_mesh_x()
y = wavefront.get_mesh_y()
x_rescaling = wavefront.get_mesh_x() * magnification_x
y_rescaling = wavefront.get_mesh_y() * magnification_y
r1sq = x**2 * (1 - magnification_x) + y**2 * (1 - magnification_y)
r2sq = x_rescaling**2 * (
(magnification_x - 1) / magnification_x) + y_rescaling**2 * (
(magnification_y - 1) / magnification_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(
magnification_x * magnification_y)
wf_propagated = GenericWavefront2D.initialize_wavefront_from_arrays(
x_array=wavefront.get_coordinate_x() * magnification_x,
y_array=wavefront.get_coordinate_y() * magnification_y,
z_array=ifft,
wavelength=wavelength)
return wf_propagated
def propagate_wavefront(cls,wavefront,propagation_distance,shuffle_interval=False,calculate_grid_only=True):
#
# Fresnel-Kirchhoff integral (neglecting inclination factor)
#
if not calculate_grid_only:
#
# 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 not shuffle_interval:
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()
output_wavefront = GenericWavefront2D.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[int(indices_x_flatten[i]),int(indices_y_flatten[i])] = (amplitude / r * numpy.exp(1.j * wavenumber * r)).sum()
output_wavefront = GenericWavefront2D.initialize_wavefront_from_arrays(x_array=x,
y_array=y,
z_array=complex_amplitude_propagated,
wavelength=wavefront.get_wavelength())
# added [email protected] 2018-03-23 to conserve energy - TODO: review method!
output_wavefront.rescale_amplitude( numpy.sqrt(wavefront.get_intensity().sum() /
output_wavefront.get_intensity().sum()))
return output_wavefront
def write_file(self):
self.setStatusMessage("")
try:
if not self.af is None:
congruence.checkDir(self.file_name)
if self.TYPE_OF_OUTPUT == 0: # ['COMSYL hdf5 with multi-mode','WOFRY hdf5 with multi-mode','WOFRY multiple files'
if self.ALL_MODES:
self.af.write_h5(self.file_name,
maximum_number_of_modes=None)
else:
self.af.write_h5(self.file_name,
maximum_number_of_modes=self.MODE_TO)
path, file_name = os.path.split(self.file_name)
self.setStatusMessage("File Out: " + file_name)
else: # WOFRY
if self.ALL_MODES == 0 and (self.MODE_TO <
self.af.number_of_modes()):
nmax = self.MODE_TO
else:
nmax = self.af.number_of_modes()
for i in range(nmax + 1):
eigenvalue = numpy.real(self.af.eigenvalue(i), )
eigenfunction = self.af.mode(i)
w = GenericWavefront2D.initialize_wavefront_from_arrays(
self.af.x_coordinates(), self.af.y_coordinates(),
eigenfunction * numpy.sqrt(eigenvalue))
w.set_photon_energy(self.af.photon_energy())
if self.TYPE_OF_OUTPUT == 1: #
file_name = self.file_name
subgroupname = "mode" + self.index_format % (i)
overwrite = False
elif self.TYPE_OF_OUTPUT == 2:
subgroupname = "mode" + self.index_format % (i)
file_name = self.file_name.split(".")[0] + "_" + (
self.index_format % i) + ".h5"
overwrite = True
if i == 0:
w.save_h5_file(file_name,
subgroupname=subgroupname,
intensity=True,
phase=True,
overwrite=True)
else:
w.save_h5_file(file_name,
subgroupname=subgroupname,
intensity=True,
phase=True,
overwrite=overwrite)
path, file_name = os.path.split(file_name)
self.setStatusMessage("File Out: " + file_name)
# path, file_name = os.path.split(self.file_name)
# self.setStatusMessage("File Out: " + file_name)
else:
QMessageBox.critical(self, "Error", "COMSYL modes not present",
QMessageBox.Ok)
except Exception as exception:
QMessageBox.critical(self, "Error", str(exception), QMessageBox.Ok)
X2 = numpy.outer(numpy.ones_like(x), x)
# plot_image(X2,x,x)
z = csd(X1, X2, sigmaI, sigmaMu)
zi = get_intensity_from_csd(z)
if do_plot:
plot_image(z,
1e6 * x,
1e6 * x,
xtitle="X1",
ytitle="X2",
title="CDS before propagation",
show=0)
plot(1e6 * x, zi, show=0)
wf = GenericWavefront2D.initialize_wavefront_from_arrays(x, x, z)
wf.set_photon_energy(23000)
# wfp = fresnel(wf,propagation_distance=propagation_distance)
magnification = 5
wfp = fresnel_zoom(wf,
propagation_distance=propagation_distance,
magnification_x=magnification,
magnification_y=magnification)
zp = wfp.get_intensity()
zpi = get_intensity_from_csd(zp)
fwhm_z, tmp, tmp = get_fwhm(zi / zi.max(), 1e6 * x)
fwhm_zp, tmp, tmp = get_fwhm(zpi / zpi.max(), 1e6 * x)