BLAST-ImpactX · cemitch99 · Feb 11, 2025 · Feb 11, 2025 · Feb 11, 2025 · Feb 11, 2025
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -66,7 +66,7 @@ repos:
 # Python: Ruff linter & formatter
 #   https://docs.astral.sh/ruff/
 - repo: https://github.com/astral-sh/ruff-pre-commit
-  rev: v0.9.4
+  rev: v0.9.6
   hooks:
     # Run the linter
     - id: ruff

diff --git a/docs/source/index.rst b/docs/source/index.rst
@@ -80,7 +80,7 @@ Usage
    usage/python
    usage/parameters
    usage/dashboard
-   usage/workflows
+   usage/howto
 
 Data Analysis
 -------------

diff --git a/docs/source/usage/examples.rst b/docs/source/usage/examples.rst
@@ -83,6 +83,7 @@ Channels & Rings
    :maxdepth: 1
 
    examples/fodo_channel/README.rst
+   examples/solenoid_restart/README.rst
 
 
 Lattice Design & Optimization

diff --git a/docs/source/usage/howto.rst b/docs/source/usage/howto.rst
@@ -0,0 +1,11 @@
+.. _usage-howto:
+
+How-To Guides
+=============
+
+This section collects typical user workflows and best practices for ImpactX.
+
+.. toctree::
+   :maxdepth: 1
+
+   howto/add_element
diff --git a/docs/source/usage/workflows/add_element.rst → docs/source/usage/howto/add_element.rst b/docs/source/usage/workflows/add_element.rst → docs/source/usage/howto/add_element.rst
@@ -1,4 +1,4 @@
-.. _usage-workflows-add-element:
+.. _usage-howto-add-element:
 
 Add New Beamline Elements
 =========================
@@ -10,27 +10,27 @@ The workflows described here apply both for thin kicks or thick elements.
 Thick elements can also use soft-edged fringe fields (see `existing soft-edged elements for implementation details <https://github.com/ECP-WarpX/impactx/tree/development/src/elements>`__).
 
 
-.. _usage-workflows-add-element-linmap:
+.. _usage-howto-add-element-linmap:
 
 Linear Map
 ----------
 
 A custom linear element can be provided by specifying the 6x6 linear transport matrix :math:`R` as an input.
-See the :ref:`example <examples-fodo-userdef>` for Python and inputs file syntax to specify a custom linear element.
+See the :ref:`FODO cell example <examples-fodo-userdef>` for a demonstration of the custom linear element.
 
 The matrix elements :math:`R(i,j)` are indexed beginning with 1, so that :math:`i,j=1,2,3,4,5,6`.
 
-The matrix :math:`R` multiplies the phase space vector :math:`(x,px,y,py,t,pt)`, where coordinates :math:`(x,y,t)` have units of m
-and momenta :math:`(px,py,pt)` are dimensionless.  So, for example, :math:`R(1,1)` is dimensionless, and :math:`R(1,2)` has units of m.
+The matrix :math:`R` multiplies the phase space vector :math:`(x,p_x,y,p_y,t,p_t)`, where coordinates :math:`(x,y,t)` have units of m and momenta :math:`(p_x,p_y,p_t)` are dimensionless.
+So, for example, :math:`R(1,1)` is dimensionless, and :math:`R(1,2)` has units of m.
 
 
 .. note::
 
    If a user-provided linear map is used, it is up to the user to ensure that the 6x6 transport matrix is symplectic.
-   If a more general form of user-defined transport is needed, the :ref:`Python Programmable Element <usage-workflows-add-element-python>` and the :ref:`C++ Element <usage-workflows-add-element-cxx>` provide a more general approach.
+   If a more general form of user-defined transport is needed, the :ref:`Python Programmable Element <usage-howto-add-element-python>` and the :ref:`C++ Element <usage-howto-add-element-cxx>` provide a more general approach.
 
 
-.. _usage-workflows-add-element-python:
+.. _usage-howto-add-element-python:
 
 Python Programmable Element
 ---------------------------
@@ -52,7 +52,7 @@ Detailed examples that show usage of the programmable element are:
 Detailed particle computing interfaces are presented in the `pyAMReX examples <https://pyamrex.readthedocs.io/en/latest/usage/compute.html#particles>`__.
 
 
