{ "cells": [ { "attachments": {}, "cell_type": "markdown", "id": "556562f3-8ece-4517-8c93-ee5e2fc29131", "metadata": {}, "source": [ "# Example-08: Multipole element factory" ] }, { "cell_type": "markdown", "id": "ed4b5f6d-27ee-435a-be3a-4e4b3eb288e4", "metadata": {}, "source": [ "In this example multipole factory is illustrated. \n", "\n", "The multipole 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, a_{s, q} + a_{s, s} + a_{s, o})\\\\\n", "& a_{s, q} = -\\frac{1}{2} k_{n, q} \\left(q_x^2 + q_y^2 \\right) + k_{s, q} q_x q_y \\\\\n", "& a_{s, s} = -\\frac{1}{2!} k_{n, s} \\left(\\frac{q_x^3}{3} - q_x q_y^2 \\right) - \\frac{1}{2!} k_{s, s} \\left(\\frac{q_y^3}{3} - q_x^2 q_y\\right) \\\\\n", "& a_{s, o} = -\\frac{1}{3!} k_{n, o} \\left(\\frac{q_x^4}{4} - 3 q_x^2 q_y^2 + \\frac{q_y^4}{4}\\right) - \\frac{1}{3!} k_{s, o} \\left(q_x q_y^3 - q_x^3 q_y \\right) \\\\\n", "& \\\\\n", "& \\varphi = 0 \\\\\n", "& t = h = 0 \\\\\n", "\\end{align}\n", "$\n", "\n", "The constructed element signature is:\n", "\n", "```python\n", "def multipole(qsps:Array, length:Array, kq_n:Array, kq_s:Array, ks_n:Array, ks_s:Array, ko_n:Array, ko_s:Array) -> Array:\n", " ...\n", "```" ] }, { "cell_type": "code", "execution_count": 1, "id": "66791263-2744-4976-bef7-1a8bf3d9c3cf", "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.multipole import multipole_factory\n", "\n", "jax.numpy.set_printoptions(linewidth=256, precision=12)" ] }, { "cell_type": "code", "execution_count": 2, "id": "01323bc3-6ed0-4155-9b99-3520b4cf1452", "metadata": {}, "outputs": [], "source": [ "# Set data type\n", "\n", "jax.config.update(\"jax_enable_x64\", True)" ] }, { "cell_type": "code", "execution_count": 3, "id": "e60584ad-4e29-40f0-aa7d-a917419b9619", "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": "cfd67fcc-16f3-4b0f-aadc-db3d28bde7d5", "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": "fc6a4d54-5a02-4c8e-8536-73aca613397f", "metadata": {}, "outputs": [], "source": [ "# Define generic multipole element\n", "\n", "gamma = 10**3\n", "element = jit(multipole_factory(beta=beta(gamma), gamma=gamma, order=2**1, iterations=100))" ] }, { "cell_type": "code", "execution_count": 6, "id": "05aef585-c73c-43e7-8b09-5ecd80cb41c4", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "[-0.010195569406 0.004634631697 0.00099927232 -0.002587384238 -0.003898816771 -0.0001 ]\n", "[-0.010195569406 0.004634631697 0.00099927232 -0.002587384238 -0.003898816771 -0.0001 ]\n", "True\n" ] } ], "source": [ "# Compare with PTC\n", "\n", "length = jax.numpy.float64(0.25)\n", "\n", "kq_n = jax.numpy.float64(-2.0)\n", "kq_s = jax.numpy.float64(+1.5)\n", "ks_n = jax.numpy.float64(-50.0)\n", "ks_s = jax.numpy.float64(+75.0)\n", "ko_n = jax.numpy.float64(-100.0)\n", "ko_s = jax.numpy.float64(+500.0)\n", "\n", "print(res := element(qsps, length, kq_n, kq_s, ks_n, ks_s, ko_n, ko_s))\n", "print(ref := ptc(qsps, 'quadrupole', {'l': float(length), 'knl': f'{{0.0,{float(kq_n*length)}, {float(ks_n*length)}, {float(ko_n*length)}}}', 'ksl': f'{{0.0,{float(kq_s*length)}, {float(ks_s*length)}, {float(ko_s*length)}}}'}, gamma=gamma))\n", "print(jax.numpy.allclose(res, ref))" ] }, { "cell_type": "code", "execution_count": 7, "id": "e6811a93-111a-4add-a5e1-4f488d14dd67", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "[[ 1.058513409511e+00 1.641931362563e-02 0.000000000000e+00 2.548728245266e-01 1.371326922699e-03 1.953577918270e-04]\n", " [ 1.642278685252e-02 9.426914513883e-01 0.000000000000e+00 1.371531111411e-03 2.452402248722e-01 3.641389680595e-04]\n", " [ 1.183449134442e-04 -1.074855507550e-04 1.000000000000e+00 2.074102528591e-04 3.548724731577e-04 1.701670012363e-06]\n", " [ 4.713919361348e-01 1.315473343451e-01 0.000000000000e+00 1.058139548857e+00 1.647371074410e-02 2.099930109674e-05]\n", " [ 1.316213051530e-01 -4.522013363610e-01 0.000000000000e+00 1.647803990033e-02 9.430571887037e-01 -3.375926930070e-05]\n", " [ 0.000000000000e+00 0.000000000000e+00 0.000000000000e+00 0.000000000000e+00 0.000000000000e+00 1.000000000000e+00]]\n", "\n", "[-2.587671336735e-03 -3.899249388134e-03 -1.095013990866e-05 -1.481596650450e-02 -1.916336194105e-02 0.000000000000e+00]\n", "\n" ] } ], "source": [ "# Differentiability\n", "\n", "print(jacrev(element)(qsps, length, kq_n, kq_s, ks_n, ks_s, ko_n, ko_s))\n", "print()\n", "\n", "print(jacrev(element, 1)(qsps, length, kq_n, kq_s, ks_n, ks_s, ko_n, ko_s))\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 }