Doc: Add Custom Elements

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
 ==================
@@ -367,9 +367,6 @@ Initial Beam Distributions
 
 This module provides particle beam distributions that can be used to initialize particle beams in an :py:class:`impactx.ParticleContainer`.
 
-.. py:module:: impactx.distribution
-   :synopsis: Particle beam distributions in ImpactX
-
 .. py:class:: impactx.distribution.Gaussian(lambdax, lambday, lambdat, lambdapx, lambdapy, lambdapt, muxpx=0.0, muypy=0.0, mutpt=0.0)
 
    A 6D Gaussian distribution.
@@ -430,9 +427,6 @@ Lattice Elements
 
 This module provides elements for the accelerator lattice.
 
-.. py:module:: impactx.elements
-   :synopsis: Accelerator lattice elements in ImpactX
-
 .. py:class:: impactx.elements.KnownElementsList
 
    An iterable, ``list``-like type of elements.
@@ -633,12 +627,26 @@ 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:: push
+
+      This is a function hook for pushing the whole particle container.
+      Either this function is implemented or ``beam_particles`` and ``ref_particle`` are needed.
+      This accepts a function or lambda with the following arguments:
+
+      .. py:method:: user_defined_function(pc: impactx.ParticleContainer, step: int)
+
+         This function is called for the particle container as it passes through the element.
+         Note that the reference particle must be updated *before* the beam particles are pushed.
+
    .. py:property:: beam_particles
 
       This is a function hook for pushing all beam particles.
       This accepts a function or lambda with the following arguments:
 
-      .. py:method:: user_defined_function(pti: ImpactXParIter, refpart: RefPart)
+      .. py:method:: user_defined_beam_function(pti: impactx.ImpactXParIter, refpart: impactx.RefPart)
 
          This function is called repeatedly for all particle tiles or boxes in the beam particle container.
          Particles can be pushed and are relative to the reference particle
@@ -648,7 +656,7 @@ This module provides elements for the accelerator lattice.
       This is a function hook for pushing the reference particle.
       This accepts a function or lambda with the following argument:
 
-      .. py:method:: another_user_defined_function(refpart: RefPart)
+      .. py:method:: user_defined_refpart_function(refpart: impactx.RefPart)
 
          This function is called for the reference particle as it passes through the element.
          The reference particle is updated *before* the beam particles are pushed.

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,91 @@
+.. _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 add new elements to ImpactX, you can pick the one that fits your needs best.
+
+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/particles/elements>`__).
+
+
+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.
+
+The Programmable element can implement a custom element in two ways:
+
+* Push the whole container, by assigning a ``push`` function or
+* Push the reference particle and beam particles in two individual functions (``beam_particles`` and ``ref_particle``).
+
+Per ImpactX convention, the reference particle is updated *before* the beam particles are pushed.
+
+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
+
+Detailed particle computing interfaces are presented in the `pyAMReX examples <https://pyamrex.readthedocs.io/en/latest/usage/compute.html#particles>`__.
+
+
+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 straight forward 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:: Python Binding: 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);
+
+Pull requests that added a new element and can be taken as examples 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>`__
+  * `Quadrupole with Soft-Edge Fringe Fields <https://github.com/ECP-WarpX/impactx/pull/322>`__
+  * 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.