{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6388403b-bb6e-4adc-b2d9-1d61ae486e1d",
   "metadata": {},
   "source": [
    "## ELETTRA-04: Tune response matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "12dbfada-17ba-472d-94d4-cd02749f992e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# In this example tune responce matrix is constructed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "1eb92ebb-cda1-4836-b0a3-149c10d2c538",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import\n",
    "\n",
    "import torch\n",
    "from torch import Tensor\n",
    "\n",
    "from pathlib import Path\n",
    "\n",
    "import matplotlib\n",
    "from matplotlib import pyplot as plt\n",
    "from matplotlib.patches import Rectangle\n",
    "matplotlib.rcParams['text.usetex'] = True\n",
    "\n",
    "from model.library.element import Element\n",
    "from model.library.line import Line\n",
    "from model.library.quadrupole import Quadrupole\n",
    "\n",
    "from model.command.external import load_lattice\n",
    "from model.command.build import build\n",
    "from model.command.tune import tune"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "cac77d3b-1d06-4f5a-8d37-dfcc0a0440a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set data type and device\n",
    "\n",
    "Element.dtype = dtype = torch.float64\n",
    "Element.device = device = torch.device('cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ac1e80f6-17bb-4f09-a87a-11a6796be8e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load lattice (ELEGANT table)\n",
    "# Note, lattice is allowed to have repeated elements\n",
    "\n",
    "path = Path('elettra.lte')\n",
    "data = load_lattice(path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "600c2cae-2c7a-4c28-9a8b-93506d5ae842",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'BPM': 168, 'Drift': 708, 'Dipole': 156, 'Quadrupole': 360, 'Marker': 12}"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Build and setup lattice\n",
    "\n",
    "ring:Line = build('RING', 'ELEGANT', data)\n",
    "\n",
    "# Flatten sublines\n",
    "\n",
    "ring.flatten()\n",
    "\n",
    "# Remove all marker elements but the ones starting with MLL (long straight section centers)\n",
    "\n",
    "ring.remove_group(pattern=r'^(?!MLL_).*', kinds=['Marker'])\n",
    "\n",
    "# Replace all sextupoles with quadrupoles\n",
    "\n",
    "def factory(element:Element) -> None:\n",
    "    table = element.serialize\n",
    "    table.pop('ms', None)\n",
    "    return Quadrupole(**table)\n",
    "\n",
    "ring.replace_group(pattern=r'', factory=factory, kinds=['Sextupole'])\n",
    "\n",
    "# Set linear dipoles\n",
    "\n",
    "def apply(element:Element) -> None:\n",
    "    element.linear = True\n",
    "\n",
    "ring.apply(apply, kinds=['Dipole'])\n",
    "\n",
    "# Merge drifts\n",
    "\n",
    "ring.merge()\n",
    "\n",
    "# Change lattice start\n",
    "\n",
    "ring.start = \"BPM_S01_01\"\n",
    "\n",
    "# Split BPMs\n",
    "\n",
    "ring.split((None, ['BPM'], None, None))\n",
    "\n",
    "# Roll lattice\n",
    "\n",
    "ring.roll(1)\n",
    "\n",
    "# Splice lattice\n",
    "\n",
    "ring.splice()\n",
    "\n",
    "# Describe\n",
    "\n",
    "ring.describe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "c24501ca-9126-415d-9ee3-4d0ab7af7a8b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set quadrupole names for global tune correction\n",
    "\n",
    "QF = [f'QF_S{i:02}_{j:02}' for j in [2, 3] for i in range(1, 12 + 1)]\n",
    "QD = [f'QD_S{i:02}_{j:02}' for j in [2, 3] for i in range(1, 12 + 1)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "3efa13ab-9cae-4c8c-90c4-9de588c64e9f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([2, 48])\n",
      "tensor([[ 0.2439,  0.2439,  0.2439,  0.2439,  0.2439,  0.2439,  0.2439,  0.2439,\n",
      "          0.2439,  0.2439,  0.2439,  0.2439,  0.2439,  0.2439,  0.2439,  0.2439,\n",
      "          0.2439,  0.2439,  0.2439,  0.2439,  0.2439,  0.2439,  0.2439,  0.2439,\n",
      "          0.0873,  0.0873,  0.0873,  0.0873,  0.0873,  0.0873,  0.0873,  0.0873,\n",
      "          0.0873,  0.0873,  0.0873,  0.0873,  0.0873,  0.0873,  0.0873,  0.0873,\n",
      "          0.0873,  0.0873,  0.0873,  0.0873,  0.0873,  0.0873,  0.0873,  0.0873],\n",
      "        [-0.1247, -0.1247, -0.1247, -0.1247, -0.1247, -0.1247, -0.1247, -0.1247,\n",
      "         -0.1247, -0.1247, -0.1247, -0.1247, -0.1247, -0.1247, -0.1247, -0.1247,\n",
      "         -0.1247, -0.1247, -0.1247, -0.1247, -0.1247, -0.1247, -0.1247, -0.1247,\n",
      "         -0.0525, -0.0525, -0.0525, -0.0525, -0.0525, -0.0525, -0.0525, -0.0525,\n",
      "         -0.0525, -0.0525, -0.0525, -0.0525, -0.0525, -0.0525, -0.0525, -0.0525,\n",
      "         -0.0525, -0.0525, -0.0525, -0.0525, -0.0525, -0.0525, -0.0525, -0.0525]],\n",
      "       dtype=torch.float64)\n",
      "tensor([[ 5.8543,  2.0964],\n",
      "        [-2.9918, -1.2602]], dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "# Compute response matrix (individual quadrupoles)\n",
    "\n",
    "def observable(knobs):\n",
    "    kn = knobs\n",
    "    return tune(ring, [kn], ('kn', None, QF + QD, None), matched=True, limit=1)\n",
    "\n",
    "knobs = torch.zeros(len(QF + QD), dtype=dtype)\n",
    "matrix = torch.func.jacrev(observable)(knobs)\n",
    "\n",
    "print(matrix.shape)\n",
    "print(matrix)\n",
    "print(matrix.reshape(2, 2, (len(QF) + len(QD)) // 2).sum(-1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "d698c502-4b3a-4e2d-a75a-6ea4e7ec65c7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([2, 2])\n",
      "tensor([[ 5.8543,  2.0964],\n",
      "        [-2.9918, -1.2602]], dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "# Compute response matrix (two quadrupole families)\n",
    "\n",
    "def observable(knobs):\n",
    "    kf, kd = knobs\n",
    "    kn = torch.stack(len(QF)*[kf] + len(QD)*[kd])\n",
    "    return tune(ring, [kn], ('kn', None, QF + QD, None), matched=True, limit=1)\n",
    "\n",
    "knobs = torch.tensor([0.0, 0.0], dtype=dtype)\n",
    "matrix = torch.func.jacrev(observable)(knobs)\n",
    "\n",
    "print(matrix.shape)\n",
    "print(matrix)"
   ]
  }
 ],
 "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.14.0"
  },
  "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
}