Skip to content

Commit

Permalink
Doc: Add Custom Elements
Browse files Browse the repository at this point in the history
Document how to add custom beamline elements to ImpactX as a user
workflow.
  • Loading branch information
ax3l committed Jun 11, 2024
1 parent 9759f25 commit 617937d
Show file tree
Hide file tree
Showing 6 changed files with 113 additions and 22 deletions.
2 changes: 1 addition & 1 deletion docs/source/install/hpc/perlmutter.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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``.

Expand Down
4 changes: 2 additions & 2 deletions docs/source/usage/how_to_run.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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``.
Expand Down
26 changes: 17 additions & 9 deletions docs/source/usage/python.rst
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
.. _usage-picmi:
.. _usage-python:

Parameters: Python
==================
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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
Expand All @@ -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.
Expand Down
14 changes: 5 additions & 9 deletions docs/source/usage/workflows.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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.
87 changes: 87 additions & 0 deletions docs/source/usage/workflows/add_element.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1,87 @@
.. _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.


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);

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
2 changes: 1 addition & 1 deletion examples/fodo_programmable/README.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand Down

0 comments on commit 617937d

Please sign in to comment.