{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "fd58df7d-32e6-4019-8bb5-44ae51af421b",
   "metadata": {},
   "source": [
    "# Example-12: Gradient (element)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e667c560-c32b-4f70-a52a-b80742e83186",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Comparison of gradient element with MADX-PTC and other features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "9005919a-5332-4889-91a0-9b6993c5ebf4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "from os import system\n",
    "\n",
    "import torch\n",
    "from model.library.gradient import Gradient"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "b945f196-7424-4e36-8582-93dd194cc5f6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.001, 7.5e-06, -0.005, -2.75e-05]\n",
      "[0.001, 7.5e-06, -0.005, -2.75e-05]\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",
    "kn = +5.0E-3\n",
    "ks = -2.5E-3\n",
    "\n",
    "dp = 0.005\n",
    "\n",
    "state = torch.tensor([0.001, 0.0, -0.005, 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",
    "mag: multipole,knl={{0.0,{kn}}},ksl={{0.0,{ks}}};\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=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",
    "G = Gradient('G', kn=kn, ks=ks, dp=dp)\n",
    "res = G(state, alignment=align, data={**G.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": "e0fc2448-8b0f-45bc-90a1-216290077395",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.000991887042015914, 0.00016412084104346932, -0.005011606734840507, 0.00023479809553215934]\n",
      "[0.0009918870420159223, 0.00016412084104346948, -0.005011606734840506, 0.00023479809553215904]\n",
      "[-8.239936510889834e-18, -1.6263032587282567e-19, -8.673617379884035e-19, 2.981555974335137e-19]\n"
     ]
    }
   ],
   "source": [
    "# Tracking (alignment)\n",
    "\n",
    "ptc = Path('ptc')\n",
    "obs = Path('track.obs0001.p0001')\n",
    "\n",
    "exact = True\n",
    "align = True\n",
    "\n",
    "kn = +5.0E-3\n",
    "ks = -2.5E-3\n",
    "\n",
    "dp = 0.005\n",
    "\n",
    "state = torch.tensor([0.001, 0.0, -0.005, 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",
    "mag: multipole,knl={{0.0,{kn}}},ksl={{0.0,{ks}}};\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=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",
    "G = Gradient('G', kn=kn, ks=ks, dp=dp)\n",
    "res = G(state, alignment=align, data={**G.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": "a13f4b72-d13a-4785-995e-cabf56a29d7b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([ 1.0000e-03,  7.5000e-06, -5.0000e-03, -2.7500e-05],\n",
      "       dtype=torch.float64)\n",
      "\n",
      "tensor([ 0.0010,  0.0000, -0.0050,  0.0000], dtype=torch.float64)\n",
      "\n",
      "tensor([ 1.0000e-03,  3.2500e-05, -5.0000e-03,  4.7500e-05],\n",
      "       dtype=torch.float64)\n",
      "\n",
      "tensor([0., 0., 0., 0.], dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "# Deviation/error variables\n",
    "\n",
    "kn = +5.0E-3\n",
    "ks = -2.5E-3\n",
    "dp = 0.005\n",
    "state = torch.tensor([0.001, 0.0, -0.005, 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",
    "G = Gradient('G', kn, ks, 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(G(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(G(state, data={**G.data(), **{'kn': -kn, 'ks': -ks}}))\n",
    "print()\n",
    "\n",
    "# In the above G.data() creates default deviation dictionary (with zero values for each deviaton)\n",
    "# {**G.data(), **{'kn': -kn, 'ks': -ks}} replaces the 'kn' and 'ks' key values \n",
    "\n",
    "# Additionaly, alignment errors are passed as deivation variables\n",
    "# They are used if alignment flag is raised\n",
    "\n",
    "print(G(state, data={**G.data(), **error}, alignment=True))\n",
    "print()\n",
    "\n",
    "\n",
    "# The following elements can be made equivalent using deviation variables\n",
    "\n",
    "GA = Gradient('GA', kn, ks, dp)\n",
    "GB = Gradient('GB', kn - 0.001, ks, dp)\n",
    "\n",
    "print(GA(state) - GB(state, data={**GB.data(), **{'kn': + 0.001}}))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "a8e1f4b6-0f1f-45ae-a80b-b15a95bde944",
   "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",
    "kn = +5.0E-3\n",
    "ks = -2.5E-3\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",
    "G = Gradient('G', kn, ks, dp)\n",
    "\n",
    "state = 1.0E-3*torch.randn((512, 4), dtype=G.dtype, device=G.device)\n",
    "\n",
    "print(torch.vmap(G)(state).shape)\n",
    "\n",
    "# To map over deviations parameters a wrapper function (or a lambda expression) can be used\n",
    "\n",
    "def wrapper(state, kn, ks):\n",
    "    return G(state, data={**G.data(), **{'kn': kn, 'ks': ks}})\n",
    "\n",
    "kn = 1.0E-3*torch.randn(512, dtype=G.dtype, device=G.device)\n",
    "ks = 1.0E-3*torch.randn(512, dtype=G.dtype, device=G.device)\n",
    "\n",
    "print(torch.vmap(wrapper)(state, kn, ks).shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "f93a9e04-1724-4eaf-95d6-3dd2889a4ac6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[ 1.0000,  0.0000,  0.0000,  0.0000],\n",
      "        [-0.0050,  1.0000, -0.0025,  0.0000],\n",
      "        [ 0.0000,  0.0000,  1.0000,  0.0000],\n",
      "        [-0.0025,  0.0000,  0.0050,  1.0000]], dtype=torch.float64)\n",
      "\n",
      "tensor([[-0.0000,  0.0000],\n",
      "        [-0.0010, -0.0050],\n",
      "        [-0.0000,  0.0000],\n",
      "        [-0.0050,  0.0010]], 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",
    "kn = +5.0E-3\n",
    "ks = -2.5E-3\n",
    "dp = 0.005\n",
    "state = torch.tensor([0.001, 0.0, -0.005, 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",
    "error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}\n",
    "\n",
    "G = Gradient('G', kn, ks, dp)\n",
    "\n",
    "# Compute derivative with respect to state\n",
    "\n",
    "print(torch.func.jacrev(G)(state))\n",
    "print()\n",
    "\n",
    "# Compute derivative with respect to a deviation variable\n",
    "\n",
    "dkn = torch.tensor(0.0, dtype=torch.float64)\n",
    "dks = torch.tensor(0.0, dtype=torch.float64)\n",
    "dk = torch.stack([dkn, dks])\n",
    "\n",
    "def wrapper(state, dk):\n",
    "    dkn, dks = dk\n",
    "    return G(state, data={**G.data(), **{'kn': dkn, 'ks': dks}})\n",
    "\n",
    "print(torch.func.jacrev(wrapper, 1)(state, dk))\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
}