<jupyter_start><jupyter_text>Import modules<jupyter_code>from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

#Importing necessary modules:
import os
import sys
wpg_path = '/gpfs/exfel/data/user/juncheng/WPG/'
sys.path.insert(0,wpg_path)

import time
import copy
import numpy as np
import pylab as plt

import wpg

#import SRW core functions
from wpg.srwlib import srwl,SRWLOptD,SRWLOptA,SRWLOptC,SRWLOptT,SRWLOptL,SRWLOptMirEl

#Wavefront plots
from wpg.wpg_uti_wf import calculate_fwhm, averaged_intensity, look_at_q_space, plot_t_wf
wpg.__file__
#import SRW helpers functions
from wpg.useful_code.srwutils import AuxTransmAddSurfHeightProfileScaled

#import some helpers functions
from wpg.useful_code.wfrutils import calculate_fwhm_x, plot_wfront, calculate_fwhm_y, print_beamline, get_mesh
from wpg.useful_code.wfrutils import propagate_wavefront

#Import base wavefront class
from wpg import Wavefront
from wpg.beamline import Beamline
from wpg.optical_elements import Empty, Use_PP
#Gaussian beam generator
from wpg.generators import build_gauss_wavefront_xy, build_gauss_wavefront
plt.ion()<jupyter_output><empty_output><jupyter_text>Auxilliary<jupyter_code>#Plotting
def plot_1d(profile, title_fig, title_x, title_y):
    plt.plot(profile[0], profile[1])
    plt.xlabel(title_x)
    plt.ylabel(title_y)
    plt.title(title_fig)
    plt.grid(True)


def plot_2d(amap, xmin, xmax, ymin, ymax, title_fig, title_x, title_y):
    plt.imshow(amap, extent=(ymin, ymax, xmin, xmax))
    plt.colorbar()
    plt.xlabel(title_x)
    plt.ylabel(title_y)
    plt.title(title_fig)
def _resample(wf, axis, data, x0, x1):
    if axis.lower()=='x':
        y = data[int(data.shape[0]/2),:]
        x = np.linspace(wf.params.Mesh.xMin, wf.params.Mesh.xMax, y.shape[0])
    elif axis.lower()=='y':
        y = data[:,int(data.shape[1]/2)]
        x = np.linspace(wf.params.Mesh.yMin, wf.params.Mesh.yMax, y.shape[0])
    else:
        raise ValueError(
            'Wrong axis {}, should be "x" or "y"'.format(axis))

    if not x0 is None:
        xmin = x0
    else:
        xmin = x[0]

    if not x1 is None:
        xmax = x1
    else:
        xmax = x[-1]

    x1 = np.linspace(xmin,xmax,len(y))
    y1 = np.interp(x1, x,y)
    return x1, y1

def intensity_cut(wf, axis, polarization, x0=None, x1=None):

    if polarization.lower()  == 'v' or polarization.lower() == 'vertical':
        pol = 'vertical'
    elif polarization.lower() == 'h' or polarization.lower() == 'horizontal':
        pol = 'horizontal'
    elif polarization.lower() == 't' or polarization.lower() == 'total':
        pol = 'total'
    else:
        raise ValueError(
            'Wrong polarization {}, should be "v" or "vertical"'+
            ' or "h" or "horizontal" or "t" or "total"'.format(polarization))

    data = wf.get_intensity(slice_number=0, polarization=pol)
    return _resample(wf, axis, data, x0, x1)

def phase_cut(wf, axis, polarization, x0=None, x1=None):

    if polarization.lower()  == 'v' or polarization.lower() == 'vertical':
        pol = 'vertical'
    elif polarization.lower() == 'h' or polarization.lower() == 'horizontal':
        pol = 'horizontal'
    else:
        raise ValueError(
            'Wrong polarization {}, should be "v" or "vertical" or "h" or "horizontal"'.format(polarization))

    data = wf.get_phase(slice_number=0, polarization=pol)
    return _resample(wf, axis, data, x0, x1)
#calculate source size from photon energy and FWHM angular divergence
def calculate_source_fwhm(ekev, theta_fwhm):
    wl = 12.39e-10/ekev
    k = 2 * np.sqrt(2*np.log(2))
    theta_sigma = theta_fwhm /k
    sigma0 = wl /(2*np.pi*theta_sigma)
    return sigma0*k
