User-Defined Linear Map (#743)

* [Draft] Transport and Covariance Matrices General structure. Co-authored-by: Chad Mitchell <[email protected]> * Update InitElement.cpp Update initialization of LinearMap element. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update LinearMap.H Finish implementation of LinearMap element push. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update InitDistribution.cpp Initialize covariance matrix from distribution parameters. * Add Drift return linear map. * Use LinearMap struct in Drift.H. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Transport/Covariane: Use `SmallMatrix` * Add Python bindings for LinearMap element. * Update elements.cpp Remove unnecessary ,. * Update elements.cpp Correct naming in elements.cpp binding for LinearMap. * Update elements.cpp Correct argument declaration for Python bindings of LinearMap in elements.cpp. * Add LinearMap to 'Named' elements. * Update LinearMap.H Document "Name" parameter in LinearMap. * Modify LinearMap Python bindings. * Update src/initialization/InitDistribution.cpp * Update src/initialization/InitDistribution.cpp * Update InitDistribution.cpp Modify zero-initialization of CovarianceMatrix cv. * Leftover Improvements from first PR * Finalize Python * More Detailed Warning Include element name in output to user. * Start Examples * Start Documentation * Fix Bugs (incl. AMReX) Use parser for inputs. Needs upstream AMReX fix. * Docs Updates * `LinearMap`: `ds` bookkeeping And reference particle pushing (drifting), if thick. * Add thin linear map example. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Document example README. * Update run_map.py Add twiss * Update run_map.py Correct name of Rmat/R. * Update run_map.py Fix naming of R again. * Update input_map.in Fix number of periods. * Apply suggestions from code review Co-authored-by: Chad Mitchell <[email protected]> * Add docs for matrix elements. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Doc Updates * Generalize CMake Args for inputs Taken over from generalized version we merged to WarpX. * Add symplectic warning in app element documentation. * Add symplectic warning in Python element documentation. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Remove runtime warning. * Add support for nonzero ds to app input. * Update examples/CMakeLists.txt * FODO Analysis: Check s and gamma * __repr__: no select entries --------- Co-authored-by: Axel Huebl <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
ECP-WarpX · Jan 13, 2025 · 5638dfd · 5638dfd
1 parent 350565b
commit 5638dfd
Show file tree

Hide file tree

Showing 24 changed files with 1,173 additions and 21 deletions.
diff --git a/docs/source/usage/examples.rst b/docs/source/usage/examples.rst
@@ -31,10 +31,11 @@ Single Particle Dynamics
    examples/aperture/README.rst
    examples/iota_lens/README.rst
    examples/achromatic_spectrometer/README.rst
+   examples/fodo_userdef/README.rst
    examples/fodo_programmable/README.rst
    examples/dogleg/README.rst
    examples/coupled_optics/README.rst
-
+   examples/linear_map/README.rst
 
 Collective Effects
 ------------------

diff --git a/docs/source/usage/parameters.rst b/docs/source/usage/parameters.rst
@@ -349,6 +349,25 @@ Lattice Elements
             * ``<element_name>.rotation`` (``float``, in degrees) rotation error in the transverse plane
             * ``<element_name>.nslice`` (``integer``) number of slices used for the application of space charge (default: ``1``)
 
+        * ``linear_map`` for a custom, linear transport matrix.
+
+          The matrix elements :math:`R(i,j)` are indexed beginning with 1, so that :math:`i,j=1,2,3,4,5,6`.
+          The transport matrix :math:`R` is defaulted to the identity matrix, so only matrix entries that differ from that need to be specified.
+
+          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 internal tracking methods used by ImpactX are symplectic.  However, if a user-defined linear map :math:`R` is provided, it is up to the user to ensure that the matrix :math:`R` is symplectic.  Otherwise, this condition may be violated.
+
+          This element requires these additional parameters:
+
+            * ``<element_name>.R(i,j)`` (``float``, ...) matrix entries
+                a 1-indexed, 6x6, linear transport map to multiply with the the phase space vector :math:`x,px,y,py,t,pt`.
+            * ``<element_name>.ds`` (``float``, in meters) length associated with a user-defined linear element (defaults to 0)
+            * ``<element_name>.dx`` (``float``, in meters) horizontal translation error
+            * ``<element_name>.dy`` (``float``, in meters) vertical translation error
+            * ``<element_name>.rotation`` (``float``, in degrees) rotation error in the transverse plane
+
         * ``multipole`` for a thin multipole element.
           This requires these additional parameters:
 

diff --git a/docs/source/usage/python.rst b/docs/source/usage/python.rst
@@ -611,6 +611,25 @@ This module provides elements for the accelerator lattice.
    :param unit: specification of units (``"dimensionless"`` in units of the magnetic rigidity of the reference particle or ``"T-m"``)
    :param name: an optional name for the element
 
+.. py::class:: impactx.elements.LinearMap(R, dx=0, dy=0, rotation=0, name=None)
+
+   A custom, linear transport matrix.
+
+   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 internal tracking methods used by ImpactX are symplectic.  However, if a user-defined linear map :math:`R` is provided, it is
+   up to the user to ensure that the matrix :math:`R` is symplectic.  Otherwise, this condition may be violated.
+
+   :param R: a linear transport map to multiply with the the phase space vector :math:`(x,px,y,py,t,pt)`.
+   :param ds: length associated with a user-defined linear element (defaults to 0), in m
+   :param dx: horizontal translation error in m
+   :param dy: vertical translation error in m
+   :param rotation: rotation error in the transverse plane [degrees]
+   :param name: an optional name for the element
+
 .. py:class:: impactx.elements.Multipole(multipole, K_normal, K_skew, dx=0, dy=0, rotation=0, name=None)
 
    A general thin multipole element.

diff --git a/docs/source/usage/workflows/add_element.rst b/docs/source/usage/workflows/add_element.rst
@@ -10,6 +10,28 @@ 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>`__).
 
 
+.. _usage-workflows-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.
+
+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.
+
+
+.. 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.
+
+
+.. _usage-workflows-add-element-python:
+
 Python Programmable Element
 ---------------------------
 
@@ -30,14 +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>`__.
 
 
-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.
-
+.. _usage-workflows-add-element-cxx:
 
 C++ Element
 -----------

diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt
@@ -55,6 +55,30 @@ function(add_impactx_test name input is_mpi analysis_script plot_script)
     # make a unique run directory
     file(MAKE_DIRECTORY ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/${name})
 
+    # get input file/script and optional command-line arguments
+    separate_arguments(INPUTS_LIST UNIX_COMMAND "${input}")
+    list(GET INPUTS_LIST 0 INPUTS_FILE)
+    list(LENGTH INPUTS_LIST INPUTS_LIST_LENGTH)
+    if(INPUTS_LIST_LENGTH GREATER 1)
+        list(SUBLIST INPUTS_LIST 1 -1 INPUTS_ARGS)
+        list(JOIN INPUTS_ARGS " " INPUTS_ARGS)
+    else()
+        set(INPUTS_ARGS "")
+    endif()
+    cmake_path(SET INPUTS_FILE "${ImpactX_SOURCE_DIR}/${INPUTS_FILE}")
+
+    # get analysis script and optional command-line arguments
+    separate_arguments(ANALYSIS_LIST UNIX_COMMAND "${analysis_script}")
+    list(GET ANALYSIS_LIST 0 ANALYSIS_FILE)
+    cmake_path(SET ANALYSIS_FILE "${CMAKE_CURRENT_SOURCE_DIR}/${ANALYSIS_FILE}")
+    list(LENGTH ANALYSIS_LIST ANALYSIS_LIST_LENGTH)
+    if(ANALYSIS_LIST_LENGTH GREATER 1)
+        list(SUBLIST ANALYSIS_LIST 1 -1 ANALYSIS_ARGS)
+        list(JOIN ANALYSIS_ARGS " " ANALYSIS_ARGS)
+    else()
+        set(ANALYSIS_ARGS "")
+    endif()
+
     # test run
     set(THIS_WORKING_DIR ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/${name})
     set(THIS_MPI_TEST_EXE)
@@ -64,10 +88,10 @@ function(add_impactx_test name input is_mpi analysis_script plot_script)
     set(THIS_Python_EXE)
     if(is_python)
         set(THIS_Python_EXE ${Python_EXECUTABLE})
-    endif()
-    if(is_python)
+        # for argparse, do not pass command-line arguments as one quoted string
+        separate_arguments(INPUTS_ARGS UNIX_COMMAND "${INPUTS_ARGS}")
         add_test(NAME ${name}.run
-                 COMMAND ${THIS_MPI_TEST_EXE} ${THIS_Python_EXE} ${ImpactX_SOURCE_DIR}/${input}
+                 COMMAND ${THIS_MPI_TEST_EXE} ${THIS_Python_EXE} ${INPUTS_FILE} ${INPUTS_ARGS}
                  WORKING_DIRECTORY ${THIS_WORKING_DIR}
         )
         # TODO:
@@ -77,12 +101,13 @@ function(add_impactx_test name input is_mpi analysis_script plot_script)
     else()
         add_test(NAME ${name}.run
                  COMMAND
-                    ${THIS_MPI_TEST_EXE} $<TARGET_FILE:app> ${ImpactX_SOURCE_DIR}/${input}
+                    ${THIS_MPI_TEST_EXE} $<TARGET_FILE:app> ${INPUTS_FILE}
                         amrex.abort_on_unused_inputs=1
                         amrex.throw_exception = 1
                         amrex.signal_handling = 0
                         impactx.always_warn_immediately=1
                         impactx.abort_on_warning_threshold=low
+                        ${INPUTS_ARGS}
                  WORKING_DIRECTORY ${THIS_WORKING_DIR}
         )
     endif()
@@ -210,6 +235,21 @@ add_impactx_test(FODO.MADX.py
     examples/fodo/plot_fodo.py
 )
 
+# Python: FODO Cell w/ custom linear element  #################################
+#
+add_impactx_test(FODO.userdef
+        examples/fodo_userdef/input_fodo_userdef.in
+        ON   # ImpactX MPI-parallel
+        examples/fodo_userdef/analysis_fodo.py
+        examples/fodo_userdef/plot_fodo.py
+)
+add_impactx_test(FODO.userdef.py
+        examples/fodo_userdef/run_fodo_userdef.py
+        OFF  # ImpactX MPI-parallel
+        examples/fodo_userdef/analysis_fodo.py
+        examples/fodo_userdef/plot_fodo.py
+)
+
 # Python: MPI-parallel FODO Cell ##############################################
 #
 add_impactx_test(FODO.py.MPI
@@ -1056,3 +1096,19 @@ add_impactx_test(linac-segment.py
     examples/linac_segment/analysis_linac_segment.py
     OFF  # no plot script yet
 )
+
+# Iteration of a linear one-turn map #########################################
+#
+# w/o space charge
+add_impactx_test(linear-map
+    examples/linear_map/input_map.in
+    ON   # ImpactX MPI-parallel
+    examples/linear_map/analysis_map.py
+    OFF  # no plot script yet
+)
+add_impactx_test(linear-map.py
+    examples/linear_map/run_map.py
+    ON  # ImpactX MPI-parallel
+    examples/linear_map/analysis_map.py
+    OFF  # no plot script yet
+)
diff --git a/examples/fodo/analysis_fodo.py b/examples/fodo/analysis_fodo.py
@@ -38,7 +38,8 @@ def get_moments(beam):
 series = io.Series("diags/openPMD/monitor.h5", io.Access.read_only)
 last_step = list(series.iterations)[-1]
 initial = series.iterations[1].particles["beam"].to_df()
-final = series.iterations[last_step].particles["beam"].to_df()
+beam_final = series.iterations[last_step].particles["beam"]
+final = beam_final.to_df()
 
 # compare number of particles
 num_particles = 10000
@@ -74,24 +75,29 @@ def get_moments(beam):
 print("")
 print("Final Beam:")
 sigx, sigy, sigt, emittance_x, emittance_y, emittance_t = get_moments(final)
+s_ref = beam_final.get_attribute("s_ref")
+gamma_ref = beam_final.get_attribute("gamma_ref")
 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}"
+    f"  emittance_x={emittance_x:e} emittance_y={emittance_y:e} emittance_t={emittance_t:e}\n"
+    f"  s_ref={s_ref:e} gamma_ref={gamma_ref:e}"
 )
 
 atol = 0.0  # ignored
 rtol = 2.2 * 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],
