def add_beam_optics(self, magnetic_field_container):
quadrupole_1 = srwlib.SRWLMagFldM(self.quadrupole_1_strength, 2, 's',
self.quadrupole_1_length,
self.quadrupole_1_edge_length, 0)
magnetic_field_container.add(
quadrupole_1, self.quadrupole_1_offset_x,
self.quadrupole_1_offset_y, self.dipole_1_exit_z +
self.drift_1_length + 0.5 * self.quadrupole_1_length)
quadrupole_2 = srwlib.SRWLMagFldM(self.quadrupole_2_strength, 2, 'n',
self.quadrupole_2_length,
self.quadrupole_2_edge_length, 0)
magnetic_field_container.add(
quadrupole_2, self.quadrupole_2_offset_x,
self.quadrupole_2_offset_y, self.dipole_1_exit_z +
self.drift_1_length + self.quadrupole_1_length +
self.drift_2_length + 0.5 * self.quadrupole_2_length)
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 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)))
