diff --git a/docs/source/_cookbook/SRWLIB_Example02.py b/docs/source/_cookbook/SRWLIB_Example02.py new file mode 100644 index 0000000..5dba6a5 --- /dev/null +++ b/docs/source/_cookbook/SRWLIB_Example02.py @@ -0,0 +1,205 @@ +""" +SRW Example #2 +************** + +Problem +------- + +Calculating electron trajectory in magnetic field of a segmented planar undulator + +Example Solution +---------------- +""" + +from __future__ import print_function #Python 2.7 compatibility +from srwlib import * +from uti_plot import * +import os + +print('SRWLIB Python Example # 2:') +print('Calculating electron trajectory in magnetic field of a segmented planar undulator with FODO lattice') + +#**********************Input Parameters +#**********************Folder and File names +strExDataFolderName = 'data_example_02' #example data sub-folder name +strFldOutFileName = 'ex02_res_fld.dat' #file name for output (tabulated) magnetic field data +strTrajOutFileName = 'ex02_res_traj.dat' #file name for output trajectory data + +#**********************Initial Conditions for Particle Trajectory Calculation +part = SRWLParticle() +part.x = 0.0001 #Initial Transverse Coordinates (initial Longitudinal Coordinate will be defined later on) [m] +part.y = 0.0001 +part.xp = 0 #Initial Transverse Velocities +part.yp = 0 +part.gamma = 3/0.51099890221e-03 #Relative Energy +part.relE0 = 1 #Electron Rest Mass +part.nq = -1 #Electron Charge + +npTraj = 20001 #Number of Points for Trajectory calculation + +arPrecPar = [1] #General Precision parameters for Trajectory calculation: +#[0]: integration method No: + #1- fourth-order Runge-Kutta (precision is driven by number of points) + #2- fifth-order Runge-Kutta +#[1],[2],[3],[4],[5]: absolute precision values for X[m],X'[rad],Y[m],Y'[rad],Z[m] (yet to be tested!!) - to be taken into account only for R-K fifth order or higher +#[6]: tolerance (default = 1) for R-K fifth order or higher +#[7]: max. number of auto-steps for R-K fifth order or higher (default = 5000) + +#**********************Magnetic Field +numSegm = 5 #Number of ID Segments +numPer = 100 #Number of Periods in one Segment (without counting for terminations) +undPer = 0.02 #Period Length [m] +xcID = 0 #Transverse Coordinates of ID Center [m] +ycID = 0 +zcID = 0 #Longitudinal Coordinate of ID Center [m] + +und = SRWLMagFldU([SRWLMagFldH(1, 'v', 1.05, 0, 1)], undPer, numPer) #Undulator Segment +qf = SRWLMagFldM(0.5, 2, 'n', 0.2) #Focusing Quad +qd = SRWLMagFldM(-0.5, 2, 'n', 0.2) #Defocusing Quad + +arZero = array('d', [0]*9) +undLen = (numPer + 2)*undPer +distBwSegm = 0.4 #Distance between Undulator Segments +undLenExt = undLen + distBwSegm +arZc = array('d', [-2*undLenExt, -1.5*undLenExt, -undLenExt, -0.5*undLenExt, 0, 0.5*undLenExt, undLenExt, 1.5*undLenExt, 2*undLenExt]) +magFldCnt = SRWLMagFldC([und, qf, und, qd, und, qf, und, qd, und], arZero, arZero, arZc) #Container of all Field Elements + +part.z = arZc[0] - 0.5*undLenExt #Initial Longitudinal Coordinate (set before the ID) + +print(' Tabulating Magnetic Field for display ... ', end='') +dispBy = array('d', [0]*npTraj) +dispMagFld3D = SRWLMagFld3D(_arBy=dispBy, _nx=1, _ny=1, _nz=npTraj, _rz=5*undLenExt) +xDisp = part.x +yDisp = part.y +dispMagFldCnt = SRWLMagFldC([dispMagFld3D], array('d', [xDisp]), array('d', [yDisp]), array('d', [0])) +srwl.CalcMagnField(dispMagFldCnt, magFldCnt) +print('done') + +#**********************Trajectory structure, where the results will be stored +partTraj = SRWLPrtTrj() +partTraj.partInitCond = part +#partTraj.allocate(npTraj) +partTraj.allocate(npTraj, True) #True ensures that field along trajectory will also be extracted +partTraj.ctStart = 0 #"Start Time" (c*t) for the calculation (0 corresponds to the time moment for which the initial conditions are defined) +partTraj.ctEnd = partTraj.ctStart + 5*undLenExt #End Time + +#**********************Trajectory Calculation (SRWLIB function call) +print(' Performing calculation ... ', end='') +partTraj = srwl.CalcPartTraj(partTraj, magFldCnt, arPrecPar) +print('done') + +#**********************Saving results +#Auxiliary function to write tabulated Magnetic Field data to ASCII file: +def AuxSaveMagFldData(dispFld3D, xc, yc, zc, filePath): + + hx = 0 if(dispFld3D.nx <= 1) else dispFld3D.rx/(dispFld3D.nx - 1) + hy = 0 if(dispFld3D.ny <= 1) else dispFld3D.ry/(dispFld3D.ny - 1) + hz = 0 if(dispFld3D.nz <= 1) else dispFld3D.rz/(dispFld3D.nz - 1) + + arXdef = False + if((dispFld3D.arX != None) and (len(dispFld3D.arX) == dispFld3D.nx)): + arXdef = True + arYdef = False + if((dispFld3D.arY != None) and (len(dispFld3D.arY) == dispFld3D.ny)): + arYdef = True + arZdef = False + if((dispFld3D.arZ != None) and (len(dispFld3D.arZ) == dispFld3D.nz)): + arZdef = True + + nTot = dispFld3D.nx*dispFld3D.ny*dispFld3D.nz + arBxDef = False + if((dispFld3D.arBx != None) and (len(dispFld3D.arBx) == nTot)): + arBxDef = True + arByDef = False + if((dispFld3D.arBy != None) and (len(dispFld3D.arBy) == nTot)): + arByDef = True + arBzDef = False + if((dispFld3D.arBz != None) and (len(dispFld3D.arBz) == nTot)): + arBzDef = True + + f = open(filePath, 'w') + resStr = '#X [m], Y [m], Z [m]' + if(arBxDef == True): + resStr += ', Bx [T]' + if(arByDef == True): + resStr += ', By [T]' + if(arBzDef == True): + resStr += ', Bz [T]' + f.write(resStr + '\n') + + i = 0 + z = zc - 0.5*dispFld3D.rz + for iz in range(dispFld3D.nz): + if(arZdef == True): + z = dispFld3D.arZ[iz] + zc + #print(z) + y = yc - 0.5*dispFld3D.ry + for iy in range(dispFld3D.ny): + if(arYdef == True): + y = dispFld3D.arY[iy] + yc + x = xc - 0.5*dispFld3D.rx + for ix in range(dispFld3D.nx): + if(arXdef == True): + x = dispFld3D.arX[ix] + xc + resStr = repr(x) + '\t' + repr(y) + '\t' + repr(z) + if(arBxDef == True): + resStr += '\t' + repr(dispFld3D.arBx[i]) + if(arByDef == True): + resStr += '\t' + repr(dispFld3D.arBy[i]) + if(arBzDef == True): + resStr += '\t' + repr(dispFld3D.arBz[i]) + f.write(resStr + '\n') + i += 1 + x += hx + y += hy + z += hz + f.close() + +#Auxiliary function to write tabulated resulting Trajectory data to ASCII file: +def AuxSaveTrajData(traj, filePath): + f = open(filePath, 'w') + resStr = '#ct [m], X [m], BetaX [rad], Y [m], BetaY [rad], Z [m], BetaZ [rad]' + if(hasattr(traj, 'arBx')): + resStr += ', Bx [T]' + if(hasattr(traj, 'arBy')): + resStr += ', By [T]' + if(hasattr(traj, 'arBz')): + resStr += ', Bz [T]' + f.write(resStr + '\n') + ctStep = 0 + if traj.np > 0: + ctStep = (traj.ctEnd - traj.ctStart)/(traj.np - 1) + ct = traj.ctStart + for i in range(traj.np): + resStr = str(ct) + '\t' + repr(traj.arX[i]) + '\t' + repr(traj.arXp[i]) + '\t' + repr(traj.arY[i]) + '\t' + repr(traj.arYp[i]) + '\t' + repr(traj.arZ[i]) + '\t' + repr(traj.arZp[i]) + if(hasattr(traj, 'arBx')): + resStr += '\t' + repr(traj.arBx[i]) + if(hasattr(traj, 'arBy')): + resStr += '\t' + repr(traj.arBy[i]) + if(hasattr(traj, 'arBz')): + resStr += '\t' + repr(traj.arBz[i]) + f.write(resStr + '\n') + ct += ctStep + f.close() + +print(' Saving trajectory data to a file ... ', end='') +AuxSaveMagFldData(dispMagFld3D, 0, 0, 0, os.path.join(os.getcwd(), strExDataFolderName, strFldOutFileName)) +AuxSaveTrajData(partTraj, os.path.join(os.getcwd(), strExDataFolderName, strTrajOutFileName)) +print('done') + +#**********************Plotting results +print(' Plotting the results (close all graph windows to proceed with the script execution) ... ', end='') +ctMesh = [partTraj.ctStart, partTraj.ctEnd, partTraj.np] +for i in range(partTraj.np): #converting from [m] to [mm] and from [rad] to [mrad] + partTraj.arXp[i] *= 1000 + partTraj.arX[i] *= 1000 + partTraj.arYp[i] *= 1000 + partTraj.arY[i] *= 1000 +uti_plot1d(partTraj.arBy, ctMesh, ['ct [m]', 'Vertical Magnetic Field [T]']) +uti_plot1d(partTraj.arXp, ctMesh, ['ct [m]', 'Horizontal Angle [mrad]']) +uti_plot1d(partTraj.arX, ctMesh, ['ct [m]', 'Horizontal Position [mm]']) +uti_plot1d(partTraj.arBx, ctMesh, ['ct [m]', 'Horizontal Magnetic Field [T]']) +uti_plot1d(partTraj.arYp, ctMesh, ['ct [m]', 'Vertical Angle [mrad]']) +uti_plot1d(partTraj.arY, ctMesh, ['ct [m]', 'Vertical Position [mm]']) +uti_plot_show() +print('done') diff --git a/docs/source/_cookbook/data_example_02/.gitignore b/docs/source/_cookbook/data_example_02/.gitignore new file mode 100644 index 0000000..e69de29