diffract-core 0.2.1__py3-none-any.whl

diffract/viz/data/ordering.py

@@ -0,0 +1,129 @@
+from __future__ import annotations
+
+from collections.abc import Callable, Sequence
+from dataclasses import dataclass
+from enum import Enum, auto
+from typing import Any
+
+import numpy as np
+from numpy.typing import NDArray
+
+
+class OrderMode(Enum):
+    """How to order categorical values."""
+
+    AS_IS = auto()
+    LEXICOGRAPHIC = auto()
+    NUMERIC = auto()
+    MAP = auto()
+    CUSTOM = auto()
+
+    @classmethod
+    def from_string(cls, name: str) -> OrderMode:
+        """Return the OrderMode matching the given name (case-insensitive)."""
+        match name.lower():
+            case "as_is":
+                return cls.AS_IS
+            case "lexicographic":
+                return cls.LEXICOGRAPHIC
+            case "numeric":
+                return cls.NUMERIC
+            case "map":
+                return cls.MAP
+            case "custom":
+                return cls.CUSTOM
+            case _:
+                raise ValueError(f"Unknown order mode: {name!r}")
+
+
+@dataclass(kw_only=True)
+class Ordering:
+    """Defines how to order values.
+
+    Returns indices (like np.argsort) so the same ordering
+    can be applied to multiple related arrays.
+    """
+
+    mode: OrderMode = OrderMode.AS_IS
+    descending: bool = False
+
+    # For MAP mode: function to extract comparison key
+    key: Callable[[Any], Any] | None = None
+
+    # For CUSTOM mode: explicit order of values
+    custom_order: Sequence[Any] | None = None
+
+    def argsort(self, values: Sequence[Any]) -> NDArray[np.intp]:
+        """Return indices that would sort unique values.
+
+        Args:
+            values: Sequence of values to order.
+
+        Returns:
+            Array of indices into unique values that produces the desired order.
+            Use np.unique(..., return_inverse=True) to map back to original data.
+        """
+        unique = list(dict.fromkeys(values))  # preserve first occurrence order
+
+        if self.mode == OrderMode.AS_IS:
+            indices = np.arange(len(unique))
+
+        elif self.mode == OrderMode.LEXICOGRAPHIC:
+            indices = np.argsort([str(v) for v in unique])
+
+        elif self.mode == OrderMode.NUMERIC:
+            indices = np.argsort(
+                [float(v) if v is not None else float("-inf") for v in unique]
+            )
+
+        elif self.mode == OrderMode.MAP:
+            if self.key is None:
+                indices = np.arange(len(unique))
+            else:
+                indices = np.argsort([self.key(v) for v in unique])
+
+        elif True:
+            if self.custom_order is None:
+                indices = np.arange(len(unique))
+            else:
+                order_map = {v: i for i, v in enumerate(self.custom_order)}
+                max_idx = len(order_map)
+                # Values not in custom_order go to the end
+                priorities = [order_map.get(v, max_idx) for v in unique]
+                indices = np.argsort(priorities)
+
+        else:
+            indices = np.arange(len(unique))
+
+        if self.descending:
+            indices = indices[::-1].copy()
+
+        return indices
+
+
+# Convenience constructors
+
+
+def as_is() -> Ordering:
+    """Keep original order from data."""
+    return Ordering(mode=OrderMode.AS_IS)
+
+
+def lexicographic(descending: bool = False) -> Ordering:
+    """Sort alphabetically by string representation."""
+    return Ordering(mode=OrderMode.LEXICOGRAPHIC, descending=descending)
+
+
+def numeric(descending: bool = False) -> Ordering:
+    """Sort by numeric value."""
+    return Ordering(mode=OrderMode.NUMERIC, descending=descending)
+
+
+def by_key(key: Callable[[Any], Any], descending: bool = False) -> Ordering:
+    """Sort by applying key function to each value."""
+    return Ordering(mode=OrderMode.MAP, key=key, descending=descending)
+
+
+def custom(order: Sequence[Any], descending: bool = False) -> Ordering:
+    """Sort by explicit order. Values not in order go to the end."""
+    return Ordering(mode=OrderMode.CUSTOM, custom_order=order, descending=descending)

