User-Defined Linear Map

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

docs/source/usage/parameters.rst

@@ -349,6 +349,22 @@ 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.
+
+          This 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>.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:
 

docs/source/usage/python.rst

@@ -611,6 +611,22 @@ 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 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.
+
+   :param R: a linear transport map to multiply with the the phase space vector :math:`(x,px,y,py,t,pt)`.
+   :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.

docs/source/usage/workflows/add_element.rst

@@ -10,6 +10,29 @@ 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 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.
+
+
+.. 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 +53,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
 -----------

examples/CMakeLists.txt

@@ -82,7 +82,7 @@ function(add_impactx_test name input is_mpi analysis_script plot_script)
                         amrex.throw_exception = 1
                         amrex.signal_handling = 0
                         impactx.always_warn_immediately=1
-                        impactx.abort_on_warning_threshold=low
+                        impactx.abort_on_warning_threshold=medium
                  WORKING_DIRECTORY ${THIS_WORKING_DIR}
         )
     endif()
@@ -210,6 +210,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 +1071,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
+)

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

examples/fodo_userdef/analysis_fodo.py

@@ -0,0 +1 @@
+../fodo/analysis_fodo.py

examples/fodo_userdef/input_fodo_userdef.in

@@ -0,0 +1,59 @@
+###############################################################################
+# Particle Beam(s)
+###############################################################################
+beam.npart = 10000
+beam.units = static
+beam.kin_energy = 2.0e3
+beam.charge = 1.0e-9
+beam.particle = electron
+beam.distribution = waterbag_from_twiss
+beam.alphaX = -1.5905003499999992
+beam.alphaY = 1.5905003499999992
+beam.alphaT = 0.0
+beam.betaX = 2.8216194100262637
+beam.betaY = 2.8216194100262637
+beam.betaT = 0.5
+beam.emittX = 2e-09
+beam.emittY = 2e-09
+beam.emittT = 2e-06
+
+
+###############################################################################
+# Beamline: lattice elements and segments
+###############################################################################
+lattice.elements = monitor drift1 monitor quad1 monitor drift2 monitor quad2 monitor drift1 monitor
+lattice.nslice = 25
+
+monitor.type = beam_monitor
+monitor.backend = h5
+
+drift1.type = linear_map
+drift1.R12 = 0.25            # ds
+drift1.R34 = 0.25            # ds
+drift1.R56 = 0.25 / 16.6464  # ds / (beta*gamma^2)
+
+quad1.type = quad
+quad1.ds = 1.0
+quad1.k = 1.0
+
+drift2.type = linear_map
+drift2.R12 = 0.5            # ds
+drift2.R34 = 0.5            # ds
+drift2.R56 = 0.5 / 16.6464  # ds / (beta*gamma^2)
+
+quad2.type = quad
+quad2.ds = 1.0
+quad2.k = -1.0
+
+
+###############################################################################
+# Algorithms
+###############################################################################
+algo.particle_shape = 2
+algo.space_charge = false
+
+
+###############################################################################
+# Diagnostics
+###############################################################################
+diag.slice_step_diagnostics = true

examples/fodo_userdef/plot_fodo.py

@@ -0,0 +1 @@
+../fodo/plot_fodo.py

examples/fodo_userdef/run_fodo_userdef.py

@@ -0,0 +1,89 @@
+#!/usr/bin/env python3
+#
+# Copyright 2022-2024 ImpactX contributors
+# Authors: Axel Huebl, Chad Mitchell, Marco Garten
+# License: BSD-3-Clause-LBNL
+#
+# -*- coding: utf-8 -*-
+
+from impactx import ImpactX, distribution, elements, twiss
+
+sim = ImpactX()
+
+# set numerical parameters and IO control
+sim.particle_shape = 2  # B-spline order
+sim.space_charge = False
+# sim.diagnostics = False  # benchmarking
+sim.slice_step_diagnostics = True
+
+# domain decomposition & space charge mesh
+sim.init_grids()
+
+# load a 2 GeV electron beam with an initial
+# unnormalized rms emittance of 2 nm
+kin_energy_MeV = 2.0e3  # reference energy
+bunch_charge_C = 1.0e-9  # used with space charge
+npart = 10000  # number of macro particles
+
+#   reference particle
+ref = sim.particle_container().ref_particle()
+ref.set_charge_qe(-1.0).set_mass_MeV(0.510998950).set_kin_energy_MeV(kin_energy_MeV)
+
+#   particle bunch
+distr = distribution.Waterbag(
+    **twiss(
+        beta_x=2.8216194100262637,
+        beta_y=2.8216194100262637,
+        beta_t=0.5,
+        emitt_x=2e-09,
+        emitt_y=2e-09,
+        emitt_t=2e-06,
+        alpha_x=-1.5905003499999992,
+        alpha_y=1.5905003499999992,
+        alpha_t=0.0,
+    )
+)
+sim.add_particles(bunch_charge_C, distr, npart)
+
+# add beam diagnostics
+monitor = elements.BeamMonitor("monitor", backend="h5")
+
+# add a user-defined, linear element for the drifts
+Iden = elements.LinearMap.Map6x6.identity()
+R1, R2 = Iden, Iden
+
+ds1 = 0.25
+R1[1, 2] = ds1
+R1[3, 4] = ds1
+R1[5, 6] = ds1 / 16.6464  # ds / (beta*gamma^2)
+drift1 = elements.LinearMap(name="drift1", R=R1, ds=ds1)
+
+ds2 = 0.5
+R2[1, 2] = ds2
+R2[3, 4] = ds2
+R2[5, 6] = ds2 / 16.6464  # ds / (beta*gamma^2)
+drift2 = elements.LinearMap(name="drift2", R=R2, ds=ds2)
+
+# design the accelerator lattice)
+ns = 25  # number of slices per ds in the element
+fodo = [
+    monitor,
+    drift1,
+    monitor,
+    elements.Quad(name="quad1", ds=1.0, k=1.0, nslice=ns),
+    monitor,
+    drift2,
+    monitor,
+    elements.Quad(name="quad2", ds=1.0, k=-1.0, nslice=ns),
+    monitor,
+    drift1,
+    monitor,
+]
+# assign a fodo segment
+sim.lattice.extend(fodo)
+
+# run simulation
+sim.evolve()
+
+# clean shutdown
+sim.finalize()