{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "ea354b5c-a05a-4e63-80e8-72432ddbc71a",
   "metadata": {},
   "source": [
    "# Example-15: BPM (element)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "944a061b-c05c-4a7c-b8ce-8ec66dbbc69e",
   "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.bpm import BPM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "738300d3-b477-466e-aacf-4f7d86292b8d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([ 0.0100, -0.0050, -0.0050,  0.0010], dtype=torch.float64)"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# BPM acts as identity transformation with calibration error\n",
    "\n",
    "state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)\n",
    "\n",
    "B = BPM('B', direction='forward')\n",
    "B(state)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d7276a68-55f7-4114-b468-6ee14328970a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'xx': tensor(0., dtype=torch.float64),\n",
       " 'xy': tensor(0., dtype=torch.float64),\n",
       " 'yx': tensor(0., dtype=torch.float64),\n",
       " 'yy': tensor(0., dtype=torch.float64),\n",
       " 'dp': tensor(0., dtype=torch.float64),\n",
       " 'dx': tensor(0., dtype=torch.float64),\n",
       " 'dy': tensor(0., dtype=torch.float64),\n",
       " 'dz': tensor(0., dtype=torch.float64),\n",
       " 'wx': tensor(0., dtype=torch.float64),\n",
       " 'wy': tensor(0., dtype=torch.float64),\n",
       " 'wz': tensor(0., dtype=torch.float64)}"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Calibration errors can be passed as deviation variables\n",
    "\n",
    "# qx -> (1 + xx) qx + xy qy\n",
    "# qy -> yx qx + (1 + yy) qy\n",
    "\n",
    "B.data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "d7ab3b4a-2423-4b6e-a359-5f978357c3d3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([ 0.0100, -0.0050, -0.0050,  0.0010], dtype=torch.float64)\n",
      "tensor([ 0.0105, -0.0048, -0.0042,  0.0011], dtype=torch.float64)\n",
      "tensor([ 0.0100, -0.0050, -0.0050,  0.0010], dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "# Transform to BPM frame and back to beam frame\n",
    "\n",
    "xx = torch.tensor(+0.05, dtype=torch.float64)\n",
    "xy = torch.tensor(+0.01, dtype=torch.float64)\n",
    "yx = torch.tensor(+0.05, dtype=torch.float64)\n",
    "yy = torch.tensor(-0.06, dtype=torch.float64)\n",
    "\n",
    "B = BPM('B', direction='forward')\n",
    "\n",
    "state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)\n",
    "print(state)\n",
    "\n",
    "state = B(state, data={**B.data(), **{'xx': xx, 'xy': xy, 'yx': yx, 'yy': yy}})\n",
    "print(state)\n",
    "\n",
    "B.direction = 'inverse'\n",
    "\n",
    "state = B(state, data={**B.data(), **{'xx': xx, 'xy': xy, 'yx': yx, 'yy': yy}})\n",
    "print(state)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "8c01a22d-a494-4c36-bc98-50e6f5e873ca",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([ 0.0100, -0.0050, -0.0050,  0.0010], dtype=torch.float64)\n",
      "tensor([ 0.0105, -0.0048, -0.0042,  0.0011], dtype=torch.float64)\n",
      "tensor([ 0.0100, -0.0050, -0.0050,  0.0010], dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "# Transform to BPM frame and back to beam frame using a pair of BPMS\n",
    "\n",
    "xx = torch.tensor(+0.05, dtype=torch.float64)\n",
    "xy = torch.tensor(+0.01, dtype=torch.float64)\n",
    "yx = torch.tensor(+0.05, dtype=torch.float64)\n",
    "yy = torch.tensor(-0.06, dtype=torch.float64)\n",
    "\n",
    "BA = BPM('B', direction='forward')\n",
    "BB = BPM('B', direction='inverse')\n",
    "\n",
    "state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)\n",
    "print(state)\n",
    "\n",
    "state = BA(state, data={**BA.data(), **{'xx': xx, 'xy': xy, 'yx': yx, 'yy': yy}})\n",
    "print(state)\n",
    "\n",
    "state = BB(state, data={**BB.data(), **{'xx': xx, 'xy': xy, 'yx': yx, 'yy': yy}})\n",
    "print(state)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "d9bc6410-fd9d-424c-bd66-e4f7360255a2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[ 5.0000e-03, -4.0000e-03,  0.0000e+00,  0.0000e+00],\n",
      "        [ 4.5880e-03, -2.4404e-04, -1.0625e-03,  5.6516e-05],\n",
      "        [ 0.0000e+00,  0.0000e+00,  5.0000e-03, -4.0000e-03],\n",
      "        [-4.8809e-05,  5.1249e-03,  1.1303e-05, -1.1868e-03]],\n",
      "       dtype=torch.float64)\n",
      "tensor([[ 0.0000e+00,  8.6736e-19,  0.0000e+00, -5.4888e-19],\n",
      "        [ 0.0000e+00,  0.0000e+00,  3.3881e-20,  0.0000e+00],\n",
      "        [-2.7105e-19,  0.0000e+00,  8.6736e-19,  0.0000e+00],\n",
      "        [ 2.1684e-19,  3.9302e-19, -3.3881e-21,  2.1684e-19]],\n",
      "       dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "# Differentiability\n",
    "\n",
    "B = BPM('B', direction='forward')\n",
    "D = Drift('D', 1.0)\n",
    "\n",
    "xx = torch.tensor(+0.05, dtype=torch.float64)\n",
    "xy = torch.tensor(+0.01, dtype=torch.float64)\n",
    "yx = torch.tensor(+0.05, dtype=torch.float64)\n",
    "yy = torch.tensor(-0.06, dtype=torch.float64)\n",
    "\n",
    "error = torch.stack([xx, xy, yx, yy])\n",
    "\n",
    "state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)\n",
    "\n",
    "def line(state, error):\n",
    "    xx, xy, yx, yy = error\n",
    "    state = D(state)\n",
    "    state = B(state, data={**B.data(), **{'xx': xx, 'xy': xy, 'yx': yx, 'yy': yy}})\n",
    "    return state\n",
    "    \n",
    "print(torch.func.jacrev(line, 1)(state, error))\n",
    "\n",
    "def line(state, error):\n",
    "    xx, xy, yx, yy = error\n",
    "    state = D(state)\n",
    "    B.direction = 'forward'\n",
    "    state = B(state, data={**B.data(), **{'xx': xx, 'xy': xy, 'yx': yx, 'yy': yy}})\n",
    "    B.direction = 'inverse'\n",
    "    state = B(state, data={**B.data(), **{'xx': xx, 'xy': xy, 'yx': yx, 'yy': yy}})\n",
    "    return state\n",
    "    \n",
    "print(torch.func.jacrev(line, 1)(state, error))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "3bfceff3-dd58-4e28-9cd1-ea5a39592d52",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([ 0.0105, -0.0048, -0.0042,  0.0011], dtype=torch.float64)\n",
      "tensor([ 0.0086, -0.0048, -0.0082,  0.0008], dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "# Alignment support\n",
    "\n",
    "xx = torch.tensor(+0.05, dtype=torch.float64)\n",
    "xy = torch.tensor(+0.01, dtype=torch.float64)\n",
    "yx = torch.tensor(+0.05, dtype=torch.float64)\n",
    "yy = torch.tensor(-0.06, 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",
    "B = BPM('B', direction='forward')\n",
    "\n",
    "state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)\n",
    "state = B(state, data={**B.data(), **{'xx': xx, 'xy': xy, 'yx': yx, 'yy': yy}, **error})\n",
    "print(state)\n",
    "\n",
    "state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)\n",
    "state = B(state, data={**B.data(), **{'xx': xx, 'xy': xy, 'yx': yx, 'yy': yy}, **error}, alignment=True)\n",
    "print(state)"
   ]
  }
 ],
 "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
}