Simplify overlay unit validation with CanvasSpec

Replace OverlayUnitKey (lossy string-based, weak cross-layer check) with
CanvasSpec — a frozen dataclass holding coord_units and value_unit as native
sc.Unit values. One validation function (validate_overlay_units) now serves
both intra-layer and cross-layer sites with full consistency checks.

ROI and annotation plotters opt out via participates_in_overlay_validation
class flag, eliminating three _extract_overlay_unit_key overrides.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: SimonHeybrock

src/ess/livedata/dashboard/plots.py

@@ -7,6 +7,7 @@
 import weakref
 from collections.abc import Callable
 from dataclasses import dataclass
+from dataclasses import field as dataclass_field
 from typing import Any, ClassVar, cast
 
 import holoviews as hv
@@ -216,117 +217,66 @@ def _compute_time_info(data: dict[str, sc.DataArray]) -> str | None:
 
 
 @dataclass(frozen=True)
-class OverlayUnitKey:
-    """Unit signature for cross-layer overlay validation.
+class CanvasSpec:
+    """Unit metadata extracted from a scipp DataArray for overlay validation.
 
-    Captures the units of all plotting dimensions and (optionally) the value
-    dimension.  Two keys are compatible when every shared kdim label has the
-    same unit string and, if both keys carry a vdim unit, those match too.
+    Captures coordinate units (keyed by dimension name) and the value unit.
+    Two overlays are compatible when they share the same coordinate dimensions
+    with matching units and matching value units.
     """
 
-    kdim_units: tuple[tuple[str, str | None], ...] = ()
-    vdim_unit: str | None = None
+    coord_units: dict[str, sc.Unit] = dataclass_field(default_factory=dict)
+    value_unit: sc.Unit = dataclass_field(default=sc.units.dimensionless)
 
     @classmethod
-    def from_data_array(
-        cls,
-        da: sc.DataArray,
-        *,
-        include_vdim: bool = True,
-    ) -> OverlayUnitKey:
-        """Build a key from a scipp DataArray using its dimension coordinates."""
-        kdim_units = tuple(
-            sorted(
-                (dim, _unit_str(da.coords[dim].unit))
-                for dim in da.dims
-                if dim in da.coords
-            )
+    def from_data_array(cls, da: sc.DataArray) -> CanvasSpec:
+        """Build from a scipp DataArray using its dimension coordinates."""
+        return cls(
+            coord_units={
+                dim: da.coords[dim].unit for dim in da.dims if dim in da.coords
+            },
+            value_unit=da.unit,
         )
-        vdim_unit = _unit_str(da.unit) if include_vdim else None
-        return cls(kdim_units=kdim_units, vdim_unit=vdim_unit)
-
-
-def _unit_str(unit: sc.Unit | None) -> str | None:
-    """Convert a scipp unit to its string representation, or None."""
-    return str(unit) if unit is not None else None
-
-
-def _validate_overlay_units(data: dict[ResultKey, sc.DataArray]) -> None:
-    """Validate that all DataArrays in an overlay share compatible units.
-
-    Raises
-    ------
-    ValueError
-        If dimension names, coordinate units, or value units do not match.
-    """
-    if len(data) < 2:
-        return
-
-    items = iter(data.items())
-    ref_key, ref = next(items)
-    ref_source = ref_key.job_id.source_name
-
-    for key, da in items:
-        source = key.job_id.source_name
-        if ref.dims != da.dims:
-            raise ValueError(
-                f"Cannot overlay '{ref_source}' and '{source}': "
-                f"dimension mismatch ({ref.dims} vs {da.dims})"
-            )
-        for dim in ref.dims:
-            if dim in ref.coords and dim in da.coords:
-                u_ref = ref.coords[dim].unit
-                u_da = da.coords[dim].unit
-                if u_ref != u_da:
-                    raise ValueError(
-                        f"Cannot overlay '{ref_source}' and '{source}': "
-                        f"unit mismatch for coordinate '{dim}' "
-                        f"({u_ref} vs {u_da})"
-                    )
-        if ref.unit != da.unit:
-            raise ValueError(
-                f"Cannot overlay '{ref_source}' and '{source}': "
-                f"value unit mismatch ({ref.unit} vs {da.unit})"
-            )
 
 
-def validate_cross_layer_units(
-    keys: list[tuple[str, OverlayUnitKey]],
+def validate_canvas_spec(
+    entries: list[tuple[str, CanvasSpec]],
 ) -> str | None:
-    """Check overlay unit keys from multiple layers for compatibility.
+    """Check that overlay entries share compatible units.
 
     Parameters
     ----------
-    keys:
-        Pairs of (layer label, unit key) for each data-carrying layer.
+    entries:
+        Pairs of (label, units) for each data-carrying element or layer.
 
     Returns
     -------
     :
         An error message if units are incompatible, or None if compatible.
     """
-    if len(keys) < 2:
+    if len(entries) < 2:
         return None
 
