diffract-core 0.2.1__py3-none-any.whl

diffract/viz/plots/heatmap.py

@@ -0,0 +1,168 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, replace
+from typing import Any
+
+import numpy as np
+import plotly.graph_objects as go
+
+from diffract.viz.data import (
+    DataType,
+    Entry,
+    FieldRef,
+)
+from diffract.viz.plots.base.axis import AxisType, SupportsAxis
+from diffract.viz.plots.base.coloraxis import SupportsColoraxis
+from diffract.viz.plots.base.plot import Plot
+from diffract.viz.styling import (
+    CategoricalPropertyResolver,
+    NumericPropertyResolver,
+)
+
+
+@dataclass(kw_only=True)
+class HeatmapPlot(
+    Plot,
+    SupportsAxis("x", AxisType.CATEGORICAL),
+    SupportsAxis("y", AxisType.CATEGORICAL),
+    SupportsColoraxis("heatmap"),
+):
+    """Heatmap plot that pivots a scalar field by two categorical dimensions.
+
+    Each cell in the heatmap represents a value of ``z`` at the intersection
+    of a ``y`` row and an ``x`` column.  When multiple entries map to the
+    same (x, y) pair the **last** value wins (no aggregation).
+
+    Coloraxis settings (colorscale, cmin/cmax, colorbar) are inherited from
+    the ``SupportsColoraxis("heatmap")`` mixin and can be controlled via
+    ``heatmap_colorscale``, ``heatmap_cmin``, ``heatmap_cmax``, etc.
+    """
+
+    z: FieldRef
+    z_rescale_range: tuple[float, float] | None = None
+
+    y: FieldRef
+    x: FieldRef
+
+    fill_value: float = float("nan")
+
+    show_text: bool = False
+    text_format: str = ".2f"
+    text_font_size: int = 10
+
+    reverse_y: bool = True
+
+    def _build_traces_data(
+        self, entries: dict[str, Entry]
+    ) -> dict[str, dict[str, Any]]:
+        """Build a single trace-data dict containing the full z-matrix."""
+        x_resolver = y_resolver = CategoricalPropertyResolver()
+        z_resolver = NumericPropertyResolver(self.z_rescale_range)
+
+        x_ref = replace(self.x, data_type=self.x.data_type or DataType.CATEGORICAL)
+        y_ref = replace(self.y, data_type=self.y.data_type or DataType.CATEGORICAL)
+        x_values = x_resolver.resolve(x_ref, entries)
+        y_values = y_resolver.resolve(y_ref, entries)
+        z_values = z_resolver.resolve(self.z, entries)
+
+        sorted_x = self._get_sorted_unique(x_values, self.x)
+        sorted_y = self._get_sorted_unique(y_values, self.y)
+
+        cell_values = self._collect_cell_values(x_values, y_values, z_values)
+        z_matrix = self._build_matrix(sorted_x, sorted_y, cell_values)
+
+        traces_data: dict[str, dict[str, Any]] = {
+            "heatmap": {
+                "x_labels": sorted_x,
+                "y_labels": sorted_y,
+                "z_matrix": z_matrix,
+            },
+        }
+        return traces_data
+
+    @staticmethod
+    def _get_sorted_unique(values: list[Any], ref: FieldRef) -> list[str]:
+        """Return unique values sorted according to the field's ordering."""
+        unique = list(dict.fromkeys(values))
+        order_indices = ref.ordering.argsort(unique)
+        return [str(unique[i]) for i in order_indices]
+
+    @staticmethod
+    def _collect_cell_values(
+        x_values: list[Any],
+        y_values: list[Any],
+        z_values: list[Any],
+    ) -> dict[tuple[str, str], float]:
+        """Map (x, y) string pairs to scalar z-values."""
+        cell_values: dict[tuple[str, str], float] = {}
+        for x_val, y_val, z_val in zip(x_values, y_values, z_values, strict=False):
+            if z_val is None:
+                continue
+            cell_values[(str(x_val), str(y_val))] = float(z_val)
+        return cell_values
+
+    def _build_matrix(
+        self,
+        sorted_x: list[str],
+        sorted_y: list[str],
+        cell_values: dict[tuple[str, str], float],
+    ) -> np.ndarray:
+        """Build a 2-D z-matrix (rows=y, cols=x) from cell values."""
+        z_matrix = np.full(
+            (len(sorted_y), len(sorted_x)),
+            self.fill_value,
+            dtype=np.float64,
+        )
+        x_idx = {label: j for j, label in enumerate(sorted_x)}
+        y_idx = {label: i for i, label in enumerate(sorted_y)}
+
+        for (x_label, y_label), value in cell_values.items():
+            j = x_idx.get(x_label)
+            i = y_idx.get(y_label)
+            if i is not None and j is not None:
+                z_matrix[i, j] = value
+
+        return z_matrix
+
+    def _build_figure(self) -> go.Figure:
+        """Build the Plotly figure with a single Heatmap trace."""
+        fig = go.Figure()
+
+        if self._traces_data is None:
+            return fig
+
+        trace_data = self._traces_data["heatmap"]
+        x_labels: list[str] = trace_data["x_labels"]
+        y_labels: list[str] = trace_data["y_labels"]
+        z_matrix: np.ndarray = trace_data["z_matrix"]
+
+        heatmap_kwargs: dict[str, Any] = {
+            "z": z_matrix,
+            "x": x_labels,
+            "y": y_labels,
+            "coloraxis": self.resolve_coloraxis(fig),
+        }
+
+        if self.show_text:
+            text = np.vectorize(
+                lambda val: "" if np.isnan(val) else format(val, self.text_format),
+            )(z_matrix)
+            heatmap_kwargs["text"] = text
+            heatmap_kwargs["texttemplate"] = "%{text}"
+            heatmap_kwargs["textfont"] = {"size": self.text_font_size}
+
+        fig.add_trace(go.Heatmap(**heatmap_kwargs))
+
+        # Title
+        if self.title:
+            fig.update_layout(title=self.title)
+        else:
+            fig.update_layout(
+                title=f"{self.z.field} by {self.y.field} x {self.x.field}"
+            )
+
+        # Reverse y-axis so the first row is at the top (matrix convention)
+        if self.reverse_y:
+            fig.update_yaxes(autorange="reversed")
+
+        return fig

