{ "cells": [ { "attachments": {}, "cell_type": "markdown", "id": "556562f3-8ece-4517-8c93-ee5e2fc29131", "metadata": {}, "source": [ "# Example-04: Drift element factory" ] }, { "cell_type": "markdown", "id": "5bbb7617-3919-4921-b060-437844459bf5", "metadata": {}, "source": [ "In this example drift factory is illustrated. \n", "\n", "The drift 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, 0) \\\\\n", "& \\varphi = 0 \\\\\n", "& t = h = 0 \\\\\n", "\\end{align}\n", "$\n", "\n", "The constructed element signature is:\n", "\n", "```python\n", "def drift(qsps:Array, length:Array) -> Array:\n", " ...\n", "```\n", "\n", "Note, by default, exact solution is used instead of hamiltonial based." ] }, { "cell_type": "code", "execution_count": 1, "id": "774153b3-086e-40b9-a830-432f554a5eb0", "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.drift import drift_factory\n", "\n", "jax.numpy.set_printoptions(linewidth=256, precision=12)" ] }, { "cell_type": "code", "execution_count": 2, "id": "4c7f05b2-08d6-4e0b-b753-426a8daceb88", "metadata": {}, "outputs": [], "source": [ "# Set data type\n", "\n", "jax.config.update(\"jax_enable_x64\", True)" ] }, { "cell_type": "code", "execution_count": 3, "id": "d7825540-dff0-45d2-b3d3-ef4d1c0eabfd", "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": "cb35d996-ba64-4bc9-8d2f-21ad896d2b59", "metadata": {}, "outputs": [], "source": [ "# Set initial condition\n", "\n", "(q_x, q_y, q_s) = qs = jax.numpy.array([0.0, 0.0, 0.01])\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": "b41a0022-587c-42e2-b2aa-a0c049c79cb0", "metadata": {}, "outputs": [], "source": [ "# Define generic drift element\n", "\n", "gamma = 10**3\n", "element = jit(drift_factory(beta=beta(gamma), gamma=gamma))" ] }, { "cell_type": "code", "execution_count": 6, "id": "e000e4b5-b81c-4aef-95c4-24c8c416fbff", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "[ 0.00100010101 0.00100010101 0.009998999698 0.001 0.001 -0.0001 ]\n", "[ 0.00100010101 0.00100010101 0.009998999698 0.001 0.001 -0.0001 ]\n", "True\n" ] } ], "source": [ "# Compare with PTC\n", "\n", "print(res := element(qsps, 1.0))\n", "print(ref := ptc(qsps, 'drift', {'l': 1.0}, gamma=gamma))\n", "print(jax.numpy.allclose(res, ref))" ] }, { "cell_type": "code", "execution_count": 7, "id": "244a9ee3-3b8b-41e2-ac54-80a0e7379d84", "metadata": {}, "outputs": [], "source": [ "# Define generic drift element using hamiltonian *)\n", "\n", "gamma = 10**3\n", "element = jit(drift_factory(exact=False, beta=beta(gamma), gamma=gamma))" ] }, { "cell_type": "code", "execution_count": 8, "id": "102f5a57-80f6-49c3-80e6-ed624f780e14", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "[ 0.00100010101 0.00100010101 0.009998999698 0.001 0.001 -0.0001 ]\n", "[ 0.00100010101 0.00100010101 0.009998999698 0.001 0.001 -0.0001 ]\n", "True\n" ] } ], "source": [ "# Compare with PTC\n", "\n", "print(res := element(qsps, 1.0))\n", "print(ref := ptc(qsps, 'drift', {'l': 1.0}, gamma=gamma))\n", "print(jax.numpy.allclose(res, ref))" ] }, { "cell_type": "code", "execution_count": 9, "id": "ceb3af1e-2ffe-4fc2-93b3-e7a77c4df411", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "[[ 1.000000000000e+00 0.000000000000e+00 0.000000000000e+00 1.000102010656e+00 1.000303061668e-06 -1.000203531514e-03]\n", " [ 0.000000000000e+00 1.000000000000e+00 0.000000000000e+00 1.000303061668e-06 1.000102010656e+00 -1.000203531514e-03]\n", " [ 0.000000000000e+00 0.000000000000e+00 1.000000000000e+00 -1.000203531514e-03 -1.000203531514e-03 3.000910185458e-06]\n", " [ 0.000000000000e+00 0.000000000000e+00 0.000000000000e+00 1.000000000000e+00 0.000000000000e+00 0.000000000000e+00]\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", "[ 1.000101010353e-03 1.000101010353e-03 -1.000302046226e-06 0.000000000000e+00 0.000000000000e+00 0.000000000000e+00]\n", "\n" ] } ], "source": [ "# Differentiability\n", "\n", "length = jax.numpy.float64(1.0)\n", "\n", "print(jacrev(element)(qsps, length))\n", "print()\n", "\n", "print(jacrev(element, -1)(qsps, length))\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 }