Radia_Ex-kickmap.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Undulator focusing and generation of a map file suitable for electron tracking\n",
    "\n",
    "\n",
    "## _Import Radia and other packages_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import numpy as np\n",
    "import scipy.constants as sc\n",
    "import matplotlib.pyplot as plt\n",
    "import time as tm\n",
    "from math import *\n",
    "from copy import *\n",
    "from array import array\n",
    "\n",
    "import radia as rad\n",
    "import ipywidgets\n",
    "from jupyter_rs_radia import radia_viewer\n",
    "\n",
    "# import an illustration of this magnet\n",
    "from IPython.display import Image\n",
    "Image(filename=('./U46HybridUnd.png'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## _Define function to build the undulator_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def undulator(lp, mp, np, cp,  lm, mm, nm, cm,  gap, gapOffset, numPer):\n",
    "    \"\"\"\n",
    "    create hybrid undulator magnet\n",
    "    arguments:\n",
    "      lp = [lpx, lpy, lpz] = dimensions of the iron poles / mm\n",
    "      mp = magnetic properties of the iron poles (M-H curve)\n",
    "      np = segmentation of the iron poles\n",
    "      cp = [r,g,b] = color for the iron poles\n",
    "      lm = [lmx, lmy, lmz] = dimensions of the magnet blocks / mm\n",
    "      mm = magnetic properties of the magnet blocks (remanent magnetization)\n",
    "      nm = segmentation of the magnet blocks\n",
    "      cm = [r,g,b] = color for the magnet blocks\n",
    "      gap = undulator gap / mm\n",
    "      gapOffset = vertical offset / mm of the magnet blocks w/rt the poles\n",
    "      numPer = number of full periods of the undulator magnetic field\n",
    "    return: Radia representations of\n",
    "      undulator group, poles, permanent magnets\n",
    "    \"\"\"\n",
    "    zer = [0,0,0]\n",
    "\n",
    "    # full magnet will be assembled into this Radia group\n",
    "    grp = rad.ObjCnt([])\n",
    "\n",
    "    # principal poles and magnet blocks in octant(+,+,–)\n",
    "    # -- half pole\n",
    "    y = lp[1]/4\n",
    "    pole = rad.ObjFullMag([lp[0]/4,    y,    -lp[2]/2 - gap/2],\n",
    "                          [lp[0]/2, lp[1]/2,  lp[2]          ], zer, np, grp, mp, cp)\n",
    "    y += lp[1]/4\n",
    "    # -- magnet and pole pairs\n",
    "    mDir = -1\n",
    "    for i in range(0, numPer):\n",
    "        initM = [0, mDir, 0]; mDir *= -1\n",
    "        y += lm[1]/2\n",
    "        magnet = rad.ObjFullMag([lm[0]/4,   y,   -lm[2]/2 - gap/2 - gapOffset],\n",
    "                                [lm[0]/2, lm[1],  lm[2]                      ], initM, nm, grp, mm, cm)\n",
    "        y += (lm[1] + lp[1])/2\n",
    "        pole = rad.ObjFullMag([lp[0]/4,   y,   -lp[2]/2 - gap/2],\n",
    "                              [lp[0]/2, lp[1],  lp[2]          ], zer, np, grp, mp, cp)\n",
    "        y += lp[1]/2\n",
    "    # -- end magnet block\n",
    "    initM = [0, mDir, 0]\n",
    "    y += lm[1]/4;\n",
    "    magnet = rad.ObjFullMag([lm[0]/4,   y,     -lm[2]/2 - gap/2 - gapOffset],\n",
    "                            [lm[0]/2, lm[1]/2,  lm[2]                      ], initM, nm, grp, mm, cm)\n",
    "\n",
    "    # use mirror symmetry to define the full undulator\n",
    "    rad.TrfZerPerp(grp, zer, [1,0,0]) # reflect in the (y,z) plane\n",
    "    rad.TrfZerPara(grp, zer, [0,0,1]) # reflect in the (x,y) plane\n",
    "    rad.TrfZerPerp(grp, zer, [0,1,0]) # reflect in the (z,x) plane\n",
    "    \n",
    "    return grp, pole, magnet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# From RadiaToTrack.m\n",
    "#hybrid\n",
    "def hyb(p,per,gap,nper,gapoffset, lp,np,matp,cp, lm,nm,matm,cm, lms,nms,matms,cms):\n",
    "    zer = [0,0,0]\n",
    "    inn = -1\n",
    "    grp = rad.ObjCnt([])\n",
    "\n",
    "    yy = lp[1]/4\n",
    "    rad.ObjFullMag(p + [lp[0]/4, yy, -lp[2]/2 - gap/2], [lp[0]/2, lp[1]/2, lp[2]], zer, np, grp, matp, cp)\n",
    "\n",
    "    yys = yy + (lms[1] - lp[1])/4\n",
    "    rad.ObjFullMag(p + [lp[0]/2 + lms[0]/2, yys, -lms[2]/2 - gap/2 - gapoffset], [lms[0], lms[1]/2, lms[2]],[inn, 0, 0], nms, grp, matms, cms)\n",
    "\n",
    "    yy += lp[1]/4\n",
    "    for i in range(1,nper+1):\n",
    "        initm = [0, (i + 1)%2 - i%2, 0]\n",
    "        yy += lm[1]/2\n",
    "        rad.ObjFullMag(p + [lm[0]/4, yy, -lm[2]/2 - gap/2 - gapoffset], [lm[0]/2, lm[1], lm[2]], initm, nm, grp, matm, cm)\n",
    "        yy += (lm[1] + lp[1])/2;\n",
    "        rad.ObjFullMag(p + [lp[0]/4, yy, -lp[2]/2 - gap/2], [lp[0]/2, lp[1], lp[2]], zer, np, grp, matp, cp)\n",
    "        inn *= -1\n",
    "        rad.ObjFullMag(p + [lp[0]/2 + lms[0]/2, yy, -lms[2]/2 - gap/2 - gapoffset], [lms[0], lms[1], lms[2]], [inn, 0, 0], nms, grp, matms, cms)\n",
    "  \n",
    "        yy += lp[1]/2\n",
    "    initm = [0, nper%2 - (nper + 1)%2, 0]\n",
    "    yy += lm[1]/4\n",
    "    rad.ObjFullMag(p + [lm[0]/4, yy, -lm[2]/2 - gap/2 - gapoffset], [lm[0]/2, lm[1]/2, lm[2]], initm, nm, grp, matm, cm)\n",
    "\n",
    "    rad.TrfZerPerp(grp, p, [1,0,0]) # reflect in the (y,z) plane\n",
    "    rad.TrfZerPara(grp, p, [0,0,1]) # reflect in the (x,y) plane\n",
    "    rad.TrfZerPerp(grp, p, [0,1,0]) # reflect in the (z,x) plane\n",
    "    \n",
    "    return grp\n",
    "    \n",
    "# Pure Permanent Magnet\n",
    "def undparts(po, wv, wh, nnp, per, br, si, axe=0.):\n",
    "    g = rad.ObjCnt([])\n",
    "    p = po - [0, nnp*per/2, 0]\n",
    "    for i in range(0,4*nnp + 1):\n",
    "        if i == 0 or i == 4*nnp: s = 0.5\n",
    "        else: s = 1.\n",
    "        if i%2 == 0: w = wv\n",
    "        else: w = wh\n",
    "        t = -(i - 1)*np.pi/2*si\n",
    "        m = np.array([np.sin(axe)*np.sin(t), np.cos(t), np.cos(axe)*np.sin(t)])*br*s\n",
    "        ma = rad.ObjRecMag(p, w, m)\n",
    "        rad.ObjAddToCnt(g, [ma])\n",
    "        p = p + [0, per/4, 0]\n",
    "    rad.ObjDrwAtr(g, [0, 0, 1])\n",
    "    return g\n",
    "\n",
    "# \"Apple II\" type undulator\n",
    "def apple2G(pos, per, gap, gapx, plr, pll, pur, pul, lx, lz, airgap, br, nper):\n",
    "    wv = [lx/2, per/4 - airgap, lz]\n",
    "    wh = wv\n",
    "    px = lx/4 + gapx/2\n",
    "    pz = gap/2 + lz/2\n",
    "    \n",
    "    g1 = undparts(pos + [px, pur, pz], wv, wh, nper, per, br, si=1)\n",
    "    g2 = undparts(pos + [-px, pul, pz], wv, wh, nper, per, br, si=1)\n",
    "    g3 = undparts(pos + [px, plr, -pz], wv, wh, nper, per, -br, si=-1)\n",
    "    g4 = undparts(pos + [-px, pll, -pz], wv, wh, nper, per, -br, si=-1)\n",
    "\n",
    "    g = rad.ObjCnt([g1, g2, g3, g4])\n",
    "    return g\n",
    "\n",
    "def apple2(pos, per, gap, gapx, phase, lx, lz, airgap, br, nper):\n",
    "    wv = [lx/2, per/4 - airgap, lz]\n",
    "    wh = wv\n",
    "    px = lx/4 + gapx/2\n",
    "    pz = gap/2 + lz/2\n",
    "    \n",
    "    g1 = undparts(pos + [px, phase/2, pz], wv, wh, nper, per, br, 1)\n",
    "    g2 = undparts(pos + [-px, -phase/2, pz], wv, wh, nper, per, br, 1)\n",
    "    g3 = undparts(pos + [px, -phase/2, -pz], wv, wh, nper, per, -br, -1)\n",
    "    g4 = undparts(pos + [-px, phase/2, -pz], wv, wh, nper, per, -br, -1)\n",
    "    g = rad.ObjCnt([g1, g2, g3, g4])\n",
    "    return g\n",
    "\n",
    "#Spring8\n",
    "def sp8(pos, per, gap, gapx, phase, lxc, lzc, colc, lxs, lzs, cols, airgap, br, nper):\n",
    "    wc = [lxc, per/4 - airgap, lzc]\n",
    "    px = 0\n",
    "    pz = gap/2 + lzc/2\n",
    "    g1 = undparts(pos + [px, -phase/2, pz], wc, wc, nper, per, br, 1)\n",
    "    rad.ObjDrwAtr(g1, colc)\n",
    "    g2 = undparts(pos + [px, -phase/2, -pz], wc, wc, nper, per, -br, -1)\n",
    "    rad.ObjDrwAtr(g2, colc)\n",
    "    wc = [lxs, per/4 - airgap, lzs]\n",
    "    px = lxc/2 + gapx + lxs/2\n",
    "    pz = gap/2 + lzs/2\n",
    "    \n",
    "    g3 = undparts(pos + [px, phase/2, pz], wc, wc, nper, per, br, 1)\n",
    "    rad.ObjDrwAtr(g3, cols)\n",
    "    g4 = undparts(pos + [px, phase/2, -pz], wc, wc, nper, per, br, -1)\n",
    "    rad.ObjDrwAtr(g4, cols)\n",
    "    g5 = undparts(pos + [-px, phase/2, pz], wc, wc, nper, per, -br, 1)\n",
    "    rad.ObjDrwAtr(g5, cols)\n",
    "    g6 = undparts(pos + [-px, phase/2, -pz], wc, wc, nper, per, -br, -1)\n",
    "    rad.ObjDrwAtr(g6, cols)\n",
    "    \n",
    "    g = rad.ObjCnt([g1, g2, g3, g4, g5, g6])\n",
    "    return g"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# From RADIA_APPLE_II_Demo.py\n",
    "def MagnetBlock(_pc, _wc, _cx, _cz, _type, _ndiv, _m):\n",
    "\n",
    "    u = rad.ObjCnt([])\n",
    "    \n",
    "    wwc = _wc\n",
    "    if(_type!=0):\n",
    "        wwc = copy(_wc)\n",
    "        wwc[0] -= 2*_cx\n",
    "        \n",
    "    b1 = rad.ObjRecMag(_pc, wwc, _m)\n",
    "    rad.ObjAddToCnt(u, [b1])\n",
    "    #ndiv2 = [1,_ndiv[1],_ndiv[2]]\n",
    "\n",
    "    if((_cx>0.01) and (_cz>0.01)):\n",
    "        if(_type==1):\n",
    "            ppc = [_pc[0]-_wc[0]/2+_cx/2,_pc[1],_pc[2]-_cz/2]\n",
    "            wwc = [_cx,_wc[1],_wc[2]-_cz]\n",
    "            b2 = rad.ObjRecMag(ppc, wwc, _m)\n",
    "\n",
    "            ppc = [_pc[0]+_wc[0]/2-_cx/2,_pc[1],_pc[2]+_cz/2]\n",
    "            wwc = [_cx,_wc[1],_wc[2]-_cz]\n",
    "            b3 = rad.ObjRecMag(ppc, wwc, _m)\n",
    "            rad.ObjAddToCnt(u, [b2,b3])\n",
    "\n",
    "        elif(_type==2):\n",
    "            ppc = [_pc[0]-_wc[0]/2+_cx/2,_pc[1],_pc[2]+_cz/2]\n",
    "            wwc = [_cx,_wc[1],_wc[2]-_cz]\n",
    "            b2 = rad.ObjRecMag(ppc, wwc, _m)\n",
    "\n",
    "            ppc = [_pc[0]+_wc[0]/2-_cx/2,_pc[1],_pc[2]-_cz/2]\n",
    "            wwc = [_cx,_wc[1],_wc[2]-_cz]\n",
    "            b3 = rad.ObjRecMag(ppc, wwc, _m)\n",
    "            rad.ObjAddToCnt(u, [b2,b3])\n",
    "\n",
    "        elif(_type==3):\n",
    "            ppc = [_pc[0]-_wc[0]/2+_cx/2,_pc[1],_pc[2]]\n",
    "            wwc = [_cx,_wc[1],_wc[2]-2*_cz]\n",
    "            b2 = rad.ObjRecMag(ppc, wwc, _m)\n",
    "\n",
    "            ppc = [_pc[0]+_wc[0]/2-_cx/2,_pc[1],_pc[2]]\n",
    "            wwc = [_cx,_wc[1],_wc[2]-2*_cz]\n",
    "            b3 = rad.ObjRecMag(ppc, wwc, _m)\n",
    "            rad.ObjAddToCnt(u, [b2,b3])\n",
    "\n",
    "    rad.ObjDivMag(u, _ndiv, 'Frame->LabTot')\n",
    "    return u\n",
    "\n",
    "#*************Magnet Array\n",
    "def MagnetArray(_per, _nper, _po, _w, _si, _type, _cx, _cz, _br, _mu, _ndiv, _bs1, _s1, _bs2, _s2, _bs3, _s3, _bs2dz=0, _qp_ind_mag=None, _qp_dz=0):\n",
    "\n",
    "    u = rad.ObjCnt([])\n",
    "\n",
    "    Le = _bs1+_s1+_bs2+_s2+_bs3+_s3\n",
    "    Lc = (_nper+0.25)*_per\n",
    "    p = [_po[0],_po[1]-(Lc/2+Le),_po[2]] #po-{0,(Lc/2+Le),0}\n",
    "\n",
    "    nMagTot = 4*_nper+7\n",
    "    iMagCen = int(nMagTot/2.) #0-based index of the central magnet\n",
    "    #print('iMagCen =', iMagCen) #DEBUG\n",
    "\n",
    "    QP_IsDef = False; QP_DispIsConst = True\n",
    "    nQP_Disp = 0\n",
    "    if(_qp_ind_mag is not None):\n",
    "        if(isinstance(_qp_ind_mag, list) or isinstance(_qp_ind_mag, array)):\n",
    "            nQP_Disp = len(_qp_ind_mag)\n",
    "            if(nQP_Disp > 0): QP_IsDef = True\n",
    "        if(isinstance(_qp_dz, list) or isinstance(_qp_dz, array)): QP_DispIsConst = False\n",
    "        elif(_qp_dz==0): QP_IsDef = False\n",
    "\n",
    "    for i in range(nMagTot):\n",
    "        wc = copy(_w)\n",
    "        \n",
    "        if(i==0):\n",
    "            p[1] += _bs1/2\n",
    "            wc[1] = _bs1\n",
    "        elif(i==1):\n",
    "            p[1] += _bs1/2+_s1+_bs2/2\n",
    "            wc[1] = _bs2\n",
    "        elif(i==2):\n",
    "            p[1] += _bs2/2+_s2+_bs3/2\n",
    "            wc[1] = _bs3\n",
    "        elif(i==3):\n",
    "            p[1] += _bs3/2+_s3+_per/8\n",
    "        elif((i>3) and (i<4*_nper+4)):\n",
    "            p[1] += _per/4\n",
    "        elif(i==4*_nper+4):\n",
    "            p[1] += _per/8+_s3+_bs3/2\n",
    "            wc[1] = _bs3\n",
    "        elif(i==4*_nper+5):\n",
    "            p[1] += _bs3/2+_s2+_bs2/2\n",
    "            wc[1] = _bs2\n",
    "        elif(i==4*_nper+6):\n",
    "            p[1] += _bs2/2+_s1+_bs1/2\n",
    "            wc[1] = _bs1\n",
    "\n",
    "        pc = copy(p)\n",
    "\n",
    "        if((i==1) or (i==4*_nper+5)):\n",
    "            if(_si==1): pc[2] += _bs2dz\n",
    "            else: pc[2] -= _bs2dz\n",
    "\n",
    "        if(QP_IsDef):\n",
    "            for iQP in range(nQP_Disp):\n",
    "                if(i == _qp_ind_mag[iQP] + iMagCen):\n",
    "                    qpdz = _qp_dz\n",
    "                    if(not QP_DispIsConst): qpdz = _qp_dz[iQP]\n",
    "                    pc[2] += qpdz\n",
    "                    #print('Abs. Ind. of Mag. to be Displaced:', i) #DEBUG\n",
    "                    break\n",
    "\n",
    "        t = -i*pi/2*_si\n",
    "        mcol = [0.0,cos(t),sin(t)]\n",
    "        m = [mcol[0],mcol[1]*_br,mcol[2]*_br]\n",
    "\n",
    "        ma = MagnetBlock(pc, wc, _cx, _cz, _type, _ndiv, m)\n",
    "\n",
    "        mcol = [0.27, 0.9*abs(mcol[1]), 0.9*abs(mcol[2])]\n",
    "        rad.ObjDrwAtr(ma, mcol, 0.0001)\n",
    "\n",
    "        rad.ObjAddToCnt(u, [ma])\n",
    "\n",
    "    mat = rad.MatLin(_mu, abs(_br))\n",
    "    rad.MatApl(u, mat)\n",
    "\n",
    "    return u \n",
    "\n",
    "#*************Undulator\n",
    "def APPLE_II(_per, _nper, _gap, _gapx, _phase, _phase_type, _lx, _lz, _cx, _cz, _air, _br, _mu, _ndiv, _bs1, _s1, _bs2, _s2, _bs3, _s3, _bs2dz, _qp_ind_mag, _qp_dz, _use_sym=False):\n",
    "\n",
    "    w = [_lx,_per/4-_air,_lz]\n",
    "    px = _lx/2+_gapx/2;\n",
    "    pz = _gap/2+_lz/2;\n",
    "\n",
    "    p1 = 0; p2 = _phase; p3 = 0; p4 = _phase\n",
    "    if(_phase_type < 0): p2 = -_phase\n",
    "\n",
    "    #print('w =', w)\n",
    "    g1 = MagnetArray(_per, _nper, _po=[px,p1,pz], _w=w, _si=1, _type=1, _cx=_cx, _cz=_cz, _br=_br, _mu=_mu, _ndiv=_ndiv,\n",
    "                     _bs1=_bs1, _s1=_s1, _bs2=_bs2, _s2=_s2, _bs3=_bs3, _s3=_s3, _bs2dz=_bs2dz, _qp_ind_mag=_qp_ind_mag, _qp_dz=_qp_dz)\n",
    "\n",
    "    g2 = MagnetArray(_per, _nper, _po=[-px,p2,pz], _w=w, _si=1, _type=2, _cx=_cx, _cz=_cz, _br=_br, _mu=_mu, _ndiv=_ndiv,\n",
    "                     _bs1=_bs1, _s1=_s1, _bs2=_bs2, _s2=_s2, _bs3=_bs3, _s3=_s3, _bs2dz=_bs2dz, _qp_ind_mag=_qp_ind_mag, _qp_dz=_qp_dz)\n",
    "\n",
    "    if(_use_sym):\n",
    "        u = rad.ObjCnt([g1,g2])\n",
    "        trf = rad.TrfCmbL(rad.TrfRot([0,0,0],[0,1,0],pi), rad.TrfInv())\n",
    "        rad.TrfMlt(u, trf, 2)\n",
    "        return u, g1, g2, 0, 0\n",
    "\n",
    "    g3 = MagnetArray(_per, _nper, _po=[-px,p3,-pz], _w=w, _si=-1, _type=1, _cx=_cx, _cz=_cz, _br=-_br, _mu=_mu, _ndiv=_ndiv,\n",
    "                     _bs1=_bs1, _s1=_s1, _bs2=_bs2, _s2=_s2, _bs3=_bs3, _s3=_s3, _bs2dz=_bs2dz, _qp_ind_mag=_qp_ind_mag, _qp_dz=_qp_dz)\n",
    "\n",
    "    g4 = MagnetArray(_per, _nper, _po=[px,p4,-pz], _w=w, _si=-1, _type=2, _cx=_cx, _cz=_cz, _br=-_br, _mu=_mu, _ndiv=_ndiv,\n",
    "                     _bs1=_bs1, _s1=_s1, _bs2=_bs2, _s2=_s2, _bs3=_bs3, _s3=_s3, _bs2dz=_bs2dz, _qp_ind_mag=_qp_ind_mag, _qp_dz=_qp_dz)\n",
    "\n",
    "    u = rad.ObjCnt([g1,g2,g3,g4])\n",
    "    return u, g1, g2, g3, g4\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## _Define functions for the magnetic materials and fields_\n",
    "\n",
    "The first of the following functions returns a Radia representation of the magnetization curves _M_ (in Tesla) vs _H_ (in Amp/m) for both the pole tips and the magnets. The material for the poles is defined by calling the function `rad.MatSatIsoTab` with the data table `ma`, which contains the paired values of _H_ and _M_, both expressed in Tesla. (The conversion of _H_ in Amp/m to Tesla is accomplished simply by multiplying _H_ by _μ<sub>o</sub>_, i.e. `4*pi/10^7`.) The material for the permanent magnets is defined by calling `rad.MatStd`, with arguments _material type_ and _remanent magnetization_. We will later plot the resulting _M_(_H_) curves for verification.\n",
    "\n",
    "The next function extracts magnetic field information from a given magnetic material defined by Radia. This allows us to plot the _M_(_H_) curves tocompare with our input data. And the third function calculates (for a _solved_ Radia magnet) the vertical field along the Y (longitudinal) axis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# magnetic constant (vacuum magnetic permeability)\n",
    "mu0 = 4 * np.pi / 1e7\n",
    "\n",
    "def materials(H, M,  smat, rm):\n",
    "    \"\"\"\n",
    "    define magnetic materials for the undulator poles and magnets\n",
    "    arguments:\n",
    "      H    = list of magnetic field values / (Amp/m)\n",
    "      M    = corresponding magnetization values / T\n",
    "      smat = material type string\n",
    "      rm   = remanent magnetization / T\n",
    "    return: Radia representations of ...\n",
    "      pole-tip material, magnet material\n",
    "    \"\"\"\n",
    "    # -- magnetic property of poles\n",
    "    #convH = 4. * np.pi / 1.0e7\n",
    "    ma = [[mu0 * H[i], M[i]] for i in range(len(H))]\n",
    "    mp = rad.MatSatIsoTab(ma)\n",
    "    # -- permanent magnet material\n",
    "    mm = rad.MatStd(smat, rm)\n",
    "\n",
    "    return mp, mm\n",
    "\n",
    "\n",
    "def getMagnMaterCompMvsH(meshH, ind, cmpnH,  cmpnM):\n",
    "    \"\"\"\n",
    "    extract magnetization M_comp vs field strength H_comp from material data\n",
    "    arguments:\n",
    "      meshH = how to mesh the H axis: [hmin, hmax, nh]\n",
    "      ind   = magnet material index\n",
    "      cmpnH = component of H ('x', 'y', or 'z')\n",
    "      cmpnM = component of M ('x', 'y', or 'z')\n",
    "    return:\n",
    "      given component of magnetization along the given mesh in H\n",
    "    \"\"\"\n",
    "    hMin = meshH[0]; hMax = meshH[1]; nh = meshH[2]\n",
    "    hStep = (hMax - hMin) / (nh - 1)\n",
    "    m = [ 0 ] * nh\n",
    "    sCmpnM = 'm' + cmpnM\n",
    "    h = hMin\n",
    "    H = [0,0,0]\n",
    "    for i in range(nh):\n",
    "        if(  cmpnH == 'x'): H[0] = h\n",
    "        elif(cmpnH == 'y'): H[1] = h\n",
    "        elif(cmpnH == 'z'): H[2] = h\n",
    "        m[i] = rad.MatMvsH(ind, sCmpnM, H)\n",
    "        h += hStep\n",
    "    return m\n",
    "\n",
    "\n",
    "def calcField(grp, per, numPer):\n",
    "    \"\"\"\n",
    "    compute vertical magnetic field vs longitudinal position (Y)\n",
    "    arguments:\n",
    "      grp    = magnet (Radia group)\n",
    "      per    = period / mm\n",
    "      numPer = number of full magnetic periods\n",
    "    return:\n",
    "      Bz(y), meshY\n",
    "    \"\"\"\n",
    "    xc = 0 # transverse X coordinate\n",
    "    zc = 0 # transverse Z coordinate\n",
    "    yMax = per * (numPer + 1) / 2.\n",
    "    yMin = -yMax\n",
    "    ny = 501\n",
    "    yStep = (yMax - yMin) / ny\n",
    "    y = np.arange(yMin, yMin + ny * yStep, yStep)\n",
    "    \n",
    "    points = [[xc, yi, zc] for yi in y]\n",
    "    BzVsY = rad.Fld(grp, 'bz', points)\n",
    "\n",
    "    return BzVsY, [yMin, yMax, ny]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## _Build and visualize an undulator_"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "###  _Vertical Field,  Pure Permanent Magnet_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "per = 35        #Undulator period in [mm]\n",
    "gap = 10        #Magnetic Gap in [mm]\n",
    "br = 1.2        #Remanent Field of the material in [T]\n",
    "nper = 3        #Number of periods\n",
    "lx = 30         #Horizontal width of the magnet blocks in [mm]\n",
    "lz = 18         #Vertical height of the magnet blocks in [mm]\n",
    "airgap = 0.5    #Airgap between the magnets in [mm]\n",
    "\n",
    "#________________  Start Computations ___________________\n",
    "wv = [lx, per/4 - airgap, lz]\n",
    "wh = wv\n",
    "p = np.array([0, 0, 0])\n",
    "g = undparts(p + [0, 0, gap/2 + lz/2], wv, wh, nper, per, br, 1)\n",
    "rad.TrfZerPara(g, p, [0, 0, 1])\n",
    "\n",
    "# set up the radia viewer and display geometry\n",
    "rvX = radia_viewer.RadiaViewer()\n",
    "rvX.add_geometry('Vertical Field,  Pure Permanent Magnet', g)\n",
    "rvX.display()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "###  _\"Hybrid\" with side Magnets_\n",
    "####  _U46 Hybrid Undulator_\n",
    "First enter and plot the data describing the magnetization curve for the iron pole tips. Then set the magnet parameters and materials, and “build the magnet”.Then set the magnet parameters and materials, and “build the magnet”."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ~iron type Va Permendur\n",
    "ironH = [  0.8,   1.5,     2.2,    3.6,    5.0,     6.8,     9.8,    18.0,\n",
    "          28.0,   37.5,   42.0,   55.0,   71.5,    80.0,    85.0,    88.0,\n",
    "          92.0,  100.0,  120.0,  150.0,  200.0,   300.0,   400.0,   600.0,\n",
    "         800.0, 1000.0, 2000.0, 4000.0, 6000.0, 10000.0, 25000.0, 40000.0]\n",
    "ironM = [0.000998995, 0.00199812, 0.00299724, 0.00499548, 0.00699372, 0.00999145, 0.0149877, 0.0299774,\n",
    "         0.0499648,   0.0799529,  0.0999472,  0.199931,   0.49991,    0.799899,   0.999893, 1.09989,\n",
    "         1.19988,     1.29987,    1.41985,    1.49981,    1.59975,    1.72962,    1.7995,   1.89925,\n",
    "         1.96899,     1.99874,    2.09749,    2.19497,    2.24246,    2.27743,    2.28958,  2.28973]\n",
    "# set parameters for this undulator\n",
    "# -- general parameters\n",
    "nper =  2    # number of full magnetic periods\n",
    "per = 46       # period / mm\n",
    "gap = 20       # gap height / mm\n",
    "gapOffset = 1  # \n",
    "\n",
    "# -- parameters for the iron poles\n",
    "lp =  [45,5,25] # dimensions / mm\n",
    "nsp = [2,2,5]   # pole-tip segmentation (try varying this setting)\n",
    "cp =  [1,0,1]   # assign color\n",
    "ll = per/2 - lp[1]\n",
    "\n",
    "# -- parameters for the magnet block\n",
    "lm  = [65,ll,45] # dimensions / mm\n",
    "nsm = [1,3,1]    # magnet-block segmentation (try varying this)\n",
    "cm  = [0,1,1]    # assign color\n",
    "\n",
    "# -- magnetic materials\n",
    "# pole tips: ~iron type Va Permendur\n",
    "# permanent magnets: NdFeB with 1.2 Tesla remanent magnetization\n",
    "mp, mm = materials(ironH, ironM, 'NdFeB', 1.2)\n",
    "\n",
    "# then build the undulator\n",
    "und, pole, magnet = undulator(lp, mp, nsp, cp, lm, mm, nsm, cm, gap, gapOffset, nper)\n",
    "\n",
    "# set up the radia viewer and display geometry\n",
    "rv = radia_viewer.RadiaViewer()\n",
    "rv.add_geometry('U46 Hybrid Undulator', und)\n",
    "rv.display()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "####  _SPEAR BL11 Wiggler_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gap = 16\n",
    "nper = 3\n",
    "per = 175\n",
    "gapoffset = 5\n",
    "br = 1.25\n",
    "\n",
    "lp = [50,42,90]\n",
    "nsp = [2,2,10]\n",
    "cp = [1,0,1]\n",
    "\n",
    "lm = [126,per/2-lp[1],110]\n",
    "nsm = [1,3,1]\n",
    "cm = [0,1,1]\n",
    "\n",
    "lms = [36,lp[1]-1,110]\n",
    "nsms = [1,1,1]\n",
    "cms = [0,1,0]\n",
    "\n",
    "matp = rad.MatStd('Steel42')\n",
    "matm = rad.MatStd('NdFeB', br)\n",
    "matms = rad.MatStd('NdFeB', br)\n",
    "\n",
    "p = np.array([0,0,0])\n",
    "SPEAR = hyb(p,per,gap,nper,gapoffset, lp,nsp,matp,cp, lm,nsm,matm,cm, lms,nsms,matms,cms)\n",
    "\n",
    "# set up the radia viewer and display geometry\n",
    "rvX = radia_viewer.RadiaViewer()\n",
    "rvX.add_geometry('SPEAR BLII Wiggler', SPEAR)\n",
    "rvX.display()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### _\"Apple II\" Type Helical/Planar_\n",
    "\n",
    "#### _ESRF  HU88 on ID8_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "per = 88                 # Undulator Period in [mm] \n",
    "lx = 25*2                 # Total Horizontal width of the magnet array in [mm]\n",
    "lz = 12                 # Vertical height of the magnet blocks\n",
    "gapx = 1                 # Horizontal gap between magnet arrays in [mm]\n",
    "gap = 16                 # Magnetic Gap in[mm]\n",
    "p1 = per*0.0                 # Antisym.  motion of the girders  : Rotate the Linear Pol.\n",
    "p2 = per*0.5                 # Symmetric L motion of the girders : Linear\\[Rule] Circular\\[Rule] Linear\n",
    "br = 1.2                 # Remanent field of the magnet blocks in [T]\n",
    "nper = 6                 # Number of periods\n",
    "airgap = 0                 # longitudinal airgap between magnets in [mm]\n",
    "\n",
    "#Start Computations\n",
    "p = np.array([0, 0, 0])\n",
    "HU88 = apple2G(p, per, gap, gapx, -p1 + p2/2, -p2/2, -p2/2, p1 + p2/2, lx, lz, airgap, br, nper)\n",
    "\n",
    "# set up the radia viewer and display geometry\n",
    "rvX = radia_viewer.RadiaViewer()\n",
    "rvX.add_geometry('ESRF  HU88 on ID8', HU88)\n",
    "rvX.display()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### _ESRF  HU38 on ID12_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "per = 38                 # Undulator Period in [mm] \n",
    "lx = 25*2                 # Total Horizontal width of the magnet array in [mm]\n",
    "lz = 15                 # Vertical height of the magnet blocks\n",
    "gapx = 1                 # Horizontal gap between magnet arrays in [mm]\n",
    "gap = 16                 # Magnetic Gap in[mm]\n",
    "p1 = per*0.0                 # Antisym.  motion of the girders  : Rotate the Linear Pol.\n",
    "p2 = per*0.5                 # Symmetric L motion of the girders : Linear\\[Rule] Circular\\[Rule] Linear\n",
    "br = 1.2                 # Remanent field of the magnet blocks in [T]\n",
    "nper = 4                 # Number of periods\n",
    "airgap = 0                 # longitudinal airgap between magnets in [mm]\n",
    "\n",
    "#Start Computations\n",
    "p = np.array([0, 0, 0])\n",
    "HU38 = apple2G(p, per, gap, gapx, -p1 + p2/2, -p2/2, -p2/2, p1 + p2/2, lx, lz, airgap, br, nper)\n",
    "\n",
    "# set up the radia viewer and display geometry\n",
    "rvX = radia_viewer.RadiaViewer()\n",
    "rvX.add_geometry('ESRF  HU38 on ID12', HU38)\n",
    "rvX.display()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### _APPLE-II type PPM Undulator_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Undulator Geometry Parameters (all lengths are in [mm])\n",
    "per = 50. #Period\n",
    "nper = 40 #38 #Number of Periods (should be EVEN here to make correct integer (M3+M5) modular structure!)\n",
    "\n",
    "gapx = 1. #Horizontal Gap between Magnet Arrays\n",
    "gap = 11.5 #Vertical Magnetic Gap\n",
    "phase = 0. #0.2923*per #0.29*per #0.28*per #per/2. #per/4. #Longitudinal Shift between Magnet Arrays\n",
    "phaseType = 1 #=1 means Parallel, -1 Anti-Parallel displacement of Magnet Arrays\n",
    "\n",
    "lx = 34. #Horizontal Magnet Size\n",
    "lz = 34. #Vertical Magnet Size\n",
    "cx = 5. #Horizontal Notch size\n",
    "cz = 5. #Vertical Notch size\n",
    "\n",
    "air = 0.05 #Air Space between Magnets in Longitudinal Direction\n",
    "\n",
    "#Terminations\n",
    "bs3 = 7.25*per/49.2 #8.43/57.2*per (*7.89*per/105.2;*) (*7.25*per/49.2;*)  (*11.2*per/80.;*) #Magnet D1\n",
    "bs2 = 6.11*per/49.2 #7.09/57.2*per #(*13.10*per/105.2;*) (*6.11*per/49.2;*) (*9.985*per/80.;*) #Magnet C\n",
    "bs1 = 3.82*per/49.2 #6.58/57.2*per #(*13.56*per/105.2;*)(*3.82*per/49.2;*) (*7.9*per/80.;*) #Magnet D2\n",
    "s1 = 1.*per/49.2 #1.41*per/57.2 #(*4.63*per/105.2;*)  (*1.*per/49.2;*)   (*2.4261*per/80.; *) #Outmost Gap G2\n",
    "s2 = 4.5*per/49.2 #3.84*per/57.2 #(*2.47*per/105.2;*) (*4.5*per/49.2;*) (*5.7595*per/80.;*) #Inmost Gap G1\n",
    "s3 = 1.*per/49.2 #0. #(*1.*per/49.2;*) (*1.;*)(*0.;*)\n",
    "bs2dz = 0. #(*0.16961;*) #Vertical displacement of vertically-magnetised termination block\n",
    "\n",
    "#Material Constants\n",
    "br = 1.27 #Remanent Magnetization\n",
    "mu = [0.05,0.15] #[0.06,0.17] #Suscettivita\n",
    "\n",
    "#Subdivision Params\n",
    "nDiv = [3,2,3] \n",
    "\n",
    "#Quasi-Periodic Perturbation\n",
    "indsMagDispQP = [-77, -63, -49, -35, -21, -7, 7, 21, 35, 49, 63, 77] #None #Indexes of magnets counting from the central magnet of the structure, which has index 0.\n",
    "vertMagDispQP = 0. \n",
    "\n",
    "#Start Computations\n",
    "und = APPLE_II(_per=per, _nper=nper, _gap=gap, _gapx=gapx, _phase=phase, _phase_type=phaseType, _lx=lx, _lz=lz, _cx=cx, _cz=cz, _air=air,\n",
    "                   _br=br, _mu=mu, _ndiv=nDiv, _bs1=bs1, _s1=s1, _bs2=bs2, _s2=s2, _bs3=bs3, _s3=s3, _bs2dz=bs2dz,\n",
    "                   _qp_ind_mag=indsMagDispQP, _qp_dz=vertMagDispQP, _use_sym=True)[0]\n",
    "\n",
    "# set up the radia viewer and display geometry\n",
    "rvX = radia_viewer.RadiaViewer()\n",
    "rvX.add_geometry('APPLE-II type PPM Undulator', und)\n",
    "rvX.display()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### _\"Spring8\" Type Helical/Planar_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "per = 88                 #Undulator Period in [mm]\n",
    "gapx = 1                 #Horizontal gap between magnet arrays in [mm]\n",
    "gap = 16                 #Magnetic Gap in[mm]\n",
    "phase = per/4            #Phase in [mm]\n",
    "lxc = 14                 #Horizontal size of the magnets of the middle array\n",
    "lzc = 12                 #Vertical size of the magnets of the middle array\n",
    "colc = [1, 0, 0]         #RGB color of the middle array\n",
    "lxs = 25                 #Horizontal size of the magnets of the side arrays\n",
    "lzs = 16                 #Vertical size of the magnets of the side arrays\n",
    "cols = [0, 1, 0]         #RGB color of the side arrays\n",
    "br = 1.2                 #Remanent field of the magnet blocks in [T]\n",
    "nper = 4                 #Number of periods\n",
    "airgap = 0.1             #longitudinal airgap between magnets in [mm]\n",
    "\n",
    "#________________  Start Computations ___________________\n",
    "# rad.UtiDelAll()\n",
    "p = np.array([0, 0, 0])\n",
    "Spring8 = sp8(p, per, gap, gapx, phase, lxc, lzc, colc, lxs, lzs, cols, airgap, br, nper)\n",
    "\n",
    "# set up the radia viewer and display geometry\n",
    "rvX = radia_viewer.RadiaViewer()\n",
    "rvX.add_geometry('\"Spring8\" Type Helical/Planar', Spring8)\n",
    "rvX.display()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## _Solve for the undulator magnetization_\n",
    "\n",
    "Here we ask Radia to solve for the undulator magnetization—_i.e._ for the magnetization within each sub-block that constitutes the undulator—and also report some statistics concerning that solution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# solve for the magnetization\n",
    "prec = 0.0003  # precision for this computation\n",
    "maxIter = 1000 # maximum allowed iterations\n",
    "t0 = tm.time()\n",
    "res = rad.Solve(und, prec, maxIter)\n",
    "t1 = tm.time()\n",
    "\n",
    "print(\"Solved for magnetization in {0:6f} s\".format(t1 - t0))\n",
    "print(\"Relaxation results\")\n",
    "print(\"  number of iterations: {0:5d}\".format(int(res[3])))\n",
    "print(\"  average stability of magnetization at last iteration: {0:.4e} T\".format(res[0]))\n",
    "print(\"  maximum absolute magnetization at last iteration: {0:.5f} T\".format(res[1]))\n",
    "print(\"  maximum H vector at last iteration: {0:.5f} T\".format(res[2]))\n",
    "print(\"Peak magnetic field: {0:.8f} T\".format(rad.Fld(und, 'bz', [0,0,0])))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now ask Radia to compute the magnetic field along the undulator’s longitudinal axis. And then plot the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute magnetic field values\n",
    "t2 = tm.time()\n",
    "BzVsY, MeshY = calcField(und, per, nper)\n",
    "t3 = tm.time()\n",
    "print(\"Computed requested field vaues in {0:6f} s\\n\".format(t3 - t2))\n",
    "\n",
    "# and plot\n",
    "y = np.linspace(MeshY[0], MeshY[1], MeshY[2])\n",
    "plt.figure(figsize=(7,4.3))\n",
    "plt.plot(y, BzVsY, linewidth = 2.0)\n",
    "plt.title(r\"Vertical Magnetic Field, $B_z$, along undulator's $y$ axis\")\n",
    "plt.xlabel('y / mm')\n",
    "plt.ylabel('$B_z$ / T')\n",
    "plt.grid()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## _First Field Integral_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute the first field integral at the point outside the undulator\n",
    "ypoint = per*(nper+1)\n",
    "Bz_int1st_point = rad.FldInt(und,'fin','ibz',[0,-1e5,0],[0,ypoint,0])\n",
    "print(\"the 1st field integral outside the undulator is: \", Bz_int1st_point)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute the first field integral along the longitudinal axis\n",
    "yy = np.linspace(-per*(nper+1)/2.,ypoint, 501)\n",
    "Bz_int1st = np.zeros_like(yy)\n",
    "for i, yi in enumerate(yy):\n",
    "    if i == 0:\n",
    "        Bz_int1st[i] = rad.FldInt(und,'fin','ibz',[0,-1e5,0],[0,yi,0])\n",
    "    else:\n",
    "        Bz_int1st[i] = Bz_int1st[i-1] + rad.FldInt(und,'fin','ibz',[0,yy[i-1],0],[0,yi,0])\n",
    "print(\"the 1st field integral outside the undulator is: \", Bz_int1st[-1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#plot the first field integral along the longitudinal axis\n",
    "plt.figure(figsize=(7,4.3))\n",
    "plt.plot(yy, Bz_int1st, linewidth = 2.0)\n",
    "plt.title(r\"1st Field integral along undulator's $y$ axis\")\n",
    "plt.xlabel('y / mm')\n",
    "plt.ylabel('$B_z$ / T')\n",
    "plt.grid()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## _Optimizing Parameters_\n",
    "\n",
    "The fundamental undulator parameter K, is given by\n",
    "$$K = \\frac{e B_0 \\lambda_u}{2\\pi m_e c}.$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def undulatorK_simple(obj, per, pf_loc = [0,0,0] ,prec = 1e-5, maxIter = 10000, lprint = True):\n",
    "    \"\"\"\n",
    "    compute undulator K value\n",
    "    arguments:\n",
    "      obj = undulator object\n",
    "      per = undulator period / m\n",
    "      pf_loc = peak field location [x, y, z]\n",
    "      prec = precision goal for this computation\n",
    "      maxIter = maximum allowed iterations\n",
    "      lprint: whether or not to print results\n",
    "    return:\n",
    "      K = (e B_0 \\lambda_u) / (2\\pi m_e c)\n",
    "    \"\"\"\n",
    "    res = rad.Solve(obj, prec, maxIter)\n",
    "    peak_field = abs(rad.Fld(obj, 'bz', pf_loc)) #peak field / T\n",
    "    k = sc.e * peak_field * per * 1e-3 / (2 * sc.pi * sc.m_e * sc.c)\n",
    "    if lprint:\n",
    "        print(\"peak field:\", peak_field, \"(calculated at given location\",pf_loc,\")\\nperiod is\", per, \"(given input)\\nk is\", k)\n",
    "\n",
    "    return k\n",
    "\n",
    "def undulatorK(obj, xyz_start, xyz_end, n = 1000, prec = 1e-5, maxIter = 10000, lsove = True, lprint = True):\n",
    "    \"\"\"\n",
    "    compute undulator K value\n",
    "    arguments:\n",
    "      obj = undulator object\n",
    "      xyz_start = start point [x, y, z] in field calculation\n",
    "      xyz_end = end point [x, y, z] in field calculation\n",
    "      n = number of points in field calculation\n",
    "      prec = precision goal for this computation\n",
    "      maxIter = maximum allowed iterations\n",
    "      lsove: if true, run rad.Solve\n",
    "      lprint: if true, print results\n",
    "    return:\n",
    "      K = (e B_0 \\lambda_u) / (2\\pi m_e c)\n",
    "    \"\"\"\n",
    "    # --- Sampling\n",
    "    x0, y0, z0 = xyz_start\n",
    "    x1, y1, z1 = xyz_end\n",
    "    # Steps\n",
    "    dx = (x1 - x0) / (n - 1)\n",
    "    dy = (y1 - y0) / (n - 1)\n",
    "    dz = (z1 - z0) / (n - 1)\n",
    "    # Positions\n",
    "    x = [x0 + k * dx for k in range(n)]\n",
    "    y = [y0 + k * dy for k in range(n)]\n",
    "    z = [z0 + k * dz for k in range(n)]\n",
    "    xyz = [[x[k], y[k], z[k]] for k in range(n)]\n",
    "    \n",
    "    if lsove: \n",
    "        res = rad.Solve(obj, prec, maxIter)\n",
    "        \n",
    "    fields = rad.Fld(obj, 'bz', xyz)\n",
    "    peak_field = max([abs(fieldb) for fieldb in fields]) #peak field / T\n",
    "    fields_fft = np.fft.fft(fields)\n",
    "    fields_fft_abs = [abs(fieldfft) for fieldfft in fields_fft]\n",
    "    fft_per = 1/(fields_fft_abs.index(max(fields_fft_abs))/dy/n)\n",
    "    k = sc.e * peak_field * fft_per * 1e-3 / (2 * sc.pi * sc.m_e * sc.c)\n",
    "    \n",
    "    if lprint:\n",
    "        print(\"peak field:\", peak_field, \"(obtained by searching maximum value)\\nperiod is\",\n",
    "              fft_per, \"(obtained from FFT)\\nk is\", k,\"(calcuated using maximum searched peak field)\\n\")\n",
    "    return k"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "undulatorK_simple(und,per)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "undulatorK(und, [0,-per*nper/2,0], [0,per*nper/2,0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## _Generating an Undulator Map File_\n",
    "An undulator (or wiggler)  must be created before executing any cell of this section. A suggestion is that: go to \"Build and visualize an undulator\" section and re-build the desired undulator. It is also a good strategy to check that the undulator has the desired geometry and that it is producing the correct field using the two previous sections before creating the tracking file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "obj = und     #object (undulator), re-build the desired undulator to ensure the correct geometry inputs\n",
    "p0 = [0,-per*nper/2,0]     # the starting point of longitudinal integration\n",
    "vl = [0,1,0]     # longitudinal integration direction\n",
    "vt = [1,0,0]     # one of the transverse direction (horizontal)\n",
    "r1 = gap*0.75     #range of the transverse grid along vt (horizontal)\n",
    "np1 = 21     #number of points in transverse direction vt (horizontal)\n",
    "r2 = gap*0.75     #range of the transverse grid along (vt cross vl, verticle)\n",
    "np2 = 21     #number of points in transverse direction (vt cross vl, verticle)\n",
    "# default paras:\n",
    "dpar = [1,8,0,0]    #[maximum number of magnetic field harmonics to treat:1,number of longitudinal points:8,steps of transverse differentiation:0,0]\n",
    "unit = 'T2m2'     #the units for the resulting 2nd order kick values T2m2 or rad or microrad\n",
    "en = 1     #eletron energy in GeV (required only if units are rad or microrad)\n",
    "oFormat = 'fix'     #the format of the output data string: fix or tab\n",
    "\n",
    "km = rad.FldFocKickPer(obj,p0,vl,per,nper,vt,r1,np1,r2,np2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### _plot the kick maps_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hKick = np.array(km[0])\n",
    "vKick = np.array(km[1])\n",
    "lmSqr = np.array(km[2])\n",
    "xmesh,zmesh= np.meshgrid(np.array(km[3]),np.array(km[4]))\n",
    "\n",
    "# horizontal kick\n",
    "hkm = plt.axes(projection='3d')\n",
    "hkm.plot_surface(xmesh, zmesh,hKick,  rstride=1, cstride=1,\n",
    "                cmap='viridis', edgecolor='none')\n",
    "hkm.set_title('horizontal kick');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# verticle kick\n",
    "vkm = plt.axes(projection='3d')\n",
    "vkm.plot_surface(xmesh, zmesh, vKick,  rstride=1, cstride=1,\n",
    "                cmap='viridis', edgecolor='none')\n",
    "vkm.set_title('vertical kick');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# longitudinally-integrated squared transverse magnetic field\n",
    "lmsqr = plt.axes(projection='3d')\n",
    "lmsqr.plot_surface(xmesh, zmesh, lmSqr,  rstride=1, cstride=1,\n",
    "                cmap='viridis', edgecolor='none')\n",
    "lmsqr.set_title('longitudinally-integrated squared transverse magnetic field');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### _save the map to text file_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"km.txt\", \"w\") as km_text:\n",
    "    km_text.write(km[5])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "py3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}