def run_bending_magnet_srw(example_index): # example_index=0 is infrared example, example_index=1 is xrays example
###################################################################################################
# Main idea: abstract definition of the setting (electron beam, radiation source, beamline)
# We want to put everything in generic classes that is independent of a specific implementation.
# These are basically the information a scientist would need to physically build the beamline.
#
# Then, we need extra information/settings to perform a calculation. And the extra settings
# vary for different programs. We provide these extra information by attaching program depended
# "settings".
###################################################################################################
#
# 1) define first the electron beam
#
if example_index == 0:
electron_beam = ElectronBeamPencil(energy_in_GeV=3.0,energy_spread=0.89e-3,current=0.5)
# electron_beam = ElectronBeam(energy_in_GeV=3.0,
# energy_spread=0.89e-03,
# current=0.5,
# electrons_per_bunch=500,
# moment_xx = (127.346e-6)**2 ,
# moment_xxp = 0. ,
# moment_xpxp = 100*(91.88e-6)**2,
# moment_yy = (92.3093e-6)**2 ,
# moment_yyp = 0 ,
# moment_ypyp = 100*(7.94e-6)**2 )
else:
#electron_beam = ElectronBeamPencil(energy_in_GeV=6.0,energy_spread=0.89e-3,current=0.2)
electron_beam = ElectronBeam(energy_in_GeV=6.0,
energy_spread=0.89e-03,
current=0.2,
electrons_per_bunch=500,
moment_xx = (77.9e-06)**2 ,
moment_xxp = 0. ,
moment_xpxp = (110.9e-06)**2,
moment_yy = (12.9e-06)**2 ,
moment_yyp = 0 ,
moment_ypyp = (0.5e-06)**2 )
#
# 2) define the magnetic structure
#
if example_index == 0:
#bending_magnet = BendingMagnet(radius=2.25,magnetic_field=0.4,length=4.0)
bending_magnet = BendingMagnet(radius=25.01,magnetic_field=0.4,length=4.0)
else:
bending_magnet = BendingMagnet(radius=23.2655,magnetic_field=0.86,length=0.5)
# Attach SRW bending magnet settings.
#TODO: angular acceptance is used to define screen size.
# NOTE: Maybe angular acceptance is generic and should move to BendingMagnet or Source class??
srw_bending_magnet_setting = SRWBendingMagnetSetting()
if example_index == 0:
horizontal_angle = 0.1
vertical_angle = 0.02
energy = 0.5*0.123984
else:
horizontal_angle = 1e-3
vertical_angle = 0.4e-3
energy = 15000.0
srw_bending_magnet_setting.set_relPrec(0.003)
srw_bending_magnet_setting.set_sampFactNxNyForProp(0.0035)
srw_bending_magnet_setting.set_acceptance_angle(horizontal_angle=horizontal_angle,
vertical_angle=vertical_angle)
bending_magnet.add_settings(srw_bending_magnet_setting)
#
# 3) define beamline containing the optical elements
# In this case, create a beamline that only has one lens attached plus an image (detector) plane
#
#
beamline = Beamline()
# First create the lens.
if example_index == 0:
lens_focal_length = 2.5
else:
lens_focal_length = 12.5
lens = LensIdeal("focus lens",
focal_x=lens_focal_length,
focal_y=lens_focal_length)
# Specify the position of the lens (could set extra parameters for: off-axis and inclination)
# lens_position=p verifies lens equation (1/F = 1/p + 1/q, and p=q for 1:1 magnification)
lens_position = BeamlinePosition(2*lens_focal_length)
# Set settings for SRW.
# These are settings that depend on the "driver" to use.
# If no special settings are set the driver will use its default settings.
lens_setting = SRWBeamlineComponentSetting()
if example_index == 0:
#for SRW experts:
#lens_setting.from_list([1, 1, 1., 0, 0, 1., 2., 1., 2., 0, 0, 0])
lens_setting.set_auto_resize_before_propagation(1) #[0]: Auto-Resize (1) or not (0) Before propagation
lens_setting.set_auto_resize_after_propagation(1) #[1]: Auto-Resize (1) or not (0) After propagation
lens_setting.set_auto_resize_relative_precision(1.) #[2]: Relative Precision for propagation with Auto-Resizing (1. is nominal)
lens_setting.set_allow_semi_analytical_phase_treatment(0) #[3]: Allow (1) or not (0) for semi-analytical treatment of the quadratic (leading) phase terms at the propagation
lens_setting.set_resize_on_ft_side(0) #[4]: Do any Resizing on Fourier side, using FFT, (1) or not (0)
lens_setting.set_resize_factor_horizontal(1.) #[5]: Horizontal Range modification factor at Resizing (1. means no modification)
lens_setting.set_resize_resolution_horizontal(2.) #[6]: Horizontal Resolution modification factor at Resizing
lens_setting.set_resize_factor_vertical(1.) #[7]: Vertical Range modification factor at Resizing
lens_setting.set_resize_resolution_vertical(2.) #[8]: Vertical Resolution modification factor at Resizing
else:
#lens_setting.from_list([0, 0, 1., 0, 0, 1., 5., 1., 8., 0, 0, 0])
lens_setting.set_auto_resize_before_propagation(0)
lens_setting.set_auto_resize_after_propagation(0)
lens_setting.set_auto_resize_relative_precision(1.)
lens_setting.set_allow_semi_analytical_phase_treatment(0)
lens_setting.set_resize_on_ft_side(0)
lens_setting.set_resize_factor_horizontal(1.)
lens_setting.set_resize_resolution_horizontal(5.)
lens_setting.set_resize_factor_vertical(1.)
lens_setting.set_resize_resolution_vertical(8.)
lens.add_settings(lens_setting)
# We could also _simultaneously_ add settings for shadow here:
# lens_setting = ShadowBeamlineComponentSetting()
# lens_setting.setSOMETHING(..)
# lens.addSettings(lens_setting)
# The lens would be configured _simultaneously_ for SRW and SHADOW.
# Attach the component at its position to the beamline.
beamline.attach_component_at(lens, lens_position)
# Second create the image plane.
plane = ImagePlane("Image screen")
plane_setting = SRWBeamlineComponentSetting()
if example_index == 0:
pass #these are default values, so no need to set
#plane_setting.from_list([1, 1, 1., 0, 0, 1., 1., 1., 1., 0, 0, 0])
#plane_setting.set_auto_resize_before_propagation(0)
#plane_setting.set_auto_resize_after_propagation(0)
#plane_setting.set_auto_resize_relative_precision(1.)
#plane_setting.set_allow_semi_analytical_phase_treatment(0)
#plane_setting.set_resize_on_ft_side(0)
#plane_setting.set_resize_factor_horizontal(1.)
#plane_setting.set_resize_resolution_horizontal(1.)
#plane_setting.set_resize_factor_vertical(1.)
#plane_setting.set_resize_resolution_vertical(1.)
else:
#define non-default settings for the propagation in the drift space
#note that although in SRW the driftSpace is a component, in the present beamline
#definition it is not necessary to be defined, as it is automatically added by the
#driver. However, we set here the settings of the drift space that is inserted upstream
#of the "plane" element
drift_space_settings = SRWBeamlineComponentSetting()
drift_space_settings.from_list([0, 0, 1., 1, 0, 1., 1., 1., 1., 0, 0, 0])
plane_setting.set_drift_space_settings(drift_space_settings)
#plane_setting.from_list([0, 0, 1., 0, 0, 4., 1.,1.5, 1., 0, 0, 0])
plane_setting.set_auto_resize_before_propagation(0)
plane_setting.set_auto_resize_after_propagation(0)
plane_setting.set_auto_resize_relative_precision(1.)
plane_setting.set_allow_semi_analytical_phase_treatment(0)
plane_setting.set_resize_on_ft_side(0)
plane_setting.set_resize_factor_horizontal(4.)
plane_setting.set_resize_resolution_horizontal(1.)
plane_setting.set_resize_factor_vertical(1.5)
plane_setting.set_resize_resolution_vertical(1.)
# Attach a screen/image plane.
# Absolute position = distance_source_lens + distance_lens_plane =
# 2*lens_focal_length + 2*lens_focal_lengh = 4*lens_focal_length
plane.add_settings(plane_setting)
plane_position = BeamlinePosition(4*lens_focal_length)
beamline.attach_component_at(plane, plane_position)
#
# Print a summary of the elements used
#
components = [electron_beam,bending_magnet,lens]
print("===========================================================================================================")
for component_index,component in enumerate(components):
tmp = component.to_dictionary()
#tmp = electron_beam.to_dictionary()
print("Component index %d:"%component_index,inspect.getmodule(component))
for i,var in enumerate(tmp):
print(" %20s = %10.5f %5s %s"%(var,tmp[var][0],tmp[var][1],tmp[var][2]))
print("===========================================================================================================")
#
# Calculate the radiation (i.e., run the codes). It returns a native SRWLWfr()
#
# Specify to use SRW.
driver = SRWDriver()
srw_wavefront = driver.calculate_radiation(electron_beam=electron_beam,
magnetic_structure=bending_magnet,
beamline=beamline,
energy_min=energy,
energy_max=energy)
#
# extract the intensity
#
intensity, dim_x, dim_y = driver.calculate_intensity(srw_wavefront)
# # Do some tests.
# # assert abs(1.7063003e+09 - intensity[10, 10])<1e+6, \
# # 'Quick verification of intensity value'
# flux = intensity.sum() * (dim_x[1]-dim_x[0]) * (dim_y[1]-dim_y[0])
# print("Total flux = %10.5e photons/s/.1%%bw"%flux)
# if example_index == 0:
# assert abs(2.40966e+08 - flux)<1e+3, \
# 'Quick verification of intensity value'
# else:
# assert abs(2.14704e+07 - flux)<1e+3, \
# 'Quick verification of intensity value'
# Calculate phases.
#phase = driver.calculate_phase(srw_wavefront)
# # Do some tests.
# checksum = np.sum( np.abs(srw_wavefront.arEx) ) + np.abs( np.sum(srw_wavefront.arEy) )
# print("checksum is: ",checksum)
#
# if example_index == 0:
# assert np.abs(checksum - 1.1845644e+10) < 1e3, "Test electric field checksum"
# else:
# assert np.abs(checksum - 1.53895e+13) < 1e8, "Test electric field checksum"
return srw_wavefront, dim_x, dim_y, intensity
def calculate_radiation(self, electron_beam, magnetic_structure, beamline, energy_min, energy_max):
"""
Calculates radiation.
:param electron_beam: ElectronBeam object
:param magnetic_structure: Source object
:param beamline: beamline object
:param energy_min: Minimal energy for the calculation
:param energy_max: Maximal energy for the calculation
:return: SRW wavefront.
"""
# Get position of the first component. We need this to know where to calculate the source radiation.
first_component = beamline.component_by_index(0)
position_first_component = beamline.position_of(first_component)
# Instanciate an adapter.
srw_adapter = SRWAdapter()
# Create srw electron beam from generic electron beam.
srw_electron_beam = srw_adapter.SRW_electron_beam(electron_beam)
# Calculate the source radiation depending on the chosen source.
# Only undulator here.
# In the real driver this should be refactored to separate functions.
if isinstance(magnetic_structure, Undulator):
undulator = magnetic_structure
srw_undulator = srw_adapter.magnetic_field_from_undulator(undulator)
max_theta = undulator.gaussianCentralConeDivergence(electron_beam.gamma()) * 2.5
z_start = undulator.length()+position_first_component.z()
grid_length = max_theta * z_start / sqrt(2.0)
wavefront = srw_adapter.create_quadratic_SRW_wavefront_single_energy(grid_size=1000,
grid_length=grid_length,
z_start=z_start,
srw_electron_beam=srw_electron_beam,
energy=int(undulator.resonanceEnergy(electron_beam.gamma(),0.0,0.0)))
# Use custom settings if present. Otherwise use default SRW settings.
if undulator.has_settings(self):
# Mind the self in the next line.
# It tells the DriverSettingsManager to use SRW settings.
undulator_settings = undulator.settings(self)
else:
undulator_settings = SRWUndulatorSetting()
print("SRW_driver.calculate_radiation calls CalcElecFieldSR (undulator)...")
t0 = time.time()
srwl.CalcElecFieldSR(wavefront, 0, srw_undulator, undulator_settings.toList())
print("done in ",round(time.time() - t0), "s")
elif isinstance(magnetic_structure, BendingMagnet):
bending_magnet = magnetic_structure
# Use custom settings if present. Otherwise use default SRW settings.
if bending_magnet.has_settings(self):
bending_magnet_settings = bending_magnet.settings(self)
else:
bending_magnet_settings = SRWBendingMagnetSetting()
srw_bending_magnet = srw_adapter.magnetic_field_from_bending_magnet(bending_magnet)
# Determine position of first optical element to calculate initial wavefront there.
z_start = position_first_component.z()
# Determine acceptances.
# Horizontal
horizontal_angle = bending_magnet_settings.horizontal_acceptance_angle()
horizontal_grid_length = 0.5*horizontal_angle*z_start
# Vertical
vertical_angle = bending_magnet_settings.vertical_acceptance_angle()
vertical_grid_length = 0.5*vertical_angle*z_start
# Create rectangular SRW wavefront.
wavefront = srw_adapter.create_rectangular_SRW_wavefront(grid_size=10,
grid_length_vertical=horizontal_grid_length,
grid_length_horizontal=vertical_grid_length,
z_start=z_start,
srw_electron_beam=srw_electron_beam,
energy_min=energy_min,
energy_max=energy_max)
# Calculate initial wavefront.
print("SRW_driver.calculate_radiation calls CalcElecFieldSR (bending magnet)...")
t0 = time.time()
srwl.CalcElecFieldSR(wavefront, 0, srw_bending_magnet, bending_magnet_settings.to_list())
print("done in ",round(time.time() - t0), "s")
else:
raise NotImplementedError
# Create the srw beamline.
srw_optical_element = list()
srw_preferences = list()
# Iterate over all beamline components and translate them.
# Translate free space between two components to drift space.
# Only lenses implemented.
# In the real driver this should be refactored to separate functions.
current_z_position = position_first_component.z()
for component in beamline:
position = beamline.position_of(component)
# Add drift space between two components.
if position.z() > current_z_position:
distance = position.z()-current_z_position
srw_optical_element.append(SRWLOptD(distance))
if component.has_settings(self):
component_settings = component.settings(self)
if component_settings.has_drift_space_settings():
drift_space_settings = component_settings._drift_space_settings
else:
# If there are no drift space settings use default settings.
drift_space_settings = SRWBeamlineComponentSetting()
#TODO: check this, it was inside else
srw_preferences.append(drift_space_settings.to_list())
current_z_position = position.z()
if isinstance(component, LensIdeal):
srw_component = SRWLOptL(_Fx=component.focalX(),
_Fy=component.focalY())
srw_optical_element.append(srw_component)
elif isinstance(component, ImagePlane):
pass
else:
raise NotImplementedError
# Use custom settings if present. Otherwise use default SRW settings.
if component.has_settings(self):
# Mind the self in the next line.
# It tells the DriverSettingsManager to use SRW settings.
component_settings = component.settings(self)
else:
component_settings = SRWBeamlineComponentSetting()
srw_preferences.append(component_settings.to_list())
# Create the srw beamline object.
srw_beamline = SRWLOptC(srw_optical_element,
srw_preferences)
# Call SRW to perform propagation.
print("SRW_driver calls PropagElecField...")
t0 = time.time()
srwl.PropagElecField(wavefront, srw_beamline)
print("done in ",round(time.time() - t0), "s")
# TODO: Decoration of SRW wavefront with glossary object. Consider to better use a "driver results" object.
wavefront._electron_beam = deepcopy(electron_beam)
return wavefront
def run_bending_magnet():
electron_beam = ElectronBeamPencil(
energy_in_GeV=3.0,
energy_spread=0.89e-3,
current=0.5,
)
bending_magnet = BendingMagnet(
radius=25.01,
magnetic_field=0.4,
length=4.0,
)
srw_bending_magnet_setting = SRWBendingMagnetSetting()
horizontal_angle = 0.1
vertical_angle = 0.02
energy = 0.5 * 0.123984
srw_bending_magnet_setting.set_acceptance_angle(
horizontal_angle=horizontal_angle,
vertical_angle=vertical_angle,
)
bending_magnet.add_settings(srw_bending_magnet_setting)
beamline = Beamline()
lens_focal_length = 2.5
lens = LensIdeal(
'focus lens',
focal_x=lens_focal_length,
focal_y=lens_focal_length,
)
lens_position = BeamlinePosition(2 * lens_focal_length)
lens_setting = SRWBeamlineComponentSetting()
lens_setting.set_auto_resize_before_propagation(1)
lens_setting.set_auto_resize_after_propagation(1)
lens_setting.set_auto_resize_relative_precision(1.)
lens_setting.set_allow_semi_analytical_phase_treatment(0)
lens_setting.set_resize_on_ft_side(0)
lens_setting.set_resize_factor_horizontal(1.)
lens_setting.set_resize_resolution_horizontal(2.)
lens_setting.set_resize_factor_vertical(1.)
lens_setting.set_resize_resolution_vertical(2.)
lens.add_settings(lens_setting)
beamline.attach_component_at(lens, lens_position)
plane = ImagePlane('Image screen')
plane_setting = SRWBeamlineComponentSetting()
plane.add_settings(plane_setting)
plane_position = BeamlinePosition(4*lens_focal_length)
beamline.attach_component_at(plane, plane_position)
driver = SRWDriver()
wavefront = driver.calculate_radiation(
electron_beam=electron_beam,
magnetic_structure=bending_magnet,
beamline=beamline,
energy_min=energy,
energy_max=energy,
)
res = driver.calculate_intensity(wavefront)
res = pkcollections.OrderedMapping(
dict(zip(('intensity', 'dim_x', 'dim_y'), res)))
res.wavefront = wavefront
return res
