def propagate_wavefront(wavefront, beamline, output_file=None):
"""
Propagate wavefront and store it in output file.
:param wavefront: Wavefront object or path to HDF5 file
:param beamline: SRWLOptC container of beamline
:param output_file: if parameter present - store propagaed wavefront to file
:return: propagated wavefront object:
"""
if not isinstance(beamline, Beamline):
bl = Beamline(beamline)
else:
bl = beamline
print(bl)
if isinstance(wavefront, Wavefront):
wfr = Wavefront(srwl_wavefront=wavefront._srwl_wf)
else:
print("*****reading wavefront from h5 file...")
wfr = Wavefront()
wfr.load_hdf5(wavefront)
print_mesh(wfr)
print("*****propagating wavefront (with resizing)...")
bl.propagate(wfr)
if output_file is not None:
print("save hdf5:", output_file)
wfr.store_hdf5(output_file)
print("done")
return wfr
def propagate_wavefront(wavefront, beamline, output_file=None):
"""
Propagate wavefront and store it in output file.
:param wavefront: wpg.Wavefront object or path to HDF5 file
:param beamline: SRWLOptC container of beamline
:param output_file: if parameter present - store propagaed wavefront to file
:return: propagated wavefront object
"""
if not isinstance(beamline, Beamline):
bl = Beamline(beamline)
else:
bl = beamline
print(bl)
if isinstance(wavefront, Wavefront):
wfr = Wavefront(srwl_wavefront=wavefront._srwl_wf)
else:
print('*****reading wavefront from h5 file...')
wfr = Wavefront()
wfr.load_hdf5(wavefront)
print_mesh(wfr)
print('*****propagating wavefront (with resizing)...')
bl.propagate(wfr)
if output_file is not None:
print('save hdf5:', output_file)
wfr.store_hdf5(output_file)
print('done')
return wfr
def extract_intensity(fname):
wfr = Wavefront()
wfr.load_hdf5(indir + fname)
np.save(intensity_dir + fname, wfr.get_intensity().sum(-1))
np.save(complex_dir + fname, wfr.as_complex_array().sum(-1))
def load_mode(self, mode_no, indir):
wfr = Wavefront()
wfr.load_hdf5(indir + "/wfr_mode_{}.hdf5".format(mode_no))
return wfr
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'
"""Multi-Electron propagation"""
v.wm = True
v.wm_pl = 'xy'
v.wm_ns = v.sm_ns = 5
"""Single-Electron propagation"""
# 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)
if v.rs_type != 'r':
# wf0 = srwl_bl.SRWLBeamline(_name=v.name, _mag_approx=mag).calc_all(v, op)
srwl_bl.SRWLBeamline(_name=v.name, _mag_approx=mag).calc_all(v, op)
wf0 = v.w_res
# Fraction of wavefront to sample for analysis
Fx = 1/4
Fy = 1/4
""" For multi-electron propagation - wavefront is pickled in srwlib.py line 7188 """
# import analyseWave as aw
# print(" ")
# print("""-----Analysing Wavefield-----""")
# aw.analyseWave('wavefield.pkl')
import pickle
Fx = 1/4
Fy = 1/4
""" For single-electron propagation """
print(" ")
print("""-----Writing Wavefield to pickle file----""")
path = "/home/jerome/WPG/" #"/data/group/vanriessenlab/project/xrnl-srw-wpg-docker/data/maskTest/"
waveName = path + "wavefield_TESTY.pkl" # Save path for wavefield pickle
pathS0 = path + "S0.png" # Save path for Stokes parameter S0 plot
pathS1 = path + "S1.png" # Save path for Stokes parameter S1 plot
pathS2 = path + "S2.png" # Save path for Stokes parameter S2 plot
pathS3 = path + "S3.png" # Save path for Stokes parameter S3 plot
pathD = path + "D.png" # Save path for Degree of Polarisation (D) plot
pathE = path + "E.png" # Save path for Ellipticity (E) plot
pathIn = path + "In.png" # Save path for Inclination (In) plot
pathCS = path + "CS.png" # Save path for Coherence plot from wfStokes()
pathCSL = path + "CSL.png" # Save path for Coherence profiles from wfStokes())
pathX = path + "X.png" # Save path for horizontal cut coherence plot
pathY = path + "Y.png" # Save path for vertical cut coherence plot
pathC = path + "C.png" # Save path for Degree of coherence (U) plot
pathCL = path + "CL.png" # Save path for Coherence line profiles plot
pathI = path + "I.png" # Save path for Intensity (I) plot
pathCm = path + "Cm.png" # Save path for mean coherence plot
pathCmL = path + "CmL.png" # Save path for mean coherence line profile plot
with open(waveName, "wb") as g:
pickle.dump(wf0, g)
print("Wavefield written to: {}".format(waveName))
print(" ")
print("""-----Analysing Wavefield-----""")
aw.analyseWave(waveName,2,1,1,Fx,Fy)#,
# pathS0, pathS1, pathS2, pathS3,
# pathD, pathE, pathIn, pathCS,
# pathCSL, pathX, pathY, pathC,
# pathCL, pathI, pathCm, pathCmL)
''' Add this block and edit as required:
***** start block A *****'''
if v.rs_type == 'r':
if v.wfr_file is None:
print ('Wavefront file not specified')
else:
print('\nSetting up beamline {} with elements:\n{}...'.format(v.name,v.opList))
BL = srwl_blx.beamline(v, _op=op)
print('Loading wavefront from file {}'.format(v.wfr_file) )
wfr = Wavefront()
wfr.load_hdf5(v.wfr_file)
BL.show(wfr._srwl_wf,v)
# Getting phase of initial wavefront and displaying
phase = wfr.get_phase()
newphase = np.squeeze(phase) #phase.reshape((phase.shape[2]*phase.shape[0]),phase.shape[1])
#newphase = newphase.transpose()
plt.imshow(newphase)
import imageio
imageio.imwrite('fileP.tif',phase)
imageio.imwrite('fileI.tif',wfr.get_intensity())
# from extensions.wfutils import modifyDimensions
# modifyDimensions(wfr,
# R=(100e-6,
# 100e-6))
# modifyDimensions(wfr,
# D=(v.op_MaskResolutionX,
# v.op_MaskResolutionY))
# BL.show(wfr._srwl_wf,v)
#
# # Getting phase of wavefront after resizing/resampling
# phase = wfr.get_phase()
# newphase = np.squeeze(phase)
# plt.imshow(newphase)
# for l in disabledList:
# BL.disable(l)
# toggle enabled state of all optical elements (BEWARE!!)
#BL.toggleEnabledState(v.opList)
# implement following (todo)
#BL.disable(optElemType = Screen)
BL.printOE()
print('\nPropagating wavefront through optical elements...')
BL.calc_part(v, wfr=wfr._srwl_wf)
BL.show(wfr._srwl_wf,v)
# Getting phase of wavefront after propagation
phase = wfr.get_phase()
newphase = np.squeeze(phase)
plt.imshow(newphase)
# drift through multiple slices.
# zRange = 200e-6
# zPlanes = 2
# outPath='/opt/WPG/extensions/data_jk/vol_test/'
# dVol=driftVol((0.0,zRange,zPlanes),
# Use_PP(),
# outPath, saveIntensity=True,
# zoomXOutput = 0.5,zoomYOutput = 0.5,
# method=1)
# dVol.propagate(wfr)
''' ***** end block A ******'''
