MPL: KnownElementsList.plot_survey()

docs/source/usage/python.rst

@@ -623,6 +623,20 @@ 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.
+
+      A positive element strength denotes horizontal focusing (e.g. for quadrupoles) and bending to the right (for dipoles).  In general, this depends on both the sign of the field and the sign of the charge.
+
+      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,132 @@
+"""
+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.
+
+    A positive element strength denotes horizontal focusing (e.g. for quadrupoles) and bending to the right (for dipoles).  In general, this depends on both the sign of the field and the sign of the charge.
+
+    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

tests/python/test_dataframe.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: