Merge remote-tracking branch 'mainline/development' into add_IPAC_3DS…

…C_benchmarks

.github/workflows/dependencies/hip-openmpi.sh

@@ -36,7 +36,8 @@ sudo apt-get install -y --no-install-recommends \
     rocm-dev        \
     rocfft-dev      \
     rocprim-dev     \
-    rocrand-dev
+    rocrand-dev     \
+    hiprand-dev
 
 # activate
 #

.pre-commit-config.yaml

@@ -73,7 +73,7 @@ repos:
 # Sorts Python imports according to PEP8
 # https://www.python.org/dev/peps/pep-0008/#imports
 - repo: https://github.com/pycqa/isort
-  rev: 5.12.0
+  rev: 5.13.2
   hooks:
   - id: isort
     name: isort (python)
@@ -91,15 +91,15 @@ repos:
 
 # Python Formatting
 - repo: https://github.com/psf/black
-  rev: 23.11.0 # Keep in sync with blacken-docs
+  rev: 23.12.0 # Keep in sync with blacken-docs
   hooks:
   - id: black
 - repo: https://github.com/asottile/blacken-docs
   rev: 1.16.0
   hooks:
   - id: blacken-docs
     additional_dependencies:
-    - black==23.11.0 # keep in sync with black hook
+    - black==23.12.0 # keep in sync with black hook
   # TODO: black-jupyter
 
 # Jupyter Notebooks: clean up all cell outputs

CMakeLists.txt

@@ -331,7 +331,7 @@ file(TO_CMAKE_PATH "${ABS_INSTALL_LIB_DIR}" ABS_INSTALL_LIB_DIR)
 
 install(CODE "file(CREATE_LINK
     $<TARGET_FILE_NAME:lib>
-    \"${ABS_INSTALL_LIB_DIR}/libImpactX$<TARGET_FILE_SUFFIX:lib>\"
+    \"${ABS_INSTALL_LIB_DIR}/libimpactx$<TARGET_FILE_SUFFIX:lib>\"
     COPY_ON_ERROR SYMBOLIC)")
 
 # CMake package file for find_package(ImpactX::ImpactX) in depending projects

docs/source/usage/examples.rst

@@ -31,6 +31,7 @@ This section allows you to **download input files** that correspond to different
    examples/thin_dipole/README.rst
    examples/aperture/README.rst
    examples/epac2004_benchmarks/README.rst
+   examples/pytorch_surrogate_model/README.rst
 
 
 Unit tests

examples/pytorch_surrogate_model/README.rst

@@ -0,0 +1,81 @@
+.. _examples-ml-surrogate:
+
+9 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.
+For more details, see:
+
+
+- Sandberg R T, Lehe R, Mitchell C E, Garten M, 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.
+- 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.
+
+   [fig:lpa_schematic] Schematic of the 9 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``.
+
+We use a 1 GeV electron beam with initial normalized rms emittance of 1 mm-mrad.
+
+In this test, the initial and final values of :math:`\sigma_x`, :math:`\sigma_y`, :math:`\sigma_t`, :math:`\epsilon_x`, :math:`\epsilon_y`, and :math:`\epsilon_t` must agree with nominal values.
+
+Run
+---
+
+This example can be **only** be run with **Python**:
+
+* **Python** script: ``python3 run_ml_surrogate.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
+   :language: python
+   :caption: You can copy this file from ``examples/pytorch_surrogate_model/run_ml_surrogate.py``.
+
+Analyze
+-------
+
+We run the following script to analyze correctness:
+
+.. dropdown:: Script ``analyze_ml_surrogate.py``
+
+   .. literalinclude:: analyze_ml_surrogate.py
+      :language: python
+      :caption: You can copy this file from ``examples/pytorch_surrogate_model/analyze_ml_surrogate.py``.
+
+Visualize
+---------
+
+You can run the following script to visualize the beam evolution over time:
+
+.. dropdown:: Script ``visualize_ml_surrogate.py``
+
+   .. literalinclude:: visualize_ml_surrogate.py
+      :language: python
+      :caption: You can copy this file from ``examples/pytorch_surrogate_model/visualize_ml_surrogate.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.
+
+   Evolution of electron beam moments through 9 stages of LPAs (via neural network surrogates).
+
+.. figure:: https://user-images.githubusercontent.com/10621396/289956805-49e0a94a-454f-4b48-b448-7ac772edf95a.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.
+
+.. 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
+
+   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.

