diff --git a/src/lib/plotting/animated_plot.py b/src/lib/plotting/animated_plot.py index 2e80696..acd596f 100644 --- a/src/lib/plotting/animated_plot.py +++ b/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) diff --git a/src/lib/plotting/labeler.py b/src/lib/plotting/labeler.py new file mode 100644 index 0000000..eb1c98c --- /dev/null +++ b/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) diff --git a/src/lib/plotting/plot.py b/src/lib/plotting/plot.py index f5d9ae9..28d35f7 100644 --- a/src/lib/plotting/plot.py +++ b/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))) diff --git a/src/lib/plotting/plot_info.py b/src/lib/plotting/plot_info.py index 4b76348..4b363b2 100644 --- a/src/lib/plotting/plot_info.py +++ b/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})" diff --git a/src/lib/plotting/renderer2.py b/src/lib/plotting/renderer2.py new file mode 100644 index 0000000..3e386f8 --- /dev/null +++ b/src/lib/plotting/renderer2.py @@ -0,0 +1,6 @@ +from abc import ABC, abstractmethod + + +class Renderer2(ABC): + @abstractmethod + def update(self): ... diff --git a/src/lib/plotting/setup_fig.py b/src/lib/plotting/setup_fig.py index 73c7f6e..4b01b08 100644 --- a/src/lib/plotting/setup_fig.py +++ b/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 diff --git a/tests/baseline/test_image_and_cuts.png b/tests/baseline/test_image_and_cuts.png new file mode 100644 index 0000000..3bc3672 Binary files /dev/null and b/tests/baseline/test_image_and_cuts.png differ diff --git a/tests/baseline/test_image_and_line.png b/tests/baseline/test_image_and_line.png index 7c5bb91..fbb971c 100644 Binary files a/tests/baseline/test_image_and_line.png and b/tests/baseline/test_image_and_line.png differ diff --git a/tests/test_plots.py b/tests/test_plots.py index 05849e5..01ca869 100644 --- a/tests/test_plots.py +++ b/tests/test_plots.py @@ -94,6 +94,12 @@ def test_image_and_line(): return make_plot("prt.i --bin y py=100 --nan0 --scale log --compute -v y py -w pfd::ey_ec -v y".split()) +@pytest.mark.mpl_image_compare(**MPL_KWARGS) +def test_image_and_cuts(): + """Chronotopogram of a field with some select cuts shown as lines.""" + return make_plot("pfd ey_ec -v t y --scale symlog#.001 --copy ey_ec -i y=0 -v t --copy ey_ec -i y=-1 -v t".split()) + + # --- Cross-dataset plots ---