MPL: KnownElementsList.plot()

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

sim.lattice.plot()
plt.show()
```

Author: ax3l

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,120 @@
+"""
+This file is part of ImpactX
+
+Copyright 2025 ImpactX contributors
+Authors: Axel Huebl
+License: BSD-3-Clause-LBNL
+"""
+
+from ..madx_to_impactx import read_lattice  # noqa
+
+
+def kel_plot(self, ax=None, legend=True, legend_ncols=5):
+    """Plot over s of all elements in the KnownElementsList."""
+    from math import copysign
+
+    import matplotlib.pyplot as plt
+    import numpy as np
+    from matplotlib.patches import Rectangle
+
+    ax = ax or plt.subplot(111)
+
+    element_lengths = [element.ds for element in self]
+    element_s = np.cumulative_sum(element_lengths, include_initial=True)
+
+    ax.hlines(0, 0, element_s[-1], color="black", linestyle="--")
+
+    # plot config
+    colors = {
+        "Quad": "tab:blue",
+        "ChrQuad": "tab:blue",
+        "ExactQuad": "tab:blue",
+        "SoftQuadrupole": "tab:blue",
+        "Multipole": "tab:orange",
+        "ExactMultipole": "tab:orange",
+        "Sbend": "tab:green",
+        "ExactSbend": "tab:green",
+        "CFbend": "tab:olive",
+        "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'
+
+    skip_names = [
+        "BeamMonitor",
+        "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
+
+        if el_type in colors:
+            color = colors[el_dict["type"]]
+        else:
+            color = "tab:gray"
+
+        height = 0.5
+        if "k" in el_dict:
+            height = copysign(0.8, el_dict["k"])
+        if "phi" in el_dict:
+            height = copysign(0.8, el_dict["phi"])
+
+        patch = Rectangle(
+            (element_s[i], 0),
+            element_lengths[i],
+            height,
+            color=color,
+            alpha=1.0,
+            zorder=2,
+        )
+        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
+
+
+def register_KnownElementsList_extension(kel):
+    """KnownElementsList helper methods"""
+    # register member functions for KnownElementsList
+    kel.load_file = lambda self, madx_file, nslice=1: self.extend(
+        read_lattice(madx_file, nslice)
+    )
+    kel.plot = kel_plot