{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "262a5ec8-2553-4237-ab62-319b6ca22089",
   "metadata": {},
   "source": [
    "# Example-40: Tune (Computation of tunes and chromaticities)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "40d4159c-3d88-4c31-8cd3-3bc34fe94f37",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import\n",
    "\n",
    "from random import random\n",
    "from pprint import pprint\n",
    "\n",
    "import torch\n",
    "from torch import Tensor\n",
    "\n",
    "from pathlib import Path\n",
    "\n",
    "from model.library.line import Line\n",
    "\n",
    "from model.command.external import load_sdds\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": 2,
   "id": "4e29b9cf-2c13-40cf-9876-657766a84ca4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load ELEGANT twiss\n",
    "\n",
    "path = Path('ic.twiss')\n",
    "parameters, columns = load_sdds(path)\n",
    "\n",
    "nu_qx:Tensor = torch.tensor(parameters['nux'] % 1, dtype=torch.float64)\n",
    "nu_qy:Tensor = torch.tensor(parameters['nuy'] % 1, dtype=torch.float64)\n",
    "\n",
    "psi_qx:Tensor = torch.tensor(parameters['dnux/dp'], dtype=torch.float64)\n",
    "psi_qy:Tensor = torch.tensor(parameters['dnuy/dp'], dtype=torch.float64)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "41eadf68-1357-481d-9350-766fa29e4270",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build and setup lattice\n",
    "\n",
    "# Load ELEGANT table\n",
    "\n",
    "path = Path('ic.lte')\n",
    "data = load_lattice(path)\n",
    "\n",
    "# Build ELEGANT table\n",
    "\n",
    "ring:Line = build('RING', 'ELEGANT', data)\n",
    "ring.flatten()\n",
    "\n",
    "# Merge drifts\n",
    "\n",
    "ring.merge()\n",
    "\n",
    "# Set exact dipoles\n",
    "\n",
    "for element in ring: \n",
    "    if element.__class__.__name__ == 'Dipole':\n",
    "        element.exact = True\n",
    "\n",
    "# Set number of integration steps and integration order in sextupoles and dipoles\n",
    "\n",
    "ring.ns    = (('Sextupole', 0.01), ('Dipole', 0.01))\n",
    "ring.order = (('Sextupole', 1), ('Dipole', 1))\n",
    "\n",
    "# Set number of elements of different kinds\n",
    "\n",
    "nb = ring.describe['BPM']\n",
    "nq = ring.describe['Quadrupole']\n",
    "ns = ring.describe['Sextupole']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "829b26dc-f98e-4a37-a92b-164baa04ea00",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n",
      "True\n"
     ]
    }
   ],
   "source": [
    "# Compute tunes (fractional part)\n",
    "\n",
    "guess = torch.tensor(4*[0.0], dtype=torch.float64)\n",
    "nuqx, nuqy = tune(ring, [], alignment=False, matched=True, guess=guess, limit=8, epsilon=1.0E-9)\n",
    "\n",
    "# Compare with elegant\n",
    "\n",
    "print(torch.allclose(nu_qx, nuqx))\n",
    "print(torch.allclose(nu_qy, nuqy))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "bd72f81a-b532-431e-bd37-4f4d3dfa3ba9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n",
      "True\n"
     ]
    }
   ],
   "source": [
    "# Compute chormaticities\n",
    "\n",
    "psiqx, psiqy = chromaticity(ring, [], alignment=False, matched=True, guess=guess, limit=1, epsilon=None)\n",
    "\n",
    "print(torch.allclose(psi_qx, psiqx))\n",
    "print(torch.allclose(psi_qy, psiqy))\n",
    "\n",
    "# Note, since closed orbit depends on dp, matched flag should be True\n",
    "# These chromaticity values depend on the number of intergation steps in sextupoles\n",
    "# Without linear flag, dipoles will also have sextupole like component in the body due to cylindrical potential\n",
    "# With exact flag in dipoles, fringes will also contribute to chromaticities\n",
    "# These result match (converge to) MADX 5.09.01 TWISS (default settings in SBENDs) and ELEGNAT SVN revision: 30188M for CSBENDs with FINT=0.0 and EDGE_ORDER=0"
   ]
  }
 ],
 "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
}