Fix labeling in multiplots

src/lib/plotting/animated_plot.py

@@ -32,6 +32,8 @@ def _initialize(self):
     def _next_frame(self, frame: int):
         for renderer in self.renderers:
             renderer.update_plot_info(frame)
+        for renderer in self.renderer2s:
+            renderer.update()
         self.post_update_fig(DrawMessage(plot_info=self.renderers[0].plot_info, axes=self.fig.axes[0], frame_data=self.renderers[0]._get_data_at_frame(frame)))
         print_progress(frame, self.n_frames)
 

src/lib/plotting/labeler.py

@@ -0,0 +1,92 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Callable
+
+from lib.plotting.plot_info import PlotInfo
+from lib.plotting.renderer2 import Renderer2
+
+
+@dataclass
+class TreeLabeler(Renderer2):
+    """Manages the labels associated with one or more datasets within a figure. A label comprises an optional subject
+    (variable name) and any number of sublabels (e.g. scalar coordinates). When multiple datasets are plotted within
+    the same figure, common label components can be "factored out" to a higher label location, e.g. from a legend to
+    an axis title. Label locations are well-described by a tree structure, where common label components propagate
+    from the leaves to the root."""
+
+    set_text: Callable[[str], None]
+    source: PlotInfo | None = None
+
+    children: list[TreeLabeler] = field(default_factory=list, init=False)
+    parent: TreeLabeler | None = field(default=None, init=False)
+
+    _subject: str | None = field(default=None, init=False)
+    _sublabels: list[str] = field(default_factory=list, init=False)
+
+    def add_child(self, child: TreeLabeler):
+        assert child.parent is None
+        child.parent = self
+        self.children.append(child)
+
+    def update(self):
+        """Propagate updates up to the root labeler, which makes sure that everyone rebuilds and then everyone updates text."""
+        if self.parent:
+            return self.parent.update()
+        else:
+            self._rebuild()
+            self._update_text()
+
+    def _rebuild(self):
+        for child in self.children:
+            child._rebuild()
+
+        child_subjects = {child._subject for child in self.children}
+        all_child_sublabels = {sublabel: None for child in self.children for sublabel in child._sublabels}  # use dict to preserve insertion order
+        common_child_sublabels = {sublabel: None for sublabel in all_child_sublabels if all(sublabel in child._sublabels for child in self.children)}  # use dict to preserve insertion order
+
+        if self.source:
+            self._subject = self.source.subject
+            self._sublabels = self.source.get_sublabels()
+
+            # only eliminate child subjects + sublabels if every child shares the root subject and all its sublabels
+            has_common_subject = {self._subject} == child_subjects
+            has_common_sublabels = set(self._sublabels) <= set(common_child_sublabels.keys())
+
+            if has_common_subject and has_common_sublabels:
+                self._eliminate_subject()
+                self._eliminate_common_sublabels()
+
+        else:
+            # no source -> lift all common subject and/or sublabels independently
+
+            if len(child_subjects) == 1:
+                self._subject = child_subjects.pop()
+                self._eliminate_subject()
+            else:
+                self._subject = None
+
+            self._sublabels = list(common_child_sublabels.keys())
+            self._eliminate_common_sublabels()
+
+    def _update_text(self):
+        # set_text intelligently checks if the text actually changes or not
+        self.set_text(self._get_label())
+        for child in self.children:
+            child._update_text()
+
+    def _get_label(self) -> str:
+        sublabels = ", ".join(self._sublabels)
+
+        if self._subject and sublabels:
+            return f"{self._subject} ({sublabels})"
+        return self._subject or sublabels
+
+    def _eliminate_subject(self):
+        for child in self.children:
+            child._subject = None
+
+    def _eliminate_common_sublabels(self):
+        for child in self.children:
+            for sublabel in self._sublabels:
+                child._sublabels.remove(sublabel)

src/lib/plotting/plot.py

@@ -28,7 +28,7 @@ def _initialize(self):
             return
         self._initialized = True
 