#calculate angular divergence using formula from CDR2011
def calculate_theta_fwhm_cdr(ekev,qnC):
    theta_fwhm = (17.2 - 6.4 * np.sqrt(qnC))*1e-6/ekev**0.85
    return theta_fwhm
#define optical path difference (OPD) from mirror profile, i.e.
#fill the struct opTrErMirr
#input:
#    mdatafile: an ascii file with mirror profile data
#    ncol:      number of columns in the file
#    delim:     delimiter between numbers in an row, can be space (' '), tab '\t', etc
#    Orient:    mirror orientation, 'x' (horizontal) or 'y' (vertical)
#    theta:     incidence angle
#    scale:     scaling factor for the mirror profile
def defineOPD(opTrErMirr, mdatafile, ncol, delim, Orient, theta, scale):
    heightProfData = np.loadtxt(mdatafile).T
    AuxTransmAddSurfHeightProfileScaled(opTrErMirr, heightProfData, Orient, theta, scale)
    plt.figure()
    plot_1d(heightProfData,'profile from ' + mdatafile,'x (m)', 'h (m)') #@todo add the func def in on top of example
def defineEFM(orient,p,q,thetaEFM,theta0,lengthEFM):
    """
    A wrapper to a SRWL function SRWLOptMirEl() for defining a plane elliptical focusing mirror propagator

    :param Orient:    mirror orientation, 'x' (horizontal) or 'y' (vertical)
    :param p:  the distance to two ellipsis centers
    :param q:  the distance to two ellipsis centers
    :param thetaEFM:  the design incidence angle in the center of the mirror
    :param theta0:    the "real" incidence angle in the center of the mirror
    :param lengthEFM: mirror length, [m]
    :return: the struct opEFM
    """
    if orient == 'x':     #horizontal plane ellipsoidal mirror
        opEFM = SRWLOptMirEl(_p=p, _q=q, _ang_graz=thetaEFM, _r_sag=1.e+40, _size_tang=lengthEFM,
                            _nvx=np.cos(theta0), _nvy=0, _nvz=-np.sin(theta0), _tvx=-np.sin(theta0), _tvy=0,
                             _x=0, _y=0, _treat_in_out=1)
    elif orient == 'y': #vertical plane ellipsoidal mirror
        opEFM = SRWLOptMirEl(_p=p, _q=q, _ang_graz=thetaEFM, _r_sag=1.e+40, _size_tang=lengthEFM,
                            _nvx=0, _nvy=np.cos(theta0), _nvz=-np.sin(theta0), _tvx=0, _tvy=-np.sin(theta0),
                             _x=0, _y=0, _treat_in_out=1)
    else:
        raise TypeError('orient should be "x" or "y"')
    return opEFM
# For SimEx output

