{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "5004095b-629e-4a91-ba79-3fdfe9148dc1",
   "metadata": {},
   "source": [
    "# Example-25: Group"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "ebffed3a-4647-427a-8517-f66a6f33334b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# In this example another wrapper construction procedure is illustraded"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d57bc6b4-1c16-494d-931e-70213d9c2ec4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import\n",
    "\n",
    "import torch\n",
    "torch.set_printoptions(linewidth=128)\n",
    "\n",
    "from twiss import twiss\n",
    "from twiss import propagate\n",
    "from twiss import wolski_to_cs\n",
    "\n",
    "from model.library.drift import Drift\n",
    "from model.library.quadrupole import Quadrupole\n",
    "from model.library.dipole import Dipole\n",
    "from model.library.line import Line\n",
    "\n",
    "from model.command.wrapper import group"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "9e1047fe-285d-4114-90f1-c198ce61522b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define simple FODO based lattice using nested lines\n",
    "\n",
    "DR = Drift('DR', 0.75)\n",
    "BM = Dipole('BM', 3.50, torch.pi/4.0)\n",
    "\n",
    "QF_A = Quadrupole('QF_A', 0.5, +0.20)\n",
    "QD_A = Quadrupole('QD_A', 0.5, -0.19)\n",
    "QF_B = Quadrupole('QF_B', 0.5, +0.20)\n",
    "QD_B = Quadrupole('QD_B', 0.5, -0.19)\n",
    "QF_C = Quadrupole('QF_C', 0.5, +0.20)\n",
    "QD_C = Quadrupole('QD_C', 0.5, -0.19)\n",
    "QF_D = Quadrupole('QF_D', 0.5, +0.20)\n",
    "QD_D = Quadrupole('QD_D', 0.5, -0.19)\n",
    "\n",
    "FODO_A = Line('FODO_A', [QF_A, DR, BM, DR, QD_A, QD_A, DR, BM, DR, QF_A], propagate=True, dp=0.0, exact=False, output=False, matrix=False)\n",
    "FODO_B = Line('FODO_B', [QF_B, DR, BM, DR, QD_B, QD_B, DR, BM, DR, QF_B], propagate=True, dp=0.0, exact=False, output=False, matrix=False)\n",
    "FODO_C = Line('FODO_C', [QF_C, DR, BM, DR, QD_C, QD_C, DR, BM, DR, QF_C], propagate=True, dp=0.0, exact=False, output=False, matrix=False)\n",
    "FODO_D = Line('FODO_D', [QF_D, DR, BM, DR, QD_D, QD_D, DR, BM, DR, QF_D], propagate=True, dp=0.0, exact=False, output=False, matrix=False)\n",
    "\n",
    "RING = Line('RING', [FODO_A, FODO_B, FODO_C, FODO_D], propagate=True, dp=0.0, exact=False, output=False, matrix=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "d42cb6ea-8009-4d10-9b0b-e09b412939da",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[(None, ['QF_A', 'QD_A', 'QF_B', 'QD_B'], 'kn')]\n",
      "\n",
      "tensor([0., 0., 0., 0.], dtype=torch.float64)\n",
      "\n",
      "Quadrupole(name=\"QF_A\", length=0.5, kn=0.200000000000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)\n",
      "Drift(name=\"DR\", length=0.75, dp=0.0, exact=False, ns=1, order=0)\n",
      "Dipole(name=\"BM\", length=3.5, angle=0.7853981633974493, e1=0.0, e2=0.0, kn=1e-15, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)\n",
      "Drift(name=\"DR\", length=0.75, dp=0.0, exact=False, ns=1, order=0)\n",
      "Quadrupole(name=\"QD_A\", length=0.5, kn=-0.189999999999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)\n",
      "Quadrupole(name=\"QD_A\", length=0.5, kn=-0.189999999999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)\n",
      "Drift(name=\"DR\", length=0.75, dp=0.0, exact=False, ns=1, order=0)\n",
      "Dipole(name=\"BM\", length=3.5, angle=0.7853981633974493, e1=0.0, e2=0.0, kn=1e-15, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)\n",
      "Drift(name=\"DR\", length=0.75, dp=0.0, exact=False, ns=1, order=0)\n",
      "Quadrupole(name=\"QF_A\", length=0.5, kn=0.200000000000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)\n",
      "Quadrupole(name=\"QF_B\", length=0.5, kn=0.200000000000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)\n",
      "Drift(name=\"DR\", length=0.75, dp=0.0, exact=False, ns=1, order=0)\n",
      "Dipole(name=\"BM\", length=3.5, angle=0.7853981633974493, e1=0.0, e2=0.0, kn=1e-15, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)\n",
      "Drift(name=\"DR\", length=0.75, dp=0.0, exact=False, ns=1, order=0)\n",
      "Quadrupole(name=\"QD_B\", length=0.5, kn=-0.189999999999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)\n",
      "Quadrupole(name=\"QD_B\", length=0.5, kn=-0.189999999999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)\n",
      "Drift(name=\"DR\", length=0.75, dp=0.0, exact=False, ns=1, order=0)\n",
      "Dipole(name=\"BM\", length=3.5, angle=0.7853981633974493, e1=0.0, e2=0.0, kn=1e-15, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)\n",
      "Drift(name=\"DR\", length=0.75, dp=0.0, exact=False, ns=1, order=0)\n",
      "Quadrupole(name=\"QF_B\", length=0.5, kn=0.200000000000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Full ring or subline can be wrapped by element kind\n",
    "\n",
    "fn, table, line = group(RING,                               # -- source line\n",
    "                        'FODO_A',                           # -- start (name or position in source line sequence)\n",
    "                        'FODO_B',                           # -- end (name or position in source line sequence)\n",
    "                        ('kn', ['Quadrupole'], None, None)) # -- groups (key:str, kinds:list[str]|None, names:list[str]|None, clean:list[str]|None\n",
    "\n",
    "# Information about deviation variables is returbed in wrapper format\n",
    "\n",
    "print(table)\n",
    "print()\n",
    "\n",
    "# Wrapped function fn can be called with deviation variables\n",
    "\n",
    "(_, names, _), *_ = table\n",
    "knobs = torch.tensor(len(names)*[0.0], dtype=torch.float64)\n",
    "state = torch.tensor([0.0, 0.0, 0.0, 0.0], dtype=torch.float64)\n",
    "print(fn(state, knobs))\n",
    "print()\n",
    "\n",
    "# Constructed line also returned\n",
    "\n",
    "print(line)\n",
    "print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "fde9f15e-577f-4ad1-aca4-25561c46a768",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[(None, ['QF_A', 'QD_A', 'QF_B', 'QD_B'], 'kn')]\n",
      "\n",
      "[(None, ['QF_A', 'QD_A', 'QF_B', 'QD_B', 'QF_C', 'QD_C', 'QF_D', 'QD_D'], 'kn')]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# By default names are excracted from created subline\n",
    "\n",
    "_, table, _ = group(RING, 'FODO_A', 'FODO_B', ('kn', ['Quadrupole'], None, None))\n",
    "\n",
    "print(table)\n",
    "print()\n",
    "\n",
    "# Use root flag to extract name from the root line\n",
    "\n",
    "\n",
    "_, table, _ = group(RING, 'FODO_A', 'FODO_B', ('kn', ['Quadrupole'], None, None), root=True)\n",
    "\n",
    "print(table)\n",
    "print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "83878426-31c3-4680-93ba-bfe70e8e2a5f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([0., 0., 0., 0.], dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "# Set transport between observation points\n",
    "\n",
    "# 0--A--1--B--2--C--3--D--4\n",
    "\n",
    "line01, *_ =  group(RING, 'FODO_A', 'FODO_A', ('kn', ['Quadrupole'], None, None), root=True)\n",
    "line12, *_ =  group(RING, 'FODO_B', 'FODO_B', ('kn', ['Quadrupole'], None, None), root=True)\n",
    "line23, *_ =  group(RING, 'FODO_C', 'FODO_C', ('kn', ['Quadrupole'], None, None), root=True)\n",
    "line34, *_ =  group(RING, 'FODO_D', 'FODO_D', ('kn', ['Quadrupole'], None, None), root=True)\n",
    "\n",
    "lines = [\n",
    "    line01, \n",
    "    line12, \n",
    "    line23, \n",
    "    line34\n",
    "]\n",
    "\n",
    "def ring(state, knobs):\n",
    "    for line in lines:\n",
    "        state = line(state, knobs)\n",
    "    return state\n",
    "\n",
    "\n",
    "state = torch.tensor(4*[0.0], dtype=torch.float64)\n",
    "knobs = torch.tensor(8*[0.0], dtype=torch.float64)\n",
    "\n",
    "print(ring(state, knobs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d41f356c-2c90-47c4-aac3-ce28c60cfc46",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([0.6951, 0.7019], dtype=torch.float64)\n",
      "tensor([[ 1.4567,  0.1055,  1.4567,  0.1055,  1.4567,  0.1055,  1.4567,  0.1055],\n",
      "        [-0.5132, -1.6271, -0.5132, -1.6271, -0.5132, -1.6271, -0.5132, -1.6271]], dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "# Compute tunes and corresponding derivatives with respect to deviation parameters\n",
    "\n",
    "def fn(knobs):\n",
    "    m = torch.func.jacfwd(ring)(state, knobs)\n",
    "    t, *_ = twiss(m)\n",
    "    return t\n",
    "\n",
    "print(fn(knobs))\n",
    "print(torch.func.jacrev(fn)(knobs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "7743b600-a29e-4fce-aee4-0ce9b0e762a7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([18.6083, 18.6083, 18.6083, 18.6083], dtype=torch.float64)\n",
      "tensor([[ 21.1373,  -1.5714,  21.1373,  -1.5714, 140.4889, -10.4446, 140.4889, -10.4446],\n",
      "        [140.4889, -10.4446,  21.1373,  -1.5714,  21.1373,  -1.5714, 140.4889, -10.4446],\n",
      "        [140.4889, -10.4446, 140.4889, -10.4446,  21.1373,  -1.5714,  21.1373,  -1.5714],\n",
      "        [ 21.1373,  -1.5714, 140.4889, -10.4446, 140.4889, -10.4446,  21.1373,  -1.5714]], dtype=torch.float64)\n",
      "\n",
      "tensor([6.3291, 6.3291, 6.3291, 6.3291], dtype=torch.float64)\n",
      "tensor([[ 10.9592,  66.8187,  10.9592,  66.8187,  -5.0152, -30.5777,  -5.0152, -30.5777],\n",
      "        [ -5.0152, -30.5777,  10.9592,  66.8187,  10.9592,  66.8187,  -5.0152, -30.5777],\n",
      "        [ -5.0152, -30.5777,  -5.0152, -30.5777,  10.9592,  66.8187,  10.9592,  66.8187],\n",
      "        [ 10.9592,  66.8187,  -5.0152, -30.5777,  -5.0152, -30.5777,  10.9592,  66.8187]], dtype=torch.float64)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Compute beta functions at observation points and corresponding derivatives with respect to deviation parameters\n",
    "\n",
    "def fn(knobs):\n",
    "\n",
    "    bxs = []\n",
    "    bys = []\n",
    "\n",
    "    m = torch.func.jacfwd(ring)(state, knobs)\n",
    "\n",
    "    *_, w = twiss(m)\n",
    "    _, bx, _, by = wolski_to_cs(w)\n",
    "\n",
    "    for line in lines:\n",
    "        w = propagate(w, torch.func.jacrev(line)(state, knobs))\n",
    "        _, bx, _, by = wolski_to_cs(w)\n",
    "        bxs.append(bx)\n",
    "        bys.append(by)\n",
    "\n",
    "    bxs = torch.stack(bxs)\n",
    "    bys = torch.stack(bys)\n",
    "\n",
    "    return bxs, bys\n",
    "\n",
    "bx, by = fn(knobs)\n",
    "dbxdk, dbydk = torch.func.jacrev(fn)(knobs)\n",
    "\n",
    "print(bx)\n",
    "print(dbxdk)\n",
    "print()\n",
    "\n",
    "print(by)\n",
    "print(dbydk)\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
}