examples/pytorch_surrogate_model/analyze_ml_surrogate.py

@@ -0,0 +1,106 @@
+#!/usr/bin/env python3
+#
+# Copyright 2022-2023 ImpactX contributors
+# Authors: Ryan Sandberg, Axel Huebl, Chad Mitchell
+# License: BSD-3-Clause-LBNL
+#
+# -*- coding: utf-8 -*-
+
+import numpy as np
+import openpmd_api as io
+from scipy.stats import moment
+
+
+def get_moments(beam):
+    """Calculate standard deviations of beam position & momenta
+    and emittance values
+
+    Returns
+    -------
+    sigx, sigy, sigt, emittance_x, emittance_y, emittance_t
+    """
+    sigx = moment(beam["position_x"], moment=2) ** 0.5  # variance -> std dev.
+    sigpx = moment(beam["momentum_x"], moment=2) ** 0.5
+    sigy = moment(beam["position_y"], moment=2) ** 0.5
+    sigpy = moment(beam["momentum_y"], moment=2) ** 0.5
+    sigt = moment(beam["position_t"], moment=2) ** 0.5
+    sigpt = moment(beam["momentum_t"], moment=2) ** 0.5
+
+    epstrms = beam.cov(ddof=0)
+    emittance_x = (
+        sigx**2 * sigpx**2 - epstrms["position_x"]["momentum_x"] ** 2
+    ) ** 0.5
+    emittance_y = (
+        sigy**2 * sigpy**2 - epstrms["position_y"]["momentum_y"] ** 2
+    ) ** 0.5
+    emittance_t = (
+        sigt**2 * sigpt**2 - epstrms["position_t"]["momentum_t"] ** 2
+    ) ** 0.5
+
+    return (sigx, sigy, sigt, emittance_x, emittance_y, emittance_t)
+
+
+# initial/final beam
+series = io.Series("diags/openPMD/monitor.bp", io.Access.read_only)
+last_step = list(series.iterations)[-1]
+initial = series.iterations[1].particles["beam"].to_df()
+final = series.iterations[last_step].particles["beam"].to_df()
+
+# compare number of particles
+num_particles = 10000
+assert num_particles == len(initial)
+assert num_particles == len(final)
+
+print("Initial Beam:")
+sigx, sigy, sigt, emittance_x, emittance_y, emittance_t = get_moments(initial)
+print(f"  sigx={sigx:e} sigy={sigy:e} sigt={sigt:e}")
+print(
+    f"  emittance_x={emittance_x:e} emittance_y={emittance_y:e} emittance_t={emittance_t:e}"
+)
+
+atol = 0.0  # ignored
+rtol = num_particles**-0.5  # from random sampling of a smooth distribution
+print(f"  rtol={rtol} (ignored: atol~={atol})")
+
+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,
+    ],
+    rtol=rtol,
+    atol=atol,
+)
+
+atol = 1.0e-6
+print(f"  atol~={atol}")
+assert np.allclose([emittance_t], [0.0], atol=atol)
+
+print("")
+print("Final Beam:")
+sigx, sigy, sigt, emittance_x, emittance_y, emittance_t = get_moments(final)
+print(f"  sigx={sigx:e} sigy={sigy:e} sigt={sigt:e}")
+print(
+    f"  emittance_x={emittance_x:e} emittance_y={emittance_y:e} emittance_t={emittance_t:e}"
+)
+
+atol = 0.0  # ignored
+rtol = num_particles**-0.5  # from random sampling of a smooth distribution
+print(f"  rtol={rtol} (ignored: atol~={atol})")
+
+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,
+    ],
+    rtol=rtol,
+    atol=atol,
+)