Feature: Initialize beam from arrays (#560)

* start beam from array example
* load from array, visualize, and test
* include figure in documentation
* clean up utility function names and run script
* review suggestions
* add test to CI, make sure it runs
* update CI test folder
* Monitor: Use HDF5 for CI
* address review comments
* Minor Updates to Docs & Strings

Co-authored-by: Axel Huebl <axel.huebl@plasma.ninja>

docs/source/usage/examples.rst

@@ -66,6 +66,8 @@ Virtual Test Stands
    examples/positron_channel/README.rst
    examples/pytorch_surrogate_model/README.rst
    examples/apochromatic/README.rst
+   examples/fodo_tune/README.rst
+   examples/initialize_from_array/README.rst
 
 
 Unit tests

examples/CMakeLists.txt

@@ -906,3 +906,14 @@ add_impactx_test(triplet.py
     examples/optimize_triplet/plot_triplet.py
 )
 label_impactx_test(triplet.py slow)
+
+# Initialize a Beam from Arrays ##############################################
+#
+# w/o space charge
+add_impactx_test(initialize_from_array.py
+    examples/initialize_from_array/run_from_array.py
+      OFF  # ImpactX MPI-parallel
+      ON   # ImpactX Python interface
+    examples/initialize_from_array/analyze_from_array.py
+    examples/initialize_from_array/visualize_from_array.py
+)

examples/initialize_from_array/README.rst

@@ -0,0 +1,90 @@
+.. _examples-from-array:
+
+Initialize a Beam from Arrays
+=============================
+
+
+.. note::
+
+   This example is an early draft of this workflow and will be simplified in future versions of ImpactX, adding direct support for NumPy and CuPy arrays.
+
+This example demonstrates how a beam can be initalized in ImpactX from array-like structures.
+This allows various applications of interest,
+such as using a beam from a different simulation,
+initializing a beam from file,
+or creating a custom distribution.
+This example includes a set of utilities for transforming the beam to the fixed-s coordinates of ImpactX.
+
+In this example, a custom beam is specified at fixed t, transformed to fixed s, and
+then loaded in ImpactX.
+The custom beam is a ring in x-y,
+with radius r=2 mm,
+radial width :math:`\sigma_r = 5\ \mu`m;
+Gaussian in :math:`p_x` and :math:`p_y` with momentum width :math:`\sigma_p=10`;
+and chirped in z-pz with bunch length :math:`\sigma_z=1` mm,
+mean energy about 10 GeV, 1% uncorrelated energy spread, and z-pz covariance of -0.18.
+
+In specifying the beam at fixed t and transforming to fixed s,
+it is assumed that the local and global coordinate frames align.
+That is, the beam transverse directions x and y are the global x and y directions
+and the beam z/t direction is the global z direction.
+The transformation utility function reproduces the t-to-s and s-to-t transformations
+done internally in ImpactX given this assumption that the beam and global coordinate systems align.
+These utility functions are provided in the following script:
+
+.. dropdown:: Script ``transformation_utilities.py``
+
+   .. literalinclude:: transformation_utilities.py
+      :language: python3
+      :caption: You can copy this file from ``examples/initialize_from_array/transformation_utilities.py``.
+
+
+Run
+---
+
+This example can **only** be run with **Python**:
+
+* **Python** script: ``python3 run_from_array.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_from_array.py
+    :language: python3
+    :caption: You can copy this file from ``examples/initialize_from_array/run_from_array.py``.
+
+Analyze
+-------
+
+We run the following script to analyze correctness:
+
+.. dropdown:: Script ``analyze_from_array.py``
+
+   .. literalinclude:: analyze_from_array.py
+      :language: python3
+      :caption: You can copy this file from ``examples/initialize_from_array/analyze_from_array.py``.
+
+This uses the following utilities to read ImpactX output:
+
+.. dropdown:: Script ``impactx_read_utilities.py``
+
+   .. literalinclude:: impactx_read_utilities.py
+      :language: python3
+      :caption: You can copy this file from ``examples/initialize_from_array/impactx_read_utilities.py``.
+
+Visualize
+---------
+
+We run the following script to visualize the ImpactX output and confirm the beam is properly initialized:
+
+.. dropdown:: Script ``visualize_from_array.py``
+
+   .. literalinclude:: visualize_from_array.py
+      :language: python3
+      :caption: You can copy this file from ``examples/initialize_from_array/visualize_from_array.py``.
+
+The resulting phase space snapshots are shown in the following figure:
+
+.. figure:: https://gist.githubusercontent.com/RTSandberg/613d16def3d025f9415d348a58bddba6/raw/f8d0dfe47f2a75063845b748df4cb3fb9f3b38bb/phase_space.png
+   :alt: [fig:custom_beam] Phase space snapshots of custom beam from arrays, showing the ring in x-y, Gaussian in px-py, and linear correlation in z-pz.
+
+   [fig:custom_beam] Phase space snapshots of custom beam from arrays, showing the ring in x-y, Gaussian in px-py, and linear correlation in z-pz.

examples/initialize_from_array/analyze_from_array.py

@@ -0,0 +1,65 @@
+#!/usr/bin/env python3
+#
+# Copyright 2022-2023 ImpactX contributors
+# Authors: Ryan Sandberg
+# License: BSD-3-Clause-LBNL
+#
+
+import impactx_read_utilities as ix_read
+import numpy as np
+import openpmd_api as io
+import transformation_utilities as coord
+
+print("Initial Beam:")
+
+beam_series = io.Series("diags/openPMD/monitor.h5", io.Access.read_only)
+ref_series = ix_read.read_time_series("diags/ref_particle.*")
+
+num_particles = int(1e5)
+millimeter = 1.0e3
+micron = 1.0e6
+step = 1
+beam_at_step = beam_series.iterations[step].particles["beam"].to_df()
+ref_part_step = ref_series.loc[step]
+ref_part = myref = coord.MyRefPart(
+    ref_part_step["x"],
+    ref_part_step["y"],
+    ref_part_step["z"],
+    ref_part_step["px"],
+    ref_part_step["py"],
+    ref_part_step["pz"],
+    ref_part_step["pt"],
+)
+dx_s = beam_at_step["position_x"]
+dy_s = beam_at_step["position_y"]
+dt = beam_at_step["position_t"]
+dpx_s = beam_at_step["momentum_x"]
+dpy_s = beam_at_step["momentum_y"]
+dpt = beam_at_step["momentum_t"]
+
+dx_t, dy_t, dz, dpx_t, dpy_t, dpz = coord.to_t_from_s(
+    ref_part, dx_s, dy_s, dt, dpx_s, dpy_s, dpt
+)
+x1, y1, z1, px1, py1, pz1 = coord.to_global_t_from_ref_part_t(
+    ref_part, dx_t, dy_t, dz, dpx_t, dpy_t, dpz
+)
+
+sigx = x1.std()
+sigy = y1.std()
+sigz = z1.std()
+sigpx = px1.std()
+sigpy = py1.std()
+sigpz = pz1.std()
+print(f"sig x={sigx:.2e}, sig y={sigy:.2e}, sig z={sigz:.2e}")
+print(f"sig px={sigpx:.2e}, sig py={sigpy:.2e}, sig pz={sigpz:.2e}")
+
+atol = 0.0  # ignored
+rtol = 2 * num_particles**-0.5  # from random sampling of a smooth distribution
+print(f"  rtol={rtol} (ignored: atol~={atol})")
+
+assert np.allclose(
+    [sigx, sigy, sigz, sigpx, sigpy, sigpz],
+    [1.46e-3, 1.46e-3, 1.0e-3, 10, 10, 200],
+    rtol=rtol,
+    atol=atol,
+)

examples/initialize_from_array/impactx_read_utilities.py

@@ -0,0 +1,28 @@
+import glob
+import re
+
+import pandas as pd
+
+
+def read_file(file_pattern):
+    for filename in glob.glob(file_pattern):
+        df = pd.read_csv(filename, delimiter=r"\s+")
+        if "step" not in df.columns:
+            step = int(re.findall(r"[0-9]+", filename)[0])
+            df["step"] = step
+        yield df
+
+
+def read_time_series(file_pattern):
+    """Read in all CSV files from each MPI rank (and potentially OpenMP
+    thread). Concatenate into one Pandas dataframe.
+
+    Returns
+    -------
+    pandas.DataFrame
+    """
+    return pd.concat(
+        read_file(file_pattern),
+        axis=0,
+        ignore_index=True,
+    )  # .set_index('id')

examples/initialize_from_array/run_from_array.py

@@ -0,0 +1,108 @@
+#!/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 amrex.space3d as amr
+import numpy as np
+import transformation_utilities as pycoord
+from impactx import Config, ImpactX, elements
+
+################
+
+N_part = int(1e5)
+beam_radius = 2e-3
+sigr = 500e-6
+sigpx = 10
+sigpy = 10
+px = np.random.normal(0, sigpx, N_part)
+py = np.random.normal(0, sigpy, N_part)
+theta = 2 * np.pi * np.random.rand(N_part)
+r = abs(np.random.normal(beam_radius, sigr, N_part))
+x = r * np.cos(theta)
+y = r * np.sin(theta)
+z_mean = 0
+pz_mean = 2e4
+z_std = 1e-3
+pz_std = 2e2
+zpz_std = -0.18
+zpz_cov_list = [[z_std**2, zpz_std], [zpz_std, pz_std**2]]
+z, pz = np.random.default_rng().multivariate_normal([0, 0], zpz_cov_list, N_part).T
+pz += pz_mean
+
+sim = ImpactX()
+
+# set numerical parameters and IO control
+sim.particle_shape = 2  # B-spline order
+sim.space_charge = False
+# sim.diagnostics = False  # benchmarking
+sim.slice_step_diagnostics = True
+
+# domain decomposition & space charge mesh
+sim.init_grids()
+
+energy_gamma = np.sqrt(1 + pz_mean**2)
+energy_MeV = 0.510998950 * energy_gamma  # reference energy
+bunch_charge_C = 10.0e-15  # used with space charge
+
+#   reference particle
+ref = sim.particle_container().ref_particle()
+ref.set_charge_qe(-1.0).set_mass_MeV(0.510998950).set_kin_energy_MeV(energy_MeV)
+qm_eev = -1.0 / 0.510998950 / 1e6  # electron charge/mass in e / eV
+ref.z = 0
+
+pc = sim.particle_container()
+
+dx, dy, dz, dpx, dpy, dpz = pycoord.to_ref_part_t_from_global_t(
+    ref, x, y, z, px, py, pz
+)
+dx, dy, dt, dpx, dpy, dpt = pycoord.to_s_from_t(ref, dx, dy, dz, dpx, dpy, dpz)
+
+if not Config.have_gpu:  # initialize using cpu-based PODVectors
+    dx_podv = amr.PODVector_real_std()
+    dy_podv = amr.PODVector_real_std()
+    dt_podv = amr.PODVector_real_std()
+    dpx_podv = amr.PODVector_real_std()
+    dpy_podv = amr.PODVector_real_std()
+    dpt_podv = amr.PODVector_real_std()
+else:  # initialize on device using arena/gpu-based PODVectors
+    dx_podv = amr.PODVector_real_arena()
+    dy_podv = amr.PODVector_real_arena()
+    dt_podv = amr.PODVector_real_arena()
+    dpx_podv = amr.PODVector_real_arena()
+    dpy_podv = amr.PODVector_real_arena()
+    dpt_podv = amr.PODVector_real_arena()
+
+for p_dx in dx:
+    dx_podv.push_back(p_dx)
+for p_dy in dy:
+    dy_podv.push_back(p_dy)
+for p_dt in dt:
+    dt_podv.push_back(p_dt)
+for p_dpx in dpx:
+    dpx_podv.push_back(p_dpx)
+for p_dpy in dpy:
+    dpy_podv.push_back(p_dpy)
+for p_dpt in dpt:
+    dpt_podv.push_back(p_dpt)
+
+pc.add_n_particles(
+    dx_podv, dy_podv, dt_podv, dpx_podv, dpy_podv, dpt_podv, qm_eev, bunch_charge_C
+)
+
+monitor = elements.BeamMonitor("monitor", backend="h5")
+sim.lattice.extend(
+    [
+        monitor,
+        elements.Drift(ds=0.01),
+        monitor,
+    ]
+)
+
+sim.evolve()
+
+# clean shutdown
+sim.finalize()