Skip to content

Commit 7dcaf8f

Browse files
ax3lax3l
committed
Merge remote-tracking branch 'mainline/development' into add_IPAC_3DSC_benchmarks
2 parents 18e738b + 7329cad commit 7dcaf8f

File tree

9 files changed

+1061
-5
lines changed

9 files changed

+1061
-5
lines changed

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,8 @@ sudo apt-get install -y --no-install-recommends \
3636
rocm-dev \
3737
rocfft-dev \
3838
rocprim-dev \
39-
rocrand-dev
39+
rocrand-dev \
40+
hiprand-dev
4041

4142
# activate
4243
#

‎.pre-commit-config.yaml‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -73,7 +73,7 @@ repos:
7373
# Sorts Python imports according to PEP8
7474
# https://www.python.org/dev/peps/pep-0008/#imports
7575
- repo: https://github.com/pycqa/isort
76-
rev: 5.12.0
76+
rev: 5.13.2
7777
hooks:
7878
- id: isort
7979
name: isort (python)
@@ -91,15 +91,15 @@ repos:
9191

9292
# Python Formatting
9393
- repo: https://github.com/psf/black
94-
rev: 23.11.0 # Keep in sync with blacken-docs
94+
rev: 23.12.0 # Keep in sync with blacken-docs
9595
hooks:
9696
- id: black
9797
- repo: https://github.com/asottile/blacken-docs
9898
rev: 1.16.0
9999
hooks:
100100
- id: blacken-docs
101101
additional_dependencies:
102-
- black==23.11.0 # keep in sync with black hook
102+
- black==23.12.0 # keep in sync with black hook
103103
# TODO: black-jupyter
104104

105105
# Jupyter Notebooks: clean up all cell outputs

‎CMakeLists.txt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -331,7 +331,7 @@ file(TO_CMAKE_PATH "${ABS_INSTALL_LIB_DIR}" ABS_INSTALL_LIB_DIR)
331331

332332
install(CODE "file(CREATE_LINK
333333
$<TARGET_FILE_NAME:lib>
334-
\"${ABS_INSTALL_LIB_DIR}/libImpactX$<TARGET_FILE_SUFFIX:lib>\"
334+
\"${ABS_INSTALL_LIB_DIR}/libimpactx$<TARGET_FILE_SUFFIX:lib>\"
335335
COPY_ON_ERROR SYMBOLIC)")
336336

337337
# CMake package file for find_package(ImpactX::ImpactX) in depending projects

‎docs/source/usage/examples.rst‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ This section allows you to **download input files** that correspond to different
3131
examples/thin_dipole/README.rst
3232
examples/aperture/README.rst
3333
examples/epac2004_benchmarks/README.rst
34+
examples/pytorch_surrogate_model/README.rst
3435

3536

3637
Unit tests

