def __init__(self,
N_e,
sigma_matrix,
weighted_fields,
x_coordinates,
y_coordinates,
k,
strategy=None):
self._N_e = N_e
self._sigma_matrix = sigma_matrix
self._density = PhaseSpaceDensity(sigma_matrix, k)
self._weighted_fields = weighted_fields.copy()
self._field_x_coordinates = x_coordinates.copy()
self._field_y_coordinates = y_coordinates.copy()
self._x_coordinates = x_coordinates.copy(
) # self._minkowskiSum(x_coordinates)
self._y_coordinates = y_coordinates.copy(
) # self._minkowskiSum(y_coordinates)
if strategy is None:
if self._density.isAlphaZero():
strategy = BuilderStrategyConvolution
else:
log("Found alpha not equal to zero. Can not use convolutions.")
strategy = BuilderStrategyPython
self.setStrategy(strategy)
self.setAllCoordinates(self._field_x_coordinates,
self._field_y_coordinates)
self.setDoNotUseConvolutions(False)
def fromJson(json_filename):
json_content = json.load(open(json_filename, "r"))
configuration = AutocorrelationSimulatorConfiguration()
def process_json_content(json_content, conf, root_key):
for key, value in json_content.items():
if isinstance(value, dict):
process_json_content(value, conf, root_key+key+"/")
else:
conf.setByName(root_key + key, value)
process_json_content(json_content, configuration, "")
if configuration.isSet("patch/source"):
patch_configuration = configuration
template_filename = configuration.byName("patch/source")
log("Patching configuration: %s" % template_filename)
initial_configuration = AutocorrelationSimulatorConfiguration.fromJson(template_filename)
for patch_setting, value in patch_configuration._setting.items():
setting_name = patch_setting.replace("patch/", "")
if setting_name == "source":
continue
log("Patching: %s = %s" %(setting_name, str(value)) )
initial_configuration.setByName(setting_name, value)
return initial_configuration
else:
return configuration
def __init__(self, x_coordinates, y_coordinates, density, field_x_coordinates, field_y_coordinates, weighted_fields):
self._x_coordinates = x_coordinates
self._y_coordinates = y_coordinates
self._density = density
self._weigthed_fields = weighted_fields
self._setActiveField(0)
self._field_tmp = np.zeros_like(self._field)
self._field_tmp_x = np.zeros(self._field.shape[0], dtype=np.complex128)
self._field_tmp_y = np.zeros(self._field.shape[1], dtype=np.complex128)
self._field_x_coordinates = field_x_coordinates
self._field_y_coordinates = field_y_coordinates
self._last_r_1_x = None
self._last_r_1_y = None
self._x_coordinates_i_min = -(self._x_coordinates.shape[0] + 1)/2
self._x_coordinates_i_max = (self._x_coordinates.shape[0] + 1)/2
self._x_coordinates_step_width = self._x_coordinates[1] - self._x_coordinates[0]
self._y_coordinates_i_min = -(self._y_coordinates.shape[0] + 1)/2
self._y_coordinates_i_max = (self._y_coordinates.shape[0] + 1)/2
self._y_coordinates_step_width = self._y_coordinates[1] - self._y_coordinates[0]
self._grid_area = self._x_coordinates_step_width * self._y_coordinates_step_width
log("Calculating density (integration part)")
self._setRhoR(self._field_x_coordinates, self._field_y_coordinates)
self._rho_phase_tmp = np.zeros_like(self._rho)
self._convolution = Convolution()
self._postInit()
def _runSingleSimulation(self, configuration):
resetLog()
self._setConfiguration(configuration)
filename = self._determineFilename()
lattice_name = configuration.latticeName()
electron_beam = latticeByName(lattice_name)
undulator = undulatorByName(configuration.undulatorName())
info = AutocorrelationInfo()
info.logStart()
info.setConfiguration(configuration)
info.setTag(configuration.tag())
af = self.calculateAutocorrelation(electron_beam=electron_beam,
undulator=undulator,
info=info)
af.save(filename)
log("File created: %s" % filename + ".npy")
log("File created: %s" % filename + ".npz")
try:
pass
# import h5py
# af.saveh5(filename+".h5")
# log("File created: %s" % filename+".h5")
except:
pass
barrier()
return filename
def fromJson(json_filename):
json_content = json.load(open(json_filename, "r"))
configuration = AutocorrelationSimulatorConfiguration()
def process_json_content(json_content, conf, root_key):
for key, value in json_content.items():
if isinstance(value, dict):
process_json_content(value, conf, root_key + key + "/")
else:
conf.setByName(root_key + key, value)
process_json_content(json_content, configuration, "")
if configuration.isSet("patch/source"):
patch_configuration = configuration
template_filename = configuration.byName("patch/source")
log("Patching configuration: %s" % template_filename)
initial_configuration = AutocorrelationSimulatorConfiguration.fromJson(
template_filename)
for patch_setting, value in patch_configuration._setting.items():
setting_name = patch_setting.replace("patch/", "")
if setting_name == "source":
continue
log("Patching: %s = %s" % (setting_name, str(value)))
initial_configuration.setByName(setting_name, value)
return initial_configuration
else:
return configuration
def _runSingleSimulation(self, configuration):
resetLog()
self._setConfiguration(configuration)
filename = self._determineFilename()
lattice_name = configuration.latticeName()
electron_beam = latticeByName(lattice_name)
undulator = undulatorByName(configuration.undulatorName())
info = AutocorrelationInfo()
info.logStart()
info.setConfiguration(configuration)
info.setTag(configuration.tag())
af = self.calculateAutocorrelation(electron_beam=electron_beam,
undulator=undulator,
info=info)
af.save(filename)
log("File created: %s" % filename+".npy")
log("File created: %s" % filename+".npz")
try:
pass
# import h5py
# af.saveh5(filename+".h5")
# log("File created: %s" % filename+".h5")
except:
pass
barrier()
return filename
def setDoNotUseConvolutions(self, do_not_use_convolutions):
self._do_not_use_convolutions = do_not_use_convolutions
if self._do_not_use_convolutions:
log("Never use convolutions")
else:
log("Using convolutions if possible")
def setDoNotUseConvolutions(self, do_not_use_convolutions):
self._do_not_use_convolutions = do_not_use_convolutions
if self._do_not_use_convolutions:
log("Never use convolutions")
else:
log("Using convolutions if possible")
def __init__(self, N_e, sigma_matrix, weighted_fields, x_coordinates, y_coordinates, k, number_modes):
log("Setting up autocorrelation operator")
self._action = 0
self._builder = AutocorrelationBuilder(N_e, sigma_matrix, weighted_fields, x_coordinates, y_coordinates, k, strategy=BuilderStrategyPython)
self._distribution_plan = DistributionPlan(communicator=mpi.COMM_WORLD, n_columns=self.dimensionSize(), n_rows=self.dimensionSize())
log("Setting up PETSc interface")
self._petSc_operator = PetScOperator(self)
self._number_modes = number_modes
mpi.COMM_WORLD.barrier()
def evaluateAllR_2_Fredholm_parallel_compare(self, v_in, v_out):
v_in_clone = v_in.clone()
self.evaluateAllR_2_Fredholm_parallel_direct(v_in, v_out)
self.evaluateAllR_2_Fredholm_parallel_convolution(v_in_clone, v_in_clone)
#from comsyl.math.utils import plotSurface
# plotSurface(self._field_x_coordinates,self._field_y_coordinates, np.abs(res))
# plotSurface(self._field_x_coordinates,self._field_y_coordinates, np.abs(v2))
# plotSurface(self._field_x_coordinates,self._field_y_coordinates, np.abs(v2-res))
res = v_out.fullData()
v2 = v_in_clone.fullData()
log("abs error, rel error: %f, %f" % (np.linalg.norm(np.abs(res)-np.abs(v2)), np.linalg.norm(np.abs(res-v2)/np.linalg.norm(np.abs(res)))))
def _runConvergenceSeries(self, configuration):
sampling_factors = np.linspace(configuration.convergenceSeriesStart(),
configuration.samplingFactor(),
configuration.convergenceSeriesLength())
for i_factor, sampling_factor in enumerate(sampling_factors):
log("Doing sampling factor %e (%i/%i)" % (sampling_factor, i_factor+1, len(sampling_factors)))
simulation = AutocorrelationSimulator()
current_configuration = configuration.clone()
current_configuration.setDoConvergenceSeries(False)
current_configuration.setSamplingFactor(str(sampling_factor))
filename = simulation.run(current_configuration)
return filename
def _setUpPhases(self):
log("Calculating density (integration phases)")
delta_r = np.array([self._x_coordinates_step_width, self._y_coordinates_step_width])
frequency = self._density.expPhasePointsPerOscillation((self._field_x_coordinates.max(), self._field_y_coordinates.max()))
oscillations_per_step = frequency * delta_r
log("Have %.3f horizontal and %.3f vertical phase oscillations per step" % (oscillations_per_step[0], oscillations_per_step[1]))
if oscillations_per_step[0] > 0.25:
raise Exception("ABORT. Need AT LEAST 4 integration points per oscillation. Have %f" % (1/oscillations_per_step[0]))
if oscillations_per_step[1] > 0.25:
raise Exception("ABORT. Need AT LEAST 4 integration points per oscillation. Have %f" % (1/oscillations_per_step[1]))
self._setRhoPhase(self._field_x_coordinates, self._field_y_coordinates)
yy, xx = np.meshgrid(self._field_y_coordinates, self._field_x_coordinates)
self._rho_phase_exp_x = np.zeros((len(self._field_x_coordinates), xx.shape[0], xx.shape[1]), dtype=np.complex128)
self._rho_phase_exp_y = np.zeros((len(self._field_y_coordinates), yy.shape[0], yy.shape[1]), dtype=np.complex128)
for r_x in self._field_x_coordinates:
i_x = self.xIndexByCoordinate(r_x)
rho_phase = self._rho_phase_x[i_x]
self._rho_phase_exp_x[i_x, :, :] = np.exp(xx * rho_phase)
for r_y in self._field_y_coordinates:
i_y = self.xIndexByCoordinate(r_y)
rho_phase = self._rho_phase_y[i_y]
self._rho_phase_exp_y[i_y, :, :] = np.exp(yy * rho_phase)
self._coordinate_map_x = np.zeros(len(self._field_x_coordinates), dtype=np.int)
self._coordinate_map_y = np.zeros(len(self._field_y_coordinates), dtype=np.int)
self._coordinate_map_minus_x = np.zeros(len(self._field_x_coordinates), dtype=np.int)
self._coordinate_map_minus_y = np.zeros(len(self._field_y_coordinates), dtype=np.int)
for i_r_x, r_x in enumerate(self._field_x_coordinates):
self._coordinate_map_x[i_r_x] = self.xIndexByCoordinate(-r_x)
self._coordinate_map_minus_x[i_r_x] = self.xIndexByCoordinate(r_x)
for i_r_y, r_y in enumerate(self._field_y_coordinates):
self._coordinate_map_y[i_r_y] = self.yIndexByCoordinate(-r_y)
self._coordinate_map_minus_y[i_r_y] = self.yIndexByCoordinate(r_y)
self._has_phases = True
def _runConvergenceSeries(self, configuration):
sampling_factors = np.linspace(configuration.convergenceSeriesStart(),
configuration.samplingFactor(),
configuration.convergenceSeriesLength())
for i_factor, sampling_factor in enumerate(sampling_factors):
log("Doing sampling factor %e (%i/%i)" %
(sampling_factor, i_factor + 1, len(sampling_factors)))
simulation = AutocorrelationSimulator()
current_configuration = configuration.clone()
current_configuration.setDoConvergenceSeries(False)
current_configuration.setSamplingFactor(str(sampling_factor))
filename = simulation.run(current_configuration)
return filename
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 evaluateAllR_2_Fredholm_parallel_compare(self, v_in, v_out):
v_in_clone = v_in.clone()
self.evaluateAllR_2_Fredholm_parallel_direct(v_in, v_out)
self.evaluateAllR_2_Fredholm_parallel_convolution(
v_in_clone, v_in_clone)
#from comsyl.math.utils import plotSurface
# plotSurface(self._field_x_coordinates,self._field_y_coordinates, np.abs(res))
# plotSurface(self._field_x_coordinates,self._field_y_coordinates, np.abs(v2))
# plotSurface(self._field_x_coordinates,self._field_y_coordinates, np.abs(v2-res))
res = v_out.fullData()
v2 = v_in_clone.fullData()
log("abs error, rel error: %f, %f" %
(np.linalg.norm(np.abs(res) - np.abs(v2)),
np.linalg.norm(np.abs(res - v2) / np.linalg.norm(np.abs(res)))))
def _createParallelMatrix(self, f_gamma):
log("Building matrix")
return self._createParallelMatrixPETSc(f_gamma)
product_coordinates=self.productCoordinates()
n_coordinates = product_coordinates.shape[0]
distribution_plan = DistributionPlan(communicator=mpi.COMM_WORLD,
n_rows=product_coordinates.shape[0],
n_columns=product_coordinates.shape[0])
matrix = ParallelMatrix(distribution_plan=distribution_plan)
if self._mode_element_wise:
for i_row in distribution_plan.localRows():
self._printProgress(n_coordinates, i_row)
r_i = product_coordinates[i_row, :]
for i_column in range(n_coordinates):
r_j = product_coordinates[i_column, :]
value = f_gamma(r_i, r_j)
# TODO
raise NotImplementedError("Can only handle entire rows")
# matrix.setElement(i_row, i_column, value)
else:
for i_row in distribution_plan.localRows():
self._printProgress(len(distribution_plan.localRows()), i_row)
r_i = product_coordinates[i_row, :]
value = f_gamma(r_i)
value = value.reshape(value.size)
matrix.setRow(global_index=i_row,
content=value)
if distribution_plan.communicator().Get_rank() == 0:
log("done")
return matrix
def build(self,
electron_beam,
xp,
yp,
z_offset,
x=0.0,
y=0.0,
initial_z=None):
adapter = SRWAdapter()
adapter.setSamplingFactor(self._sampling_factor)
max_theta_x = self._undulator.gaussian_central_cone_aperture(
electron_beam.gamma(), n=1) * 3.0
z = self._undulator.length() + z_offset
min_dimension_x_theta = 1.0 * self._min_dimension_x / z * np.sqrt(2.0)
max_dimension_x_theta = 1.0 * self._max_dimension_x / z * np.sqrt(2.0)
min_dimension_y_theta = 1.0 * self._min_dimension_y / z * np.sqrt(2.0)
max_dimension_y_theta = 1.0 * self._max_dimension_y / z * np.sqrt(2.0)
max_theta_x = self._applyLimits(max_theta_x, min_dimension_x_theta,
max_dimension_x_theta)
max_theta_y = self._applyLimits(max_theta_x / 1.5,
min_dimension_y_theta,
max_dimension_y_theta)
self._setAdapterInitialZ(adapter)
log("Using initial z_0 for initial conditions: %e" %
adapter._initial_z)
calc_wavefront = adapter.wavefrontRectangularForSingleEnergy(
electron_beam,
self._undulator,
z,
max_theta_x,
max_theta_y,
self._photon_energy,
x=x,
xp=xp,
y=y,
yp=yp)
return calc_wavefront
def evaluateCutY(self):
log("Calculating vertical cut")
i_x = self.xFieldIndexByCoordinate(0.0)
result = np.zeros((self._field.shape[1], self._field.shape[1]), dtype=np.complex128)
for r_y in self._field_y_coordinates:
i_y_1 = self.yFieldIndexByCoordinate(r_y)
r1 = np.array([0.0, r_y])
result[i_y_1, :] = self.evaluateAllR_2(r1)[i_x, :]
log("done")
result2 = np.zeros_like(result)
for y1 in self._field_y_coordinates:
for y2 in self._field_y_coordinates:
result2[self.yFieldIndexByCoordinate(y1), self.yFieldIndexByCoordinate(y2)] = result[self.yIndexByCoordinate(y1), self.yIndexByCoordinate(y2)]
return result2
def evaluateCutX(self):
log("Calculating horizontal cut")
i_y = self.yFieldIndexByCoordinate(0.0)
result = np.zeros((self._field.shape[0], self._field.shape[0]), dtype=np.complex128)
for r_x in self._field_x_coordinates:
i_x_1 = self.xFieldIndexByCoordinate(r_x)
r1 = np.array([r_x, 0.0])
result[i_x_1, :] = self.evaluateAllR_2(r1)[:, i_y]
log("done")
result2 = np.zeros_like(result)
for x1 in self._field_x_coordinates:
for x2 in self._field_x_coordinates:
result2[self.xFieldIndexByCoordinate(x1), self.xFieldIndexByCoordinate(x2)] = result[self.xIndexByCoordinate(x1), self.xIndexByCoordinate(x2)]
return result2
def createReferenceWavefrontAtVirtualSource(self, Rx, dRx, Ry, dRy,
configuration, source_position,
wavefront):
adapter = SRWAdapter()
if source_position == VIRTUAL_SOURCE_CENTER:
z = -1.0 * self._undulator.length()
elif source_position == VIRTUAL_SOURCE_ENTRANCE:
z = -1.5 * self._undulator.length(
) #- 2 * self._undulator.periodLength()
else:
raise NotImplementedError("Source position %s" % source_position)
log("Using source position: %s with z=%.02f" % (source_position, z))
wavefront = adapter.propagate(wavefront, Rx, dRx, Ry, dRy, z)
x_min = -configuration.sourceWavefrontMaximalSizeHorizontal()
x_max = -x_min
y_min = -configuration.sourceWavefrontMaximalSizeVertical()
y_max = -y_min
if x_min > wavefront.minimal_x_coodinate() or x_max < wavefront.maximal_x_coodinate() or \
y_min > wavefront.minimal_y_coodinate() or y_max < wavefront.maximal_y_coodinate():
dim_x = int((x_max - x_min) / wavefront.x_stepwidth())
dim_y = int((y_max - y_min) / wavefront.y_stepwidth())
divisor_x = configuration.samplingFactorDivisorHorizontal()
if divisor_x == "":
divisor_x = 1.0
divisor_y = configuration.samplingFactorDivisorVertical()
if divisor_y == "":
divisor_y = 1.0
wavefront = wavefront.onDomain(x_min, x_max,
int(dim_x / divisor_x), y_min,
y_max, int(dim_y / divisor_y))
#wavefront = wavefront.zeroPadded(zero_padding_factor_x=1.0, zero_padding_factor_y=1.0)
return wavefront
def fromWavefront(self, wavefront, info):
e_field = wavefront.E_field_as_numpy()[0, :, :, 0]
x_coordinates = wavefront.absolute_x_coordinates()
y_coordinates = wavefront.absolute_y_coordinates()
x_index = np.abs(x_coordinates).argmin()
y_index = np.abs(y_coordinates).argmin()
x_fwhm = getFWHM(x_coordinates, np.abs(e_field[:, y_index])**1)
y_fwhm = getFWHM(y_coordinates, np.abs(e_field[x_index, :])**1)
sigma_x = x_fwhm / (2 * np.sqrt(2 * np.log(2) ))
sigma_y = y_fwhm / (2 * np.sqrt(2 * np.log(2) ))
#sigma_x /= np.sqrt(2.0)
#sigma_y /= np.sqrt(2.0)
log(">>>")
log(">>>Creating Gaussian wavefront using sigma_x: %e and sigma_y: %e" % (sigma_x, sigma_y))
log(">>>")
info.set("gaussian_wavefront_sigma_x", str(sigma_x))
info.set("gaussian_wavefront_sigma_y", str(sigma_y))
new_e_field = np.zeros_like(e_field)
new_e_field[:, :] = createGaussian2D(sigma_x, sigma_y, x_coordinates, y_coordinates)
new_wavefront = wavefront.toNumpyWavefront().clone()
new_wavefront._e_field[0, :, :, 0] = new_e_field[:, :]
return new_wavefront
def buildOnGrid(
self,
reference_wavefront,
electron_beam,
z_offset,
xp,
yp,
x=0.0,
y=0.0,
):
adapter = SRWAdapter()
adapter.setSamplingFactor(self._sampling_factor)
z = self._undulator.length() + z_offset
grid_length_x = reference_wavefront.absolute_x_coordinates().max()
grid_length_y = reference_wavefront.absolute_y_coordinates().max()
energy = self._photon_energy
self._setAdapterInitialZ(adapter)
log("Using initial z_0 for initial conditions: %e" %
adapter._initial_z)
calc_wavefront = adapter.wavefrontByCoordinates(
electron_beam=electron_beam,
undulator=self._undulator,
z_start=z,
grid_length_x=grid_length_x,
grid_length_y=grid_length_y,
energy_number=1,
energy_start=energy,
energy_end=energy,
x=x,
xp=xp,
y=y,
yp=yp)
return calc_wavefront
def _createParallelMatrixPETSc(self, f_gamma):
product_coordinates=self.productCoordinates()
n_coordinates = product_coordinates.shape[0]
n_rows=product_coordinates.shape[0]
n_columns=product_coordinates.shape[0]
petsc_matrix = PETSc.Mat().createDense([n_rows, n_columns])
petsc_matrix.setUp()
matrix = ParallelMatrixPETSc(petsc_matrix)
distribution_plan = matrix.distributionPlan()
if self._mode_element_wise:
for i_row in distribution_plan.localRows():
self._printProgress(n_coordinates, i_row)
r_i = product_coordinates[i_row, :]
for i_column in range(n_coordinates):
r_j = product_coordinates[i_column, :]
value = f_gamma(r_i, r_j)
# TODO
#raise NotImplementedError("Can only handle entire rows")
petsc_matrix[i_row, i_column] = value
else:
for i_row in distribution_plan.localRows():
self._printProgress(len(distribution_plan.localRows()), i_row)
r_i = product_coordinates[i_row, :]
value = f_gamma(r_i)
value = value.reshape(value.size).conj()
matrix.setRow(global_index=i_row,
content=value)
log("Waiting for others")
barrier()
log("done")
log("PETSc matrix assembling")
matrix.assemble()
# matrix.transpose()
log("done")
return matrix
def evaluateCutX(self):
log("Calculating horizontal cut")
i_y = self.yFieldIndexByCoordinate(0.0)
result = np.zeros((self._field.shape[0], self._field.shape[0]),
dtype=np.complex128)
for r_x in self._field_x_coordinates:
i_x_1 = self.xFieldIndexByCoordinate(r_x)
r1 = np.array([r_x, 0.0])
result[i_x_1, :] = self.evaluateAllR_2(r1)[:, i_y]
log("done")
result2 = np.zeros_like(result)
for x1 in self._field_x_coordinates:
for x2 in self._field_x_coordinates:
result2[self.xFieldIndexByCoordinate(x1),
self.xFieldIndexByCoordinate(x2)] = result[
self.xIndexByCoordinate(x1),
self.xIndexByCoordinate(x2)]
return result2
def evaluateCutY(self):
log("Calculating vertical cut")
i_x = self.xFieldIndexByCoordinate(0.0)
result = np.zeros((self._field.shape[1], self._field.shape[1]),
dtype=np.complex128)
for r_y in self._field_y_coordinates:
i_y_1 = self.yFieldIndexByCoordinate(r_y)
r1 = np.array([0.0, r_y])
result[i_y_1, :] = self.evaluateAllR_2(r1)[i_x, :]
log("done")
result2 = np.zeros_like(result)
for y1 in self._field_y_coordinates:
for y2 in self._field_y_coordinates:
result2[self.yFieldIndexByCoordinate(y1),
self.yFieldIndexByCoordinate(y2)] = result[
self.yIndexByCoordinate(y1),
self.yIndexByCoordinate(y2)]
return result2
def __init__(self, x_coordinates, y_coordinates, density,
field_x_coordinates, field_y_coordinates, weighted_fields):
self._x_coordinates = x_coordinates
self._y_coordinates = y_coordinates
self._density = density
self._weigthed_fields = weighted_fields
self._setActiveField(0)
self._field_tmp = np.zeros_like(self._field)
self._field_tmp_x = np.zeros(self._field.shape[0], dtype=np.complex128)
self._field_tmp_y = np.zeros(self._field.shape[1], dtype=np.complex128)
self._field_x_coordinates = field_x_coordinates
self._field_y_coordinates = field_y_coordinates
self._last_r_1_x = None
self._last_r_1_y = None
self._x_coordinates_i_min = -(self._x_coordinates.shape[0] + 1) / 2
self._x_coordinates_i_max = (self._x_coordinates.shape[0] + 1) / 2
self._x_coordinates_step_width = self._x_coordinates[
1] - self._x_coordinates[0]
self._y_coordinates_i_min = -(self._y_coordinates.shape[0] + 1) / 2
self._y_coordinates_i_max = (self._y_coordinates.shape[0] + 1) / 2
self._y_coordinates_step_width = self._y_coordinates[
1] - self._y_coordinates[0]
self._grid_area = self._x_coordinates_step_width * self._y_coordinates_step_width
log("Calculating density (integration part)")
self._setRhoR(self._field_x_coordinates, self._field_y_coordinates)
self._rho_phase_tmp = np.zeros_like(self._rho)
self._convolution = Convolution()
self._postInit()
def createReferenceWavefrontAtVirtualSource(self, Rx, dRx, Ry, dRy, configuration, source_position, wavefront):
adapter = SRWAdapter()
if source_position == VIRTUAL_SOURCE_CENTER:
z = -1.0 * self._undulator.length()
elif source_position == VIRTUAL_SOURCE_ENTRANCE:
z = -2.0 * self._undulator.length()
else:
raise NotImplementedError("Source position %s" % source_position)
log("Using source position: %s with z=%.02f" % (source_position, z))
wavefront = adapter.propagate(wavefront, Rx, dRx, Ry, dRy, z)
x_min = -configuration.sourceWavefrontMaximalSizeHorizontal()
x_max = -x_min
y_min = -configuration.sourceWavefrontMaximalSizeVertical()
y_max = -y_min
if x_min > wavefront.minimal_x_coodinate() or x_max < wavefront.maximal_x_coodinate() or \
y_min > wavefront.minimal_y_coodinate() or y_max < wavefront.maximal_y_coodinate():
dim_x = int((x_max-x_min)/wavefront.x_stepwidth())
dim_y = int((y_max-y_min)/wavefront.y_stepwidth())
divisor_x = configuration.samplingFactorDivisorHorizontal()
if divisor_x == "":
divisor_x = 1.0
divisor_y = configuration.samplingFactorDivisorVertical()
if divisor_y == "":
divisor_y = 1.0
wavefront = wavefront.onDomain(x_min, x_max, int(dim_x/divisor_x),
y_min, y_max, int(dim_y/divisor_y))
return wavefront
def setStrategy(self, strategy):
log("Setting autocorrelation strategy: %s" % str(strategy.__name__))
self._strategy = strategy(self._x_coordinates, self._y_coordinates,
self._density, self._field_x_coordinates,
self._field_y_coordinates,
self._weighted_fields)
def log(self, log_string):
log(log_string)
def __init__(self, x_coordinates, y_coordinates, intensity,
eigenvalues_spatial, eigenvectors_parallel,
phase_space_density, method):
self._method = method
if self._method not in ["accurate", "quick"]:
raise Exception("Unknown divergence action %s" % self._method)
communicator = mpi.COMM_WORLD
n_vectors = eigenvalues_spatial.size
self._intensity = intensity
eigenvalues = eigenvalues_spatial
self._number_modes = n_vectors
self._x_coordinates = x_coordinates
self._y_coordinates = y_coordinates
self._petSc_operator = PetScOperatorDivergence(self)
self._my_distribution_plan = DistributionPlan(
communicator=communicator,
n_rows=n_vectors,
n_columns=self.dimensionSize())
self._prepareEigenvectors(communicator, eigenvectors_parallel)
self._my_eigenvalues = eigenvalues[
self._my_distribution_plan.localRows()]
self._my_eigenvectors_conjugated = self._my_eigenvectors.conj()
self._my_eigenvectors_times_eigenvalues = self._my_eigenvectors
self._my_eigenvectors = None
for i_e, e in enumerate(self._my_eigenvalues):
self._my_eigenvectors_times_eigenvalues[i_e, :, :] *= e
self._phase_space_density = phase_space_density
self._sigma_p_x = phase_space_density.divergencePartSigmaX()
self._sigma_p_y = phase_space_density.divergencePartSigmaY()
self._prefactor = phase_space_density.normalizationConstant()
log("Divergence action sigma x/y: %e %e" %
(self._sigma_p_x, self._sigma_p_y))
x_coordinates_weights = x_coordinates[
(x_coordinates > -5 * self._sigma_p_x)
& (x_coordinates < 5 * self._sigma_p_x)]
y_coordinates_weights = y_coordinates[
(y_coordinates > -5 * self._sigma_p_y)
& (y_coordinates < 5 * self._sigma_p_y)]
log("Calculating phase space density xy")
weight_function = np.zeros(
(x_coordinates_weights.shape[0], y_coordinates_weights.shape[0]),
dtype=np.complex128)
for i_x, x in enumerate(x_coordinates_weights):
for i_y, y in enumerate(y_coordinates_weights):
weight_function[i_x, i_y] = phase_space_density.staticPart(
np.array([x, y]))
weight_function_horizontal = np.zeros((x_coordinates.shape[0]),
dtype=np.complex128)
weight_function_vertical = np.zeros((y_coordinates.shape[0]),
dtype=np.complex128)
log("Calculating phase space density x")
for i_x, x in enumerate(x_coordinates_weights):
weight_function_horizontal[i_x] = phase_space_density.staticPart(
np.array([x, 0.0]))
log("Calculating phase space density y")
for i_y, y in enumerate(y_coordinates_weights):
weight_function_vertical[i_y] = phase_space_density.staticPart(
np.array([0.0, y]))
#plot(x_coordinates, weight_function_horizontal)
#plot(y_coordinates, weight_function_vertical)
self._weight_function = weight_function
self._weight_function_horizontal = weight_function_horizontal
self._weight_function_vertical = weight_function_vertical
self._i_action = 0
self._convolution = Convolution()
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 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 _setUpPhases(self):
log("Calculating density (integration phases)")
delta_r = np.array(
[self._x_coordinates_step_width, self._y_coordinates_step_width])
frequency = self._density.expPhasePointsPerOscillation(
(self._field_x_coordinates.max(), self._field_y_coordinates.max()))
oscillations_per_step = frequency * delta_r
log("Have %.3f horizontal and %.3f vertical phase oscillations per step"
% (oscillations_per_step[0], oscillations_per_step[1]))
if oscillations_per_step[0] > 0.25:
raise Exception(
"ABORT. Need AT LEAST 4 integration points per oscillation. Have %f"
% (1 / oscillations_per_step[0]))
if oscillations_per_step[1] > 0.25:
raise Exception(
"ABORT. Need AT LEAST 4 integration points per oscillation. Have %f"
% (1 / oscillations_per_step[1]))
self._setRhoPhase(self._field_x_coordinates, self._field_y_coordinates)
yy, xx = np.meshgrid(self._field_y_coordinates,
self._field_x_coordinates)
self._rho_phase_exp_x = np.zeros(
(len(self._field_x_coordinates), xx.shape[0], xx.shape[1]),
dtype=np.complex128)
self._rho_phase_exp_y = np.zeros(
(len(self._field_y_coordinates), yy.shape[0], yy.shape[1]),
dtype=np.complex128)
for r_x in self._field_x_coordinates:
i_x = self.xIndexByCoordinate(r_x)
rho_phase = self._rho_phase_x[i_x]
self._rho_phase_exp_x[i_x, :, :] = np.exp(xx * rho_phase)
for r_y in self._field_y_coordinates:
i_y = self.xIndexByCoordinate(r_y)
rho_phase = self._rho_phase_y[i_y]
self._rho_phase_exp_y[i_y, :, :] = np.exp(yy * rho_phase)
self._coordinate_map_x = np.zeros(len(self._field_x_coordinates),
dtype=np.int)
self._coordinate_map_y = np.zeros(len(self._field_y_coordinates),
dtype=np.int)
self._coordinate_map_minus_x = np.zeros(len(self._field_x_coordinates),
dtype=np.int)
self._coordinate_map_minus_y = np.zeros(len(self._field_y_coordinates),
dtype=np.int)
for i_r_x, r_x in enumerate(self._field_x_coordinates):
self._coordinate_map_x[i_r_x] = self.xIndexByCoordinate(-r_x)
self._coordinate_map_minus_x[i_r_x] = self.xIndexByCoordinate(r_x)
for i_r_y, r_y in enumerate(self._field_y_coordinates):
self._coordinate_map_y[i_r_y] = self.yIndexByCoordinate(-r_y)
self._coordinate_map_minus_y[i_r_y] = self.yIndexByCoordinate(r_y)
self._has_phases = True
def log(self, log_string):
log(log_string)
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 _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