diffract/viz/data/provider.py

@@ -0,0 +1,99 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from .extraction import get_field_data
+from .filtering import apply_filter
+from .types import DataShape, DataType, Entry
+
+if TYPE_CHECKING:
+    from diffract.session import Session
+
+
+class DataProvider:
+    """Data provider for visualization.
+
+    Fetches data from session and provides typed field access.
+    Always fetches aggregates merged with entries.
+    """
+
+    def __init__(self, session: Session) -> None:
+        self._session = session
+        self._entries: dict[str, Entry] = {}
+        self._cached_fields: set[str] = set()
+
+    def fetch(
+        self,
+        fields: list[str],
+        *,
+        value_filter: dict[str, tuple[str, Any]] | None = None,
+    ) -> dict[str, Entry]:
+        """Fetch data from session with aggregates merged.
+
+        Uses unified export API that handles aggregate merge internally.
+        Always expands contextual fields into entries.
+
+        Args:
+            fields: Computed field names to fetch.
+            value_filter: Optional filter {field: (op, threshold)}.
+
+        Returns:
+            Entries dict: {uid: Entry}
+        """
+        fields_to_fetch = set(fields)
+        if value_filter:
+            fields_to_fetch.update(value_filter.keys())
+
+        missing_fields = fields_to_fetch - self._cached_fields
+        if missing_fields:
+            update = self._session.results.export(
+                *sorted(missing_fields),
+                sources="all",
+                export_format="dict",
+                expand_contextual=True,
+            )
+            self._merge_entries(update)
+            self._cached_fields.update(missing_fields)
+
+        if value_filter:
+            return apply_filter(self._entries, value_filter)
+
+        return self._entries
+
+    def get_field_data(self, field: str) -> tuple[list[Any], DataType, DataShape]:
+        """Fetch a field and return its values, data type, and data shape.
+
+        Args:
+            field: Computed field name to fetch.
+
+        Returns:
+            Tuple of (values, data type, data shape).
+        """
+        return get_field_data(self.fetch(fields=[field]), field)
+
+    def _merge_entries(self, update: dict[str, Any]) -> None:
+        for uid, item in update.items():
+            fields = item.get("fields")
+            metadata = item.get("metadata")
+
+            if not isinstance(fields, dict):
+                raise TypeError(
+                    f"Expected 'fields' mapping for entry '{uid}', "
+                    f"got {type(fields).__name__}"
+                )
+            if not isinstance(metadata, dict):
+                raise TypeError(
+                    f"Expected 'metadata' mapping for entry '{uid}', "
+                    f"got {type(metadata).__name__}"
+                )
+
+            merged_fields = dict(fields)
+            merged_fields.update(metadata)
+
+            existing = self._entries.get(uid)
+            if existing is not None:
+                existing_fields = existing.get("fields", {})
+                if isinstance(existing_fields, dict):
+                    merged_fields = dict(existing_fields) | merged_fields
+
+            self._entries[uid] = Entry(fields=merged_fields)

diffract/viz/data/types.py

