diffract-core 0.2.1__py3-none-any.whl

diffract/session/namespaces/viz/__init__.py

@@ -0,0 +1,91 @@
+"""Visualization namespace for Session."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+from diffract.session.session import Session, SessionContext
+from diffract.viz.renderer import render, render_from_config
+
+from .box import box
+from .grid import bound_grid, grid
+from .heatmap import heatmap
+from .scatter import scatter
+from .sparkline import sparkline
+from .violin import violin
+
+if TYPE_CHECKING:  # pragma: no cover
+    from diffract.viz.styling import Theme
+
+
+class VizNamespace:
+    """Visualization API for Session.
+
+    Provides multiple ways to create visualizations:
+
+    1. Simple methods (recommended for quick exploration):
+        >>> session.viz.box(y="stable_rank", x="model_id")
+        >>> session.viz.scatter(x="frob_norm", y="stable_rank")
+
+    2. Plot objects (for Hydra configs and advanced customization):
+        >>> from diffract.viz.plots.boxplot import BoxPlot
+        >>> session.viz.draw(plot=BoxPlot(y="stable_rank", x="model_id"))
+
+    3. Config files (for reproducible workflows):
+        >>> session.viz.draw(config_path="plots/boxplot.yaml")
+
+    4. Grid helpers (for subplot dashboards and bound grids):
+        >>> session.viz.grid(subplots=[...])
+        >>> session.viz.bound_grid(plot_template=..., row=..., col=...)
+    """
+
+    def __init__(
+        self,
+        session_or_context: Session | SessionContext,
+    ) -> None:
+        self.__session_or_context = session_or_context
+
+    def draw(
+        self,
+        *,
+        plot: Any = None,
+        config_path: str | Path | None = None,
+        overrides: list[str] | None = None,
+        theme: Theme | None = None,
+        theme_path: str | Path | None = None,
+    ) -> Any:
+        """Render a Plotly figure using the viz module.
+
+        Provide either `plot` (a Plot object) or `config_path` (a Hydra YAML file).
+
+        Args:
+            plot: A Plot instance to render.
+            config_path: Path to a Hydra YAML config file.
+            overrides: Hydra overrides to apply (only with config_path).
+            theme: A Theme instance to apply.
+            theme_path: Path to a YAML file with theme config.
+        """
+        with self.__session_or_context:
+            if (plot is None) == (config_path is None):
+                raise ValueError("Provide exactly one of: plot=... or config_path=...")
+
+            if config_path is not None:
+                return render_from_config(
+                    session=self.__session_or_context,
+                    config_path=config_path,
+                    overrides=overrides,
+                    theme=theme,
+                    theme_path=theme_path,
+                )
+
+            return render(plot, session=self.__session_or_context, theme=theme)
+
+    box = box
+    grid = grid
+    bound_grid = bound_grid
+    violin = violin
+    scatter = scatter
+    heatmap = heatmap
+    sparkline = sparkline
+    line = sparkline

diffract/session/namespaces/viz/_utils.py

@@ -0,0 +1,14 @@
+"""Shared helpers for viz namespace wrappers."""
+
+from __future__ import annotations
+
+from typing import Any
+
+
+def _to_field_ref(value: str | Any) -> Any:
+    """Convert str to FieldRef for plot constructor; pass through otherwise."""
+    from diffract.viz.data import FieldRef
+
+    if isinstance(value, str):
+        return FieldRef(field=value)
+    return value

diffract/session/namespaces/viz/box.py

@@ -0,0 +1,207 @@
+"""Box plot wrapper for Session.viz."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Literal
+
+from diffract.session.namespaces.viz._utils import _to_field_ref
+from diffract.viz.data import FieldRef
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+    import plotly.graph_objects as go  # type: ignore[import-not-found]
+
+    from diffract.viz.styling import Theme
+
+
+def box(
+    self: Any,
+    *,
+    y: str | FieldRef,
+    x: str | FieldRef,
+    title: str | None = None,
+    value_filter: dict[str, tuple[str, Any]] | None = None,
+    box_width: float = 0.5,
+    boxpoints: Literal["all", "outliers", False] = "outliers",
+    jitter_enabled: bool = False,
+    jitter_width: float = 0.12,
+    jitter_offset: float = -0.35,
+    jitter_seed: int = 42,
+    jitter_density_scale: bool = True,
+    jitter_color: str | FieldRef | None = None,
+    jitter_colorscale: str = "Viridis",
+    x_title: str | None = None,
+    x_showticklabels: bool = True,
+    x_tickangle: int | None = None,
+    x_tickfont_size: int | None = None,
+    x_tickfont_family: str | None = None,
+    x_tickfont_color: str | None = None,
+    x_showgrid: bool = True,
+    x_gridcolor: str | None = None,
+    x_showline: bool = True,
+    x_linecolor: str | None = None,
+    y_title: str | None = None,
+    y_showticklabels: bool = True,
+    y_tickangle: int | None = None,
+    y_tickfont_size: int | None = None,
+    y_tickfont_family: str | None = None,
+    y_tickfont_color: str | None = None,
+    y_showgrid: bool = True,
+    y_gridcolor: str | None = None,
+    y_showline: bool = True,
+    y_linecolor: str | None = None,
+    y_range: tuple[float, float] | None = None,
+    y_dtick: float | str | None = None,
+    y_tick0: float | None = None,
+    y_tickformat: str | None = None,
+    y_zeroline: bool = False,
+    y_zerolinecolor: str | None = None,
+    marker_coloraxis_id: int | None = None,
+    marker_colorscale: str = "Viridis",
+    marker_showscale: bool = True,
+    marker_cmin: float | None = None,
+    marker_cmax: float | None = None,
+    marker_colorbar_title: str | None = None,
+    marker_coloraxis_override: bool = False,
+    marker_size: str | FieldRef | float | None = 6,
+    marker_size_range: tuple[float, float] | None = None,
+    marker_opacity: str | FieldRef | float | None = 0.7,
+    marker_opacity_range: tuple[float, float] | None = None,
+    marker_color: str | FieldRef | None = None,
+    marker_symbol: str | FieldRef | None = None,
+    theme: Theme | None = None,
+    theme_path: str | Path | None = None,
+) -> go.Figure:
+    """Create a box plot.
+
+    Args:
+        self: The bound viz namespace providing the draw() method.
+        y: Field for vertical values (field name or FieldRef).
+        x: Field for categories (field name or FieldRef).
+        title: Figure title.
+        value_filter: Filter entries by field values.
+        box_width: Width of boxes.
+        boxpoints: Show points: "all", "outliers", or False.
+        jitter_enabled: Enable the jitter overlay of individual points.
+        jitter_width: Horizontal spread of the jitter overlay.
+        jitter_offset: Horizontal offset of the jitter overlay from the box.
+        jitter_seed: Random seed for the jitter overlay.
+        jitter_density_scale: Scale jitter spread by local point density.
+        jitter_color: Field or color for jitter points.
+        jitter_colorscale: Colorscale for jitter points.
+        x_title: X-axis title.
+        x_showticklabels: Show x-axis tick labels.
+        x_tickangle: Angle of x-axis tick labels.
+        x_tickfont_size: Font size of x-axis tick labels.
+        x_tickfont_family: Font family of x-axis tick labels.
+        x_tickfont_color: Font color of x-axis tick labels.
+        x_showgrid: Show x-axis grid lines.
+        x_gridcolor: Color of x-axis grid lines.
+        x_showline: Show the x-axis line.
+        x_linecolor: Color of the x-axis line.
+        y_title: Y-axis title.
+        y_showticklabels: Show y-axis tick labels.
+        y_tickangle: Angle of y-axis tick labels.
+        y_tickfont_size: Font size of y-axis tick labels.
+        y_tickfont_family: Font family of y-axis tick labels.
+        y_tickfont_color: Font color of y-axis tick labels.
+        y_showgrid: Show y-axis grid lines.
+        y_gridcolor: Color of y-axis grid lines.
+        y_showline: Show the y-axis line.
+        y_linecolor: Color of the y-axis line.
+        y_range: Y-axis value range.
+        y_dtick: Y-axis tick step.
+        y_tick0: Y-axis starting tick.
+        y_tickformat: Y-axis tick format string.
+        y_zeroline: Show the y-axis zero line.
+        y_zerolinecolor: Color of the y-axis zero line.
+        marker_coloraxis_id: Shared color axis id for markers.
+        marker_colorscale: Colorscale for markers.
+        marker_showscale: Show the marker colorbar.
+        marker_cmin: Minimum value for marker color mapping.
+        marker_cmax: Maximum value for marker color mapping.
+        marker_colorbar_title: Title of the marker colorbar.
+        marker_coloraxis_override: Override the shared marker color axis.
+        marker_size: Field or fixed size for markers.
+        marker_size_range: Range to scale marker sizes into.
+        marker_opacity: Field or fixed opacity for markers.
+        marker_opacity_range: Range to scale marker opacity into.
+        marker_color: Field or color for markers.
+        marker_symbol: Field or symbol for markers.
+        theme: Theme to apply to the figure.
+        theme_path: Path to a theme file to apply to the figure.
+
+    Returns:
+        Plotly Figure.
+
+    Example:
+        >>> session.viz.box(y="stable_rank", x="model_id", marker_color="layer_id")
+    """
+    from diffract.viz.plots.boxplot import BoxPlot
+
+    plot = BoxPlot(
+        y=_to_field_ref(y),
+        x=_to_field_ref(x),
+        title=title,
+        value_filter=value_filter,
+        box_width=box_width,
+        boxpoints=boxpoints,
+        jitter_enabled=jitter_enabled,
+        jitter_width=jitter_width,
+        jitter_offset=jitter_offset,
+        jitter_seed=jitter_seed,
+        jitter_density_scale=jitter_density_scale,
+        jitter_color=_to_field_ref(jitter_color)
+        if isinstance(jitter_color, str)
+        else jitter_color,
+        jitter_colorscale=jitter_colorscale,
+        x_title=x_title,
+        x_showticklabels=x_showticklabels,
+        x_tickangle=x_tickangle,
+        x_tickfont_size=x_tickfont_size,
+        x_tickfont_family=x_tickfont_family,
+        x_tickfont_color=x_tickfont_color,
+        x_showgrid=x_showgrid,
+        x_gridcolor=x_gridcolor,
+        x_showline=x_showline,
+        x_linecolor=x_linecolor,
+        y_title=y_title,
+        y_showticklabels=y_showticklabels,
+        y_tickangle=y_tickangle,
+        y_tickfont_size=y_tickfont_size,
+        y_tickfont_family=y_tickfont_family,
+        y_tickfont_color=y_tickfont_color,
+        y_showgrid=y_showgrid,
+        y_gridcolor=y_gridcolor,
+        y_showline=y_showline,
+        y_linecolor=y_linecolor,
+        y_range=y_range,
+        y_dtick=y_dtick,
+        y_tick0=y_tick0,
+        y_tickformat=y_tickformat,
+        y_zeroline=y_zeroline,
+        y_zerolinecolor=y_zerolinecolor,
+        marker_coloraxis_id=marker_coloraxis_id,
+        marker_colorscale=marker_colorscale,
+        marker_showscale=marker_showscale,
+        marker_cmin=marker_cmin,
+        marker_cmax=marker_cmax,
+        marker_colorbar_title=marker_colorbar_title,
+        marker_coloraxis_override=marker_coloraxis_override,
+        marker_size=_to_field_ref(marker_size)
+        if isinstance(marker_size, str)
+        else marker_size,
+        marker_size_range=marker_size_range,
+        marker_opacity=_to_field_ref(marker_opacity)
+        if isinstance(marker_opacity, str)
+        else marker_opacity,
+        marker_opacity_range=marker_opacity_range,
+        marker_color=_to_field_ref(marker_color)
+        if isinstance(marker_color, str)
+        else marker_color,
+        marker_symbol=_to_field_ref(marker_symbol)
+        if isinstance(marker_symbol, str)
+        else marker_symbol,
+    )
+    return self.draw(plot=plot, theme=theme, theme_path=theme_path)

