{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "262a5ec8-2553-4237-ab62-319b6ca22089",
   "metadata": {},
   "source": [
    "# Example-58: Frequency (parametric derivatives for linear system)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "21defbfe-5f6e-4e91-94dd-a8267d6144d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# In this example frequencies for linear model are computed from trajectory data\n",
    "# Frequency values are compared with ones obtained from one-turn matrix\n",
    "# Derivatives with respect to parameters are also computed and compared"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "6302e8bb-59b6-47fc-a81d-717dadda6e06",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import\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\n",
    "from model.command.trajectory import trajectory\n",
    "from model.command.frequency import filter\n",
    "from model.command.frequency import frequency_factory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "93add4df-b6ac-475f-b641-d00cf202ea3a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load ELEGANT twiss\n",
    "\n",
    "path = Path('ic.twiss')\n",
    "parameters, columns = load_sdds(path)\n",
    "\n",
    "NUX:Tensor = torch.tensor(parameters['nux'] % 1, dtype=torch.float64)\n",
    "NUY:Tensor = torch.tensor(parameters['nuy'] % 1, dtype=torch.float64)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "4aa78283-6227-47bd-a8d0-816e2bf3bc03",
   "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",
    "# 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",
    "# Split BPMs\n",
    "\n",
    "ring.split((None, ['BPM'], None, None))\n",
    "\n",
    "# Roll lattice start\n",
    "\n",
    "ring.roll(1)\n",
    "\n",
    "# Split lattice into lines by BPMs\n",
    "\n",
    "ring.splice()\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": 5,
   "id": "59f8426b-30da-4540-b98c-d43720b3e0e8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(3.1086e-15, dtype=torch.float64)\n",
      "tensor(5.5511e-16, dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "# Compute tunes (one-turn matrix)\n",
    "\n",
    "nux, nuy = tune(ring, [], alignment=False, matched=True, limit=8, epsilon=1.0E-12)\n",
    "\n",
    "# Compare with elegant\n",
    "\n",
    "print((NUX - nux).abs())\n",
    "print((NUY - nuy).abs())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "bef63178-4ba8-47e2-81dd-5cc6b1f5e926",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(3.3061e-12, dtype=torch.float64)\n",
      "tensor(3.1350e-10, dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "# Compute tunes (trajectory)\n",
    "\n",
    "# Set trajectory generator\n",
    "\n",
    "generator = trajectory(ring, [0], matched=True)\n",
    "\n",
    "# Set initial condition\n",
    "\n",
    "state = torch.tensor([+1.0E-9, 0.0, -1.0E-9, 0.0], dtype=torch.float64)\n",
    "\n",
    "# Set window data\n",
    "\n",
    "window = filter(2**10, 1.0, dtype=ring.dtype, device=ring.device)\n",
    "\n",
    "# Set frequency generator\n",
    "\n",
    "frequency = frequency_factory(generator)\n",
    "\n",
    "# Compute frequencies\n",
    "\n",
    "nux, nuy = frequency(window, state)\n",
    "\n",
    "# Compare with elegant\n",
    "\n",
    "print((NUX - nux).abs())\n",
    "print((NUY - nuy).abs())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "43d44857-d79a-469d-8f38-b167ad480efa",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[-7.8819],\n",
      "        [-3.9483]], dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "# Derivative with respect to momentum deviation (one-turn matrix)\n",
    "\n",
    "dp = torch.tensor([0.0], dtype=torch.float64)\n",
    "\n",
    "print(torch.func.jacrev(lambda dp: tune(ring, [dp], ('dp', None, None, None), matched=True, limit=1, epsilon=None))(dp))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "827913f7-2278-40c2-b88d-d41758c3192d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[-7.8819],\n",
      "        [-3.9486]], dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "# Derivative with respect to momentum deviation (trajectory)\n",
    "\n",
    "# Set parametric trajectory generator\n",
    "\n",
    "generator = trajectory(ring, [0], ('dp', None, None, None), matched=True)\n",
    "\n",
    "# Set initial state and momentum deviation\n",
    "# Note, state should not be equal to zero, since zero is a fixed point,\n",
    "\n",
    "state = torch.tensor([+1.0E-9, 0.0, -1.0E-9, 0.0], dtype=torch.float64)\n",
    "dp = torch.tensor([0.0], dtype=torch.float64)\n",
    "\n",
    "# Set window data\n",
    "\n",
    "window = filter(2**10, 1.0, dtype=ring.dtype, device=ring.device)\n",
    "\n",
    "# Set frequency generator\n",
    "\n",
    "frequency = frequency_factory(generator)\n",
    "\n",
    "# Compute derivative\n",
    "\n",
    "print(torch.func.jacrev(lambda dp: frequency(window, state, dp), chunk_size=256)(dp))"
   ]
  }
 ],
 "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
}