-    ref_label, ref_key = keys[0]
-    ref_kdims = dict(ref_key.kdim_units)
-
-    for label, key in keys[1:]:
-        other_kdims = dict(key.kdim_units)
-        for dim_label, ref_unit in ref_kdims.items():
-            if dim_label in other_kdims and ref_unit != other_kdims[dim_label]:
-                return (
-                    f"Cannot overlay layers '{ref_label}' and '{label}': "
-                    f"unit mismatch for coordinate '{dim_label}' "
-                    f"({ref_unit} vs {other_kdims[dim_label]})"
-                )
-        if ref_key.vdim_unit is not None and key.vdim_unit is not None:
-            if ref_key.vdim_unit != key.vdim_unit:
+    ref_label, ref = entries[0]
+    for label, other in entries[1:]:
+        if ref.coord_units.keys() != other.coord_units.keys():
+            return (
+                f"Cannot overlay '{ref_label}' and '{label}': "
+                f"dimension mismatch "
+                f"({set(ref.coord_units)} vs {set(other.coord_units)})"
+            )
+        for dim, ref_unit in ref.coord_units.items():
+            if ref_unit != other.coord_units[dim]:
                 return (
-                    f"Cannot overlay layers '{ref_label}' and '{label}': "
-                    f"value unit mismatch "
-                    f"({ref_key.vdim_unit} vs {key.vdim_unit})"
+                    f"Cannot overlay '{ref_label}' and '{label}': "
+                    f"unit mismatch for coordinate '{dim}' "
+                    f"({ref_unit} vs {other.coord_units[dim]})"
                 )
+        if ref.value_unit != other.value_unit:
+            return (
+                f"Cannot overlay '{ref_label}' and '{label}': "
+                f"value unit mismatch ({ref.value_unit} vs {other.value_unit})"
+            )
     return None
 
 
@@ -338,6 +288,8 @@ class Plotter:
     This enables efficient polling-based update detection.
     """
 
+    participates_in_overlay_validation: ClassVar[bool] = True
+
     def __init__(
         self,
         *,
@@ -361,7 +313,7 @@ def __init__(
         """
         self._normalize_to_rate = normalize_to_rate
         self._cached_state: Any | None = None
-        self._overlay_unit_key: OverlayUnitKey | None = None
+        self.canvas_spec: CanvasSpec | None = None
         self._presenters: weakref.WeakSet[PresenterBase] = weakref.WeakSet()
         self.autoscaler_kwargs = kwargs
         self.autoscalers: dict[ResultKey, Autoscaler] = {}
