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 main():
v = srwl_bl.srwl_uti_parse_options(varParam, use_sys_argv=True)
source_type, mag = srwl_bl.setup_source(v)
op = set_optics(v)
v.ws = True
v.ws_pl = 'xy'
srwl_bl.SRWLBeamline(_name=v.name, _mag_approx=mag).calc_all(v, op)
def main():
v = srwl_bl.srwl_uti_parse_options(varParam, use_sys_argv=False)
source_type, mag = srwl_bl.setup_source(v)
op = set_optics(v)
v.wm_na = v.sm_na = 5
# Number of "iterations" per save is best set to num processes
v.wm_ns = v.sm_ns = 4
srwl_bl.SRWLBeamline(_name=v.name).calc_all(v, op)
def run_background(cfg_dir):
with pkio.save_chdir(cfg_dir):
fn = 'run_background.py'
cmd = [sys.executable or 'python', fn]
script = pkio.read_text('srw_parameters.py')
p = dict(pkcollections.map_items(cfg))
if cfg.slave_processes > 1:
cmd[0:0] = [
'mpiexec',
'-n',
# SRW includes a master process so 2 really needs 3 processes
str(cfg.slave_processes + 1),
]
script += '''
from mpi4py import MPI
if MPI.COMM_WORLD.Get_rank():
import signal
signal.signal(signal.SIGTERM, lambda x, y: MPI.COMM_WORLD.Abort(1))
'''
else:
# In interactive (dev) mode, output as frequently as possible
p['particles_per_slave'] = 1
script += '''
import srwl_bl
v = srwl_bl.srwl_uti_parse_options(get_srw_params(), use_sys_argv=False)
source_type, mag = setup_source(v)
v.wm = True
v.wm_nm = {total_particles}
v.wm_na = {particles_per_slave}
# Number of "iterations" per save is best set to num processes
v.wm_ns = {slave_processes}
op = get_beamline_optics()
srwl_bl.SRWLBeamline(_name=v.name).calc_all(v, op)
'''.format(**p)
pkio.write_text(fn, script)
try:
p = subprocess.Popen(
cmd,
stdin=open(os.devnull),
stdout=open('run_background.out', 'w'),
stderr=subprocess.STDOUT,
)
signal.signal(signal.SIGTERM, lambda x, y: p.terminate())
rc = p.wait()
if rc != 0:
p = None
raise RuntimeError('child terminated: retcode={}'.format(rc))
finally:
if not p is None:
p.terminate()
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 main():
v = srwl_bl.srwl_uti_parse_options(varParam, use_sys_argv=True)
source_type, mag = srwl_bl.setup_source(v)
setup_magnetic_measurement_files("magn_meas_srx.zip", v)
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'
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 _run_srw():
run_dir = os.getcwd()
with open('in.json') as f:
data = json.load(f)
#TODO(pjm): need to properly escape data values, untrusted from client
# this defines the get_srw_params() and get_beamline_optics() functions
exec(pkio.read_text('srw_parameters.py'), locals(), locals())
v = srwl_bl.srwl_uti_parse_options(get_srw_params(), use_sys_argv=False)
source_type, mag = setup_source(v)
op = None
if data['report'] == 'intensityReport':
v.ss = True
outfile = v.ss_fn
elif data['report'] == 'fluxReport':
v.sm = True
outfile = v.sm_fn
elif data['report'] == 'powerDensityReport':
v.pw = True
outfile = v.pw_fn
elif data['report'] == 'initialIntensityReport' or data[
'report'] == 'sourceIntensityReport':
v.si = True
outfile = v.si_fn
elif data['report'] == 'mirrorReport':
_process_output(_mirror_plot(data), data)
return
elif re.search('^watchpointReport', data['report']):
op = get_beamline_optics()
v.ws = True
outfile = v.ws_fni
else:
raise Exception('unknown report: {}'.format(data['report']))
if isinstance(mag, srwlib.SRWLGsnBm):
mag = None
srwl_bl.SRWLBeamline(_name=v.name, _mag_approx=mag).calc_all(v, op)
_process_output(outfile, data)
['ueb', 'i', 0, 'Use user defined beam'],
['ueb_e', 'f', 3.0, 'energy [GeV]'],
['ueb_sig_e', 'f', 0.00089, 'RMS energy spread'],
['ueb_beam_definition', 's', 't', 'definition of the beam using Twiss Parameters (t) or Moments (m)'],
['ueb_emit_x', 'f', 9e-10, 'horizontal emittance [m]'],
['ueb_beta_x', 'f', 2.02, 'horizontal beta-function [m]'],
['ueb_alpha_x', 'f', 0.0, 'horizontal alpha-function [rad]'],
['ueb_eta_x', 'f', 0.0, 'horizontal dispersion function [m]'],
['ueb_eta_x_pr', 'f', 0.0, 'horizontal dispersion function derivative [rad]'],
['ueb_emit_y', 'f', 8e-12, 'vertical emittance [m]'],
['ueb_beta_y', 'f', 1.06, 'vertical beta-function [m]'],
['ueb_alpha_y', 'f', 0.0, 'vertical alpha-function [rad]'],
['ueb_eta_y', 'f', 0.0, 'vertical dispersion function [m]'],
['ueb_eta_y_pr', 'f', 0.0, 'vertical dispersion function derivative [rad]'],
['ueb_rms_size_x', 'f', 0.000372612, "horizontal RMS size [m]"],
['ueb_rms_diverg_x', 'f', 1.04666e-05, "horizontal RMS divergence [rad]"],
['ueb_xxpr_x', 'f', 0.0, "horizontal <(x-<x>)(x'-<x'>)> [m]"],
['ueb_rms_size_y', 'f', 9.87421e-06, "vertical RMS size [m]"],
['ueb_rms_diverg_y', 'f', 3.94968e-06, "vertical RMS divergence [rad]"],
['ueb_xxpr_y', 'f', 0.0, "vertical <(x-<x>)(x'-<x'>)> [m]"],
])
if __name__ == '__main__':
v = srwl_bl.srwl_uti_parse_options(varParam)
source_type, mag = srwl_bl.setup_source(v)
op = None
op = set_optics()
v.ws = True
v.ws_pl = 'xy'
srwl_bl.SRWLBeamline(_name=v.name, _mag_approx=mag).calc_all(v, op)
['ueb_beta_y', 'f', 1.06, 'vertical beta-function [m]'],
['ueb_alpha_y', 'f', 0.0, 'vertical alpha-function [rad]'],
['ueb_eta_y', 'f', 0.0, 'vertical dispersion function [m]'],
[
'ueb_eta_y_pr', 'f', 0.0,
'vertical dispersion function derivative [rad]'
],
['ueb_rms_size_x', 'f', 4.26380112107e-05, "horizontal RMS size [m]"],
[
'ueb_rms_diverg_x', 'f', 2.11079263419e-05,
"horizontal RMS divergence [rad]"
],
['ueb_xxpr_x', 'f', 0.0, "horizontal <(x-<x>)(x'-<x'>)> [m]"],
['ueb_rms_size_y', 'f', 2.91204395571e-06, "vertical RMS size [m]"],
[
'ueb_rms_diverg_y', 'f', 2.74721127897e-06,
"vertical RMS divergence [rad]"
],
['ueb_xxpr_y', 'f', 0.0, "vertical <(x-<x>)(x'-<x'>)> [m]"],
])
import srwl_bl
v = srwl_bl.srwl_uti_parse_options(varParam, use_sys_argv=True)
source_type, mag = srwl_bl.setup_source(v)
v.wm_na = v.sm_na = 5
# Number of "iterations" per save is best set to num processes
v.wm_ns = v.sm_ns = 2
op = set_optics()
srwl_bl.SRWLBeamline(_name=v.name).calc_all(v, op)
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)))