{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d8c1bf6a-ba82-4a37-8476-e4eef0606e76",
   "metadata": {},
   "source": [
    "## ELETTRA-37: ID linear optics distortion (one-turn matrix factorization)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c9c1e661-3dc9-4b81-a8b3-b2aaa4fb9810",
   "metadata": {},
   "outputs": [],
   "source": [
    "# In this example one-turn matrix factorization is illustrated\n",
    "# The one-turn matrix can be expressed as M exp(S A) where Aij = Aji is (a block matrix if there is no coupling that) describes cumulative effect of perturbation\n",
    "# Here the matrix A is constructed perturbatively using Magnus expansion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d1df2aaf-f327-4fd1-a0cd-1c16e696e481",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import\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",
    "from matplotlib.patches import Rectangle\n",
    "matplotlib.rcParams['text.usetex'] = True\n",
    "\n",
    "from model.library.element import Element\n",
    "from model.library.line import Line\n",
    "from model.library.quadrupole import Quadrupole\n",
    "from model.library.marker import Marker\n",
    "from model.library.matrix import Matrix\n",
    "\n",
    "from model.command.external import load_lattice\n",
    "from model.command.build import build"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "f3eb59a5-7a96-4487-beb3-8301eaa1fcdf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set data type and device\n",
    "\n",
    "Element.dtype = dtype = torch.float64\n",
    "Element.device = device = torch.device('cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "582f559c-b049-4880-a605-75b46fa7d3c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load lattice (ELEGANT table)\n",
    "# Note, lattice is allowed to have repeated elements\n",
    "\n",
    "path = Path('elettra.lte')\n",
    "data = load_lattice(path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "3a7bc32a-810a-4dc0-b85f-55aa64616d16",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'BPM': 168, 'Drift': 708, 'Dipole': 156, 'Quadrupole': 360, 'Marker': 12}"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Build and setup lattice\n",
    "\n",
    "ring:Line = build('RING', 'ELEGANT', data)\n",
    "\n",
    "# Flatten sublines\n",
    "\n",
    "ring.flatten()\n",
    "\n",
    "# Remove all marker elements but the ones starting with MLL (long straight section centers)\n",
    "\n",
    "ring.remove_group(pattern=r'^(?!MLL_).*', kinds=['Marker'])\n",
    "\n",
    "# Replace all sextupoles with quadrupoles\n",
    "\n",
    "def factory(element:Element) -> None:\n",
    "    table = element.serialize\n",
    "    table.pop('ms', None)\n",
    "    return Quadrupole(**table)\n",
    "\n",
    "ring.replace_group(pattern=r'', factory=factory, kinds=['Sextupole'])\n",
    "\n",
    "# Set linear dipoles\n",
    "\n",
    "def apply(element:Element) -> None:\n",
    "    element.linear = True\n",
    "\n",
    "ring.apply(apply, kinds=['Dipole'])\n",
    "\n",
    "# Merge drifts\n",
    "\n",
    "ring.merge()\n",
    "\n",
    "# Change lattice start\n",
    "\n",
    "ring.start = \"BPM_S01_01\"\n",
    "\n",
    "# Describe\n",
    "\n",
    "ring.describe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "1c38b03d-435c-4df8-a66f-604fd024b375",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[-0.5038,  3.9513,  0.0000,  0.0000],\n",
      "        [-0.2394, -0.1073,  0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000,  0.1784,  1.9752],\n",
      "        [ 0.0000,  0.0000, -0.4264,  0.8845]], dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "# Unperturbed one-turn matrix\n",
    "\n",
    "state = torch.tensor(4*[0.0], dtype=dtype)\n",
    "M0 = torch.func.jacrev(ring)(state)\n",
    "print(M0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "7725f8df-e92c-4fa7-a264-61c84e87a53c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define IDs\n",
    "\n",
    "# The first ID is 'fake' and acts as an identity\n",
    "# It is also located at the lattice start\n",
    "\n",
    "ca, cb, cc, cd = -0.034441907232402175, -0.04458009513208418, 0.056279356423643276, 0.08037110220505986\n",
    "A = torch.tensor([[ca, 0.0, 0.0, 0.0], [0.0, cb, 0.0, 0.0], [0.0, 0.0, cc, 0.0], [0.0, 0.0, 0.0, cd]], dtype=dtype)\n",
    "mask = torch.triu(torch.ones_like(A, dtype=torch.bool))\n",
    "\n",
    "ID1 = Matrix('ID1', length=0.0, A=(0.0*A[mask]).tolist())\n",
    "ID2 = Matrix('ID2', length=0.0, A=(1.0*A[mask]).tolist())\n",
    "ID3 = Matrix('ID3', length=0.0, A=(1.1*A[mask]).tolist())\n",
    "ID4 = Matrix('ID4', length=0.0, A=(0.8*A[mask]).tolist())\n",
    "ID5 = Matrix('ID5', length=0.0, A=(0.5*A[mask]).tolist())\n",
    "ID6 = Matrix('ID6', length=0.0, A=(1.2*A[mask]).tolist())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "d095b2c2-3795-4ae6-9e53-82e3e5ec3b51",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'BPM': 168,\n",
       " 'Matrix': 6,\n",
       " 'Drift': 708,\n",
       " 'Dipole': 156,\n",
       " 'Quadrupole': 360,\n",
       " 'Marker': 12}"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Insert IDs\n",
    "\n",
    "# Each ID (or other perturbation) is a thin insertion matrix\n",
    "# The matrices are inserted after given markers\n",
    "\n",
    "elements = [ID1, ID2, ID3, ID4, ID5, ID6]\n",
    "markers = ['BPM_S01_01', 'MLL_S01', 'MLL_S02', 'MLL_S03', 'MLL_S05', 'MLL_S08']\n",
    "\n",
    "error = ring.clone()\n",
    "for element, marker in zip(elements, markers):\n",
    "    error.insert(element, marker)\n",
    "    \n",
    "# Describe\n",
    "\n",
    "error.describe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "c0f96ba5-d2b7-482c-b66e-eec772682a5b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[-0.5038,  3.9513,  0.0000,  0.0000],\n",
      "        [-0.2394, -0.1073,  0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000,  0.1784,  1.9752],\n",
      "        [ 0.0000,  0.0000, -0.4264,  0.8845]], dtype=torch.float64)\n",
      "tensor([[ 0.4728,  4.4862,  0.0000,  0.0000],\n",
      "        [-0.1703,  0.4987,  0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000, -0.1576,  2.3014],\n",
      "        [ 0.0000,  0.0000, -0.4772,  0.6233]], dtype=torch.float64)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Perturbed one-turn matrix\n",
    "\n",
    "state = torch.tensor(4*[0.0], dtype=dtype)\n",
    "M1 = torch.func.jacrev(error)(state)\n",
    "\n",
    "print(M0)\n",
    "print(M1)\n",
    "print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4eaa217f-8bc1-49cc-8715-533cc1fee3b3",
   "metadata": {},
   "source": [
    "$\\hat M = \\prod_i T_{i, i + 1} K_i $"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "b5fb7dbc-9361-4ace-b1de-917395db3e7c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[True, True, True, True, True, True]\n",
      "\n",
      "tensor([[-0.5038,  3.9513,  0.0000,  0.0000],\n",
      "        [-0.2394, -0.1073,  0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000,  0.1784,  1.9752],\n",
      "        [ 0.0000,  0.0000, -0.4264,  0.8845]], dtype=torch.float64)\n",
      "tensor([[-0.5038,  3.9513,  0.0000,  0.0000],\n",
      "        [-0.2394, -0.1073,  0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000,  0.1784,  1.9752],\n",
      "        [ 0.0000,  0.0000, -0.4264,  0.8845]], dtype=torch.float64)\n",
      "True\n",
      "\n",
      "tensor([[ 0.4728,  4.4862,  0.0000,  0.0000],\n",
      "        [-0.1703,  0.4987,  0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000, -0.1576,  2.3014],\n",
      "        [ 0.0000,  0.0000, -0.4772,  0.6233]], dtype=torch.float64)\n",
      "tensor([[ 0.4728,  4.4862,  0.0000,  0.0000],\n",
      "        [-0.1703,  0.4987,  0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000, -0.1576,  2.3014],\n",
      "        [ 0.0000,  0.0000, -0.4772,  0.6233]], dtype=torch.float64)\n",
      "True\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Skew identity matrix\n",
    "\n",
    "S = torch.tensor([[0, 1, 0, 0], [-1, 0, 0, 0], [0, 0, 0, 1], [0, 0, -1, 0]], dtype=dtype)\n",
    "\n",
    "# Compute perturbations\n",
    "\n",
    "Ais = []\n",
    "Kis = []\n",
    "for element in elements:\n",
    "    Ai = torch.zeros((4, 4), dtype=dtype)\n",
    "    Ai[mask] = element.A\n",
    "    Ki = torch.linalg.matrix_exp(S @ Ai)\n",
    "    Ais.append(Ai)\n",
    "    Kis.append(Ki)\n",
    "    \n",
    "Ais = torch.stack(Ais)\n",
    "Kis = torch.stack(Kis)\n",
    "\n",
    "print([torch.allclose(Ki, torch.func.jacrev(element)(state)) for Ki, element in zip(Kis, elements)])\n",
    "print()\n",
    "\n",
    "# Compute matrices between IDs (using markers in the unperturbed lattice)\n",
    "\n",
    "Tis = []\n",
    "for i in range(len(markers)):\n",
    "    _, *sequence, _ = ring[markers[i] : markers[(i + 1) % len(markers)]]\n",
    "    line = Line('', sequence=sequence)\n",
    "    Tis.append(torch.func.jacrev(line)(state))\n",
    "Tis = torch.stack(Tis)\n",
    "\n",
    "# Construct one-turn matrix\n",
    "\n",
    "hM = torch.eye(4, dtype=dtype)\n",
    "for Ti in Tis:\n",
    "    hM = Ti @ hM\n",
    "\n",
    "print(M0)\n",
    "print(hM)\n",
    "print(torch.allclose(M0, hM))\n",
    "print()\n",
    "\n",
    "# Construct one-turn matrix\n",
    "\n",
    "hM = torch.eye(4, dtype=dtype)\n",
    "for Ti, Ki in zip(Tis, Kis):\n",
    "    hM = Ti @ Ki @ hM\n",
    "\n",
    "print(M1)\n",
    "print(hM)\n",
    "print(torch.allclose(M1, hM))\n",
    "print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96dad8fe-2901-4a38-a0b2-4ceb2a47c39e",
   "metadata": {},
   "source": [
    "$\\hat M = \\prod_i T_{i, i + 1} K_i  = M \\prod_i \\hat K_i$\n",
    "\n",
    "$\\hat K_i = T_i^{-1} K_i T_i $"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "15f11abf-6619-4ec5-a29d-8db1adf0327d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[ 0.4728,  4.4862,  0.0000,  0.0000],\n",
      "        [-0.1703,  0.4987,  0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000, -0.1576,  2.3014],\n",
      "        [ 0.0000,  0.0000, -0.4772,  0.6233]], dtype=torch.float64)\n",
      "tensor([[ 0.4728,  4.4862,  0.0000,  0.0000],\n",
      "        [-0.1703,  0.4987,  0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000, -0.1576,  2.3014],\n",
      "        [ 0.0000,  0.0000, -0.4772,  0.6233]], dtype=torch.float64)\n",
      "True\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Compute unperturbed transport matrices from lattice start to each ID\n",
    "\n",
    "Tis = []\n",
    "for marker in markers:\n",
    "    *sequence, _ = ring[ring.start : marker]\n",
    "    line = Line('', sequence=sequence)\n",
    "    Tis.append(torch.func.jacrev(line)(state))\n",
    "Tis = torch.stack(Tis)\n",
    "\n",
    "# Construct one-turn matrix\n",
    "\n",
    "hM = torch.eye(4, dtype=dtype)\n",
    "for Ti, Ki in zip(Tis, Kis):\n",
    "    hKi = Ti.inverse() @ Ki @ Ti\n",
    "    hM =  hKi @ hM\n",
    "hM = M0 @ hM\n",
    "\n",
    "print(M1)\n",
    "print(hM)\n",
    "print(torch.allclose(M1, hM))\n",
    "print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd0efba6-cc59-4b88-bb4c-34d4b456af62",
   "metadata": {},
   "source": [
    "$\\hat M = \\prod_i T_{i, i + 1} K_i  = M \\prod_i \\hat K_i$\n",
    "\n",
    "$\\hat K_i = T_i^{-1} K_i T_i = \\exp ( T_i^{-1} S A_i T_i ) = \\exp ( S \\hat A_i )$\n",
    "\n",
    "$\\hat A_i = T_i^{T} A_i T_i$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "2b1daf10-a68e-4b1b-bb54-a2cc524c5037",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[ 0.4728,  4.4862,  0.0000,  0.0000],\n",
      "        [-0.1703,  0.4987,  0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000, -0.1576,  2.3014],\n",
      "        [ 0.0000,  0.0000, -0.4772,  0.6233]], dtype=torch.float64)\n",
      "tensor([[ 0.4728,  4.4862,  0.0000,  0.0000],\n",
      "        [-0.1703,  0.4987,  0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000, -0.1576,  2.3014],\n",
      "        [ 0.0000,  0.0000, -0.4772,  0.6233]], dtype=torch.float64)\n",
      "True\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Construct shifted exponents\n",
    "\n",
    "hAis = []\n",
    "for Ti, Ai in zip(Tis, Ais):\n",
    "    hAi = Ti.T @ Ai @ Ti\n",
    "    hAis.append(hAi)\n",
    "hAis = torch.stack(hAis)\n",
    "\n",
    "# Construct one-turn matrix\n",
    "\n",
    "hM = torch.eye(4, dtype=dtype)\n",
    "for hAi in hAis:\n",
    "    hM =  torch.linalg.matrix_exp(S @ hAi) @ hM\n",
    "hM = M0 @ hM\n",
    "\n",
    "print(M1)\n",
    "print(hM)\n",
    "print(torch.allclose(M1, hM))\n",
    "print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "261fcce1-1af0-4ebc-b01d-68d545f1dd6c",
   "metadata": {},
   "source": [
    "$\\hat M = M \\prod_i \\exp(S \\hat A_i) = M \\exp(\\varepsilon \\hat A^{(1)} + \\varepsilon^2 \\hat A^{(2)} + \\dots)$\n",
    "\n",
    "$\\hat A^{(1)} = \\sum_i \\hat A_i$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "67e295af-a86e-459e-be52-91cf1f9957de",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[ 0.4728,  4.4862,  0.0000,  0.0000],\n",
      "        [-0.1703,  0.4987,  0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000, -0.1576,  2.3014],\n",
      "        [ 0.0000,  0.0000, -0.4772,  0.6233]], dtype=torch.float64)\n",
      "tensor([[-0.5038,  3.9513,  0.0000,  0.0000],\n",
      "        [-0.2394, -0.1073,  0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000,  0.1784,  1.9752],\n",
      "        [ 0.0000,  0.0000, -0.4264,  0.8845]], dtype=torch.float64)\n",
      "tensor(1.3791, dtype=torch.float64)\n",
      "\n",
      "tensor([[ 0.4728,  4.4862,  0.0000,  0.0000],\n",
      "        [-0.1703,  0.4987,  0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000, -0.1576,  2.3014],\n",
      "        [ 0.0000,  0.0000, -0.4772,  0.6233]], dtype=torch.float64)\n",
      "tensor([[ 0.2006,  4.7250,  0.0000,  0.0000],\n",
      "        [-0.1836,  0.6597,  0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000, -0.1595,  2.2923],\n",
      "        [ 0.0000,  0.0000, -0.4798,  0.6261]], dtype=torch.float64)\n",
      "tensor(0.3967, dtype=torch.float64)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Construct 1st order product approximation\n",
    "\n",
    "hA1 = torch.zeros((4, 4), dtype=dtype)\n",
    "for hAi in hAis:\n",
    "    hA1 += hAi\n",
    "\n",
    "# Construct one-turn matrix with perturbation\n",
    "\n",
    "hM = M0 @ torch.linalg.matrix_exp(S @ hA1)\n",
    "\n",
    "print(M1)\n",
    "print(M0)\n",
    "print((M1 - M0).norm())\n",
    "print()\n",
    "\n",
    "print(M1)\n",
    "print(hM)\n",
    "print((M1 - hM).norm())\n",
    "print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d1cf9b1-c7e7-4c05-9dbc-e156258951b4",
   "metadata": {},
   "source": [
    "$\\hat M = M \\prod_i \\exp(S \\hat A_i) = M \\exp(\\varepsilon \\hat A^{(1)} + \\varepsilon^2 \\hat A^{(2)} + \\dots)$\n",
    "\n",
    "$S \\hat A^{(1)} = S \\sum_i \\hat A_i$\n",
    "\n",
    "$S \\hat A^{(2)} = \\sum_{i \\ge j} \\frac{1}{1 + \\delta_{i, j}} \\frac{1}{2} \\{ S \\hat A_i, S \\hat A_j \\} = S \\sum_{i \\ge j} \\frac{1}{1 + \\delta_{i, j}} \\frac{1}{2} (\\hat A_i S \\hat A_j - \\hat A_j S \\hat A_i)$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "07901400-a4f9-4717-8bf6-283c020bb771",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define bracket (commutator)\n",
    "\n",
    "def bracket(X, Y):\n",
    "    return X @ Y - Y @ X"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "3f94820b-0e53-46eb-87e9-d29b6fac82c3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[ 0.4728,  4.4862,  0.0000,  0.0000],\n",
      "        [-0.1703,  0.4987,  0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000, -0.1576,  2.3014],\n",
      "        [ 0.0000,  0.0000, -0.4772,  0.6233]], dtype=torch.float64)\n",
      "tensor([[-0.5038,  3.9513,  0.0000,  0.0000],\n",
      "        [-0.2394, -0.1073,  0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000,  0.1784,  1.9752],\n",
      "        [ 0.0000,  0.0000, -0.4264,  0.8845]], dtype=torch.float64)\n",
      "tensor(1.3791, dtype=torch.float64)\n",
      "\n",
      "tensor([[ 0.4728,  4.4862,  0.0000,  0.0000],\n",
      "        [-0.1703,  0.4987,  0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000, -0.1576,  2.3014],\n",
      "        [ 0.0000,  0.0000, -0.4772,  0.6233]], dtype=torch.float64)\n",
      "tensor([[ 0.2006,  4.7250,  0.0000,  0.0000],\n",
      "        [-0.1836,  0.6597,  0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000, -0.1595,  2.2923],\n",
      "        [ 0.0000,  0.0000, -0.4798,  0.6261]], dtype=torch.float64)\n",
      "tensor(0.3967, dtype=torch.float64)\n",
      "\n",
      "tensor([[ 0.4728,  4.4862,  0.0000,  0.0000],\n",
      "        [-0.1703,  0.4987,  0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000, -0.1576,  2.3014],\n",
      "        [ 0.0000,  0.0000, -0.4772,  0.6233]], dtype=torch.float64)\n",
      "tensor([[ 0.3914,  4.5444,  0.0000,  0.0000],\n",
      "        [-0.1772,  0.4978,  0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000, -0.1571,  2.3014],\n",
      "        [ 0.0000,  0.0000, -0.4770,  0.6229]], dtype=torch.float64)\n",
      "tensor(0.1004, dtype=torch.float64)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Construct 2nd order product approximation\n",
    "\n",
    "hA2 = torch.zeros((4, 4), dtype=dtype)\n",
    "\n",
    "for i in range(len(elements)):\n",
    "    for j in range(i + 1):\n",
    "        dij = (i == j)\n",
    "        factor = 1/(1 + dij)\n",
    "        Xi = S @ hAis[i]\n",
    "        Xj = S @ hAis[j]\n",
    "        hA2 += factor*1/2*bracket(Xi, Xj)\n",
    "        \n",
    "hA2 = - S @ hA2\n",
    "\n",
    "# Construct one-turn matrix\n",
    "\n",
    "print(M1)\n",
    "print(M0)\n",
    "print((M1 - M0).norm())\n",
    "print()\n",
    "\n",
    "hM = M0 @ torch.linalg.matrix_exp(S @ hA1)\n",
    "\n",
    "print(M1)\n",
    "print(hM)\n",
    "print((M1 - hM).norm())\n",
    "print()\n",
    "\n",
    "hM = M0 @ torch.linalg.matrix_exp(S @ (hA1 + hA2))\n",
    "\n",
    "print(M1)\n",
    "print(hM)\n",
    "print((M1 - hM).norm())\n",
    "print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49d85bc1-278d-4e51-8ad9-1b56b0a30eef",
   "metadata": {},
   "source": [
    "$\\hat M = M \\prod_i \\exp(S \\hat A_i) = M \\exp(\\varepsilon \\hat A^{(1)} + \\varepsilon^2 \\hat A^{(2)} + \\dots)$\n",
    "\n",
    "$S \\hat A^{(1)} = S \\sum_i \\hat A_i$\n",
    "\n",
    "$S \\hat A^{(2)} = \\sum_{i \\ge j} \\frac{1}{1 + \\delta_{i, j}} \\frac{1}{2} \\{ S \\hat A_i, S \\hat A_j \\} = S \\sum_{i \\ge j} \\frac{1}{1 + \\delta_{i, j}} \\frac{1}{2} (\\hat A_i S \\hat A_j - \\hat A_j S \\hat A_i)$\n",
    "\n",
    "$S\\,\\hat A^{(3)}=\\sum_{i \\ge j \\ge k} \\frac{1}{1+\\delta_{ij}} \\frac{1}{1+\\delta_{ik}+\\delta_{jk}} (\\frac{1}{4} \\{\\{S\\hat A_i,S\\hat A_j\\},S\\hat A_k\\}-\\frac{1}{12}\\{\\{S\\hat A_i,S\\hat A_k\\},S\\hat A_j\\}-\\frac{1}{12}\\{\\{S\\hat A_j,S\\hat A_k\\},S\\hat A_i\\})$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "d14ba5fe-54ca-4578-8cf5-208acae62fc0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[ 0.4728,  4.4862,  0.0000,  0.0000],\n",
      "        [-0.1703,  0.4987,  0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000, -0.1576,  2.3014],\n",
      "        [ 0.0000,  0.0000, -0.4772,  0.6233]], dtype=torch.float64)\n",
      "tensor([[-0.5038,  3.9513,  0.0000,  0.0000],\n",
      "        [-0.2394, -0.1073,  0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000,  0.1784,  1.9752],\n",
      "        [ 0.0000,  0.0000, -0.4264,  0.8845]], dtype=torch.float64)\n",
      "tensor(1.3791, dtype=torch.float64)\n",
      "\n",
      "tensor([[ 0.4728,  4.4862,  0.0000,  0.0000],\n",
      "        [-0.1703,  0.4987,  0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000, -0.1576,  2.3014],\n",
      "        [ 0.0000,  0.0000, -0.4772,  0.6233]], dtype=torch.float64)\n",
      "tensor([[ 0.2006,  4.7250,  0.0000,  0.0000],\n",
      "        [-0.1836,  0.6597,  0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000, -0.1595,  2.2923],\n",
      "        [ 0.0000,  0.0000, -0.4798,  0.6261]], dtype=torch.float64)\n",
      "tensor(0.3967, dtype=torch.float64)\n",
      "\n",
      "tensor([[ 0.4728,  4.4862,  0.0000,  0.0000],\n",
      "        [-0.1703,  0.4987,  0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000, -0.1576,  2.3014],\n",
      "        [ 0.0000,  0.0000, -0.4772,  0.6233]], dtype=torch.float64)\n",
      "tensor([[ 0.3914,  4.5444,  0.0000,  0.0000],\n",
      "        [-0.1772,  0.4978,  0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000, -0.1571,  2.3014],\n",
      "        [ 0.0000,  0.0000, -0.4770,  0.6229]], dtype=torch.float64)\n",
      "tensor(0.1004, dtype=torch.float64)\n",
      "\n",
      "tensor([[ 0.4728,  4.4862,  0.0000,  0.0000],\n",
      "        [-0.1703,  0.4987,  0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000, -0.1576,  2.3014],\n",
      "        [ 0.0000,  0.0000, -0.4772,  0.6233]], dtype=torch.float64)\n",
      "tensor([[ 0.4452,  4.5011,  0.0000,  0.0000],\n",
      "        [-0.1722,  0.5048,  0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000, -0.1575,  2.3013],\n",
      "        [ 0.0000,  0.0000, -0.4772,  0.6232]], dtype=torch.float64)\n",
      "tensor(0.0320, dtype=torch.float64)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Construct 3rd order product approximation\n",
    "\n",
    "hA3 = torch.zeros((4, 4), dtype=dtype)\n",
    "\n",
    "for i in range(len(elements)):\n",
    "    for j in range(i + 1):\n",
    "        for k in range(j + 1):\n",
    "            dij = (i == j)\n",
    "            dik = (i == k)\n",
    "            djk = (j == k)\n",
    "            factor = (1.0/(1.0 + dij))*(1.0/(1.0 + dik + djk))\n",
    "            Xi = S @ hAis[i]\n",
    "            Xj = S @ hAis[j]\n",
    "            Xk = S @ hAis[k]\n",
    "            hA3 += factor*(\n",
    "                + 1/4*bracket(bracket(Xi, Xj), Xk) \n",
    "                - (1/12)*bracket(bracket(Xi, Xk), Xj) \n",
    "                - (1/12)*bracket(bracket(Xj, Xk), Xi)\n",
    "            )\n",
    "            \n",
    "hA3 = - S @ hA3\n",
    "\n",
    "# Construct one-turn matrix\n",
    "\n",
    "print(M1)\n",
    "print(M0)\n",
    "print((M1 - M0).norm())\n",
    "print()\n",
    "\n",
    "hM = M0 @ torch.linalg.matrix_exp(S @ hA1)\n",
    "\n",
    "print(M1)\n",
    "print(hM)\n",
    "print((M1 - hM).norm())\n",
    "print()\n",
    "\n",
    "hM = M0 @ torch.linalg.matrix_exp(S @ (hA1 + hA2))\n",
    "\n",
    "print(M1)\n",
    "print(hM)\n",
    "print((M1 - hM).norm())\n",
    "print()\n",
    "\n",
    "hM = M0 @ torch.linalg.matrix_exp(S @ (hA1 + hA2 + hA3))\n",
    "\n",
    "print(M1)\n",
    "print(hM)\n",
    "print((M1 - hM).norm())\n",
    "print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07faffbf-1ad0-4b50-b2c9-e943f12f4776",
   "metadata": {},
   "source": [
    "$\\hat M = M \\prod_i \\exp(S \\hat A_i) = M \\exp(\\varepsilon \\hat A^{(1)} + \\varepsilon^2 \\hat A^{(2)} + \\dots)$\n",
    "\n",
    "$S \\hat A^{(1)} = S \\sum_i \\hat A_i$\n",
    "\n",
    "$S \\hat A^{(2)} = \\sum_{i \\ge j} \\frac{1}{1 + \\delta_{i, j}} \\frac{1}{2} \\{ S \\hat A_i, S \\hat A_j \\} = S \\sum_{i \\ge j} \\frac{1}{1 + \\delta_{i, j}} \\frac{1}{2} (\\hat A_i S \\hat A_j - \\hat A_j S \\hat A_i)$\n",
    "\n",
    "$S \\hat A^{(3)}=\\sum_{i \\ge j \\ge k} \\frac{1}{1+\\delta_{ij}} \\frac{1}{1+\\delta_{ik}+\\delta_{jk}} (\\frac{1}{4} \\{\\{S\\hat A_i,S\\hat A_j\\},S\\hat A_k\\}-\\frac{1}{12}\\{\\{S\\hat A_i,S\\hat A_k\\},S\\hat A_j\\}-\\frac{1}{12}\\{\\{S\\hat A_j,S\\hat A_k\\},S\\hat A_i\\})$\n",
    "\n",
    "$\n",
    "S\\hat A^{(4)} =\n",
    "\\sum_{i \\ge j \\ge k \\ge l}\n",
    "\\frac{1}{1+\\delta_{ij}}\\;\n",
    "\\frac{1}{1+\\delta_{ik}+\\delta_{jk}}\\;\n",
    "\\frac{1}{1+\\delta_{il}+\\delta_{jl}+\\delta_{kl}}\n",
    "\\frac{1}{12} (\n",
    "\\{\\{\\{S \\hat A_i,S \\hat A_j\\},S \\hat A_k\\},S \\hat A_l\\}\n",
    "+\\{S \\hat A_i,\\{\\{S \\hat A_j,S \\hat A_k\\},S \\hat A_l\\}\\}\n",
    "+\\{S \\hat A_i,\\{S \\hat A_j,\\{S \\hat A_k,S \\hat A_l\\}\\}\\}\n",
    "+\\{S \\hat A_j,\\{S \\hat A_k,\\{S \\hat A_l,S \\hat A_i\\}\\}\\}\n",
    ")\n",
    "$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "6256ff4c-39fd-4e5a-a948-2b1167da4171",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[ 0.4728,  4.4862,  0.0000,  0.0000],\n",
      "        [-0.1703,  0.4987,  0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000, -0.1576,  2.3014],\n",
      "        [ 0.0000,  0.0000, -0.4772,  0.6233]], dtype=torch.float64)\n",
      "tensor([[-0.5038,  3.9513,  0.0000,  0.0000],\n",
      "        [-0.2394, -0.1073,  0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000,  0.1784,  1.9752],\n",
      "        [ 0.0000,  0.0000, -0.4264,  0.8845]], dtype=torch.float64)\n",
      "tensor(1.3791, dtype=torch.float64)\n",
      "\n",
      "tensor([[ 0.4728,  4.4862,  0.0000,  0.0000],\n",
      "        [-0.1703,  0.4987,  0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000, -0.1576,  2.3014],\n",
      "        [ 0.0000,  0.0000, -0.4772,  0.6233]], dtype=torch.float64)\n",
      "tensor([[ 0.2006,  4.7250,  0.0000,  0.0000],\n",
      "        [-0.1836,  0.6597,  0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000, -0.1595,  2.2923],\n",
      "        [ 0.0000,  0.0000, -0.4798,  0.6261]], dtype=torch.float64)\n",
      "tensor(0.3967, dtype=torch.float64)\n",
      "\n",
      "tensor([[ 0.4728,  4.4862,  0.0000,  0.0000],\n",
      "        [-0.1703,  0.4987,  0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000, -0.1576,  2.3014],\n",
      "        [ 0.0000,  0.0000, -0.4772,  0.6233]], dtype=torch.float64)\n",
      "tensor([[ 0.3914,  4.5444,  0.0000,  0.0000],\n",
      "        [-0.1772,  0.4978,  0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000, -0.1571,  2.3014],\n",
      "        [ 0.0000,  0.0000, -0.4770,  0.6229]], dtype=torch.float64)\n",
      "tensor(0.1004, dtype=torch.float64)\n",
      "\n",
      "tensor([[ 0.4728,  4.4862,  0.0000,  0.0000],\n",
      "        [-0.1703,  0.4987,  0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000, -0.1576,  2.3014],\n",
      "        [ 0.0000,  0.0000, -0.4772,  0.6233]], dtype=torch.float64)\n",
      "tensor([[ 0.4452,  4.5011,  0.0000,  0.0000],\n",
      "        [-0.1722,  0.5048,  0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000, -0.1575,  2.3013],\n",
      "        [ 0.0000,  0.0000, -0.4772,  0.6232]], dtype=torch.float64)\n",
      "tensor(0.0320, dtype=torch.float64)\n",
      "\n",
      "tensor([[ 0.4728,  4.4862,  0.0000,  0.0000],\n",
      "        [-0.1703,  0.4987,  0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000, -0.1576,  2.3014],\n",
      "        [ 0.0000,  0.0000, -0.4772,  0.6233]], dtype=torch.float64)\n",
      "tensor([[ 0.4635,  4.4927,  0.0000,  0.0000],\n",
      "        [-0.1711,  0.4991,  0.0000,  0.0000],\n",
      "        [ 0.0000,  0.0000, -0.1576,  2.3014],\n",
      "        [ 0.0000,  0.0000, -0.4772,  0.6233]], dtype=torch.float64)\n",
      "tensor(0.0114, dtype=torch.float64)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Construct 4th order product approximation\n",
    "\n",
    "hA4 = torch.zeros((4, 4), dtype=dtype)\n",
    "\n",
    "for i in range(len(elements)):\n",
    "    for j in range(i + 1):\n",
    "        for k in range(j + 1):\n",
    "            for l in range(k + 1):\n",
    "                dij = (i == j)\n",
    "                dik = (i == k)\n",
    "                djk = (j == k)\n",
    "                dil = (i == l)\n",
    "                djl = (j == l)\n",
    "                dkl = (k == l)\n",
    "                factor = (1.0 / (1.0 + dij)) * (1.0 / (1.0 + dik + djk)) * (1.0 / (1.0 + dil + djl + dkl))\n",
    "                Xi = S @ hAis[i]\n",
    "                Xj = S @ hAis[j]\n",
    "                Xk = S @ hAis[k]\n",
    "                Xl = S @ hAis[l]\n",
    "                hA4 += factor*(\n",
    "                    1/12*bracket(bracket(bracket(Xi, Xj), Xk), Xl) +\n",
    "                    1/12*bracket(Xi, bracket(bracket(Xj, Xk), Xl)) +\n",
    "                    1/12*bracket(Xi, bracket(Xj, bracket(Xk, Xl))) +\n",
    "                    1/12*bracket(Xj, bracket(Xk, bracket(Xl, Xi)))\n",
    "                )\n",
    "\n",
    "hA4 = - S @ hA4\n",
    "\n",
    "# Construct one-turn matrix\n",
    "\n",
    "print(M1)\n",
    "print(M0)\n",
    "print((M1 - M0).norm())\n",
    "print()\n",
    "\n",
    "hM = M0 @ torch.linalg.matrix_exp(S @ hA1)\n",
    "\n",
    "print(M1)\n",
    "print(hM)\n",
    "print((M1 - hM).norm())\n",
    "print()\n",
    "\n",
    "hM = M0 @ torch.linalg.matrix_exp(S @ (hA1 + hA2))\n",
    "\n",
    "print(M1)\n",
    "print(hM)\n",
    "print((M1 - hM).norm())\n",
    "print()\n",
    "\n",
    "hM = M0 @ torch.linalg.matrix_exp(S @ (hA1 + hA2 + hA3))\n",
    "\n",
    "print(M1)\n",
    "print(hM)\n",
    "print((M1 - hM).norm())\n",
    "print()\n",
    "\n",
    "hM = M0 @ torch.linalg.matrix_exp(S @ (hA1 + hA2 + hA3 + hA4))\n",
    "\n",
    "print(M1)\n",
    "print(hM)\n",
    "print((M1 - hM).norm())\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.14.0"
  },
  "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
}