Doc: Add Custom Elements

Document how to add custom beamline elements to ImpactX as a user
workflow.

docs/source/install/hpc/perlmutter.rst

@@ -187,7 +187,7 @@ Use the following :ref:`cmake commands <building-cmake>` to compile the applicat
          cmake -S . -B build_pm_cpu_py -DImpactX_COMPUTE=OMP -DImpactX_APP=OFF -DImpactX_PYTHON=ON
          cmake --build build_pm_cpu_py -j 16 --target pip_install
 
-Now, you can :ref:`submit Perlmutter compute jobs <running-cpp-perlmutter>` for ImpactX :ref:`Python (PICMI) scripts <usage-picmi>` (:ref:`example scripts <usage-examples>`).
+Now, you can :ref:`submit Perlmutter compute jobs <running-cpp-perlmutter>` for ImpactX :ref:`Python scripts <usage-python>` (:ref:`example scripts <usage-examples>`).
 Or, you can use the ImpactX executables to submit Perlmutter jobs (:ref:`example inputs <usage-examples>`).
 For executables, you can reference their location in your :ref:`job script <running-cpp-perlmutter>` or copy them to a location in ``$PSCRATCH``.
 

docs/source/usage/how_to_run.rst

@@ -69,10 +69,10 @@ or
 
 .. code-block:: bash
 
-   # run with a PICMI input script:
+   # run with a Python input script:
    mpirun -np <n_ranks> python <python_script>
 
-Here, ``<n_ranks>`` is the number of MPI ranks used, and ``<input_file>`` is the name of the input file (``<python_script>`` is the name of the :ref:`PICMI <usage-picmi>` script).
+Here, ``<n_ranks>`` is the number of MPI ranks used, and ``<input_file>`` is the name of the input file (``<python_script>`` is the name of the :ref:`Python <usage-python>` script).
 Note that the actual executable might have a longer name, depending on build options.
 
 We used the copied executable in the current directory (``./``); if you installed with a package manager, skip the ``./`` because ImpactX is in your ``PATH``.

docs/source/usage/python.rst

@@ -1,4 +1,4 @@
-.. _usage-picmi:
+.. _usage-python:
 
 Parameters: Python
 ==================
@@ -633,6 +633,9 @@ This module provides elements for the accelerator lattice.
 
    This element can be programmed to receive callback hooks into Python functions.
 
+   :param ds: Segment length in m.
+   :param nslice: number of slices used for the application of space charge
+
    .. py:property:: beam_particles
 
       This is a function hook for pushing all beam particles.

docs/source/usage/workflows.rst

@@ -5,15 +5,11 @@ Workflows
 
 This section collects typical user workflows and best practices for ImpactX.
 
-.. note::
+.. toctree::
+   :maxdepth: 1
 
-   TODO: Add more workflows as in https://warpx.readthedocs.io/en/latest/usage/workflows.html
+   workflows/add_element
 
-.. toctree::
-   :maxdepth: 2
+.. note::
 
-..   workflows/parallelization
-..   workflows/profiling
-   workflows/debugging
-..   workflows/libensemble
-..   workflows/plot_distribution_mapping
+   We will add more `workflows as in WarpX <https://warpx.readthedocs.io/en/latest/usage/workflows.html>`__ over time.

docs/source/usage/workflows/add_element.rst

