{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "fbb64279-2e1f-4d77-9911-7d59b1e362b9",
   "metadata": {},
   "source": [
    "# Example-06: Quadrupole (element)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "eeea8ded-1087-45ce-aadb-8e6a0006f228",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Comparison of quadrupole element with MADX-PTC and other features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "77130758-5915-437a-8c60-54b43124d525",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "from os import system\n",
    "\n",
    "import torch\n",
    "from model.library.drift import Drift\n",
    "from model.library.quadrupole import Quadrupole"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "1145aee7-dfa1-4c92-8407-51e1f4cf067e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.012268608165994052, 0.012991610983278109, 0.005825218798687177, 0.01752224400608683]\n",
      "[0.012268608165994056, 0.012991610983278081, 0.005825218798687173, 0.017522244006086804]\n",
      "[-3.469446951953614e-18, 2.7755575615628914e-17, 4.336808689942018e-18, 2.42861286636753e-17]\n"
     ]
    }
   ],
   "source": [
    "# Tracking (paraxial)\n",
    "\n",
    "ptc = Path('ptc')\n",
    "obs = Path('track.obs0001.p0001')\n",
    "\n",
    "exact = False\n",
    "align = False\n",
    "\n",
    "kn = - 2.0\n",
    "ks = + 1.5\n",
    "dp = 0.005\n",
    "length = 1.0\n",
    "state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)\n",
    "qx, px, qy, py = state.tolist()\n",
    "\n",
    "dx = align*torch.tensor(0.05, dtype=torch.float64)\n",
    "dy = align*torch.tensor(-0.02, dtype=torch.float64)\n",
    "dz = align*torch.tensor(0.05, dtype=torch.float64)\n",
    "\n",
    "wx = align*torch.tensor(0.005, dtype=torch.float64)\n",
    "wy = align*torch.tensor(-0.005, dtype=torch.float64)\n",
    "wz = align*torch.tensor(0.1, dtype=torch.float64)\n",
    "\n",
    "error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}\n",
    "\n",
    "code = f\"\"\"\n",
    "mag:quadrupole,l={length},k1={kn},k1s={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",
    "Q = Quadrupole('Q', length=length, kn=kn, ks=ks, dp=dp, exact=exact)\n",
    "res = Q(state, alignment=align, data={**Q.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": "946050b0-95a6-4c86-a151-1e5375fba9b3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.012269208914523159, 0.012992157908766264, 0.005826335256074007, 0.017521822791072554]\n",
      "[0.012269208914522952, 0.012992157908766015, 0.005826335256074073, 0.017521822791072776]\n",
      "[2.0643209364124004e-16, 2.498001805406602e-16, -6.591949208711867e-17, -2.220446049250313e-16]\n"
     ]
    }
   ],
   "source": [
    "# Tracking (exact)\n",
    "\n",
    "ptc = Path('ptc')\n",
    "obs = Path('track.obs0001.p0001')\n",
    "\n",
    "exact = True\n",
    "align = False\n",
    "\n",
    "kn = - 2.0\n",
    "ks = + 1.5\n",
    "dp = 0.005\n",
    "length = 1.0\n",
    "state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)\n",
    "qx, px, qy, py = state.tolist()\n",
    "\n",
    "dx = align*torch.tensor(0.05, dtype=torch.float64)\n",
    "dy = align*torch.tensor(-0.02, dtype=torch.float64)\n",
    "dz = align*torch.tensor(0.05, dtype=torch.float64)\n",
    "\n",
    "wx = align*torch.tensor(0.005, dtype=torch.float64)\n",
    "wy = align*torch.tensor(-0.005, dtype=torch.float64)\n",
    "wz = align*torch.tensor(0.1, dtype=torch.float64)\n",
    "\n",
    "error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}\n",
    "\n",
    "code = f\"\"\"\n",
    "mag:quadrupole,l={length},k1={kn},k1s={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",
    "Q = Quadrupole('Q', length=length, kn=kn, ks=ks, dp=dp, exact=exact, order=5, ns=5)\n",
    "res = Q(state, alignment=align, data={**Q.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": "0287bf48-152d-4c38-ac1a-474baddf5f1f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[-0.022075488016794924, -0.09165584224601611, -0.04570124622656498, -0.08629975808408008]\n",
      "[-0.02207548801679271, -0.09165584224601468, -0.04570124622656417, -0.08629975808408101]\n",
      "[-2.213507155346406e-15, -1.429412144204889e-15, -8.118505867571457e-16, 9.298117831235686e-16]\n"
     ]
    }
   ],
   "source": [
    "# Tracking (exact, alignment)\n",
    "\n",
    "ptc = Path('ptc')\n",
    "obs = Path('track.obs0001.p0001')\n",
    "\n",
    "exact = True\n",
    "align = True\n",
    "\n",
    "kn = - 2.0\n",
    "ks = + 1.5\n",
    "dp = 0.005\n",
    "length = 1.0\n",
    "state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)\n",
    "qx, px, qy, py = state.tolist()\n",
    "\n",
    "dx = align*torch.tensor(0.05, dtype=torch.float64)\n",
    "dy = align*torch.tensor(-0.02, dtype=torch.float64)\n",
    "dz = align*torch.tensor(0.05, dtype=torch.float64)\n",
    "\n",
    "wx = align*torch.tensor(0.005, dtype=torch.float64)\n",
    "wy = align*torch.tensor(-0.005, dtype=torch.float64)\n",
    "wz = align*torch.tensor(0.1, dtype=torch.float64)\n",
    "\n",
    "error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}\n",
    "\n",
    "code = f\"\"\"\n",
    "mag:quadrupole,l={length},k1={kn},k1s={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",
    "Q = Quadrupole('Q', length=length, kn=kn, ks=ks, dp=dp, exact=exact, order=5, ns=5)\n",
    "res = Q(state, alignment=align, data={**Q.data(), **error})\n",
    "\n",
    "print(ref.tolist())\n",
    "print(res.tolist())\n",
    "print((ref - res).tolist())\n",
    "\n",
    "ptc.unlink()\n",
    "obs.unlink()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "eda600a3-24d9-4ed5-85d1-822d87f98183",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([0.0242, 0.0380, 0.0152, 0.0200], dtype=torch.float64)\n",
      "\n",
      "tensor([ 0.0100, -0.0050, -0.0050,  0.0010], dtype=torch.float64)\n",
      "\n",
      "tensor([-0.0908, -0.2335, -0.0963, -0.1316], dtype=torch.float64)\n",
      "\n",
      "tensor([0., 0., 0., 0.], dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "# Deviation/error variables\n",
    "\n",
    "kn = - 2.0\n",
    "ks = + 1.5\n",
    "dp = 0.005\n",
    "length = 1.5\n",
    "state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)\n",
    "\n",
    "dx = torch.tensor(0.05, dtype=torch.float64)\n",
    "dy = torch.tensor(-0.02, dtype=torch.float64)\n",
    "dz = torch.tensor(0.05, dtype=torch.float64)\n",
    "\n",
    "wx = torch.tensor(0.005, dtype=torch.float64)\n",
    "wy = torch.tensor(-0.005, dtype=torch.float64)\n",
    "wz = torch.tensor(0.1, dtype=torch.float64)\n",
    "\n",
    "error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}\n",
    "\n",
    "Q = Quadrupole('Q', length, 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(Q(state))\n",
    "print()\n",
    "\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(Q(state, data={**Q.data(), **{'dl': -Q.length}}))\n",
    "print()\n",
    "\n",
    "# In the above Q.data() creates default deviation dictionary (with zero values for each deviaton)\n",
    "# {**Q.data(), **{'dl': -Q.length}} replaces the 'dl' key value \n",
    "\n",
    "# Additionaly, alignment errors are passed as deivation variables\n",
    "# They are used if alignment flag is raised\n",
    "\n",
    "print(Q(state, data={**Q.data(), **error}, alignment=True))\n",
    "print()\n",
    "\n",
    "# The following elements can be made equivalent using deviation variables\n",
    "\n",
    "QA = Quadrupole('QA', length, kn, ks, dp)\n",
    "QB = Quadrupole('QB', length - 0.1, kn, ks, dp)\n",
    "\n",
    "print(QA(state) - QB(state, data={**QB.data(), **{'dl': torch.tensor(+0.1, dtype=QB.dtype)}}))\n",
    "\n",
    "# Note, while in some cases float values can be passed as values to deviation variables\n",
    "# The correct behaviour in guaranteed only for tensors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "7c0e4ae8-44d8-47b1-8456-05912bf68f35",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([-5.2042e-18,  1.0408e-17, -3.4694e-18, -3.4694e-18],\n",
      "       dtype=torch.float64)\n",
      "tensor([-0.0002,  0.0037,  0.0003,  0.0031], dtype=torch.float64)\n",
      "tensor([-2.2924e-06, -3.9787e-06, -9.4215e-07,  2.1943e-07],\n",
      "       dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "# Insertion element\n",
    "\n",
    "# In this mode elements are treated as thin insertions (at the center)\n",
    "# Using parameters specified on initialization, transport two matrices are computed\n",
    "# These matrices are used to insert the element\n",
    "# Input state is transformed from the element center to its entrance\n",
    "# Next, transformation from the entrance frame to the exit frame is performed\n",
    "# This transformation can contain errors\n",
    "# The final step is to transform state from the exit frame back to the element center\n",
    "# Without errors, this results in identity transformation for linear elements\n",
    "\n",
    "kn = - 2.0\n",
    "ks = + 1.5\n",
    "dp = 0.005\n",
    "length = 1.5\n",
    "state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)\n",
    "\n",
    "dx = torch.tensor(0.05, dtype=torch.float64)\n",
    "dy = torch.tensor(-0.02, dtype=torch.float64)\n",
    "dz = torch.tensor(0.05, dtype=torch.float64)\n",
    "\n",
    "wx = torch.tensor(0.005, dtype=torch.float64)\n",
    "wy = torch.tensor(-0.005, dtype=torch.float64)\n",
    "wz = torch.tensor(0.1, dtype=torch.float64)\n",
    "\n",
    "error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}\n",
    "\n",
    "Q = Quadrupole('Q', length, kn, ks, dp, exact=False, insertion=True)\n",
    "\n",
    "# Identity transformation without errors\n",
    "\n",
    "print(Q(state) - state)\n",
    "\n",
    "# Represents effect of an error\n",
    "\n",
    "print(Q(state, data={**Q.data(), **{'dl': 0.1, 'kn': -0.1}}) - state)\n",
    "\n",
    "# Exact tracking corresponds to inclusion of kinematic term as errors\n",
    "\n",
    "Q = Quadrupole('Q', length, kn, ks, dp, exact=True, insertion=True, ns=100, order=1)\n",
    "\n",
    "print(Q(state) - state)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "07ebda5c-b22b-4a94-8f2c-0d75eac6e74d",
   "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 = - 2.0\n",
    "ks = + 1.5\n",
    "dp = 0.0\n",
    "length = 1.5\n",
    "\n",
    "dx = torch.tensor(0.05, dtype=torch.float64)\n",
    "dy = torch.tensor(-0.02, dtype=torch.float64)\n",
    "dz = torch.tensor(0.05, dtype=torch.float64)\n",
    "\n",
    "wx = torch.tensor(0.005, dtype=torch.float64)\n",
    "wy = torch.tensor(-0.005, dtype=torch.float64)\n",
    "wz = torch.tensor(0.1, dtype=torch.float64)\n",
    "\n",
    "error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}\n",
    "\n",
    "Q = Quadrupole('Q', length, kn, ks, dp, exact=True)\n",
    "\n",
    "state = 1.0E-3*torch.randn((512, 4), dtype=Q.dtype, device=Q.device)\n",
    "\n",
    "print(torch.vmap(Q)(state).shape)\n",
    "\n",
    "# To map over deviations parameters a wrapper function (or a lambda expression) can be used\n",
    "\n",
    "def wrapper(state, dp):\n",
    "    return Q(state, data={**Q.data(), **{'dp': dp}})\n",
    "\n",
    "dp = 1.0E-3*torch.randn(512, dtype=Q.dtype, device=Q.device)\n",
    "\n",
    "print(torch.vmap(wrapper)(state, dp).shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "635690c6-3b37-411a-b12b-686f2a9148d9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[ 4.7923,  3.0672,  1.8367,  0.8757],\n",
      "        [ 7.4479,  4.7923,  2.8495,  1.8367],\n",
      "        [ 1.8367,  0.8757, -0.1057,  0.7321],\n",
      "        [ 2.8495,  1.8367, -0.1507, -0.1057]], dtype=torch.float64)\n",
      "\n",
      "tensor([-0.0175, -0.0354, -0.0029, -0.0029], dtype=torch.float64)\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([-12.7901,  12.7901], dtype=torch.float64)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "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 = - 2.0\n",
    "ks = + 1.5\n",
    "dp = 0.0\n",
    "length = 1.5\n",
    "state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)\n",
    "\n",
    "dx = torch.tensor(0.05, dtype=torch.float64)\n",
    "dy = torch.tensor(-0.02, dtype=torch.float64)\n",
    "dz = torch.tensor(0.05, dtype=torch.float64)\n",
    "\n",
    "wx = torch.tensor(0.005, dtype=torch.float64)\n",
    "wy = torch.tensor(-0.005, dtype=torch.float64)\n",
    "wz = torch.tensor(0.1, dtype=torch.float64)\n",
    "\n",
    "error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}\n",
    "\n",
    "Q = Quadrupole('Q', length, kn, ks, dp, exact=False)\n",
    "\n",
    "# Compute derivative with respect to state\n",
    "\n",
    "print(torch.func.jacrev(Q)(state))\n",
    "print()\n",
    "\n",
    "# Compute derivative with respect to a deviation variable\n",
    "\n",
    "kn = torch.tensor(0.0, dtype=torch.float64)\n",
    "\n",
    "def wrapper(state, kn):\n",
    "    return Q(state, data={**Q.data(), **{'kn': kn}})\n",
    "\n",
    "print(torch.func.jacrev(wrapper, 1)(state, kn))\n",
    "print()\n",
    "\n",
    "# Compositional derivative (compute derivative of jacobian trace with respect quadrupole strength)\n",
    "\n",
    "length = 0.5\n",
    "knf = +0.2\n",
    "knd = -0.2\n",
    "\n",
    "QF = Quadrupole('QF', length, knf)\n",
    "QD = Quadrupole('QD', length, knd)\n",
    "DR = Drift('DR', 5.0)\n",
    "\n",
    "dknf = torch.tensor(0.0, dtype=torch.float64)\n",
    "dknd = torch.tensor(0.0, dtype=torch.float64)\n",
    "dkn = torch.stack([dknf, dknd])\n",
    "\n",
    "def fodo(state, dkn):\n",
    "    dknf, dknd = dkn\n",
    "    state = QF(state, data={**QF.data(), **{'kn': dknf}})\n",
    "    state = DR(state)\n",
    "    state = QD(state, data={**QD.data(), **{'kn': dknd}})\n",
    "    state = QD(state, data={**QD.data(), **{'kn': dknd}})\n",
    "    state = DR(state)\n",
    "    state = QF(state, data={**QF.data(), **{'kn': dknf}})\n",
    "    return state\n",
    "\n",
    "state = torch.tensor([0.0, 0.0, 0.0, 0.0], dtype=torch.float64)\n",
    "\n",
    "def trace(dkn):\n",
    "    return (torch.func.jacrev(fodo)(state, dkn)).trace()\n",
    "\n",
    "torch.func.jacrev(trace)(dkn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "36f2f5c2-8f46-452f-bc3d-92d64f7bbf8b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([0.0243, 0.0381, 0.0153, 0.0200], dtype=torch.float64)\n",
      "torch.Size([10, 4])\n",
      "torch.Size([10, 4, 4])\n",
      "torch.Size([100, 4])\n",
      "torch.Size([100, 4, 4])\n"
     ]
    }
   ],
   "source": [
    "# Output at each step\n",
    "\n",
    "# It is possible to collect output of state or tangent matrix at each integration step\n",
    "# Number of integratin steps is controlled by ns parameter on initialization\n",
    "# Alternatively, desired integration step length can be passed\n",
    "# Number of integration steps is computed as ceil(length/ds)\n",
    "\n",
    "kn = - 2.0\n",
    "ks = + 1.5\n",
    "dp = 0.0\n",
    "length = 1.5\n",
    "state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)\n",
    "\n",
    "dx = torch.tensor(0.05, dtype=torch.float64)\n",
    "dy = torch.tensor(-0.02, dtype=torch.float64)\n",
    "dz = torch.tensor(0.05, dtype=torch.float64)\n",
    "\n",
    "wx = torch.tensor(0.005, dtype=torch.float64)\n",
    "wy = torch.tensor(-0.005, dtype=torch.float64)\n",
    "wz = torch.tensor(0.1, dtype=torch.float64)\n",
    "\n",
    "error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}\n",
    "\n",
    "Q = Quadrupole('Q', length, kn, ks, dp, exact=False, ns=10, output=True, matrix=True)\n",
    "\n",
    "# Final state is still returned\n",
    "\n",
    "print(Q(state))\n",
    "\n",
    "# Data is added to special attributes (state and tangent matrix)\n",
    "\n",
    "print(Q.container_output.shape)\n",
    "print(Q.container_matrix.shape)\n",
    "\n",
    "# Number of integration steps can be changed\n",
    "\n",
    "Q.ns = 100\n",
    "\n",
    "Q(state)\n",
    "print(Q.container_output.shape)\n",
    "print(Q.container_matrix.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "81e3b3ec-79db-4d28-884d-88b4ef05c99c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.024284271092022615, 0.03811354319280181, 0.015254854998694, 0.01995832080452666]\n",
      "[0.024286764143785347, 0.038112599406681526, 0.015255420211007705, 0.01995809327180016]\n",
      "[-2.4930517627322346e-06, 9.437861202832298e-07, -5.652123137040582e-07, 2.2753272650027911e-07]\n",
      "\n",
      "[0, 1, 0, 1, 0, 1, 0]\n",
      "[0.6756035959798289, 1.3512071919596578, -0.17560359597982877, -1.7024143839193153, -0.17560359597982877, 1.3512071919596578, 0.6756035959798289]\n"
     ]
    }
   ],
   "source": [
    "# Integration order is set on initialization (default value is zero)\n",
    "# This order is related to difference order as 2n + 2\n",
    "# Thus, zero corresponds to second order difference method\n",
    "\n",
    "kn = - 2.0\n",
    "ks = + 1.5\n",
    "dp = 0.0\n",
    "length = 1.5\n",
    "state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)\n",
    "\n",
    "dx = torch.tensor(0.05, dtype=torch.float64)\n",
    "dy = torch.tensor(-0.02, dtype=torch.float64)\n",
    "dz = torch.tensor(0.05, dtype=torch.float64)\n",
    "\n",
    "wx = torch.tensor(0.005, dtype=torch.float64)\n",
    "wy = torch.tensor(-0.005, dtype=torch.float64)\n",
    "wz = torch.tensor(0.1, dtype=torch.float64)\n",
    "\n",
    "error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}\n",
    "\n",
    "Q = Quadrupole('Q', length, kn, ks, dp, order=1, exact=True)\n",
    "\n",
    "# For quadrupole integration is performed only with exact flag\n",
    "# In this case, kinematic term error is added\n",
    "\n",
    "Q.ns = 1\n",
    "ref = Q(state)\n",
    "\n",
    "Q.ns = 10\n",
    "res = Q(state)\n",
    "\n",
    "print(ref.tolist())\n",
    "print(res.tolist())\n",
    "print((ref - res).tolist())\n",
    "print()\n",
    "\n",
    "# Integrator parameters are stored in data attribute (if integration is actually performed)\n",
    "\n",
    "maps, weights = Q._data\n",
    "print(maps)\n",
    "print(weights)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "94b99173-1c93-4f71-b7b3-e89773376fc1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([0.2107, 0.1703], dtype=torch.float64)\n",
      "tensor([0.2126, 0.1681], dtype=torch.float64)\n",
      "tensor([0.2126, 0.1681], dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "# Derivatives of twiss parameters\n",
    "\n",
    "# pip install git+https://github.com/i-a-morozov/twiss.git@main\n",
    "\n",
    "from twiss import twiss\n",
    "\n",
    "length = 0.5\n",
    "knf = +0.21\n",
    "knd = -0.19\n",
    "\n",
    "QF = Quadrupole('QF', length, knf)\n",
    "QD = Quadrupole('QD', length, knd)\n",
    "DR = Drift('DR', 5.0)\n",
    "\n",
    "dknf = torch.tensor(0.0, dtype=torch.float64)\n",
    "dknd = torch.tensor(0.0, dtype=torch.float64)\n",
    "dkn = torch.stack([dknf, dknd])\n",
    "\n",
    "def fodo(state, dkn):\n",
    "    dknf, dknd = dkn\n",
    "    state = QF(state, data={**QF.data(), **{'kn': dknf}})\n",
    "    state = DR(state)\n",
    "    state = QD(state, data={**QD.data(), **{'kn': dknd}})\n",
    "    state = QD(state, data={**QD.data(), **{'kn': dknd}})\n",
    "    state = DR(state)\n",
    "    state = QF(state, data={**QF.data(), **{'kn': dknf}})\n",
    "    return state\n",
    "\n",
    "state = torch.tensor([0.0, 0.0, 0.0, 0.0], dtype=torch.float64)\n",
    "\n",
    "def tune(dkn):\n",
    "    matrix = torch.func.jacrev(fodo)(state, dkn)\n",
    "    tune, *_ = twiss(matrix)\n",
    "    return tune\n",
    "\n",
    "# Compute tunes and jacobian\n",
    "\n",
    "values = tune(dkn)\n",
    "jacobian = torch.func.jacrev(tune)(dkn)\n",
    "\n",
    "# Test jacobiant\n",
    "\n",
    "print(values)\n",
    "print(tune(dkn + 1.0E-3))\n",
    "print(values + jacobian @ (dkn + 1.0E-3))"
   ]
  }
 ],
 "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
}