diffract/session/namespaces/viz/grid.py

@@ -0,0 +1,160 @@
+"""Grid plot wrappers for Session.viz."""
+
+from __future__ import annotations
+
+from dataclasses import replace
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+from diffract.session.namespaces.viz._utils import _to_field_ref
+
+if TYPE_CHECKING:
+    import plotly.graph_objects as go  # type: ignore[import-not-found]
+
+    from diffract.viz.plots.base.plot import Plot
+    from diffract.viz.plots.subplots import GridAxisBind, GridCellRule, SubplotSpec
+    from diffract.viz.styling import Theme
+
+
+def grid(
+    self: Any,
+    *,
+    subplots: list[SubplotSpec],
+    make_subplots_kwargs: dict[str, Any] | None = None,
+    theme: Theme | None = None,
+    theme_path: str | Path | None = None,
+) -> go.Figure:
+    """Create and render a `GridPlot` from explicit subplot specifications."""
+    from diffract.viz.plots.subplots import GridPlot
+
+    plot = GridPlot(
+        subplots=subplots,
+        make_subplots_kwargs=dict(make_subplots_kwargs or {}),
+    )
+    return self.draw(plot=plot, theme=theme, theme_path=theme_path)
+
+
+def bound_grid(
+    self: Any,
+    *,
+    plot_template: Plot,
+    row: GridAxisBind | None = None,
+    col: GridAxisBind | None = None,
+    cell_rules: list[GridCellRule] | None = None,
+    title_template: str | None = None,
+    make_subplots_kwargs: dict[str, Any] | None = None,
+    base_session_filter: dict[str, Any] | None = None,
+    base_value_filter: dict[str, tuple[str, Any]] | None = None,
+    theme: Theme | None = None,
+    theme_path: str | Path | None = None,
+) -> go.Figure:
+    """Create and render a bound grid generated from row/column axis binds.
+
+    For `plot` binds targeting `FieldRef` plot attributes, plain string bind
+    values are automatically converted into `FieldRef(field=...)`. The same
+    conversion is applied to `cell_rules[*].plot` overrides.
+    """
+    from diffract.viz.plots.subplots import build_bound_grid
+
+    resolved_row = _coerce_plot_bind_field_refs(plot_template=plot_template, bind=row)
+    resolved_col = _coerce_plot_bind_field_refs(plot_template=plot_template, bind=col)
+    resolved_rules = _coerce_rule_plot_field_refs(
+        plot_template=plot_template,
+        rules=cell_rules,
+    )
+
+    plot = build_bound_grid(
+        plot_template=plot_template,
+        row=resolved_row,
+        col=resolved_col,
+        cell_rules=resolved_rules,
+        title_template=title_template,
+        make_subplots_kwargs=make_subplots_kwargs,
+        base_session_filter=base_session_filter,
+        base_value_filter=base_value_filter,
+    )
+    return self.draw(plot=plot, theme=theme, theme_path=theme_path)
+
+
+def _coerce_plot_bind_field_refs(
+    *,
+    plot_template: Plot,
+    bind: GridAxisBind | None,
+) -> GridAxisBind | None:
+    """Convert string plot-bind values to FieldRef when template field is FieldRef."""
+    if bind is None or bind.target != "plot":
+        return bind
+
+    if not hasattr(plot_template, bind.key):
+        return bind
+
+    from diffract.viz.data import FieldRef
+
+    template_value = getattr(plot_template, bind.key)
+    if not isinstance(template_value, FieldRef):
+        return bind
+
+    changed = False
+    converted_values: list[Any] = []
+    for value in bind.values:
+        if isinstance(value, str):
+            converted_values.append(_to_field_ref(value))
+            changed = True
+            continue
+        converted_values.append(value)
+
+    if not changed:
+        return bind
+
+    return replace(bind, values=tuple(converted_values))
+
+
+def _coerce_rule_plot_field_refs(
+    *,
+    plot_template: Plot,
+    rules: list[GridCellRule] | None,
+) -> list[GridCellRule] | None:
+    """Convert string `cell_rules[*].plot` values to FieldRef when needed."""
+    if not rules:
+        return rules
+
+    changed = False
+    converted_rules: list[GridCellRule] = []
+    for rule in rules:
+        converted_rule, rule_changed = _coerce_single_rule_plot_field_refs(
+            plot_template=plot_template,
+            rule=rule,
+        )
+        converted_rules.append(converted_rule)
+        changed = changed or rule_changed
+
+    if not changed:
+        return rules
+    return converted_rules
+
+
+def _coerce_single_rule_plot_field_refs(
+    *,
+    plot_template: Plot,
+    rule: GridCellRule,
+) -> tuple[GridCellRule, bool]:
+    """Convert one rule's `plot` overrides where template fields are FieldRef."""
+    if not rule.plot:
+        return rule, False
+
+    from diffract.viz.data import FieldRef
+
+    changed = False
+    converted_plot = dict(rule.plot)
+    for key, value in rule.plot.items():
+        if not hasattr(plot_template, key):
+            continue
+
+        template_value = getattr(plot_template, key)
+        if isinstance(template_value, FieldRef) and isinstance(value, str):
+            converted_plot[key] = _to_field_ref(value)
+            changed = True
+
+    if not changed:
+        return rule, False
+    return replace(rule, plot=converted_plot), True

diffract/session/namespaces/viz/heatmap.py

@@ -0,0 +1,164 @@
+"""Heatmap plot wrapper for Session.viz."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from diffract.session.namespaces.viz._utils import _to_field_ref
+from diffract.viz.data import FieldRef
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+    import plotly.graph_objects as go  # type: ignore[import-not-found]
+
+    from diffract.viz.styling import Theme
+
+
+def heatmap(
+    self: Any,
+    *,
+    z: str | FieldRef,
+    x: str | FieldRef,
+    y: str | FieldRef,
+    title: str | None = None,
+    value_filter: dict[str, tuple[str, Any]] | None = None,
+    fill_value: float = float("nan"),
+    show_text: bool = False,
+    text_format: str = ".2f",
+    text_font_size: int = 10,
+    reverse_y: bool = True,
+    x_title: str | None = None,
+    x_showticklabels: bool = True,
+    x_tickangle: int | None = None,
+    x_tickfont_size: int | None = None,
+    x_tickfont_family: str | None = None,
+    x_tickfont_color: str | None = None,
+    x_showgrid: bool = True,
+    x_gridcolor: str | None = None,
+    x_showline: bool = True,
+    x_linecolor: str | None = None,
+    x_categoryorder: str | None = None,
+    x_categoryarray: list[str] | None = None,
+    y_title: str | None = None,
+    y_showticklabels: bool = True,
+    y_tickangle: int | None = None,
+    y_tickfont_size: int | None = None,
+    y_tickfont_family: str | None = None,
+    y_tickfont_color: str | None = None,
+    y_showgrid: bool = True,
+    y_gridcolor: str | None = None,
+    y_showline: bool = True,
+    y_linecolor: str | None = None,
+    y_categoryorder: str | None = None,
+    y_categoryarray: list[str] | None = None,
+    heatmap_coloraxis_id: int | None = None,
+    heatmap_colorscale: str = "Viridis",
+    heatmap_showscale: bool = True,
+    heatmap_cmin: float | None = None,
+    heatmap_cmax: float | None = None,
+    heatmap_colorbar_title: str | None = None,
+    heatmap_coloraxis_override: bool = False,
+    theme: Theme | None = None,
+    theme_path: str | Path | None = None,
+) -> go.Figure:
+    """Create a heatmap by pivoting a scalar field.
+
+    Args:
+        self: The bound viz namespace providing the draw() method.
+        z: Field for cell values (field name or FieldRef).
+        x: Field for columns (field name or FieldRef).
+        y: Field for rows (field name or FieldRef).
+        title: Figure title.
+        value_filter: Filter entries by field values.
+        fill_value: Value for missing (x, y) cells.
+        show_text: Show cell values as text.
+        text_format: Format string for text (e.g. ".2f").
+        text_font_size: Font size for cell text.
+        reverse_y: Reverse y-axis so first row is at top.
+        x_title: X-axis title.
+        x_showticklabels: Show x-axis tick labels.
+        x_tickangle: Angle of x-axis tick labels.
+        x_tickfont_size: Font size of x-axis tick labels.
+        x_tickfont_family: Font family of x-axis tick labels.
+        x_tickfont_color: Font color of x-axis tick labels.
+        x_showgrid: Show x-axis grid lines.
+        x_gridcolor: Color of x-axis grid lines.
+        x_showline: Show the x-axis line.
+        x_linecolor: Color of the x-axis line.
+        x_categoryorder: Ordering strategy for x-axis categories.
+        x_categoryarray: Explicit ordering of x-axis categories.
+        y_title: Y-axis title.
+        y_showticklabels: Show y-axis tick labels.
+        y_tickangle: Angle of y-axis tick labels.
+        y_tickfont_size: Font size of y-axis tick labels.
+        y_tickfont_family: Font family of y-axis tick labels.
+        y_tickfont_color: Font color of y-axis tick labels.
+        y_showgrid: Show y-axis grid lines.
+        y_gridcolor: Color of y-axis grid lines.
+        y_showline: Show the y-axis line.
+        y_linecolor: Color of the y-axis line.
+        y_categoryorder: Ordering strategy for y-axis categories.
+        y_categoryarray: Explicit ordering of y-axis categories.
+        heatmap_coloraxis_id: Shared color axis id for the heatmap.
+        heatmap_colorscale: Colorscale for the heatmap.
+        heatmap_showscale: Show the heatmap colorbar.
+        heatmap_cmin: Minimum value for heatmap color mapping.
+        heatmap_cmax: Maximum value for heatmap color mapping.
+        heatmap_colorbar_title: Title of the heatmap colorbar.
+        heatmap_coloraxis_override: Override the shared heatmap color axis.
+        theme: Theme to apply to the figure.
+        theme_path: Path to a theme file to apply to the figure.
+
+    Returns:
+        Plotly Figure.
+
+    Example:
+        >>> session.viz.heatmap(z="stable_rank", x="head_id", y="layer_id")
+    """
+    from diffract.viz.plots.heatmap import HeatmapPlot
+
+    plot = HeatmapPlot(
+        z=_to_field_ref(z),
+        x=_to_field_ref(x),
+        y=_to_field_ref(y),
+        title=title,
+        value_filter=value_filter,
+        fill_value=fill_value,
+        show_text=show_text,
+        text_format=text_format,
+        text_font_size=text_font_size,
+        reverse_y=reverse_y,
+        x_title=x_title,
+        x_showticklabels=x_showticklabels,
+        x_tickangle=x_tickangle,
+        x_tickfont_size=x_tickfont_size,
+        x_tickfont_family=x_tickfont_family,
+        x_tickfont_color=x_tickfont_color,
+        x_showgrid=x_showgrid,
+        x_gridcolor=x_gridcolor,
+        x_showline=x_showline,
+        x_linecolor=x_linecolor,
+        x_categoryorder=x_categoryorder,
+        x_categoryarray=x_categoryarray,
+        y_title=y_title,
+        y_showticklabels=y_showticklabels,
+        y_tickangle=y_tickangle,
+        y_tickfont_size=y_tickfont_size,
+        y_tickfont_family=y_tickfont_family,
+        y_tickfont_color=y_tickfont_color,
+        y_showgrid=y_showgrid,
+        y_gridcolor=y_gridcolor,
+        y_showline=y_showline,
+        y_linecolor=y_linecolor,
+        y_categoryorder=y_categoryorder,
+        y_categoryarray=y_categoryarray,
+        heatmap_coloraxis_id=heatmap_coloraxis_id,
+        heatmap_colorscale=heatmap_colorscale,
+        heatmap_showscale=heatmap_showscale,
+        heatmap_cmin=heatmap_cmin,
+        heatmap_cmax=heatmap_cmax,
+        heatmap_colorbar_title=heatmap_colorbar_title,
+        heatmap_coloraxis_override=heatmap_coloraxis_override,
+    )
+    return self.draw(plot=plot, theme=theme, theme_path=theme_path)