@@ -0,0 +1,98 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from enum import Enum, auto
+from typing import TYPE_CHECKING, Any, TypedDict
+
+from .ordering import as_is
+
+if TYPE_CHECKING:
+    from .ordering import Ordering
+
+
+class Entry(TypedDict, total=False):
+    """Single data entry from session results."""
+
+    fields: dict[str, Any]
+
+
+class DataType(Enum):
+    """Whether a field's values are numeric or categorical."""
+
+    NUMERIC = auto()
+    CATEGORICAL = auto()
+
+    @classmethod
+    def from_string(cls, name: str) -> DataType:
+        """Return the DataType matching the given name (case-insensitive)."""
+        match name.lower():
+            case "numeric":
+                return DataType.NUMERIC
+            case "categorical":
+                return DataType.CATEGORICAL
+            case _:
+                raise ValueError(f"Unknown data type: {name!r}")
+
+
+class DataShape(Enum):
+    """Whether a field holds a scalar or a vector per entry."""
+
+    SCALAR = auto()
+    VECTOR = auto()
+
+    @classmethod
+    def from_string(cls, name: str) -> DataShape:
+        """Return the DataShape matching the given name (case-insensitive)."""
+        match name.lower():
+            case "scalar":
+                return DataShape.SCALAR
+            case "vector":
+                return DataShape.VECTOR
+            case _:
+                raise ValueError(f"Unknown data shape: {name!r}")
+
+
+@dataclass
+class FieldRef:
+    """Reference to a data field with optional ordering.
+
+    Attributes:
+        field: Name of the field to reference.
+        ordering: How to order values of this field. Defaults to AS_IS (no reordering).
+    """
+
+    field: str
+    data_type: DataType | None = field(default=None)
+    data_shape: DataShape | None = field(default=None)
+    ordering: Ordering = field(default_factory=as_is)
+
+
+@dataclass
+class FieldMeta:
+    """Metadata about a field's type and shape."""
+
+    data_type: DataType
+    data_shape: DataShape
+
+
+@dataclass
+class EntryContext:
+    """Entry fields plus extracted model and parameter context."""
+
+    fields: dict[str, Any]
+    model_id: Any
+    parameter_name: Any
+
+    @classmethod
+    def from_entry(cls, entry: Entry) -> EntryContext:
+        """Create EntryContext from a session entry."""
+        entry_fields = entry.get("fields", {})
+        if not isinstance(entry_fields, dict):
+            raise TypeError(
+                f"Entry 'fields' must be a mapping, got {type(entry_fields).__name__}"
+            )
+        return cls(
+            fields=entry_fields,
+            model_id=entry_fields.get("model_id"),
+            parameter_name=entry_fields.get("name"),
+        )

diffract/viz/plots/__init__.py

@@ -0,0 +1,37 @@
+"""Built-in plots for `diffract.viz`.
+
+This module mirrors (approximately) the public surface of `diffract.viz.plots`
+for the refactored plotting stack.
+"""
+
+try:
+    import plotly
+except ImportError as e:
+    _msg = (
+        "diffract.viz requires the viz extra. "
+        'Install it with: pip install "diffract-core[viz]"'
+    )
+    raise ImportError(_msg) from e
+
+from .base import UpdateFigure
+from .boxplot import BoxPlot
+from .cluster import ClusterBarChart
+from .heatmap import HeatmapPlot
+from .scatter import ScatterPlot
+from .sparkline import SparklinePlot
+from .subplots import GridPlot, SubplotSpec
+from .violin import ViolinPlot
+
+Sparkline = SparklinePlot
+
+__all__ = [
+    "BoxPlot",
+    "ClusterBarChart",
+    "GridPlot",
+    "HeatmapPlot",
+    "ScatterPlot",
+    "Sparkline",
+    "SparklinePlot",
+    "SubplotSpec",
+    "ViolinPlot",
+]

diffract/viz/plots/base/__init__.py

@@ -0,0 +1,20 @@
+"""Base plot components and mixins."""
+
+from .axis import AxisType, SupportsAxis
+from .coloraxis import SupportsColoraxis
+from .configurator import Configurator
+from .jitter import SupportsJitter, density_scaled_jitter
+from .marker import SupportsMarker
+from .plot import Plot
+from .update import UpdateFigure
+
+__all__ = [
+    "AxisType",
+    "Configurator",
+    "Plot",
+    "SupportsAxis",
+    "SupportsColoraxis",
+    "SupportsJitter",
+    "SupportsMarker",
+    "density_scaled_jitter",
+]

diffract/viz/plots/base/axis.py

