Calculate Linear Transfer Map

Add a `transfer_map(ref, order="linear", fallback_identity_map=False)`
method to the element list, supporting the calculation of
a linear transfer map of all elements in the list.

Author: ax3l

docs/source/usage/python.rst

@@ -769,6 +769,16 @@ This module provides elements and methods for the accelerator lattice.
       :return: True if at least one element of the specified kind exists
       :rtype: bool
 
+   .. py::method:: transfer_map(ref, order="linear", fallback_identity_map=False)
+
+      Calculate the transfer map of the elements in the list.
+
+      :param ref: A reference particle.
+      :param order: So far, only the calculation of linear transfer maps are supported in this function.
+      :param fallback_identity_map: For elements with an undefined transfer map in lattice, assume the identity matrix.
+      :return: The transfer map map of all elements in the list.
+      :rtype: Map6x6
+
    .. py:method:: plot_survey(ref=None, ax=None, legend=True, legend_ncols=5)
 
       Plot over s of all elements in the KnownElementsList.
@@ -777,7 +787,6 @@ This module provides elements and methods for the accelerator lattice.
 
       Either populates the matplotlib axes in ax or creates a new axes containing the plot.
 
-      :param self: The KnownElementsList class in ImpactX
       :param ref: A reference particle, checked for the charge sign to plot focusing/defocusing strength directions properly.
       :param ax: A plotting area in matplotlib (called axes there).
       :param legend: Plot a legend if true.

src/python/elements.cpp

@@ -9,6 +9,7 @@
 #include <elements/All.H>
 #include <elements/mixin/lineartransport.H>
 #include <elements/transformation/Insert.H>
+#include <particles/CovarianceMatrix.H>
 
 #include <AMReX_Enum.H>
 #include <AMReX_REAL.H>
@@ -20,7 +21,6 @@
 #include <utility>
 #include <variant>
 #include <vector>
-#include <particles/transformation/CoordinateTransformation.H>
 
 namespace py = pybind11;
 using namespace impactx;
@@ -313,7 +313,7 @@ void init_elements(py::module& m)
         )
     ;
 
-    /*
+    /* TODO
     py::class_<elements::mixin::LinearTransport>(mx, "LinearTransport")
         // type of map
         .def_property_readonly_static("Map6x6",
@@ -2615,6 +2615,49 @@ void init_elements(py::module& m)
             auto it = std::next(v.begin(), index);
             return *it;  // return by reference
         }, py::return_value_policy::reference_internal)
+
+        .def(
+            "transfer_map",
+            [](
+                KnownElementsList &v,
+                RefPart const & ref,
+                std::string order,
+                bool fallback_identity_map
+            )
+            {
+                if (order != "linear") {
+                    throw std::runtime_error("So far, only the calculation of linear transfer maps are supported in this function.");
+                }
+                Map6x6 result = Map6x6::Identity();
+                for (auto & el_v : v) {
+                    std::visit([&result, &ref, &fallback_identity_map](auto const & el) {
+                        using Element = std::decay_t<decltype(el)>;
+                        std::string not_impl_msg = "Undefined transfer map in lattice for element ";
+                        if (el.has_name()) not_impl_msg += el.name() + " ";
+                        not_impl_msg += std::string("of type ") + Element::type;
+
+                        if constexpr (std::is_base_of_v<elements::mixin::LinearTransport<Element>, Element>) {
+                            try {
+                                result = result * el.transport_map(ref);
+                            } catch (std::exception const & e) {
+                                if (!fallback_identity_map) {
+                                    throw std::runtime_error(not_impl_msg);
+                                }
+                            }
+                        } else {
+                            if (!fallback_identity_map) {
+                                throw std::runtime_error(not_impl_msg);
+                            }
+                        }
+                    }, el_v);
+                }
+                return result;
+            },
+            py::arg("ref"),
+            py::arg("order") = "linear",
+            py::arg("fallback_identity_map") = false,
+            "Calculate the transfer map of the elements in the list."
+            )
     ;
 
 

tests/python/test_lattice_optics.py

@@ -0,0 +1,63 @@
+#!/usr/bin/env python3
+#
+# Copyright 2022-2024 The ImpactX Community
+#
+# Authors: Axel Huebl, Chad Mitchell
+# License: BSD-3-Clause-LBNL
+#
+# -*- coding: utf-8 -*-
+
+import numpy as np
+import pytest
+
+from impactx import RefPart, elements
+
+
+def test_lattice_linear_map():
+    """Calculate the linear transfer map of a lattice."""
+
+    # Create reference particle
+    ref = RefPart()
+    ref.set_charge_qe(-1.0).set_mass_MeV(0.510998950).set_kin_energy_MeV(1.0e3)
+
+    # Create a valid test lattice (all elements define a linear transfer map)
+    lattice = elements.KnownElementsList()
+    lattice.extend(
+        [
+            elements.Drift(name="drift1", ds=1.0),
+            elements.Quad(name="quad1", ds=0.5, k=1.0),
+            elements.Drift(name="drift2", ds=2.0),
+            elements.Sbend(name="bend1", ds=1.0, rc=10.0),
+        ]
+    )
+
+    # Expected result (matrix multiplication)
+    R_expected = np.array(
+        [
+            [3.74670546e-01, 2.54005827e00, 0, 0, 0, -2.34864805e-01],
+            [-4.76219076e-01, -5.59489407e-01, 0, 0, 0, 3.20646253e-02],
+            [0, 0, 1.64872127e00, 6.59488508e00, 0, 0],
+            [0, 0, 5.21095305e-01, 2.69091188e00, 0, 0],
+            [9.98334297e-02, 4.99583537e-02, 0, 0, 1.00000000e00, -1.66466013e-03],
+            [0, 0, 0, 0, 0, 1.00000000e00],
+        ]
+    )
+
+    # Calculate Linear Transfer Map
+    R = lattice.transfer_map(ref)
+    assert np.allclose(R.to_numpy(), R_expected)
+
+    # Check unexpected/unsupported options
+    with pytest.raises(RuntimeError):
+        lattice.transfer_map(ref, order="invalid")
+
+    # Create a lattice with an element that does not define a linear transfer map
+    lattice.append(elements.TaperedPL(k=0, taper=0, unit=0))
+
+    # Ensure that the calculation asserts
+    with pytest.raises(RuntimeError):
+        lattice.transfer_map(ref)
+
+    # Now the user explicitly assumes that undefined maps are identify maps
+    R = lattice.transfer_map(ref, fallback_identity_map=True)
+    assert np.allclose(R.to_numpy(), R_expected)