def as_srw(self, gap="min", energy=None, harmonic=1, **kwargs):
use_srw = kwargs.get("use_srw", False)
if energy is not None:
pars = self.find_harmonic_and_gap(energy,
sort_harmonics=True,
use_srw=use_srw)[0]
gap = pars["gap"]
harmonic = pars["harmonic"]
if isinstance(gap, str) and gap == "min":
gap = self.min_gap
import srwlib
ebeam = beam.srw_ebeam(self.ebeam)
harmB = srwlib.SRWLMagFldH() # magnetic field harmonic
harmB.n = harmonic # harmonic number
harmB.h_or_v = "v" # magnetic field plane: horzontal ('h') or vertical ('v')
harmB.B = self.field(gap=gap) # magnetic field amplitude [T]
und = srwlib.SRWLMagFldU([harmB])
und.per = self.period / 1e3 # period length [m]
und.nPer = self.N # number of periods (will be rounded to integer)
magFldCnt = srwlib.SRWLMagFldC(
[und],
srwlib.array("d", [0]),
srwlib.array("d", [0]),
srwlib.array("d", [0]),
) # Container of all magnetic field elements
return ds(ebeam=ebeam, und=und, field=magFldCnt)
def main():
v = srwl_bl.srwl_uti_parse_options(varParam, use_sys_argv=True)
op = set_optics(v)
v.ss = True
v.ss_pl = 'e'
v.sm = True
v.sm_pl = 'e'
v.pw = True
v.pw_pl = 'xy'
v.si = True
v.si_pl = 'xy'
v.tr = True
v.tr_pl = 'xz'
v.ws = True
v.ws_pl = 'xy'
mag = None
if v.rs_type == 'm':
mag = srwlib.SRWLMagFldC()
mag.arXc.append(0)
mag.arYc.append(0)
mag.arMagFld.append(
srwlib.SRWLMagFldM(v.mp_field, v.mp_order, v.mp_distribution,
v.mp_len))
mag.arZc.append(v.mp_zc)
srwl_bl.SRWLBeamline(_name=v.name, _mag_approx=mag).calc_all(v, op)
def setup(self):
"""
Defines various helpful attributes, builds the fields, and aligns the
reference trajectory.
"""
self.dipole_separation = self.get_dipole_separation()
super(Dogleg, self).setup()
# define helpful attributes
self.reference_gamma = self.reference_energy_mev / 0.5109989461
self.reference_beta = np.sqrt(1 - self.reference_gamma ** -2)
self.dipole_1_rho = self.reference_gamma * self.reference_beta * 299792458 / (1.758820024e11 * self.dipole_1_strength)
self.dipole_1_angle = self.dipole_1_arc_length / self.dipole_1_rho
self.dipole_1_length_z = self.dipole_1_rho * np.sin(self.dipole_1_angle)
self.dipole_1_entrance_z = self.dipole_1_pad_drift_length * np.cos(self.dipole_1_angle)
self.dipole_1_exit_z = self.dipole_1_entrance_z + self.dipole_1_length_z
self.dipole_2_rho = self.reference_gamma * self.reference_beta * 299792458 / (1.758820024e11 * self.dipole_2_strength)
self.dipole_2_angle = self.dipole_2_arc_length / self.dipole_2_rho
self.dipole_2_length_z = self.dipole_2_rho * np.sin(self.dipole_2_angle)
self.dipole_2_entrance_z = self.dipole_1_exit_z + self.dipole_separation
self.dipole_2_exit_z = self.dipole_2_entrance_z + self.dipole_1_length_z
self.detector_z = self.dipole_2_entrance_z + self.detector_distance
self.z_at_end = self.dipole_2_exit_z + self.dipole_2_pad_drift_length * np.cos(self.dipole_2_angle)
# build fields
self.__magnetic_field_regular = srwlib.SRWLMagFldC()
self.__magnetic_field_no_dipole_1 = srwlib.SRWLMagFldC()
self.__magnetic_field_no_dipole_2 = srwlib.SRWLMagFldC()
self.__magnetic_field_only_dipole_1 = srwlib.SRWLMagFldC()
self.__magnetic_field_regular.allocate(0)
self.__magnetic_field_no_dipole_1.allocate(0)
self.__magnetic_field_no_dipole_2.allocate(0)
self.__magnetic_field_only_dipole_1.allocate(0)
# dipole 1
dipole_1 = srwlib.SRWLMagFldM(self.dipole_1_polarity * self.dipole_1_strength, 1, 'n', self.dipole_1_length_z, self.dipole_1_edge_length, 0)
self.__magnetic_field_regular.add(dipole_1, 0, 0, self.dipole_1_entrance_z + 0.5 * self.dipole_1_length_z)
self.__magnetic_field_no_dipole_2.add(dipole_1, 0, 0, self.dipole_1_entrance_z + 0.5 * self.dipole_1_length_z)
self.__magnetic_field_only_dipole_1.add(dipole_1, 0, 0, self.dipole_1_entrance_z + 0.5 * self.dipole_1_length_z)
# beam optics between dipoles
self.add_beam_optics(self.__magnetic_field_regular)
self.add_beam_optics(self.__magnetic_field_no_dipole_1)
self.add_beam_optics(self.__magnetic_field_no_dipole_2)
# dipole 2
dipole_2 = srwlib.SRWLMagFldM(self.dipole_2_polarity * self.dipole_2_strength, 1, 'n', self.dipole_2_length_z, self.dipole_2_edge_length, 0)
self.__magnetic_field_regular.add(dipole_2, 0, 0, self.dipole_2_entrance_z + 0.5 * self.dipole_2_length_z)
self.__magnetic_field_no_dipole_1.add(dipole_2, 0, 0, self.dipole_2_entrance_z + 0.5 * self.dipole_2_length_z)
# align reference trajectory
self.align_reference_trajectory()
def main():
v = srwl_bl.srwl_uti_parse_options(varParam, use_sys_argv=True)
setup_magnetic_measurement_files(v, "configurations/magn_meas_u20_hxn.zip")
op = set_optics(v)
v.ws = True
v.ws_pl = 'xy'
mag = None
if v.rs_type == 'm':
mag = srwlib.SRWLMagFldC()
mag.arXc.append(0)
mag.arYc.append(0)
mag.arMagFld.append(srwlib.SRWLMagFldM(v.mp_field, v.mp_order, v.mp_distribution, v.mp_len))
mag.arZc.append(v.mp_zc)
srwl_bl.SRWLBeamline(_name=v.name, _mag_approx=mag).calc_all(v, op)
def run_experiment(other_params, propagation_params, set_up_funcs, plot=False):
"""
names - [list] the name list of instruments.
setting_params - [list] the meta information of the experiment.
physics_params - [list] .
propagation_params - [list] .
"""
varParam = srwl_bl.srwl_uti_ext_options(other_params + propagation_params)
v = srwl_bl.srwl_uti_parse_options(varParam, use_sys_argv=True)
op = None
for n, func in enumerate(set_up_funcs):
value = func(v)
if func.__name__ == 'set_optics':
op = value
if op is None:
raise ValueError(
"set_optics() function should be included in set_up_funcs")
# this part is different for different beamline?
v.ws = True
if plot:
v.ws_pl = 'xy'
mag = None
if v.rs_type == 'm':
mag = srwlib.SRWLMagFldC()
mag.arXc.append(0)
mag.arYc.append(0)
mag.arMagFld.append(
srwlib.SRWLMagFldM(v.mp_field, v.mp_order, v.mp_distribution,
v.mp_len))
mag.arZc.append(v.mp_zc)
srwl_bl.SRWLBeamline(_name=v.name, _mag_approx=mag).calc_all(v, op)
# remove saved image.
if plot:
plt.show()
else:
for file in glob.glob('*.png'):
os.remove(file)
plt.close('all')
return
def calc2d_srw(self, photon_energy, photon_energy_step, scanning_data):
Kv = scanning_data.get_additional_parameter("Kv")
Kh = scanning_data.get_additional_parameter("Kh")
period_id = scanning_data.get_additional_parameter("period_id")
n_periods = scanning_data.get_additional_parameter("n_periods")
B0v = Kv/period_id/(codata.e/(2*numpy.pi*codata.electron_mass*codata.c))
B0h = Kh/period_id/(codata.e/(2*numpy.pi*codata.electron_mass*codata.c))
eBeam = srwlib.SRWLPartBeam()
eBeam.Iavg = scanning_data.get_additional_parameter("electron_current")
eBeam.partStatMom1.gamma = scanning_data.get_additional_parameter("electron_energy") / (codata_mee * 1e-3)
eBeam.partStatMom1.relE0 = 1.0
eBeam.partStatMom1.nq = -1
eBeam.partStatMom1.x = 0.0
eBeam.partStatMom1.y = 0.0
eBeam.partStatMom1.z = -0.5*period_id*n_periods + 4
eBeam.partStatMom1.xp = 0.0
eBeam.partStatMom1.yp = 0.0
eBeam.arStatMom2[ 0] = scanning_data.get_additional_parameter("electron_beam_size_h") ** 2
eBeam.arStatMom2[ 1] = 0.0
eBeam.arStatMom2[ 2] = scanning_data.get_additional_parameter("electron_beam_divergence_h") ** 2
eBeam.arStatMom2[ 3] = scanning_data.get_additional_parameter("electron_beam_size_v") ** 2
eBeam.arStatMom2[ 4] = 0.0
eBeam.arStatMom2[ 5] = scanning_data.get_additional_parameter("electron_beam_divergence_v") ** 2
eBeam.arStatMom2[10] = scanning_data.get_additional_parameter("electron_energy_spread") ** 2
gap_h = scanning_data.get_additional_parameter("gap_h")
gap_v = scanning_data.get_additional_parameter("gap_v")
mesh = srwlib.SRWLRadMesh(photon_energy,
photon_energy,
1,
-gap_h / 2, gap_h / 2, scanning_data.get_additional_parameter("h_slits_points"),
-gap_v / 2, gap_v / 2, scanning_data.get_additional_parameter("v_slits_points"),
scanning_data.get_additional_parameter("distance"))
srw_magnetic_fields = []
if B0v > 0: srw_magnetic_fields.append(srwlib.SRWLMagFldH(1, "v", B0v))
if B0h > 0: srw_magnetic_fields.append(srwlib.SRWLMagFldH(1, "h", B0h))
magnetic_structure = srwlib.SRWLMagFldC([srwlib.SRWLMagFldU(srw_magnetic_fields, period_id, n_periods)],
srwlib.array("d", [0]), srwlib.array("d", [0]), srwlib.array("d", [0]))
wfr = srwlib.SRWLWfr()
wfr.mesh = mesh
wfr.partBeam = eBeam
wfr.allocate(mesh.ne, mesh.nx, mesh.ny)
srwlib.srwl.CalcElecFieldSR(wfr, 0, magnetic_structure, [1, 0.01, 0, 0, 50000, 1, 0])
mesh_out = wfr.mesh
h_array=numpy.linspace(mesh_out.xStart, mesh_out.xFin, mesh_out.nx)*1e3 # in mm
v_array=numpy.linspace(mesh_out.yStart, mesh_out.yFin, mesh_out.ny)*1e3 # in mm
intensity_array = numpy.zeros((h_array.size, v_array.size))
arI0 = srwlib.array("f", [0]*mesh_out.nx*mesh_out.ny) #"flat" array to take 2D intensity data
srwlib.srwl.CalcIntFromElecField(arI0, wfr, 6, 1, 3, photon_energy, 0, 0)
data = numpy.ndarray(buffer=arI0, shape=(mesh_out.ny, mesh_out.nx),dtype=arI0.typecode)
for ix in range(h_array.size):
for iy in range(v_array.size):
intensity_array[ix, iy] = data[iy,ix]
return self.calculate_power(h_array, v_array, intensity_array, photon_energy_step)
def track_drift(trajectory):
dummy_field = srwlib.SRWLMagFldM(0, 1, 'n', 0.1, 0, 0)
container = srwlib.SRWLMagFldC()
container.add(dummy_field, 0, 0, 0.05)
return srwlib.srwl.CalcPartTraj(trajectory, container, [1])
def run_srw_simulation(task_cut, hOffsetIdx, vOffsetIdx, hOffsetIdx2,
vOffsetIdx2, file_idx, nvx_idx, nvy_idx, nvz_idx,
nvx2_idx, nvy2_idx, nvz2_idx, wp_idx, process_number):
""" Runs the number of SRW simulations as determined by the number of rows in task_cut. task_cut holds the parameters that will be updated to varParam to introduce offsets
task_cut:
i: task number
offset_notation: binary values for which offsets/rotations are applied: offsets_mirror1, offsets_mirror2, rotations_mirror1, rotations_mirror2
dx1: the horizontal offset to mirror 1
dy1: the vertical offset to mirror 1
dx2: the horizontal offset to mirror 2
dy2: the vertical offset to mirror 2
thetax1: thetax for rotation matrix to be applied to mirror 1
thetay1: thetay for rotation matrix to be applied to mirror 1
thetaz1: thetaz for rotation matrix to be applied to mirror 1
thetax2: thetax for rotation matrix to be applied to mirror 2
thetay2: thetay for rotation matrix to be applied to mirror 2
thetaz2: thetaz for rotation matrix to be applied to mirror 2
wp: watchpoint location for final screen
hOffsetIdx: index of where to update horizontal offset in varParams for mirror 1
vOffsetIdx: index of where to update vertical offset in varParams for mirror 1
hOffsetIdx2: index of where to update horizontal offset in varParams for mirror 2
vOffsetIdx2: index of where to update vertical offset in varParams for mirror
file_idx: the index of the output beam file
nvx_idx: index of where to update thetax in varParams for mirror 1
nvy_idx: index of where to update thetay in varParams for mirror 1
nvz_idx: index of where to update thetaz in varParams for mirror 1
nvx2_idx: index of where to update thetax in varParams for mirror 2
nvy2_idx: index of where to update thetay in varParams for mirror 2
nvz2_idx: index of where to update thetaz in varParams for mirror 2
wp_idx: index of where to update the final screen
process_number: the process number running this task cut
"""
print('Process ' + str(process_number) + ' to complete ' +
str(len(task_cut)) + ' tasks')
for t in range(len(task_cut)):
varParam = get_reset_varParam()
#print(len(task_cut))
#print(task_cut)
#### get task parameters based on which combo
i, offset_notation, dx1, dy1, dx2, dy2, thetax1, thetay1, thetaz1, thetax2, thetay2, thetaz2, wp = task_cut[
t]
offsets_mirror1, offsets_mirror2, rotations_mirror1, rotations_mirror2, watchpoint_pos = offset_notation
task_params = []
#### update params
if offsets_mirror1:
varParam[hOffsetIdx][2] = dx1
task_params.append(dx1.reshape(1, 1))
varParam[vOffsetIdx][2] = dy1
if offsets_mirror2:
varParam[hOffsetIdx2][2] = dx2
varParam[vOffsetIdx2][2] = dy2
task_params.append(dy2.reshape(1, 1))
if rotations_mirror1:
vx = varParam[nvx_idx][2]
vy = varParam[nvy_idx][2]
vz = varParam[nvz_idx][2]
Rx = np.array([[1, 0, 0], [0, np.cos(thetax1), -np.sin(thetax1)],
[0, np.sin(thetax1),
np.cos(thetax1)]])
Ry = np.array([[np.cos(thetay1), 0,
np.sin(thetay1)], [0, 1, 0],
[-np.sin(thetay1), 0,
np.cos(thetay1)]])
Rz = np.array([[np.cos(thetaz1), -np.sin(thetaz1), 0],
[np.sin(thetaz1),
np.cos(thetaz1), 0], [0, 0, 1]])
Rxy = np.dot(Rx, Ry)
R_tot = np.dot(Rxy, Rz)
v = np.array([vx, vy, vz]).reshape(3, 1)
rtot_v = np.dot(R_tot, v)
varParam[nvx_idx][2] = rtot_v[0]
varParam[nvy_idx][2] = rtot_v[1]
varParam[nvz_idx][2] = rtot_v[2]
#task_params.append(thetax1.reshape(1,1))
task_params.append(thetay1.reshape(1, 1))
task_params.append(thetaz1.reshape(1, 1))
if rotations_mirror2:
vx2 = varParam[nvx2_idx][2]
vy2 = varParam[nvy2_idx][2]
vz2 = varParam[nvz2_idx][2]
Rx2 = np.array([[1, 0, 0], [0,
np.cos(thetax2), -np.sin(thetax2)],
[0, np.sin(thetax2),
np.cos(thetax2)]])
Ry2 = np.array([[np.cos(thetay2), 0,
np.sin(thetay2)], [0, 1, 0],
[-np.sin(thetay2), 0,
np.cos(thetay2)]])
Rz2 = np.array([[np.cos(thetaz2), -np.sin(thetaz2), 0],
[np.sin(thetaz2),
np.cos(thetaz2), 0], [0, 0, 1]])
Rxy2 = np.dot(Rx2, Ry2)
R_tot2 = np.dot(Rxy2, Rz2)
v2 = np.array([vx2, vy2, vz2]).reshape(3, 1)
rtot_v2 = np.dot(R_tot2, v2)
varParam[nvx2_idx][2] = rtot_v2[0]
varParam[nvy2_idx][2] = rtot_v2[1]
varParam[nvz2_idx][2] = rtot_v2[2]
task_params.append(thetax2.reshape(1, 1))
#task_params.append(thetay2.reshape(1,1))
task_params.append(thetaz2.reshape(1, 1))
if watchpoint_pos:
print('updating with ' + str(wp))
varParam[wp_idx][2] = wp
task_params.append(wp.reshape(1, 1))
save_dat = 'dat_files/res_int_se_' + str(i) + '.dat'
varParam[file_idx][2] = save_dat
v = srwl_bl.srwl_uti_parse_options(varParam, use_sys_argv=True)
op = set_optics(v)
v.si = True
v.si_pl = ''
v.ws = True
v.ws_pl = ''
mag = None
if v.rs_type == 'm':
mag = srwlib.SRWLMagFldC()
mag.arXc.append(0)
mag.arYc.append(0)
mag.arMagFld.append(
srwlib.SRWLMagFldM(v.mp_field, v.mp_order, v.mp_distribution,
v.mp_len))
mag.arZc.append(v.mp_zc)
srwl_bl.SRWLBeamline(_name=v.name, _mag_approx=mag).calc_all(v, op)
beam = read_srw_file(save_dat)
h = beam.shape[0]
w = beam.shape[1]
beam = beam.reshape(1, h, w)
task_params = np.concatenate(task_params, axis=1)
np.save('beams/beam_' + str(int(i)) + '.npy', beam)
np.save('parameters/offsets_' + str(int(i)) + '.npy', task_params)
#time.sleep(0.5)
print('Process ' + str(process_number) + ' finished task ' +
str(int(i)))
def _srw_undulators(und, Xc, Yc, Zc): #for the moment only one works
cnt = sl.SRWLMagFldC([und], array.array('d', [Xc]),
array.array('d', [Yc]), array.array('d', [Zc]))
return cnt