.. note:: :class: sphx-glr-download-link-note Click :ref:`here ` to download the full example code .. rst-class:: sphx-glr-example-title .. _sphx_glr_cookbook_SRWLIB_Example02.py: SRW Example #2 ************** Problem ------- Calculating electron trajectory in magnetic field of a segmented planar undulator Example Solution ---------------- .. rst-class:: sphx-glr-horizontal * .. image:: /cookbook/images/sphx_glr_SRWLIB_Example02_001.png :class: sphx-glr-multi-img * .. image:: /cookbook/images/sphx_glr_SRWLIB_Example02_002.png :class: sphx-glr-multi-img * .. image:: /cookbook/images/sphx_glr_SRWLIB_Example02_003.png :class: sphx-glr-multi-img * .. image:: /cookbook/images/sphx_glr_SRWLIB_Example02_004.png :class: sphx-glr-multi-img * .. image:: /cookbook/images/sphx_glr_SRWLIB_Example02_005.png :class: sphx-glr-multi-img * .. image:: /cookbook/images/sphx_glr_SRWLIB_Example02_006.png :class: sphx-glr-multi-img .. rst-class:: sphx-glr-script-out Out: .. code-block:: none SRWLIB Python Example # 2: Calculating electron trajectory in magnetic field of a segmented planar undulator with FODO lattice Tabulating Magnetic Field for display ... done Performing calculation ... done Saving trajectory data to a file ... done Plotting the results (close all graph windows to proceed with the script execution) ... done | .. code-block:: default from __future__ import print_function #Python 2.7 compatibility from srwpy.srwlib import * from srwpy.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') .. rst-class:: sphx-glr-timing **Total running time of the script:** ( 0 minutes 1.374 seconds) .. _sphx_glr_download_cookbook_SRWLIB_Example02.py: .. only :: html .. container:: sphx-glr-footer :class: sphx-glr-footer-example .. container:: sphx-glr-download :download:`Download Python source code: SRWLIB_Example02.py ` .. container:: sphx-glr-download :download:`Download Jupyter notebook: SRWLIB_Example02.ipynb ` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_