{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "556562f3-8ece-4517-8c93-ee5e2fc29131",
   "metadata": {},
   "source": [
    "# Example-11: Alignment errors (straight layout)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf53d7da-c3d7-4c27-b8c8-9aa2e4f55fce",
   "metadata": {},
   "source": [
    "In this example alignment errors for straight layout are illustrated.\n",
    "\n",
    "Alignment errors (translations and rotations) are defined with respect to the element entrance frame.\n",
    "\n",
    "The full transformations sequence:\n",
    "\n",
    "```python\n",
    "# forward translations and rotations\n",
    "qsps = tx(qsps, +dx)\n",
    "qsps = ty(qsps, +dy)\n",
    "qsps = tz(qsps, +dz, beta, constant)\n",
    "qsps = rx(qsps, +wx, beta, constant)\n",
    "qsps = ry(qsps, +wy, beta, constant)\n",
    "qsps = rz(qsps, +wz)\n",
    "\n",
    "# element body transformation\n",
    "qsps = element(qsps, length, ...)\n",
    "\n",
    "# finite lenght correction\n",
    "qsps = tz(qsps, -length, beta=beta, constant=constant)\n",
    "\n",
    "# inverse translation and rotations\n",
    "qsps = rz(qsps, -wz)\n",
    "qsps = ry(qsps, -wy, beta, constant)\n",
    "qsps = rx(qsps, -wx, beta, constant)\n",
    "qsps = tz(qsps, -dz, beta, constant)\n",
    "qsps = ty(qsps, -dy)\n",
    "qsps = tx(qsps, -dx)\n",
    "\n",
    "# finite lenght correction\n",
    "qsps = tz(qsps, +length, beta, constant)\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "3a490b58-d056-4d5a-9e63-5b4314b71a46",
   "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.quadrupole import quadrupole_factory\n",
    "from elementary.alignment import alignment_factory\n",
    "\n",
    "jax.numpy.set_printoptions(linewidth=256, precision=12)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "8e9f4a53-0f8a-4d8c-93c2-3841c5487e8d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set data type\n",
    "\n",
    "jax.config.update(\"jax_enable_x64\", True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "7df9a0e7-2aa2-47e2-9e9e-10493d2ddd12",
   "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": "da0581c2-d4bf-4eba-8cee-889e37c411c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set initial condition\n",
    "\n",
    "(q_x, q_y, q_s) = qs = jax.numpy.array([-0.01, 0.005, 0.001])\n",
    "(p_x, p_y, p_s) = ps = jax.numpy.array([0.001, 0.001, -0.0001])\n",
    "qsps = jax.numpy.hstack([qs, ps])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "789c1ea3-37f0-42b8-91c1-90c5f28060a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define generic quadrupole element with alignment errors\n",
    "\n",
    "gamma = 10**3\n",
    "\n",
    "body = quadrupole_factory(beta=beta(gamma), gamma=gamma, order=2**1, iterations=200)\n",
    "xyz_entrance, xyz_exit = alignment_factory(beta=beta(gamma), gamma=gamma, flag=False)\n",
    "\n",
    "@jit\n",
    "def element(x, length, kn, ks, dx, dy, dz, wx, wy, wz):\n",
    "    x = xyz_entrance(x, dx, dy, dz, wx, wy, wz)\n",
    "    x = body(x, length, kn, ks)\n",
    "    x = xyz_exit(x, dx, dy, dz, wx, wy, wz, length)\n",
    "    return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b2bdba8e-b5a0-4e3f-973d-5bb43fe1aa3b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set alignment errors\n",
    "\n",
    "dx, dy, dz = jax.numpy.array([0.05, -0.02, 0.05])\n",
    "wx, wy, wz = jax.numpy.array([0.005, -0.005, 0.1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "da17a076-fb9a-4dcc-8d52-f4640d4ea991",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[-4.532564066213e-02 -5.683873885073e-02 -2.649160041102e-03 -1.148938967200e-01 -1.263031959206e-01 -1.000000000000e-04]\n",
      "[-4.532564066213e-02 -5.683873885073e-02 -2.649160041114e-03 -1.148938967200e-01 -1.263031959206e-01 -1.000000000000e-04]\n",
      "True\n"
     ]
    }
   ],
   "source": [
    "# Compare with PTC\n",
    "\n",
    "length = jax.numpy.float64(1.0)\n",
    "kn = jax.numpy.float64(-2.0)\n",
    "ks = jax.numpy.float64(+1.5)\n",
    "\n",
    "print(res := element(qsps, length, kn, ks, dx, dy, dz, wx, wy, wz))\n",
    "print(ref := ptc(qsps, 'quadrupole', {'l': float(length), 'k1': float(kn), 'k1s': float(ks)}, gamma=gamma, tx=float(dx), ty=float(dy), tz=float(dz), rx=float(wx), ry=float(wy), rz=float(wz)))\n",
    "print(jax.numpy.allclose(res, ref))"
   ]
  }
 ],
 "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
}