{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "29099120-c743-4ec3-81ed-568ea2b7cdb3",
   "metadata": {},
   "source": [
    "# Example-28: Transformation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "80fe3bed-c3d7-489a-ba86-732b7d26983b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# In this example another wrappers are used to construct parametric transformations between elements\n",
    "# Given two element a transformation can be constructed from the first element enterence frame to the second element exit frame\n",
    "# If the first element appears after the second one in the line, inverse transformation is constructed\n",
    "# Note, these transformations are given around initial reference orbit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "4177c76a-5132-47f7-b750-81301aaae5b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import\n",
    "\n",
    "import torch\n",
    "\n",
    "import matplotlib\n",
    "from matplotlib import pyplot as plt\n",
    "matplotlib.rcParams['text.usetex'] = True\n",
    "\n",
    "from twiss import twiss\n",
    "\n",
    "from ndmap.signature import chop\n",
    "from ndmap.evaluate import evaluate\n",
    "from ndmap.pfp import parametric_fixed_point\n",
    "\n",
    "from model.library.drift import Drift\n",
    "from model.library.quadrupole import Quadrupole\n",
    "from model.library.sextupole import Sextupole\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": "bca43646-5dc4-40ff-96cf-9ad1e2cad74c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define simple FODO based lattice using nested lines\n",
    "# Note, all elements have unique names\n",
    "\n",
    "QF_A = Quadrupole('QF_A', 1.0, +0.20)\n",
    "QD_A = Quadrupole('QD_A', 1.0, -0.19)\n",
    "QF_B = Quadrupole('QF_B', 1.0, +0.20)\n",
    "QD_B = Quadrupole('QD_B', 1.0, -0.19)\n",
    "QF_C = Quadrupole('QF_C', 1.0, +0.20)\n",
    "QD_C = Quadrupole('QD_C', 1.0, -0.19)\n",
    "QF_D = Quadrupole('QF_D', 1.0, +0.20)\n",
    "QD_D = Quadrupole('QD_D', 1.0, -0.19)\n",
    "\n",
    "SF1_A = Sextupole('SF1_A', 0.25, 0.00)\n",
    "SD1_A = Sextupole('SD1_A', 0.25, 0.00)\n",
    "SF2_A = Sextupole('SF2_A', 0.25, 0.00)\n",
    "SD2_A = Sextupole('SD2_A', 0.25, 0.00)\n",
    "SF1_B = Sextupole('SF1_B', 0.25, 0.00)\n",
    "SD1_B = Sextupole('SD1_B', 0.25, 0.00)\n",
    "SF2_B = Sextupole('SF2_B', 0.25, 0.00)\n",
    "SD2_B = Sextupole('SD2_B', 0.25, 0.00)\n",
    "SF1_C = Sextupole('SF1_C', 0.25, 0.00)\n",
    "SD1_C = Sextupole('SD1_C', 0.25, 0.00)\n",
    "SF2_C = Sextupole('SF2_C', 0.25, 0.00)\n",
    "SD2_C = Sextupole('SD2_C', 0.25, 0.00)\n",
    "SF1_D = Sextupole('SF1_D', 0.25, 0.00)\n",
    "SD1_D = Sextupole('SD1_D', 0.25, 0.00)\n",
    "SF2_D = Sextupole('SF2_D', 0.25, 0.00)\n",
    "SD2_D = Sextupole('SD2_D', 0.25, 0.00)\n",
    "\n",
    "BM1_A = Dipole('BM1_A', 3.50, torch.pi/4.0)\n",
    "BM2_A = Dipole('BM2_A', 3.50, torch.pi/4.0)\n",
    "BM1_B = Dipole('BM1_B', 3.50, torch.pi/4.0)\n",
    "BM2_B = Dipole('BM2_B', 3.50, torch.pi/4.0)\n",
    "BM1_C = Dipole('BM1_C', 3.50, torch.pi/4.0)\n",
    "BM2_C = Dipole('BM2_C', 3.50, torch.pi/4.0)\n",
    "BM1_D = Dipole('BM1_D', 3.50, torch.pi/4.0)\n",
    "BM2_D = Dipole('BM2_D', 3.50, torch.pi/4.0)\n",
    "\n",
    "DR1_A = Drift('DR1_A', 0.25)\n",
    "DR2_A = Drift('DR2_A', 0.25)\n",
    "DR3_A = Drift('DR3_A', 0.25)\n",
    "DR4_A = Drift('DR4_A', 0.25)\n",
    "DR5_A = Drift('DR5_A', 0.25)\n",
    "DR6_A = Drift('DR6_A', 0.25)\n",
    "DR7_A = Drift('DR7_A', 0.25)\n",
    "DR1_B = Drift('DR1_B', 0.25)\n",
    "DR2_B = Drift('DR2_B', 0.25)\n",
    "DR3_B = Drift('DR3_B', 0.25)\n",
    "DR4_B = Drift('DR4_B', 0.25)\n",
    "DR5_B = Drift('DR5_B', 0.25)\n",
    "DR6_B = Drift('DR6_B', 0.25)\n",
    "DR7_B = Drift('DR7_B', 0.25)\n",
    "DR1_C = Drift('DR1_C', 0.25)\n",
    "DR2_C = Drift('DR2_C', 0.25)\n",
    "DR3_C = Drift('DR3_C', 0.25)\n",
    "DR4_C = Drift('DR4_C', 0.25)\n",
    "DR5_C = Drift('DR5_C', 0.25)\n",
    "DR6_C = Drift('DR6_C', 0.25)\n",
    "DR7_C = Drift('DR7_C', 0.25)\n",
    "DR1_D = Drift('DR1_D', 0.25)\n",
    "DR2_D = Drift('DR2_D', 0.25)\n",
    "DR3_D = Drift('DR3_D', 0.25)\n",
    "DR4_D = Drift('DR4_D', 0.25)\n",
    "DR5_D = Drift('DR5_D', 0.25)\n",
    "DR6_D = Drift('DR6_D', 0.25)\n",
    "DR7_D = Drift('DR7_D', 0.25)\n",
    "\n",
    "FODO_A = Line('FODO_A', [QF_A, DR1_A, SF1_A, DR2_A, BM1_A, DR3_A, SD1_A, DR3_A, QD_A, DR4_A, SD2_A, DR5_A, BM2_A, DR6_A, SF2_A, DR7_A], propagate=True)\n",
    "FODO_B = Line('FODO_B', [QF_B, DR1_B, SF1_B, DR2_B, BM1_B, DR3_B, SD1_B, DR3_B, QD_B, DR4_B, SD2_B, DR5_B, BM2_B, DR6_B, SF2_B, DR7_B], propagate=True)\n",
    "FODO_C = Line('FODO_C', [QF_C, DR1_C, SF1_C, DR2_C, BM1_C, DR3_C, SD1_C, DR3_C, QD_C, DR4_C, SD2_C, DR5_C, BM2_C, DR6_C, SF2_C, DR7_C], propagate=True)\n",
    "FODO_D = Line('FODO_D', [QF_D, DR1_D, SF1_D, DR2_D, BM1_D, DR3_D, SD1_D, DR3_D, QD_D, DR4_D, SD2_D, DR5_D, BM2_D, DR6_D, SF2_D, DR7_D], propagate=True)\n",
    "\n",
    "RING = Line('RING', [FODO_A, FODO_B, FODO_C, FODO_D], propagate=True, dp=0.0, exact=False, output=False, matrix=False)\n",
    "\n",
    "RING.flatten()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f11efb1f-2471-4362-983d-a7e8217aa51b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create parametric transformation from one element to another and its inverse\n",
    "\n",
    "probe = 'SD2_A'\n",
    "other = 'SF1_D'\n",
    "\n",
    "forward, *_ = group(RING, probe, other, ('kn', ['Quadrupole'], None, None), ('ms', ['Sextupole'], None, None), ('dp', None, None, None), root=True, alignment=False)\n",
    "inverse, *_ = group(RING, other, probe, ('kn', ['Quadrupole'], None, None), ('ms', ['Sextupole'], None, None), ('dp', None, None, None), root=True, alignment=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "5301c64b-603b-4723-b8a0-b9e60cc95118",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([ 0.0010,  0.0050, -0.0050,  0.0010], dtype=torch.float64)\n",
      "tensor([-0.0041,  0.0022,  0.0047, -0.0001], dtype=torch.float64)\n",
      "tensor([ 0.0010,  0.0050, -0.0050,  0.0010], dtype=torch.float64)\n",
      "tensor([ 4.3368e-19, -8.6736e-19, -4.3368e-18,  2.1684e-19],\n",
      "       dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "# Test propagation and inverse transformation\n",
    "\n",
    "state = torch.tensor([0.001, 0.005, -0.005, 0.001], dtype=torch.float64)\n",
    "\n",
    "kn = 1.0E-3*torch.randn( 8, dtype=torch.float64)\n",
    "ms = 1.0E-3*torch.randn(16, dtype=torch.float64)\n",
    "dp = torch.tensor([0.001], dtype=torch.float64)\n",
    "\n",
    "print(local := state.clone())\n",
    "print(local := forward(local, kn, ms, dp))\n",
    "print(local := inverse(local, kn, ms, dp))\n",
    "print(state - local)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "e412d7ed-afbf-4d64-9864-7f5d08b23a53",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[ 0.3917,  0.5359,  0.0000,  0.0000],\n",
      "        [ 0.4481,  3.1656,  0.0000,  0.0000],\n",
      "        [-0.0000, -0.0000, -1.1626, -3.2206],\n",
      "        [ 0.0000,  0.0000,  0.2044, -0.2939]], dtype=torch.float64)\n",
      "tensor([[ 0.3917,  0.5359,  0.0000,  0.0000],\n",
      "        [ 0.4481,  3.1656,  0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000, -1.1626, -3.2206],\n",
      "        [ 0.0000,  0.0000,  0.2044, -0.2939]], dtype=torch.float64)\n",
      "\n",
      "tensor([ 0.0000,  0.0000,  2.6498, 39.4992, 30.9318, 39.5209,  2.0561,  0.0000],\n",
      "       dtype=torch.float64)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Test derivatives\n",
    "\n",
    "state = torch.tensor([0.0, 0.0, 0.0, 0.0], dtype=torch.float64)\n",
    "\n",
    "kn = torch.zeros( 8, dtype=torch.float64)\n",
    "ms = torch.zeros(16, dtype=torch.float64)\n",
    "dp = torch.tensor([0.0], dtype=torch.float64)\n",
    "\n",
    "# Transport matrix\n",
    "\n",
    "print(torch.func.jacrev(forward)(state, kn, ms, dp).inverse())\n",
    "print(torch.func.jacrev(inverse)(state, kn, ms, dp))\n",
    "print()\n",
    "\n",
    "# Derivatives of transport matrix trace with respect to quadrupole deviations\n",
    "\n",
    "def matrix(kn, ms, dp):\n",
    "    return torch.func.jacrev(forward)(state, kn, ms, dp).trace()\n",
    "\n",
    "print(torch.func.jacrev(matrix)(kn, ms, dp))\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
}