@@ -374,25 +326,6 @@ def __init__(
         # with responsive mode in Panel containers (upstream bug).
         self._sizing_opts: dict[str, Any] = {'responsive': True}
 
-    @property
-    def overlay_unit_key(self) -> OverlayUnitKey | None:
-        """Unit signature from the last compute(), for cross-layer validation."""
-        return self._overlay_unit_key
-
-    def _extract_overlay_unit_key(
-        self, data: dict[ResultKey, sc.DataArray]
-    ) -> OverlayUnitKey | None:
-        """Extract an overlay unit key from input data.
-
-        Default implementation uses dimension coordinates and value unit from the
-        first DataArray. Subclasses with non-standard coordinate structures (e.g.,
-        ROI plotters) should override this method.
-        """
-        if not data:
-            return None
-        da = next(iter(data.values()))
-        return OverlayUnitKey.from_data_array(da)
-
     @staticmethod
     def _make_tick_opts(tick_params: TickParams | None) -> dict[str, Any]:
         """
@@ -566,13 +499,19 @@ def compute(
         if self._normalize_to_rate:
             data = {key: _normalize_to_rate(da) for key, da in data.items()}
 
-        self._overlay_unit_key = self._extract_overlay_unit_key(data)
-
         resolver = title_resolver or TitleResolver()
         plots: list[hv.Element] = []
         try:
-            if self.layout_params.combine_mode == 'overlay':
-                _validate_overlay_units(data)
+            if self.participates_in_overlay_validation and data:
+                entries = [
+                    (key.job_id.source_name, CanvasSpec.from_data_array(da))
+                    for key, da in data.items()
+                ]
+                if self.layout_params.combine_mode == 'overlay':
+                    error = validate_canvas_spec(entries)
+                    if error is not None:
+                        raise ValueError(error)
+                self.canvas_spec = entries[0][1]
             for data_key, da in data.items():
                 label = resolver.get_legend_label(
                     data_key.job_id.source_name, data_key.output_name

src/ess/livedata/dashboard/roi_readback_plots.py

@@ -21,23 +21,7 @@
 )
 from ess.livedata.config.workflow_spec import ResultKey
 
-from .plots import OverlayUnitKey, Plotter, _unit_str
-
-
-def _roi_overlay_unit_key(da: sc.DataArray) -> OverlayUnitKey:
-    """Build an overlay unit key from ROI coordinate metadata.
-
-    ROI DataArrays use named coords ('x', 'y') rather than dimension
-    coordinates, and have no meaningful value unit.
-    """
-    kdim_units = tuple(
-        sorted(
-            (name, _unit_str(da.coords[name].unit))
-            for name in ('x', 'y')
-            if name in da.coords
-        )
-    )
-    return OverlayUnitKey(kdim_units=kdim_units, vdim_unit=None)
+from .plots import Plotter
 
 
 class ROIReadbackStyle(pydantic.BaseModel):
@@ -87,6 +71,8 @@ class RectanglesReadbackPlotter(Plotter):
     rectangles with per-shape colors based on the ROI index.
     """
 
+    participates_in_overlay_validation = False
+
     def __init__(self, params: RectanglesReadbackParams) -> None:
         super().__init__()
         self._params = params
@@ -97,14 +83,6 @@ def from_params(cls, params: RectanglesReadbackParams) -> RectanglesReadbackPlot
         """Create plotter from params."""
         return cls(params)
 
-    def _extract_overlay_unit_key(
-        self, data: dict[ResultKey, sc.DataArray]
-    ) -> OverlayUnitKey | None:
-        if not data:
-            return None
-        da = next(iter(data.values()))
-        return _roi_overlay_unit_key(da)
-
     def plot(
         self, data: sc.DataArray, data_key: ResultKey, *, label: str = '', **kwargs
     ) -> hv.Rectangles:
@@ -201,6 +179,8 @@ class PolygonsReadbackPlotter(Plotter):
     polygons with per-shape colors based on the ROI index.
     """
 
+    participates_in_overlay_validation = False
+
     def __init__(self, params: PolygonsReadbackParams) -> None:
         super().__init__()
         self._params = params
@@ -211,14 +191,6 @@ def from_params(cls, params: PolygonsReadbackParams) -> PolygonsReadbackPlotter:
         """Create plotter from params."""
         return cls(params)
 
-    def _extract_overlay_unit_key(
-        self, data: dict[ResultKey, sc.DataArray]
-    ) -> OverlayUnitKey | None:
-        if not data:
-            return None
-        da = next(iter(data.values()))
-        return _roi_overlay_unit_key(da)
-
     def plot(
         self, data: sc.DataArray, data_key: ResultKey, *, label: str = '', **kwargs
     ) -> hv.Polygons:

src/ess/livedata/dashboard/roi_request_plots.py

@@ -48,7 +48,7 @@
     get_roi_mapper,
 )
 
-from .plots import OverlayUnitKey, Plotter, PresenterBase
+from .plots import Plotter, PresenterBase
 from .static_plots import Color, LineDash, RectanglesCoordinates
 
 if TYPE_CHECKING:
@@ -533,6 +533,8 @@ class BaseROIRequestPlotter(Plotter, ABC, Generic[ROIType, ParamsType, Converter
     skip logic, publishing) is handled via a closure-based edit callback.
     """
 
+    participates_in_overlay_validation = False
+
     def __init__(
         self,
         params: ParamsType,
@@ -626,16 +628,6 @@ def compute(
             else None
         )
 
-        # Build overlay unit key from ROI coordinate units
-        kdim_units = tuple(
-            sorted(
-                (name, unit)
-                for name, unit in [('x', self._x_unit), ('y', self._y_unit)]
-                if unit is not None
-            )
-        )
-        self._overlay_unit_key = OverlayUnitKey(kdim_units=kdim_units, vdim_unit=None)
-
         # Forward data (presenter may use in future)
         self._set_cached_state(data)
         return data

src/ess/livedata/dashboard/widgets/plot_grid_tabs.py

@@ -35,7 +35,7 @@
     SubscriptionId,
 )
 from ..plot_params import PlotAspectType, StretchMode
-from ..plots import OverlayUnitKey, validate_cross_layer_units
+from ..plots import CanvasSpec, validate_canvas_spec
 from ..save_filename import build_save_filename_from_cell, make_save_filename_hook
 from ..session_layer import SessionLayer
 from ..session_updater import SessionUpdater
@@ -941,7 +941,7 @@ def _get_session_composed_plot(
             Composed plot from session DMaps/elements, or None if none available.
         """
         plots = []
-        unit_keys: list[tuple[str, OverlayUnitKey]] = []
+        overlay_entries: list[tuple[str, CanvasSpec]] = []
         has_layout = False
         for layer in cell.layers:
             layer_id = layer.layer_id
@@ -969,17 +969,18 @@ def _get_session_composed_plot(
                     has_layout = True
                 plots.append(dmap)
 
-                # Collect unit key for cross-layer validation
+                # Collect overlay units for cross-layer validation
+                # (only data-carrying plotters set canvas_spec)
                 if state is not None and state.plotter is not None:
-                    key = state.plotter.overlay_unit_key
-                    if key is not None:
-                        unit_keys.append((layer.config.plot_name, key))
+                    units = state.plotter.canvas_spec
+                    if units is not None:
+                        overlay_entries.append((layer.config.plot_name, units))
 
         if not plots:
             return None
 
         # Validate cross-layer unit compatibility before composing
-        error = validate_cross_layer_units(unit_keys)
+        error = validate_canvas_spec(overlay_entries)
         if error is not None:
             return hv.Text(0.5, 0.5, f"Error: {error}").opts(
                 text_align='center', text_baseline='middle'