merge smoothen buffer

.github/workflows/cuda.yml

@@ -115,7 +115,7 @@ jobs:
         which nvcc || echo "nvcc not in PATH!"
 
         git clone https://github.com/AMReX-Codes/amrex.git ../amrex
-        cd ../amrex && git checkout --detach 24.03 && cd -
+        cd ../amrex && git checkout --detach 155f1f4c3f4fa03528d38d2769beb4b35b38042c && cd -
         make COMP=gcc QED=FALSE USE_MPI=TRUE USE_GPU=TRUE USE_OMP=FALSE USE_PSATD=TRUE USE_CCACHE=TRUE -j 4
 
         ccache -s

.pre-commit-config.yaml

@@ -80,6 +80,7 @@ repos:
   hooks:
   - id: isort
     name: isort (python)
+    args: ['--profile black']
 
 # Python: Flake8 (checks only, does this support auto-fixes?)
 #- repo: https://github.com/PyCQA/flake8

.zenodo.json

@@ -20,6 +20,11 @@
         "name": "Andriyash, Igor",
         "orcid": "0000-0003-0313-4496"
       },
+      {
+        "affiliation": "Lawrence Livermore National Laboratory",
+        "name": "Angus, Justin Ray",
+        "orcid": "0000-0003-1474-0002"
+      },
       {
         "affiliation": "Lawrence Berkeley National Laboratory",
         "name": "Belkin, Daniel",
@@ -155,19 +160,29 @@
       },
       {
         "affiliation": "Lawrence Berkeley National Laboratory",
-        "name": "Sandberg, Ryan T.",
+        "name": "Sandberg, Ryan Thor",
         "orcid": "0000-0001-7680-8733"
       },
       {
         "affiliation": "Modern Electron",
         "name": "Scherpelz, Peter",
         "orcid": "0000-0001-8185-3387"
       },
+      {
+        "affiliation": "Laboratory for Laser Energetics, University of Rochester",
+        "name": "Weichman, Kale",
+        "orcid": "0000-0002-3487-7922"
+      },
       {
         "affiliation": "Lawrence Berkeley National Laboratory",
         "name": "Yang, Eloise",
         "orcid": "0000-0002-9319-4216"
       },
