{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "240c12c4-8fe9-4d7d-a941-e2fe79d1190e",
   "metadata": {},
   "source": [
    "# Example-13: Kick (element)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "093201fe-a8ec-41e3-8b44-f557a093b695",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Comparison of kick element with MADX-PTC and other features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "1d032d6b-f215-4fc8-885d-cf41ca47eaed",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "from os import system\n",
    "\n",
    "import torch\n",
    "from model.library.kick import Kick"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ee0fe349-80f3-4ca1-89cd-4220cc2d113d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.001, 0.00011938333333333334, -0.005, -5.091666666666667e-05]\n",
      "[0.001, 0.00011938333333333334, -0.005, -5.0916666666666666e-05]\n",
      "[0.0, 0.0, 0.0, -6.776263578034403e-21]\n"
     ]
    }
   ],
   "source": [
    "# Tracking\n",
    "\n",
    "ptc = Path('ptc')\n",
    "obs = Path('track.obs0001.p0001')\n",
    "\n",
    "exact = True\n",
    "align = False\n",
    "\n",
    "ms = +10.0\n",
    "mo = -50.0\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,0.0,{ms},{mo}}};\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",
    "K = Kick('K', ms=ms, mo=mo, dp=dp)\n",
    "res = K(state, alignment=align, data={**K.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": "fa33a45c-27ca-4881-ab50-9a45b0c422b7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.0014217329730612119, -0.00853058669498686, -0.004440545938333818, -0.011316334457474285]\n",
      "[0.0014217329730612169, -0.008530586694986856, -0.004440545938333821, -0.011316334457474281]\n",
      "[-4.9873299934333204e-18, -3.469446951953614e-18, 2.6020852139652106e-18, -3.469446951953614e-18]\n"
     ]
    }
   ],
   "source": [
    "# Tracking (alignment)\n",
    "\n",
    "ptc = Path('ptc')\n",
    "obs = Path('track.obs0001.p0001')\n",
    "\n",
    "exact = True\n",
    "align = True\n",
    "\n",
    "ms = +10.0\n",
    "mo = -50.0\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,0.0,{ms},{mo}}};\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",
    "K = Kick('K', ms=ms, mo=mo, dp=dp)\n",
    "res = K(state, alignment=align, data={**K.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": "45f40daa-9d13-4ce9-8607-f49e2d13ad00",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([ 1.0000e-03,  1.1938e-04, -5.0000e-03, -5.0917e-05],\n",
      "       dtype=torch.float64)\n",
      "\n",
      "tensor([ 0.0010,  0.0000, -0.0050,  0.0000], dtype=torch.float64)\n",
      "\n",
      "tensor([ 0.0010,  0.0003, -0.0050,  0.0004], dtype=torch.float64)\n",
      "\n",
      "tensor([0., 0., 0., 0.], dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "# Deviation/error variables\n",
    "\n",
    "ms = +10.0\n",
    "mo = -50.0\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",
    "K = Kick('K', 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(K(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(K(state, data={**K.data(), **{'ms': -ms, 'mo': -mo}}))\n",
    "print()\n",
    "\n",
    "# In the above K.data() creates default deviation dictionary (with zero values for each deviaton)\n",
    "# {**K.data(), **{'ms': -ms, 'mo': -mo}} replaces the 'ms' and 'mo' key values \n",
    "\n",
    "# Additionaly, alignment errors are passed as deivation variables\n",
    "# They are used if alignment flag is raised\n",
    "\n",
    "print(K(state, data={**K.data(), **error}, alignment=True))\n",
    "print()\n",
    "\n",
    "\n",
    "# The following elements can be made equivalent using deviation variables\n",
    "\n",
    "KA = Kick('KA', ms, mo, dp)\n",
    "KB = Kick('KB', ms - 0.1, mo, dp)\n",
    "\n",
    "print(KA(state) - KB(state, data={**KB.data(), **{'ms': + 0.1}}))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "cd9ed885-d813-48c5-ae6e-eb52e9116657",
   "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",
    "ms = +10.0\n",
    "mo = -50.0\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",
    "K = Kick('K', ms, mo, dp)\n",
    "\n",
    "state = 1.0E-3*torch.randn((512, 4), dtype=K.dtype, device=K.device)\n",
    "\n",
    "print(torch.vmap(K)(state).shape)\n",
    "\n",
    "# To map over deviations parameters a wrapper function (or a lambda expression) can be used\n",
    "\n",
    "def wrapper(state, ms, mo):\n",
    "    return K(state, data={**K.data(), **{'ms': ms, 'mo': mo}})\n",
    "\n",
    "kn = 1.0E-3*torch.randn(512, dtype=K.dtype, device=K.device)\n",
    "ks = 1.0E-3*torch.randn(512, dtype=K.dtype, device=K.device)\n",
    "\n",
    "print(torch.vmap(wrapper)(state, kn, ks).shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "4f0505ce-aa7e-40fd-9ab1-a34161b0cabf",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[ 1.0000,  0.0000,  0.0000,  0.0000],\n",
      "        [-0.0106,  1.0000, -0.0498,  0.0000],\n",
      "        [ 0.0000,  0.0000,  1.0000,  0.0000],\n",
      "        [-0.0498,  0.0000,  0.0106,  1.0000]], dtype=torch.float64)\n",
      "\n",
      "tensor([[ 0.0000e+00,  0.0000e+00],\n",
      "        [ 1.2000e-05,  1.2333e-08],\n",
      "        [ 0.0000e+00,  0.0000e+00],\n",
      "        [-5.0000e-06,  1.8333e-08]], 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",
    "ms = +10.0\n",
    "mo = -50.0\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",
    "K = Kick('K', ms, mo, dp)\n",
    "\n",
    "# Compute derivative with respect to state\n",
    "\n",
    "print(torch.func.jacrev(K)(state))\n",
    "print()\n",
    "\n",
    "# Compute derivative with respect to a deviation variable\n",
    "\n",
    "dms = torch.tensor(0.0, dtype=torch.float64)\n",
    "dmo = torch.tensor(0.0, dtype=torch.float64)\n",
    "dm = torch.stack([dms, dmo])\n",
    "\n",
    "def wrapper(state, dm):\n",
    "    dms, dmo = dm\n",
    "    return K(state, data={**K.data(), **{'ms': dms, 'mo': dmo}})\n",
    "\n",
    "print(torch.func.jacrev(wrapper, 1)(state, dm))\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
}