refactor draw elements

eleganttools/plot.py

@@ -1,5 +1,7 @@
-import numpy as np
+from itertools import groupby
+
 import matplotlib.pyplot as plt
+import numpy as np
 from matplotlib.offsetbox import AnchoredOffsetbox, TextArea, VPacker
 
 DOUBLET_NAMES = [
@@ -42,83 +44,83 @@
 }
 
 
-def draw_elements(ax, data, *, s_lim=None, labels=True):
+def draw_elements(ax, data, *, labels=True):
     """Draw lattice on matplotlib axes."""
-    s = np.array(data["s"], dtype=np.float64)
-    element_type = data["ElementType"]
-    element_name = data["ElementName"]
-    s0, s1 = s_lim if s_lim is not None else (s[0], s[-1])
-    on_top = {"QUAD", "KQUAD"}
-
-    # get space for labels and rectangles
+    x_min, x_max = ax.get_xlim()
     y_min, y_max = ax.get_ylim()
     rect_height = 0.05 * (y_max - y_min)
+    y_max += rect_height
+    ax.set_ylim(y_min, y_max)
+
+    positions = data["s"]
+    element_types = data["ElementType"]
+    element_names = data["ElementName"]
+
+    sign = -1
+    i = -1
+    start = end = 0
+    for element_name, group in groupby(element_names):
+        start = end
+        i += len(list(group))
+        end = positions[i]
+        if end <= x_min:
+            continue
+        elif start >= x_max:
+            break
 
-    i0 = np.argmax(s >= s0)
-    i1 = np.argmax(s >= s1)
-    start = s0
-    for i in range(i0, i1 + 1):
-        if i > i0:  # save start if previous element was something else
-            if element_type[i] != element_type[i - 1]:
-                start = s[i - 1]
-        if i < i1:  # skip if next element of same type
-            if element_type[i] == element_type[i + 1]:
-                continue
-
-        end = np.min((s[i], s1))
-        length = end - start
-        face_color = COLOR_MAP.get(element_type[i])
-        if face_color is None:
+        try:
+            color = COLOR_MAP[element_types[i]]
+        except KeyError:
             continue
 
-        retangle = plt.Rectangle(
-            (start, y_max - 0.5 * rect_height),
-            length,
-            rect_height,
-            # ec="k",
-            facecolor=face_color,
-            clip_on=False,
-            zorder=10,
+        ax.add_patch(
+            plt.Rectangle(
+                (start, y_max - 0.5 * rect_height),
+                min(end, x_max) - max(start, x_min),
+                rect_height,
+                facecolor=color,
+                clip_on=False,
+                zorder=10,
+            )
         )
-        ax.add_patch(retangle)
         if labels:
-            sign = ((element_type[i] in on_top) << 1) - 1
+            sign = -sign
             plt.annotate(
-                element_name[i],
-                xy=((end + start) / 2, y_max + sign * rect_height),
+                element_name,
+                xy=((start + end) / 2, y_max + sign * rect_height),
                 fontsize=5,
                 va="center",
                 ha="center",
                 annotation_clip=False,
-                zorder=100,
+                zorder=11,
             )
 
 
-def axis_labels(ax, *, yscale=1, eta_x_scale=10):
+def axis_labels(ax, *, eta_x_scale=10):
     plt.xlabel("s / m")
-    ybox1 = TextArea(
-        "       $\\eta_x / {0}".format(int(100 / eta_x_scale)) + "\\mathrm{cm}$",
-        textprops=dict(color=green, rotation=90, ha="left", va="center"),
-    )
-    ybox2 = TextArea(
-        "  $\\beta_y / \\mathrm{m}$",
-        textprops=dict(color=blue, rotation=90, ha="left", va="center"),
-    )
-    ybox3 = TextArea(
-        "$\\beta_x / \\mathrm{m}$",
-        textprops=dict(color=red, rotation=90, ha="left", va="center"),
-    )
-    ybox = VPacker(children=[ybox1, ybox2, ybox3], align="bottom", pad=0, sep=5)
-    anchored_ybox = AnchoredOffsetbox(
-        loc=8,
-        child=ybox,
-        pad=0.0,
-        frameon=False,
-        bbox_to_anchor=(-0.08 * yscale, 0.15),
-        bbox_transform=plt.gca().transAxes,
-        borderpad=0.0,
+    text_areas = [
+        TextArea(
+            rf"{eta_x_scale} $\eta_x$ / m",
+            textprops=dict(color=green, rotation=90),
+        ),
+        TextArea(
+            r"$\beta_y$ / m",
+            textprops=dict(color=blue, rotation=90),
+        ),
+        TextArea(
+            r"$\beta_x$ / m",
+            textprops=dict(color=red, rotation=90),
+        ),
+    ]
+    ax.add_artist(
+        AnchoredOffsetbox(
+            child=VPacker(children=text_areas, align="bottom", pad=0, sep=20),
+            loc="center left",
+            bbox_to_anchor=(-0.125, 0, 1.125, 1),
+            bbox_transform=ax.transAxes,
+            frameon=False,
+        )
     )
-    ax.add_artist(anchored_ybox)
 
 
 def plot_bessy2_section(data, section_name, ax=None):
@@ -149,6 +151,6 @@ def plot_bessy2_section(data, section_name, ax=None):
     ax.annotate(label, ((s1 + s0) / 2.0, y1 - y_span * 0.1), fontsize=10, ha="center")
     ax.yaxis.grid(alpha=0.3, zorder=0)
 
-    draw_elements(ax, data, s_lim=(s0, s1))
     axis_labels(ax)
     ax.set_xlim(s0, s1)
+    draw_elements(ax, data)