{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "262a5ec8-2553-4237-ab62-319b6ca22089",
   "metadata": {},
   "source": [
    "# Example-60: Chromatic matrix (element)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "17bd55fa-6965-4d18-9e8c-a4fd6ec68b2d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Can be used to model 4x4 linear elements with leading order chromatic effects\n",
    "# Transport matrix is constructed given elements of the symmetric matrix\n",
    "# M = exp(S A) exp(dp S B)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "a6ab7bee-d4b1-4ea9-9a32-4cbc5f124da9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.linalg import logm\n",
    "\n",
    "import torch\n",
    "from model.library.drift import Drift\n",
    "from model.library.matrix import Matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "f3a2b2f3-fc6f-4066-b0d6-34d23d839a19",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Drift element\n",
    "\n",
    "D = Drift('D', length=1.5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "88eee13c-83d3-4581-9e0b-0faff3d6d243",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[1.0000, 1.5000, 0.0000, 0.0000],\n",
      "        [0.0000, 1.0000, 0.0000, 0.0000],\n",
      "        [0.0000, 0.0000, 1.0000, 1.5000],\n",
      "        [0.0000, 0.0000, 0.0000, 1.0000]], dtype=torch.float64)\n",
      "\n",
      "tensor([[1.0000, 1.4985, 0.0000, 0.0000],\n",
      "        [0.0000, 1.0000, 0.0000, 0.0000],\n",
      "        [0.0000, 0.0000, 1.0000, 1.4985],\n",
      "        [0.0000, 0.0000, 0.0000, 1.0000]], dtype=torch.float64)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Parametric drift matrix\n",
    "\n",
    "def matrix(state, dp):\n",
    "    return torch.func.jacrev(lambda state: D(state, data={**D.data(), **{'dp': dp}}))(state)\n",
    "\n",
    "state = torch.tensor([0.0, 0.0, 0.0, 0.0], dtype=torch.float64)\n",
    "\n",
    "print(matrix(state, 0.0))\n",
    "print()\n",
    "\n",
    "print(matrix(state, 0.001))\n",
    "print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "2a554a4b-fa95-4f03-84d6-124b3df7002d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[1.0000, 1.4985, 0.0000, 0.0000],\n",
      "        [0.0000, 1.0000, 0.0000, 0.0000],\n",
      "        [0.0000, 0.0000, 1.0000, 1.4985],\n",
      "        [0.0000, 0.0000, 0.0000, 1.0000]], dtype=torch.float64)\n",
      "\n",
      "tensor([[1.0000, 1.4985, 0.0000, 0.0000],\n",
      "        [0.0000, 1.0000, 0.0000, 0.0000],\n",
      "        [0.0000, 0.0000, 1.0000, 1.4985],\n",
      "        [0.0000, 0.0000, 0.0000, 1.0000]], dtype=torch.float64)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Construct A and B matrices\n",
    "\n",
    "S = torch.tensor([[0.,  1.,  0.,  0.], [-1., 0.,  0.,  0.], [0.,  0.,  0.,  1.], [0.,  0., -1.,  0.]], dtype=torch.float64)\n",
    "\n",
    "state = torch.tensor([0.0, 0.0, 0.0, 0.0], dtype=torch.float64)\n",
    "dp = torch.tensor(0.0, dtype=torch.float64)\n",
    "\n",
    "m = matrix(state, dp)\n",
    "dmdp = torch.func.jacrev(matrix, 1)(state, dp)\n",
    "\n",
    "A = (torch.linalg.inv(S) @ torch.tensor(logm(m), dtype=torch.float64))\n",
    "B = - S @ m @ dmdp\n",
    "\n",
    "print(matrix(state, 0.001))\n",
    "print()\n",
    "\n",
    "print(torch.matrix_exp(S @ A) @ torch.matrix_exp(0.001 * S @ B))\n",
    "print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "1375e8b2-b0a1-4d0e-8809-83be4b95e8c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define matrix element\n",
    "\n",
    "M = Matrix('M', \n",
    "           length=1.5, \n",
    "           A=A[torch.triu(torch.ones_like(A, dtype=torch.bool))].tolist(), \n",
    "           B=B[torch.triu(torch.ones_like(B, dtype=torch.bool))].tolist())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "99ed416b-e257-4dde-8f5f-beb3f675c66e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[1.0000, 1.5000, 0.0000, 0.0000],\n",
      "        [0.0000, 1.0000, 0.0000, 0.0000],\n",
      "        [0.0000, 0.0000, 1.0000, 1.5000],\n",
      "        [0.0000, 0.0000, 0.0000, 1.0000]], dtype=torch.float64)\n",
      "\n",
      "tensor([[1.0000, 1.4985, 0.0000, 0.0000],\n",
      "        [0.0000, 1.0000, 0.0000, 0.0000],\n",
      "        [0.0000, 0.0000, 1.0000, 1.4985],\n",
      "        [0.0000, 0.0000, 0.0000, 1.0000]], dtype=torch.float64)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Compare with drift\n",
    "\n",
    "def matrix(state, dp):\n",
    "    return torch.func.jacrev(lambda state: M(state, data={**M.data(), **{'dp': dp}}))(state)\n",
    "\n",
    "state = torch.tensor([0.0, 0.0, 0.0, 0.0], dtype=torch.float64)\n",
    "\n",
    "print(matrix(state, 0.0))\n",
    "print()\n",
    "\n",
    "print(matrix(state, 0.001))\n",
    "print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3eee112d-80bd-45b5-975e-6afdaa34e068",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}