@@ -0,0 +1,219 @@
+from __future__ import annotations
+
+from dataclasses import field, make_dataclass
+from enum import Flag, auto
+from typing import Any, Literal
+
+import plotly.graph_objects as go
+
+from .configurator import Configurator
+
+
+class AxisType(Flag):
+    """Bit flags describing which axis modes an axis supports."""
+
+    NUMERIC = auto()
+    CATEGORICAL = auto()
+
+
+AxisMode = Literal["numeric", "categorical"]
+
+
+def SupportsAxis(axis_name: str, axis_type: AxisType) -> type[Configurator]:  # noqa: N802  factory returns a class
+    """Factory that creates axis configurator mixins."""
+    has_numeric = AxisType.NUMERIC in axis_type
+    has_categorical = AxisType.CATEGORICAL in axis_type
+    has_both = has_numeric and has_categorical
+
+    fields = _common_fields(axis_name)
+
+    if has_numeric:
+        fields += _numeric_fields(axis_name)
+    if has_categorical:
+        fields += _categorical_fields(axis_name)
+    if has_both:
+        fields += _mode_fields(axis_name)
+
+    def configure(self: Any, fig: go.Figure) -> None:
+        kwargs: dict[str, Any] = {}
+
+        _configure_common(self, axis_name, kwargs)
+
+        if has_both:
+            mode = getattr(self, f"{axis_name}_axis_mode")
+            if mode is None:
+                mode = getattr(self, f"_{axis_name}_resolved_mode", "categorical")
+            if mode == "numeric":
+                _configure_numeric(self, axis_name, kwargs)
+            else:
+                _configure_categorical(self, axis_name, kwargs)
+        elif has_numeric:
+            _configure_numeric(self, axis_name, kwargs)
+        elif has_categorical:
+            _configure_categorical(self, axis_name, kwargs)
+
+        fig.update_layout(**{f"{axis_name}axis": kwargs})
+
+    return make_dataclass(
+        f"SupportsAxis_{axis_name}",
+        fields=fields,
+        bases=(Configurator,),
+        namespace={
+            "axis_name": axis_name,
+            "axis_type": axis_type,
+            "configure": configure,
+        },
+        kw_only=True,
+    )
+
+
+# --- Field definitions ---
+
+
+def _common_fields(axis_name: str) -> list[tuple[str, type, Any]]:
+    return [
+        (f"{axis_name}_title", str | None, field(default=None)),
+        (f"{axis_name}_showticklabels", bool, field(default=True)),
+        (f"{axis_name}_tickangle", int | None, field(default=None)),
+        (f"{axis_name}_tickfont_size", int | None, field(default=None)),
+        (f"{axis_name}_tickfont_family", str | None, field(default=None)),
+        (f"{axis_name}_tickfont_color", str | None, field(default=None)),
+        (f"{axis_name}_showgrid", bool, field(default=True)),
+        (f"{axis_name}_gridcolor", str | None, field(default=None)),
+        (f"{axis_name}_showline", bool, field(default=True)),
+        (f"{axis_name}_linecolor", str | None, field(default=None)),
+    ]
+
+
+def _numeric_fields(axis_name: str) -> list[tuple[str, type, Any]]:
+    return [
+        (f"{axis_name}_range", tuple[float, float] | None, field(default=None)),
+        (f"{axis_name}_dtick", float | str | None, field(default=None)),
+        (f"{axis_name}_tick0", float | None, field(default=None)),
+        (f"{axis_name}_tickformat", str | None, field(default=None)),
+        (f"{axis_name}_tickmode", str | None, field(default=None)),
+        (f"{axis_name}_tickvals", list[float] | None, field(default=None)),
+        (f"{axis_name}_ticktext", list[str] | None, field(default=None)),
+        (f"{axis_name}_zeroline", bool, field(default=False)),
+        (f"{axis_name}_zerolinecolor", str | None, field(default=None)),
+    ]
+
+
+def _categorical_fields(axis_name: str) -> list[tuple[str, type, Any]]:
+    return [
+        (f"{axis_name}_categoryorder", str | None, field(default=None)),
+        (f"{axis_name}_categoryarray", list[str] | None, field(default=None)),
+    ]
+
+
+def _mode_fields(axis_name: str) -> list[tuple[str, type, Any]]:
+    # None means "infer from the data"; the plot stores the outcome in
+    # _{axis}_resolved_mode before configuration runs.
+    return [
+        (f"{axis_name}_axis_mode", AxisMode | None, field(default=None)),
+    ]
+
+
+# --- Configure helpers ---
+
+
+def _configure_common(self: Any, axis_name: str, kwargs: dict[str, Any]) -> None:
+    def _get(suffix: str) -> Any:
+        return getattr(self, f"{axis_name}_{suffix}")
+
+    # Get theme axes style for defaults
+    theme = getattr(self, "_theme", None)
+    axes_style = theme.axes if theme else None
+
+    title = _get("title")
+    if title is not None:
+        kwargs["title"] = title
+
+    kwargs["showticklabels"] = _get("showticklabels")
+
+    # showgrid: explicit value or theme default
+    show_grid = _get("showgrid")
+    if show_grid is not None:
+        kwargs["showgrid"] = show_grid
+    elif axes_style:
+        kwargs["showgrid"] = axes_style.show_grid
+
+    # showline: explicit value or theme default
+    show_line = _get("showline")
+    if show_line is not None:
+        kwargs["showline"] = show_line
+    elif axes_style:
+        kwargs["showline"] = axes_style.show_line
+
+    # mirror: from theme if available
+    if axes_style:
+        kwargs["mirror"] = axes_style.mirror
+
+    tickangle = _get("tickangle")
+    if tickangle is not None:
+        kwargs["tickangle"] = tickangle
+
+    # gridcolor: explicit value or theme default
+    gridcolor = _get("gridcolor")
+    if gridcolor is not None:
+        kwargs["gridcolor"] = gridcolor
+    elif axes_style:
+        kwargs["gridcolor"] = axes_style.grid_color
+
+    # linecolor: explicit value or theme default
+    linecolor = _get("linecolor")
+    if linecolor is not None:
+        kwargs["linecolor"] = linecolor
+    elif axes_style:
+        kwargs["linecolor"] = axes_style.line_color
+
+    tickfont: dict[str, Any] = {}
+    if (size := _get("tickfont_size")) is not None:
+        tickfont["size"] = size
+    if (family := _get("tickfont_family")) is not None:
+        tickfont["family"] = family
+    if (color := _get("tickfont_color")) is not None:
+        tickfont["color"] = color
+    if tickfont:
+        kwargs["tickfont"] = tickfont
+
+
+def _configure_numeric(self: Any, axis_name: str, kwargs: dict[str, Any]) -> None:
+    def _get(suffix: str) -> Any:
+        return getattr(self, f"{axis_name}_{suffix}")
+
+    if (rng := _get("range")) is not None:
+        kwargs["range"] = rng
+
+    if (dtick := _get("dtick")) is not None:
+        kwargs["dtick"] = dtick
+
+    if (tick0 := _get("tick0")) is not None:
+        kwargs["tick0"] = tick0
+
+    if (tickmode := _get("tickmode")) is not None:
+        kwargs["tickmode"] = tickmode
+
+    if (tickvals := _get("tickvals")) is not None:
+        kwargs["tickvals"] = tickvals
+
+    if (ticktext := _get("ticktext")) is not None:
+        kwargs["ticktext"] = ticktext
+
+    if (fmt := _get("tickformat")) is not None:
+        kwargs["tickformat"] = fmt
+
+    kwargs["zeroline"] = _get("zeroline")
+    if (zcolor := _get("zerolinecolor")) is not None:
+        kwargs["zerolinecolor"] = zcolor
+
+
+def _configure_categorical(self: Any, axis_name: str, kwargs: dict[str, Any]) -> None:
+    def _get(suffix: str) -> Any:
+        return getattr(self, f"{axis_name}_{suffix}")
+
+    if (order := _get("categoryorder")) is not None:
+        kwargs["categoryorder"] = order
+
+    if (arr := _get("categoryarray")) is not None:
+        kwargs["categoryarray"] = arr