{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "262a5ec8-2553-4237-ab62-319b6ca22089",
   "metadata": {},
   "source": [
    "# Example-41: Tune (Tune and chormaticity correction)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a0651ae0-b59d-4120-87aa-9f99bf176156",
   "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": "31c6df34-e96e-46db-a050-b190e622ecbc",
   "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": "206944c0-8999-447e-8334-2d1e7901eb22",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build and setup lattice\n",
    "\n",
    "# Note, sextupoles are turned off and dipoles are linear \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",
    "# Turn off sextupoles and set linear dipoles\n",
    "\n",
    "for element in ring: \n",
    "    if element.__class__.__name__ == 'Sextupole':\n",
    "        element.ms = 0.0\n",
    "    if element.__class__.__name__ == 'Dipole':\n",
    "        element.linear = True\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": "6270642e-f562-413b-994c-5a14f288d57c",
   "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": "1fbdbfcf-2997-4f90-8c88-8e33d44ebbac",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[ 0.0187,  0.0217,  0.0461,  0.0601,  0.0484,  0.0297,  0.0299,  0.0488,\n",
      "          0.0597,  0.0457,  0.0214,  0.0192,  0.0669,  0.0667,  0.0191,  0.0216,\n",
      "          0.0461,  0.0597,  0.0483,  0.0296,  0.0297,  0.0485,  0.0595,  0.0459,\n",
      "          0.0209,  0.0196,  0.0683,  0.0647],\n",
      "        [-0.0668, -0.0283, -0.0136, -0.0078, -0.0141, -0.0205, -0.0202, -0.0137,\n",
      "         -0.0079, -0.0137, -0.0282, -0.0661, -0.0250, -0.0251, -0.0663, -0.0282,\n",
      "         -0.0137, -0.0079, -0.0138, -0.0203, -0.0205, -0.0140, -0.0078, -0.0137,\n",
      "         -0.0288, -0.0659, -0.0250, -0.0258]], dtype=torch.float64)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Compute tunes derivative with respect to deviations in quadrupole settings\n",
    "\n",
    "kn = torch.zeros(nq, dtype=torch.float64)\n",
    "\n",
    "pprint(torch.func.jacrev(lambda kn: tune(ring, [kn], ('kn', ['Quadrupole'], None, None), matched=True, limit=1, epsilon=None))(kn))\n",
    "print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "4fc6e27c-9298-40d6-93ae-c3d5991d36e0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[ 0.0051,  0.0144,  0.0143,  0.0050,  0.0049,  0.0142,  0.0143,  0.0051,\n",
      "          0.0050,  0.0143,  0.0142,  0.0050,  0.0049,  0.0141,  0.0142,  0.0050],\n",
      "        [-0.0047, -0.0037, -0.0038, -0.0049, -0.0049, -0.0038, -0.0038, -0.0049,\n",
      "         -0.0049, -0.0038, -0.0038, -0.0049, -0.0049, -0.0038, -0.0037, -0.0048]],\n",
      "       dtype=torch.float64)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Compute chromaticities derivatives \n",
    "\n",
    "ms = torch.zeros(ns, dtype=torch.float64)\n",
    "\n",
    "pprint(torch.func.jacrev(lambda ms: chromaticity(ring, [ms], ('ms', ['Sextupole'], None, None), matched=True, limit=1, epsilon=None))(ms))\n",
    "print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "46a3f83d-641b-4135-b565-77a3a84ee0f1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([2, 28])\n",
      "\n",
      "tensor(0.0471, dtype=torch.float64)\n",
      "tensor(0.0007, dtype=torch.float64)\n",
      "tensor(1.9547e-05, dtype=torch.float64)\n",
      "tensor(5.4013e-07, dtype=torch.float64)\n",
      "\n",
      "tensor([0.7500, 0.8500], dtype=torch.float64)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Tune correction\n",
    "\n",
    "# [nux, nuy] = M [..., kn_i, ..., ks_i, ...]\n",
    "\n",
    "kn = torch.zeros(nq, dtype=torch.float64)\n",
    "response = torch.func.jacrev(lambda kn: tune(ring, [kn], ('kn', ['Quadrupole'], None, None), matched=True, limit=1, epsilon=None))(kn)\n",
    "print(response.shape)\n",
    "print()\n",
    "\n",
    "# Set target tunes\n",
    "\n",
    "nuqx_target = torch.tensor(0.75, dtype=torch.float64)\n",
    "nuqy_target = torch.tensor(0.85, dtype=torch.float64)\n",
    "\n",
    "# Perform correction (model to experiment)\n",
    "\n",
    "lr = 1.0\n",
    "\n",
    "kn = torch.zeros(nq, dtype=torch.float64)\n",
    "\n",
    "for _ in range(4):\n",
    "    nuqx, nuqy = tune(ring, [kn], ('kn', ['Quadrupole'], None, None), matched=True, limit=1, epsilon=None)\n",
    "    dkn = (- lr*torch.linalg.lstsq(response, torch.stack([nuqx, nuqy]) - torch.stack([nuqx_target, nuqy_target]), driver='gels').solution)\n",
    "    kn += dkn\n",
    "    print((torch.stack([nuqx, nuqy]) - torch.stack([nuqx_target, nuqy_target])).norm())\n",
    "print()\n",
    "\n",
    "print(tune(ring, [kn], ('kn', ['Quadrupole'], None, None), matched=True, limit=1, epsilon=None))\n",
    "print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "eb164d28-fa33-42a9-8952-75e43e819337",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([2, 16])\n",
      "\n",
      "tensor(10.1673, dtype=torch.float64)\n",
      "\n",
      "tensor([1.0000, 1.0000], dtype=torch.float64)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Chromaticity correction\n",
    "# Note, one iteration is sufficient for chromaticity correction\n",
    "\n",
    "# [psix, psiy] = M [..., ms_i, ..., ms_i, ...]\n",
    "\n",
    "ms = torch.zeros(ns, dtype=torch.float64)\n",
    "response = torch.func.jacrev(lambda ms: chromaticity(ring, [ms], ('ms', ['Sextupole'], None, None), matched=True, limit=1, epsilon=None))(ms)\n",
    "print(response.shape)\n",
    "print()\n",
    "\n",
    "# Set target tunes\n",
    "\n",
    "psiqx_target = torch.tensor(1.0, dtype=torch.float64)\n",
    "psiqy_target = torch.tensor(1.0, dtype=torch.float64)\n",
    "\n",
    "# Perform correction (model to experiment)\n",
    "\n",
    "lr = 1.0\n",
    "\n",
    "ks = torch.zeros(ns, dtype=torch.float64)\n",
    "\n",
    "for _ in range(1):\n",
    "    psiqx, psiqy = chromaticity(ring, [ms], ('ms', ['Sextupole'], None, None), matched=True, limit=1, epsilon=None)\n",
    "    dms = (- lr*torch.linalg.lstsq(response, torch.stack([psiqx, psiqy]) - torch.stack([psiqx_target, psiqy_target]), driver='gels').solution)\n",
    "    ms += dms\n",
    "    print((torch.stack([psiqx, psiqy]) - torch.stack([psiqx_target, psiqy_target])).norm())\n",
    "print()\n",
    "\n",
    "print(chromaticity(ring, [ms], ('ms', ['Sextupole'], None, None), matched=True, limit=1, epsilon=None))\n",
    "print()"
   ]
  }
 ],
 "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
}