diffract-core 0.2.1__py3-none-any.whl

diffract/viz/plots/subplots/grid.py

@@ -0,0 +1,214 @@
+"""Grid plot composition for `viz` Plot objects."""
+
+from __future__ import annotations
+
+import inspect
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
+from plotly.subplots import make_subplots
+
+from diffract.viz.plots.base.plot import Plot
+from diffract.viz.styling import Theme, apply_theme
+
+from .coloraxis import (
+    ColoraxisRegistry,
+    apply_coloraxis_configs,
+    distribute_colorbars,
+    extract_coloraxis_share_keys,
+    extract_legacy_coloraxis_name,
+    remap_coloraxis,
+)
+from .layout import add_figure_to_subplot
+from .spec import VALID_SESSION_FILTER_KEYS, VALUE_FILTER_SENTINEL, SubplotSpec
+
+if TYPE_CHECKING:
+    import plotly.graph_objects as go
+
+    from diffract.session import Session
+
+
+@dataclass(kw_only=True)
+class GridPlot:
+    """Compose multiple plots into a grid layout."""
+
+    subplots: list[SubplotSpec]
+    make_subplots_kwargs: dict[str, Any]
+
+    def render(self, session: Session, theme: Theme | None = None) -> go.Figure:
+        """Render all subplot specs into a single Plotly grid figure."""
+        _validate_subplots(self.subplots)
+
+        total_rows = max((spec.row for spec in self.subplots), default=1)
+        total_cols = max((spec.col for spec in self.subplots), default=1)
+        sorted_specs = sorted(self.subplots, key=lambda spec: (spec.row, spec.col))
+        titles = _build_subplot_titles(sorted_specs, total_rows, total_cols)
+
+        kwargs = dict(self.make_subplots_kwargs or {})
+        for forbidden_key in ("rows", "cols", "subplot_titles"):
+            kwargs.pop(forbidden_key, None)
+
+        grid = make_subplots(
+            rows=total_rows,
+            cols=total_cols,
+            subplot_titles=titles,
+            **kwargs,
+        )
+
+        seen_legend_names: set[str] = set()
+        coloraxis_registry = ColoraxisRegistry()
+
+        for spec in sorted_specs:
+            render_session = session
+            resolved_session_filter = _resolve_session_filter(spec)
+            if resolved_session_filter:
+                render_session = _apply_session_filter(session, resolved_session_filter)
+
+            child_figure = _render_subplot_plot(
+                plot=spec.plot,
+                session=render_session,
+                theme=theme,
+                extra_value_filter=spec.value_filter,
+            )
+
+            remap_coloraxis(
+                child_figure,
+                coloraxis_registry,
+                explicit_share_keys=extract_coloraxis_share_keys(spec.plot),
+                legacy_share_name=extract_legacy_coloraxis_name(spec.plot),
+            )
+
+            add_figure_to_subplot(
+                grid,
+                child_figure,
+                row=spec.row,
+                col=spec.col,
+                transfer_layout=True,
+                seen_legend_names=seen_legend_names,
+            )
+
+        apply_coloraxis_configs(grid, coloraxis_registry)
+
+        if theme is not None:
+            apply_theme(grid, theme)
+
+        distribute_colorbars(grid)
+
+        return grid
+
+
+def _build_subplot_titles(specs: list[SubplotSpec], rows: int, cols: int) -> list[str]:
+    titles = [""] * (rows * cols)
+    for spec in specs:
+        idx = (spec.row - 1) * cols + (spec.col - 1)
+        existing = titles[idx]
+        if existing and spec.title and existing != spec.title:
+            raise ValueError(
+                "Conflicting subplot titles for the same cell "
+                f"(row={spec.row}, col={spec.col}): '{existing}' vs '{spec.title}'."
+            )
+        if not existing and spec.title:
+            titles[idx] = spec.title
+    return titles
+
+
+def _validate_subplots(subplots: list[SubplotSpec]) -> None:
+    for spec in subplots:
+        if spec.row < 1 or spec.col < 1:
+            raise ValueError(
+                "Subplot row/col must be 1-indexed positive integers; "
+                f"got row={spec.row}, col={spec.col}."
+            )
+
+        filter_dict = _resolve_session_filter(spec)
+        if filter_dict:
+            invalid_keys = sorted(set(filter_dict) - VALID_SESSION_FILTER_KEYS)
+            if invalid_keys:
+                valid_keys = sorted(VALID_SESSION_FILTER_KEYS)
+                raise ValueError(
+                    f"Invalid filter keys {invalid_keys} in subplot filter. "
+                    f"Valid keys are: {valid_keys}"
+                )
+
+
+def _resolve_session_filter(spec: SubplotSpec) -> dict[str, Any] | None:
+    if spec.session_filter is None:
+        return spec.filter
+    if spec.filter is None:
+        return spec.session_filter
+
+    merged = dict(spec.filter)
+    merged.update(spec.session_filter)
+    return merged
+
+
+def _render_subplot_plot(
+    *,
+    plot: Plot,
+    session: Session,
+    theme: Theme | None,
+    extra_value_filter: dict[str, tuple[str, Any]] | None,
+) -> go.Figure:
+    original_value_filter = getattr(plot, "value_filter", VALUE_FILTER_SENTINEL)
+
+    if original_value_filter is not VALUE_FILTER_SENTINEL:
+        merged = _merge_value_filters(original_value_filter, extra_value_filter)
+        plot.value_filter = merged
+
+    try:
+        return _render_plot_with_optional_theme(plot, session, theme)
+    finally:
+        if original_value_filter is not VALUE_FILTER_SENTINEL:
+            plot.value_filter = original_value_filter
+
+
+def _render_plot_with_optional_theme(
+    plot: Plot,
+    session: Session,
+    theme: Theme | None,
+) -> go.Figure:
+    if theme is None:
+        return plot.render(session)
+
+    if _plot_supports_theme_kwarg(plot):
+        return plot.render(session, theme=theme)
+    return plot.render(session)
+
+
+def _plot_supports_theme_kwarg(plot: Plot) -> bool:
+    try:
+        signature = inspect.signature(plot.render)
+    except (TypeError, ValueError):
+        return True
+
+    for parameter in signature.parameters.values():
+        if parameter.kind == inspect.Parameter.VAR_KEYWORD:
+            return True
+        if parameter.name == "theme":
+            return True
+
+    return False
+
+
+def _merge_value_filters(
+    base: dict[str, tuple[str, Any]] | None,
+    extra: dict[str, tuple[str, Any]] | None,
+) -> dict[str, tuple[str, Any]] | None:
+    if base is None and extra is None:
+        return None
+
+    merged: dict[str, tuple[str, Any]] = {}
+    if base:
+        merged.update(base)
+    if extra:
+        merged.update(extra)
+    return merged
+
+
+def _apply_session_filter(session: Session, filter_dict: dict[str, Any]) -> Session:
+    return session.filter(
+        param_ids=filter_dict.get("param_ids"),
+        param_names=filter_dict.get("param_names"),
+        param_types=filter_dict.get("param_types"),
+        model_ids=filter_dict.get("model_ids"),
+    )

