06_27_2025: cleanup and docs update

Author: i-a-morozov

docs/source/examples/example-01.ipynb

@@ -11,18 +11,17 @@
   {
    "cell_type": "code",
    "execution_count": 1,
-   "id": "9aca811d-d252-4872-b1d2-84eb2ffb50ef",
+   "id": "e5b0ff81-f09d-4772-b3b5-a5fa89183572",
    "metadata": {},
    "outputs": [],
    "source": [
-    "# In this example non-autonomous hamiltonial is integrated using midpoint and tao integrators\n",
-    "# For explicitly defined hamiltonian function and factory generated one"
+    "# In this example non-autonomous hamiltonial is integrated"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": 2,
-   "id": "936385d7-e524-4ceb-b98b-781cd68e9bb6",
+   "id": "8d36ea78-e74a-4b7c-bea9-b3fb3e912f2d",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -31,7 +30,7 @@
     "import jax\n",
     "from jax import Array\n",
     "from jax import jit\n",
-    "from jax import vmap\n",
+    "from jax import jacrev\n",
     "\n",
     "from elementary import fold\n",
     "from elementary import nest\n",
@@ -48,7 +47,7 @@
   {
    "cell_type": "code",
    "execution_count": 3,
-   "id": "702bdf71-e4d6-49b8-bd42-42a68d52cb21",
+   "id": "b8fc02ee-9f08-4456-92fe-d561a276a2c0",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -60,7 +59,7 @@
   {
    "cell_type": "code",
    "execution_count": 4,
-   "id": "fa15f4d1-0702-4969-bafd-63fb23793950",
+   "id": "dc2e0cd9-8be4-4f2c-a40a-335c41935958",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -73,7 +72,7 @@
   {
    "cell_type": "code",
    "execution_count": 5,
-   "id": "bb56ad88-26ed-465f-b141-d581cb5cd324",
+   "id": "fd191c35-b7ce-4038-86dd-2a0f1a56e475",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -87,7 +86,7 @@
   {
    "cell_type": "code",
    "execution_count": 6,
-   "id": "43d3146a-ae63-4ad0-9971-80c40c8d7986",
+   "id": "d0bb3c3f-eea3-4894-8a7a-79dcd729b369",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -101,7 +100,7 @@
   {
    "cell_type": "code",
    "execution_count": 7,
-   "id": "3b57db3d-8b24-4220-bd76-34394cf77dcb",
+   "id": "d957559c-df5d-430f-b080-74e2327a15c7",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -121,7 +120,7 @@
   {
    "cell_type": "code",
    "execution_count": 8,
-   "id": "3cdc8a98-2453-4378-b79d-aeb4a5405f25",
+   "id": "ec7439c5-33c0-4b90-841d-5cb91de3e615",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -135,7 +134,7 @@
   {
    "cell_type": "code",
    "execution_count": 9,
-   "id": "b1a8cb34-eb7b-4e2e-9b97-99dbee504c36",
+   "id": "7bf4d7b0-8f7f-4f56-8a37-da51aaa8a8f9",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -147,7 +146,7 @@
   {
    "cell_type": "code",
    "execution_count": 10,
-   "id": "4c99d4c6-da3b-4330-a48e-bf7f6799895e",
+   "id": "2cddbce9-6bdc-4775-b704-77a8751fe0d5",
    "metadata": {},
    "outputs": [
     {
@@ -169,7 +168,7 @@
   {
    "cell_type": "code",
    "execution_count": 11,
-   "id": "a6bfe0bf-dc77-41f7-9553-85c683c42df5",
+   "id": "19999d5f-4812-49a5-947d-d8ab2dd23921",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -181,7 +180,7 @@
   {
    "cell_type": "code",
    "execution_count": 12,
-   "id": "a0290e60-2679-457e-932d-49e018ad22ff",
+   "id": "0197ddbc-306e-473e-8efb-ef7350151b65",
    "metadata": {},
    "outputs": [
     {
@@ -203,13 +202,13 @@
   {
    "cell_type": "code",
    "execution_count": 13,
-   "id": "f0f85ac6-0edc-4888-83f3-f13aa9fa654e",
+   "id": "bed740e9-90ae-41cd-9497-4515a1357de8",
    "metadata": {},
    "outputs": [],
    "source": [
     "# Define non-autonomous and extended hamiltonian (factory)\n",
     "\n",
-    "def vector(qs:Array, s:Array, kn:Array, *args:Array) -> Array:\n",
+    "def vector(qs:Array, s:Array, kn:Array, *args:Array) -> tuple[Array, Array, Array]:\n",
     "    q_x, q_y, q_s = qs\n",
     "    a_x, a_y, a_s = jax.numpy.zeros_like(qs)\n",
     "    a_s = - 1/2*kn*(1 + jax.numpy.cos(s))*(q_x**2 + q_y**2)\n",
@@ -227,7 +226,7 @@
   {
    "cell_type": "code",
    "execution_count": 14,
-   "id": "7281a8c6-e48e-412c-8a7a-4171568cfc3d",
+   "id": "e0d482f4-339d-447b-b9e4-bf4ea9affe9a",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -241,7 +240,7 @@
   {
    "cell_type": "code",
    "execution_count": 15,
-   "id": "0f0d5d1c-e319-4903-870b-5a5c889e01ca",
+   "id": "7c4b69cc-a971-496c-b06a-60e5a5523bf4",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -253,7 +252,7 @@
   {
    "cell_type": "code",
    "execution_count": 16,
-   "id": "7ee34eb9-5831-499d-acb2-2d7bb47845ff",
+   "id": "da3b5a65-4622-48ad-98bf-5aeb105efe36",
    "metadata": {},
    "outputs": [
     {
@@ -275,7 +274,7 @@
   {
    "cell_type": "code",
    "execution_count": 17,
-   "id": "f3e51caf-9d1d-4e4f-9c13-f091eddc705f",
+   "id": "82028949-c38d-44a5-951f-adb343d576e3",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -287,7 +286,7 @@
   {
    "cell_type": "code",
    "execution_count": 18,
-   "id": "a7bfdfe6-9287-4301-9db8-b28ab3482e88",
+   "id": "2fe48996-c406-4f12-af48-bcb6ea190787",
    "metadata": {},
    "outputs": [
     {
@@ -309,7 +308,7 @@
   {
    "cell_type": "code",
    "execution_count": 19,
-   "id": "3d4ed639-5ede-4c3c-bd0a-29faca01499e",
+   "id": "951dec84-88db-49fd-980d-3712ee53dac9",
    "metadata": {},
    "outputs": [
     {
@@ -333,13 +332,13 @@
    "source": [
     "# Diffirentiability (initial condition)\n",
     "\n",
-    "jax.jacrev(element)(X, ds, si, kn)"
+    "jacrev(element)(X, ds, si, kn)"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": 20,
-   "id": "a0944868-7d08-4841-9a51-3d71fdec59b5",
+   "id": "6d84bd97-6ce3-418f-89be-6fb842d1c58d",
    "metadata": {},
    "outputs": [
     {
@@ -356,7 +355,7 @@
    "source": [
     "# Diffirentiability (parameter)\n",
     "\n",
-    "jax.jacrev(element, argnums=-1)(X, ds, si, kn)"
+    "jacrev(element, argnums=-1)(X, ds, si, kn)"
    ]
   }
  ],

docs/source/index.rst

@@ -1,7 +1,56 @@
 Welcome to elementary's documentation!
 ===================================
 
-Generic differentiable accelerator elements modeling in JAX
+Generic differentiable accelerator elements modeling in JAX.
+Single particle hamiltonian function:
+
+.. math::
+
+   H(q_x, q_y, q_s, p_x, p_y, p_s; s) = p_s/\beta - t(s)(q_x p_y - q_y p_x) - (1 + h(s) q_x) (\sqrt{(p_s + 1/beta - \varphi(qs, s))^2 - (p_x - a_x(qs, s))^2 - (p_y - a_y(qs, s))^2 - 1/(\beta \gamma)^2} - a_s)
+
+where :math:`\beta` and :math:`\gamma` are the relativistic factors, :math:`h(s)` is the reference trajectory curvature and :math:`t(s)` is the reference trajectory torsion, :math:`a_x(qs, s)`, :math:`a_y(qs, s)` and :math:`a_s(qs, s)` are the scaled vector potential components, and :math:`\varphi(qs, s)` is the scaled scalar potential.
+Additionaly, longitudinal coordinate and momentum are given by:
+
+.. math::
+   q_s = \frac{s}{\beta} - c t
+   p_s = \frac{E}{c P} - 1/beta
+
+Common predefined elements are available or you can create your own by specifying scaled potentials and reference trajectory parameters (curvature and torsion).
+All but vector potentials arguments are optional for hamiltonian and element construction.
+Vector and scalar potentials are assumed to have matching signatures.
+
+.. code-block:: python
+   def vector(qs:Array, s:Array, *args:Array) -> tuple[Array, Array, Array]:
+      q_x, q_y, q_s = qs
+      ...
+
+   def scalar(qs:Array, s:Array, *args:Array) -> Array:
+      q_x, q_y, q_s = qs
+      ...
+
+This is also the case for curvature and torsion functions.
+Note, same extra arguments as in vector and scalar functions should be passed.
+
+.. code-block:: python
+   def curvature(s:Array, *args:Array) -> Array:
+      ...
+
+   def torsion(s:Array, *args:Array) -> Array:
+      ...
+
+The resulting hamiltonian and element signatures are:
+
+.. code-block:: python
+   def hamiltonian(qs: Array, ps: Array, s: Array, *args: Array) -> Array:
+      q_x, q_y, q_s = qs
+      p_x, p_y, p_s = ps
+      ...
+
+   def element(qsps:Array, length:Array, start:Array, *args:Array) -> Array:
+      qs, ps = jax.numpy.reshape(qsps, (2, -1))
+      q_x, q_y, q_s = qs
+      p_x, p_y, p_s = ps
+      ...
 
 .. toctree::
    :caption: Examples:

elementary/dipole.py

@@ -90,7 +90,8 @@ def vector(qs:Array,
                               final=final)
     def dipole(qsps, length, angle, *args):
         r = jax.numpy.abs(length)/angle
-        return element(qsps, length, 0.0, r, *args)
+        start = jax.numpy.zeros_like(length)
+        return element(qsps, length, start, r, *args)
     return dipole
 
 

elementary/hamiltonian.py

@@ -1,8 +1,8 @@
-"""
+r"""
 Hamiltonian
 -----------
 
-Generic single particle Hamiltonian factory
+Generic single particle accelerator Hamiltonian factory
 
 """
 from typing import Callable
@@ -21,13 +21,6 @@ def hamiltonian_factory(vector:Callable[..., tuple[Array, Array, Array]],
     """
     Generic single particle Hamiltonian factory
 
-    H = H(qs, ps; s) = H(q_x, q_y, q_s, p_x, p_y, p_s; s)
-    H = p_s/beta - torsion(s)*(q_x*p_y - q_y*p_x) - (1 + curvature(s)*q_x)*(root + a_s(qs, s))
-    root = sqrt(P_s**2 - P_x**2 - P_y**2 - 1/(beta*gamma)**2)
-    P_s = p_s + 1/beta - scalar(qs, s)
-    P_x = p_x - a_x(qs, s)
-    P_y = p_y - a_y(qs, s)
-
     Parameters
     ----------
     vector: Callable[..., tuple[Array, Array, Array]]