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 testCalculateRhoPhase(self):
sigma_matrix = SigmaWaist(sigma_x=3e-6,
sigma_y=1e-6,
sigma_x_prime=5e-6,
sigma_y_prime=5e-6)
x_coordinates = np.linspace(-1e-5, 1e-5, 50)
y_coordinates = np.linspace(-1e-5, 1e-5, 50)
wavenumber = 1e+11
density = PhaseSpaceDensity(sigma_matrix, wavenumber)
e_field = createGaussian2D(sigma_x=1.0e-6,
sigma_y=1.0e-6,
x_coordinates=x_coordinates,
y_coordinates=y_coordinates)
e_field = e_field + 0j
strategy = BuilderStrategyPython(x_coordinates, y_coordinates, density,
x_coordinates, y_coordinates,
e_field[np.newaxis, :, :])
for x in x_coordinates[::5]:
for y in y_coordinates[::2]:
norm_diff = np.linalg.norm(
strategy.calculateRhoPhase((x, y)) -
strategy.calculateRhoPhaseSlow((x, y)))
self.assertLess(norm_diff, 1e-10)
def testStaticPartFixedR1(self):
sigma_matrix = SigmaWaist(sigma_x=3e-6,
sigma_y=1e-6,
sigma_x_prime=5e-6,
sigma_y_prime=4e-6)
wavenumber = 1e+11
density = PhaseSpaceDensity(sigma_matrix, wavenumber)
x_coordinates = np.linspace(-10e-6, 10e-6, 100)
y_coordinates = np.linspace(-10e-6, 10e-6, 100)
r_1 = np.array([0.0, 0.0])
density.setAllStaticPartCoordinates(x_coordinates, y_coordinates)
values = density.staticPartFixedR1(r_1)
for i_x, x in enumerate(x_coordinates):
for i_y, y in enumerate(y_coordinates):
dr = r_1 - np.array([x, y])
diff = np.abs(1 - values[i_x, i_y] / density.staticPart(dr))
self.assertLess(diff, 1e-12)
def testIntegrationPartGaussianWithAlpha(self):
sigma_matrix = SigmaMatrixFromCovariance(xx=3e-6,
yy=1e-6,
xxp=5e-6,
yyp=4e-6,
xpxp=0.5 * 3e-6,
ypyp=0.5 * 1e-6)
wavenumber = 1e+11
density = PhaseSpaceDensity(sigma_matrix, wavenumber)
x_coordinates = np.linspace(-10e-6, 10e-6, 100)
y_coordinates = np.linspace(-10e-6, 10e-6, 100)
for dx in x_coordinates:
dr = np.array([dx, 0.0])
for x in x_coordinates:
diff = density.integrationPartOscillation(delta_r=dr,
x=x,
y=0.0)
print(2 * wavenumber * dr[0] *
(sigma_matrix.element('x', 'xp') /
sigma_matrix.element('xp', 'xp')) * x)
diff -= np.exp(2j * wavenumber * dr[0] *
(sigma_matrix.element('x', 'xp') /
sigma_matrix.element('xp', 'xp')) * x)
self.assertLess(np.abs(diff), 1e-12)
for dy in y_coordinates:
dr = np.array([0.0, dy])
for y in y_coordinates:
diff = density.integrationPartOscillation(delta_r=dr,
x=0.0,
y=y)
diff -= np.exp(2j * wavenumber * dr[1] *
(sigma_matrix.element('y', 'yp') /
sigma_matrix.element('yp', 'yp')) * y)
self.assertLess(np.abs(diff), 1e-12)
def testStaticElectronDensity(self):
sigma_x = 3e-6
sigma_y = 1e-6
sigma_matrix = SigmaWaist(sigma_x=sigma_x,
sigma_y=sigma_y,
sigma_x_prime=5e-6,
sigma_y_prime=4e-6)
wavenumber = 1e+11
density = PhaseSpaceDensity(sigma_matrix, wavenumber)
x_coordinates = np.linspace(-10e-6, 10e-6, 100)
y_coordinates = np.linspace(-10e-6, 10e-6, 100)
prefactor = 1.0 / (2 * np.pi * sigma_matrix.sigma_x()**2 *
sigma_matrix.sigma_y()**2 *
sigma_matrix.sigma_d()**2)
#TODO prefactor not correct!
prefactor = density.staticPart(np.array([0, 0]))
diff_prefactor = prefactor / density.staticPart(np.array([0, 0]))
self.assertLess(np.abs(1 - diff_prefactor), 1e-12)
dy = 0.0
for dx in x_coordinates:
dr = np.array([dx, dy])
diff = density.staticPart(delta_r=dr) / \
(prefactor * np.exp(-(dx-dy)**2 * (0.5*wavenumber**2*sigma_matrix.sigma_x_prime()**2)))
self.assertLess(np.abs(1 - diff), 1e-12)
dx = 0.0
for dy in y_coordinates:
dr = np.array([dx, dy])
diff = density.staticPart(delta_r=dr) / \
(prefactor * np.exp(-(dx-dy)**2 * (0.5*wavenumber**2*sigma_matrix.sigma_y_prime()**2)))
self.assertLess(np.abs(1 - diff), 1e-12)
def testIntegrationPartGaussian(self):
sigma_x = 3e-6
sigma_y = 1e-6
sigma_matrix = SigmaWaist(sigma_x=sigma_x,
sigma_y=sigma_y,
sigma_x_prime=5e-6,
sigma_y_prime=4e-6)
wavenumber = 1e+11
density = PhaseSpaceDensity(sigma_matrix, wavenumber)
x_coordinates = np.linspace(-15e-6, 15e-6, 200)
y_coordinates = np.linspace(-7e-6, 7e-6, 200)
for dx in x_coordinates:
diff = density.integrationPartGaussian(
delta=0.0, z=0.0, x=dx, y=0.0) - np.exp(
-dx**2 / (2 * sigma_matrix.sigma_x()**2))
self.assertLess(diff, 1e-12)
for dy in y_coordinates:
diff = density.integrationPartGaussian(
delta=0.0, z=0.0, x=0.0, y=dy) - np.exp(
-dy**2 / (2 * sigma_matrix.sigma_y()**2))
self.assertLess(diff, 1e-12)
rho = np.zeros((x_coordinates.size, y_coordinates.size),
dtype=np.complex128)
for i_x, dx in enumerate(x_coordinates):
for i_y, dy in enumerate(y_coordinates):
rho[i_x, i_y] = density.integrationPartGaussian(delta=0.0,
z=0.0,
x=dx,
y=dy)
norm_1 = np.trapz(np.trapz(np.abs(rho), y_coordinates), x_coordinates)
self.assertLess(1 - norm_1 / (2 * np.pi * sigma_x * sigma_y), 1e-6)
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 phaseSpaceDensity(self):
return PhaseSpaceDensity(sigma_matrix=self._sigma_matrix,
k=self._wavefront.wavenumbers()[0])