-        self.fig = setup_fig([r.plot_info for r in self.renderers])
+        self.fig, self.renderer2s = setup_fig([r.plot_info for r in self.renderers])
         # TODO hooks should be per-renderer; for now, just apply them to the 1st one
         self.post_init_fig(DrawMessage(plot_info=self.renderers[0].plot_info, axes=self.fig.axes[0], frame_data=self.renderers[0]._get_data_at_frame(0)))
 

src/lib/plotting/plot_info.py

@@ -50,8 +50,11 @@ def get_coord_label(self, dim: DimKey) -> Latex:
         maybe_space = "\\ " if unit else ""
         return Latex(f"{display} = {coord_val:.3f}{maybe_space}{unit}")
 
+    def get_sublabels(self) -> list[str]:
+        return [f"${self.get_coord_label(dim)}$" for dim in self.scalar_coord_values]
+
     def get_title(self) -> str:
-        coord_labels_str = ", ".join(f"${self.get_coord_label(dim)}$" for dim in self.scalar_coord_values)
+        coord_labels_str = ", ".join(self.get_sublabels())
 
         if self.subject and coord_labels_str:
             return f"{self.subject} ({coord_labels_str})"

src/lib/plotting/renderer2.py

@@ -0,0 +1,6 @@
+from abc import ABC, abstractmethod
+
+
+class Renderer2(ABC):
+    @abstractmethod
+    def update(self): ...

src/lib/plotting/setup_fig.py

@@ -11,7 +11,9 @@
 from matplotlib.text import Text
 
 from lib.plotting import plt_util
-from lib.plotting.plot_info import DimKey, ImageInfo, LineInfo, PlotInfo, PlotInfo2D, PolarMeshInfo, ScatterInfo
+from lib.plotting.labeler import TreeLabeler
+from lib.plotting.plot_info import ImageInfo, LineInfo, PlotInfo, PlotInfo2D, PolarMeshInfo, ScatterInfo
+from lib.plotting.renderer2 import Renderer2
 
 type AxesIdx = tuple[int, int]
 
@@ -93,6 +95,9 @@ def __call__(self, *_):
 
 
 class AxesManager(ABC):
+    @abstractmethod
+    def get_renderers(self) -> list[Renderer2]: ...
+
     @abstractmethod
     def setup(self): ...
 
@@ -118,12 +123,11 @@ def __init__(self, ax: A, info: PI):
         self.info = info
 
     def setup_title(self):
-        update_title = UpdateText(self.ax.title, self.info)
-        self.info._setter_callbacks["subject"] = update_title
-        self.info._setter_callbacks["dim_displays"] = update_title
-        self.info._setter_callbacks["dim_units"] = update_title
-        self.info._setter_callbacks["scalar_coord_values"] = update_title
-        update_title()
+        self.labeler = TreeLabeler(self.ax.title.set_text, self.info)
+        self.labeler.update()
+
+    def get_renderers(self):
+        return [self.labeler]
 
 
 class AxesManagerSingle2D[PI2D: PlotInfo2D](AxesManagerSingle[Axes, PI2D]):
@@ -235,76 +239,25 @@ class AxesManagerMultiLine(AxesManager):
     def __init__(self, ax: Axes, infos: list[LineInfo]):
         self.ax = ax
         self.infos = infos
-
-        self.common_coord_dims: set[DimKey] = set()
-        self.unique_coord_dimss: list[set[DimKey]] = [set() for _ in self.infos]
-        self._update_common_scalar_coordinates()
-
-        self.common_subject: str | None = None
-        self._update_common_subject()
-
         self.lines: list[Line2D] = []  # populated later
 
