Merge pull request #120 from ChristopherMayes/update_docs_with_wavefront

Add wavefront notebooks to docs

Author: ChristopherMayes

docs/examples/wavefront/advanced_drift.ipynb

@@ -22,7 +22,6 @@
     "    drift_wavefront,\n",
     "    drift_wavefront_advanced,\n",
     ")\n",
-    "from pmd_beamphysics.wavefront.gaussian import add_gaussian\n",
     "from pmd_beamphysics.units import Z0, c_light\n",
     "import numpy as np\n",
     "\n",
@@ -44,19 +43,20 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "W0 = Wavefront(\n",
-    "    Ex=np.zeros((101, 101, 1368)),\n",
+    "wavelength = 1.2227e-10\n",
+    "w0 = 50e-6\n",
+    "sigma0 = w0 / 2  # RMS beam size at waist\n",
+    "zR = np.pi * w0**2 / wavelength\n",
+    "\n",
+    "W = Wavefront.from_gaussian(\n",
+    "    shape=(101, 101, 1368),\n",
     "    dx=6e-6,\n",
     "    dy=6e-6,\n",
     "    dz=1.46724e-09,\n",
-    "    wavelength=1.2227e-10,\n",
+    "    wavelength=wavelength,\n",
+    "    sigma0=sigma0,\n",
+    "    energy=1.2345e-6,\n",
     ")\n",
-    "w0 = 50e-6\n",
-    "zR = np.pi * w0**2 / W0.wavelength\n",
-    "\n",
-    "\n",
-    "W = W0.copy()\n",
-    "add_gaussian(W, z=0, w0=w0, energy=1.2345e-6)\n",
     "\n",
     "W.plot_fluence()"
    ]
@@ -315,6 +315,14 @@
     "plt.legend()"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "f34daabc",
+   "metadata": {},
+   "source": [
+    "# Check energy conservation"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -346,7 +354,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
+   "display_name": "pmd_beamphysics-dev",
    "language": "python",
    "name": "python3"
   },

…amples/wavefront/pytorch_wavefront.ipynb …ont/experimental_pytorch_wavefront.ipynbdocs/examples/wavefront/pytorch_wavefront.ipynb renamed to docs/examples/wavefront/experimental_pytorch_wavefront.ipynb docs/examples/wavefront/pytorch_wavefront.ipynb renamed to docs/examples/wavefront/experimental_pytorch_wavefront.ipynb

@@ -20,6 +20,10 @@ nav:
           - examples/fields/field_conversion.ipynb
           - examples/fields/corrector_modeling.ipynb
           - examples/fields/solenoid_modeling.ipynb
+      - Wavefront:
+          - examples/wavefront/wavefront.ipynb
+          - examples/wavefront/advanced_drift.ipynb
+          - examples/wavefront/experimental_pytorch_wavefront.ipynb
       - Utilities:
           - examples/units.ipynb
           - examples/labels.ipynb

mkdocs.yml

@@ -190,6 +190,10 @@ def __init__(self, h5=None, data=None):
                     assert len(pp) == 1, f"Number of particle paths in {h5}: {len(pp)}"
                     data = load_bunch_data(hh5[pp[0]])
 
+            elif isinstance(h5, File):
+                pp = particle_paths(h5)
+                assert len(pp) == 1, f"Number of particle paths in {h5}: {len(pp)}"
+                data = load_bunch_data(h5[pp[0]])
             else:
                 # Try dict
                 data = load_bunch_data(h5)

pmd_beamphysics/particles.py

@@ -1,7 +1,7 @@
 #!/bin/bash
 
 # Define an array of patterns to skip
-SKIP_LIST=("parallel")
+SKIP_LIST=("parallel", "experimental")
 
 
 export PYDEVD_DISABLE_FILE_VALIDATION=1