Update PyTorch surrogate example (#547)

* fix old typo

* update pytorch surrogate model to 15 stages

* remove old example files

* trying to include new figures

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* update PASC reference

* update figures, clean code

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* clean out old files

* address pre-commit asserts

* more pre-commit recommendations

* more cleaning

* pre-commit cleaning try 2

* try 3

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* try 4

* remove unused figure reference

* code scanning suggestions

* revert one change, make delta naming consistent

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

docs/source/acknowledge_us.rst

@@ -42,6 +42,11 @@ If your project leads to a scientific publication, please consider citing the pa
 Further ImpactX References
 **************************
 
+- Sandberg R T, Lehe R, Mitchell C E, Garten M, Myers A, Qiang J, Vay J-L and Huebl A.
+  **Synthesizing Particle-in-Cell Simulations Through Learning and GPU Computing for Hybrid Particle Accelerator Beamlines**.
+  Proc. of Platform for Advanced Scientific Computing (PASC'24), *submitted*, 2024.
+  `arXiv:2402.17248 <http://arxiv.org/abs/2402.17248>`__
+
 - Sandberg R T, Lehe R, Mitchell C E, Garten M, Qiang J, Vay J-L and Huebl A.
   **Hybrid Beamline Element ML-Training for Surrogates in the ImpactX Beam-Dynamics Code**.
   14th International Particle Accelerator Conference (IPAC'23), WEPA101, 2023.

examples/pytorch_surrogate_model/README.rst

@@ -1,30 +1,33 @@
 .. _examples-ml-surrogate:
 
-9 Stage Laser-Plasma Accelerator Surrogate
-==========================================
+15 Stage Laser-Plasma Accelerator Surrogate
+===========================================
 
-This example models an electron beam accelerated through nine stages of laser-plasma accelerators with ideal plasma lenses providing the focusing between stages.
+This example models an electron beam accelerated through fifteen stages of laser-plasma accelerators with ideal plasma lenses providing the focusing between stages.
 For more details, see:
 
-
-- Sandberg R T, Lehe R, Mitchell C E, Garten M, Qiang J, Vay J-L and Huebl A.
+- Sandberg R T, Lehe R, Mitchell C E, Garten M, Myers A, Qiang J, Vay J-L and Huebl A.
   **Synthesizing Particle-in-Cell Simulations Through Learning and GPU Computing for Hybrid Particle Accelerator Beamlines**.
   Proc. of Platform for Advanced Scientific Computing (PASC'24), *submitted*, 2024.
+  `arXiv:2402.17248 <http://arxiv.org/abs/2402.17248>`__
+
 - Sandberg R T, Lehe R, Mitchell C E, Garten M, Qiang J, Vay J-L and Huebl A.
   **Hybrid Beamline Element ML-Training for Surrogates in the ImpactX Beam-Dynamics Code**.
   14th International Particle Accelerator Conference (IPAC'23), WEPA101, 2023.
   `DOI:10.18429/JACoW-IPAC2023-WEPA101 <https://doi.org/10.18429/JACoW-IPAC2023-WEPA101>`__
 
 A schematic with more information can be seen in the figure below:
 
-.. figure:: https://user-images.githubusercontent.com/10621396/289956389-c7463b99-fb56-490a-8511-22c43f45cdf8.png
-   :alt: [fig:lpa_schematic] Schematic of the 9 stages of laser-plasma accelerators.
+.. figure:: https://gist.githubusercontent.com/RTSandberg/cf3f6193b3da12e7fd815f69789fd6f2/raw/2308e412d7482d55811afa2e9a0c6fa97627fc2f/schema_15_stages.png
+   :alt: Schematic of the 15 stages of laser-plasma accelerators.
 
-   [fig:lpa_schematic] Schematic of the 9 stages of laser-plasma accelerators.
+   Schematic of the 15 stages of laser-plasma accelerators.
 
-The laser-plasma accelerator elements are modeled with neural network surrogates,
-previously trained and included in ``models``.
-The neural networks require normalized input data; the normalizations can be found in ``datasets``.
+The laser-plasma accelerator elements are modeled with neural networks as surrogates.
+These networks are trained beforehand.
+In this example, pre-trained neural networks are downloaded from `a Zenodo archive <https://zenodo.org/records/10368972>`__ and saved in the ``models`` directory.
+For more about how these neural network surrogate models were created,
+see this `description of a workflow for training neural networks from WarpX simulation data <https://warpx.readthedocs.io/en/latest/usage/workflows/ml_dataset_training.html>`__.
 
 We use a 1 GeV electron beam with initial normalized rms emittance of 1 mm-mrad.
 
@@ -33,49 +36,58 @@ In this test, the initial and final values of :math:`\sigma_x`, :math:`\sigma_y`
 Run
 ---
 
-This example can be **only** be run with **Python**:
+This example can **only** be run with **Python**:
 
-* **Python** script: ``python3 run_ml_surrogate.py``
+* **Python** script: ``python3 run_ml_surrogate_15_stage.py``
 
 For `MPI-parallel <https://www.mpi-forum.org>`__ runs, prefix these lines with ``mpiexec -n 4 ...`` or ``srun -n 4 ...``, depending on the system.
 
-.. literalinclude:: run_ml_surrogate.py
+.. literalinclude:: run_ml_surrogate_15_stage.py
    :language: python
-   :caption: You can copy this file from ``examples/pytorch_surrogate_model/run_ml_surrogate.py``.
+   :caption: You can copy this file from ``examples/pytorch_surrogate_model/run_ml_surrogate_15_stage.py``.
+
+
+This script requires some utility code for using the neural networks that is provided here:
+
+.. dropdown:: Script ``surrogate_model_definitions.py``
+
+   .. literalinclude:: surrogate_model_definitions.py
+      :language: python
+      :caption: You can copy this file from ``examples/pytorch_surrogate_model/surrogate_model_definitions.py``.
 
 Analyze
 -------
 
 We run the following script to analyze correctness:
 
-.. dropdown:: Script ``analyze_ml_surrogate.py``
+.. dropdown:: Script ``analyze_ml_surrogate_15_stage.py``
 
-   .. literalinclude:: analyze_ml_surrogate.py
+   .. literalinclude:: analyze_ml_surrogate_15_stage.py
       :language: python
-      :caption: You can copy this file from ``examples/pytorch_surrogate_model/analyze_ml_surrogate.py``.
+      :caption: You can copy this file from ``examples/pytorch_surrogate_model/analyze_ml_surrogate_15_stage.py``.
 
 Visualize
 ---------
 
 You can run the following script to visualize the beam evolution over time:
 
-.. dropdown:: Script ``visualize_ml_surrogate.py``
+.. dropdown:: Script ``visualize_ml_surrogate_15_stage.py``
 
-   .. literalinclude:: visualize_ml_surrogate.py
+   .. literalinclude:: visualize_ml_surrogate_15_stage.py
       :language: python
-      :caption: You can copy this file from ``examples/pytorch_surrogate_model/visualize_ml_surrogate.py``.
+      :caption: You can copy this file from ``examples/pytorch_surrogate_model/visualize_ml_surrogate_15_stage.py``.
 
-.. figure:: https://user-images.githubusercontent.com/10621396/289976300-6f861d19-5a5c-42eb-9435-9f57bd2010bf.png
-   :alt: Evolution of beam moments through 9 stage LPA via neural network surrogates.
+.. figure:: https://gist.githubusercontent.com/RTSandberg/cf3f6193b3da12e7fd815f69789fd6f2/raw/2308e412d7482d55811afa2e9a0c6fa97627fc2f/lpa_ml_surrogate_moments.png
+   :alt: Evolution of beam moments through 15 stage LPA via neural network surrogates.
 
-   Evolution of electron beam moments through 9 stages of LPAs (via neural network surrogates).
+   Evolution of electron beam moments through 15 stages of LPAs (via neural network surrogates).
 
-.. figure:: https://user-images.githubusercontent.com/10621396/289956805-49e0a94a-454f-4b48-b448-7ac772edf95a.png
+.. figure:: https://gist.githubusercontent.com/RTSandberg/cf3f6193b3da12e7fd815f69789fd6f2/raw/2308e412d7482d55811afa2e9a0c6fa97627fc2f/initial_phase_spaces.png
    :alt: Initial phase space projections
 
-   Initial phase space projections going into 9 stage LPA (via neural network surrogates) simulation. Top row: spatial projections, middle row: momentum projections, bottom row: phase spaces.
+   Initial phase space projections going into 15 stage LPA (via neural network surrogates) simulation. Top row: spatial projections, middle row: momentum projections, bottom row: phase spaces.
 
-.. figure:: https://user-images.githubusercontent.com/10621396/289975961-7d389864-9106-4446-8556-b0ea4bb28145.png
-   :alt: Final phase space projections after 9 stage LPA (via neural network surrogates) simulation
+.. figure:: https://gist.githubusercontent.com/RTSandberg/cf3f6193b3da12e7fd815f69789fd6f2/raw/2308e412d7482d55811afa2e9a0c6fa97627fc2f/stage_14_phase_spaces.png
+   :alt: Final phase space projections after 15 stage LPA (via neural network surrogates) simulation
 
-   Final phase space projections after 9 stage LPA (via neural network surrogates) simulation. Top row: spatial projections, middle row: momentum projections, bottom row: phase spaces.
+   Final phase space projections after 15 stage LPA (via neural network surrogates) simulation. Top row: spatial projections, middle row: momentum projections, bottom row: phase spaces.

examples/pytorch_surrogate_model/analyze_ml_surrogate.py → examples/pytorch_surrogate_model/analyze_ml_surrogate_15_stage.py

@@ -41,7 +41,7 @@ def get_moments(beam):
 final = series.iterations[last_step].particles["beam"].to_df()
 
 # compare number of particles
-num_particles = 10000
+num_particles = 100000
 assert num_particles == len(initial)
 assert num_particles == len(final)
 
@@ -59,11 +59,11 @@ def get_moments(beam):
 assert np.allclose(
     [sigx, sigy, sigt, emittance_x, emittance_y],
     [
-        7.488319e-07,
-        7.501963e-07,
-        9.996533e-08,
-        5.052374e-10,
-        5.130370e-10,
+        7.494325e-07,
+        7.478916e-07,
+        9.976192e-08,
+        5.070297e-10,
+        5.080007e-10,
     ],
     rtol=rtol,
     atol=atol,
@@ -88,12 +88,12 @@ def get_moments(beam):
 assert np.allclose(
     [sigx, sigy, sigt, emittance_x, emittance_y, emittance_t],
     [
-        3.062763e-07,
-        2.873031e-07,
-        1.021142e-07,
-        9.090898e-12,
-        9.579053e-12,
-        2.834852e-11,
+        1.590999e-07,
+        1.634865e-07,
+        1.030930e-07,
+        5.031797e-12,
+        5.242205e-12,
+        2.049623e-11,
     ],
     rtol=rtol,
     atol=atol,