-    def _update_common_scalar_coordinates(self):
-        self.common_coord_dims.clear()
-        for unique_dims in self.unique_coord_dimss:
-            unique_dims.clear()
-
-        all_dims: set[DimKey] = {dim for info in self.infos for dim in info.scalar_coord_values}
-
-        for dim in all_dims:
-            if all(dim in info.scalar_coord_values for info in self.infos) and len({info.get_coord_label(dim) for info in self.infos}) == 1:
-                self.common_coord_dims.add(dim)
-            else:
-                for info, unique_dims in zip(self.infos, self.unique_coord_dimss):
-                    if dim in info.scalar_coord_values:
-                        unique_dims.add(dim)
-
-    def _update_common_subject(self):
-        self.common_subject = _one_or_none(info.subject for info in self.infos)
-
-    def _get_title(self) -> str:
-        # by definition, it shouldn't matter which info we use to construct this string
-        coord_labels_str = ", ".join(self.infos[0].get_coord_label(dim).maybe_with_dollars() for dim in self.common_coord_dims)
-
-        if self.common_subject and coord_labels_str:
-            return f"{self.common_subject} ({coord_labels_str})"
-        return self.common_subject or coord_labels_str
-
-    def _get_legend_labels(self) -> list[str]:
-        legend_labels: list[str] = []
-
-        for info, unique_dims in zip(self.infos, self.unique_coord_dimss):
-            coord_labels_str = ", ".join(info.get_coord_label(dim).maybe_with_dollars() for dim in unique_dims)
-
-            if self.common_subject:
-                legend_labels.append(coord_labels_str)
-            elif info.subject and coord_labels_str:
-                legend_labels.append(f"{info.subject} ({coord_labels_str})")
-            else:
-                legend_labels.append(info.subject or coord_labels_str)
-
-        return legend_labels
-
-    def _update_title_and_legend(self, *_):
-        self._update_common_scalar_coordinates()
-        self._update_common_subject()
-        self.ax.set_title(self._get_title())
-        for line, label in zip(self.lines, self._get_legend_labels()):
-            line.set_label(label)
+    def get_renderers(self):
+        return [self.labeler]
 
     def setup(self):
-        self.setup_title()
         self.setup_labels()
         self.setup_data()
+        self.setup_title()  # after data, to make sure lines is populated
         self.setup_scales()
         self.setup_bounds()
 
-        for info in self.infos:
-            info._setter_callbacks["scalar_coord_values"] = self._update_title_and_legend
-
     def setup_title(self):
-        self.ax.set_title(self._get_title())
+        self.labeler = TreeLabeler(self.ax.title.set_text)
+        for info, line in zip(self.infos, self.lines):
+            line_labeler = TreeLabeler(line.set_label, info)
+            self.labeler.add_child(line_labeler)
+        self.labeler.update()
+        self.ax.legend()
 
     def setup_labels(self):
         x_labels = [info.get_dim_label(info.x_dim) for info in self.infos]
@@ -340,15 +293,13 @@ def setup_bounds(self):
         self.ax.set_ybound(*find_widest_bounds(info.dim_bounds[info.y_dim] for info in self.infos))
 
     def setup_data(self):
-        for info, label in zip(self.infos, self._get_legend_labels()):
-            [line] = self.ax.plot(info.x_data, info.y_data, linestyle=info.line_style, scalex=False, scaley=False, label=label)
+        for info in self.infos:
+            [line] = self.ax.plot(info.x_data, info.y_data, linestyle=info.line_style, scalex=False, scaley=False)
             info._setter_callbacks["x_data"] = line.set_xdata
             info._setter_callbacks["y_data"] = line.set_ydata
             info._setter_callbacks["line_style"] = line.set_linestyle
             self.lines.append(line)
 
-        self.ax.legend()
-
 
 class AxesManagerImageAndLines(AxesManager):
     def __init__(self, ax: Axes, image_info: ImageInfo, line_infos: list[LineInfo]):
@@ -357,67 +308,27 @@ def __init__(self, ax: Axes, image_info: ImageInfo, line_infos: list[LineInfo]):
         self.image_info = image_info
         self.line_infos = line_infos
         self.infos: list[PlotInfo2D] = [image_info, *line_infos]
-
-        self.common_coord_dims: set[DimKey] = set()
-        self.unique_coord_dimss: list[set[DimKey]] = [set() for _ in self.line_infos]
-        self._update_common_scalar_coordinates()
-
         self.lines: list[Line2D] = []  # populated later
 