‎examples/pytorch_surrogate_model/README.rst‎

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,81 @@
1+
.. _examples-ml-surrogate:
2+
3+
9 Stage Laser-Plasma Accelerator Surrogate
4+
==========================================
5+
6+
This example models an electron beam accelerated through nine stages of laser-plasma accelerators with ideal plasma lenses providing the focusing between stages.
7+
For more details, see:
8+
9+
10+
- Sandberg R T, Lehe R, Mitchell C E, Garten M, Qiang J, Vay J-L and Huebl A.
11+
**Synthesizing Particle-in-Cell Simulations Through Learning and GPU Computing for Hybrid Particle Accelerator Beamlines**.
12+
Proc. of Platform for Advanced Scientific Computing (PASC'24), *submitted*, 2024.
13+
- Sandberg R T, Lehe R, Mitchell C E, Garten M, Qiang J, Vay J-L and Huebl A.
14+
**Hybrid Beamline Element ML-Training for Surrogates in the ImpactX Beam-Dynamics Code**.
15+
14th International Particle Accelerator Conference (IPAC'23), WEPA101, 2023.
16+
`DOI:10.18429/JACoW-IPAC2023-WEPA101 <https://doi.org/10.18429/JACoW-IPAC2023-WEPA101>`__
17+
18+
A schematic with more information can be seen in the figure below:
19+
20+
.. figure:: https://user-images.githubusercontent.com/10621396/289956389-c7463b99-fb56-490a-8511-22c43f45cdf8.png
21+
:alt: [fig:lpa_schematic] Schematic of the 9 stages of laser-plasma accelerators.
22+
23+
[fig:lpa_schematic] Schematic of the 9 stages of laser-plasma accelerators.
24+
25+
The laser-plasma accelerator elements are modeled with neural network surrogates,
26+
previously trained and included in ``models``.
27+
The neural networks require normalized input data; the normalizations can be found in ``datasets``.
28+
29+
We use a 1 GeV electron beam with initial normalized rms emittance of 1 mm-mrad.
30+
31+
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.
32+
33+
Run
34+
---
35+
36+
This example can be **only** be run with **Python**:
37+
38+
* **Python** script: ``python3 run_ml_surrogate.py``
39+
40+
For `MPI-parallel <https://www.mpi-forum.org>`__ runs, prefix these lines with ``mpiexec -n 4 ...`` or ``srun -n 4 ...``, depending on the system.
41+
42+
.. literalinclude:: run_ml_surrogate.py
43+
:language: python
44+
:caption: You can copy this file from ``examples/pytorch_surrogate_model/run_ml_surrogate.py``.
45+
46+
Analyze
47+
-------
48+
49+
We run the following script to analyze correctness:
50+
51+
.. dropdown:: Script ``analyze_ml_surrogate.py``
52+
53+
.. literalinclude:: analyze_ml_surrogate.py
54+
:language: python
55+
:caption: You can copy this file from ``examples/pytorch_surrogate_model/analyze_ml_surrogate.py``.
56+
57+
Visualize
58+
---------
59+
60+
You can run the following script to visualize the beam evolution over time:
61+
62+
.. dropdown:: Script ``visualize_ml_surrogate.py``
63+
64+
.. literalinclude:: visualize_ml_surrogate.py
65+
:language: python
66+
:caption: You can copy this file from ``examples/pytorch_surrogate_model/visualize_ml_surrogate.py``.
67+
68+
.. figure:: https://user-images.githubusercontent.com/10621396/289976300-6f861d19-5a5c-42eb-9435-9f57bd2010bf.png
69+
:alt: Evolution of beam moments through 9 stage LPA via neural network surrogates.
70+
71+
Evolution of electron beam moments through 9 stages of LPAs (via neural network surrogates).
72+
73+
.. figure:: https://user-images.githubusercontent.com/10621396/289956805-49e0a94a-454f-4b48-b448-7ac772edf95a.png
74+
:alt: Initial phase space projections
75+
76+
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.
77+
78+
.. figure:: https://user-images.githubusercontent.com/10621396/289975961-7d389864-9106-4446-8556-b0ea4bb28145.png
79+
:alt: Final phase space projections after 9 stage LPA (via neural network surrogates) simulation
80+
81+
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‎

Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,106 @@
1+
#!/usr/bin/env python3
2+
#
3+
# Copyright 2022-2023 ImpactX contributors
4+
# Authors: Ryan Sandberg, Axel Huebl, Chad Mitchell
5+
# License: BSD-3-Clause-LBNL
6+
#
7+
# -*- coding: utf-8 -*-
8+
9+
import numpy as np
10+
import openpmd_api as io
11+
from scipy.stats import moment
12+
13+
14+
def get_moments(beam):
15+
"""Calculate standard deviations of beam position & momenta
16+
and emittance values
17+
18+
Returns
19+
-------
20+
sigx, sigy, sigt, emittance_x, emittance_y, emittance_t
21+
"""
22+
sigx = moment(beam["position_x"], moment=2) ** 0.5 # variance -> std dev.
23+
sigpx = moment(beam["momentum_x"], moment=2) ** 0.5
24+
sigy = moment(beam["position_y"], moment=2) ** 0.5
25+
sigpy = moment(beam["momentum_y"], moment=2) ** 0.5
26+
sigt = moment(beam["position_t"], moment=2) ** 0.5
27+
sigpt = moment(beam["momentum_t"], moment=2) ** 0.5
28+
29+
epstrms = beam.cov(ddof=0)
30+
emittance_x = (
31+
sigx**2 * sigpx**2 - epstrms["position_x"]["momentum_x"] ** 2
32+
) ** 0.5
33+
emittance_y = (
34+
sigy**2 * sigpy**2 - epstrms["position_y"]["momentum_y"] ** 2
35+
) ** 0.5
36+
emittance_t = (
37+
sigt**2 * sigpt**2 - epstrms["position_t"]["momentum_t"] ** 2
38+
) ** 0.5
39+
40+
return (sigx, sigy, sigt, emittance_x, emittance_y, emittance_t)
41+
42+
43+
# initial/final beam
44+
series = io.Series("diags/openPMD/monitor.bp", io.Access.read_only)
45+
last_step = list(series.iterations)[-1]
46+
initial = series.iterations[1].particles["beam"].to_df()
47+
final = series.iterations[last_step].particles["beam"].to_df()
48+
49+
# compare number of particles
50+
num_particles = 10000
51+
assert num_particles == len(initial)
52+
assert num_particles == len(final)
53+
54+
print("Initial Beam:")
55+
sigx, sigy, sigt, emittance_x, emittance_y, emittance_t = get_moments(initial)
56+
print(f" sigx={sigx:e} sigy={sigy:e} sigt={sigt:e}")
57+
print(
58+
f" emittance_x={emittance_x:e} emittance_y={emittance_y:e} emittance_t={emittance_t:e}"
59+
)
60+
61+
atol = 0.0 # ignored
62+
rtol = num_particles**-0.5 # from random sampling of a smooth distribution
63+
print(f" rtol={rtol} (ignored: atol~={atol})")
64+
65+
assert np.allclose(
66+
[sigx, sigy, sigt, emittance_x, emittance_y],
67+
[
68+
7.488319e-07,
69+
7.501963e-07,
70+
9.996533e-08,
71+
5.052374e-10,
72+
5.130370e-10,
73+
],
74+
rtol=rtol,
75+
atol=atol,
76+
)
77+
78+
atol = 1.0e-6
79+
print(f" atol~={atol}")
80+
assert np.allclose([emittance_t], [0.0], atol=atol)
81+
82+
print("")
83+
print("Final Beam:")
84+
sigx, sigy, sigt, emittance_x, emittance_y, emittance_t = get_moments(final)
85+
print(f" sigx={sigx:e} sigy={sigy:e} sigt={sigt:e}")
86+
print(
87+
f" emittance_x={emittance_x:e} emittance_y={emittance_y:e} emittance_t={emittance_t:e}"
88+
)
89+
90+
atol = 0.0 # ignored
91+
rtol = num_particles**-0.5 # from random sampling of a smooth distribution
92+
print(f" rtol={rtol} (ignored: atol~={atol})")
93+
94+
assert np.allclose(
95+
[sigx, sigy, sigt, emittance_x, emittance_y, emittance_t],
96+
[
97+
3.062763e-07,
98+
2.873031e-07,
99+
1.021142e-07,
100+
9.090898e-12,
101+
9.579053e-12,
102+
2.834852e-11,
103+
],
104+
rtol=rtol,
105+
atol=atol,
106+
)

0 commit comments

Comments
 (0)