"""
WOLightSourcePySRU — wofry undulator light source using pySRU for coherent-mode-decomposition wavefront generation.
"""
import numpy
from syned.storage_ring.light_source import LightSource
from syned.storage_ring.electron_beam import ElectronBeam
from syned.storage_ring.magnetic_structures.undulator import Undulator
from wofry.beamline.decorators import LightSourceDecorator
from wofryimpl.propagator.util.undulator_coherent_mode_decomposition_1d import UndulatorCoherentModeDecomposition1D
from wofry.propagator.wavefront2D.generic_wavefront import GenericWavefront2D
from pySRU.Simulation import create_simulation
from pySRU.ElectronBeam import ElectronBeam as PySruElectronBeam
from pySRU.MagneticStructureUndulatorPlane import MagneticStructureUndulatorPlane
from pySRU.TrajectoryFactory import TrajectoryFactory
from pySRU.RadiationFactory import TRAJECTORY_METHOD_ANALYTIC, TRAJECTORY_METHOD_ODE
from pySRU.RadiationFactory import RadiationFactory
from pySRU.RadiationFactory import RADIATION_METHOD_NEAR_FIELD, RADIATION_METHOD_APPROX_FARFIELD
[docs]
class WOPySRULightSource(LightSource, LightSourceDecorator):
def __init__(self,
name = "Undefined",
electron_beam = None,
magnetic_structure = None,
number_of_trajectory_points = 5000,
traj_method = TRAJECTORY_METHOD_ODE,
rad_method = RADIATION_METHOD_APPROX_FARFIELD,
):
LightSource.__init__(self, name=name, electron_beam=electron_beam, magnetic_structure=magnetic_structure)
self.__source_wavefront_parameters = {
'distance' : 30.0,
'gapH' : 0.003,
'gapV' : 0.003,
'photon_energy' : 7000.0,
'h_slit_points' : 51,
'v_slit_points' : 51,
'flag_send_wavefront_dimension' : 0, # 0=2D wavefront, 1=H wawefront, 2=V wavefront (for python script)
'number_of_trajectory_points' : number_of_trajectory_points,
'traj_method' : traj_method,
'rad_method' : rad_method,
}
self._dimension = 2
self._number_of_trajectory_points = number_of_trajectory_points
self._traj_method = traj_method
self._rad_method = rad_method
self._set_support_text([
# ("name" , "to define ", "" ),
("dimension" , "dimension ", "" ),
("number_of_trajectory_points", "number_of_trajectory_points ", ""),
("traj_method", "traj_method ", ""),
("rad_method" , "rad_method ", "" ),
] )
[docs]
@classmethod
def initialize_from_keywords(cls,
name="Undefined",
energy_in_GeV=6.04,
current=0.2,
K_vertical=1.68,
period_length=0.018,
number_of_periods=222,
distance=10.0,
gapH=0.003,
gapV=0.003,
photon_energy=7000.0,
h_slit_points=51,
v_slit_points=51,
flag_send_wavefront_dimension=0,
number_of_trajectory_points=5000,
traj_method=TRAJECTORY_METHOD_ODE,
rad_method=RADIATION_METHOD_APPROX_FARFIELD,
):
out = WOPySRULightSource(name=name,
electron_beam=ElectronBeam(
energy_in_GeV=energy_in_GeV,
current=current),
magnetic_structure=Undulator(
K_vertical=K_vertical,
period_length=period_length,
number_of_periods=number_of_periods),
number_of_trajectory_points=number_of_trajectory_points,
traj_method=traj_method,
rad_method=rad_method,
)
out.set_source_wavefront_parameters(distance=distance,
gapH=gapH,
gapV=gapV,
photon_energy=photon_energy,
h_slit_points=h_slit_points,
v_slit_points=v_slit_points,
flag_send_wavefront_dimension=flag_send_wavefront_dimension)
return out
[docs]
def set_source_wavefront_parameters(self, distance=None, gapH=None, gapV=None, photon_energy=None,
h_slit_points=None, v_slit_points=None,
flag_send_wavefront_dimension=None,):
if distance is not None: self.__source_wavefront_parameters['distance'] = distance
if gapH is not None: self.__source_wavefront_parameters['gapH'] = gapH
if gapV is not None: self.__source_wavefront_parameters['gapV'] = gapV
if photon_energy is not None: self.__source_wavefront_parameters['photon_energy'] = photon_energy
if h_slit_points is not None: self.__source_wavefront_parameters['h_slit_points'] = h_slit_points
if v_slit_points is not None: self.__source_wavefront_parameters['v_slit_points'] = v_slit_points
if flag_send_wavefront_dimension is not None:
self.__source_wavefront_parameters['flag_send_wavefront_dimension'] = flag_send_wavefront_dimension
[docs]
def get_source_wavefront_parameters(self):
return self.__source_wavefront_parameters
[docs]
def get_dimension(self): # overwrite this method to export 1D wavefronts if wanted
if self.__source_wavefront_parameters['flag_send_wavefront_dimension'] == 0:
return 2
elif self.__source_wavefront_parameters['flag_send_wavefront_dimension'] == 1:
return 1
elif self.__source_wavefront_parameters['flag_send_wavefront_dimension'] == 2:
return 1
# from Wofry Decorator
[docs]
def get_wavefront(self):
E, x, y, H, V, I, II = self.calculate_undulator_emission()
# normalize E
distance = self.__source_wavefront_parameters['distance']
# pySRU results are ph/s/0.1bw/mrad^2/distance^2 so I multiply E by sqrt(distance**2)
wf = GenericWavefront2D.initialize_wavefront_from_arrays(x, y,
E[:,:,0].copy() * distance ,
z_array_pi=E[:,:,1].copy() * distance,
wavelength=1e-10)
wf.set_photon_energy(self.__source_wavefront_parameters['photon_energy'])
return wf #, I, x, y, II
[docs]
def to_python_code(self, do_plot=True, add_import_section=False):
txt = ""
txt += "#"
txt += "\n# create output_wavefront\n#"
txt += "\n#"
txt += "\nfrom wofryimpl.propagator.light_source_pysru import WOPySRULightSource"
ebeam = self.get_electron_beam()
und = self.get_magnetic_structure()
txt += "\nlight_source = WOPySRULightSource.initialize_from_keywords("
txt += "\n energy_in_GeV=%g," % ebeam.energy()
txt += "\n current=%g," % ebeam.current()
txt += "\n K_vertical=%g," % und.K_vertical()
txt += "\n period_length=%g," % und.period_length()
txt += "\n number_of_periods=%g," % und.number_of_periods()
txt += "\n distance=%g," % self.__source_wavefront_parameters['distance']
txt += "\n gapH=%g," % self.__source_wavefront_parameters['gapH']
txt += "\n gapV=%g," % self.__source_wavefront_parameters['gapV']
txt += "\n photon_energy=%g," % self.__source_wavefront_parameters['photon_energy']
txt += "\n h_slit_points=%d," % self.__source_wavefront_parameters['h_slit_points']
txt += "\n v_slit_points=%d," % self.__source_wavefront_parameters['v_slit_points']
txt += "\n number_of_trajectory_points=%d," % self.__source_wavefront_parameters['number_of_trajectory_points']
txt += "\n traj_method=%d, # 0=TRAJECTORY_METHOD_ANALYTIC, 1=TRAJECTORY_METHOD_ODE" % self.__source_wavefront_parameters['traj_method']
txt += "\n rad_method=%d, # 0=RADIATION_METHOD_NEAR_FIELD, 1= RADIATION_METHOD_APPROX, 2=RADIATION_METHOD_APPROX_FARFIELD" % self.__source_wavefront_parameters['rad_method']
txt += "\n )"
if self.__source_wavefront_parameters['flag_send_wavefront_dimension'] == 0:
txt += "\n\noutput_wavefront = light_source.get_wavefront()"
elif self.__source_wavefront_parameters['flag_send_wavefront_dimension'] == 1:
txt += "\n\noutput_wavefront = light_source.get_wavefront().get_Wavefront1D_from_profile(0, 0.0)"
elif self.__source_wavefront_parameters['flag_send_wavefront_dimension'] == 2:
txt += "\n\noutput_wavefront = light_source.get_wavefront().get_Wavefront1D_from_profile(1, 0.0)"
return txt
[docs]
def calculate_undulator_emission(self):
distance = self.__source_wavefront_parameters['distance']
gapH = self.__source_wavefront_parameters['gapH']
gapV = self.__source_wavefront_parameters['gapV']
photon_energy = self.__source_wavefront_parameters['photon_energy']
zero_emittance = False
h_slit_points = self.__source_wavefront_parameters['h_slit_points']
v_slit_points = self.__source_wavefront_parameters['v_slit_points']
print('Running pySRU', distance, gapV, gapH, photon_energy, v_slit_points, h_slit_points)
hArray = numpy.linspace(-0.5 * gapH, 0.5 * gapH, h_slit_points)
vArray = numpy.linspace(-0.5 * gapV, 0.5 * gapV, v_slit_points)
H = numpy.outer(hArray, numpy.ones_like(vArray))
V = numpy.outer(numpy.ones_like(hArray), vArray)
myBeam = PySruElectronBeam(Electron_energy = self.get_electron_beam().energy(),
I_current = self.get_electron_beam().current())
myUndulator = MagneticStructureUndulatorPlane(K = self.get_magnetic_structure().K_vertical(),
period_length = self.get_magnetic_structure().period_length(),
length = self.get_magnetic_structure().length())
simulation_test = create_simulation(magnetic_structure=myUndulator, electron_beam=myBeam,
magnetic_field=None, photon_energy=photon_energy,
traj_method=TRAJECTORY_METHOD_ODE, Nb_pts_trajectory=None,
rad_method=RADIATION_METHOD_NEAR_FIELD, Nb_pts_radiation=None,
initial_condition=None, distance=distance, XY_are_list=False,
X=hArray, Y=vArray)
intensArray = simulation_test.radiation.intensity.copy()
electric_field = simulation_test.electric_field
E = electric_field._electrical_field.copy()
II = numpy.sum(numpy.abs(E) ** 2, axis=-1)
# grid in m
return (E, hArray, vArray, H, V, intensArray, II)
if __name__ == "__main__":
pp = WOPySRULightSource.initialize_from_keywords(
name="",
energy_in_GeV=6.04,
current=0.2,
K_vertical=1.68,
period_length=0.018,
number_of_periods=222,
distance=10.0,
gapH=0.003,
gapV=0.003,
photon_energy=7000.0,
h_slit_points=51,
v_slit_points=51,
flag_send_wavefront_dimension=1,
)
wf = pp.get_wavefront()
print(">>>>> Dimension: ", pp.get_dimension())
from srxraylib.plot.gol import plot_image
plot_image(wf.get_intensity(), wf.get_coordinate_x(), wf.get_coordinate_y())
# plot_image(iint, x, y, title="pySRU")
# plot_image(II, x, y, title="Me")
print(pp.to_python_code())