{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "96686a69-20a6-4bf4-ae7e-8725c634ec1c",
   "metadata": {},
   "source": [
    "# Example-10: Dipole (element)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "2c38c277-00e8-49f7-a045-9e580c2af148",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Comparison of dipole element with MADX-PTC and other features\n",
    "\n",
    "# Note, cylindrical multipole is included upto octupole order\n",
    "# Potential is not truncated in paraxial case, only sqrt is expanded \n",
    "# In exact case effects of multipoles are not accounted in wedges (cases with e1 or e2 not equal to zero)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "0ceaf51f-dfec-4437-8760-770e85e1445b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "from os import system\n",
    "\n",
    "import torch\n",
    "from model.library.drift import Drift\n",
    "from model.library.quadrupole import Quadrupole\n",
    "from model.library.sextupole import Sextupole\n",
    "from model.library.dipole import Dipole"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "deb2db37-6623-411c-8d64-adcd2b4a4324",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[5.160257777686003e-05, -0.004956273150572646, -0.0030019980019983517, 0.001]\n",
      "[5.1602577776445975e-05, -0.004956273150572915, -0.0030019980019981656, 0.0010000000000000434]\n",
      "[4.1405680967221414e-16, 2.688821387764051e-16, -1.8604909279851256e-16, -4.336808689942018e-17]\n"
     ]
    }
   ],
   "source": [
    "# Tracking (paraxial)\n",
    "\n",
    "ptc = Path('ptc')\n",
    "obs = Path('track.obs0001.p0001')\n",
    "\n",
    "exact = False\n",
    "align = False\n",
    "\n",
    "length = 2.0\n",
    "angle = 0.05\n",
    "e1 = 0.0\n",
    "e2 = 0.0\n",
    "kn = 0.0\n",
    "ks = 0.0\n",
    "ms = 0.0\n",
    "mo = 0.0\n",
    "dp = 0.001\n",
    "state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)\n",
    "qx, px, qy, py = state.tolist()\n",
    "\n",
    "dx = align*torch.tensor(0.05, dtype=torch.float64)\n",
    "dy = align*torch.tensor(-0.02, dtype=torch.float64)\n",
    "dz = align*torch.tensor(0.05, dtype=torch.float64)\n",
    "\n",
    "wx = align*torch.tensor(0.005, dtype=torch.float64)\n",
    "wy = align*torch.tensor(-0.005, dtype=torch.float64)\n",
    "wz = align*torch.tensor(0.1, dtype=torch.float64)\n",
    "\n",
    "error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}\n",
    "\n",
    "code = f\"\"\"\n",
    "mag:sbend,l={length},angle={angle},e1={e1},e2={e2},knl={{0.0,{kn*length},{ms*length},{mo*length}}},ksl={{0.0,{ks*length}}},kill_ent_fringe=false,kill_exi_fringe=false;\n",
    "map:line=(mag) ;\n",
    "beam,energy=1.0E+6,particle=electron ;\n",
    "set,format=\"20.20f\",\"-20s\" ;\n",
    "use,period=map ;\n",
    "select,flag=error,pattern=\"mag\" ;\n",
    "ealign,dx={dx.item()},dy={dy.item()},ds={dz.item()},dphi={wx.item()},dtheta={wy.item()},dpsi={wz.item()} ;\n",
    "ptc_create_universe,sector_nmul_max=5,sector_nmul=5 ;\n",
    "ptc_create_layout,model=1,method=6,nst=1000,exact={str(exact).lower()} ;\n",
    "ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=true ;\n",
    "ptc_align ;\n",
    "ptc_start,x={qx},px={px},y={qy},py={py},pt={dp},t=0.0 ;\n",
    "ptc_track,icase=5,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}} ;\n",
    "ptc_track_end ;\n",
    "ptc_end ;\n",
    "\"\"\" \n",
    "\n",
    "with ptc.open('w') as stream:\n",
    "    stream.write(code)\n",
    "    \n",
    "system(f'madx < {str(ptc)} > /dev/null')\n",
    "\n",
    "with obs.open('r') as stream:\n",
    "    for line in stream:\n",
    "        continue\n",
    "    _, _, qx, px, qy, py, *_ = line.split()\n",
    "    \n",
    "ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)\n",
    "\n",
    "D = Dipole('D', length=length, angle=angle, e1=e1, e2=e2, kn=kn, ks=ks, ms=ms, mo=mo, dp=dp, exact=exact, order=2, ns=25)\n",
    "res = D(state, alignment=align, data={**D.data(), **error})\n",
    "\n",
    "print(ref.tolist())\n",
    "print(res.tolist())\n",
    "print((ref - res).tolist())\n",
    "\n",
    "ptc.unlink()\n",
    "obs.unlink()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ce07fc8e-d746-44da-a26b-892b69152ccb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[6.408749411377232e-05, -0.00494998958983017, -0.002995752944647053, 0.0010049983869644134]\n",
      "[6.408749411335537e-05, -0.004949989589830473, -0.0029957529446465857, 0.0010049983869644566]\n",
      "[4.169434979564568e-16, 3.0270924655795284e-16, -4.670742959067553e-16, -4.3151246464923076e-17]\n"
     ]
    }
   ],
   "source": [
    "# Tracking (paraxial, e1 & e2)\n",
    "\n",
    "ptc = Path('ptc')\n",
    "obs = Path('track.obs0001.p0001')\n",
    "\n",
    "exact = False\n",
    "align = False\n",
    "\n",
    "length = 2.0\n",
    "angle = 0.05\n",
    "e1 = 0.025\n",
    "e2 = 0.025\n",
    "kn = 0.0\n",
    "ks = 0.0\n",
    "ms = 0.0\n",
    "mo = 0.0\n",
    "dp = 0.001\n",
    "state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)\n",
    "qx, px, qy, py = state.tolist()\n",
    "\n",
    "dx = align*torch.tensor(0.05, dtype=torch.float64)\n",
    "dy = align*torch.tensor(-0.02, dtype=torch.float64)\n",
    "dz = align*torch.tensor(0.05, dtype=torch.float64)\n",
    "\n",
    "wx = align*torch.tensor(0.005, dtype=torch.float64)\n",
    "wy = align*torch.tensor(-0.005, dtype=torch.float64)\n",
    "wz = align*torch.tensor(0.1, dtype=torch.float64)\n",
    "\n",
    "error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}\n",
    "\n",
    "code = f\"\"\"\n",
    "mag:sbend,l={length},angle={angle},e1={e1},e2={e2},knl={{0.0,{kn*length},{ms*length},{mo*length}}},ksl={{0.0,{ks*length}}},kill_ent_fringe=false,kill_exi_fringe=false;\n",
    "map:line=(mag) ;\n",
    "beam,energy=1.0E+6,particle=electron ;\n",
    "set,format=\"20.20f\",\"-20s\" ;\n",
    "use,period=map ;\n",
    "select,flag=error,pattern=\"mag\" ;\n",
    "ealign,dx={dx.item()},dy={dy.item()},ds={dz.item()},dphi={wx.item()},dtheta={wy.item()},dpsi={wz.item()} ;\n",
    "ptc_create_universe,sector_nmul_max=5,sector_nmul=5 ;\n",
    "ptc_create_layout,model=1,method=6,nst=1000,exact={str(exact).lower()} ;\n",
    "ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=true ;\n",
    "ptc_align ;\n",
    "ptc_start,x={qx},px={px},y={qy},py={py},pt={dp},t=0.0 ;\n",
    "ptc_track,icase=5,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}} ;\n",
    "ptc_track_end ;\n",
    "ptc_end ;\n",
    "\"\"\" \n",
    "\n",
    "with ptc.open('w') as stream:\n",
    "    stream.write(code)\n",
    "    \n",
    "system(f'madx < {str(ptc)} > /dev/null')\n",
    "\n",
    "with obs.open('r') as stream:\n",
    "    for line in stream:\n",
    "        continue\n",
    "    _, _, qx, px, qy, py, *_ = line.split()\n",
    "    \n",
    "ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)\n",
    "\n",
    "D = Dipole('D', length=length, angle=angle, e1=e1, e2=e2, kn=kn, ks=ks, ms=ms, mo=mo, dp=dp, exact=exact, order=2, ns=25)\n",
    "res = D(state, alignment=align, data={**D.data(), **error})\n",
    "\n",
    "print(ref.tolist())\n",
    "print(res.tolist())\n",
    "print((ref - res).tolist())\n",
    "\n",
    "ptc.unlink()\n",
    "obs.unlink()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "6e82c64a-9aa6-498a-b545-5bdc7f06307d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.0027250502286323983, -0.004697460778410744, -0.007781317874336121, -0.004015687563474919]\n",
      "[0.0027250502286311207, -0.004697460778410097, -0.007781317874338802, -0.004015687563476177]\n",
      "[1.2776238400569184e-15, -6.47051856539349e-16, 2.6810151321221554e-15, 1.2576745200831851e-15]\n"
     ]
    }
   ],
   "source": [
    "# Tracking (paraxial, e1 & e2, alignment)\n",
    "\n",
    "ptc = Path('ptc')\n",
    "obs = Path('track.obs0001.p0001')\n",
    "\n",
    "exact = False\n",
    "align = True\n",
    "\n",
    "length = 2.0\n",
    "angle = 0.05\n",
    "e1 = 0.025\n",
    "e2 = 0.025\n",
    "kn = 0.0\n",
    "ks = 0.0\n",
    "ms = 0.0\n",
    "mo = 0.0\n",
    "dp = 0.001\n",
    "state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)\n",
    "qx, px, qy, py = state.tolist()\n",
    "\n",
    "dx = align*torch.tensor(0.05, dtype=torch.float64)\n",
    "dy = align*torch.tensor(-0.02, dtype=torch.float64)\n",
    "dz = align*torch.tensor(0.05, dtype=torch.float64)\n",
    "\n",
    "wx = align*torch.tensor(0.005, dtype=torch.float64)\n",
    "wy = align*torch.tensor(-0.005, dtype=torch.float64)\n",
    "wz = align*torch.tensor(0.1, dtype=torch.float64)\n",
    "\n",
    "error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}\n",
    "\n",
    "code = f\"\"\"\n",
    "mag:sbend,l={length},angle={angle},e1={e1},e2={e2},knl={{0.0,{kn*length},{ms*length},{mo*length}}},ksl={{0.0,{ks*length}}},kill_ent_fringe=false,kill_exi_fringe=false;\n",
    "map:line=(mag) ;\n",
    "beam,energy=1.0E+6,particle=electron ;\n",
    "set,format=\"20.20f\",\"-20s\" ;\n",
    "use,period=map ;\n",
    "select,flag=error,pattern=\"mag\" ;\n",
    "ealign,dx={dx.item()},dy={dy.item()},ds={dz.item()},dphi={wx.item()},dtheta={wy.item()},dpsi={wz.item()} ;\n",
    "ptc_create_universe,sector_nmul_max=5,sector_nmul=5 ;\n",
    "ptc_create_layout,model=1,method=6,nst=1000,exact={str(exact).lower()} ;\n",
    "ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=true ;\n",
    "ptc_align ;\n",
    "ptc_start,x={qx},px={px},y={qy},py={py},pt={dp},t=0.0 ;\n",
    "ptc_track,icase=5,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}} ;\n",
    "ptc_track_end ;\n",
    "ptc_end ;\n",
    "\"\"\" \n",
    "\n",
    "with ptc.open('w') as stream:\n",
    "    stream.write(code)\n",
    "    \n",
    "system(f'madx < {str(ptc)} > /dev/null')\n",
    "\n",
    "with obs.open('r') as stream:\n",
    "    for line in stream:\n",
    "        continue\n",
    "    _, _, qx, px, qy, py, *_ = line.split()\n",
    "    \n",
    "ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)\n",
    "\n",
    "D = Dipole('D', length=length, angle=angle, e1=e1, e2=e2, kn=kn, ks=ks, ms=ms, mo=mo, dp=dp, exact=exact, order=2, ns=25)\n",
    "res = D(state, alignment=align, data={**D.data(), **error})\n",
    "\n",
    "print(ref.tolist())\n",
    "print(res.tolist())\n",
    "print((ref - res).tolist())\n",
    "\n",
    "ptc.unlink()\n",
    "obs.unlink()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "c3c4baf6-e4ee-46ac-bb9e-5f43090d3109",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[-0.002043638956596164, -0.006566087436123452, -0.00259447159458907, 0.0010855376916213359]\n",
      "[-0.0020440852924812135, -0.006566346702412346, -0.002594440402484201, 0.001085525276668085]\n",
      "[4.46335885049675e-07, 2.592662888935629e-07, -3.119210486906762e-08, 1.2414953250768773e-08]\n"
     ]
    }
   ],
   "source": [
    "# Tracking (paraxial, kn & ks)\n",
    "\n",
    "# Note, in paraxial case, MADX seems to truncate cylindrical potential\n",
    "\n",
    "ptc = Path('ptc')\n",
    "obs = Path('track.obs0001.p0001')\n",
    "\n",
    "exact = False\n",
    "align = False\n",
    "\n",
    "length = 2.0\n",
    "angle = 0.05\n",
    "e1 = 0.0\n",
    "e2 = 0.0\n",
    "kn = 0.1\n",
    "ks = 0.1\n",
    "ms = 0.0\n",
    "mo = 0.0\n",
    "dp = 0.001\n",
    "state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)\n",
    "qx, px, qy, py = state.tolist()\n",
    "\n",
    "dx = align*torch.tensor(0.05, dtype=torch.float64)\n",
    "dy = align*torch.tensor(-0.02, dtype=torch.float64)\n",
    "dz = align*torch.tensor(0.05, dtype=torch.float64)\n",
    "\n",
    "wx = align*torch.tensor(0.005, dtype=torch.float64)\n",
    "wy = align*torch.tensor(-0.005, dtype=torch.float64)\n",
    "wz = align*torch.tensor(0.1, dtype=torch.float64)\n",
    "\n",
    "error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}\n",
    "\n",
    "code = f\"\"\"\n",
    "mag:sbend,l={length},angle={angle},e1={e1},e2={e2},knl={{0.0,{kn*length},{ms*length},{mo*length}}},ksl={{0.0,{ks*length}}},kill_ent_fringe=false,kill_exi_fringe=false;\n",
    "map:line=(mag) ;\n",
    "beam,energy=1.0E+6,particle=electron ;\n",
    "set,format=\"20.20f\",\"-20s\" ;\n",
    "use,period=map ;\n",
    "select,flag=error,pattern=\"mag\" ;\n",
    "ealign,dx={dx.item()},dy={dy.item()},ds={dz.item()},dphi={wx.item()},dtheta={wy.item()},dpsi={wz.item()} ;\n",
    "ptc_create_universe,sector_nmul_max=5,sector_nmul=5 ;\n",
    "ptc_create_layout,model=1,method=6,nst=1000,exact={str(exact).lower()} ;\n",
    "ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=true ;\n",
    "ptc_align ;\n",
    "ptc_start,x={qx},px={px},y={qy},py={py},pt={dp},t=0.0 ;\n",
    "ptc_track,icase=5,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}} ;\n",
    "ptc_track_end ;\n",
    "ptc_end ;\n",
    "\"\"\" \n",
    "\n",
    "with ptc.open('w') as stream:\n",
    "    stream.write(code)\n",
    "    \n",
    "system(f'madx < {str(ptc)} > /dev/null')\n",
    "\n",
    "with obs.open('r') as stream:\n",
    "    for line in stream:\n",
    "        continue\n",
    "    _, _, qx, px, qy, py, *_ = line.split()\n",
    "    \n",
    "ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)\n",
    "\n",
    "D = Dipole('D', length=length, angle=angle, e1=e1, e2=e2, kn=kn, ks=ks, ms=ms, mo=mo, dp=dp, exact=exact, order=2, ns=25)\n",
    "res = D(state, alignment=align, data={**D.data(), **error})\n",
    "\n",
    "print(ref.tolist())\n",
    "print(res.tolist())\n",
    "print((ref - res).tolist())\n",
    "\n",
    "ptc.unlink()\n",
    "obs.unlink()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "37b829ae-2de0-40fd-bdca-504e2949c3d9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.013167193439103305, 0.005511119368028012, -0.010781508233295324, -0.007329921037793256]\n",
      "[0.013155910631910853, 0.0055012372192967465, -0.010779780613239438, -0.00732819304969001]\n",
      "[1.128280719245138e-05, 9.882148731265271e-06, -1.727620055886822e-06, -1.7279881032459046e-06]\n"
     ]
    }
   ],
   "source": [
    "# Tracking (paraxial, kn & ks, alignment)\n",
    "\n",
    "# Note, in paraxial case, MADX seems to truncate cylindrical potential\n",
    "\n",
    "ptc = Path('ptc')\n",
    "obs = Path('track.obs0001.p0001')\n",
    "\n",
    "exact = False\n",
    "align = True\n",
    "\n",
    "length = 2.0\n",
    "angle = 0.05\n",
    "e1 = 0.0\n",
    "e2 = 0.0\n",
    "kn = 0.1\n",
    "ks = 0.1\n",
    "ms = 0.0\n",
    "mo = 0.0\n",
    "dp = 0.001\n",
    "state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)\n",
    "qx, px, qy, py = state.tolist()\n",
    "\n",
    "dx = align*torch.tensor(0.05, dtype=torch.float64)\n",
    "dy = align*torch.tensor(-0.02, dtype=torch.float64)\n",
    "dz = align*torch.tensor(0.05, dtype=torch.float64)\n",
    "\n",
    "wx = align*torch.tensor(0.005, dtype=torch.float64)\n",
    "wy = align*torch.tensor(-0.005, dtype=torch.float64)\n",
    "wz = align*torch.tensor(0.1, dtype=torch.float64)\n",
    "\n",
    "error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}\n",
    "\n",
    "code = f\"\"\"\n",
    "mag:sbend,l={length},angle={angle},e1={e1},e2={e2},knl={{0.0,{kn*length},{ms*length},{mo*length}}},ksl={{0.0,{ks*length}}},kill_ent_fringe=false,kill_exi_fringe=false;\n",
    "map:line=(mag) ;\n",
    "beam,energy=1.0E+6,particle=electron ;\n",
    "set,format=\"20.20f\",\"-20s\" ;\n",
    "use,period=map ;\n",
    "select,flag=error,pattern=\"mag\" ;\n",
    "ealign,dx={dx.item()},dy={dy.item()},ds={dz.item()},dphi={wx.item()},dtheta={wy.item()},dpsi={wz.item()} ;\n",
    "ptc_create_universe,sector_nmul_max=5,sector_nmul=5 ;\n",
    "ptc_create_layout,model=1,method=6,nst=1000,exact={str(exact).lower()} ;\n",
    "ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=true ;\n",
    "ptc_align ;\n",
    "ptc_start,x={qx},px={px},y={qy},py={py},pt={dp},t=0.0 ;\n",
    "ptc_track,icase=5,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}} ;\n",
    "ptc_track_end ;\n",
    "ptc_end ;\n",
    "\"\"\" \n",
    "\n",
    "with ptc.open('w') as stream:\n",
    "    stream.write(code)\n",
    "    \n",
    "system(f'madx < {str(ptc)} > /dev/null')\n",
    "\n",
    "with obs.open('r') as stream:\n",
    "    for line in stream:\n",
    "        continue\n",
    "    _, _, qx, px, qy, py, *_ = line.split()\n",
    "    \n",
    "ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)\n",
    "\n",
    "D = Dipole('D', length=length, angle=angle, e1=e1, e2=e2, kn=kn, ks=ks, ms=ms, mo=mo, dp=dp, exact=exact, order=2, ns=25)\n",
    "res = D(state, alignment=align, data={**D.data(), **error})\n",
    "\n",
    "print(ref.tolist())\n",
    "print(res.tolist())\n",
    "print((ref - res).tolist())\n",
    "\n",
    "ptc.unlink()\n",
    "obs.unlink()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "b413af29-7203-4107-a89d-f601cc8dd214",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.013072849472264409, 0.005416718324594963, -0.010908836090706359, -0.007452351251381879]\n",
      "[0.013061618858281962, 0.0054068762860350235, -0.010906993838412838, -0.007450515116390934]\n",
      "[1.1230613982447443e-05, 9.842038559939453e-06, -1.842252293521307e-06, -1.8361349909453567e-06]\n"
     ]
    }
   ],
   "source": [
    "# Tracking (paraxial, kn & ks, ms & mo, alignment)\n",
    "\n",
    "# Note, in paraxial case, MADX seems to truncate cylindrical potential\n",
    "\n",
    "ptc = Path('ptc')\n",
    "obs = Path('track.obs0001.p0001')\n",
    "\n",
    "exact = False\n",
    "align = True\n",
    "\n",
    "length = 2.0\n",
    "angle = 0.05\n",
    "e1 = 0.0\n",
    "e2 = 0.0\n",
    "kn = 0.1\n",
    "ks = 0.1\n",
    "ms = 0.1\n",
    "mo = 0.1\n",
    "dp = 0.001\n",
    "state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)\n",
    "qx, px, qy, py = state.tolist()\n",
    "\n",
    "dx = align*torch.tensor(0.05, dtype=torch.float64)\n",
    "dy = align*torch.tensor(-0.02, dtype=torch.float64)\n",
    "dz = align*torch.tensor(0.05, dtype=torch.float64)\n",
    "\n",
    "wx = align*torch.tensor(0.005, dtype=torch.float64)\n",
    "wy = align*torch.tensor(-0.005, dtype=torch.float64)\n",
    "wz = align*torch.tensor(0.1, dtype=torch.float64)\n",
    "\n",
    "error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}\n",
    "\n",
    "code = f\"\"\"\n",
    "mag:sbend,l={length},angle={angle},e1={e1},e2={e2},knl={{0.0,{kn*length},{ms*length},{mo*length}}},ksl={{0.0,{ks*length}}},kill_ent_fringe=false,kill_exi_fringe=false;\n",
    "map:line=(mag) ;\n",
    "beam,energy=1.0E+6,particle=electron ;\n",
    "set,format=\"20.20f\",\"-20s\" ;\n",
    "use,period=map ;\n",
    "select,flag=error,pattern=\"mag\" ;\n",
    "ealign,dx={dx.item()},dy={dy.item()},ds={dz.item()},dphi={wx.item()},dtheta={wy.item()},dpsi={wz.item()} ;\n",
    "ptc_create_universe,sector_nmul_max=5,sector_nmul=5 ;\n",
    "ptc_create_layout,model=1,method=6,nst=1000,exact={str(exact).lower()} ;\n",
    "ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=true ;\n",
    "ptc_align ;\n",
    "ptc_start,x={qx},px={px},y={qy},py={py},pt={dp},t=0.0 ;\n",
    "ptc_track,icase=5,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}} ;\n",
    "ptc_track_end ;\n",
    "ptc_end ;\n",
    "\"\"\" \n",
    "\n",
    "with ptc.open('w') as stream:\n",
    "    stream.write(code)\n",
    "    \n",
    "system(f'madx < {str(ptc)} > /dev/null')\n",
    "\n",
    "with obs.open('r') as stream:\n",
    "    for line in stream:\n",
    "        continue\n",
    "    _, _, qx, px, qy, py, *_ = line.split()\n",
    "    \n",
    "ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)\n",
    "\n",
    "D = Dipole('D', length=length, angle=angle, e1=e1, e2=e2, kn=kn, ks=ks, ms=ms, mo=mo, dp=dp, exact=exact, order=2, ns=25)\n",
    "res = D(state, alignment=align, data={**D.data(), **error})\n",
    "\n",
    "print(ref.tolist())\n",
    "print(res.tolist())\n",
    "print((ref - res).tolist())\n",
    "\n",
    "ptc.unlink()\n",
    "obs.unlink()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "e2d46534-c296-482a-9a33-9a250aa80160",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.013030972087034147, 0.005379072261796265, -0.010940636831296861, -0.007479587493733525]\n",
      "[0.013019736470996354, 0.005369216085022485, -0.0109387887264908, -0.007477745021700929]\n",
      "[1.1235616037793758e-05, 9.85617677377957e-06, -1.8481048060618732e-06, -1.8424720325954658e-06]\n"
     ]
    }
   ],
   "source": [
    "# Tracking (paraxial, e1 & e2, kn & ks, ms & mo, alignment)\n",
    "\n",
    "# Note, in paraxial case, MADX seems to truncate cylindrical potential\n",
    "# Note, model dipole doesn't account for multipoles in wedges\n",
    "\n",
    "ptc = Path('ptc')\n",
    "obs = Path('track.obs0001.p0001')\n",
    "\n",
    "exact = False\n",
    "align = True\n",
    "\n",
    "length = 2.0\n",
    "angle = 0.05\n",
    "e1 = 0.025\n",
    "e2 = 0.025\n",
    "kn = 0.1\n",
    "ks = 0.1\n",
    "ms = 0.1\n",
    "mo = 0.1\n",
    "dp = 0.001\n",
    "state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)\n",
    "qx, px, qy, py = state.tolist()\n",
    "\n",
    "dx = align*torch.tensor(0.05, dtype=torch.float64)\n",
    "dy = align*torch.tensor(-0.02, dtype=torch.float64)\n",
    "dz = align*torch.tensor(0.05, dtype=torch.float64)\n",
    "\n",
    "wx = align*torch.tensor(0.005, dtype=torch.float64)\n",
    "wy = align*torch.tensor(-0.005, dtype=torch.float64)\n",
    "wz = align*torch.tensor(0.1, dtype=torch.float64)\n",
    "\n",
    "error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}\n",
    "\n",
    "code = f\"\"\"\n",
    "mag:sbend,l={length},angle={angle},e1={e1},e2={e2},knl={{0.0,{kn*length},{ms*length},{mo*length}}},ksl={{0.0,{ks*length}}},kill_ent_fringe=false,kill_exi_fringe=false;\n",
    "map:line=(mag) ;\n",
    "beam,energy=1.0E+6,particle=electron ;\n",
    "set,format=\"20.20f\",\"-20s\" ;\n",
    "use,period=map ;\n",
    "select,flag=error,pattern=\"mag\" ;\n",
    "ealign,dx={dx.item()},dy={dy.item()},ds={dz.item()},dphi={wx.item()},dtheta={wy.item()},dpsi={wz.item()} ;\n",
    "ptc_create_universe,sector_nmul_max=5,sector_nmul=5 ;\n",
    "ptc_create_layout,model=1,method=6,nst=1000,exact={str(exact).lower()} ;\n",
    "ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=true ;\n",
    "ptc_align ;\n",
    "ptc_start,x={qx},px={px},y={qy},py={py},pt={dp},t=0.0 ;\n",
    "ptc_track,icase=5,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}} ;\n",
    "ptc_track_end ;\n",
    "ptc_end ;\n",
    "\"\"\" \n",
    "\n",
    "with ptc.open('w') as stream:\n",
    "    stream.write(code)\n",
    "    \n",
    "system(f'madx < {str(ptc)} > /dev/null')\n",
    "\n",
    "with obs.open('r') as stream:\n",
    "    for line in stream:\n",
    "        continue\n",
    "    _, _, qx, px, qy, py, *_ = line.split()\n",
    "    \n",
    "ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)\n",
    "\n",
    "D = Dipole('D', length=length, angle=angle, e1=e1, e2=e2, kn=kn, ks=ks, ms=ms, mo=mo, dp=dp, exact=exact, order=2, ns=25)\n",
    "res = D(state, alignment=align, data={**D.data(), **error})\n",
    "\n",
    "print(ref.tolist())\n",
    "print(res.tolist())\n",
    "print((ref - res).tolist())\n",
    "\n",
    "ptc.unlink()\n",
    "obs.unlink()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "663ace91-5321-42ce-83ef-252ca3166142",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.002770798305588665, -0.004651371270674763, -0.007745680488823617, -0.0039921009984161026]\n",
      "[0.0027707983054956944, -0.004651371270673946, -0.007745680488833436, -0.003992100998417344]\n",
      "[9.2970769971501e-14, -8.161873954470877e-16, 9.819402235766717e-15, 1.2411946470614055e-15]\n"
     ]
    }
   ],
   "source": [
    "# Tracking (exact, alignment)\n",
    "\n",
    "ptc = Path('ptc')\n",
    "obs = Path('track.obs0001.p0001')\n",
    "\n",
    "exact = True\n",
    "align = True\n",
    "\n",
    "length = 2.0\n",
    "angle = 0.05\n",
    "e1 = 0.0\n",
    "e2 = 0.0\n",
    "kn = 0.0\n",
    "ks = 0.0\n",
    "ms = 0.0\n",
    "mo = 0.0\n",
    "dp = 0.001\n",
    "state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)\n",
    "qx, px, qy, py = state.tolist()\n",
    "\n",
    "dx = align*torch.tensor(0.05, dtype=torch.float64)\n",
    "dy = align*torch.tensor(-0.02, dtype=torch.float64)\n",
    "dz = align*torch.tensor(0.05, dtype=torch.float64)\n",
    "\n",
    "wx = align*torch.tensor(0.005, dtype=torch.float64)\n",
    "wy = align*torch.tensor(-0.005, dtype=torch.float64)\n",
    "wz = align*torch.tensor(0.1, dtype=torch.float64)\n",
    "\n",
    "error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}\n",
    "\n",
    "code = f\"\"\"\n",
    "mag:sbend,l={length},angle={angle},e1={e1},e2={e2},knl={{0.0,{kn*length},{ms*length},{mo*length}}},ksl={{0.0,{ks*length}}},kill_ent_fringe=false,kill_exi_fringe=false;\n",
    "map:line=(mag) ;\n",
    "beam,energy=1.0E+6,particle=electron ;\n",
    "set,format=\"20.20f\",\"-20s\" ;\n",
    "use,period=map ;\n",
    "select,flag=error,pattern=\"mag\" ;\n",
    "ealign,dx={dx.item()},dy={dy.item()},ds={dz.item()},dphi={wx.item()},dtheta={wy.item()},dpsi={wz.item()} ;\n",
    "ptc_create_universe,sector_nmul_max=5,sector_nmul=5 ;\n",
    "ptc_create_layout,model=1,method=6,nst=1000,exact={str(exact).lower()} ;\n",
    "ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=true ;\n",
    "ptc_align ;\n",
    "ptc_start,x={qx},px={px},y={qy},py={py},pt={dp},t=0.0 ;\n",
    "ptc_track,icase=5,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}} ;\n",
    "ptc_track_end ;\n",
    "ptc_end ;\n",
    "\"\"\" \n",
    "\n",
    "with ptc.open('w') as stream:\n",
    "    stream.write(code)\n",
    "    \n",
    "system(f'madx < {str(ptc)} > /dev/null')\n",
    "\n",
    "with obs.open('r') as stream:\n",
    "    for line in stream:\n",
    "        continue\n",
    "    _, _, qx, px, qy, py, *_ = line.split()\n",
    "    \n",
    "ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)\n",
    "\n",
    "D = Dipole('D', length=length, angle=angle, e1=e1, e2=e2, kn=kn, ks=ks, ms=ms, mo=mo, dp=dp, exact=exact, order=2, ns=25)\n",
    "res = D(state, alignment=align, data={**D.data(), **error})\n",
    "\n",
    "print(ref.tolist())\n",
    "print(res.tolist())\n",
    "print((ref - res).tolist())\n",
    "\n",
    "ptc.unlink()\n",
    "obs.unlink()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "393a9dbe-6eb2-4cee-8cab-a2f3be6d59fd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.013055723507830729, 0.005405610520344896, -0.01089583005089035, -0.007447552679263773]\n",
      "[0.013055723646617522, 0.005405610850033292, -0.010895828250285611, -0.007447551055280752]\n",
      "[-1.387867928220876e-10, -3.2968839637492753e-10, -1.8006047382279622e-09, -1.6239830209763273e-09]\n"
     ]
    }
   ],
   "source": [
    "# Tracking (exact, kn & ks, ms & mo, alignment)\n",
    "\n",
    "# Note, in model dipole cylindrical potential is truncated at octupole\n",
    "\n",
    "ptc = Path('ptc')\n",
    "obs = Path('track.obs0001.p0001')\n",
    "\n",
    "exact = True\n",
    "align = True\n",
    "\n",
    "length = 2.0\n",
    "angle = 0.05\n",
    "e1 = 0.0\n",
    "e2 = 0.0\n",
    "kn = 0.1\n",
    "ks = 0.1\n",
    "ms = 0.1\n",
    "mo = 0.1\n",
    "dp = 0.001\n",
    "state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)\n",
    "qx, px, qy, py = state.tolist()\n",
    "\n",
    "dx = align*torch.tensor(0.05, dtype=torch.float64)\n",
    "dy = align*torch.tensor(-0.02, dtype=torch.float64)\n",
    "dz = align*torch.tensor(0.05, dtype=torch.float64)\n",
    "\n",
    "wx = align*torch.tensor(0.005, dtype=torch.float64)\n",
    "wy = align*torch.tensor(-0.005, dtype=torch.float64)\n",
    "wz = align*torch.tensor(0.1, dtype=torch.float64)\n",
    "\n",
    "error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}\n",
    "\n",
    "code = f\"\"\"\n",
    "mag:sbend,l={length},angle={angle},e1={e1},e2={e2},knl={{0.0,{kn*length},{ms*length},{mo*length}}},ksl={{0.0,{ks*length}}},kill_ent_fringe=false,kill_exi_fringe=false;\n",
    "map:line=(mag) ;\n",
    "beam,energy=1.0E+6,particle=electron ;\n",
    "set,format=\"20.20f\",\"-20s\" ;\n",
    "use,period=map ;\n",
    "select,flag=error,pattern=\"mag\" ;\n",
    "ealign,dx={dx.item()},dy={dy.item()},ds={dz.item()},dphi={wx.item()},dtheta={wy.item()},dpsi={wz.item()} ;\n",
    "ptc_create_universe,sector_nmul_max=5,sector_nmul=5 ;\n",
    "ptc_create_layout,model=1,method=6,nst=1000,exact={str(exact).lower()} ;\n",
    "ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=true ;\n",
    "ptc_align ;\n",
    "ptc_start,x={qx},px={px},y={qy},py={py},pt={dp},t=0.0 ;\n",
    "ptc_track,icase=5,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}} ;\n",
    "ptc_track_end ;\n",
    "ptc_end ;\n",
    "\"\"\" \n",
    "\n",
    "with ptc.open('w') as stream:\n",
    "    stream.write(code)\n",
    "    \n",
    "system(f'madx < {str(ptc)} > /dev/null')\n",
    "\n",
    "with obs.open('r') as stream:\n",
    "    for line in stream:\n",
    "        continue\n",
    "    _, _, qx, px, qy, py, *_ = line.split()\n",
    "    \n",
    "ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)\n",
    "\n",
    "D = Dipole('D', length=length, angle=angle, e1=e1, e2=e2, kn=kn, ks=ks, ms=ms, mo=mo, dp=dp, exact=exact, order=2, ns=25)\n",
    "res = D(state, alignment=align, data={**D.data(), **error})\n",
    "\n",
    "print(ref.tolist())\n",
    "print(res.tolist())\n",
    "print((ref - res).tolist())\n",
    "\n",
    "ptc.unlink()\n",
    "obs.unlink()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "530eced8-8ef0-4afe-9829-665262e75cdb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.013018589129527029, 0.005372361621197678, -0.010919187705983761, -0.007468383469439608]\n",
      "[0.013013881710648583, 0.0053679558070361885, -0.010927596957624465, -0.007474796553899775]\n",
      "[4.707418878445446e-06, 4.4058141614898225e-06, 8.409251640703955e-06, 6.413084460166543e-06]\n"
     ]
    }
   ],
   "source": [
    "# Tracking (exact, e1 & e2, kn & ks, ms & mo, alignment)\n",
    "\n",
    "# Note, in model dipole cylindrical potential is truncated at octupole\n",
    "# Note, model dipole doesn't account for multipoles in wedges\n",
    "\n",
    "ptc = Path('ptc')\n",
    "obs = Path('track.obs0001.p0001')\n",
    "\n",
    "exact = True\n",
    "align = True\n",
    "\n",
    "length = 2.0\n",
    "angle = 0.05\n",
    "e1 = 0.025\n",
    "e2 = 0.025\n",
    "kn = 0.1\n",
    "ks = 0.1\n",
    "ms = 0.1\n",
    "mo = 0.1\n",
    "dp = 0.001\n",
    "state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)\n",
    "qx, px, qy, py = state.tolist()\n",
    "\n",
    "dx = align*torch.tensor(0.05, dtype=torch.float64)\n",
    "dy = align*torch.tensor(-0.02, dtype=torch.float64)\n",
    "dz = align*torch.tensor(0.05, dtype=torch.float64)\n",
    "\n",
    "wx = align*torch.tensor(0.005, dtype=torch.float64)\n",
    "wy = align*torch.tensor(-0.005, dtype=torch.float64)\n",
    "wz = align*torch.tensor(0.1, dtype=torch.float64)\n",
    "\n",
    "error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}\n",
    "\n",
    "code = f\"\"\"\n",
    "mag:sbend,l={length},angle={angle},e1={e1},e2={e2},knl={{0.0,{kn*length},{ms*length},{mo*length}}},ksl={{0.0,{ks*length}}},kill_ent_fringe=false,kill_exi_fringe=false;\n",
    "map:line=(mag) ;\n",
    "beam,energy=1.0E+6,particle=electron ;\n",
    "set,format=\"20.20f\",\"-20s\" ;\n",
    "use,period=map ;\n",
    "select,flag=error,pattern=\"mag\" ;\n",
    "ealign,dx={dx.item()},dy={dy.item()},ds={dz.item()},dphi={wx.item()},dtheta={wy.item()},dpsi={wz.item()} ;\n",
    "ptc_create_universe,sector_nmul_max=5,sector_nmul=5 ;\n",
    "ptc_create_layout,model=1,method=6,nst=1000,exact={str(exact).lower()} ;\n",
    "ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=true ;\n",
    "ptc_align ;\n",
    "ptc_start,x={qx},px={px},y={qy},py={py},pt={dp},t=0.0 ;\n",
    "ptc_track,icase=5,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}} ;\n",
    "ptc_track_end ;\n",
    "ptc_end ;\n",
    "\"\"\" \n",
    "\n",
    "with ptc.open('w') as stream:\n",
    "    stream.write(code)\n",
    "    \n",
    "system(f'madx < {str(ptc)} > /dev/null')\n",
    "\n",
    "with obs.open('r') as stream:\n",
    "    for line in stream:\n",
    "        continue\n",
    "    _, _, qx, px, qy, py, *_ = line.split()\n",
    "    \n",
    "ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)\n",
    "\n",
    "D = Dipole('D', length=length, angle=angle, e1=e1, e2=e2, kn=kn, ks=ks, ms=ms, mo=mo, dp=dp, exact=exact, order=2, ns=25)\n",
    "res = D(state, alignment=align, data={**D.data(), **error})\n",
    "\n",
    "print(ref.tolist())\n",
    "print(res.tolist())\n",
    "print((ref - res).tolist())\n",
    "\n",
    "ptc.unlink()\n",
    "obs.unlink()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "602fe0b1-e79e-4193-aa15-2a434e9a9fdf",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([-0.0020, -0.0066, -0.0026,  0.0011], dtype=torch.float64)\n",
      "\n",
      "tensor([ 0.0044, -0.0061, -0.0038,  0.0013], dtype=torch.float64)\n",
      "\n",
      "tensor([ 0.0130,  0.0054, -0.0109, -0.0075], dtype=torch.float64)\n",
      "\n",
      "tensor([0., 0., 0., 0.], dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "# Deviation/error variables\n",
    "\n",
    "length = 2.0\n",
    "angle = 0.05\n",
    "e1 = 0.025\n",
    "e2 = 0.025\n",
    "kn = 0.1\n",
    "ks = 0.1\n",
    "ms = 0.1\n",
    "mo = 0.1\n",
    "dp = 0.001\n",
    "state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)\n",
    "\n",
    "dx = torch.tensor(0.05, dtype=torch.float64)\n",
    "dy = torch.tensor(-0.02, dtype=torch.float64)\n",
    "dz = torch.tensor(0.05, dtype=torch.float64)\n",
    "\n",
    "wx = torch.tensor(0.005, dtype=torch.float64)\n",
    "wy = torch.tensor(-0.005, dtype=torch.float64)\n",
    "wz = torch.tensor(0.1, dtype=torch.float64)\n",
    "\n",
    "error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}\n",
    "\n",
    "D = Dipole('D', length, angle, e1, e2, kn, ks, ms, mo, dp)\n",
    "\n",
    "# Each element has two variant of a call method\n",
    "# In the first case only state is passed, it is transformed using parameters specified on initializaton\n",
    "\n",
    "print(D(state))\n",
    "print()\n",
    "\n",
    "# Deviation errors can be also passed to call method\n",
    "# These variables are added to corresponding parameters specified on initializaton\n",
    "# For example, element lenght can changed\n",
    "\n",
    "print(D(state, data={**D.data(), **{'dl': -0.5*D.length}}))\n",
    "print()\n",
    "\n",
    "# In the above D.data() creates default deviation dictionary (with zero values for each deviaton)\n",
    "# {**D.data(), **{'dl': -D.length}} replaces the 'dl' key value \n",
    "\n",
    "# Additionaly, alignment errors are passed as deivation variables\n",
    "# They are used if alignment flag is raised\n",
    "\n",
    "print(D(state, data={**D.data(), **error}, alignment=True))\n",
    "print()\n",
    "\n",
    "# The following elements can be made equivalent using deviation variables\n",
    "\n",
    "DA = Dipole('DA', length, angle, e1, e2, kn, ks, ms, mo, dp)\n",
    "DB = Dipole('DB', length - 0.1, angle, e1, e2, kn, ks, ms, mo, dp)\n",
    "\n",
    "print(DA(state) - DB(state, data={**DB.data(), **{'dl': torch.tensor(+0.1, dtype=DB.dtype)}}))\n",
    "\n",
    "# Note, while in some cases float values can be passed as values to deviation variables\n",
    "# The correct behaviour in guaranteed only for tensors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "c357a592-e172-4463-8937-0097d5c626ca",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([ 1.0408e-17, -4.3368e-18, -4.3368e-18,  8.6736e-19],\n",
      "       dtype=torch.float64)\n",
      "tensor([ 9.9380e-06,  4.9995e-05, -1.9980e-06,  6.5052e-19],\n",
      "       dtype=torch.float64)\n",
      "tensor([-2.3859e-06, -5.8600e-06, -7.2738e-07, -2.4947e-07],\n",
      "       dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "# Insertion element\n",
    "\n",
    "import torch\n",
    "from model.library.dipole import Dipole\n",
    "\n",
    "# In this mode elements are treated as thin insertions (at the center)\n",
    "# Using parameters specified on initialization, transport two matrices are computed\n",
    "# These matrices are used to insert the element\n",
    "# Input state is transformed from the element center to its entrance\n",
    "# Next, transformation from the entrance frame to the exit frame is performed\n",
    "# This transformation can contain errors\n",
    "# The final step is to transform state from the exit frame back to the element center\n",
    "# Without errors, this results in identity transformation for linear elements\n",
    "\n",
    "length = 2.0\n",
    "angle = 0.05\n",
    "e1 = 0.025\n",
    "e2 = 0.025\n",
    "kn = 0.0\n",
    "ks = 0.0\n",
    "ms = 0.0\n",
    "mo = 0.0\n",
    "dp = 0.0\n",
    "state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)\n",
    "\n",
    "dx = torch.tensor(0.05, dtype=torch.float64)\n",
    "dy = torch.tensor(-0.02, dtype=torch.float64)\n",
    "dz = torch.tensor(0.05, dtype=torch.float64)\n",
    "\n",
    "wx = torch.tensor(0.005, dtype=torch.float64)\n",
    "wy = torch.tensor(-0.005, dtype=torch.float64)\n",
    "wz = torch.tensor(0.1, dtype=torch.float64)\n",
    "\n",
    "error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}\n",
    "\n",
    "D = Dipole('D', length, angle, e1, e2, kn, ks, ms, mo, dp, exact=False, insertion=True)\n",
    "\n",
    "# For dipole insertion is an identity transformation in the following is true\n",
    "# dp = 0 (chomatic orbit change)\n",
    "# kn & ks and ms & mo are all equal to zero (nonlinearity of cylindrical potential)\n",
    "\n",
    "print(D(state) - state)\n",
    "\n",
    "# Represents effect of an error (any nonzero value of strengh or a change in other parameter)\n",
    "\n",
    "print(D(state, data={**D.data(), **{'dp': 0.001}}) - state)\n",
    "\n",
    "# Exact tracking corresponds to inclusion of kinematic term as errors\n",
    "\n",
    "D = Dipole('D', length, angle, e1, e2, kn, ks, ms, mo, dp, exact=True, insertion=True)\n",
    "\n",
    "print(D(state) - state)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "317dbb7d-0b73-4f36-bbf1-51b530380753",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([512, 4])\n",
      "torch.Size([512, 4])\n"
     ]
    }
   ],
   "source": [
    "# Mapping over a set of initial conditions\n",
    "\n",
    "# Call method can be used to map over a set of initial conditions\n",
    "# Note, device can be set to cpu or gpu via base element classvariables\n",
    "\n",
    "length = 2.0\n",
    "angle = 0.05\n",
    "e1 = 0.025\n",
    "e2 = 0.025\n",
    "kn = 0.5\n",
    "ks = 0.5\n",
    "ms = 1.0\n",
    "mo = 1.0\n",
    "dp = 0.001\n",
    "\n",
    "dx = torch.tensor(0.05, dtype=torch.float64)\n",
    "dy = torch.tensor(-0.02, dtype=torch.float64)\n",
    "dz = torch.tensor(0.05, dtype=torch.float64)\n",
    "\n",
    "wx = torch.tensor(0.005, dtype=torch.float64)\n",
    "wy = torch.tensor(-0.005, dtype=torch.float64)\n",
    "wz = torch.tensor(0.1, dtype=torch.float64)\n",
    "\n",
    "error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}\n",
    "\n",
    "D = Dipole('D', length, angle, e1, e2, kn, ks, ms, mo, dp)\n",
    "\n",
    "state = 1.0E-3*torch.randn((512, 4), dtype=D.dtype, device=D.device)\n",
    "\n",
    "print(torch.vmap(D)(state).shape)\n",
    "\n",
    "# To map over deviations parameters a wrapper function (or a lambda expression) can be used\n",
    "\n",
    "def wrapper(state, dp):\n",
    "    return D(state, data={**D.data(), **{'dp': dp}})\n",
    "\n",
    "dp = 1.0E-3*torch.randn(512, dtype=D.dtype, device=D.device)\n",
    "\n",
    "print(torch.vmap(wrapper)(state, dp).shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "2c21e064-e2af-4433-8da5-1d81a0cc841d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[ 0.3083,  1.4545,  1.0125,  0.6652],\n",
      "        [-0.3993,  0.3083,  1.0602,  1.0168],\n",
      "        [ 1.0168,  0.6652,  2.3568,  2.8065],\n",
      "        [ 1.0602,  1.0125,  1.7383,  2.3568]], dtype=torch.float64)\n",
      "\n",
      "tensor([0.0004, 0.0002, 0.0005, 0.0005], dtype=torch.float64)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Differentiability\n",
    "\n",
    "# Both call methods are differentiable\n",
    "# Derivative with respect to state can be computed directly\n",
    "# For deviation variables, wrapping is required\n",
    "\n",
    "length = 2.0\n",
    "angle = 0.05\n",
    "e1 = 0.025\n",
    "e2 = 0.025\n",
    "kn = 0.5\n",
    "ks = 0.5\n",
    "ms = 1.0\n",
    "mo = 1.0\n",
    "dp = 0.001\n",
    "state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)\n",
    "\n",
    "dx = torch.tensor(0.05, dtype=torch.float64)\n",
    "dy = torch.tensor(-0.02, dtype=torch.float64)\n",
    "dz = torch.tensor(0.05, dtype=torch.float64)\n",
    "\n",
    "wx = torch.tensor(0.005, dtype=torch.float64)\n",
    "wy = torch.tensor(-0.005, dtype=torch.float64)\n",
    "wz = torch.tensor(0.1, dtype=torch.float64)\n",
    "\n",
    "error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}\n",
    "\n",
    "D = Dipole('D', length, angle, e1, e2, kn, ks, ms, mo, dp)\n",
    "\n",
    "# Compute derivative with respect to state\n",
    "\n",
    "print(torch.func.jacrev(D)(state))\n",
    "print()\n",
    "\n",
    "# Compute derivative with respect to a deviation variable\n",
    "\n",
    "e1 = torch.tensor(0.0, dtype=torch.float64)\n",
    "\n",
    "def wrapper(state, e1):\n",
    "    return D(state, data={**D.data(), **{'e1': e1}})\n",
    "\n",
    "print(torch.func.jacrev(wrapper, 1)(state, e1))\n",
    "print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "43d9d831-edd0-41f1-aea5-415c1709882b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([-0.0086, -0.0098, -0.0022, -0.0008], dtype=torch.float64)\n",
      "torch.Size([10, 4])\n",
      "torch.Size([10, 4, 4])\n",
      "torch.Size([100, 4])\n",
      "torch.Size([100, 4, 4])\n"
     ]
    }
   ],
   "source": [
    "# Output at each step\n",
    "\n",
    "# It is possible to collect output of state or tangent matrix at each integration step\n",
    "# Number of integratin steps is controlled by ns parameter on initialization\n",
    "# Alternatively, desired integration step length can be passed\n",
    "# Number of integration steps is computed as ceil(length/ds)\n",
    "\n",
    "length = 2.0\n",
    "angle = 0.05\n",
    "e1 = 0.025\n",
    "e2 = 0.025\n",
    "kn = 0.5\n",
    "ks = 0.5\n",
    "ms = 1.0\n",
    "mo = 1.0\n",
    "dp = 0.001\n",
    "state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)\n",
    "\n",
    "dx = torch.tensor(0.05, dtype=torch.float64)\n",
    "dy = torch.tensor(-0.02, dtype=torch.float64)\n",
    "dz = torch.tensor(0.05, dtype=torch.float64)\n",
    "\n",
    "wx = torch.tensor(0.005, dtype=torch.float64)\n",
    "wy = torch.tensor(-0.005, dtype=torch.float64)\n",
    "wz = torch.tensor(0.1, dtype=torch.float64)\n",
    "\n",
    "error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}\n",
    "\n",
    "D = Dipole('D', length, angle, e1, e2, kn, ks, ms, mo, dp, exact=False, ns=10, output=True, matrix=True)\n",
    "\n",
    "# Final state is still returned\n",
    "\n",
    "print(D(state))\n",
    "\n",
    "# Data is added to special attributes (state and tangent matrix)\n",
    "\n",
    "print(D.container_output.shape)\n",
    "print(D.container_matrix.shape)\n",
    "\n",
    "# Number of integration steps can be changed\n",
    "\n",
    "D.ns = 100\n",
    "\n",
    "D(state)\n",
    "print(D.container_output.shape)\n",
    "print(D.container_matrix.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "fe4a7227-f2a4-4986-ac44-f22e2d55a69a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[-0.00858751580977462, -0.009819760091907982, -0.0021662845096666727, -0.0008078063877268448]\n",
      "[-0.008587660168877631, -0.00982000432362944, -0.0021667806526843485, -0.0008081645802833406]\n",
      "[1.443591030117486e-07, 2.442317214579326e-07, 4.961430176758305e-07, 3.5819255649581234e-07]\n",
      "\n",
      "[0, 1, 2, 3, 4, 3, 2, 1, 0]\n",
      "[0.5, 0.5, 0.5, 0.5, 1.0, 0.5, 0.5, 0.5, 0.5]\n"
     ]
    }
   ],
   "source": [
    "# Integration order is set on initialization (default value is zero)\n",
    "# This order is related to difference order as 2n + 2\n",
    "# Thus, zero corresponds to second order difference method\n",
    "\n",
    "length = 2.0\n",
    "angle = 0.05\n",
    "e1 = 0.025\n",
    "e2 = 0.025\n",
    "kn = 0.5\n",
    "ks = 0.5\n",
    "ms = 1.0\n",
    "mo = 1.0\n",
    "dp = 0.001\n",
    "state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)\n",
    "\n",
    "dx = torch.tensor(0.05, dtype=torch.float64)\n",
    "dy = torch.tensor(-0.02, dtype=torch.float64)\n",
    "dz = torch.tensor(0.05, dtype=torch.float64)\n",
    "\n",
    "wx = torch.tensor(0.005, dtype=torch.float64)\n",
    "wy = torch.tensor(-0.005, dtype=torch.float64)\n",
    "wz = torch.tensor(0.1, dtype=torch.float64)\n",
    "\n",
    "error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}\n",
    "\n",
    "D = Dipole('D', length, angle, e1, e2, kn, ks, ms, mo, dp, exact=True, ns=10)\n",
    "\n",
    "# For dipole integration is always performed\n",
    "# In exact case, kinematic term error is added\n",
    "\n",
    "D.ns = 10\n",
    "ref = D(state)\n",
    "\n",
    "D.ns = 100\n",
    "res = D(state)\n",
    "\n",
    "print(ref.tolist())\n",
    "print(res.tolist())\n",
    "print((ref - res).tolist())\n",
    "print()\n",
    "\n",
    "# Integrator parameters are stored in data attribute (if integration is actually performed)\n",
    "\n",
    "maps, weights = D._data\n",
    "print(maps)\n",
    "print(weights)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "eb6d377d-4e6a-4a73-aec8-a95c711b25ad",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([0.2210, 0.1703], dtype=torch.float64)\n",
      "tensor([-0.2310, -0.2107], dtype=torch.float64)\n",
      "tensor([[ 1.6013,  0.3613],\n",
      "        [-0.7593, -1.2452]], dtype=torch.float64)\n",
      "tensor([-1.1102e-16,  1.3878e-16], dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "# Derivatives of twiss parameters (chromaticity)\n",
    "\n",
    "# pip install git+https://github.com/i-a-morozov/twiss.git@main\n",
    "# pip install git+https://github.com/i-a-morozov/ndmap.git@main\n",
    "\n",
    "from twiss import twiss\n",
    "\n",
    "from ndmap.pfp import parametric_fixed_point\n",
    "from ndmap.evaluate import evaluate\n",
    "\n",
    "QF = Quadrupole('QF', 0.5, +0.21)\n",
    "QD = Quadrupole('QD', 0.5, -0.19)\n",
    "SF = Sextupole('SF', 0.25)\n",
    "SD = Sextupole('SD', 0.25)\n",
    "DR = Drift('DR', 0.25)\n",
    "BM = Dipole('DR', 3.50, 0.15)\n",
    "\n",
    "def fodo(state, dp, ms):\n",
    "    dp, *_ = dp\n",
    "    msf, msd, *_ = ms\n",
    "    state = QF(state, data={**QF.data(), **{'dp': dp}})\n",
    "    state = DR(state, data={**DR.data(), **{'dp': dp}})\n",
    "    state = SF(state, data={**SF.data(), **{'dp': dp, 'ms': msf}})\n",
    "    state = DR(state, data={**DR.data(), **{'dp': dp}})\n",
    "    state = BM(state, data={**BM.data(), **{'dp': dp}})\n",
    "    state = DR(state, data={**DR.data(), **{'dp': dp}})\n",
    "    state = SD(state, data={**SD.data(), **{'dp': dp, 'ms': msd}})\n",
    "    state = DR(state, data={**DR.data(), **{'dp': dp}})\n",
    "    state = QD(state, data={**QD.data(), **{'dp': dp}})\n",
    "    state = QD(state, data={**QD.data(), **{'dp': dp}})\n",
    "    state = DR(state, data={**DR.data(), **{'dp': dp}})\n",
    "    state = SD(state, data={**SD.data(), **{'dp': dp, 'ms': msd}})\n",
    "    state = DR(state, data={**DR.data(), **{'dp': dp}})\n",
    "    state = BM(state, data={**BM.data(), **{'dp': dp}})\n",
    "    state = DR(state, data={**DR.data(), **{'dp': dp}})\n",
    "    state = SF(state, data={**SF.data(), **{'dp': dp, 'ms': msf}})\n",
    "    state = DR(state, data={**DR.data(), **{'dp': dp}})\n",
    "    state = QF(state, data={**QF.data(), **{'dp': dp}})\n",
    "    return state\n",
    "\n",
    "# Set deviation parameters\n",
    "\n",
    "msf = torch.tensor(0.0, dtype=torch.float64)\n",
    "msd = torch.tensor(0.0, dtype=torch.float64)\n",
    "ms = torch.stack([msf, msd])\n",
    "dp = torch.tensor([0.0], dtype=torch.float64)\n",
    "\n",
    "# Set fixed point\n",
    "\n",
    "fp = torch.tensor([0.0, 0.0, 0.0, 0.0], dtype=torch.float64)\n",
    "\n",
    "\n",
    "# Compute parametrix fixed point (first order in momentum deviation)\n",
    "# Note, all parameters must be vectors\n",
    "\n",
    "pfp, *_ = parametric_fixed_point((1, ), fp, [dp], fodo, ms)\n",
    "\n",
    "# Define transformation around fixed point\n",
    "\n",
    "def pfp_fodo(state, dp, ms):\n",
    "    return fodo(state + evaluate(pfp, [dp]), dp, ms) - evaluate(pfp, [dp])\n",
    "\n",
    "# Tune\n",
    "\n",
    "def tune(dp, ms):\n",
    "    matrix = torch.func.jacrev(pfp_fodo)(fp, dp, ms)\n",
    "    tune, *_ = twiss(matrix)\n",
    "    return tune\n",
    "\n",
    "# Chromaticity\n",
    "\n",
    "def chromaticity(ms):\n",
    "    return torch.func.jacrev(tune)(dp, ms)\n",
    "\n",
    "# Compute tunes\n",
    "\n",
    "tunes = tune(dp, ms)\n",
    "print(tunes)\n",
    "\n",
    "\n",
    "# Compute chromaticity\n",
    "\n",
    "chromaticities = chromaticity(ms).squeeze()\n",
    "print(chromaticities)\n",
    "\n",
    "# Compute derivative of chromaticities \n",
    "\n",
    "jacobian = torch.func.jacrev(chromaticity)(ms).squeeze()\n",
    "print(jacobian)\n",
    "\n",
    "# Correct chomaticity\n",
    "\n",
    "print((chromaticity(ms - torch.linalg.pinv(jacobian) @ chromaticities)).squeeze())"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [
    "myt0_gMIOq7b",
    "5d97819c"
   ],
   "name": "03_frequency.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "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.12.1"
  },
  "latex_envs": {
   "LaTeX_envs_menu_present": true,
   "autoclose": false,
   "autocomplete": true,
   "bibliofile": "biblio.bib",
   "cite_by": "apalike",
   "current_citInitial": 1,
   "eqLabelWithNumbers": true,
   "eqNumInitial": 1,
   "hotkeys": {
    "equation": "Ctrl-E",
    "itemize": "Ctrl-I"
   },
   "labels_anchors": false,
   "latex_user_defs": false,
   "report_style_numbering": false,
   "user_envs_cfg": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}