{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3f3418c9-1c41-4dec-8c4f-b4fec4e76eb9",
   "metadata": {},
   "source": [
    "## ELETTRA-44: ID effect on chromaticity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "23341b14-0379-4230-a0df-4471b0a4b72b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# In this example effect of ID on natural chomaticity change is illustrated using an apple-ii ID type\n",
    "# The chromaticities are computed using leading order chromatic trasport matrix of an ID obtained from numerical tracking"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "f2f9b81f-2aa4-4899-a6e9-5e38ab2c2e51",
   "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",
    "from model.library.matrix import Matrix\n",
    "\n",
    "from model.command.external import load_lattice\n",
    "from model.command.build import build\n",
    "from model.command.tune import tune\n",
    "from model.command.tune import chromaticity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "53b2d155-a3d6-478b-8f81-d3ff7ad16570",
   "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": "397498c5-cc82-4647-95ee-79453aace902",
   "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": "1e55c1e4-78a0-421d-b8df-7b86c226719c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'BPM': 168, 'Drift': 720, 'Dipole': 156, 'Quadrupole': 360, 'Marker': 24}"
      ]
     },
     "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 or MSS (long and short straight section centers)\n",
    "\n",
    "ring.remove_group(pattern=r'^(?!MSS_)(?!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": "73067e15-eb04-4af1-a3a5-c5848b5a5962",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute tunes (fractional part)\n",
    "\n",
    "nux, nuy = tune(ring, [], matched=True, limit=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "9993f52c-51bb-40a6-93bd-9f02a6819401",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute chromaticity (natural)\n",
    "\n",
    "psix, psiy = chromaticity(ring, [])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "a4695c67-0e4c-4e0a-a1dc-c087efbcb5b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ID\n",
    "\n",
    "A = torch.tensor([[-0.03484222052711237, 1.0272120741819959E-7, -4.698931299341201E-9, 0.0015923185492594811],\n",
    "                  [1.0272120579834892E-7, -0.046082787920135176, 0.0017792061173117564, 3.3551298301095784E-8],\n",
    "                  [-4.6989312853101E-9, 0.0017792061173117072, 0.056853750760983084, -1.5929605363332683E-7],\n",
    "                  [0.0015923185492594336, 3.3551298348653296E-8, -1.5929605261642905E-7, 0.08311631737263032]], dtype=dtype)\n",
    "\n",
    "B = torch.tensor([[0.03649353186115209, 0.0015448347221877217, 0.00002719892025520868, -0.0033681183134964482],\n",
    "                  [0.0015448347221877217, 0.13683886657005795, -0.0033198692682377406, 0.00006140578258682469],\n",
    "                  [0.00002719892025520868, -0.0033198692682377406, -0.05260095308967722, 0.005019907688182885],\n",
    "                  [-0.0033681183134964482, 0.00006140578258682469, 0.005019907688182885, -0.2531573249456863]], dtype=dtype)\n",
    "\n",
    "ID = Matrix('ID', \n",
    "                   length=0.0, \n",
    "                   A=A[torch.triu(torch.ones_like(A, dtype=torch.bool))].tolist(), \n",
    "                   B=B[torch.triu(torch.ones_like(B, dtype=torch.bool))].tolist())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "1fc57833-765b-4a7d-abb6-52ee079e402f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'BPM': 168,\n",
       " 'Drift': 720,\n",
       " 'Dipole': 156,\n",
       " 'Quadrupole': 360,\n",
       " 'Matrix': 1,\n",
       " 'Marker': 23}"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Insert ID into the existing lattice\n",
    "\n",
    "\n",
    "elements = [ID]\n",
    "markers = ['MLL_S01']\n",
    "\n",
    "error = ring.clone()\n",
    "error.flatten()\n",
    "for element, marker in zip(elements, markers):\n",
    "    error.insert(element, marker, position=0.0)\n",
    "error.splice()\n",
    "\n",
    "error.describe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "d1288051-abb3-411a-81d8-3e447d851639",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute tunes (fractional part)\n",
    "\n",
    "nux_id, nuy_id = tune(error, [], matched=True, limit=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "91ac3574-6ef7-49a8-ab09-e5885eb4e605",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute chromaticity (natural)\n",
    "\n",
    "psix_id, psiy_id = chromaticity(error, [])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "13e99aa2-96de-4e66-b540-fdfd0b9c7972",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.0260, dtype=torch.float64)\n",
      "tensor(-0.0114, dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "# Tune shifts\n",
    "\n",
    "print((nux - nux_id))\n",
    "print((nuy - nuy_id))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "be09a697-fbed-4405-be77-933b97241bb5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(1.0504, dtype=torch.float64)\n",
      "tensor(-0.0106, dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "# Chromaticity change\n",
    "\n",
    "print(psix - psix_id)\n",
    "print(psiy - psiy_id)"
   ]
  }
 ],
 "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
}