-    def _update_common_scalar_coordinates(self):
-        self.common_coord_dims.clear()
-        for unique_dims in self.unique_coord_dimss:
-            unique_dims.clear()
-
-        all_dims: set[DimKey] = {dim for info in self.infos for dim in info.scalar_coord_values}
-
-        for dim in all_dims:
-            if all(dim in info.scalar_coord_values for info in self.infos) and len({info.get_coord_label(dim) for info in self.infos}) == 1:
-                self.common_coord_dims.add(dim)
-            else:
-                for info, unique_dims in zip(self.infos, self.unique_coord_dimss):
-                    if dim in info.scalar_coord_values:
-                        unique_dims.add(dim)
-
-    def _get_title(self) -> str:
-        # by definition, it shouldn't matter which info we use to construct this string
-        coord_labels_str = ", ".join(self.line_infos[0].get_coord_label(dim).maybe_with_dollars() for dim in self.common_coord_dims)
-
-        if self.image_info.subject and coord_labels_str:
-            return f"{self.image_info.subject} ({coord_labels_str})"
-        return self.image_info.subject or coord_labels_str
-
-    def _get_legend_labels(self) -> list[str]:
-        legend_labels: list[str] = []
-
-        for info, unique_dims in zip(self.line_infos, self.unique_coord_dimss):
-            coord_labels_str = ", ".join(info.get_coord_label(dim).maybe_with_dollars() for dim in unique_dims)
-
-            if info.subject and coord_labels_str:
-                legend_labels.append(f"{info.subject} ({coord_labels_str})")
-            else:
-                legend_labels.append(info.subject or coord_labels_str)
-
-        return legend_labels
-
-    def _update_title_and_legend(self, *_):
-        self._update_common_scalar_coordinates()
-        self.image_ax.set_title(self._get_title())
-        for line, label in zip(self.lines, self._get_legend_labels()):
-            line.set_label(label)
+    def get_renderers(self):
+        return [self.labeler]
 
     def setup(self):
-        self.setup_title()
         self.setup_labels()
         self.setup_data()
+        self.setup_title()  # after data to get line info and cbar
         self.setup_scales()
         self.setup_bounds()
 
-        for info in self.infos:
-            info._setter_callbacks["scalar_coord_values"] = self._update_title_and_legend
-
     def setup_title(self):
-        self.image_ax.set_title(self._get_title())
+        self.labeler = TreeLabeler(self.image_ax.title.set_text)
+
+        self.labeler.add_child(TreeLabeler(self.cbar.set_label, self.image_info))
+        for info, line in zip(self.line_infos, self.lines):
+            self.labeler.add_child(TreeLabeler(line.set_label, info))
+
+        self.labeler.update()
+        self.line_ax.legend()
 
     def setup_labels(self):
         x_labels = [info.get_dim_label(info.x_dim) for info in self.infos]
@@ -468,22 +379,21 @@ def setup_data(self):
         )
         self.image_info._setter_callbacks["data"] = image.set_data
 
-        self.image_ax.figure.colorbar(image)
+        self.cbar = self.image_ax.figure.colorbar(image)
         data_lower, data_upper = self.image_info.dim_bounds[self.image_info.color_dim]
         plt_util.update_cbar(image, data_min_override=data_lower, data_max_override=data_upper)
 
-        for info, label in zip(self.line_infos, self._get_legend_labels()):
-            [line] = self.line_ax.plot(info.x_data, info.y_data, linestyle=info.line_style, scalex=False, scaley=False, label=label)
+        for info in self.line_infos:
+            [line] = self.line_ax.plot(info.x_data, info.y_data, linestyle=info.line_style, scalex=False, scaley=False)
             info._setter_callbacks["x_data"] = line.set_xdata
             info._setter_callbacks["y_data"] = line.set_ydata
             info._setter_callbacks["line_style"] = line.set_linestyle
             self.lines.append(line)
 
-        self.line_ax.legend()
-
 
-def setup_fig(plot_infos: list[PlotInfo]) -> Figure:
+def setup_fig(plot_infos: list[PlotInfo]) -> tuple[Figure, list[Renderer2]]:
     figure = plt.figure(layout="constrained")
+    renderers = []
 
     for ax, infos in _setup_axes(figure, plot_infos).values():
         manager: AxesManager
@@ -510,5 +420,6 @@ def setup_fig(plot_infos: list[PlotInfo]) -> Figure:
                 raise NotImplementedError("don't yet support multiple non-line plots per axes")
 
         manager.setup()
+        renderers += manager.get_renderers()
 
-    return figure
+    return figure, renderers