{ "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 }