+    [sigx, sigy, sigt, emittance_x, emittance_y, emittance_t, s_ref, gamma_ref],
     [
         7.4790118496224206e-005,
         7.5357525169680140e-005,
         9.9775879288128088e-004,
         1.9959539836392703e-009,
         2.0175014668882125e-009,
         2.0013820380883801e-006,
+        3.000000,
+        3.914902e003,
     ],
     rtol=rtol,
     atol=atol,

diff --git a/examples/fodo_userdef/README.rst b/examples/fodo_userdef/README.rst
@@ -0,0 +1,85 @@
+.. _examples-fodo-userdef:
+
+User-Defined Linear 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 linear elements* by providing a custom matrix.
+
+.. note::
+
+   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>`.
+
+The matched Twiss parameters at entry are:
+
+* :math:`\beta_\mathrm{x} = 2.82161941` m
+* :math:`\alpha_\mathrm{x} = -1.59050035`
+* :math:`\beta_\mathrm{y} = 2.82161941` m
+* :math:`\alpha_\mathrm{y} = 1.59050035`
+
+We use a 2 GeV electron beam with initial unnormalized rms emittance of 2 nm.
+
+The second moments of the particle distribution after the FODO cell should coincide with the second moments of the particle distribution before the FODO cell, to within the level expected due to noise due to statistical sampling.
+
+In this test, the initial and final values of :math:`\lambda_x`, :math:`\lambda_y`, :math:`\lambda_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_fodo_userdef.py`` or
+* ImpactX **executable** using an input file: ``impactx input_fodo_userdef.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_fodo_userdef.py
+          :language: python3
+          :caption: You can copy this file from ``examples/fodo_userdef/run_fodo_userdef.py``.
+
+   .. tab-item:: Executable: Input File
+
+       .. literalinclude:: input_fodo_userdef.in
+          :language: ini
+          :caption: You can copy this file from ``examples/fodo_userdef/input_fodo_userdef.in``.
+
+
+Analyze
+-------
+
+We run the following script to analyze correctness:
+
+.. dropdown:: Script ``analysis_fodo.py``
+
+   .. literalinclude:: analysis_fodo.py
+      :language: python3
+      :caption: You can copy this file from ``examples/fodo_userdef/analysis_fodo.py``.
+
+
+Visualize
+---------
+
+You can run the following script to visualize the beam evolution over time:
+
+.. dropdown:: Script ``plot_fodo.py``
+
+   .. literalinclude:: plot_fodo.py
+      :language: python3
+      :caption: You can copy this file from ``examples/fodo_userdef/plot_fodo.py``.
+
+.. figure:: https://user-images.githubusercontent.com/1353258/180287840-8561f6fd-278f-4856-abd8-04fbdb78c8ff.png
+   :alt: focusing, defocusing and preserved emittance in our FODO cell benchmark.
+
+   FODO transversal beam width and emittance evolution
+
+.. figure:: https://user-images.githubusercontent.com/1353258/180287845-eb0210a7-2500-4aa9-844c-67fb094329d3.png
+   :alt: focusing, defocusing and phase space rotation in our FODO cell benchmark.
+
+   FODO transversal beam width and phase space evolution
diff --git a/examples/fodo_userdef/analysis_fodo.py b/examples/fodo_userdef/analysis_fodo.py
@@ -0,0 +1 @@
+../fodo/analysis_fodo.py