{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "02f34857-7739-4c9f-a7b2-194afe494250",
   "metadata": {},
   "source": [
    "# Example-11: Corrector (element)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "841f07c5-c013-4694-84aa-72789744bfa3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Comparison of corrector element with MADX-PTC and other features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "6d66c48e-0458-4ced-a00d-4b0b6472e071",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "from os import system\n",
    "\n",
    "import torch\n",
    "from model.library.corrector import Corrector"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "2cd6ac8a-1b01-4796-a566-ce1f86391ca8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.0, 0.001, 0.0, -0.005]\n",
      "[0.0, 0.001, 0.0, -0.005]\n",
      "[0.0, 0.0, 0.0, 0.0]\n"
     ]
    }
   ],
   "source": [
    "# Tracking\n",
    "\n",
    "ptc = Path('ptc')\n",
    "obs = Path('track.obs0001.p0001')\n",
    "\n",
    "exact = True\n",
    "align = False\n",
    "\n",
    "cx = +0.001\n",
    "cy = -0.005\n",
    "\n",
    "dp = 0.005\n",
    "\n",
    "state = torch.tensor([0.0, 0.0, 0.0, 0.0], 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",
    "cx:hkicker,l=0.0,kick={cx};\n",
    "cy:vkicker,l=0.0,kick={cy};\n",
    "map:line=(cx, cy) ;\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=10,sector_nmul=10 ;\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",
    "C = Corrector('C', cx=cx, cy=cy, dp=dp)\n",
    "res = C(state, alignment=align, data={**C.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": "bff47922-c9ed-46bf-94e4-f03c8a89cb83",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.0, 0.0014941712485121669, 0.0, -0.004875187409743301]\n",
      "[-6.938893903907228e-18, 0.0014941712485121667, 0.0, -0.004875187409743301]\n",
      "[6.938893903907228e-18, 2.168404344971009e-19, 0.0, 0.0]\n"
     ]
    }
   ],
   "source": [
    "# Tracking (alignment)\n",
    "\n",
    "# Only dx and dy alignment errors seems to work as expected in MADX\n",
    "# Also, wz rotation matches tilt\n",
    "\n",
    "ptc = Path('ptc')\n",
    "obs = Path('track.obs0001.p0001')\n",
    "\n",
    "exact = True\n",
    "align = True\n",
    "\n",
    "cx = +0.001\n",
    "cy = -0.005\n",
    "dp = 0.005\n",
    "state = torch.tensor([0.0, 0.0, 0.0, 0.0], 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.0, dtype=torch.float64)\n",
    "\n",
    "wx = align*torch.tensor(0.0, dtype=torch.float64)\n",
    "wy = align*torch.tensor(0.0, 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",
    "cx:hkicker,l=0.0,kick={cx},tilt={wz.item()};\n",
    "cy:vkicker,l=0.0,kick={cy},tilt={wz.item()};\n",
    "map:line=(cx, cy) ;\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()} ;\n",
    "ptc_create_universe,sector_nmul_max=10,sector_nmul=10 ;\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",
    "C = Corrector('C', cx=cx, cy=cy, dp=dp)\n",
    "res = C(state, alignment=align, data={**C.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": "7a88a494-9094-429c-b2b3-7633f76add7a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([ 0.0000,  0.0010,  0.0000, -0.0050], dtype=torch.float64)\n",
      "\n",
      "tensor([0., 0., 0., 0.], dtype=torch.float64)\n",
      "\n",
      "tensor([ 0.0000, -0.0010,  0.0000,  0.0050], dtype=torch.float64)\n",
      "\n",
      "tensor([0., 0., 0., 0.], dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "# Deviation/error variables\n",
    "\n",
    "cx = +0.001\n",
    "cy = -0.005\n",
    "dp = 0.005\n",
    "state = torch.tensor([0.0, 0.0, 0.0, 0.0], dtype=torch.float64)\n",
    "\n",
    "dx = torch.tensor(+0.01, dtype=torch.float64)\n",
    "dy = torch.tensor(-0.01, dtype=torch.float64)\n",
    "dz = torch.tensor(0.0, dtype=torch.float64)\n",
    "\n",
    "wx = torch.tensor(0.0, dtype=torch.float64)\n",
    "wy = torch.tensor(0.0, dtype=torch.float64)\n",
    "wz = torch.tensor(torch.pi, dtype=torch.float64)\n",
    "\n",
    "error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}\n",
    "\n",
    "C = Corrector('C', cx, cy, 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(C(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(C(state, data={**C.data(), **{'cx': -cx, 'cy': -cy}}))\n",
    "print()\n",
    "\n",
    "# In the above C.data() creates default deviation dictionary (with zero values for each deviaton)\n",
    "# {**C.data(), **{'cx': -cx, 'cy': -cy}} replaces the 'cx' and 'cy' key values \n",
    "\n",
    "# Additionaly, alignment errors are passed as deivation variables\n",
    "# They are used if alignment flag is raised\n",
    "\n",
    "print(C(state, data={**C.data(), **error}, alignment=True))\n",
    "print()\n",
    "\n",
    "\n",
    "# The following elements can be made equivalent using deviation variables\n",
    "\n",
    "CA = Corrector('CA', cx, cy, dp)\n",
    "CB = Corrector('CB', cx - 0.001, cy, dp)\n",
    "\n",
    "print(CA(state) - CB(state, data={**CB.data(), **{'cx': + 0.001}}))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "8f6b0cbf-af1c-4694-9ca6-9c67573711ed",
   "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",
    "cx = +0.001\n",
    "cy = -0.005\n",
    "dp = 0.005\n",
    "\n",
    "dx = torch.tensor(+0.01, dtype=torch.float64)\n",
    "dy = torch.tensor(-0.01, dtype=torch.float64)\n",
    "dz = torch.tensor(0.0, dtype=torch.float64)\n",
    "\n",
    "wx = torch.tensor(0.0, dtype=torch.float64)\n",
    "wy = torch.tensor(0.0, dtype=torch.float64)\n",
    "wz = torch.tensor(torch.pi, dtype=torch.float64)\n",
    "\n",
    "error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}\n",
    "\n",
    "C = Corrector('C', cx, cy, dp)\n",
    "\n",
    "state = 1.0E-3*torch.randn((512, 4), dtype=C.dtype, device=C.device)\n",
    "\n",
    "print(torch.vmap(C)(state).shape)\n",
    "\n",
    "# To map over deviations parameters a wrapper function (or a lambda expression) can be used\n",
    "\n",
    "def wrapper(state, cx, cy):\n",
    "    return C(state, data={**C.data(), **{'cx': cx, 'cy': cy}})\n",
    "\n",
    "cx = 1.0E-3*torch.randn(512, dtype=C.dtype, device=C.device)\n",
    "cy = 1.0E-3*torch.randn(512, dtype=C.dtype, device=C.device)\n",
    "\n",
    "print(torch.vmap(wrapper)(state, cx, cy).shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ec33fb7c-c532-4a92-8b84-ba6c8d0b1060",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[1., 0., 0., 0.],\n",
      "        [0., 1., 0., 0.],\n",
      "        [0., 0., 1., 0.],\n",
      "        [0., 0., 0., 1.]], dtype=torch.float64)\n",
      "\n",
      "tensor([[0., 0.],\n",
      "        [1., 0.],\n",
      "        [0., 0.],\n",
      "        [0., 1.]], 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",
    "cx = +0.001\n",
    "cy = -0.005\n",
    "dp = 0.005\n",
    "state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)\n",
    "\n",
    "dx = torch.tensor(+0.01, dtype=torch.float64)\n",
    "dy = torch.tensor(-0.01, dtype=torch.float64)\n",
    "dz = torch.tensor(0.0, dtype=torch.float64)\n",
    "\n",
    "wx = torch.tensor(0.0, dtype=torch.float64)\n",
    "wy = torch.tensor(0.0, dtype=torch.float64)\n",
    "wz = torch.tensor(torch.pi, dtype=torch.float64)\n",
    "error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}\n",
    "\n",
    "C = Corrector('C', cx, cy, dp)\n",
    "\n",
    "# Compute derivative with respect to state\n",
    "\n",
    "print(torch.func.jacrev(C)(state))\n",
    "print()\n",
    "\n",
    "# Compute derivative with respect to a deviation variable\n",
    "\n",
    "dcx = torch.tensor(0.0, dtype=torch.float64)\n",
    "dcy = torch.tensor(0.0, dtype=torch.float64)\n",
    "dc = torch.stack([dcx, dcy])\n",
    "\n",
    "def wrapper(state, dc):\n",
    "    dcx, dcy = dc\n",
    "    return C(state, data={**C.data(), **{'cx': dcx, 'cy': dcy}})\n",
    "\n",
    "print(torch.func.jacrev(wrapper, 1)(state, dc))\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
}