From 617937dfa782c9dc69b3c743dfc9749eaf680a0b Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Fri, 24 May 2024 13:12:43 -0500 Subject: [PATCH] Doc: Add Custom Elements Document how to add custom beamline elements to ImpactX as a user workflow. --- docs/source/install/hpc/perlmutter.rst | 2 +- docs/source/usage/how_to_run.rst | 4 +- docs/source/usage/python.rst | 26 +++--- docs/source/usage/workflows.rst | 14 ++-- docs/source/usage/workflows/add_element.rst | 87 +++++++++++++++++++++ examples/fodo_programmable/README.rst | 2 +- 6 files changed, 113 insertions(+), 22 deletions(-) create mode 100644 docs/source/usage/workflows/add_element.rst diff --git a/docs/source/install/hpc/perlmutter.rst b/docs/source/install/hpc/perlmutter.rst index 9cc3f3b01..1e52c774d 100644 --- a/docs/source/install/hpc/perlmutter.rst +++ b/docs/source/install/hpc/perlmutter.rst @@ -187,7 +187,7 @@ Use the following :ref:`cmake commands ` 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 ` for ImpactX :ref:`Python (PICMI) scripts ` (:ref:`example scripts `). +Now, you can :ref:`submit Perlmutter compute jobs ` for ImpactX :ref:`Python scripts ` (:ref:`example scripts `). Or, you can use the ImpactX executables to submit Perlmutter jobs (:ref:`example inputs `). For executables, you can reference their location in your :ref:`job script ` or copy them to a location in ``$PSCRATCH``. diff --git a/docs/source/usage/how_to_run.rst b/docs/source/usage/how_to_run.rst index 5f83805e4..921274597 100644 --- a/docs/source/usage/how_to_run.rst +++ b/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 python -Here, ```` is the number of MPI ranks used, and ```` is the name of the input file (```` is the name of the :ref:`PICMI ` script). +Here, ```` is the number of MPI ranks used, and ```` is the name of the input file (```` is the name of the :ref:`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``. diff --git a/docs/source/usage/python.rst b/docs/source/usage/python.rst index ab517ed6e..08f3e17ae 100644 --- a/docs/source/usage/python.rst +++ b/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. diff --git a/docs/source/usage/workflows.rst b/docs/source/usage/workflows.rst index ebbfbaccc..efdb7a523 100644 --- a/docs/source/usage/workflows.rst +++ b/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 `__ over time. diff --git a/docs/source/usage/workflows/add_element.rst b/docs/source/usage/workflows/add_element.rst new file mode 100644 index 000000000..c1027bddf --- /dev/null +++ b/docs/source/usage/workflows/add_element.rst @@ -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 `, 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 `: implements a user-defined drift +* :ref:`15 stage laser-plasma accelerator `: implements a user-defined LPA accelerator element using a neural network surrogate via PyTorch + +Detailed particle computing interfaces are presented in the `pyAMReX examples `__. + + +Linear Map +---------- + +.. note:: + + We plan to add a simple, linear map element that can be configured in user input. + Follow `issue #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/ `__. + +Let's take a look at an example, the `Drift `__ implementation. +To simplify the logic, we use so-called `mixin classes `__, which provide commonly used logic for `parallelization, thin/thick elements, alignment error support, etc `__. + +.. 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 `__. + +.. 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_ 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 `__ + * `Thin-Kick Dipole `__ + * `Chromatic Elements for Drift, Quad, Uniform Focusing+Solenoid `__ + * other pull requests under the `component: elements `__ label diff --git a/examples/fodo_programmable/README.rst b/examples/fodo_programmable/README.rst index 3321a36b4..48b58d555 100644 --- a/examples/fodo_programmable/README.rst +++ b/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 `. -However, in the example here we define *additional user-defined, custom elements* (:py:class:`impactx.elements.Programmable`) from the :ref:`ImpactX Python APIs `. +However, in the example here we define *additional user-defined, custom elements* (:py:class:`impactx.elements.Programmable`) from the :ref:`ImpactX Python APIs `. More generally, ImpactX exposes all data structures through `pyAMReX for adding additional computation `__, enabling rapid prototyping of new elements on both CPU and GPU.