MPL: KnownElementsList.plot_survey()

Example:
```py
import matplotlib.pyplot as plt

sim.lattice.plot_survey(ref)
plt.show()
```

Author: ax3l

docs/source/usage/python.rst

@@ -623,6 +623,21 @@ This module provides elements and methods for the accelerator lattice.
       :param madx_file: file name to MAD-X file with beamline elements
       :param nslice: number of slices used for the application of space charge
 
+   .. py:method:: plot_survey(ref=None, ax=None, legend=True, legend_ncols=5)
+
+      Plot over s of all elements in the KnownElementsList.
+
+      The signs of element strengths are determined by the sign of the charge of the reference particle.
+      The projection of all element strengths is s-x ("vertical").
+
+      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.
+      :param legend_ncols: Number of columns for lattice element types in the legend.
+
 .. py:class:: impactx.elements.CFbend(ds, rc, k, dx=0, dy=0, rotation=0, aperture_x=0, aperture_y=0, nslice=1, name=None)
 
    A combined function bending magnet.  This is an ideal Sbend with a normal quadrupole field component.

src/elements/Sbend.H

@@ -71,6 +71,17 @@ namespace impactx::elements
         {
         }
 
+        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+        amrex::ParticleReal
+        rc ([[maybe_unused]] RefPart const & refpart) const
+        {
+            using namespace amrex::literals; // for _rt and _prt
+
+            // TODO: as in ExactSbend
+            // return m_B != 0_prt ? refpart.rigidity_Tm() / m_B : m_ds / m_phi;
+            return m_rc;
+        }
+
         /** Push all particles */
         using BeamOptic::operator();
 

src/python/elements.cpp

@@ -1170,6 +1170,10 @@ void init_elements(py::module& m)
              py::arg("name") = py::none(),
              "An ideal sector bend using the exact nonlinear map.  When B = 0, the reference bending radius is defined by r0 = length / (angle in rad), corresponding to a magnetic field of B = rigidity / r0; otherwise the reference bending radius is defined by r0 = rigidity / B."
         )
+        .def("rc", &ExactSbend::rc,
+            py::arg("ref"),
+            "Radius of curvature in m"
+        )
         .def_property("phi",
             [](ExactSbend & exact_sbend) { return exact_sbend.m_phi; },
             [](ExactSbend & exact_sbend, amrex::ParticleReal phi) { exact_sbend.m_phi = phi; },
@@ -1659,9 +1663,8 @@ void init_elements(py::module& m)
              py::arg("name") = py::none(),
              "An ideal sector bend."
         )
