{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "4604a99e-2f05-4693-85f2-f0ea61312fd6",
   "metadata": {},
   "source": [
    "# Example-04: Transformations benchmark (PTC)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "080a6c07-fab4-4a55-96dd-6276bb36237a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# In this example various symplectic transformations are compared with corresponding MADX-PTC transformations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "bdb8e52a-634f-42c3-9628-603df039e0e5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.01, -0.005, -0.01, 0.005]\n",
      "[0.010000000000000002, -0.005, -0.010000000000000002, 0.005]\n",
      "[-1.734723475976807e-18, 0.0, 1.734723475976807e-18, 0.0]\n"
     ]
    }
   ],
   "source": [
    "# calibration\n",
    "\n",
    "import torch\n",
    "\n",
    "from model.library.transformations import calibration_forward\n",
    "from model.library.transformations import calibration_inverse\n",
    "\n",
    "gxx = torch.tensor(1.005, dtype=torch.float64)\n",
    "gxy = torch.tensor(0.001, dtype=torch.float64)\n",
    "gyx = torch.tensor(0.005, dtype=torch.float64)\n",
    "gyy = torch.tensor(0.955, dtype=torch.float64)\n",
    "\n",
    "state = torch.tensor([0.01, -0.005, -0.01, 0.005], dtype=torch.float64)\n",
    "\n",
    "res = calibration_forward(state, gxx, gxy, gyx, gyy)\n",
    "res = calibration_inverse(res, gxx, gxy, gyx, gyy)\n",
    "\n",
    "print(state.tolist()) \n",
    "print(res.tolist())\n",
    "print((state - res).tolist()) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "4d5b5046-cafc-45b0-824d-4ab2d787ac68",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.005024875621890547, -0.005, -0.005024875621890547, 0.005]\n",
      "[0.005024875621890547, -0.005, -0.005024875621890547, 0.005]\n",
      "[0.0, 0.0, 0.0, 0.0]\n"
     ]
    }
   ],
   "source": [
    "# drift\n",
    "\n",
    "from pathlib import Path\n",
    "from os import system\n",
    "\n",
    "import torch\n",
    "from model.library.transformations import drift\n",
    "\n",
    "ptc = Path('ptc')\n",
    "obs = Path('track.obs0001.p0001')\n",
    "\n",
    "dp = torch.tensor(0.005, dtype=torch.float64)\n",
    "length = torch.tensor(1.0, dtype=torch.float64)\n",
    "state = torch.tensor([0.01, -0.005, -0.01, 0.005], dtype=torch.float64)\n",
    "\n",
    "qx, px, qy, py = state.tolist()\n",
    "\n",
    "code = f\"\"\"\n",
    "mag:drift,l={length.item()} ;\n",
    "map:line=(mag) ;\n",
    "beam,energy=1.0E+6,particle=electron ;\n",
    "set,format=\"20.20f\",\"-20s\" ;\n",
    "use,period=map  ;\n",
    "ptc_create_universe,sector_nmul_max=10,sector_nmul=10 ;\n",
    "ptc_create_layout,model=1,method=6,nst=1000,exact=false ;\n",
    "ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=false ;\n",
    "ptc_align;\n",
    "ptc_start,x={qx},px={px},y={qy},py={py},pt={dp.item()},t=0.0 ;\n",
    "ptc_track,icase=6,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",
    "res = drift(state, dp, length)\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": "9b7001fe-8e1e-49e1-9eae-31ab093ca099",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.005024752473723395, -0.005, -0.005024752473723395, 0.005]\n",
      "[0.005024752473723394, -0.005, -0.005024752473723394, 0.005]\n",
      "[8.673617379884035e-19, 0.0, -8.673617379884035e-19, 0.0]\n"
     ]
    }
   ],
   "source": [
    "# drift (with kinematic)\n",
    "\n",
    "from pathlib import Path\n",
    "from os import system\n",
    "\n",
    "import torch\n",
    "from model.library.transformations import drift\n",
    "from model.library.transformations import kinematic\n",
    "\n",
    "ptc = Path('ptc')\n",
    "obs = Path('track.obs0001.p0001')\n",
    "\n",
    "dp = torch.tensor(0.005, dtype=torch.float64)\n",
    "length = torch.tensor(1.0, dtype=torch.float64)\n",
    "state = torch.tensor([0.01, -0.005, -0.01, 0.005], dtype=torch.float64)\n",
    "\n",
    "qx, px, qy, py = state.tolist()\n",
    "\n",
    "code = f\"\"\"\n",
    "mag:drift,l={length.item()} ;\n",
    "map:line=(mag) ;\n",
    "beam,energy=1.0E+6,particle=electron ;\n",
    "set,format=\"20.20f\",\"-20s\" ;\n",
    "use,period=map  ;\n",
    "ptc_create_universe,sector_nmul_max=10,sector_nmul=10 ;\n",
    "ptc_create_layout,model=1,method=6,nst=1000,exact=true ;\n",
    "ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=false ;\n",
    "ptc_align;\n",
    "ptc_start,x={qx},px={px},y={qy},py={py},pt={dp.item()},t=0.0 ;\n",
    "ptc_track,icase=6,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",
    "state = drift(state, dp, length)\n",
    "state = kinematic(state, dp, length)\n",
    "res = state\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": "91fd9e07-7966-4ca9-8c0f-2da9b966c091",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.0, -0.1, 0.0, 0.1]\n",
      "[0.0, -0.1, 0.0, 0.1]\n",
      "[0.0, 0.0, 0.0, 0.0]\n"
     ]
    }
   ],
   "source": [
    "# corrector\n",
    "\n",
    "from pathlib import Path\n",
    "from os import system\n",
    "\n",
    "import torch\n",
    "from model.library.transformations import corrector\n",
    "\n",
    "ptc = Path('ptc')\n",
    "obs = Path('track.obs0001.p0001')\n",
    "\n",
    "dp = torch.tensor(0.005, dtype=torch.float64)\n",
    "kx = -torch.tensor(0.1, dtype=torch.float64)\n",
    "ky = +torch.tensor(0.1, dtype=torch.float64)\n",
    "state = torch.tensor([0.0, 0.0, 0.0, 0.0], dtype=torch.float64)\n",
    "\n",
    "qx, px, qy, py = state.tolist()\n",
    "\n",
    "code = f\"\"\"\n",
    "magx:hkicker,l=0.0,kick={kx.item()};\n",
    "magy:vkicker,l=0.0,kick={ky.item()};\n",
    "map:line=(magx, magy) ;\n",
    "beam,energy=1.0E+6,particle=electron ;\n",
    "set,format=\"20.20f\",\"-20s\" ;\n",
    "use,period=map  ;\n",
    "ptc_create_universe,sector_nmul_max=10,sector_nmul=10 ;\n",
    "ptc_create_layout,model=1,method=6,nst=1000,exact=false ;\n",
    "ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=false ;\n",
    "ptc_align;\n",
    "ptc_start,x={qx},px={px},y={qy},py={py},pt={dp.item()},t=0.0 ;\n",
    "ptc_track,icase=6,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",
    "res = corrector(state, kx, ky)\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": "70c04db2-1d10-4d61-b82c-0a0c7e294ef9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.0012338134845463512, -0.011134185772061586, -0.009559363509475737, -0.004042070986490389]\n",
      "[0.0012338134845463642, -0.01113418577206157, -0.009559363509475702, -0.0040420709864903495]\n",
      "[-1.3010426069826053e-17, -1.5612511283791264e-17, -3.469446951953614e-17, -3.9898639947466563e-17]\n"
     ]
    }
   ],
   "source": [
    "# focusing quadrupole\n",
    "\n",
    "from pathlib import Path\n",
    "from os import system\n",
    "\n",
    "ptc = Path('ptc')\n",
    "obs = Path('track.obs0001.p0001')\n",
    "\n",
    "import torch\n",
    "from model.library.transformations import fquad\n",
    "\n",
    "kn = torch.tensor(1.0, dtype=torch.float64)\n",
    "dp = torch.tensor(0.005, dtype=torch.float64)\n",
    "length = torch.tensor(1.0, dtype=torch.float64)\n",
    "state = torch.tensor([0.01, -0.005, -0.01, 0.005], dtype=torch.float64)\n",
    "\n",
    "qx, px, qy, py = state.tolist()\n",
    "\n",
    "code = f\"\"\"\n",
    "mag:quadrupole,l={length.item()}, k1={kn.item()};\n",
    "map:line=(mag) ;\n",
    "beam,energy=1.0E+9,particle=electron ;\n",
    "set,format=\"20.20f\",\"-20s\" ;\n",
    "use,period=map  ;\n",
    "ptc_create_universe,sector_nmul_max=10,sector_nmul=10 ;\n",
    "ptc_create_layout,model=1,method=6,nst=1000,exact=false ;\n",
    "ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=false ;\n",
    "ptc_align;\n",
    "ptc_start,x={qx},px={px},y={qy},py={py},pt={dp.item()},t=0.0 ;\n",
    "ptc_track,icase=6,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",
    "res = fquad(state, kn.abs(), dp, length)\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": "b2856de4-d7ff-4b08-9a71-2ff691c14770",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.009559363509475737, 0.004042070986490389, -0.0012338134845463512, 0.011134185772061586]\n",
      "[0.009559363509475702, 0.0040420709864903495, -0.0012338134845463642, 0.01113418577206157]\n",
      "[3.469446951953614e-17, 3.9898639947466563e-17, 1.3010426069826053e-17, 1.5612511283791264e-17]\n"
     ]
    }
   ],
   "source": [
    "# defocusing quadrupole\n",
    "\n",
    "from pathlib import Path\n",
    "from os import system\n",
    "\n",
    "import torch\n",
    "from model.library.transformations import dquad\n",
    "\n",
    "ptc = Path('ptc')\n",
    "obs = Path('track.obs0001.p0001')\n",
    "\n",
    "kn = torch.tensor(-1.0, dtype=torch.float64)\n",
    "dp = torch.tensor(0.005, dtype=torch.float64)\n",
    "length = torch.tensor(1.0, dtype=torch.float64)\n",
    "state = torch.tensor([0.01, -0.005, -0.01, 0.005], dtype=torch.float64)\n",
    "\n",
    "qx, px, qy, py = state.tolist()\n",
    "\n",
    "code = f\"\"\"\n",
    "mag:quadrupole,l={length.item()}, k1={kn.item()};\n",
    "map:line=(mag) ;\n",
    "beam,energy=1.0E+9,particle=electron ;\n",
    "set,format=\"20.20f\",\"-20s\" ;\n",
    "use,period=map  ;\n",
    "ptc_create_universe,sector_nmul_max=10,sector_nmul=10 ;\n",
    "ptc_create_layout,model=1,method=6,nst=1000,exact=false ;\n",
    "ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=false ;\n",
    "ptc_align;\n",
    "ptc_start,x={qx},px={px},y={qy},py={py},pt={dp.item()},t=0.0 ;\n",
    "ptc_track,icase=6,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",
    "res = dquad(state, kn.abs(), dp, length)\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": "79f546f6-73cc-4e29-88d8-28571c7ab15c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.00576876084434182, -0.0020884471766207963, 0.002582224887596925, 0.017416187766128223]\n",
      "[0.005768760844341825, -0.0020884471766207725, 0.0025822248875969544, 0.0174161877661282]\n",
      "[-5.204170427930421e-18, -2.3852447794681098e-17, -2.949029909160572e-17, 2.42861286636753e-17]\n"
     ]
    }
   ],
   "source": [
    "# generic quadrupole\n",
    "\n",
    "from pathlib import Path\n",
    "from os import system\n",
    "\n",
    "import torch\n",
    "from model.library.transformations import quadrupole\n",
    "\n",
    "ptc = Path('ptc')\n",
    "obs = Path('track.obs0001.p0001')\n",
    "\n",
    "kn = torch.tensor(-1.0, dtype=torch.float64)\n",
    "ks = torch.tensor(+1.0, dtype=torch.float64)\n",
    "dp = torch.tensor(0.005, dtype=torch.float64)\n",
    "length = torch.tensor(1.0, dtype=torch.float64)\n",
    "state = torch.tensor([0.01, -0.005, -0.01, 0.005], dtype=torch.float64)\n",
    "\n",
    "qx, px, qy, py = state.tolist()\n",
    "\n",
    "code = f\"\"\"\n",
    "mag: quadrupole, l={length.item()},k1={kn.item()},k1s={ks.item()};\n",
    "map:line=(mag) ;\n",
    "beam,energy=1.0E+9,particle=electron ;\n",
    "set,format=\"20.20f\",\"-20s\" ;\n",
    "use,period=map  ;\n",
    "ptc_create_universe,sector_nmul_max=10,sector_nmul=10 ;\n",
    "ptc_create_layout,model=1,method=6,nst=1000,exact=false ;\n",
    "ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=false ;\n",
    "ptc_align;\n",
    "ptc_start,x={qx},px={px},y={qy},py={py},pt={dp.item()},t=0.0 ;\n",
    "ptc_track,icase=6,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",
    "res = quadrupole(state, kn, ks, dp, length)\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": "7d079fa0-3bf2-4b11-b44e-e7cc7abb4b6b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.005768760844341825, -0.0020884471766207725, 0.0025822248875969544, 0.0174161877661282]\n",
      "[0.005768760844341825, -0.0020884471766207725, 0.0025822248875969544, 0.0174161877661282]\n",
      "[0.0, 0.0, 0.0, 0.0]\n"
     ]
    }
   ],
   "source": [
    "# generic linear transformation\n",
    "\n",
    "import torch\n",
    "from model.library.transformations import quadrupole\n",
    "from model.library.transformations import linear\n",
    "\n",
    "kn = torch.tensor(-1.0, dtype=torch.float64)\n",
    "ks = torch.tensor(+1.0, dtype=torch.float64)\n",
    "dp = torch.tensor(0.005, dtype=torch.float64)\n",
    "length = torch.tensor(1.0, dtype=torch.float64)\n",
    "state = torch.tensor([0.01, -0.005, -0.01, 0.005], dtype=torch.float64)\n",
    "\n",
    "vector = torch.zeros_like(state)\n",
    "matrix = torch.func.jacrev(quadrupole)(0.0*state, kn, ks, dp, length)\n",
    "\n",
    "ref = quadrupole(state, kn, ks, dp, length)\n",
    "res = linear(state, vector, matrix)\n",
    "\n",
    "print(ref.tolist())\n",
    "print(res.tolist())\n",
    "print((ref - res).tolist())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "3a088826-da54-4bd5-9df5-f7a88d4b0c1e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.01, -0.07, -0.05, -0.06400000000000002]\n",
      "[0.01, -0.07, -0.05, -0.06400000000000002]\n",
      "[0.0, 0.0, 0.0, 0.0]\n"
     ]
    }
   ],
   "source": [
    "# quadrupole kick\n",
    "\n",
    "from pathlib import Path\n",
    "from os import system\n",
    "\n",
    "import torch\n",
    "from model.library.transformations import gradient\n",
    "\n",
    "ptc = Path('ptc')\n",
    "obs = Path('track.obs0001.p0001')\n",
    "\n",
    "kn = torch.tensor(1.5, dtype=torch.float64)\n",
    "ks = torch.tensor(1.0, dtype=torch.float64)\n",
    "dp = torch.tensor(0.005, dtype=torch.float64)\n",
    "length = torch.tensor(1.0, dtype=torch.float64)\n",
    "state = torch.tensor([0.01, -0.005, -0.05, 0.001], dtype=torch.float64)\n",
    "\n",
    "qx, px, qy, py = state.tolist()\n",
    "\n",
    "code = f\"\"\"\n",
    "mag: multipole,knl={{0.0,{(kn*length).item()}}},ksl={{0.0,{(ks*length).item()}}};\n",
    "map:line=(mag) ;\n",
    "beam,energy=1.0E+9,particle=electron ;\n",
    "set,format=\"20.20f\",\"-20s\" ;\n",
    "use,period=map  ;\n",
    "ptc_create_universe,sector_nmul_max=10,sector_nmul=10 ;\n",
    "ptc_create_layout,model=1,method=6,nst=1000,exact=false ;\n",
    "ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=false ;\n",
    "ptc_start,x={qx},px={px},y={qy},py={py},pt={dp.item()},t=0.0 ;\n",
    "ptc_track,icase=6,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",
    "res = gradient(state, kn, ks, length)\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": "5295eaba-09a7-45b8-9a89-b821b837c760",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.01, -0.0044, -0.05, 0.00075]\n",
      "[0.01, -0.0044, -0.05, 0.00075]\n",
      "[0.0, 0.0, 0.0, 0.0]\n"
     ]
    }
   ],
   "source": [
    "# sextupole kick\n",
    "\n",
    "from pathlib import Path\n",
    "from os import system\n",
    "\n",
    "import torch\n",
    "from model.library.transformations import sextupole\n",
    "\n",
    "ptc = Path('ptc')\n",
    "obs = Path('track.obs0001.p0001')\n",
    "\n",
    "ks = torch.tensor(0.5, dtype=torch.float64)\n",
    "dp = torch.tensor(0.005, dtype=torch.float64)\n",
    "length = torch.tensor(1.0, dtype=torch.float64)\n",
    "state = torch.tensor([0.01, -0.005, -0.05, 0.001], dtype=torch.float64)\n",
    "\n",
    "qx, px, qy, py = state.tolist()\n",
    "\n",
    "code = f\"\"\"\n",
    "mag: multipole,knl={{0.0,0.0,{(ks*length).item()}}};\n",
    "map:line=(mag) ;\n",
    "beam,energy=1.0E+9,particle=electron ;\n",
    "set,format=\"20.20f\",\"-20s\" ;\n",
    "use,period=map  ;\n",
    "ptc_create_universe,sector_nmul_max=10,sector_nmul=10 ;\n",
    "ptc_create_layout,model=1,method=6,nst=1000,exact=false ;\n",
    "ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=false ;\n",
    "ptc_start,x={qx},px={px},y={qy},py={py},pt={dp.item()},t=0.0 ;\n",
    "ptc_track,icase=6,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",
    "res = sextupole(state, ks, length)\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": "be76c431-62ec-4d08-b33f-dcc3840ae08a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.01, -0.004938333333333334, -0.05, 0.0010916666666666668]\n",
      "[0.01, -0.004938333333333334, -0.05, 0.0010916666666666668]\n",
      "[0.0, 0.0, 0.0, 0.0]\n"
     ]
    }
   ],
   "source": [
    "# octupole kick\n",
    "\n",
    "import torch\n",
    "from model.library.transformations import octupole\n",
    "\n",
    "from pathlib import Path\n",
    "from os import system\n",
    "\n",
    "ptc = Path('ptc')\n",
    "obs = Path('track.obs0001.p0001')\n",
    "\n",
    "ko = torch.tensor(5.0, dtype=torch.float64)\n",
    "dp = torch.tensor(0.005, dtype=torch.float64)\n",
    "length = torch.tensor(1.0, dtype=torch.float64)\n",
    "state = torch.tensor([0.01, -0.005, -0.05, 0.001], dtype=torch.float64)\n",
    "\n",
    "qx, px, qy, py = state.tolist()\n",
    "\n",
    "code = f\"\"\"\n",
    "mag: multipole,knl={{0.0,0.0,0.0,{(ko*length).item()}}};\n",
    "map:line=(mag) ;\n",
    "beam,energy=1.0E+9,particle=electron ;\n",
    "set,format=\"20.20f\",\"-20s\" ;\n",
    "use,period=map  ;\n",
    "ptc_create_universe,sector_nmul_max=10,sector_nmul=10 ;\n",
    "ptc_create_layout,model=1,method=6,nst=1000,exact=false ;\n",
    "ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=false ;\n",
    "ptc_start,x={qx},px={px},y={qy},py={py},pt={dp.item()},t=0.0 ;\n",
    "ptc_track,icase=6,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",
    "res = octupole(state, ko, length)\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": "f5b562d5-b7ca-4fea-b3ab-0103a626c69c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.005231962156363519, -0.004575808017791941, -0.005024875621890682, 0.005]\n",
      "[0.005231962156363559, -0.004575808017791928, -0.005024875621890547, 0.005]\n",
      "[-3.9898639947466563e-17, -1.3010426069826053e-17, -1.3530843112619095e-16, 0.0]\n"
     ]
    }
   ],
   "source": [
    "# pure dipole (no edge effects if e1 = e2 = 0)\n",
    "\n",
    "from pathlib import Path\n",
    "from os import system\n",
    "\n",
    "import torch\n",
    "from model.library.transformations import dipole\n",
    "\n",
    "ptc = Path('ptc')\n",
    "obs = Path('track.obs0001.p0001')\n",
    "\n",
    "angle = torch.tensor(0.1, dtype=torch.float64 )\n",
    "dp = torch.tensor(0.005, dtype=torch.float64)\n",
    "length = torch.tensor(1.0, dtype=torch.float64)\n",
    "state = torch.tensor([0.01, -0.005, -0.01, 0.005], dtype=torch.float64)\n",
    "\n",
    "qx, px, qy, py = state.tolist()\n",
    "\n",
    "code = f\"\"\"\n",
    "mag: sbend, l={length.item()}, angle={angle.item()},k1=0.0,k1s=0.0,e1=0.0,e2=0.0,kill_ent_fringe=false,kill_exi_fringe=false;\n",
    "map:line=(mag) ;\n",
    "beam,energy=1.0E+9,particle=electron ;\n",
    "set,format=\"20.20f\",\"-20s\" ;\n",
    "use,period=map  ;\n",
    "ptc_create_universe,sector_nmul_max=10,sector_nmul=10 ;\n",
    "ptc_create_layout,model=1,method=6,nst=1000,exact=false ;\n",
    "ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=false ;\n",
    "ptc_align;\n",
    "ptc_start,x={qx},px={px},y={qy},py={py},pt={dp.item()},t=0.0 ;\n",
    "ptc_track,icase=6,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",
    "res = dipole(state, length/angle, dp, length)\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": 14,
   "id": "98063db1-6987-4852-9732-3e4983fae16b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.007795176039169286, 0.0011097346020520168, 0.0007677684542948365, 0.01462661777023533]\n",
      "[0.0077951760391692755, 0.0011097346020520064, 0.0007677684542948021, 0.014626617770235346]\n",
      "[1.0408340855860843e-17, 1.0408340855860843e-17, 3.436920886779049e-17, -1.5612511283791264e-17]\n"
     ]
    }
   ],
   "source": [
    "# combined function dipole (no edge effects if e1 = e2 = 0)\n",
    "\n",
    "from pathlib import Path\n",
    "from os import system\n",
    "\n",
    "import torch\n",
    "from model.library.transformations import bend\n",
    "\n",
    "ptc = Path('ptc')\n",
    "obs = Path('track.obs0001.p0001')\n",
    "\n",
    "angle = torch.tensor(0.1, dtype=torch.float64)\n",
    "kn = torch.tensor(-1.0, dtype=torch.float64 )\n",
    "ks = torch.tensor(0.5, dtype=torch.float64 )\n",
    "dp = torch.tensor(0.005, dtype=torch.float64)\n",
    "length = torch.tensor(1.0, dtype=torch.float64)\n",
    "state = torch.tensor([0.01, -0.005, -0.01, 0.005], dtype=torch.float64)\n",
    "\n",
    "qx, px, qy, py = state.tolist()\n",
    "\n",
    "code = f\"\"\"\n",
    "mag: sbend, l={length.item()}, angle={angle.item()},k1={kn.item()},k1s={ks.item()},e1=0.0,e2=0.0,kill_ent_fringe=false,kill_exi_fringe=false;\n",
    "map:line=(mag) ;\n",
    "beam,energy=1.0E+9,particle=electron ;\n",
    "set,format=\"20.20f\",\"-20s\" ;\n",
    "use,period=map  ;\n",
    "ptc_create_universe,sector_nmul_max=10,sector_nmul=10 ;\n",
    "ptc_create_layout,model=1,method=6,nst=1000,exact=false ;\n",
    "ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=false ;\n",
    "ptc_align;\n",
    "ptc_start,x={qx},px={px},y={qy},py={py},pt={dp.item()},t=0.0 ;\n",
    "ptc_track,icase=6,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",
    "res = bend(state, length/angle, kn, ks, dp, length)\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": 15,
   "id": "d3d88c90-405b-45b9-b21b-3bb9cac4fec7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.02505664248885726, 0.02888471799554683, 0.01948347950480536, 0.007752605073815676]\n",
      "[0.025056642488857035, 0.028884717995546608, 0.019483479504805307, 0.007752605073815616]\n",
      "[2.255140518769849e-16, 2.220446049250313e-16, 5.204170427930421e-17, 5.984795992119984e-17]\n"
     ]
    }
   ],
   "source": [
    "# combined function dipole with edge effects\n",
    "\n",
    "from pathlib import Path\n",
    "from os import system\n",
    "\n",
    "import torch\n",
    "from model.library.transformations import bend\n",
    "from model.library.transformations import wedge\n",
    "\n",
    "ptc = Path('ptc')\n",
    "obs = Path('track.obs0001.p0001')\n",
    "\n",
    "angle = torch.tensor(0.1, dtype=torch.float64)\n",
    "kn = torch.tensor(-1.0, dtype=torch.float64)\n",
    "ks = torch.tensor(0.5, dtype=torch.float64)\n",
    "dp = torch.tensor(0.005, dtype=torch.float64)\n",
    "e1 = torch.tensor(0.02, dtype=torch.float64)\n",
    "e2 = torch.tensor(-0.03, dtype=torch.float64)\n",
    "length = torch.tensor(2.5, dtype=torch.float64)\n",
    "state = torch.tensor([0.01, -0.005, -0.01, 0.005], dtype=torch.float64)\n",
    "\n",
    "qx, px, qy, py = state.tolist()\n",
    "\n",
    "code = f\"\"\"\n",
    "mag: sbend, l={length.item()}, angle={angle.item()},k1={kn.item()},k1s={ks.item()},e1={e1.item()},e2={e2.item()},kill_ent_fringe=false,kill_exi_fringe=false;\n",
    "map:line=(mag) ;\n",
    "beam,energy=1.0E+9,particle=electron ;\n",
    "set,format=\"20.20f\",\"-20s\" ;\n",
    "use,period=map  ;\n",
    "ptc_create_universe,sector_nmul_max=10,sector_nmul=10 ;\n",
    "ptc_create_layout,model=1,method=6,nst=1000,exact=false ;\n",
    "ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=false ;\n",
    "ptc_align;\n",
    "ptc_start,x={qx},px={px},y={qy},py={py},pt={dp.item()},t=0.0 ;\n",
    "ptc_track,icase=6,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",
    "state = wedge(state, e1, length/angle)\n",
    "state = bend(state, length/angle, kn + 1.0E-16, ks + 1.0E-16, dp, length)\n",
    "res = wedge(state, e2, length/angle)\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": 16,
   "id": "fd03af43-7b65-45a6-bea1-8306c1338494",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[-0.08749190399027759, -0.05149961687573391, -0.05640828383591964, -0.019912883990595223]\n",
      "[-0.08749190399027745, -0.05149961687573378, -0.05640828383591953, -0.01991288399059518]\n",
      "[-1.3877787807814457e-16, -1.249000902703301e-16, -1.1102230246251565e-16, -4.163336342344337e-17]\n"
     ]
    }
   ],
   "source": [
    "# translations (exact alignment, straight layout, act on a thin representation at the entrance frame)\n",
    "\n",
    "from pathlib import Path\n",
    "from os import system\n",
    "\n",
    "import torch\n",
    "from model.library.transformations import drift\n",
    "from model.library.transformations import quadrupole\n",
    "from model.library.transformations import tx, ty, tz\n",
    "\n",
    "ptc = Path('ptc')\n",
    "obs = Path('track.obs0001.p0001')\n",
    "\n",
    "angle = torch.tensor(0.1, dtype=torch.float64)\n",
    "kn = torch.tensor(-1.0, dtype=torch.float64)\n",
    "ks = torch.tensor(0.5, dtype=torch.float64)\n",
    "dp = torch.tensor(0.005, dtype=torch.float64)\n",
    "length = torch.tensor(2.5, dtype=torch.float64)\n",
    "state = torch.tensor([0.01, -0.005, -0.015, 0.0025], dtype=torch.float64)\n",
    "\n",
    "dx = torch.tensor(0.01, dtype=torch.float64)\n",
    "dy = torch.tensor(-0.02, dtype=torch.float64)\n",
    "dz = torch.tensor(0.05, dtype=torch.float64)\n",
    "\n",
    "qx, px, qy, py = state.tolist()\n",
    "\n",
    "code = f\"\"\"\n",
    "dr: drift, l=1.0;\n",
    "mag: quadrupole, l={length.item()},k1={kn.item()},k1s={ks.item()};\n",
    "map:line=(dr, mag, dr) ;\n",
    "beam,energy=1.0E+9,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()};\n",
    "ptc_create_universe,sector_nmul_max=10,sector_nmul=10 ;\n",
    "ptc_create_layout,model=1,method=6,nst=1000,exact=false ;\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.item()},t=0.0 ;\n",
    "ptc_track,icase=6,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",
    "state = drift(state, dp, 1.0)\n",
    "state = tx(state, +dx)\n",
    "state = ty(state, +dy)\n",
    "state = tz(state, +dz, dp)\n",
    "state = quadrupole(state, kn, ks, dp, length)\n",
    "state = tz(state, -dz, dp)\n",
    "state = ty(state, -dy)\n",
    "state = tx(state, -dx)\n",
    "state = drift(state, dp, 1.0)\n",
    "res = state\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": 17,
   "id": "46617dad-22ce-413d-8c8e-cad2bbe267a0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[-0.19014967849021078, -0.2611883141418841, -0.10367399923762718, -0.09101974096085057]\n",
      "[-0.19014967849022751, -0.2611883141418846, -0.10367399923763249, -0.09101974096085078]\n",
      "[1.6736612096224235e-14, 4.996003610813204e-16, 5.3013149425851225e-15, 2.0816681711721685e-16]\n"
     ]
    }
   ],
   "source": [
    "# translations + rotations (exact alignment, straight layout, act on a thin representation at the entrance frame)\n",
    "\n",
    "from pathlib import Path\n",
    "from os import system\n",
    "\n",
    "import torch\n",
    "from model.library.transformations import quadrupole\n",
    "from model.library.transformations import tx, ty, tz\n",
    "from model.library.transformations import rx, ry, rz\n",
    "\n",
    "ptc = Path('ptc')\n",
    "obs = Path('track.obs0001.p0001')\n",
    "\n",
    "kn = torch.tensor(-1.0, dtype=torch.float64)\n",
    "ks = torch.tensor(0.5, dtype=torch.float64)\n",
    "dp = torch.tensor(0.005, dtype=torch.float64)\n",
    "\n",
    "length = torch.tensor(2.5, dtype=torch.float64)\n",
    "state = torch.tensor([0.01, -0.005, -0.015, 0.0025], 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",
    "qx, px, qy, py = state.tolist()\n",
    "\n",
    "code = f\"\"\"\n",
    "mag:quadrupole,l={length.item()},k1={kn.item()},k1s={ks.item()};\n",
    "map:line=(mag) ;\n",
    "beam,energy=1.0E+9,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=10,sector_nmul=10 ;\n",
    "ptc_create_layout,model=1,method=6,nst=1000,exact=false ;\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.item()},t=0.0 ;\n",
    "ptc_track,icase=6,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",
    "state = tx(state, +dx)\n",
    "state = ty(state, +dy)\n",
    "state = tz(state, +dz, dp)\n",
    "\n",
    "state = rx(state, +wx, dp)\n",
    "state = ry(state, +wy, dp)\n",
    "state = rz(state, +wz)\n",
    "\n",
    "state = quadrupole(state, kn + 1.0E-16, ks + 1.0E-16, dp, length)\n",
    "\n",
    "state = tz(state, -length, dp)\n",
    "\n",
    "state = rz(state, -wz)\n",
    "state = ry(state, -wy, dp)\n",
    "state = rx(state, -wx, dp)\n",
    "\n",
    "state = tz(state, -dz, dp)\n",
    "state = ty(state, -dy)\n",
    "state = tx(state, -dx)\n",
    "\n",
    "state = tz(state, +length, dp)\n",
    "\n",
    "res = state\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": 18,
   "id": "5a207d27-c76f-4739-b615-8644cc8b9a4f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[-0.20185873553633701, -0.27629369448112884, -0.08354566211146026, -0.06880534321809641]\n",
      "[-0.20185873553633754, -0.2762936944811289, -0.08354566211146025, -0.06880534321809621]\n",
      "[5.273559366969494e-16, 5.551115123125783e-17, -1.3877787807814457e-17, -1.942890293094024e-16]\n"
     ]
    }
   ],
   "source": [
    "# translations + rotations (exact alignment, curved layout, act on a thin representation at the entrance frame)\n",
    "\n",
    "from pathlib import Path\n",
    "from os import system\n",
    "\n",
    "import torch\n",
    "from model.library.transformations import bend\n",
    "from model.library.transformations import wedge\n",
    "from model.library.transformations import tx, ty, tz\n",
    "from model.library.transformations import rx, ry, rz\n",
    "\n",
    "ptc = Path('ptc')\n",
    "obs = Path('track.obs0001.p0001')\n",
    "\n",
    "angle = torch.tensor(0.1, dtype=torch.float64)\n",
    "kn = torch.tensor(-1.0, dtype=torch.float64)\n",
    "ks = torch.tensor(0.5, dtype=torch.float64)\n",
    "dp = torch.tensor(0.005, dtype=torch.float64)\n",
    "e1 = torch.tensor(0.005, dtype=torch.float64)\n",
    "e2 = torch.tensor(0.005, dtype=torch.float64)\n",
    "\n",
    "length = torch.tensor(2.5, dtype=torch.float64)\n",
    "state = torch.tensor([0.01, -0.005, -0.015, 0.0025], 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.01, dtype=torch.float64)\n",
    "\n",
    "qx, px, qy, py = state.tolist()\n",
    "\n",
    "code = f\"\"\"\n",
    "mag:sbend,l={length.item()},angle={angle.item()},k1={kn.item()},k1s={ks.item()},e1={e1.item()},e2={e2.item()},kill_ent_fringe=false,kill_exi_fringe=false;\n",
    "map:line=(mag) ;\n",
    "beam,energy=1.0E+9,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=10,sector_nmul=10 ;\n",
    "ptc_create_layout,model=1,method=6,nst=1000,exact=false ;\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.item()},t=0.0 ;\n",
    "ptc_track,icase=6,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",
    "state = tx(state, +dx)\n",
    "state = ty(state, +dy)\n",
    "state = tz(state, +dz, dp)\n",
    "\n",
    "state = rx(state, +wx, dp)\n",
    "state = ry(state, +wy, dp)\n",
    "state = rz(state, +wz)\n",
    "\n",
    "state = wedge(state, e1, length/angle)\n",
    "state = bend(state, length/angle, kn + 1.0E-16, ks + 1.0E-16, dp, length)\n",
    "state = wedge(state, e2, length/angle)\n",
    "\n",
    "state = ry(state, +angle/2, dp)\n",
    "state = tz(state, -2.0*length/angle*(angle/2.0).sin(), dp)\n",
    "state = ry(state, +angle/2, dp)\n",
    "\n",
    "state = rz(state, -wz)\n",
    "state = ry(state, -wy, dp)\n",
    "state = rx(state, -wx, dp)\n",
    "\n",
    "state = tz(state, -dz, dp)\n",
    "state = ty(state, -dy)\n",
    "state = tx(state, -dx)\n",
    "\n",
    "state = ry(state, -angle/2, dp)\n",
    "state = tz(state, +2.0*length/angle*(angle/2.0).sin(), dp)\n",
    "state = ry(state, -angle/2, dp)\n",
    "\n",
    "res = state\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": 19,
   "id": "554e17dc-f44b-4934-9d37-9960a1d13ba2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.00932966343120269, -3.829320024990772e-05, -0.008755742622854638, 0.0005]\n",
      "[0.009329663431209667, -3.829320024994158e-05, -0.008755742622854576, 0.0005]\n",
      "[-6.977057820378718e-15, 3.386098909943791e-17, -6.245004513516506e-17, 0.0]\n"
     ]
    }
   ],
   "source": [
    "# exact sector bend (without fringe)\n",
    "\n",
    "from pathlib import Path\n",
    "from os import system\n",
    "\n",
    "import torch\n",
    "from model.library.transformations import sector_bend\n",
    "\n",
    "ptc = Path('ptc')\n",
    "obs = Path('track.obs0001.p0001')\n",
    "\n",
    "length = 2.5\n",
    "angle = 0.1\n",
    "e1 = 0.0\n",
    "e2 = 0.0\n",
    "dp = 0.005\n",
    "\n",
    "state = torch.tensor([0.01, -0.0005, -0.01, 0.0005], dtype=torch.float64)\n",
    "qx, px, qy, py = state.tolist()\n",
    "\n",
    "code = f\"\"\"\n",
    "mag:sbend,l={length},angle={angle},e1={e1},e2={e2},kill_ent_fringe=true,kill_exi_fringe=true;\n",
    "map:line=(mag) ;\n",
    "beam,energy=1.0E+6,particle=electron ;\n",
    "set,format=\"20.20f\",\"-20s\" ;\n",
    "use,period=map  ;\n",
    "ptc_create_universe,sector_nmul_max=10,sector_nmul=10 ;\n",
    "ptc_create_layout,model=1,method=6,nst=1000,exact=true ;\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=6,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",
    "length = torch.tensor(length, dtype=torch.float64)\n",
    "angle = torch.tensor(angle, dtype=torch.float64)\n",
    "e1 = torch.tensor(e1, dtype=torch.float64)\n",
    "e2 = torch.tensor(e2, dtype=torch.float64)\n",
    "dp = torch.tensor(dp, dtype=torch.float64)\n",
    "\n",
    "state = sector_bend(state, length/angle, dp, length)\n",
    "res = state\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": 20,
   "id": "0aac734b-90ea-4f8c-a0bf-8e49979b4999",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.00933011773950498, -3.830113730822001e-05, -0.008756237750526722, 0.0004998143432367524]\n",
      "[0.009330117739498277, -3.8301137308233764e-05, -0.008756237750526754, 0.0004998143432367524]\n",
      "[6.702971511174383e-15, 1.3755815063409838e-17, 3.122502256758253e-17, 0.0]\n"
     ]
    }
   ],
   "source": [
    "# exact sector bend (with fringe)\n",
    "\n",
    "from pathlib import Path\n",
    "from os import system\n",
    "\n",
    "import torch\n",
    "from model.library.transformations import sector_bend\n",
    "from model.library.transformations import sector_bend_fringe\n",
    "\n",
    "ptc = Path('ptc')\n",
    "obs = Path('track.obs0001.p0001')\n",
    "\n",
    "length = 2.5\n",
    "angle = 0.1\n",
    "e1 = 0.0\n",
    "e2 = 0.0\n",
    "dp = 0.005\n",
    "\n",
    "state = torch.tensor([0.01, -0.0005, -0.01, 0.0005], dtype=torch.float64)\n",
    "qx, px, qy, py = state.tolist()\n",
    "\n",
    "code = f\"\"\"\n",
    "mag:sbend,l={length},angle={angle},e1={e1},e2={e2},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",
    "ptc_create_universe,sector_nmul_max=10,sector_nmul=10 ;\n",
    "ptc_create_layout,model=1,method=6,nst=1000,exact=true ;\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=6,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",
    "length = torch.tensor(length, dtype=torch.float64)\n",
    "angle = torch.tensor(angle, dtype=torch.float64)\n",
    "e1 = torch.tensor(e1, dtype=torch.float64)\n",
    "e2 = torch.tensor(e2, dtype=torch.float64)\n",
    "dp = torch.tensor(dp, dtype=torch.float64)\n",
    "\n",
    "state = sector_bend_fringe(state, +length/angle, dp)\n",
    "state = sector_bend(state, length/angle, dp, length)\n",
    "state = sector_bend_fringe(state, -length/angle, dp)\n",
    "res = state\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": 21,
   "id": "e50aff00-749c-440d-b1c4-ead947d7e271",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.009340060895391944, -3.805697197357787e-05, -0.00874628326389246, 0.0005003157283896672]\n",
      "[0.009340060895401722, -3.805697197360909e-05, -0.00874628326389242, 0.0005003157283896674]\n",
      "[-9.778636234081262e-15, 3.1218246304004493e-17, -3.9898639947466563e-17, -1.0842021724855044e-19]\n"
     ]
    }
   ],
   "source": [
    "# exact sector bend (with fringe and wedges)\n",
    "\n",
    "from pathlib import Path\n",
    "from os import system\n",
    "\n",
    "import torch\n",
    "from model.library.transformations import sector_bend\n",
    "from model.library.transformations import sector_bend_fringe\n",
    "from model.library.transformations import sector_bend_wedge\n",
    "from model.library.transformations import polar\n",
    "\n",
    "ptc = Path('ptc')\n",
    "obs = Path('track.obs0001.p0001')\n",
    "\n",
    "length = 2.5\n",
    "angle = 0.1\n",
    "e1 = 0.01\n",
    "e2 = -0.01\n",
    "dp = 0.005\n",
    "\n",
    "state = torch.tensor([0.01, -0.0005, -0.01, 0.0005], dtype=torch.float64)\n",
    "qx, px, qy, py = state.tolist()\n",
    "\n",
    "code = f\"\"\"\n",
    "mag:sbend,l={length},angle={angle},e1={e1},e2={e2},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",
    "ptc_create_universe,sector_nmul_max=10,sector_nmul=10 ;\n",
    "ptc_create_layout,model=1,method=6,nst=1000,exact=true ;\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=6,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",
    "length = torch.tensor(length, dtype=torch.float64)\n",
    "angle = torch.tensor(angle, dtype=torch.float64)\n",
    "e1 = torch.tensor(e1, dtype=torch.float64)\n",
    "e2 = torch.tensor(e2, dtype=torch.float64)\n",
    "dp = torch.tensor(dp, dtype=torch.float64)\n",
    "\n",
    "state = polar(state, e1, dp)\n",
    "state = sector_bend_fringe(state, +length/angle, dp)\n",
    "state = sector_bend_wedge(state, -e1, length/angle, dp)\n",
    "state = sector_bend(state, length/angle, dp, length)\n",
    "state = sector_bend_wedge(state, -e2, length/angle, dp)\n",
    "state = sector_bend_fringe(state, -length/angle, dp)\n",
    "state = polar(state, e2, dp)\n",
    "res = state\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": 22,
   "id": "718f06d9-f5af-4c43-a377-cd6e08afc8c7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[ 1.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00],\n",
      "        [ 4.0001e-04,  1.0000e+00,  0.0000e+00,  0.0000e+00],\n",
      "        [ 0.0000e+00,  0.0000e+00,  1.0000e+00,  0.0000e+00],\n",
      "        [ 0.0000e+00,  0.0000e+00, -4.0001e-04,  1.0000e+00]],\n",
      "       dtype=torch.float64)\n",
      "tensor([[ 1.0000e+00,  5.5508e-17,  0.0000e+00,  0.0000e+00],\n",
      "        [ 4.0001e-04,  1.0000e+00,  0.0000e+00,  0.0000e+00],\n",
      "        [ 0.0000e+00,  0.0000e+00,  1.0000e+00,  0.0000e+00],\n",
      "        [ 0.0000e+00,  0.0000e+00, -4.0001e-04,  1.0000e+00]],\n",
      "       dtype=torch.float64)\n",
      "\n",
      "tensor([[ 1.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00],\n",
      "        [-4.0001e-04,  1.0000e+00,  0.0000e+00,  0.0000e+00],\n",
      "        [ 0.0000e+00,  0.0000e+00,  1.0000e+00,  0.0000e+00],\n",
      "        [ 0.0000e+00,  0.0000e+00,  4.0001e-04,  1.0000e+00]],\n",
      "       dtype=torch.float64)\n",
      "tensor([[ 1.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00],\n",
      "        [-4.0001e-04,  1.0000e+00,  0.0000e+00,  0.0000e+00],\n",
      "        [ 0.0000e+00,  0.0000e+00,  1.0000e+00,  5.5234e-17],\n",
      "        [ 0.0000e+00,  0.0000e+00,  4.0001e-04,  1.0000e+00]],\n",
      "       dtype=torch.float64)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# linear wedge matrix\n",
    "\n",
    "from model.library.transformations import wedge\n",
    "from model.library.transformations import polar\n",
    "from model.library.transformations import sector_bend_fringe\n",
    "from model.library.transformations import sector_bend_wedge\n",
    "\n",
    "length = 2.5\n",
    "angle = 0.1\n",
    "e1 = 0.01\n",
    "e2 = -0.01\n",
    "dp = 0.005\n",
    "\n",
    "state = torch.tensor([0.0, 0.0, 0.0, 0.0], dtype=torch.float64)\n",
    "\n",
    "length = torch.tensor(length, dtype=torch.float64)\n",
    "angle = torch.tensor(angle, dtype=torch.float64)\n",
    "e1 = torch.tensor(e1, dtype=torch.float64)\n",
    "e2 = torch.tensor(e2, dtype=torch.float64)\n",
    "dp = torch.tensor(dp, dtype=torch.float64)\n",
    "\n",
    "def wedge_entrance(state):\n",
    "    state = polar(state, e1, dp)\n",
    "    state = sector_bend_fringe(state, +length/angle, dp)\n",
    "    state = sector_bend_wedge(state, -e1, length/angle, dp)\n",
    "    return state\n",
    "\n",
    "def wedge_exit(state):\n",
    "    state = sector_bend_wedge(state, -e2, length/angle, dp)\n",
    "    state = sector_bend_fringe(state, -length/angle, dp)\n",
    "    state = polar(state, e2, dp)\n",
    "    return state\n",
    "\n",
    "print(torch.func.jacrev(wedge)(state, e1, length/angle))\n",
    "print(torch.func.jacrev(wedge_entrance)(state))\n",
    "print()\n",
    "\n",
    "print(torch.func.jacrev(wedge)(state, e2, length/angle))\n",
    "print(torch.func.jacrev(wedge_exit)(state))\n",
    "print()"
   ]
  }
 ],
 "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
}