@@ -0,0 +1,80 @@
+.. _usage-workflows-add-element:
+
+Add New Beamline Elements
+=========================
+
+In ImpactX, one can easily add new beamline elements as a user.
+There are multiple ways to achieve this, which are introduced in this section.
+
+
+Python Programmable Element
+---------------------------
+
+Using the :ref:`ImpactX Python interface <usage-python>`, a custom element named :py:class:`impactx.elements.Programmable` can be defined to advance particles using NumPy, CuPy, Numba, PyTorch or any other compatible Python library.
+
+...
+
+Detailed examples that show usage of the programmable element are:
+
+* :ref:`FODO cell <examples-fodo-programmable>`: implements a user-defined drift
+* :ref:`15 stage laser-plasma accelerator <examples-ml-surrogate>`: implements a user-defined LPA accelerator element using a neural network surrogate via PyTorch
+
+
+Linear Map
+----------
+
+.. note::
+
+   We plan to add a simple, linear map element that can be configured in user input.
+   Follow `issue #538 <https://github.com/ECP-WarpX/impactx/issues/538>`__ for progress.
+
+
+C++ Element
+-----------
+
+Adding a new beamline element directly to the C++ code base of ImpactX is easy and described in the following.
+
+We store all beamline elements under `src/particles/elements/ <https://github.com/ECP-WarpX/impactx/tree/development/src/particles/elements>`__.
+
+Let's take a look at an example, the `Drift <https://impactx.readthedocs.io/en/latest/_static/doxyhtml/structimpactx_1_1_drift.html>`__ implementation.
+To simplify the logic, we use so-called `mixin classes <https://en.wikipedia.org/wiki/Mixin>`__, which provide commonly used logic for `parallelization, thin/thick elements, alignment error support, etc <https://impactx.readthedocs.io/en/latest/_static/doxyhtml/namespaceimpactx_1_1elements.html>`__.
+
+.. literalinclude:: ../../../src/particles/elements/Drift.H
+   :language: cpp
+   :dedent: 4
+   :start-at: struct Drift
+   :end-at: static constexpr auto name = "Drift";
+
+After this brief boilerplate, our beamline elements implement three simple parts:
+
+* a constructor: storing element options
+* a single-particle operator: pushing the beam particles
+* a reference-particle operator: pushing the reference particle
+
+.. dropdown:: Example Element: Drift.H
+   :color: light
+   :icon: info
+   :animate: fade-in-slide-down
+
+   .. literalinclude:: ../../../src/particles/elements/Drift.H
+      :language: cpp
+
+As a last step, we expose our C++ beamline elements to Python in `src/python/elements.cpp <https://github.com/ECP-WarpX/impactx/blob/development/src/python/elements.cpp>`__.
+
+.. dropdown:: Example Python Binding for Drift
+   :color: light
+   :icon: info
+   :animate: fade-in-slide-down
+
+   .. literalinclude:: ../../../src/python/elements.cpp
+      :language: cpp
+      :dedent: 8
+      :start-at: py::class_<Drift, elements::Thick, elements::Alignment> py_Drift(me, "Drift");
+      :end-at: register_beamoptics_push(py_Drift);
+
+Example pull requests that added a new element and can be taken as an example are:
+
+  * `Chromatic Plasma Lens <https://github.com/ECP-WarpX/impactx/pull/514>`__
+  * `Thin-Kick Dipole <https://github.com/ECP-WarpX/impactx/pull/472>`__
+  * `Chromatic Elements for Drift, Quad, Uniform Focusing+Solenoid <https://github.com/ECP-WarpX/impactx/pull/356>`__
+  * other pull requests under the `component: elements <https://github.com/ECP-WarpX/impactx/pulls?q=is%3Apr+label%3A%22component%3A+elements%22+is%3Amerged+>`__ label

examples/fodo_programmable/README.rst

@@ -4,7 +4,7 @@ FODO Cell, Programmable Element
 ===============================
 
 This implements the same FODO cell as the :ref:`stable FODO cell example <examples-fodo>`.
-However, in the example here we define *additional user-defined, custom elements* (:py:class:`impactx.elements.Programmable`) from the :ref:`ImpactX Python APIs <usage-picmi>`.
+However, in the example here we define *additional user-defined, custom elements* (:py:class:`impactx.elements.Programmable`) from the :ref:`ImpactX Python APIs <usage-python>`.
 
 More generally, ImpactX exposes all data structures through `pyAMReX for adding additional computation <https://pyamrex.readthedocs.io/en/latest/usage/compute.html>`__, enabling rapid prototyping of new elements on both CPU and GPU.