-.. _usage-workflows-add-element-cxx:
+.. _usage-howto-add-element-cxx:
 
 C++ Element
 -----------
@@ -72,10 +72,12 @@ To simplify the logic, we use so-called `mixin classes <https://en.wikipedia.org
 
 After this brief boilerplate, our beamline elements implement these parts:
 
-* a constructor: storing element options
-* a single-particle operator: pushing the beam particles
-* a reference-particle operator: pushing the reference particle
-* a linear transport map for the envelope tracking mode
+#. a constructor: storing element options
+#. for particle tracking:
+
+   * a single-particle operator: pushing the beam particles
+   * a reference-particle operator: pushing the reference particle
+#. for envelope tracking: a linear transport map
 
 .. dropdown:: Example Element: Drift.H
    :color: light

diff --git a/docs/source/usage/parameters.rst b/docs/source/usage/parameters.rst
@@ -627,6 +627,20 @@ This requires these additional parameters:
 * ``<element_name>.nslice`` (``integer``) number of slices used for the application of space charge (default: ``1``)
 
 
+``source``
+^^^^^^^^^^^
+
+``source`` for a particle source.
+Typically at the beginning of a beam line.
+
+Currently, this only supports openPMD files from our ``beam_monitor``.
+
+* ``<element_name>.distribution`` (``string``)
+  Distribution type of particles in the source. currently, only ``"openPMD"`` is supported
+* ``<element_name>.openpmd_path`` (``string``)
+  path to the openPMD series
+
+
 ``tapered_pl``
 ^^^^^^^^^^^^^^
 

diff --git a/docs/source/usage/python.rst b/docs/source/usage/python.rst
@@ -738,6 +738,15 @@ This module provides elements for the accelerator lattice.
 
       Scale factor (in meters^(1/2)) of the IOTA nonlinear magnetic insert element used for computing H and I.
 
+.. py:class:: impactx.elements.Source(distribution, openpmd_path, name)
+
+   A particle source.
+   Currently, this only supports openPMD files from our :py:class:`impactx.elements.BeamMonitor`
+
+   :param distribution: Distribution type of particles in the source. currently, only ``"openPMD"`` is supported
+   :param openpmd_path: path to the openPMD series
+   :param name: an optional name for the element
+
 .. py:class:: impactx.elements.Programmable(ds=0.0, nslice=1, name=None)
 
    A programmable beam optics element.

diff --git a/docs/source/usage/workflows.rst b/docs/source/usage/workflows.rst
diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt
@@ -624,6 +624,28 @@ add_impactx_test(solenoid.MADX.py
 )
 
 
+# Ideal, Hard-Edge Solenoid (Restart + 270 degrees) ###########################
+#
+# w/o space charge
+add_impactx_test(solenoid.restart
+    examples/solenoid_restart/input_solenoid.in
+    OFF  # ImpactX MPI-parallel
+    examples/solenoid_restart/analysis_solenoid.py
+    OFF  # no plot script yet
+)
+set_property(TEST solenoid.restart.run APPEND PROPERTY DEPENDS "solenoid.run")
+
+add_impactx_test(solenoid.restart.py
+    examples/solenoid_restart/run_solenoid.py
+    OFF  # ImpactX MPI-parallel
+    examples/solenoid_restart/analysis_solenoid.py
+    OFF  # no plot script yet
+)
+if(ImpactX_PYTHON)
+    set_property(TEST solenoid.restart.py.run APPEND PROPERTY DEPENDS "solenoid.py.run")
+endif()
+
+
 # Pole-face rotation ##########################################################
 #
 # w/o space charge

diff --git a/examples/fodo_userdef/README.rst b/examples/fodo_userdef/README.rst
@@ -10,7 +10,7 @@ However, in the example here we define *additional user-defined, custom linear e
 
    Note that generally, if a user-provided linear map is used, the beam transport may not be symplectic.
    For an even more general, user-defined element, see :ref:`the FODO Cell example that uses a Programmable Element <examples-fodo-programmable>`.
-   For more details, see :ref:`this section <usage-workflows-add-element>`.
+   For more details, see :ref:`this section <usage-howto-add-element>`.
 
 The matched Twiss parameters at entry are:
 

diff --git a/examples/solenoid_restart/README.rst b/examples/solenoid_restart/README.rst
@@ -0,0 +1,50 @@
+.. _examples-solenoid-restart:
+
+Solenoid channel (Restart)
+==========================
+
+This is the same example as the :ref:`proton beam undergoing 90 deg X-Y rotation in an ideal solenoid channel <examples-solenoid>`, but it restarts the resulting beam and rotates it another 3 channel periods to the initial X-Y conditions.
+
+The solenoid magnetic field corresponds to B = 2 T.
+
+The second moments of the transverse particle distribution after the solenoid channel are rotated by 90 (start) + 270 = 360 degrees:  the final transverse moments should coincide with the
+initial transverse moments to within the level expected due to noise due to statistical sampling.
+
+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 run **either** as:
+
+* **Python** script: ``python3 run_solenoid.py`` or
+* ImpactX **executable** using an input file: ``impactx input_solenoid.in``
+
+For `MPI-parallel <https://www.mpi-forum.org>`__ runs, prefix these lines with ``mpiexec -n 4 ...`` or ``srun -n 4 ...``, depending on the system.
+
+.. tab-set::
+
+   .. tab-item:: Python: Script
+
+       .. literalinclude:: run_solenoid.py
+          :language: python3
+          :caption: You can copy this file from ``examples/solenoid_restart/run_solenoid.py``.
+
+   .. tab-item:: Executable: Input File
+
+       .. literalinclude:: input_solenoid.in
+          :language: ini
+          :caption: You can copy this file from ``examples/solenoid_restart/input_solenoid.in``.
+
+
+Analyze
+-------
+
+We run the following script to analyze correctness:
+
+.. dropdown:: Script ``analysis_solenoid.py``
+
+   .. literalinclude:: analysis_solenoid.py
+      :language: python3
+      :caption: You can copy this file from ``examples/solenoid_restart/analysis_solenoid.py``.
diff --git a/examples/solenoid_restart/analysis_solenoid.py b/examples/solenoid_restart/analysis_solenoid.py
@@ -0,0 +1,96 @@
+#!/usr/bin/env python3
+#
+# Copyright 2022-2023 ImpactX contributors
+# Authors: Axel Huebl, Chad Mitchell
+# License: BSD-3-Clause-LBNL
+#
+
+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.h5", io.Access.read_only)
+first_step = list(series.iterations)[0]
+last_step = list(series.iterations)[-1]
+initial = series.iterations[first_step].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 = 1.3 * 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],
+    [
+        2.205510139392e-3,
+        1.559531175539e-3,
+        6.404930308742e-3,
+        2.0e-6,
+        1.0e-6,
+        1.0e-6,
+    ],
+    rtol=rtol,
+    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 = 2.3 * num_particles**-0.5  # from random sampling of a smooth distribution
+print(f"  rtol={rtol} (ignored: atol~={atol})")
+
+assert np.allclose(
+    [sigx, sigy, emittance_x, emittance_y, emittance_t],
+    [
+        1.559531175539e-3,
+        2.205510139392e-3,
+        1.0e-6,
+        2.0e-6,
+        1.0e-6,
+    ],
+    rtol=rtol,
+    atol=atol,
+)
diff --git a/examples/solenoid_restart/input_solenoid.in b/examples/solenoid_restart/input_solenoid.in
@@ -0,0 +1,44 @@
+###############################################################################
+# Reference Particle
+###############################################################################
+beam.kin_energy = 250.0
+beam.particle = proton
+
+# ignored
+beam.charge = 1.0e-9
+beam.units = static
+beam.distribution = empty
+
+
+###############################################################################
+# Beamline: lattice elements and segments
+###############################################################################
+lattice.elements = source monitor sols monitor
+lattice.nslice = 1
+
+source.type = source
+source.distribution = openPMD
+source.openpmd_path = "../solenoid/diags/openPMD/monitor.h5"
+
+monitor.type = beam_monitor
+monitor.backend = h5
+
+sol1.type = solenoid
+sol1.ds = 3.820395
+sol1.ks = 0.8223219329893234
+
+sols.type = line
+sols.elements = sol1
+sols.repeat = 3
+
+###############################################################################
+# Algorithms
+###############################################################################
+algo.particle_shape = 2
+algo.space_charge = false
+
+
+###############################################################################
+# Diagnostics
+###############################################################################
+diag.slice_step_diagnostics = true