+      {
+        "affiliation": "LIDYL, CEA-Universit\u00e9 Paris-Saclay, CEA Saclay",
+        "name": "Zaim, Ne\u00efl",
+        "orcid": "0000-0003-0313-4496"
+      },
       {
         "affiliation": "Lawrence Berkeley National Laboratory",
         "name": "Zhang, Weiqun",
@@ -187,11 +202,6 @@
         "affiliation": "Lawrence Berkeley National Laboratory",
         "name": "Zoni, Edoardo",
         "orcid": "0000-0001-5662-4646"
-      },
-      {
-        "affiliation": "LIDYL, CEA-Universit\u00e9 Paris-Saclay, CEA Saclay",
-        "name": "Zaim, Ne\u00efl",
-        "orcid": "0000-0003-0313-4496"
       }
     ],
     "contributors": [

Docs/source/acknowledge_us.rst

@@ -53,6 +53,11 @@ Prior WarpX references
 
 If your project uses a specific algorithm or component, please consider citing the respective publications in addition.
 
+- 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.
+  `preprint <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.

Docs/source/conf.py

@@ -31,8 +31,8 @@
 import urllib.request
 
 import pybtex.plugin
-from pybtex.style.formatting.unsrt import Style as UnsrtStyle
 import sphinx_rtd_theme
+from pybtex.style.formatting.unsrt import Style as UnsrtStyle
 
 sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), '../../Regression/Checksum'))
 

Docs/source/highlights.rst

@@ -24,6 +24,11 @@ Scientific works in laser-plasma and beam-plasma acceleration.
    Phys. Rev. Research **5**, 033112, 2023
    `DOI:10.1103/PhysRevResearch.5.033112 <https://doi.org/10.1103/PhysRevResearch.5.033112>`__
 
+#. 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.
+  `preprint <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.
@@ -96,6 +101,11 @@ Particle Accelerator & Beam Physics
 
 Scientific works in particle and beam modeling.
 
+#. 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.
+  `preprint <http://arxiv.org/abs/2402.17248>__`
+
 #. Sandberg R T, Lehe R, Mitchell C E, Garten M, Qiang J, Vay J-L, Huebl A.
    **Hybrid Beamline Element ML-Training for Surrogates in the ImpactX Beam-Dynamics Code**.
    14th International Particle Accelerator Conference (IPAC'23), WEPA101, *in print*, 2023.
@@ -128,6 +138,7 @@ Microelectronics
    **Characterization of Transmission Lines in Microelectronic Circuits Using the ARTEMIS Solver**.
    IEEE Journal on Multiscale and Multiphysics Computational Techniques, vol. 8, pp. 31-39, 2023.
    `DOI:10.1109/JMMCT.2022.3228281 <https://doi.org/10.1109/JMMCT.2022.3228281>`__
+
 #. Kumar P, Nonaka A, Jambunathan R, Pahwa G and Salahuddin S, Yao Z.
    **FerroX: A GPU-accelerated, 3D Phase-Field Simulation Framework for Modeling Ferroelectric Devices**.
    arXiv preprint, 2022.

Docs/source/install/dependencies.rst

@@ -63,8 +63,8 @@ Conda (Linux/macOS/Windows)
 
    .. code-block:: bash
 
-      conda update -n base conda
-      conda install -n base conda-libmamba-solver
+      conda update -y -n base conda
+      conda install -y -n base conda-libmamba-solver
       conda config --set solver libmamba
 
    We recommend to deactivate that conda self-activates its ``base`` environment.

Docs/source/install/users.rst

@@ -45,12 +45,12 @@ A package for WarpX is available via the `Conda <https://conda.io>`_ package man
 
 .. tip::
 
-   We recommend to configure your conda to use the faster `libmamba` `dependency solver <https://www.anaconda.com/blog/a-faster-conda-for-a-growing-community>`__.
+   We recommend to configure your conda to use the faster ``libmamba`` `dependency solver <https://www.anaconda.com/blog/a-faster-conda-for-a-growing-community>`__.
 
    .. code-block:: bash
 
-      conda update -n base conda
-      conda install -n base conda-libmamba-solver
+      conda update -y -n base conda
+      conda install -y -n base conda-libmamba-solver
       conda config --set solver libmamba
 
    We recommend to deactivate that conda self-activates its ``base`` environment.

Docs/source/refs.bib

@@ -207,13 +207,15 @@ @article{Roedel2010
 
 @misc{SandbergPASC24,
 address = {Zuerich, Switzerland},
-author = {Ryan Sandberg and Remi Lehe and Chad E Mitchell and Marco Garten and Ji Qiang and Jean-Luc Vay and Axel Huebl},
+author = {Ryan Sandberg and Remi Lehe and Chad E Mitchell and Marco Garten and Andrew Myers and Ji Qiang and Jean-Luc Vay and Axel Huebl},
 booktitle = {Proc. of PASC24},
-note = {submitted},
+note = {accepted},
 series = {PASC'24 - Platform for Advanced Scientific Computing},
 title = {{Synthesizing Particle-in-Cell Simulations Through Learning and GPU Computing for Hybrid Particle Accelerator Beamlines}},
 venue = {Zuerich, Switzerland},
-year = {2024}
+year = {2024},
+doi = {10.48550/arXiv.2402.17248},
+url = {https://arxiv.org/abs/2402.17248}
 }
 
 @inproceedings{SandbergIPAC23,

Docs/source/usage/parameters.rst

@@ -395,12 +395,18 @@ Domain Boundary Conditions
     * ``damped``: This is the recommended option in the moving direction when using the spectral solver with moving window (currently only supported along z). This boundary condition applies a damping factor to the electric and magnetic fields in the outer half of the guard cells, using a sine squared profile. As the spectral solver is by nature periodic, the damping prevents fields from wrapping around to the other end of the domain when the periodicity is not desired. This boundary condition is only valid when using the spectral solver.
 
     * ``pec``: This option can be used to set a Perfect Electric Conductor at the simulation boundary. Please see the :ref:`PEC theory section <theory-bc-pec>` for more details. Note that PEC boundary is invalid at `r=0` for the RZ solver. Please use ``none`` option. This boundary condition does not work with the spectral solver.
-      If an electrostatic field solve is used the boundary potentials can also be set through ``boundary.potential_lo_x/y/z`` and ``boundary.potential_hi_x/y/z`` (default `0`).
 
     * ``none``: No boundary condition is applied to the fields with the electromagnetic solver. This option must be used for the RZ-solver at `r=0`.
 
     * ``neumann``: For the electrostatic solver, a Neumann boundary condition (with gradient of the potential equal to 0) will be applied on the specified boundary.
 
+* ``boundary.potential_lo_x/y/z`` and ``boundary.potential_hi_x/y/z`` (default `0`)
+    Gives the value of the electric potential at the boundaries, for ``pec`` boundaries. With electrostatic solvers
+    (i.e., with ``warpx.do_electrostatic = ...``), this is used in order to compute the potential
+    in the simulation volume at each timestep. When using other solvers (e.g. Maxwell solver),
+    setting these variables will trigger an electrostatic solve at ``t=0``, to compute the initial
+    electric field produced by the boundaries.
+
 * ``boundary.particle_lo`` and ``boundary.particle_hi`` (`2 strings` for 2D, `3 strings` for 3D, `absorbing` by default)
     Options are:
 
@@ -473,9 +479,12 @@ Embedded Boundary Conditions
     the interior of the embeddded boundary is where the function value is positive.
 
 * ``warpx.eb_potential(x,y,z,t)`` (`string`)
-    Only used when ``warpx.do_electrostatic=labframe``. Gives the value of
-    the electric potential at the surface of the embedded boundary,
-    as a function of  `x`, `y`, `z` and time. This function is also evaluated
+    Gives the value of the electric potential at the surface of the embedded boundary,
+    as a function of  `x`, `y`, `z` and `t`. With electrostatic solvers (i.e., with
+    ``warpx.do_electrostatic = ...``), this is used in order to compute the potential
+    in the simulation volume at each timestep. When using other solvers (e.g. Maxwell solver),
+    setting this variable will trigger an electrostatic solve at ``t=0``, to compute the initial
+    electric field produced by the boundaries. Note that this function is also evaluated
     inside the embedded boundary. For this reason, it is important to define
     this function in such a way that it is constant inside the embedded boundary.
 

Docs/source/usage/workflows/ml_dataset_training.rst

@@ -14,42 +14,42 @@ For example, a simulation determined by the following input script
     .. literalinclude:: ml_materials/run_warpx_training.py
        :language: python
 
-In this section we walk through a workflow for data processing and model training.
+In this section we walk through a workflow for data processing and model training, using data from this input script as an example.
+The simulation output is stored in an online `Zenodo archive <https://zenodo.org/records/10368972>`__, in the ``lab_particle_diags`` directory.
+In the example scripts provided here, the data is downloaded from the Zenodo archive, properly formatted, and used to train a neural network.
 This workflow was developed and first presented in :cite:t:`ml-SandbergIPAC23,ml-SandbergPASC24`.
-
-This assumes you have an up-to-date environment with PyTorch and openPMD.
+It assumes you have an up-to-date environment with PyTorch and openPMD.
 
 Data Cleaning
 -------------
 
-It is important to inspect the data for artifacts to
+It is important to inspect the data for artifacts, to
 check that input/output data make sense.
-If we plot the final phase space for beams 1-8,
-the particle data is distributed in a single blob,
-as shown by :numref:`fig_phase_space_beam_1` for beam 1.
-This is as we expect and what is optimal for training neural networks.
+If we plot the final phase space of the particle beam,
+shown in :numref:`fig_unclean_phase_space`.
+we see outlying particles.
+Looking closer at the z-pz space, we see that some particles were not trapped in the accelerating region of the wake and have much less energy than the rest of the beam.
+
+.. _fig_unclean_phase_space:
 
-.. _fig_phase_space_beam_1:
+.. figure:: https://gist.githubusercontent.com/RTSandberg/649a81cc0e7926684f103729483eff90/raw/095ac2daccbcf197fa4e18a8f8505711b27e807a/unclean_stage_0.png
+   :alt: Plot showing the final phase space projections of a particle beam through a laser-plasma acceleration element where some beam particles were not accelerated.
 
-.. figure:: https://user-images.githubusercontent.com/10621396/290010209-c55baf1c-dd98-4d56-a675-ad3729481eee.png
-   :alt: Plot showing the final phase space projections for beam 1 of the training data, for a surrogate to stage 1.
+   The final phase space projections of a particle beam through a laser-plasma acceleration element where some beam particles were not accelerated.
 
-   The final phase space projections for beam 1 of the training data, for a surrogate to stage 1.
+To assist our neural network in learning dynamics of interest, we filter out these particles.
+It is sufficient for our purposes to select particles that are not too far back, setting
+``particle_selection={'z':[0.280025, None]}``.
+After filtering, we can see in :numref:`fig_clean_phase_space` that the beam phase space projections are much cleaner -- this is the beam we want to train on.
 
-.. _fig_phase_space_beam_0:
+.. _fig_clean_phase_space:
 
-.. figure:: https://user-images.githubusercontent.com/10621396/290010282-40560ac4-8509-4599-82ca-167bb1739cff.png
-   :alt: Plot showing the final phase space projections for beam 0 of the training data, for a surrogate to stage 0.
+.. figure:: https://gist.githubusercontent.com/RTSandberg/649a81cc0e7926684f103729483eff90/raw/095ac2daccbcf197fa4e18a8f8505711b27e807a/clean_stage_0.png
+   :alt: Plot showing the final phase space projections of a particle beam through a laser-plasma acceleration element after filtering out outlying particles.
 
-   The final phase space projections for beam 0 of the training data, for a surrogate to stage 0
+   The final phase space projections of a particle beam through a laser-plasma acceleration element after filtering out outlying particles.
 
-On the other hand, the final phase space for beam 0, shown in :numref:`fig_phase_space_beam_1`,
-has a halo of outlying particles.
-Looking closer at the z-pz space, we see that some particles got caught in a decelerating
-region of the wake, have slipped back and are much slower than the rest of the beam.
-To assist our neural network in learning dynamics of interest, we filter out these particles.
-It is sufficient for our purposes to select particles that are not too far back, setting
-``particle_selection={'z':[0.28002, None]}``. Then a particle tracker is set up to make sure
+A particle tracker is set up to make sure
 we consistently filter out these particles from both the initial and final data.
 
 .. literalinclude:: ml_materials/create_dataset.py
@@ -58,6 +58,9 @@ we consistently filter out these particles from both the initial and final data.
    :start-after: # Manual: Particle tracking START
    :end-before: # Manual: Particle tracking END
 
+This data cleaning ensures that the particle data is distributed in a single blob,
+as is optimal for training neural networks.
+
 Create Normalized Dataset
 -------------------------
 
@@ -119,7 +122,12 @@ This data are converted to an :math:`N\times 6` numpy array and then to a PyTorc
 Save Normalizations and Normalized Data
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
-With the data properly normalized, it and the normalizations are saved to file for
+The data is split into training and testing subsets.
+We take most of the data (70%) for training, meaning that data is used to update
+the neural network parameters.
+The testing data is reserved to determine how well the neural network generalizes;
+that is, how well the neural network performs on data that wasn't used to update the neural network parameters.
+With the data split and properly normalized, it and the normalizations are saved to file for
 use in training and inference.
 
 .. literalinclude:: ml_materials/create_dataset.py
@@ -131,13 +139,22 @@ use in training and inference.
 Neural Network Structure
 ------------------------
 
-It was found in :cite:t:`ml-SandbergPASC24` that reasonable surrogate models are obtained with
-shallow feedforward neural networks consisting of fewer than 10 hidden layers and
-just under 1000 nodes per layer.
+It was found in :cite:t:`ml-SandbergPASC24` that a reasonable surrogate model is obtained with
+shallow feedforward neural networks consisting of about 5 hidden layers and 700-900 nodes per layer.
 The example shown here uses 3 hidden layers and 20 nodes per layer
 and is trained for 10 epochs.
 
+Some utility functions for creating neural networks are provided in the script below.
+These are mostly convenience wrappers and utilities for working with `PyTorch <https://pytorch.org/>`__ neural network objects.
+This script is imported in the training scripts shown later.
+
+.. dropdown:: Python neural network class definitions
+   :color: light
+   :icon: info
+   :animate: fade-in-slide-down
 
+    .. literalinclude:: ml_materials/neural_network_classes.py
+       :language: python3
 
 Train and Save Neural Network
 -----------------------------
@@ -188,8 +205,8 @@ which is later divided by the size of the dataset in the training loop.
    :start-after: # Manual: Test function START
    :end-before: # Manual: Test function END
 
-Train Loop
-^^^^^^^^^^
+Training Loop
+^^^^^^^^^^^^^
 
 The full training loop performs ``n_epochs`` number of iterations.
 At each iteration the training and testing functions are called,
@@ -228,22 +245,22 @@ When the test-loss starts to trend flat or even upward, the neural network is no
 
 .. _fig_train_test_loss:
 
-.. figure:: https://user-images.githubusercontent.com/10621396/290010428-f83725ab-a08f-494c-b075-314b0d26cb9a.png
+.. figure:: https://gist.githubusercontent.com/RTSandberg/649a81cc0e7926684f103729483eff90/raw/095ac2daccbcf197fa4e18a8f8505711b27e807a/beam_stage_0_training_testing_error.png
    :alt: Plot of training and testing loss curves versus number of training epochs.
 
    Training (in blue) and testing (in green) loss curves versus number of training epochs.
 
 .. _fig_train_evaluation:
 
-.. figure:: https://user-images.githubusercontent.com/10621396/290010486-4a3541e7-e0be-4cf1-b33b-57d5e5985196.png
+.. figure:: https://gist.githubusercontent.com/RTSandberg/649a81cc0e7926684f103729483eff90/raw/095ac2daccbcf197fa4e18a8f8505711b27e807a/beam_stage_0_model_evaluation.png
    :alt: Plot comparing model prediction with simulation output.
 
    A comparison of model prediction (yellow-red dots, colored by mean-squared error) with simulation output (black dots).
 
 A visual inspection of the model prediction can be seen in :numref:`fig_train_evaluation`.
 This plot compares the model prediction, with dots colored by mean-square error, on the testing data with the actual simulation output in black.
 The model obtained with the hyperparameters chosen here trains quickly but is not very accurate.
-A more accurate model is obtained with 5 hidden layers and 800 nodes per layer,
+A more accurate model is obtained with 5 hidden layers and 900 nodes per layer,
 as discussed in :cite:t:`ml-SandbergPASC24`.
 
 These figures can be generated with the following Python script.
@@ -261,7 +278,7 @@ Surrogate Usage in Accelerator Physics
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
 A neural network such as the one we trained here can be incorporated in other BLAST codes.
-`Consider the example using neural networks in ImpactX <https://impactx.readthedocs.io/en/latest/usage/examples/pytorch_surrogate_model/README.html>`__.
+Consider this `example using neural network surrogates of WarpX simulations in ImpactX <https://impactx.readthedocs.io/en/latest/usage/examples/pytorch_surrogate_model/README.html>`__.
 
 .. bibliography::
    :keyprefix: ml-