def initialize_from_file(cls,filename):
filename_extension = filename.split('.')[-1]
try:
if filename_extension == "h5":
return cls.initialize_from_h5_file(filename)
elif filename_extension == "npz":
data_dict = AutocorrelationFunctionIO.load(filename)
af = AutocorrelationFunction.fromDictionary(data_dict)
af._io._setWasFileLoaded(filename)
return CompactAFReader(af,data_dict,filename)
elif filename_extension == "npy":
filename_without_extension = ('.').join(filename.split('.')[:-1])
filename_with_npz_extension = filename_without_extension+".npz"
data_dict = AutocorrelationFunctionIO.load(filename_with_npz_extension)
af = AutocorrelationFunction.fromDictionary(data_dict)
af._io._setWasFileLoaded(filename)
return CompactAFReader(af,data_dict,filename_with_npz_extension)
else:
raise FileExistsError("Please enter a file with .npy, .npz or .h5 extension")
except:
raise FileExistsError("Error reading file")
def testPropagation(self):
filename = "../calculations/new_s2_z0_0.npz"
af = AutocorrelationFunction.load(filename)
srw_beamline = createBeamline()
propagator = AutocorrelationFunctionPropagator(srw_beamline)
propagator.setMaximumMode(30)
propagator.propagate(af, slow_low_memory=True)
af.save("afptest.npz")
srw_beamline = createBeamline2()
propagator2 = AutocorrelationFunctionPropagator(srw_beamline)
propagator2.propagate(af)
af.save("afptest2.npz")
def calculateAutocorrelation(self, electron_beam, undulator, info):
configuration = self.configuration()
# electron_beam = ElectronBeam(energy_in_GeV=6.04,
# energy_spread=0,
# average_current=0.2000,
# electrons=1)
sigma_matrix = SigmaMatrixFromCovariance(
xx=electron_beam._moment_xx,
xxp=electron_beam._moment_xxp,
xpxp=electron_beam._moment_xpxp,
yy=electron_beam._moment_yy,
yyp=electron_beam._moment_yyp,
ypyp=electron_beam._moment_ypyp,
sigma_dd=electron_beam._energy_spread,
)
resonance_energy = int(
undulator.resonance_energy(electron_beam.gamma(), 0, 0))
energy = resonance_energy * configuration.detuningParameter()
if configuration.virtualSourcePosition() == "":
if sigma_matrix.isAlphaZero():
source_position = VIRTUAL_SOURCE_CENTER
else:
source_position = VIRTUAL_SOURCE_ENTRANCE
else:
source_position = configuration.virtualSourcePosition()
wavefront_builder = WavefrontBuilder(
undulator=undulator,
sampling_factor=configuration.samplingFactor(),
min_dimension_x=configuration.
exitSlitWavefrontMinimalSizeHorizontal(),
max_dimension_x=configuration.
exitSlitWavefrontMaximalSizeHorizontal(),
min_dimension_y=configuration.exitSlitWavefrontMinimalSizeVertical(
),
max_dimension_y=configuration.exitSlitWavefrontMaximalSizeVertical(
),
energy=energy,
source_position=source_position)
# from comsyl.waveoptics.WavefrontBuilderPySRU import WavefrontBuilderPySRU
# wavefront_builder = WavefrontBuilderPySRU(undulator=undulator,
# sampling_factor=configuration.samplingFactor(),
# min_dimension_x=configuration.exitSlitWavefrontMinimalSizeHorizontal(),
# max_dimension_x=configuration.exitSlitWavefrontMaximalSizeHorizontal(),
# min_dimension_y=configuration.exitSlitWavefrontMinimalSizeVertical(),
# max_dimension_y=configuration.exitSlitWavefrontMaximalSizeVertical(),
# energy=energy)
info.setSourcePosition(source_position)
e_0, sigma_e, beam_energies = self._determineBeamEnergies(
electron_beam, sigma_matrix, configuration.beamEnergies())
# determineZ(electron_beam, wavefront_builder, sigma_matrix.sigma_x(), sigma_matrix.sigma_x_prime(),
# sigma_matrix.sigma_y(), sigma_matrix.sigma_y_prime())
sorted_beam_energies = beam_energies[np.abs(beam_energies).argsort()]
field_x_coordinates = None
field_y_coordinates = None
exit_slit_wavefront = None
first_wavefront = None
weighted_fields = None
for i_beam_energy, beam_energy in enumerate(sorted_beam_energies):
# TDOO: recheck normalization, especially for delta coupled beams.
if len(sorted_beam_energies) > 1:
stepwidth_beam = np.diff(beam_energies)[0]
weight = (1 / (2 * np.pi * sigma_e**2)**0.5) * np.exp(
-beam_energy**2 / (2 * sigma_e**2)) * stepwidth_beam
else:
weight = 1.0
electron_beam._energy = e_0 + beam_energy
log("%i/%i: doing energy: %e with weight %e" %
(i_beam_energy + 1, len(beam_energies), electron_beam.energy(),
weight))
# Prepare e_field
if field_x_coordinates is None or field_y_coordinates is None:
wavefront = wavefront_builder.build(electron_beam,
xp=0.0,
yp=0.0,
z_offset=0.0)
reference_wavefront = wavefront.toNumpyWavefront()
else:
wavefront = wavefront_builder.buildOnGrid(reference_wavefront,
electron_beam,
xp=0.0,
yp=0.0,
z_offset=0.0)
try:
Rx, dRx, Ry, dRy = wavefront.instantRadii()
except AttributeError:
Rx, dRx, Ry, dRy = 0, 0, 0, 0
energy = wavefront.energies()[0]
wavefront = wavefront.toNumpyWavefront()
if exit_slit_wavefront is None:
exit_slit_wavefront = wavefront.clone()
wavefront = wavefront_builder.createReferenceWavefrontAtVirtualSource(
Rx, dRx, Ry, dRy, configuration, source_position, wavefront)
if self.configuration().useGaussianWavefront() == "true":
gaussian_wavefront_builder = GaussianWavefrontBuilder()
wavefront = gaussian_wavefront_builder.fromWavefront(
wavefront, info)
if field_x_coordinates is None or field_y_coordinates is None:
wavefront = wavefront.asCenteredGrid(resample_x=1.0,
resample_y=1.0)
field_x_coordinates = np.array(
wavefront.absolute_x_coordinates())
field_y_coordinates = np.array(
wavefront.absolute_y_coordinates())
else:
wavefront = wavefront.asCenteredGrid(field_x_coordinates,
field_y_coordinates)
size_matrix = self._estimateMemoryConsumption(wavefront)
if self.adjustMemoryConsumption():
self._performMemoryConsumptionAdjustment(
sigma_matrix, undulator, info, size_matrix)
exit(0)
if first_wavefront is None:
first_wavefront = wavefront.clone()
if weighted_fields is None:
weighted_fields = np.zeros(
(len(sorted_beam_energies), len(field_x_coordinates),
len(field_y_coordinates)),
dtype=np.complex128)
weighted_fields[i_beam_energy, :, :] = np.sqrt(
weight) * wavefront.E_field_as_numpy()[0, :, :, 0].copy()
log("Broadcasting electrical fields")
weighted_fields = mpi.COMM_WORLD.bcast(weighted_fields, root=0)
static_electron_density, work_matrix = self.calculateAutocorrelationForEnergy(
wavefront, weighted_fields, sigma_matrix)
electron_beam._energy = e_0
if isinstance(work_matrix, AutocorrelationOperator):
log("Setting up eigenmoder")
eigenmoder = Eigenmoder(field_x_coordinates, field_y_coordinates)
log("Starting eigenmoder")
eigenvalues_spatial, eigenvectors_parallel = eigenmoder.eigenmodes(
work_matrix, work_matrix.numberModes(), do_not_gather=True)
if isMaster():
total_spatial_mode_intensity = eigenvalues_spatial.sum(
) * work_matrix._builder._density.normalizationConstant(
) / np.trapz(
np.trapz(work_matrix.trace(), field_y_coordinates),
field_x_coordinates)
info.setTotalSpatialModeIntensity(total_spatial_mode_intensity)
log("Total spatial mode intensity: %e" %
total_spatial_mode_intensity.real)
if configuration.twoStepDivergenceMethod() == "":
divergence_method = "accurate"
else:
divergence_method = configuration.twoStepDivergenceMethod()
divergence_action = DivergenceAction(
x_coordinates=field_x_coordinates,
y_coordinates=field_y_coordinates,
intensity=work_matrix.trace(),
eigenvalues_spatial=eigenvalues_spatial,
eigenvectors_parallel=eigenvectors_parallel,
phase_space_density=PhaseSpaceDensity(
sigma_matrix,
wavefront.wavenumbers()[0]),
method=divergence_method)
twoform = divergence_action.apply(
number_modes=configuration.numberModes())
elif isinstance(work_matrix, Twoform):
twoform = work_matrix
else:
eigenmoder = Eigenmoder(field_x_coordinates, field_y_coordinates)
twoform = eigenmoder.eigenmodes(work_matrix,
configuration.numberModes())
total_beam_energies = e_0 + beam_energies
info.setEndTime()
autocorrelation_function = AutocorrelationFunction(
sigma_matrix=sigma_matrix,
undulator=undulator,
detuning_parameter=configuration.detuningParameter(),
energy=energy,
electron_beam_energy=electron_beam.energy(),
wavefront=first_wavefront,
exit_slit_wavefront=exit_slit_wavefront,
srw_wavefront_rx=Rx,
srw_wavefront_drx=dRx,
srw_wavefront_ry=Ry,
srw_wavefront_dry=dRy,
sampling_factor=configuration.samplingFactor(),
minimal_size=configuration.exitSlitWavefrontMinimalSizeVertical(),
beam_energies=total_beam_energies,
weighted_fields=weighted_fields,
static_electron_density=static_electron_density,
twoform=twoform,
info=info)
return autocorrelation_function
import numpy
# from srwlib import *
import sys
import comsyl
print(">>>>>>>", comsyl.__file__)
from comsyl.autocorrelation.AutocorrelationFunction import AutocorrelationFunction
from comsyl.autocorrelation.AutocorrelationFunctionPropagator import AutocorrelationFunctionPropagator
from comsyl.parallel.utils import isMaster, barrier
from comsyl.utils.Logger import log
from comsyl.waveoptics.Wavefront import NumpyWavefront, SRWWavefront
if __name__ == "__main__":
filename = "/users/srio/OASYS1.1e/paper-hierarchical-resources/comsyl/propagation_wofry_EBS/propagated_beamline.npz"
filename_out = "/users/srio/OASYS1.1e/paper-hierarchical-resources/comsyl/propagation_wofry_EBS/rediagonalized.npz"
af_name = filename.split("/")[-1].replace(".npz", "")
autocorrelation_function = AutocorrelationFunction.load(filename)
print(autocorrelation_function.eigenvalues())
autocorrelation_function.diagonalizeModes()
print(autocorrelation_function.eigenvalues())
autocorrelation_function.save(filename_out)
print("File written to disk: %s"%filename)
plt.show()
if __name__ == "__main__":
from comsyl.autocorrelation.AutocorrelationFunction import AutocorrelationFunction
from srxraylib.plot.gol import plot_image
import numpy
filename = "rediagonalized.npz"
af_name = filename.split("/")[-1].replace(".npz", "")
af = AutocorrelationFunction.load(filename)
print(af.eigenvalues())
x = af.xCoordinates()
y = af.yCoordinates()
for mode in [0,1]:
sd0 = numpy.abs(af.coherentMode(mode))**2
plot1(sd0,1e9*x,1e9*y,xrange=[-50, 50],yrange=[-50, 50],xtitle="H [nm]",ytitle="V [nm]",filename="/tmp/final_mode%d.png"%mode)
def fromDictionary(data_dict):
sigma_matrix = SigmaMatrix.fromNumpyArray(data_dict["sigma_matrix"])
undulator = undulator_from_numpy_array(data_dict["undulator"])
detuning_parameter = data_dict["detuning_parameter"][0]
energy = data_dict["energy"][0]
electron_beam_energy = data_dict["electron_beam_energy"][0]
np_wavefront_0=data_dict["wavefront_0"]
np_wavefront_1=data_dict["wavefront_1"]
np_wavefront_2=data_dict["wavefront_2"]
wavefront = NumpyWavefront.fromNumpyArray(np_wavefront_0, np_wavefront_1, np_wavefront_2)
try:
np_exit_slit_wavefront_0=data_dict["exit_slit_wavefront_0"]
np_exit_slit_wavefront_1=data_dict["exit_slit_wavefront_1"]
np_exit_slit_wavefront_2=data_dict["exit_slit_wavefront_2"]
exit_slit_wavefront = NumpyWavefront.fromNumpyArray(np_exit_slit_wavefront_0, np_exit_slit_wavefront_1, np_exit_slit_wavefront_2)
except:
exit_slit_wavefront = wavefront.clone()
try:
weighted_fields = data_dict["weighted_fields"]
except:
weighted_fields = None
srw_wavefront_rx=data_dict["srw_wavefront_rx"][0]
srw_wavefront_ry=data_dict["srw_wavefront_ry"][0]
srw_wavefront_drx = data_dict["srw_wavefront_drx"][0]
srw_wavefront_dry = data_dict["srw_wavefront_dry"][0]
info_string = str(data_dict["info"])
info = AutocorrelationInfo.fromString(info_string)
sampling_factor=data_dict["sampling_factor"][0]
minimal_size=data_dict["minimal_size"][0]
beam_energies = data_dict["beam_energies"]
static_electron_density = data_dict["static_electron_density"]
coordinates_x = data_dict["twoform_0"]
coordinates_y = data_dict["twoform_1"]
diagonal_elements = data_dict["twoform_2"]
eigenvalues = data_dict["twoform_3"]
# do not read the big array with modes
twoform_vectors = None # data_dict["twoform_4"]
twoform = Twoform(coordinates_x, coordinates_y, diagonal_elements, eigenvalues, twoform_vectors)
eigenvector_errors = data_dict["twoform_5"]
twoform.setEigenvectorErrors(eigenvector_errors)
af = AutocorrelationFunction(sigma_matrix, undulator, detuning_parameter,energy,electron_beam_energy,
wavefront,exit_slit_wavefront,srw_wavefront_rx, srw_wavefront_drx, srw_wavefront_ry, srw_wavefront_dry,
sampling_factor,minimal_size, beam_energies, weighted_fields,
static_electron_density, twoform,
info)
af._x_coordinates = coordinates_x
af._y_coordinates = coordinates_y
af._intensity = diagonal_elements.reshape(len(coordinates_x), len(coordinates_y))
return af