def propagate_af(self, autocorrelation_function,
directory_name="propagation_srw",
af_output_file_root=None,
maximum_mode=None,
python_to_be_used="/users/srio/OASYS1.1/miniconda3/bin/python"):
propagator = AutocorrelationFunctionPropagator(self._srw_beamline)
if maximum_mode is None:
mode_distribution=autocorrelation_function.modeDistribution()
maximum_mode = mode_distribution[abs(mode_distribution)>0.00005].shape[0]
propagator.setMaximumMode(maximum_mode)
data_directory = "%s/%s" % (directory_name, "wavefronts")
if isMaster():
if not os.path.exists(directory_name):
os.mkdir(directory_name)
if not os.path.exists(data_directory):
os.mkdir(data_directory)
barrier()
propagated_filename = "%s/%s.npz" % (data_directory, "wavefront")
af = propagator.propagate(autocorrelation_function, propagated_filename,method="SRW",
python_to_be_used=python_to_be_used)
barrier()
if isMaster():
if af_output_file_root is not None:
af.save("%s.npz" % (af_output_file_root))
return af
def propagate_af(self, autocorrelation_function,
directory_name="propagation_wofry",
af_output_file_root=None,
maximum_mode=None, # this is indeed the NUMBER OF MODES (maximum index plus one)
python_to_be_used="/users/srio/OASYS1.1/miniconda3/bin/python"):
propagator = AutocorrelationFunctionPropagator(self)
if maximum_mode is None:
mode_distribution = autocorrelation_function.modeDistribution()
maximum_mode = mode_distribution[abs(mode_distribution) > 0.00005].shape[0]
propagator.setMaximumMode(maximum_mode)
data_directory = "%s/%s" % (directory_name, "wavefronts")
if isMaster():
if not os.path.exists(directory_name):
os.mkdir(directory_name)
if not os.path.exists(data_directory):
os.mkdir(data_directory)
barrier()
# propagated_filename = "%s/%s.npz" % (data_directory, af_name)
propagated_filename = "%s/%s.npz" % (data_directory, "wavefront")
af = propagator.propagate(autocorrelation_function, propagated_filename, method="WOFRY",
python_to_be_used=python_to_be_used)
barrier()
if isMaster():
if af_output_file_root is not None:
af.save("%s.npz" % (af_output_file_root))
logAll("File written to disk: %s.npz" % (af_output_file_root))
return af
def _estimateMemoryConsumption(self, wavefront):
configuration = self.configuration()
size_x = wavefront.absolute_x_coordinates().size
size_y = wavefront.absolute_y_coordinates().size
size_wavefront = size_x * size_y
size_matrix = size_wavefront**2
size_matrix_per_core = size_matrix / mpi.COMM_WORLD.size
size_modes = size_wavefront * configuration.numberModes()
size_solver_worst = size_wavefront * configuration.numberModes() * 12
size_solver_worst_per_core = size_solver_worst / mpi.COMM_WORLD.size
if isMaster():
log("Using sampling factor: %.02f" %
configuration.samplingFactor())
log("Using %i x %i points per plane" % (size_x, size_y))
log("A single wavefront needs %.2f mb" %
(size_wavefront * 16 / 1024.**2))
log("The total matrix needs %.2f gb" %
(size_matrix * 16 / 1024.**3))
log("Per core the matrix needs %.2f gb" %
(size_matrix_per_core * 16 / 1024.**3))
log("The modes will take %.2f mb" % (size_modes * 16 / 1024.**2))
log("The solver needs in the worst case about %.2f gb" %
(size_solver_worst * 16 / 1024.**3))
log("Per core the solver needs in the worst case about %.2f gb" %
(size_solver_worst_per_core * 16 / 1024.**3))
return size_matrix
def save(self, filename):
file_array_shape = self._eigenvectors.shape
if isMaster():
fp = np.memmap(filename+".npy", dtype=np.complex128, mode='w+', shape=file_array_shape)
for i_vector in range(self.numberVectors()):
logProgress(self.numberVectors(), i_vector, "Writing vectors")
vector = self.vector(i_vector)
if isMaster():
fp[i_vector, :, :] = vector
log("Flushing")
if isMaster():
del fp
log("done")
def saveh5(self, filename, af, maximum_number_of_modes=None):
if has_h5py == False:
raise ImportError("h5py not available")
if maximum_number_of_modes is None:
maximum_number_of_modes = af.Twoform().numberVectors() # af.numberModes()
if maximum_number_of_modes > af.numberModes():
maximum_number_of_modes = af.numberModes()
file_array_shape = (maximum_number_of_modes,
len(af.Twoform().xCoordinates()),
len(af.Twoform().yCoordinates()))
if isMaster():
f = h5py.File(filename, 'w')
bigdataset = f.create_dataset("twoform_4",file_array_shape, dtype=np.complex)
for i_vector in range(maximum_number_of_modes):
logProgress(af.Twoform().numberVectors(), i_vector, "Writing vectors")
vector = af.Twoform().vector(i_vector)
if isMaster():
bigdataset[i_vector,:,:] = vector
if isMaster():
sys.stdout.flush()
data_dict = af.asDictionary()
for key in data_dict.keys():
if (key =="twoform_4"): # alreary done
pass
elif (key == "twoform_3"):
if (data_dict[key] is not None):
f[key] = (data_dict[key])[0:maximum_number_of_modes]
else:
if (data_dict[key] is not None):
f[key] = data_dict[key]
f.close()
def propagate_af(
self,
autocorrelation_function,
directory_name="propagation_wofry",
af_output_file_root=None,
maximum_mode=None, # this is indeed the NUMBER OF MODES (maximum index plus one)
python_to_be_used="/users/srio/OASYS1.1/miniconda3/bin/python"):
propagator = AutocorrelationFunctionPropagator(self)
if maximum_mode is None:
mode_distribution = autocorrelation_function.modeDistribution()
maximum_mode = mode_distribution[
abs(mode_distribution) > 0.00005].shape[0]
propagator.setMaximumMode(maximum_mode)
data_directory = "%s/%s" % (directory_name, "wavefronts")
if isMaster():
if not os.path.exists(directory_name):
os.mkdir(directory_name)
if not os.path.exists(data_directory):
os.mkdir(data_directory)
barrier()
# propagated_filename = "%s/%s.npz" % (data_directory, af_name)
propagated_filename = "%s/%s.npz" % (data_directory, "wavefront")
af = propagator.propagate(autocorrelation_function,
propagated_filename,
method="WOFRY",
python_to_be_used=python_to_be_used)
barrier()
if isMaster():
if af_output_file_root is not None:
af.save("%s.npz" % (af_output_file_root))
logAll("File written to disk: %s.npz" % (af_output_file_root))
return af
def propagateWavefront(srw_beamline,
wavefront,
rx,
drx,
ry,
dry,
rescale_x,
rescale_y,
i_mode,
python_to_be_used="python"):
if isMaster():
if not os.path.exists("tmp"):
os.mkdir("tmp")
s_id = str(mpi.COMM_WORLD.Get_rank()) + "_" + gethostname()
wavefront.save("./tmp/tmp%s_in" % s_id)
parameter_lines = "rx=%f\ndrx=%f\nry=%f\ndry=%f\nrescale_x=%f\nrescale_y=%f\ns_id=\"%s\"" % (
rx, drx, ry, dry, rescale_x, rescale_y, s_id)
pickle.dump(srw_beamline, open("./tmp/tmp%s_beamline.p" % s_id, "wb"))
lines = """
import pickle
from srwlib import *
from comsyl.waveoptics.SRWAdapter import SRWAdapter
from comsyl.waveoptics.Wavefront import NumpyWavefront, SRWWavefront
wavefront = NumpyWavefront.load("./tmp/tmp%s_in.npz"%s_id)
adapter = SRWAdapter()
wfr = adapter.SRWWavefrontFromWavefront(wavefront,
rx,
drx,
ry,
dry,rescale_x,rescale_y)
#print("Doing propagation in external call")
srw_beamline = pickle.load(open("./tmp/tmp%s_beamline.p"%s_id,"rb"))
srwl.PropagElecField(wfr, srw_beamline)
tmp = SRWWavefront(wfr).toNumpyWavefront()
wfr = None
tmp.save("./tmp/tmp%s_out" % s_id)
"""
file = open("./tmp/tmp%s.py" % s_id, "w")
file.writelines(parameter_lines)
file.writelines("\ni_mode=%d\n" % i_mode) # added srio
file.writelines(lines)
file.close()
os.system(python_to_be_used + " ./tmp/tmp%s.py" % s_id)
return NumpyWavefront.load("./tmp/tmp%s_out.npz" % s_id)
def propagate_af(
self,
autocorrelation_function,
directory_name="propagation_srw",
af_output_file_root=None,
maximum_mode=None,
python_to_be_used="/users/srio/OASYS1.1/miniconda3/bin/python"):
propagator = AutocorrelationFunctionPropagator(self._srw_beamline)
if maximum_mode is None:
mode_distribution = autocorrelation_function.modeDistribution()
maximum_mode = mode_distribution[
abs(mode_distribution) > 0.00005].shape[0]
propagator.setMaximumMode(maximum_mode)
data_directory = "%s/%s" % (directory_name, "wavefronts")
if isMaster():
if not os.path.exists(directory_name):
os.mkdir(directory_name)
if not os.path.exists(data_directory):
os.mkdir(data_directory)
barrier()
propagated_filename = "%s/%s.npz" % (data_directory, "wavefront")
af = propagator.propagate(autocorrelation_function,
propagated_filename,
method="SRW",
python_to_be_used=python_to_be_used)
barrier()
if isMaster():
if af_output_file_root is not None:
af.save("%s.npz" % (af_output_file_root))
return af
def save(self, filename, af):
af.Twoform().saveVectors(filename.replace(".npz", ""))
if isMaster():
sys.stdout.flush()
data_dict = af.asDictionary()
save_dict = dict()
for key in data_dict.keys():
if key !="twoform_4":
save_dict["np_"+key] = data_dict[key]
else:
eigenvectors = data_dict["twoform_4"]
filename_npz = filename.replace(".npz", "")
np.savez_compressed(filename_npz, **save_dict)
def propagateWavefront(srw_beamline, wavefront, rx, drx, ry, dry, rescale_x, rescale_y, i_mode, python_to_be_used="python"):
if isMaster():
if not os.path.exists("tmp"):
os.mkdir("tmp")
s_id=str(mpi.COMM_WORLD.Get_rank())+"_"+gethostname()
wavefront.save("./tmp/tmp%s_in"%s_id)
parameter_lines = "rx=%f\ndrx=%f\nry=%f\ndry=%f\nrescale_x=%f\nrescale_y=%f\ns_id=\"%s\"" % (rx, drx, ry, dry, rescale_x, rescale_y, s_id)
pickle.dump(srw_beamline, open("./tmp/tmp%s_beamline.p"%s_id,"wb"))
lines ="""
import pickle
from srwlib import *
from comsyl.waveoptics.SRWAdapter import SRWAdapter
from comsyl.waveoptics.Wavefront import NumpyWavefront, SRWWavefront
wavefront = NumpyWavefront.load("./tmp/tmp%s_in.npz"%s_id)
adapter = SRWAdapter()
wfr = adapter.SRWWavefrontFromWavefront(wavefront,
rx,
drx,
ry,
dry,rescale_x,rescale_y)
#print("Doing propagation in external call")
srw_beamline = pickle.load(open("./tmp/tmp%s_beamline.p"%s_id,"rb"))
srwl.PropagElecField(wfr, srw_beamline)
tmp = SRWWavefront(wfr).toNumpyWavefront()
wfr = None
tmp.save("./tmp/tmp%s_out" % s_id)
"""
file = open("./tmp/tmp%s.py"%s_id, "w")
file.writelines(parameter_lines)
file.writelines("\ni_mode=%d\n"%i_mode) # added srio
file.writelines(lines)
file.close()
os.system(python_to_be_used+" ./tmp/tmp%s.py"%s_id)
return NumpyWavefront.load("./tmp/tmp%s_out.npz"%s_id)
def _estimateMemoryConsumption(self, wavefront):
configuration = self.configuration()
size_x = wavefront.absolute_x_coordinates().size
size_y = wavefront.absolute_y_coordinates().size
size_wavefront = size_x * size_y
size_matrix = size_wavefront ** 2
size_matrix_per_core = size_matrix / mpi.COMM_WORLD.size
size_modes = size_wavefront * configuration.numberModes()
size_solver_worst = size_wavefront * configuration.numberModes() * 12
size_solver_worst_per_core = size_solver_worst / mpi.COMM_WORLD.size
if isMaster():
log("Using sampling factor: %.02f" % configuration.samplingFactor())
log("Using %i x %i points per plane" % (size_x, size_y))
log("A single wavefront needs %.2f mb" % (size_wavefront * 16 / 1024.**2))
log("The total matrix needs %.2f gb" % (size_matrix * 16 / 1024.**3))
log("Per core the matrix needs %.2f gb" % (size_matrix_per_core * 16 / 1024.**3))
log("The modes will take %.2f mb" % (size_modes * 16 / 1024.**2))
log("The solver needs in the worst case about %.2f gb" % (size_solver_worst * 16 / 1024.**3))
log("Per core the solver needs in the worst case about %.2f gb" % (size_solver_worst_per_core * 16 / 1024.**3))
return size_matrix
def propagate(self, autocorrelation_function, filename, method='SRW', python_to_be_used="python"):
source_filename = autocorrelation_function._io.fromFile()
try:
source_uid = autocorrelation_function.info().uid()
except:
source_uid = "None"
autocorrelation_function.info().logStart()
logAll("Propagating %s (%s)" % (source_filename, source_uid))
if self._maximum_mode is None:
number_modes = autocorrelation_function.numberModes()
else:
number_modes = self._maximum_mode
if isMaster():
if not os.path.exists("tmp"):
os.mkdir("tmp")
distribution_plan = DistributionPlan(mpi.COMM_WORLD, n_rows=number_modes, n_columns=1)
n_rank = mpi.COMM_WORLD.Get_rank()
x_coordinates = []
y_coordinates = []
for i_mode in distribution_plan.localRows():
for i in range(1):
logAll("%i doing mode index: %i/%i (max mode index: %i)" % (n_rank, i_mode, max(distribution_plan.localRows()), number_modes-1))
if n_rank == 0:
sys.stdout.flush()
wavefront = autocorrelation_function.coherentModeAsWavefront(i_mode)
#wavefront._e_field[np.abs(wavefront._e_field)<0.000001]=0.0
if method == 'SRW':
# CHANGE THIS FOR WOFRY
srw_wavefront = propagateWavefront(self.__srw_beamline,
wavefront,
autocorrelation_function.SRWWavefrontRx(),
autocorrelation_function.SRWWavefrontDRx(),
autocorrelation_function.SRWWavefrontRy(),
autocorrelation_function.SRWWavefrontDRy(), 1.0, 1.0, i_mode,
python_to_be_used=python_to_be_used)
elif method == 'WOFRY':
srw_wavefront = propagateWavefrontWofry(self.__srw_beamline,wavefront,i_mode,python_to_be_used=python_to_be_used)
else:
raise Exception("Method not known: %s"%method)
# norm_mode = trapez2D( np.abs(srw_wavefront.E_field_as_numpy()[0,:,:,0])**2, 1, 1)**0.5
# if norm_mode > 1e2 or np.isnan(norm_mode):
# print("TRY %i AFTER PROPAGATION:" % i, i_mode,norm_mode)
# sys.stdout.flush()
#else:
# break
#if i==19:
# exit()
# if np.any(norm_srw_wavefront > 10):
# exit()
#
# if np.any(norm_wavefront > 10):
# exit()
adjusted_wavefront = self._adjustWavefrontSize(srw_wavefront)
# norm_mode = trapez2D( np.abs(adjusted_wavefront.E_field_as_numpy()[0,:,:,0])**2, 1, 1)**0.5
# if norm_mode > 1e2 or np.isnan(norm_mode):
# print("TRY %i AFTER ADJUSTMENT:" % i, i_mode,norm_mode)
# sys.stdout.flush()
# exit()
# writes a file for every wavefront
TwoformVectorsWavefronts.pushWavefront(filename, adjusted_wavefront, index=i_mode)
#print("Saving wavefront %i" % i_mode)
x_coordinates.append(adjusted_wavefront.absolute_x_coordinates().copy())
y_coordinates.append(adjusted_wavefront.absolute_y_coordinates().copy())
mpi.COMM_WORLD.barrier()
# replace the wavefronts bu the propagated ones
af = self._saveAutocorrelation(autocorrelation_function, number_modes, x_coordinates, y_coordinates, filename)
# convert from one file per wavefront to one big array
af.Twoform().convertToTwoformVectorsEigenvectors()
af.info().setEndTime()
filelist = glob.glob(filename+"*")
for f in filelist:
os.remove(f)
return af
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 mpi4py.MPI as mpi
from comsyl.utils.Logger import logAll, log
from comsyl.parallel.utils import isMaster, barrier
from socket import gethostname
from comsyl.parallel.DistributionPlan import DistributionPlan
import os
if isMaster():
print("Hello from master")
if isMaster():
if not os.path.exists("tmp"):
os.mkdir("tmp")
logAll("Using LogAll")
log("Using Log")
s_id = str(mpi.COMM_WORLD.Get_rank()) + "_" + gethostname()
print("s_id: ", s_id)
print("str(mpi.COMM_WORLD.Get_rank()): ", str(mpi.COMM_WORLD.Get_rank()))
number_modes = 1000
distribution_plan = DistributionPlan(mpi.COMM_WORLD,
n_rows=number_modes,
n_columns=1)
print(distribution_plan)
f = open("./tmp/TMP%s_in" % s_id, 'w')
f.write(">>>>>>>>>")
f.close
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
def log(self, log_string):
if isMaster():
self.logAll(log_string)
def propagate(self,
autocorrelation_function,
filename,
method='SRW',
python_to_be_used="python"):
source_filename = autocorrelation_function._io.fromFile()
try:
source_uid = autocorrelation_function.info().uid()
except:
source_uid = "None"
autocorrelation_function.info().logStart()
logAll("Propagating %s (%s)" % (source_filename, source_uid))
if self._maximum_mode is None:
number_modes = autocorrelation_function.numberModes()
else:
number_modes = self._maximum_mode
if isMaster():
if not os.path.exists("tmp"):
os.mkdir("tmp")
distribution_plan = DistributionPlan(mpi.COMM_WORLD,
n_rows=number_modes,
n_columns=1)
n_rank = mpi.COMM_WORLD.Get_rank()
x_coordinates = []
y_coordinates = []
for i_mode in distribution_plan.localRows():
for i in range(1):
logAll("%i doing mode index: %i/%i (max mode index: %i)" %
(n_rank, i_mode, max(
distribution_plan.localRows()), number_modes - 1))
if n_rank == 0:
sys.stdout.flush()
wavefront = autocorrelation_function.coherentModeAsWavefront(
i_mode)
#wavefront._e_field[np.abs(wavefront._e_field)<0.000001]=0.0
if method == 'SRW':
# CHANGE THIS FOR WOFRY
srw_wavefront = propagateWavefront(
self.__srw_beamline,
wavefront,
autocorrelation_function.SRWWavefrontRx(),
autocorrelation_function.SRWWavefrontDRx(),
autocorrelation_function.SRWWavefrontRy(),
autocorrelation_function.SRWWavefrontDRy(),
1.0,
1.0,
i_mode,
python_to_be_used=python_to_be_used)
elif method == 'WOFRY':
srw_wavefront = propagateWavefrontWofry(
self.__srw_beamline,
wavefront,
i_mode,
python_to_be_used=python_to_be_used)
else:
raise Exception("Method not known: %s" % method)
# norm_mode = trapez2D( np.abs(srw_wavefront.E_field_as_numpy()[0,:,:,0])**2, 1, 1)**0.5
# if norm_mode > 1e2 or np.isnan(norm_mode):
# print("TRY %i AFTER PROPAGATION:" % i, i_mode,norm_mode)
# sys.stdout.flush()
#else:
# break
#if i==19:
# exit()
# if np.any(norm_srw_wavefront > 10):
# exit()
#
# if np.any(norm_wavefront > 10):
# exit()
adjusted_wavefront = self._adjustWavefrontSize(srw_wavefront)
# norm_mode = trapez2D( np.abs(adjusted_wavefront.E_field_as_numpy()[0,:,:,0])**2, 1, 1)**0.5
# if norm_mode > 1e2 or np.isnan(norm_mode):
# print("TRY %i AFTER ADJUSTMENT:" % i, i_mode,norm_mode)
# sys.stdout.flush()
# exit()
# writes a file for every wavefront
TwoformVectorsWavefronts.pushWavefront(filename,
adjusted_wavefront,
index=i_mode)
#print("Saving wavefront %i" % i_mode)
x_coordinates.append(
adjusted_wavefront.absolute_x_coordinates().copy())
y_coordinates.append(
adjusted_wavefront.absolute_y_coordinates().copy())
mpi.COMM_WORLD.barrier()
# replace the wavefronts bu the propagated ones
af = self._saveAutocorrelation(autocorrelation_function, number_modes,
x_coordinates, y_coordinates, filename)
# convert from one file per wavefront to one big array
af.Twoform().convertToTwoformVectorsEigenvectors()
af.info().setEndTime()
filelist = glob.glob(filename + "*")
for f in filelist:
os.remove(f)
return af