{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3a839e30-e425-4ae8-8939-8ba78758b62b",
   "metadata": {},
   "source": [
    "## Example-64: AT ID"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a25842d5-eb68-46e5-9971-3fb3059664bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import\n",
    "\n",
    "import torch\n",
    "from torch import Tensor\n",
    "\n",
    "from pathlib import Path\n",
    "from tqdm import tqdm\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.corrector import Corrector\n",
    "from model.library.quadrupole import Quadrupole\n",
    "from model.library.matrix import Matrix\n",
    "from model.library.kickmap import KM\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\n",
    "from model.command.orbit import dispersion\n",
    "from model.command.orbit import ORM\n",
    "from model.command.twiss import twiss\n",
    "from model.command.advance import advance\n",
    "from model.command.coupling import coupling\n",
    "\n",
    "import at\n",
    "from model.interface.at import convert"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "880538da-2a0d-4815-bfaa-1ba4e41240e3",
   "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": 3,
   "id": "af35d0f6-feb7-4344-b64e-e0e9ba5b00ef",
   "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": 4,
   "id": "95df0c74-8a37-430e-ac24-628497f7f461",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'BPM': 168,\n",
       " 'Drift': 744,\n",
       " 'Dipole': 156,\n",
       " 'Sextupole': 240,\n",
       " 'Quadrupole': 120,\n",
       " 'Corrector': 24,\n",
       " 'Marker': 24}"
      ]
     },
     "execution_count": 4,
     "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'^(?!MSS_)(?!MLL_).*', kinds=['Marker'])\n",
    "\n",
    "\n",
    "# Set sextupole integration order and step size\n",
    "\n",
    "ring.order = (('Sextupole', 1), )\n",
    "ring.ns = (('Sextupole', 0.01), )\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",
    "# Insert correctors\n",
    "\n",
    "for name, *_ in ring.layout():\n",
    "    if name.startswith('CH'):\n",
    "        corrector = Corrector(f'{name}_CXY', factor=1)\n",
    "        ring.split((1 + 1, None, [name], None), paste=[corrector])\n",
    "\n",
    "# Merge drifts\n",
    "\n",
    "ring.merge()\n",
    "\n",
    "# Change lattice start 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\n",
    "\n",
    "ring.splice()\n",
    "\n",
    "# Describe\n",
    "\n",
    "ring.describe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "e4a5aa55-2c04-471b-aff6-0ffe1f1dc800",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute tunes (fractional part)\n",
    "\n",
    "nux, nuy = tune(ring, [], matched=True, limit=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "fb09d44a-0b89-4929-9c89-cd41e4330573",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute dispersion\n",
    "\n",
    "orbit = torch.tensor(4*[0.0], dtype=dtype, device=device)\n",
    "etaqx, etapx, etaqy, etapy = dispersion(ring, orbit, [], limit=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "0b7bfe54-508a-4a63-845d-8740ec350b3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute twiss parameters\n",
    "\n",
    "ax, bx, ay, by = twiss(ring, [], matched=True, advance=True, full=False).T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "cb470980-9354-48fb-9f19-f6d226a318f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute phase advances\n",
    "\n",
    "mux, muy = advance(ring, [], alignment=False, matched=True).T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "7502f7c8-5bdc-4fbb-89c2-bf0f7c42e3f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute coupling\n",
    "\n",
    "c = coupling(ring, [])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "68297900-e449-4641-89e9-6eb878bec464",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute chromaticity\n",
    "\n",
    "psi = chromaticity(ring, [], matched=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "6821ca5e-2ec5-44ba-a154-522852b6573d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define ID (linear model)\n",
    "\n",
    "A = torch.tensor([[-0.0344386, 0., 0., 0.], [0., -0.0445673, 0., 0.], [0., 0., 0.056303, 0.], [0., 0., 0., 0.0804237]], dtype=dtype)\n",
    "\n",
    "CENTER = 0.0\n",
    "ID = Matrix('ID', length=0.0, A=A[torch.triu(torch.ones_like(A, dtype=torch.bool))].tolist())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "4119afbf-40a5-4e3f-8e02-563f34327f15",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'BPM': 168,\n",
       " 'Drift': 745,\n",
       " 'Dipole': 156,\n",
       " 'Sextupole': 240,\n",
       " 'Quadrupole': 120,\n",
       " 'Corrector': 24,\n",
       " 'Marker': 24,\n",
       " 'Matrix': 1}"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Insert ID (linear model)\n",
    "\n",
    "error = ring.clone()\n",
    "error.flatten()\n",
    "error.insert(ID, error.next('MLL_S01').name, position=CENTER*1.0E-3)\n",
    "error.splice()\n",
    "error.describe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "57dcc2f7-415e-4a35-9416-48e7bab9414a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate AT lattice\n",
    "\n",
    "lattice = convert(error, energy=2.4, alignment=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "0bbf0795-5ee1-4150-abe4-f4b4106c975c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "168\n"
     ]
    }
   ],
   "source": [
    "# Compute optics\n",
    "\n",
    "refpts = lattice.uint32_refpts(at.Monitor)\n",
    "print(len(refpts))\n",
    "_, ringdata, elemdata = lattice.get_optics(refpts=refpts, get_chrom=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "ff662562-5b6e-4799-aab5-0ec3afcaafb4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set optics at BPMs\n",
    "\n",
    "nux_id_lm_at, nuy_id_lm_at = ringdata['tune']\n",
    "psi_id_lm_at = ringdata['chromaticity']\n",
    "\n",
    "ax_id_lm_at, ay_id_lm_at = elemdata['alpha'].T\n",
    "bx_id_lm_at, by_id_lm_at = elemdata['beta'].T\n",
    "mux_id_lm_at, muy_id_lm_at = elemdata['mu'].T\n",
    "etaqx_id_lm_at, etapx_id_lm_at, etaqy_id_lm_at, etapy_id_lm_at = elemdata['dispersion'].T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "32250bd0-abc9-441e-aadd-aa746faac32d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute tunes (fractional part)\n",
    "\n",
    "nux_id_lm, nuy_id_lm = tune(error, [], matched=True, limit=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "2e95283f-04c3-437e-971d-689ba92803f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute dispersion\n",
    "\n",
    "orbit = torch.tensor(4*[0.0], dtype=dtype)\n",
    "etaqx_id_lm, etapx_id_lm, etaqy_id_lm, etapy_id_lm = dispersion(error, orbit, [], limit=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "91aae882-8891-428d-a881-410f77725786",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute twiss parameters\n",
    "\n",
    "ax_id_lm, bx_id_lm, ay_id_lm, by_id_lm = twiss(error, [], matched=True, advance=True, full=False).T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "40acdcfc-2ade-4d1a-9eab-5dd7f16b70b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute phase advances\n",
    "\n",
    "mux_id_lm, muy_id_lm = advance(error, [], alignment=False, matched=True).T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "bea65a67-7ae1-4dd2-bd46-49c88133e182",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute coupling\n",
    "\n",
    "c_id_lm = coupling(error, [])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "38cd6692-5bf7-45ab-8734-b7e69fc6cfe7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute chromaticity\n",
    "\n",
    "psi_id_lm = chromaticity(error, [], matched=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "c51add9c-48b9-4ec0-bcc0-cbd649a89e66",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.02566042669989943\n",
      "-0.011228413019111483\n",
      "\n",
      "0.025711104422837916\n",
      "-0.011228053550816186\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Tune shifts\n",
    "\n",
    "print((nux - nux_id_lm).numpy())\n",
    "print((nuy - nuy_id_lm).numpy())\n",
    "print()\n",
    "\n",
    "print((nux - nux_id_lm_at).numpy())\n",
    "print((nuy - nuy_id_lm_at).numpy())\n",
    "print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "d5d883a4-5788-4e52-b9bb-3bae2e03c48a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2.02960315 2.01312355]\n",
      "[1.71483742 2.01653988]\n",
      "[1.71445621 2.01656829]\n"
     ]
    }
   ],
   "source": [
    "# Chromaticity\n",
    "\n",
    "print(psi.numpy())\n",
    "print(psi_id_lm.numpy())\n",
    "print(psi_id_lm_at)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "eb2e6d73-e2e7-421f-87c6-410bfa8c30b1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ 2.40439332e-10, -1.21810204e-37,  1.71831614e-06])"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# AT coupling\n",
    "\n",
    "index, *_ = lattice.get_uint32_index(\"MSS_S10\")\n",
    "rf = at.RFCavity(\"CAV\", 0.0, 2_000_000.0, 499_654_096.666667, 432, 2_400_000_000.0, PassMethod=\"CavityPass\")\n",
    "lattice.insert(int(index) + 1, rf)\n",
    "lattice.enable_6d()\n",
    "_, data, _ = lattice.ohmi_envelope(refpts=refpts)\n",
    "data.mode_emittances"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "2446dd29-2ab3-40a9-bc02-9dc51144d676",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define ID (load kick map)\n",
    "\n",
    "CENTER = 0.0\n",
    "ID = KM(name='ID', path=Path('id.mat'), energy=2.4, count=40, insertion=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "fb3f6519-3004-4974-a6d9-6c03755a8a12",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'BPM': 168,\n",
       " 'Drift': 745,\n",
       " 'Dipole': 156,\n",
       " 'Sextupole': 240,\n",
       " 'Quadrupole': 120,\n",
       " 'Corrector': 24,\n",
       " 'Marker': 24,\n",
       " 'KM': 1}"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Insert ID (kick map)\n",
    "\n",
    "error = ring.clone()\n",
    "error.flatten()\n",
    "error.insert(ID, error.next('MLL_S01').name, position=CENTER*1.0E-3)\n",
    "error.splice()\n",
    "error.describe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "4b078716-110f-4874-8349-28c9a73ce568",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate AT lattice\n",
    "\n",
    "lattice = convert(error, energy=2.4, alignment=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "6b210700-f115-4e93-944e-559ebcc3d500",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "168\n"
     ]
    }
   ],
   "source": [
    "# Compute optics\n",
    "\n",
    "refpts = lattice.uint32_refpts(at.Monitor)\n",
    "print(len(refpts))\n",
    "_, ringdata, elemdata = lattice.get_optics(refpts=refpts, get_chrom=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "50a2dd2c-044d-4153-8e1a-c63e4996f6bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set optics at BPMs\n",
    "\n",
    "nux_id_km_at, nuy_id_km_at = ringdata['tune']\n",
    "psi_id_km_at = ringdata['chromaticity']\n",
    "\n",
    "ax_id_km_at, ay_id_km_at = elemdata['alpha'].T\n",
    "bx_id_km_at, by_id_km_at = elemdata['beta'].T\n",
    "mux_id_km_at, muy_id_km_at = elemdata['mu'].T\n",
    "etaqx_id_km_at, etapx_id_km_at, etaqy_id_km_at, etapy_id_km_at = elemdata['dispersion'].T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "19760f7a-45be-4e50-9b37-96dd88e6cc42",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute tunes (fractional part)\n",
    "\n",
    "nux_id_km, nuy_id_km = tune(error, [], matched=True, limit=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "9eba1822-9524-468b-ae3c-2a386785c871",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute dispersion\n",
    "\n",
    "orbit = torch.tensor(4*[0.0], dtype=dtype)\n",
    "etaqx_id_km, etapx_id_km, etaqy_id_km, etapy_id_km = dispersion(error, orbit, [], limit=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "b496168a-7af4-450b-a054-be36d0c9f5b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute twiss parameters\n",
    "\n",
    "ax_id_km, bx_id_km, ay_id_km, by_id_km = twiss(error, [], matched=True, advance=True, full=False).T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "6a070d5d-9edc-4a41-a1c0-3766c01f8ee1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute phase advances\n",
    "\n",
    "mux_id_km, muy_id_km = advance(error, [], alignment=False, matched=True).T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "792d8b02-fae6-4b61-833d-886356e9ec50",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute coupling\n",
    "\n",
    "c_id_km = coupling(error, [])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "d2b5bedb-fda9-4b55-a778-c0f60562d5fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute chromaticity\n",
    "\n",
    "psi_id_km = chromaticity(error, [], matched=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "94caa7b3-ba64-4aee-98d0-93477d9a7df1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.025636790558097755\n",
      "-0.011262894479845187\n",
      "\n",
      "0.02568748173391394\n",
      "-0.011262534209166342\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Tune shifts\n",
    "\n",
    "print((nux - nux_id_km).numpy())\n",
    "print((nuy - nuy_id_km).numpy())\n",
    "print()\n",
    "\n",
    "print((nux - nux_id_km_at).numpy())\n",
    "print((nuy - nuy_id_km_at).numpy())\n",
    "print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "9dc593b9-5b79-4f53-b9c6-b596e768a04a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2.02960315 2.01312355]\n",
      "[1.74225348 1.99775283]\n",
      "[1.74187601 1.99778124]\n"
     ]
    }
   ],
   "source": [
    "# Chromaticity\n",
    "\n",
    "print(psi.numpy())\n",
    "print(psi_id_km.numpy())\n",
    "print(psi_id_km_at)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "0c340471-9e5e-448a-83fd-ab9ca2eacc2a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([2.40431968e-10, 2.77736999e-25, 1.71831568e-06])"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# AT coupling\n",
    "\n",
    "index, *_ = lattice.get_uint32_index(\"MSS_S10\")\n",
    "rf = at.RFCavity(\"CAV\", 0.0, 2_000_000.0, 499_654_096.666667, 432, 2_400_000_000.0, PassMethod=\"CavityPass\")\n",
    "lattice.insert(int(index) + 1, rf)\n",
    "lattice.enable_6d()\n",
    "_, data, _ = lattice.ohmi_envelope(refpts=refpts)\n",
    "data.mode_emittances"
   ]
  }
 ],
 "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
}