{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "262a5ec8-2553-4237-ab62-319b6ca22089",
   "metadata": {},
   "source": [
    "# Example-53: Coupling (Minimal tune distance computation)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "41b7abeb-9659-4fd3-9539-46b25814e969",
   "metadata": {},
   "outputs": [],
   "source": [
    "# In this example minimal tune distance computation is illustrated"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "aa403f46-016e-4d41-8a58-6979fe5af620",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import\n",
    "\n",
    "from pprint import pprint\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",
    "\n",
    "from model.library.line import Line\n",
    "from model.library.gradient import Gradient\n",
    "\n",
    "from model.command.external import load_sdds\n",
    "from model.command.external import load_lattice\n",
    "\n",
    "from model.command.build import build\n",
    "\n",
    "from model.command.mapping import matrix\n",
    "from model.command.tune import tune\n",
    "from model.command.twiss import twiss\n",
    "from model.command.coupling import coupling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "9a9db05c-5916-4601-913e-60eaf89f000f",
   "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 linear dipoles\n",
    "\n",
    "for element in ring: \n",
    "    if element.__class__.__name__ == 'Dipole':\n",
    "        element.linear = True\n",
    "\n",
    "# Add gradient element to the lattice end\n",
    "\n",
    "G = Gradient('G')\n",
    "ring.append(G)\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": "073506ad-a66b-4137-8267-5ae81417d52d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([1.0000e-03, 4.0000e-05, 5.0000e-03, 6.0000e-05], dtype=torch.float64)\n",
      "tensor([[ 1.0000,  0.0000,  0.0000,  0.0000],\n",
      "        [-0.0100,  1.0000,  0.0100,  0.0000],\n",
      "        [ 0.0000,  0.0000,  1.0000,  0.0000],\n",
      "        [ 0.0100,  0.0000,  0.0100,  1.0000]], dtype=torch.float64)\n",
      "\n",
      "tensor([1.0000e-03, 4.0000e-05, 5.0000e-03, 6.0000e-05], dtype=torch.float64)\n",
      "tensor([[ 1.0000,  0.0000,  0.0000,  0.0000],\n",
      "        [-0.0100,  1.0000,  0.0100,  0.0000],\n",
      "        [ 0.0000,  0.0000,  1.0000,  0.0000],\n",
      "        [ 0.0100,  0.0000,  0.0100,  1.0000]], dtype=torch.float64)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Gradient element is equivalent to the following linear transformation\n",
    "\n",
    "def gradient(state, kn, ks):\n",
    "    (qx, px, qy, py), kn, ks = state, kn, ks\n",
    "    return torch.stack([qx, px - kn*qx + ks*qy, qy, py + ks*qx + kn*qy])\n",
    "\n",
    "state = torch.tensor([0.001, 0.0, 0.005, 0.0], dtype=torch.float64)\n",
    "\n",
    "kn = torch.tensor(0.01, dtype=torch.float64)\n",
    "ks = torch.tensor(0.01, dtype=torch.float64)\n",
    "\n",
    "print(gradient(state, kn, ks))\n",
    "print(torch.func.jacrev(gradient)(state, kn, ks))\n",
    "print()\n",
    "\n",
    "print(G(state, data={'kn': kn, 'ks': ks, 'dp': 0.0}))\n",
    "print(torch.func.jacrev(lambda state: G(state, data={'kn': kn, 'ks': ks, 'dp': 0.0}))(state))\n",
    "print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c1633454-48d7-4932-9d14-99fc5e5d6797",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[-4.4769, -2.6375,  0.0000,  0.0000],\n",
      "        [ 8.9592,  5.0549,  0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000,  3.9054, -2.0383],\n",
      "        [ 0.0000,  0.0000,  5.5965, -2.6648]], dtype=torch.float64)\n",
      "\n",
      "tensor([[-4.4769, -2.6375,  0.0000,  0.0000],\n",
      "        [ 8.9592,  5.0549,  0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000,  3.9054, -2.0383],\n",
      "        [ 0.0000,  0.0000,  5.5965, -2.6648]], dtype=torch.float64)\n",
      "\n",
      "tensor([[-4.4769, -2.6375,  0.0000,  0.0000],\n",
      "        [ 8.9592,  5.0549,  0.0391, -0.0204],\n",
      "        [ 0.0000,  0.0000,  3.9054, -2.0383],\n",
      "        [-0.0448, -0.0264,  5.5965, -2.6648]], dtype=torch.float64)\n",
      "\n",
      "tensor([[-4.4769, -2.6375,  0.0000,  0.0000],\n",
      "        [ 8.9592,  5.0549,  0.1953, -0.1019],\n",
      "        [ 0.0000,  0.0000,  3.9054, -2.0383],\n",
      "        [-0.2238, -0.1319,  5.5965, -2.6648]], dtype=torch.float64)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Compute one-turn matrix without and with skew gradient error\n",
    "\n",
    "state = torch.tensor(4*[0.0], dtype=torch.float64)\n",
    "ks = torch.tensor([0.01], dtype=torch.float64)\n",
    "\n",
    "fmt, *_ = matrix(ring, 0, len(ring) - 1, ('ks', None, ['G'], None), matched=False)\n",
    "\n",
    "\n",
    "print(torch.func.jacrev(ring)(state))\n",
    "print()\n",
    "\n",
    "print(fmt(state, 0.0*ks))\n",
    "print()\n",
    "\n",
    "print(fmt(state, 1.0*ks))\n",
    "print()\n",
    "\n",
    "print(torch.func.jacrev(lambda state: G(state, data={'kn': 0.0, 'ks': 0.05, 'dp': 0.0}))(state) @ fmt(state, 0.0*ks))\n",
    "print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "87803ab6-48b6-4a43-8a42-1fe8945e4b64",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(6.7364e-05, dtype=torch.float64)\n",
      "tensor(-8.2241e-05, dtype=torch.float64)\n",
      "\n",
      "tensor(0.0071, dtype=torch.float64)\n",
      "tensor(0.0037, dtype=torch.float64)\n",
      "tensor(-0.0039, dtype=torch.float64)\n",
      "tensor(0.0021, dtype=torch.float64)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Compute anjd compare twiss parameters at lattice start\n",
    "\n",
    "ks = torch.tensor([0.01], dtype=torch.float64)\n",
    "\n",
    "nux_a, nuy_a = tune(ring, [0.0*ks], ('ks', None, ['G'], None), matched=False)\n",
    "ax_a, bx_a, ay_a, by_a = twiss(ring, [0.0*ks], ('ks', None, ['G'], None), matched=False)\n",
    "\n",
    "nux_b, nuy_b = tune(ring, [1.0*ks], ('ks', None, ['G'], None), matched=False)\n",
    "ax_b, bx_b, ay_b, by_b = twiss(ring, [1.0*ks], ('ks', None, ['G'], None), matched=False)\n",
    "\n",
    "print(nux_a - nux_b)\n",
    "print(nuy_a - nuy_b)\n",
    "print()\n",
    "\n",
    "print(ax_a - ax_b)\n",
    "print(bx_a - bx_b)\n",
    "print(ay_a - ay_b)\n",
    "print(by_a - by_b)\n",
    "print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "8e8acf1d-745d-4cb4-a029-99b135912f41",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0.0043, dtype=torch.float64)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# dQmin (Edwards & Shyphers, first order in amplitude and unperturbed tune differenct)\n",
    "\n",
    "(ks.squeeze()/(2.0*torch.pi)*(bx_a*by_a).sqrt()).abs()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "c583d015-1142-4ec3-9d35-fe1a69210709",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0.0042, dtype=torch.float64)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# dQmin (first order in amplitude)\n",
    "# Note, tunes in [0, 1/2] are assumed\n",
    "\n",
    "mux = 2.0*torch.pi*(nux_a % 0.5)\n",
    "muy = 2.0*torch.pi*(nuy_a % 0.5)\n",
    "\n",
    "(ks.squeeze()/(torch.pi)*(bx_a*by_a).sqrt()*(mux.sin()*muy.sin()).abs().sqrt()/(mux.sin() + muy.sin())).abs()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "71075621-a042-4e7d-8dee-338a55d3b504",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0., dtype=torch.float64)\n",
      "tensor(0.0042, dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "# dQmin (TEAPOT manual, appendix G, 1996)\n",
    "\n",
    "print(coupling(ring, [0.0*ks], ('ks', None, ['G'], None), matched=False))\n",
    "print(coupling(ring, [1.0*ks], ('ks', None, ['G'], None), matched=False))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "baa542b9-a481-41ae-b400-d8d1e9672c0a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([0.4239], dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "# Derivative\n",
    "# Note, not defined for zero skew amplitudes\n",
    "\n",
    "print(torch.func.jacrev(lambda ks: coupling(ring, [ks], ('ks', None, ['G'], None), matched=False))(ks))"
   ]
  }
 ],
 "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
}