def calculate_fwhm(wfr):
    """
    Calculate FWHM of the beam calculating number of point bigger then max/2 throuhgt center of the image

    :param wfr:  wavefront
    :return: {'fwhm_x':fwhm_x, 'fwhm_y': fwhm_y} in [m]
    """
    intens = wfr.get_intensity(polarization='total').sum(axis=-1);


    mesh = wfr.params.Mesh
    dx = (mesh.xMax-mesh.xMin)/mesh.nx
    dy = (mesh.yMax-mesh.yMin)/mesh.ny

    x_center = intens[intens.shape[0]//2,:]
    fwhm_x = len(x_center[x_center>x_center.max()/2])*dx

    y_center = intens[:,intens.shape[1]//2]
    fwhm_y = len(y_center[y_center>y_center.max()/2])*dy
    return {'fwhm_x':fwhm_x, 'fwhm_y': fwhm_y}

def add_custom_data(wf, bl0):
    """ Utility to store additional info on the wavefront . """

    # Add more metadata.
    fwhm = calculate_fwhm(wf)


    wf.custom_fields['/misc/xFWHM'] = fwhm['fwhm_x']
    wf.custom_fields['/misc/yFWHM'] = fwhm['fwhm_y']
    wf.custom_fields['/params/beamline/printout'] = str(bl0)

    wf.custom_fields['/info/contact'] = [
        'Name: <NAME>', 'Email: <EMAIL>',
        'Name: <NAME>', 'Email: <EMAIL>']
    wf.custom_fields['/info/data_description'] = 'This dataset contains infromation about wavefront propagated through beamline (WPG and SRW frameworks).'
    wf.custom_fields['/info/method_description'] = """WPG, WaveProperGator (http://github.com/samoylv/WPG)is an interactive simulation framework for coherent X-ray wavefront propagation.\nSRW, Synchrotron Radiation Workshop (http://github.com/ochubar/SRW),  is a physical optics computer code  for simulation of the radiation wavefront propagation through optical systems of beamlines as well as  detailed characteristics of Synchrotron Radiation (SR) generated by relativistic electrons in magnetic fields of arbitrary configuration."""
    wf.custom_fields['/info/package_version'] = '2014.1'<jupyter_output><empty_output><jupyter_text>Initial wavefront<jupyter_code># basic parameters used for defining the beamline and calculating photon source
qnC = 0.1                    # e-bunch charge, [nC]
ekev = 4.0                   # photon energy

thetaOM = 7.e-3    # angle of incidence for the offset mirror
d2m1 = 90.66       # distance to the first offset mirror
d2hkb = 119.75     #distance to the first HFM KB
thetaKB = 7.e-3     # angle of incidence for the KB mirrorsr
dhkb_foc = 3.5     # nominal focal length for HFM KB
dvkb_foc = 2.5      # nominal focal length for VFM KB
dhkb_vkb = dhkb_foc - dvkb_foc       # distance between centers of HFM and VFM
lengthOM=0.8        #length of the offset mirror
lengthKB=0.8        # length of the KB mirror
# #**********************Input Wavefront Structure and Parameters
print('*****defining initial wavefront and writing electric field data to h5-file...')
strInputDataFolder = '../wpg-samples/data_common/'  # input data sub-folder name
strOutputDataFolder = './'  # output data sub-folder name

#init Gauusian beam parameters

z1 = d2m1
theta_fwhm = calculate_theta_fwhm_cdr(ekev,qnC)
k = 2*np.sqrt(2*np.log(2))
sigX = 12.4e-10*k/(ekev*4*np.pi*theta_fwhm)
print('waist_fwhm [um], theta_fwhms [urad]:', sigX*k*1e6, theta_fwhm*1e6)
#define limits
range_xy = theta_fwhm/k*z1*7. # sigma*7 beam size
npoints=400

#define unique filename for storing results
ip = np.floor(ekev)
frac = np.floor((ekev - ip)*1e3)
fname0 = 'g' + str(int(ip))+'_'+str(int(frac))+'kev'
print('save hdf5: '+fname0+'.h5')
ifname = os.path.join(strOutputDataFolder,fname0+'.h5')<jupyter_output>*****defining initial wavefront and writing electric field data to h5-file...
waist_fwhm [um], theta_fwhms [urad]: 29.285785725549683 4.6710057014212625
save hdf5: g4_0kev.h5<jupyter_text>2D Gussian beam parameters<jupyter_code>#tiltX = theta_rms
#build SRW gauusian wavefront

# number of time slices
nslices = 20

# tau – Pulse duration [s]
# Numbers are taken from https://wpg.readthedocs.io/en/latest/beamlines/S1_SPB_CRL_simplified/S1_SPB_CRL_simplified.html
coh_time = 0.8e-15
pulse_duration = coh_time/np.sqrt(2)
# pulse_duration = 15e-15

wfr0=build_gauss_wavefront_xy(nx=npoints, ny=npoints, ekev=ekev,
                              xMin=-range_xy/2, xMax=range_xy/2,
                              yMin=-range_xy/2, yMax=range_xy/2,
                              sigX=sigX, sigY=sigX, d2waist=z1,
                              xoff=0, yoff=0, tiltX=0, tiltY=0)

# wfr0 = build_gauss_wavefront(nx=npoints, ny=npoints, nz=nslices,
#                              ekev = ekev,
#                              xMin = -range_xy/2, xMax = range_xy/2,
#                              yMin = -range_xy/2, yMax = range_xy/2,
#                              tau = pulse_duration,
#                              sigX = sigX, sigY = sigX, d2waist = z1
#                             )

#init WPG Wavefront helper class
mwf = Wavefront(wfr0)

#store wavefront to HDF5 file
mwf.store_hdf5(ifname)
#draw wavefront with common functions
plt.subplot(1,2,1)
plt.imshow(mwf.get_intensity(slice_number=0))
plt.subplot(1,2,2)
plt.imshow(mwf.get_phase(slice_number=0,polarization='horizontal'))
plt.show()

#draw wavefront with cuts
plot_wfront(mwf, title_fig='at '+str(z1)+' m',
            isHlog=False, isVlog=False,
            i_x_min=1e-5, i_y_min=1e-5, orient='x', onePlot=True)

plt.set_cmap('bone') #set color map, 'bone', 'hot', 'jet', etc
fwhm_x = calculate_fwhm_x(mwf)
print('FWHMx [mm], theta_fwhm [urad]:',fwhm_x*1e3,fwhm_x/z1*1e6)
print('*****Defining optical beamline from first offset mirror (M1) to exp hall')

z2 = d2hkb - d2m1

DriftM1_KB = SRWLOptD(z2) #Drift from first offset mirror (M1) to exp hall
horApM1 = lengthOM*thetaOM
opApM1 = SRWLOptA('r', 'a', horApM1, range_xy)  # clear aperture of the Offset Mirror(s)
horApKB = lengthKB * thetaKB # Aperture of the KB system, CA 0.8 m

opApKB = SRWLOptA('r', 'a', horApKB, horApKB)  # clear aperture of the Offset Mirror(s)

#Wavefront Propagation Parameters:
#[0]:  Auto-Resize (1) or not (0) Before propagation
#[1]:  Auto-Resize (1) or not (0) After propagation
#[2]:  Relative Precision for propagation with Auto-Resizing (1. is nominal)
#[3]:  Allow (1) or not (0) for semi-analytical treatment of quadratic phase terms at propagation
#[4]:  Do any Resizing on Fourier side, using FFT, (1) or not (0)
#[5]:  Horizontal Range modification factor at Resizing (1. means no modification)
#[6]:  Horizontal Resolution modification factor at Resizing
#[7]:  Vertical Range modification factor at Resizing
#[8]:  Vertical Resolution modification factor at Resizing
#[9]:  Type of wavefront Shift before Resizing (not yet implemented)
#[10]: New Horizontal wavefront Center position after Shift (not yet implemented)
#[11]: New Vertical wavefront Center position after Shift (not yet implemented)
#                 [ 0] [1] [2]  [3] [4] [5]  [6]  [7]  [8]  [9] [10] [11]
ppM1 =            [ 0,  0, 1.0,  0,  0, 1.0, 1.0, 1.0, 1.0,  0,  0,   0]
ppTrErM1 =        [ 0,  0, 1.0,  0,  0, 1.0, 1.0, 1.0, 1.0,  0,  0,   0]
ppDriftM1_KB =    [ 0,  0, 1.0,  1,  0, 2.4, 1.8, 2.4, 1.8,  0,  0,   0]
ppApKB =          [ 0,  0, 1.0,  0,  0, 0.6, 8.0, 0.6, 4.0,  0,  0,   0]
ppHKB =           [ 0,  0, 1.0,  1,  0, 1.0, 1.0, 1.0, 1.0,  0,  0,   0]
ppTrErHKB =       [ 0,  0, 1.0,  0,  0, 1.0, 1.0, 1.0, 1.0,  0,  0,   0]
ppDrift_HKB_foc = [ 0,  0, 1.0,  1,  0, 1.0, 1.0, 1.0, 1.0,  0,  0,   0]
ppDrift_KB =      [ 0,  0, 1.0,  1,  0, 1.0, 1.0, 1.0, 1.0,  0,  0,   0]
ppVKB =           [ 0,  0, 1.0,  0,  0, 1.0, 1.0, 1.0, 1.0,  0,  0,   0]
ppTrErVKB =       [ 0,  0, 1.0,  0,  0, 1.0, 1.0, 1.0, 1.0,  0,  0,   0]
ppDrift_foc =     [ 0,  0, 1.0,  1,  0, 1.0, 1.0, 1.0, 1.0,  0,  0,   0]
#ppFin  =          [ 0,  0, 1.0,  0,  0, 0.05,5.0, 0.05,5.0,  0,  0,   0]
ppFin =           [ 0,  0, 1.0,  0,  1, .01, 20.0, .01, 20.0,  0,  0,   0]

optBL0 = SRWLOptC([opApM1,  DriftM1_KB],
                    [ppM1,ppDriftM1_KB])

scale = 2     #5 mirror profile scaling factor
print('*****HOM1 data for BL1 beamline ')
opTrErM1 = SRWLOptT(1500, 100, horApM1, range_xy)
#defineOPD(opTrErM1, os.path.join(strInputDataFolder,'mirror1.dat'), 2, '\t', 'x',  thetaOM, scale)
defineOPD(opTrErM1, os.path.join(strInputDataFolder,'mirror2.dat'), 2, ' ', 'x',  thetaOM, scale)
opdTmp=np.array(opTrErM1.arTr)[1::2].reshape(opTrErM1.mesh.ny,opTrErM1.mesh.nx)
plt.figure()
plot_2d(opdTmp, opTrErM1.mesh.xStart*1e3,opTrErM1.mesh.xFin*1e3,opTrErM1.mesh.yStart*1e3,opTrErM1.mesh.yFin*1e3,
        'OPD [m]', 'x (mm)', 'y (mm)')
print('*****Collimating CRL and ideal mirror')
bPlotted = False
isHlog = False
isVlog = False
bSaved = True
optBL = optBL0
strBL = 'bl0'
pos_title = 'at exp hall wall'
print('*****setting-up optical elements, beamline:'+ strBL)

bl = Beamline(optBL)
print(bl)

if bSaved:
    out_file_name = os.path.join(strOutputDataFolder, fname0+'_'+strBL+'_2D.h5')
    print('save hdf5:'+ out_file_name)
else:
    out_file_name = None

startTime = time.time()

# For SimEx output
mwf = Wavefront()
mwf.load_hdf5(ifname)
bl.propagate(mwf)
add_custom_data(mwf,bl)
mwf.store_hdf5(out_file_name)

# mwf = propagate_wavefront(ifname, bl,out_file_name)
print('propagation lasted: {} min'.format(round((time.time() - startTime) / 6.) / 10.))
print('*****Ideal mirror: HOM as an aperture')
plot_wfront(mwf, 'at '+str(z1+z2)+' m',False, False, 1e-5,1e-5,'x', True)
#plt.set_cmap('bone') #set color map, 'bone', 'hot', 'jet', etc
plt.axis('tight')
print('FWHMx [mm], theta_fwhm [urad]:',calculate_fwhm_x(mwf)*1e3,calculate_fwhm_x(mwf)/(z1+z2)*1e6)
print('FWHMy [mm], theta_fwhm [urad]:',calculate_fwhm_y(mwf)*1e3,calculate_fwhm_y(mwf)/(z1+z2)*1e6)<jupyter_output>*****Ideal mirror: HOM as an aperture
FWHMx [mm]: 0.5565926864868518
FWHMy [mm]: 0.5531939904292605
Coordinates of center, [mm]: 0.008376813007327085 0.006515299005698796
stepX, stepY [um]: 1.861514001628267 1.8626060283813484 

R-space
FWHMx [mm], theta_fwhm [urad]: 0.5565926864868518 4.647955628282688
FWHMy [mm], theta_fwhm [urad]: 0.5531939904292605 4.619574032812197
{'fraction_non_alphanumeric': 0.1362215553888472, 'fraction_numerical': 0.04826206551637909, 'mean_word_length': 3.4859786876051597, 'pattern_counts': {'":': 0, '<': 19, '<?xml version=': 0, '>': 21, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '389118', 'license': 'Apache-2.0', 'n_tokens_mistral': 6762, 'n_tokens_neox': 6230, 'n_words': 1541, 'path': 'models/wpg/1_ideal_mirror/1_Ideal_Mirrors_2D.ipynb', 'repo_name': 'LUME-SIMEX/lume-simex-demo-1', 'size': '16043'}