{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "556562f3-8ece-4517-8c93-ee5e2fc29131",
   "metadata": {},
   "source": [
    "# Example-09: Dipole element factory"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1403df6c-3bb8-488e-836b-efea26970468",
   "metadata": {},
   "source": [
    "In this example dipole factory is illustrated. \n",
    "\n",
    "The dipole hamiltonian is:\n",
    "\n",
    "$\n",
    "\\begin{align}\n",
    "& H(q_x, q_y, q_s, p_x, p_y, p_s; s) = \\frac{p_s}{\\beta} - t(s)(q_x p_y - q_y p_x) - (1 + h(s) q_x) \\left(\\sqrt{P_s^2 - P_x^2 - P_y^2 - \\frac{1}{\\beta^2 \\gamma^2}} + a_s(q_x, q_y, q_s; s)\\right)  \\\\\n",
    "& \\\\\n",
    "& P_s = p_s + 1/\\beta - \\varphi(q_x, q_y, q_s; s)  \\\\\n",
    "& P_x = p_x - a_x(q_x, q_y, q_s; s)  \\\\\n",
    "& P_y = p_y - a_y(q_x, q_y, q_s; s) \\\\\n",
    "\\\\\n",
    "& (a_x, a_y, a_s) = (0, 0, -\\frac{1}{1 + q_x/\\rho}\\left(\\frac{q_x}{ \\rho} + \\frac{q_x^2}{2 \\rho^2}   \\right))\\\\\n",
    "& \\varphi = 0 \\\\\n",
    "& t = 0 \\\\\n",
    "& h = \\frac{1}{\\rho} = \\frac{\\alpha}{l}\n",
    "\\end{align}\n",
    "$\n",
    "\n",
    "The constructed element signature is:\n",
    "\n",
    "```python\n",
    "def dipole(qsps:Array, length:Array, angle:Array) -> Array:\n",
    "    ...\n",
    "```\n",
    "\n",
    "Note, no fringe effects are icluded.\n",
    "\n",
    "By default, exact solution is used to transfrom initial conditions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "0b0c1fd3-73d5-4764-a3f3-2cda5cec9cda",
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax\n",
    "from jax import jit\n",
    "from jax import jacrev\n",
    "\n",
    "from elementary.util import ptc\n",
    "from elementary.util import beta\n",
    "from elementary.dipole import dipole_factory\n",
    "\n",
    "jax.numpy.set_printoptions(linewidth=256, precision=12)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e71271a3-de50-4f46-bc10-5158b466c940",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set data type\n",
    "\n",
    "jax.config.update(\"jax_enable_x64\", True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "b1cdb59f-67b0-4df7-a237-9e058ae05cd4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set device\n",
    "\n",
    "device, *_ = jax.devices('cpu')\n",
    "jax.config.update('jax_default_device', device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "68b9ba1b-3778-45f5-aa69-ea14a1e2f099",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set initial condition\n",
    "\n",
    "(q_x, q_y, q_s) = qs = jax.numpy.array([-0.01, 0.005, 0.002])\n",
    "(p_x, p_y, p_s) = ps = jax.numpy.array([0.001, 0.001, -0.0005])\n",
    "qsps = jax.numpy.hstack([qs, ps])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "6489a075-661f-4225-af55-129d6a34f67e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define generic dipole element\n",
    "\n",
    "gamma = 10**3\n",
    "element = jit(dipole_factory(beta=beta(gamma), gamma=gamma, order=2**1, iterations=100))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "d00e7b6f-44b8-4e73-880f-f07c7cbcb460",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[-0.008012832402  0.007000552059  0.00244821723   0.000986205451  0.001          -0.0005        ]\n",
      "[-0.008012832402  0.007000552059  0.00244821723   0.000986205451  0.001          -0.0005        ]\n",
      "True\n"
     ]
    }
   ],
   "source": [
    "# Compare with PTC\n",
    "\n",
    "length = jax.numpy.float64(2.0)\n",
    "angle = jax.numpy.float64(0.05)\n",
    "\n",
    "print(res := element(qsps, length, angle))\n",
    "print(ref := ptc(qsps, 'sbend', {'l': float(length), 'angle': float(angle), 'kill_ent_fringe': 'true', 'kill_exi_fringe': 'true'}, gamma=gamma))\n",
    "print(jax.numpy.allclose(res, ref))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "37020a27-e850-49d2-ba98-2b71841dcb3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define generic dipole element using hamiltonian\n",
    "\n",
    "gamma = 10**3\n",
    "element = jit(dipole_factory(exact=False, beta=beta(gamma), gamma=gamma, order=2**1, iterations=100))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "cfcb472b-ac68-4de3-bfbc-acc2c387883a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[-0.008012832402  0.007000552059  0.002448217234  0.000986205451  0.001          -0.0005        ]\n",
      "[-0.008012832402  0.007000552059  0.00244821723   0.000986205451  0.001          -0.0005        ]\n",
      "True\n"
     ]
    }
   ],
   "source": [
    "# Compare with PTC\n",
    "\n",
    "length = jax.numpy.float64(2.0)\n",
    "angle = jax.numpy.float64(0.05)\n",
    "\n",
    "print(res := element(qsps, length, angle))\n",
    "print(ref := ptc(qsps, 'sbend', {'l': float(length), 'angle': float(angle), 'kill_ent_fringe': 'true', 'kill_exi_fringe': 'true'}, gamma=gamma))\n",
    "print(jax.numpy.allclose(res, ref))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "f9ddc810-3fa0-4cd7-bddb-1ebdbdb7ce23",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[ 9.987995748301e-01  0.000000000000e+00  0.000000000000e+00  1.999720113524e+00 -4.804052661206e-05  4.801653036904e-02]\n",
      " [ 5.000422073929e-05  1.000000000000e+00  0.000000000000e+00  5.201635774981e-05  2.000554060243e+00 -2.000724168894e-03]\n",
      " [-4.997924363105e-02  0.000000000000e+00  1.000000000000e+00 -5.199037557913e-02 -2.000724168894e-03 -8.272513473157e-04]\n",
      " [-1.249479231765e-03  0.000000000000e+00  0.000000000000e+00  9.987002561736e-01 -5.000422142308e-05  4.997924431449e-02]\n",
      " [ 0.000000000000e+00  0.000000000000e+00  0.000000000000e+00  0.000000000000e+00  1.000000000000e+00  0.000000000000e+00]\n",
      " [ 0.000000000000e+00  0.000000000000e+00  0.000000000000e+00  0.000000000000e+00  0.000000000000e+00  1.000000000000e+00]]\n",
      "\n",
      "[ 9.875816733888e-04  1.000526050363e-03 -2.578760124653e-05 -6.247396158822e-06  0.000000000000e+00  0.000000000000e+00]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Differentiability\n",
    "\n",
    "print(jacrev(element)(qsps, length, angle))\n",
    "print()\n",
    "\n",
    "print(jacrev(element, 1)(qsps, length, angle))\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
}