def propagate(self, wfr):
"""
Propagate wavefront through beamline.
:param wfr: Input wavefront (will be re-writed after propagation)
:type wfr: wpg.wavefront.Wavefront
"""
for propagation_option in self.propagation_options:
if all([
isinstance(o, srwlib.SRWLOpt)
for o in propagation_option['optical_elements']
]):
srwl_beamline = srwlib.SRWLOptC(
propagation_option['optical_elements'],
propagation_option['propagation_parameters'])
srwl.PropagElecField(wfr._srwl_wf, srwl_beamline)
# fixing wf._srwl_wf.mesh.zStart bug for Drift
# TODO: Try to fix it with _treat parameter
for opt_element in propagation_option['optical_elements']:
if isinstance(opt_element, srwlib.SRWLOptD):
wfr.params.Mesh.zCoord = wfr.params.Mesh.zCoord + \
opt_element.L
else:
raise ValueError('Unknown type of propagators')
def propagate_sequential(self, wfr, outdir=None):
"""
Propagate sequentially through each optical element in beamline.
:param wfr: Input wavefront (will be re-writed after propagation)
:param outdir: save directory
"""
plotIntensity(wfr)
print(outdir)
if outdir is not None:
if os.path.exists(outdir) == False:
assert (ValueError, "Output Directory Could Not Be Found")
else:
pass
if outdir is not None:
wfr.write(outdir + "initialSource")
plotIntensity(wfr, save=outdir + "source")
for itr in range(len(self.propagation_options[0]['optical_elements'])):
oe = self.propagation_options[0]['optical_elements'][itr]
pp = self.propagation_options[0]['propagation_parameters'][itr]
bl = srwlib.SRWLOptC([oe], [pp])
print(oe.name)
srwl.PropagElecField(wfr._srwl_wf, bl)
if oe.name is None:
oe.name = "el{}".format(itr)
if outdir is not None:
plotIntensity(wfr, save=outdir + oe.name)
with open(outdir + "beamline_parameters.json", 'w') as f:
json.dump(self.params, f)
elif outdir is None:
plotIntensity(wfr)
def propagate_sequential(self, wfr, return_intensity = False, return_mesh = False, savedir = "", checkpoints = []):
"""
Propagate sequentially through each optical element in beamline.
:param wfr: Input wavefront (will be re-writed after propagation)
:param outdir: save directory
"""
if return_intensity:
intensity = []
if return_mesh:
mesh = []
for itr in range(len(self.propagation_options[0]['optical_elements'])):
oe = self.propagation_options[0]['optical_elements'][itr]
pp = self.propagation_options[0]['propagation_parameters'][itr]
bl = srwlib.SRWLOptC([oe],[pp])
try:
print(oe.name)
except:
print(oe)
srwl.PropagElecField(wfr._srwl_wf, bl)
if return_intensity:
intensity.append(wfr.get_intensity().sum(-1))
if return_mesh:
mesh.append(wfr.get_mesh())
plot_intensity_map(wfr)
if return_intensity:
if return_mesh:
return intensity, mesh
else:
return intensity
def set_optics(v=None):
el = []
# 1,2 White_B: aperture 14.0m
el.append(srwlib.SRWLOptA("r", "a", 0.0006, 0.0006, 0.0, 0.0)) # .0004
el.append(srwlib.SRWLOptD(1.0))
# 3,4 HFM: sphericalMirror 15.0m
el.append(
srwlib.SRWLOptMirSph(
_r=8871.45,
_size_tang=0.95,
_size_sag=0.005,
_nvx=0.999902524484,
_nvy=0.0,
_nvz=-0.0139621463437,
_tvx=0.0139621463437,
_tvy=0.0,
_x=0.0,
_y=0.0,
))
el.append(srwlib.SRWLOptD(1.0))
# 5 DCM1: crystal 16.0m
opCr = srwlib.SRWLOptCryst(
_d_sp=3.13557135638,
_psi0r=-1.43534717388e-05,
_psi0i=3.16676479565e-07,
_psi_hr=-7.5926201741e-06,
_psi_hi=2.21095173107e-07,
_psi_hbr=-7.5926201741e-06,
_psi_hbi=2.21095173107e-07,
_tc=0.01,
_ang_as=0.0,
)
# Set crystal orientation:
opCr.set_orient(
-0.970946542231,
6.59748146207e-09,
-0.239296494187,
-0.239296494187,
1.62599496025e-09,
)
el.append(opCr)
# 6,7 DCM2: crystal 16.0m
opCr = srwlib.SRWLOptCryst(
_d_sp=3.13557135638,
_psi0r=-1.43534717388e-05,
_psi0i=3.16676479565e-07,
_psi_hr=-7.5926201741e-06,
_psi_hi=2.21095173107e-07,
_psi_hbr=-7.5926201741e-06,
_psi_hbi=2.21095173107e-07,
_tc=0.01,
_ang_as=0.0,
)
# Set crystal orientation:
opCr.set_orient(
-3.48549712965e-09,
-0.970946542231,
-0.239296494187,
-8.59024886897e-10,
-0.239296494187,
)
el.append(opCr)
el.append(srwlib.SRWLOptD(2.0))
# 8,9 Exit_S: aperture 18.0m
el.append(srwlib.SRWLOptA("r", "a", 0.0006, 0.0006, 0.0, 0.0)) # .00025
el.append(srwlib.SRWLOptD(12.0))
# 10 Sec_Src_Ap: aperture 30.0m
el.append(srwlib.SRWLOptA("r", "a", 2e-05, 2e-05, 0.0, 0.0))
# 11 Sec_SRC_W: watch 30.0m
el.append(srwlib.SRWLOptD(4.99)) # 4.9
# 12 KBM_S: aperture 34.9m
el.append(srwlib.SRWLOptA("r", "a", 2e-05, 2e-05, 0.0, 0.0))
# el.append(srwlib.SRWLOptA("r", "a", 1000e-05, 1000e-05, 0.0, 0.0))
# 13 KBM_S_W: watch 34.9m
el.append(srwlib.SRWLOptD(0.01)) # 0.1
# 14, 15 KBMV: ellipsoidMirror 35.0m
el.append(
srwlib.SRWLOptMirEl(
_p=35.0,
_q=0.47,
_ang_graz=0.003,
_size_tang=0.22,
_size_sag=0.22,
_nvx=0.999995500003,
_nvy=0.0,
_nvz=-0.0029999955,
_tvx=0.0029999955,
_tvy=0.0,
_x=0.0,
_y=0.0,
))
el.append(srwlib.SRWLOptD(0.22))
# 16,17 KBMH: ellipsoidMirror 35.22m
el.append(
srwlib.SRWLOptMirEl(
_p=5.22,
_q=0.25,
_ang_graz=0.003,
_size_tang=0.22,
_size_sag=0.22,
_nvx=0.0,
_nvy=0.999995500003,
_nvz=-0.0029999955,
_tvx=0.0,
_tvy=0.0029999955,
_x=0.0,
_y=0.0,
))
el.append(srwlib.SRWLOptD(0.25))
return srwlib.SRWLOptC(el)
def set_optics(mirror_profile=False):
el = []
el.append(srwlib.SRWLOptD(0.71))
# M2: ellipsoidMirror 170.0m
el.append(
srwlib.SRWLOptMirEl(_p=170.0,
_q=1.6,
_ang_graz=0.01385,
_size_tang=0.184,
_size_sag=0.01,
_nvx=-0.999904090283,
_nvy=0.0,
_nvz=-0.013849557214,
_tvx=-0.013849557214,
_tvy=0.0,
_x=0.0,
_y=0.0))
ifnElMirror1 = mirror_profile
if ifnElMirror1:
assert os.path.isfile(
ifnElMirror1
), "Missing input file mirror_2d.dat, required by M2 beamline element"
hProfDataElMirror1 = srwlib.srwl_uti_read_data_cols(ifnElMirror1, "\t")
el.append(
srwlib.srwl_opt_setup_surf_height_2d(hProfDataElMirror1,
_dim="x",
_ang=0.01385,
_amp_coef=1.0))
el.append(srwlib.SRWLOptD(0.5))
# M3: ellipsoidMirror 170.5m
el.append(
srwlib.SRWLOptMirEl(_p=170.5,
_q=1.1,
_ang_graz=0.01385,
_size_tang=0.184,
_size_sag=0.01,
_nvx=0.0,
_nvy=0.999904090283,
_nvz=-0.013849557214,
_tvx=0.0,
_tvy=0.013849557214,
_x=0.0,
_y=0.0))
ifnElMirror2 = mirror_profile
if ifnElMirror2:
assert os.path.isfile(
ifnElMirror2
), "Missing input file mirror_2d.dat, required by M3 beamline element"
hProfDataElMirror2 = srwlib.srwl_uti_read_data_cols(ifnElMirror2, "\t")
el.append(
srwlib.srwl_opt_setup_surf_height_2d(hProfDataElMirror2,
_dim="y",
_ang=0.01385,
_amp_coef=1.0))
el.append(srwlib.SRWLOptD(1.1))
# Sample: watch 171.6m
pp = []
pp.append([
0, 0, 1.0, 1, 0, 1.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
0.0
])
# M2
pp.append([
0, 0, 1.0, 1, 0, 1.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
0.0
])
if ifnElMirror1:
pp.append([0, 0, 1.0, 0, 0, 1.0, 1.0, 1.0, 1.0])
pp.append([
0, 0, 1.0, 1, 0, 1.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
0.0
])
# M3
pp.append([
0, 0, 1.0, 1, 0, 1.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
0.0
])
if ifnElMirror2:
pp.append([0, 0, 1.0, 0, 0, 1.0, 1.0, 1.0, 1.0])
pp.append([
0, 0, 1.0, 1, 0, 1.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
0.0
])
# Sample
# final post-propagation
pp.append([
0, 0, 1.0, 0, 0, 0.05, 20.0, 0.05, 20.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
0.0, 0.0
])
return srwlib.SRWLOptC(el, pp)