diffract/viz/plots/subplots/layout.py

@@ -0,0 +1,370 @@
+"""Layout transfer helpers for subplot composition."""
+
+from __future__ import annotations
+
+import re
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    import plotly.graph_objects as go
+
+_AXIS_REF_RE = re.compile(r"^(?P<axis>[xy])(?P<id>\d+)?(?P<domain>\s+domain)?$")
+
+
+def add_figure_to_subplot(
+    fig_parent: go.Figure,
+    child_fig: go.Figure,
+    *,
+    row: int,
+    col: int,
+    transfer_layout: bool,
+    seen_legend_names: set[str] | None = None,
+) -> None:
+    """Add traces from a child figure into a specific subplot cell."""
+    if seen_legend_names is None:
+        seen_legend_names = set()
+
+    axis_refs = _get_subplot_axis_refs(fig_parent, row=row, col=col)
+    overlay_xaxis_map: dict[str, str] = {}
+    if axis_refs is not None:
+        x_ref, y_ref = axis_refs
+        overlay_xaxis_map = _register_overlay_xaxes(
+            fig_parent,
+            child_fig,
+            x_ref=x_ref,
+            y_ref=y_ref,
+        )
+
+    for trace in child_fig.data:
+        source_xaxis = getattr(trace, "xaxis", None)
+        source_yaxis = getattr(trace, "yaxis", None)
+
+        name = getattr(trace, "name", None)
+        if name and name in seen_legend_names:
+            trace.showlegend = False
+            if not getattr(trace, "legendgroup", None):
+                trace.legendgroup = name
+        elif name:
+            seen_legend_names.add(name)
+            if not getattr(trace, "legendgroup", None):
+                trace.legendgroup = name
+
+        fig_parent.add_trace(trace, row=row, col=col)
+
+        if axis_refs is None:
+            continue
+
+        mapped_xaxis = _remap_trace_axis_reference(
+            source_xaxis,
+            axis="x",
+            primary_axis=x_ref,
+            overlay_map=overlay_xaxis_map,
+        )
+        mapped_yaxis = _remap_trace_axis_reference(
+            source_yaxis,
+            axis="y",
+            primary_axis=y_ref,
+        )
+
+        fig_parent.data[-1].update(xaxis=mapped_xaxis, yaxis=mapped_yaxis)
+
+    if transfer_layout:
+        transfer_layout_to_subplot(fig_parent, child_fig, row=row, col=col)
+
+
+def transfer_layout_to_subplot(
+    fig_parent: go.Figure,
+    child_fig: go.Figure,
+    *,
+    row: int,
+    col: int,
+) -> None:
+    """Transfer axis/layout overlays from child to the target subplot cell."""
+    axis_refs = _get_subplot_axis_refs(fig_parent, row=row, col=col)
+    if axis_refs is None:
+        return
+    x_ref, y_ref = axis_refs
+
+    _transfer_axis_props(fig_parent, child_fig, row=row, col=col)
+    _transfer_shapes(fig_parent, child_fig, x_ref=x_ref, y_ref=y_ref)
+    _transfer_annotations(fig_parent, child_fig, x_ref=x_ref, y_ref=y_ref)
+    _transfer_images(fig_parent, child_fig, x_ref=x_ref, y_ref=y_ref)
+
+
+def _transfer_axis_props(
+    fig_parent: go.Figure,
+    child_fig: go.Figure,
+    *,
+    row: int,
+    col: int,
+) -> None:
+    axis_props = [
+        "tickmode",
+        "tickvals",
+        "ticktext",
+        "tickangle",
+        "tickformat",
+        "tickprefix",
+        "ticksuffix",
+        "showticklabels",
+        "ticks",
+        "tickfont",
+        "tickcolor",
+        "ticklen",
+        "tickwidth",
+        "showgrid",
+        "gridcolor",
+        "gridwidth",
+        "zeroline",
+        "zerolinecolor",
+        "zerolinewidth",
+        "type",
+        "range",
+        "autorange",
+        "categoryorder",
+        "categoryarray",
+        "title",
+        "titlefont",
+        "showline",
+        "linecolor",
+        "linewidth",
+        "mirror",
+        "dtick",
+        "tick0",
+        "rangemode",
+        "fixedrange",
+        "automargin",
+    ]
+
+    for axis_type, child_axis in (
+        ("x", child_fig.layout.xaxis),
+        ("y", child_fig.layout.yaxis),
+    ):
+        if child_axis is None:
+            continue
+
+        props = {
+            key: getattr(child_axis, key)
+            for key in axis_props
+            if getattr(child_axis, key, None) is not None
+        }
+        if not props:
+            continue
+
+        if axis_type == "x":
+            fig_parent.update_xaxes(row=row, col=col, **props)
+        else:
+            fig_parent.update_yaxes(row=row, col=col, **props)
+
+
+def _transfer_shapes(
+    fig_parent: go.Figure,
+    child_fig: go.Figure,
+    *,
+    x_ref: str,
+    y_ref: str,
+) -> None:
+    if not child_fig.layout.shapes:
+        return
+
+    new_shapes: list[dict[str, Any]] = []
+    for shape in child_fig.layout.shapes:
+        shape_dict = shape.to_plotly_json()
+        _remap_xy_refs(shape_dict, x_ref=x_ref, y_ref=y_ref)
+        new_shapes.append(shape_dict)
+
+    current_shapes = list(fig_parent.layout.shapes or [])
+    fig_parent.update_layout(shapes=current_shapes + new_shapes)
+
+
+def _transfer_annotations(
+    fig_parent: go.Figure,
+    child_fig: go.Figure,
+    *,
+    x_ref: str,
+    y_ref: str,
+) -> None:
+    if not child_fig.layout.annotations:
+        return
+
+    new_annotations: list[dict[str, Any]] = []
+    for annotation in child_fig.layout.annotations:
+        annotation_dict = annotation.to_plotly_json()
+        _remap_xy_refs(annotation_dict, x_ref=x_ref, y_ref=y_ref)
+        new_annotations.append(annotation_dict)
+
+    current_annotations = list(fig_parent.layout.annotations or [])
+    fig_parent.update_layout(annotations=current_annotations + new_annotations)
+
+
+def _transfer_images(
+    fig_parent: go.Figure,
+    child_fig: go.Figure,
+    *,
+    x_ref: str,
+    y_ref: str,
+) -> None:
+    if not child_fig.layout.images:
+        return
+
+    new_images: list[dict[str, Any]] = []
+    for image in child_fig.layout.images:
+        image_dict = image.to_plotly_json()
+        _remap_xy_refs(image_dict, x_ref=x_ref, y_ref=y_ref)
+        new_images.append(image_dict)
+
+    current_images = list(fig_parent.layout.images or [])
+    fig_parent.update_layout(images=current_images + new_images)
+
+
+def _remap_xy_refs(item: dict[str, Any], *, x_ref: str, y_ref: str) -> None:
+    x_value = item.get("xref")
+    y_value = item.get("yref")
+    ax_value = item.get("axref")
+    ay_value = item.get("ayref")
+
+    remapped_x = _remap_axis_reference(x_value, axis="x", target_axis=x_ref)
+    remapped_y = _remap_axis_reference(y_value, axis="y", target_axis=y_ref)
+    remapped_ax = _remap_axis_reference(ax_value, axis="x", target_axis=x_ref)
+    remapped_ay = _remap_axis_reference(ay_value, axis="y", target_axis=y_ref)
+
+    if remapped_x is not None:
+        item["xref"] = remapped_x
+    if remapped_y is not None:
+        item["yref"] = remapped_y
+    if remapped_ax is not None:
+        item["axref"] = remapped_ax
+    if remapped_ay is not None:
+        item["ayref"] = remapped_ay
+
+
+def _remap_axis_reference(
+    value: Any,
+    *,
+    axis: str,
+    target_axis: str,
+) -> str | None:
+    if not isinstance(value, str):
+        return None
+
+    match = _AXIS_REF_RE.fullmatch(value.strip())
+    if match is None or match.group("axis") != axis:
+        return value
+
+    domain = match.group("domain") or ""
+    return f"{target_axis}{domain}"
+
+
+def _get_subplot_axis_refs(
+    fig_parent: go.Figure,
+    *,
+    row: int,
+    col: int,
+) -> tuple[str, str] | None:
+    """Get parent subplot axis references (for example ``('x3', 'y3')``)."""
+    grid_refs = fig_parent._grid_ref[row - 1][col - 1]
+    if not grid_refs:
+        return None
+
+    grid_ref = grid_refs[0]
+    remap_ref = getattr(grid_ref, "trace_kwargs", None)
+    if not isinstance(remap_ref, dict):
+        return None
+
+    x_ref = remap_ref.get("xaxis")
+    y_ref = remap_ref.get("yaxis")
+    if not isinstance(x_ref, str) or not isinstance(y_ref, str):
+        return None
+
+    return x_ref, y_ref
+
+
+def _register_overlay_xaxes(
+    fig_parent: go.Figure,
+    child_fig: go.Figure,
+    *,
+    x_ref: str,
+    y_ref: str,
+) -> dict[str, str]:
+    """Create parent overlay x-axes for child overlay axes (x2, x3, ...)."""
+    layout_dict = child_fig.layout.to_plotly_json()
+    overlay_map: dict[str, str] = {}
+
+    for layout_key, axis_cfg in layout_dict.items():
+        if not layout_key.startswith("xaxis") or layout_key == "xaxis":
+            continue
+        if not isinstance(axis_cfg, dict):
+            continue
+
+        suffix = layout_key[5:]
+        child_axis = "x" if not suffix else f"x{suffix}"
+        overlaying = axis_cfg.get("overlaying")
+        if overlaying not in {"x", "x1"}:
+            continue
+
+        parent_axis = _next_axis_ref(fig_parent, axis="x")
+        parent_layout_key = _axis_layout_key(parent_axis)
+        parent_cfg = dict(axis_cfg)
+        parent_cfg["overlaying"] = x_ref
+        parent_cfg["anchor"] = y_ref
+        parent_cfg.pop("domain", None)
+
+        fig_parent.update_layout(**{parent_layout_key: parent_cfg})
+        overlay_map[child_axis] = parent_axis
+
+    return overlay_map
+
+
+def _remap_trace_axis_reference(
+    value: Any,
+    *,
+    axis: str,
+    primary_axis: str,
+    overlay_map: dict[str, str] | None = None,
+) -> str:
+    """Map a child trace axis reference to parent subplot axis reference."""
+    if not isinstance(value, str):
+        return primary_axis
+
+    match = _AXIS_REF_RE.fullmatch(value.strip())
+    if match is None or match.group("axis") != axis:
+        return primary_axis
+
+    axis_id = match.group("id") or ""
+    child_axis = f"{axis}{axis_id}"
+
+    if overlay_map is not None and child_axis in overlay_map:
+        return overlay_map[child_axis]
+
+    if child_axis in {axis, f"{axis}1"}:
+        return primary_axis
+
+    return child_axis
+
+
+def _next_axis_ref(fig_parent: go.Figure, *, axis: str) -> str:
+    """Get next available axis reference in parent layout."""
+    max_id = 0
+    layout_dict = fig_parent.layout.to_plotly_json()
+    axis_prefix = f"{axis}axis"
+
+    for layout_key in layout_dict:
+        if layout_key == axis_prefix:
+            max_id = max(max_id, 1)
+            continue
+        if not layout_key.startswith(axis_prefix):
+            continue
+
+        suffix = layout_key[len(axis_prefix) :]
+        if suffix.isdigit():
+            max_id = max(max_id, int(suffix))
+
+    next_id = max_id + 1
+    return axis if next_id == 1 else f"{axis}{next_id}"
+
+
+def _axis_layout_key(axis_ref: str) -> str:
+    """Convert axis ref like ``x2`` to layout key like ``xaxis2``."""
+    axis = axis_ref[0]
+    suffix = axis_ref[1:]
+    return f"{axis}axis{suffix}"

diffract/viz/plots/subplots/spec.py

@@ -0,0 +1,39 @@
+"""Subplot specifications and shared constants."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+from diffract.viz.plots.base.plot import Plot
+
+VALID_SESSION_FILTER_KEYS = frozenset(
+    {"param_ids", "param_names", "param_types", "model_ids"}
+)
+VALUE_FILTER_SENTINEL = object()
+
+
+@dataclass(slots=True)
+class SubplotSpec:
+    """Specification for a single subplot in a `GridPlot`.
+
+    Args:
+        row: Row position (1-indexed).
+        col: Column position (1-indexed).
+        title: Title for this subplot.
+        plot: Plot object to render.
+        session_filter: Optional dict forwarded to `Session.filter(...)`.
+            Valid keys: param_ids, param_names, param_types, model_ids.
+        value_filter: Optional value filter merged into `plot.value_filter` for
+            this subplot only: {field: (op, threshold)}.
+        filter: Backward-compatible alias for `session_filter`.
+    """
+
+    row: int
+    col: int
+    title: str
+    plot: Plot
+
+    session_filter: dict[str, Any] | None = None
+    value_filter: dict[str, tuple[str, Any]] | None = None
+    filter: dict[str, Any] | None = None