diffract/viz/plots/scatter.py

@@ -0,0 +1,233 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+import plotly.graph_objects as go
+
+from diffract.viz.data import (
+    Entry,
+    FieldRef,
+    get_field_values,
+)
+from diffract.viz.plots.base.axis import AxisType, SupportsAxis
+from diffract.viz.plots.base.marker import SupportsMarker
+from diffract.viz.plots.base.plot import Plot
+from diffract.viz.styling import (
+    NumericPropertyResolver,
+    ResolvedColor,
+)
+
+
+@dataclass(kw_only=True)
+class ScatterPlot(
+    Plot,
+    SupportsAxis("x", AxisType.NUMERIC),
+    SupportsAxis("y", AxisType.NUMERIC),
+    SupportsMarker,
+):
+    """Scatter plot for two numeric fields."""
+
+    y: FieldRef
+    y_rescale_range: tuple[float, float] | None = None
+    y_rescale_traces_separately: bool = False
+
+    x: FieldRef
+    x_rescale_range: tuple[float, float] | None = None
+    x_rescale_traces_separately: bool = False
+
+    group_by: FieldRef | None = None
+
+    def _build_traces_data(
+        self, entries: dict[str, Entry]
+    ) -> dict[str, dict[str, Any]]:
+        """Build per-group trace data dicts."""
+        x_resolver = NumericPropertyResolver(self.x_rescale_range)
+        y_resolver = NumericPropertyResolver(self.y_rescale_range)
+
+        x_values = x_resolver.resolve(self.x, entries)
+        y_values = y_resolver.resolve(self.y, entries)
+
+        groups, group_indices = self._group_entries(entries, x_values, y_values)
+
+        # 1) Resolve marker properties globally for all entries
+        resolved_color = self._resolve_marker_color(entries, self._theme)
+        resolved_symbol = self._resolve_marker_symbol(entries, self._theme)
+        resolved_size = self._resolve_marker_size(entries, self._theme)
+        resolved_opacity = self._resolve_marker_opacity(entries, self._theme)
+
+        group_keys = self._get_sorted_group_keys(groups)
+
+        # 2) Build per-group traces
+        traces_data: dict[str, dict[str, Any]] = {}
+
+        for idx, group in enumerate(group_keys):
+            points = groups[group]
+            indices = group_indices[group]
+
+            trace_x = [p[0] for p in points]
+            if self.x_rescale_traces_separately:
+                trace_x = x_resolver._normalize(trace_x)
+
+            trace_y = [p[1] for p in points]
+            if self.y_rescale_traces_separately:
+                trace_y = y_resolver._normalize(trace_y)
+
+            trace_id = f"scatter_{group}"
+            trace_data: dict[str, Any] = {
+                "group": group,
+                "group_idx": idx,
+                "x_values": trace_x,
+                "y_values": trace_y,
+            }
+
+            # 3) Pick per-group values and add to trace via mixin methods
+            self._add_marker_color_to_trace(
+                trace_data,
+                self._pick_group_color(indices, resolved_color),
+            )
+            self._add_marker_symbol_to_trace(
+                trace_data,
+                self._pick_group_scalar(indices, resolved_symbol),
+            )
+            self._add_marker_size_to_trace(
+                trace_data,
+                self._pick_group_values(indices, resolved_size),
+            )
+            self._add_marker_opacity_to_trace(
+                trace_data,
+                self._pick_group_values(indices, resolved_opacity),
+            )
+
+            traces_data[trace_id] = trace_data
+
+        return traces_data
+
+    def _group_entries(
+        self,
+        entries: dict[str, Entry],
+        x_values: list[Any],
+        y_values: list[Any],
+    ) -> tuple[dict[str, list[tuple[float, float]]], dict[str, list[int]]]:
+        """Group (x, y) pairs by the ``group_by`` field.
+
+        Returns:
+            groups: mapping from group key to list of (x, y) pairs.
+            group_indices: mapping from group key to original entry indices
+                (only entries with non-None x *and* y).
+        """
+        if self.group_by is not None:
+            group_values = get_field_values(entries, self.group_by.field)
+        else:
+            group_values = ["all"] * len(x_values)
+
+        groups: dict[str, list[tuple[float, float]]] = {}
+        group_indices: dict[str, list[int]] = {}
+
+        for i, (g_val, x_val, y_val) in enumerate(
+            zip(group_values, x_values, y_values, strict=False)
+        ):
+            if x_val is None or y_val is None:
+                continue
+            key = str(g_val) if g_val is not None else "null"
+            groups.setdefault(key, []).append((float(x_val), float(y_val)))
+            group_indices.setdefault(key, []).append(i)
+
+        return groups, group_indices
+
+    def _get_sorted_group_keys(self, groups: dict[str, list[Any]]) -> list[str]:
+        """Return group keys sorted according to ``group_by`` ordering."""
+        group_keys = list(groups.keys())
+        if self.group_by is None:
+            return group_keys
+        order_indices = self.group_by.ordering.argsort(group_keys)
+        return [str(group_keys[i]) for i in order_indices]
+
+    @staticmethod
+    def _pick_group_color(
+        indices: list[int],
+        resolved_color: ResolvedColor,
+    ) -> ResolvedColor:
+        """Build a per-group ResolvedColor from a global ResolvedColor.
+
+        For continuous ``values``, slices to entries belonging to this group.
+        For categorical per-entry colors, picks the first entry's color
+        (within a group all entries share the same categorical color).
+        For a single color string, passes through as-is.
+        """
+        if resolved_color.values is not None:
+            filtered = [resolved_color.values[i] for i in indices]
+            return ResolvedColor(values=filtered)
+
+        if resolved_color.color is not None:
+            if isinstance(resolved_color.color, list):
+                if not resolved_color.color or not indices:
+                    return ResolvedColor()
+                return ResolvedColor(color=resolved_color.color[indices[0]])
+            return resolved_color
+
+        return resolved_color
+
+    @staticmethod
+    def _pick_group_scalar(
+        indices: list[int],
+        resolved: str | list[str] | None,
+    ) -> str | None:
+        """Pick a single categorical value for a group (first match).
+
+        If resolved is scalar or None — return as-is.
+        If resolved is a per-entry list — return the first entry's value.
+        """
+        if resolved is None or not isinstance(resolved, list):
+            return resolved
+        if not resolved or not indices:
+            return None
+        return resolved[indices[0]]
+
+    @staticmethod
+    def _pick_group_values(
+        indices: list[int],
+        resolved: float | list[float] | None,
+    ) -> float | list[float] | None:
+        """Slice a resolved numeric property for a group.
+
+        If resolved is scalar or None — return as-is.
+        If resolved is a per-entry list — filter to entries in this group.
+        """
+        if resolved is None or not isinstance(resolved, list):
+            return resolved
+        return [resolved[i] for i in indices]
+
+    def _build_figure(self) -> go.Figure:
+        """Build the Plotly figure with scatter traces."""
+        fig = go.Figure()
+
+        if self._traces_data is None:
+            return fig
+
+        for trace_id, trace_data in self._traces_data.items():
+            group = trace_data["group"]
+            x_vals = trace_data["x_values"]
+            y_vals = trace_data["y_values"]
+
+            if not x_vals:
+                continue
+
+            fig.add_trace(
+                go.Scatter(
+                    x=x_vals,
+                    y=y_vals,
+                    name=group,
+                    mode="markers",
+                    legendgroup=group,
+                    meta={"trace_id": trace_id},
+                )
+            )
+
+        # Title
+        if self.title:
+            fig.update_layout(title=self.title)
+        else:
+            fig.update_layout(title=f"{self.y.field} vs {self.x.field}")
+
+        return fig