-        .def_property("rc",
-            [](Sbend & sbend) { return sbend.m_rc; },
-            [](Sbend & sbend, amrex::ParticleReal rc) { sbend.m_rc = rc; },
+        .def("rc", &Sbend::rc,
+            py::arg("ref") = py::none(),
             "Radius of curvature in m"
         )
     ;

src/python/impactx/__init__.py

@@ -18,25 +18,27 @@
 
 # import core bindings to C++
 from . import impactx_pybind as cxx
-from .distribution_input_helpers import twiss  # noqa
-from .extensions.ImpactXParticleContainer import (
-    register_ImpactXParticleContainer_extension,
-)
 from .impactx_pybind import *  # noqa
-from .madx_to_impactx import read_beam, read_lattice  # noqa
+from .madx_to_impactx import read_beam  # noqa
 
 __version__ = cxx.__version__
 __doc__ = cxx.__doc__
 __license__ = cxx.__license__
 __author__ = cxx.__author__
 
+from .distribution_input_helpers import twiss  # noqa
+from .extensions.KnownElementsList import (
+    register_KnownElementsList_extension,
+)
+from .extensions.ImpactXParticleContainer import (
+    register_ImpactXParticleContainer_extension,
+)
+
 # at this place we can enhance Python classes with additional methods written
 # in pure Python or add some other Python logic
 
 # MAD-X file reader for beamline lattice elements
-elements.KnownElementsList.load_file = lambda self, madx_file, nslice=1: self.extend(
-    read_lattice(madx_file, nslice)
-)  # noqa
+register_KnownElementsList_extension(cxx.elements.KnownElementsList)
 
 # MAD-X file reader for reference particle
 RefPart.load_file = read_beam  # noqa

src/python/impactx/extensions/KnownElementsList.py

@@ -0,0 +1,19 @@
+"""
+This file is part of ImpactX
+
+Copyright 2025 ImpactX contributors
+Authors: Axel Huebl
+License: BSD-3-Clause-LBNL
+"""
+
+
+def register_KnownElementsList_extension(kel):
+    """KnownElementsList helper methods"""
+    from ..madx_to_impactx import read_lattice
+    from ..plot.Survey import plot_survey
+
+    # register member functions for KnownElementsList
+    kel.load_file = lambda self, madx_file, nslice=1: self.extend(
+        read_lattice(madx_file, nslice)
+    )
+    kel.plot_survey = plot_survey

src/python/impactx/plot/ElementColors.py

@@ -0,0 +1,56 @@
+"""
+This file is part of ImpactX
+
+Copyright 2025 ImpactX contributors
+Authors: Axel Huebl
+License: BSD-3-Clause-LBNL
+"""
+
+
+def get_element_color_palette(palette="cern-lhc", plot_library="mpl"):
+    """Return a dictionary with colors for all elements.
+
+    The key is a regex that can be matched against the element type string. TODO TODO
+    """
+    color_palette = {
+        "cern-lhc": {
+            "Quad": "tab:blue",
+            "Multipole": "tab:orange",
+            "Sbend": "tab:green",
+            "CFbend": "tab:olive",  # TODO: improve and plot as two on top of each other
+            "ConstF": "tab:red",
+            "ChrPlasmaLens": "tab:red",
+            "SoftSolenoid": "tab:red",
+            "TaperedPL": "tab:red",
+            "RFCavity": "tab:brown",
+            "ShortRF": "tab:brown",
+            "Buncher": "tab:purple",
+            "Aperture": "black",
+            "Kicker": "tab:pink",
+            # 'tab:cyan'
+            "other": "tab:gray",
+        }
+    }
+
+    colors = color_palette[palette]
+
+    if plot_library != "mpl":
+        # remove "tab:" prefix
+        for k, v in colors.items():
+            colors[k] = v[4:]
+
+    return colors
+
+
+def get_element_color(element_kind: str, palette="cern-lhc", plot_library="mpl"):
+    """Get the color for a given element type string."""
+    color_palette = get_element_color_palette(palette, plot_library)
+
+    # sub-string matching of keys
+    found_keys = [key for key in color_palette.keys() if key in element_kind]
+
+    if found_keys:
+        first_found = found_keys[0]
+        return color_palette[first_found]
+    else:
+        return color_palette["other"]

src/python/impactx/plot/Survey.py

@@ -0,0 +1,133 @@
+"""
+This file is part of ImpactX
+
+Copyright 2025 ImpactX contributors
+Authors: Axel Huebl
+License: BSD-3-Clause-LBNL
+"""
+
+
+def plot_survey(
+    self, ref=None, ax=None, legend=True, legend_ncols=5, palette="cern-lhc"
+):
+    """Plot over s of all elements in the KnownElementsList.
+
+    The signs of element strengths are determined by the sign of the charge of the reference particle.
+    The projection of all element strengths is s-x ("vertical").
+
+    Parameters
+    ----------
+    self : ImpactXParticleContainer_*
+        The KnownElementsList class in ImpactX
+    ref : RefPart
+        A reference particle, checked for the charge sign to plot focusing/defocusing strength directions properly.
+    ax : matplotlib axes
+        A plotting area in matplotlib (called axes there).
+    legend: bool
+        Plot a legend if true.
+    legend_ncols: int
+        Number of columns for lattice element types in the legend.
+    palette: string
+        Color palette.
+
+    Returns
+    -------
+    Either populates the matplotlib axes in ax or creates a new axes containing the plot.
+    """
+    from math import copysign
+
+    import matplotlib.pyplot as plt
+    import numpy as np
+    from matplotlib.patches import Rectangle
+
+    from .ElementColors import get_element_color
+
+    charge_qe = 1.0 if ref is None else ref.charge_qe
+
+    ax = ax or plt.subplot(111)
+
+    element_lengths = [element.ds for element in self]
+
+    # NumPy 2.1+ (i.e. Python 3.10+):
+    # element_s = np.cumulative_sum(element_lengths, include_initial=True)
+    # backport:
+    element_s = np.insert(np.cumsum(element_lengths), 0, 0)
+
+    ax.hlines(0, 0, element_s[-1], color="black", linestyle="--")
+
+    # plot config
+    skip_names = [
+        "Drift",
+        "ChrDrift",
+        "ExactDrift",
+        "Empty",
+        "Marker",
+        "Source",
+    ]
+
+    handles = {}
+
+    for i, element in enumerate(self):
+        el_dict = element.to_dict()
+        el_type = el_dict["type"]
+        if el_type in skip_names:
+            continue
+
+        color = get_element_color(el_type, palette=palette)
+
+        y0 = 0  # default start in y for unspecified elements
+        height = 0.5  # default height for unspecified elements
+
+        # note the sub-string matching for el_type
+        if el_type == "BeamMonitor":
+            y0 = -0.5
+            height = 1.0
+        if "Quad" in el_type:
+            height = copysign(0.8, el_dict["k"] * charge_qe)
+        if "Sbend" in el_type:
+            if ref is None:
+                height = copysign(0.8, element.rc(ref))
+            else:  # guess
+                height = copysign(0.8, el_dict["phi"])
+        # TODO: sign dependent, read m_p_scale
+        # if el_type == "Kicker":
+        #    height = copysign(0.8, el_dict["xkick"])
+
+        # plot thin elements on top of thick elements
+        zorder = 2
+        if element.ds == 0:
+            zorder = 3
+
+        patch = Rectangle(
+            (element_s[i], y0),
+            element_lengths[i],
+            height,
+            color=color,
+            alpha=0.8,
+            zorder=zorder,
+        )
+        ax.add_patch(patch)
+
+        handles[el_type] = patch
+
+    if legend:
+        labels = list(handles.keys())
+        values = list(handles.values())
+        ax.legend(
+            handles=values,
+            labels=labels,
+            bbox_to_anchor=(0.0, 1.02, 1.0, 0.102),
+            loc="lower left",
+            ncols=legend_ncols,
+            mode="expand",
+            borderaxespad=0.0,
+        )
+
+    ax.set_xlabel(r"$s$ [m]")
+
+    ax.set_ylim(-1, 1)
+    ax.set_yticks([])
+
+    ax.set_aspect(1 / 1.618)  # golden ratio
+
+    return ax

src/python/impactx/plot/__init__.py

@@ -70,6 +70,15 @@ def test_df_pandas(save_png=True):
     ]
     sim.lattice.extend(fodo)
 
+    # plot lattice survey
+    if amr.ParallelDescriptor.IOProcessor():
+        sim.lattice.plot_survey(ref=ref)
+        if save_png:
+            plt.gcf().savefig("lattice_survey.png")
+            plt.close(plt.gcf())
+        else:
+            plt.show()
+
     # simulate
     sim.track_particles()
 
@@ -80,7 +89,8 @@ def test_df_pandas(save_png=True):
         print(beam_moments)
         plt.plot(beam_moments.s, beam_moments.beta_x)
         if save_png:
-            plt.savefig("beam_moments.png")
+            plt.gcf().savefig("beam_moments.png")
+            plt.close(plt.gcf())
         else:
             plt.show()
 
@@ -113,7 +123,6 @@ def test_df_pandas(save_png=True):
 
     #   note: figure data available on MPI rank zero
     if fig is not None:
-        fig.savefig("phase_space.png")
         if save_png:
             fig.savefig("phase_space.png")
         else: