Thick elements and labels from list in FloorPlot

docs/usage.md

@@ -22,6 +22,7 @@ Xsuite is not an explicit dependency, rather an API assumption on available attr
 
 ```python
 import xplt
+xplt.apply_style()  # use our matplotlib style sheet
 import numpy as np
 import pandas as pd
 

xplt/__init__.py

@@ -38,3 +38,7 @@
     hooks.register_pint_options()
 except:
     pass
+
+
+def apply_style():
+    mpl.style.use("xplt.xplt")

xplt/hooks.py

@@ -40,18 +40,6 @@ def register_matplotlib_options():
         cmap_r.name = cmap.name + "_r"
         mpl.cm.register_cmap(cmap=cmap_r)
 
-    # set rcParams
-    mpl.rcParams.update(
-        {
-            "figure.constrained_layout.use": True,
-            "legend.fontsize": "x-small",
-            "legend.title_fontsize": "small",
-            "grid.color": "#DDD",
-            "axes.prop_cycle": mpl.cycler(color=petroff_colors),
-            # 'image.cmap': cmap_petroff_gradient,
-        }
-    )
-
 
 def register_pint_options():
     """Register default options for pint"""

xplt/line.py

@@ -205,10 +205,10 @@ def __init__(
             projection (str): The projection to use: A pair of coordinates ('XZ', 'ZY' etc.)
             line (xtrack.Line): Line data with additional information about elements.
                 Use this to have colored boxes of correct size etc.
-            boxes (None | bool | str | dict): Config option for showing colored boxes for elements. See below.
+            boxes (None | bool | str | iterable | dict): Config option for showing colored boxes for elements. See below.
                 Detailed options can be "length" and all options suitable for a patch,
                 such as "color", "alpha", etc.
-            labels (None | bool | str | dict): Config option for showing labels for elements. See below.
+            labels (None | bool | str | iterable | dict): Config option for showing labels for elements. See below.
                 Detailed options can be "text" (e.g. "Dipole {name}" where name will be
                 replaced with the element name) and all options suitable for an annotation,
                 such as "color", "alpha", etc.
@@ -221,6 +221,7 @@ def __init__(
             - None: Use good defaults.
             - A bool: En-/disable option for all elements (except drifts).
             - A str (regex): Filter by element name.
+            - A list, tuple or numpy array: Filter by any of the given element names
             - A dict: Detailed options to apply for each element in the form of
               `{"regex": {...}}`. For each matching element name, the options are used.
 
@@ -247,6 +248,11 @@ def __init__(
         self.ignore = [ignore] if isinstance(ignore, str) else ignore
         self.element_width = element_width
 
+        if isinstance(self.boxes, (list, tuple, np.ndarray)):
+            self.boxes = "|".join(["^" + ll + "$" for ll in self.boxes])
+        if isinstance(self.labels, (list, tuple, np.ndarray)):
+            self.labels = "|".join(["^" + ll + "$" for ll in self.labels])
+
         # Create plot
         self.ax.set(
             xlabel=self.label_for(self.projection[0]), ylabel=self.label_for(self.projection[1])
@@ -272,7 +278,7 @@ def update(self, survey, line=None, autoscale=False):
 
         Args:
             survey (Any): Survey data.
-            line (xtrack.Line): Line data.
+            line (None | xtrack.Line): Line data.
             autoscale (bool): Whether or not to perform autoscaling on all axes
 
         Returns:
@@ -332,7 +338,8 @@ def update(self, survey, line=None, autoscale=False):
             legend_entries = []
             for i, (x, y, rt, name, arc) in enumerate(zip(X, Y, R, NAME, BEND)):
                 drift_length = get(survey, "drift_length", None)
-                if drift_length is not None and drift_length[i] > 0:
+                is_thick = line is not None and name in line.element_dict and line[name].isthick
+                if drift_length is not None and drift_length[i] > 0 and not is_thick:
                     continue  # skip drift spaces
                 if self.ignore is not None:
                     if np.any([re.match(pattern, name) is not None for pattern in self.ignore]):
@@ -346,6 +353,9 @@ def update(self, survey, line=None, autoscale=False):
                 element = line.element_dict.get(name) if line is not None else None
                 order = get(element, "order", None)
                 order = get(survey, "order", {i: order})[i]
+                if line is not None and name in line.element_dict:
+                    order = ORDER_NAMED_ELEMENTS.get(line[name].__class__.__name__, order)
+
                 length = get(element, "length", None)
                 length = get(survey, "length", {i: length})[i]
                 if length is not None:
@@ -381,6 +391,21 @@ def update(self, survey, line=None, autoscale=False):
                     else:
                         legend_entries.append(box_style.get("label"))
 
+                    # Handle thick elements
+                    if is_thick and i + 1 < len(X):
+                        # Find the center of the arc
+                        x_mid = 0.5 * (x + X[i + 1])
+                        y_mid = 0.5 * (y + Y[i + 1])
+                        dr = np.array([X[i + 1] - x, Y[i + 1] - y, 0])
+                        dn = np.cross(dr, [0, 0, 1])
+                        dn /= np.linalg.norm(dn)
+                        d = np.linalg.norm(dr) / 2
+                        sin_theta = np.abs(d * arc / length)
+                        dh = d * sin_theta
+                        p_center = np.array([x_mid, y_mid, 0]) - helicity * dh * dn
+                        x = p_center[0]
+                        y = p_center[1]
+
                     if length > 0 and arc:
                         # bending elements as wedge
                         rho = length / arc
@@ -515,6 +540,14 @@ def _get_config(config, name, **default):
             return default
 
 
+# Known class names from xtrack and their order
+ORDER_NAMED_ELEMENTS = {
+    "Bend": 0,
+    "Quadrupole": 1,
+    "Sextupole": 2,
+    "Octupole": 3,
+}
+
 ## Restrict star imports to local namespace
 __all__ = [
     name

xplt/xplt.mplstyle

@@ -0,0 +1,5 @@
+figure.constrained_layout.use: True
+legend.fontsize: x-small
+legend.title_fontsize: small
+grid.color: "#DDD"
+axes.prop_cycle: cycler('color', ['pet0', 'pet1', 'pet2', 'pet3', 'pet4', 'pet5', 'pet6', 'pet7', 'pet8', 'pet9'])