reader-workbench 1.0.0__py3-none-any.whl

reader_workbench/domains/plate_reader/plots/snapshot_barplot/planning.py

@@ -0,0 +1,168 @@
+from __future__ import annotations
+
+from typing import Literal
+
+import numpy as np
+import pandas as pd
+
+from reader_workbench.domains.plate_reader.ordering import order_levels
+from reader_workbench.domains.plate_reader.plots.grouping import GroupMatch, resolve_groups
+
+_PanelBy = Literal["channel", "x", "group"]
+
+
+def build_figure_groups(
+    *,
+    stats: pd.DataFrame,
+    group_col: str | None,
+    panel_by: _PanelBy,
+    pool_sets: list[dict[str, list[str]]] | None,
+    pool_match: GroupMatch,
+) -> list[tuple[str, list[str | None]]]:
+    if not group_col:
+        return [("all", [None])]
+
+    universe = order_levels(stats[group_col].astype(str).unique().tolist())
+    if panel_by == "group":
+        if pool_sets:
+            resolved = resolve_groups(universe, pool_sets, match=pool_match)
+            ordered: list[str] = []
+            seen: set[str] = set()
+            for _, values in resolved:
+                for value in values:
+                    if value in universe and value not in seen:
+                        ordered.append(value)
+                        seen.add(value)
+            return [("all", ordered or universe)]
+        return [("all", universe)]
+
+    if pool_sets:
+        resolved = resolve_groups(universe, pool_sets, match=pool_match)
+        union: list[str] = []
+        seen: set[str] = set()
+        for _, values in resolved:
+            for value in values:
+                if value in universe and value not in seen:
+                    union.append(value)
+                    seen.add(value)
+        return [(group, [group]) for group in (union or universe)]
+    return [(group, [group]) for group in universe]
+
+
+def resolve_panel_configuration(
+    *,
+    panel_by: _PanelBy,
+    members: list[str | None],
+    y_list: list[str],
+    group_col: str | None,
+    channel_select: str | None,
+    selected_channel_default: str | None,
+    ch_for_file: str | None,
+    stats: pd.DataFrame,
+    x_col: str,
+) -> tuple[list[str], str | None]:
+    if panel_by == "channel":
+        return y_list, None
+    if panel_by == "group":
+        if not group_col:
+            raise ValueError("panel_by='group' requires 'group_on'")
+        selected_channel = (
+            (str(ch_for_file) if ch_for_file is not None else None) or channel_select or selected_channel_default
+        )
+        if not selected_channel:
+            raise ValueError(
+                "panel_by='group' requires an explicit channel: set 'channel_select' or pass a single string to 'y'."
+            )
+        return [str(member) for member in members], str(selected_channel)
+
+    selected_channel = channel_select or selected_channel_default
+    if not selected_channel:
+        raise ValueError(
+            "panel_by='x' requires an explicit channel: set 'channel_select' or pass a single string to 'y'."
+        )
+    sub_stats = stats[stats["channel"].astype(str) == str(selected_channel)]
+    x_levels = sub_stats[x_col].astype(str).unique().tolist()
+    return order_levels(x_levels), str(selected_channel)
+
+
+def compute_shared_ylim(
+    *,
+    stats: pd.DataFrame,
+    snapped: pd.DataFrame,
+    panels: list[str],
+    panel_by: _PanelBy,
+    selected_channel: str,
+    group_col: str | None,
+    x_col: str,
+    agg: str,
+    dispersion: str,
+) -> tuple[float | None, float | None]:
+    if panel_by not in {"group", "x"} or len(panels) <= 1:
+        return None, None
+
+    def _collect_limits_for_subset(
+        sbar_sub: pd.DataFrame,
+        srep_sub: pd.DataFrame,
+    ) -> tuple[float | None, float | None]:
+        values = (
+            pd.to_numeric(sbar_sub[agg], errors="coerce").dropna() if not sbar_sub.empty else pd.Series([], dtype=float)
+        )
+        value_min = float(values.min()) if not values.empty else None
+        value_max = float(values.max()) if not values.empty else None
+        if dispersion == "sd" and "std" in sbar_sub.columns:
+            center = pd.to_numeric(sbar_sub[agg], errors="coerce")
+            spread = pd.to_numeric(sbar_sub["std"], errors="coerce")
+            bottom = (center - spread).dropna()
+            top = (center + spread).dropna()
+            if not bottom.empty:
+                value_min = min(value_min if value_min is not None else np.inf, float(bottom.min()))
+            if not top.empty:
+                value_max = max(value_max or -np.inf, float(top.max()))
+        elif dispersion == "iqr" and {"q1", "q3"}.issubset(sbar_sub.columns):
+            if agg == "median":
+                top = pd.to_numeric(sbar_sub["q3"], errors="coerce").dropna()
+                if not top.empty:
+                    value_max = max(value_max or -np.inf, float(top.max()))
+            else:
+                half = 0.5 * (
+                    pd.to_numeric(sbar_sub["q3"], errors="coerce") - pd.to_numeric(sbar_sub["q1"], errors="coerce")
+                )
+                top = (pd.to_numeric(sbar_sub[agg], errors="coerce") + half).dropna()
+                if not top.empty:
+                    value_max = max(value_max or -np.inf, float(top.max()))
+        if not srep_sub.empty and "value" in srep_sub.columns:
+            observation_values = pd.to_numeric(srep_sub["value"], errors="coerce").dropna()
+            if not observation_values.empty:
+                value_min = min(value_min if value_min is not None else float("inf"), float(observation_values.min()))
+                value_max = max(value_max if value_max is not None else float("-inf"), float(observation_values.max()))
+        return value_min, value_max
+
+    low_values: list[float] = []
+    high_values: list[float] = []
+    for panel in panels:
+        if panel_by == "group":
+            sbar_sub = stats[
+                (stats["channel"].astype(str) == str(selected_channel)) & (stats[group_col].astype(str) == str(panel))
+            ]
+            srep_sub = snapped[
+                (snapped["channel"].astype(str) == str(selected_channel))
+                & (snapped[group_col].astype(str) == str(panel))
+            ]
+        else:
+            sbar_sub = stats[
+                (stats["channel"].astype(str) == str(selected_channel)) & (stats[x_col].astype(str) == str(panel))
+            ]
+            srep_sub = snapped[
+                (snapped["channel"].astype(str) == str(selected_channel)) & (snapped[x_col].astype(str) == str(panel))
+            ]
+        value_min, value_max = _collect_limits_for_subset(sbar_sub, srep_sub)
+        if value_min is not None:
+            low_values.append(value_min)
+        if value_max is not None:
+            high_values.append(value_max)
+    if not high_values:
+        return None, None
+    y_lo = min(0.0, float(np.nanmin(low_values))) if low_values else 0.0
+    y_hi = float(np.nanmax(high_values))
+    pad = 0.05 * (y_hi - y_lo if y_hi > y_lo else max(1.0, y_hi or 1.0))
+    return (y_lo - pad if y_lo < 0.0 else 0.0), y_hi + pad

reader_workbench/domains/plate_reader/plots/snapshot_heatmap/__init__.py

@@ -0,0 +1,205 @@
+"""Snapshot heatmap with tidy aggregation and square cells."""
+
+from __future__ import annotations
+
+import hashlib
+import logging
+from contextlib import suppress
+from typing import Any
+
+import matplotlib.colors as mcolors
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import polars as pl
+from matplotlib.colors import Colormap, LinearSegmentedColormap
+
+from reader_workbench.plotting.sinks import PlotFigure
+from reader_workbench.plotting.style import new_fig_ax, use_style
+
+from ...analysis.timepoints import choose_nearest_time
+from ..common import alias_column, plot_figure, pretty_name, warn_if_empty
+from .inputs import (
+    SnapshotHeatmapInputs as SnapshotHeatmapInputs,
+)
+from .inputs import (
+    prepare_snapshot_heatmap_inputs as prepare_snapshot_heatmap_inputs,
+)
+
+__all__ = ["SnapshotHeatmapInputs", "plot_snapshot_heatmap", "prepare_snapshot_heatmap_inputs"]
+
+
+def _ensure(df: pd.DataFrame, cols: list[str]) -> None:
+    missing = [column for column in cols if column not in df.columns]
+    if missing:
+        raise ValueError(f"snapshot_heatmap: required columns missing: {missing}")
+
+
+def _apply_explicit_axis_order(
+    pivot: pd.DataFrame,
+    *,
+    axis: str,
+    labels: list[str] | None,
+) -> pd.DataFrame:
+    if labels is None:
+        return pivot
+    if not labels:
+        raise ValueError(f"snapshot_heatmap: {axis} order must not be empty")
+
+    available = list(map(str, pivot.columns if axis == "x" else pivot.index))
+    missing = [label for label in labels if label not in available]
+    if missing:
+        raise ValueError(f"snapshot_heatmap: {axis} order includes missing label(s): {missing}")
+
+    if axis == "x":
+        return pivot[labels]
+    return pivot.loc[labels]
+
+
+def _resolve_cmap(spec: Any) -> Colormap | None:
+    if spec is None:
+        return None
+    if isinstance(spec, mcolors.Colormap):
+        return spec
+    if isinstance(spec, list | tuple):
+        return LinearSegmentedColormap.from_list("custom", list(spec))
+    if isinstance(spec, str):
+        try:
+            return plt.get_cmap(spec)
+        except Exception:
+            return LinearSegmentedColormap.from_list("custom", ["#FFFFFF", spec])
+    return None
+
+
+def plot_snapshot_heatmap(
+    *,
+    df: pd.DataFrame,
+    channel: str,
+    time: float,
+    x: str = "treatment",
+    y: str = "design_id",
+    order_x: list[str] | None = None,
+    order_y: list[str] | None = None,
+    square: bool = True,
+    vmin: float | None = None,
+    vmax: float | None = None,
+    fig_kwargs: dict[str, Any],
+    filename: str | None,
+    logger: logging.Logger | None = None,
+) -> list[PlotFigure]:
+    x_col = alias_column(df, x)
+    y_col = alias_column(df, y)
+    _ensure(df, ["time", "channel", "value", x_col, y_col])
+
+    work_pl = pl.from_pandas(df)
+    time_expr = pl.col("time").cast(pl.Float64, strict=True)
+    value_expr = pl.col("value").cast(pl.Float64, strict=False)
+    value_expr = pl.when(value_expr.is_nan()).then(None).otherwise(value_expr)
+    work_pl = work_pl.with_columns(time_expr.alias("time"), value_expr.alias("value"))
+
+    rc = (fig_kwargs or {}).get("rc", {})
+    tol = (fig_kwargs or {}).get("time_tolerance", rc.get("time_tolerance", None))
+    times = (
+        work_pl.filter(pl.col("channel").cast(pl.Utf8) == str(channel))
+        .select(pl.col("time").fill_null(float("nan")).alias("time"))
+        .to_numpy()
+        .ravel()
+    )
+    tsel = choose_nearest_time(
+        times,
+        target_time=float(time),
+        tol=(None if tol is None else float(tol)),
+        where="snapshot_heatmap",
+        logger=logger,
+    )
+
+    snap_pl = work_pl.filter(
+        (pl.col("channel").cast(pl.Utf8) == str(channel)) & ((pl.col("time") - float(tsel)).abs() <= 1e-9)
+    )
+    if snap_pl.is_empty():
+        raise ValueError("snapshot_heatmap: no rows after (channel,time) selection")
+
+    pivot_pl = (
+        snap_pl.group_by([y_col, x_col])
+        .agg(pl.col("value").median().alias("value"))
+        .sort([y_col, x_col])
+        .pivot(values="value", index=y_col, on=x_col, aggregate_function="first")
+    )
+    pivot = pivot_pl.to_pandas(use_pyarrow_extension_array=False)
+    if y_col in pivot.columns:
+        pivot = pivot.set_index(y_col)
+
+    pivot = _apply_explicit_axis_order(pivot, axis="x", labels=order_x)
+    pivot = _apply_explicit_axis_order(pivot, axis="y", labels=order_y)
+
+    if pivot.empty or pivot.shape[0] == 0 or pivot.shape[1] == 0:
+        warn_if_empty(pivot, where="snapshot_heatmap", detail="after pivot")
+        return []
+
+    with use_style(rc=(fig_kwargs or {}).get("rc")):
+        fig, ax = new_fig_ax(fig_kwargs)
+        arr = pivot.to_numpy(dtype=float)
+        if not np.isfinite(arr).any():
+            raise ValueError("snapshot_heatmap: selected snapshot contains no finite values")
+        cmap = _resolve_cmap((fig_kwargs or {}).get("cmap"))
+        resolved_vmin = float(np.nanmin(arr)) if vmin is None else float(vmin)
+        resolved_vmax = float(np.nanmax(arr)) if vmax is None else float(vmax)
+
+        im = ax.imshow(
+            arr,
+            aspect=("equal" if square else "auto"),
+            vmin=resolved_vmin,
+            vmax=resolved_vmax,
+            cmap=cmap,
+        )
+
+        ax.grid(False)
+        with suppress(Exception):
+            ax.set_facecolor("white")
+
+        ax.set_xticks(range(arr.shape[1]))
+        ax.set_xticklabels(list(map(str, pivot.columns)), rotation=45, ha="right")
+        ax.set_yticks(range(arr.shape[0]))
+        ax.set_yticklabels(list(map(str, pivot.index)))
+        ax.set_xlabel(pretty_name(str(x_col)).title())
+        ax.set_ylabel(pretty_name(str(y_col)).title())
+        ax.set_title(f"{channel} @ t={tsel:g} h")
+
+        ax.set_xticks(np.arange(-0.5, arr.shape[1], 1), minor=True)
+        ax.set_yticks(np.arange(-0.5, arr.shape[0], 1), minor=True)
+        ax.grid(which="minor", color="white", linestyle="-", linewidth=0.5)
+        ax.tick_params(which="minor", bottom=False, left=False)
+
+        shrink = float((fig_kwargs or {}).get("cbar_shrink", 1.0))
+        cbar_label: str | None = (fig_kwargs or {}).get("cbar_label")
+        fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04, shrink=shrink, label=cbar_label)
+
+        base = filename or f"snapshot_heatmap__{channel}__t{tsel:g}h"
+        y_levels = list(map(str, pivot.index))
+        n_geno = len(y_levels)
+
+        def _short_id(value: str) -> str:
+            return hashlib.blake2b(value.encode("utf-8"), digest_size=4).hexdigest()
+
+        geno_id = _short_id("|".join(sorted(y_levels)))
+        x_levels = list(map(str, pivot.columns))
+        cmap_name = (
+            cmap.name if isinstance(cmap, Colormap) and hasattr(cmap, "name") else "custom" if cmap else "default"
+        )
+        tol = (fig_kwargs or {}).get("time_tolerance", None)
+        fp_payload = "|".join(
+            [
+                f"ch={channel}",
+                f"t={tsel:g}",
+                "x=" + ",".join(x_levels),
+                "y=" + ",".join(y_levels),
+                f"cmap={cmap_name}",
+                f"vmin={'nan' if vmin is None else float(vmin)}",
+                f"vmax={'nan' if vmax is None else float(vmax)}",
+                f"square={bool(square)}",
+                f"tol={'' if tol is None else float(tol)}",
+            ]
+        )
+        fp_id = _short_id(fp_payload)
+        stub = f"{base}__gy{n_geno}-{geno_id}__fp{fp_id}"
+        return [plot_figure(fig=fig, filename=stub, fig_kwargs=fig_kwargs)]

reader_workbench/domains/plate_reader/plots/snapshot_heatmap/inputs.py

@@ -0,0 +1,103 @@
+from __future__ import annotations
+
+import logging
+from collections.abc import Mapping
+from dataclasses import dataclass
+from typing import Any
+
+import numpy as np
+import pandas as pd
+
+from reader_workbench.domains.plate_reader.analysis.timepoints import choose_nearest_time
+
+
+@dataclass(frozen=True)
+class SnapshotHeatmapInputs:
+    """Prepared dataframe and rendering values for a snapshot heatmap."""
+
+    frame: pd.DataFrame
+    filename: str | None
+    fig_kwargs: dict[str, Any]
+
+
+def auto_cbar_label(channel: str, value_transform: str | None) -> str:
+    transform = (value_transform or "none").lower()
+    if channel.startswith("log2FC_"):
+        return f"log2FC ({channel.split('_', 1)[1]})"
+    if channel.startswith("FC_"):
+        return f"FC ({channel.split('_', 1)[1]})"
+    if transform == "log2":
+        return f"log2({channel})"
+    if transform == "log10":
+        return f"log10({channel})"
+    return channel
+
+
+def _transform_positive_values(values: pd.Series, *, transform: str) -> pd.Series:
+    numeric = pd.to_numeric(values, errors="coerce")
+    result = pd.Series(np.nan, index=numeric.index, dtype=float)
+    positive = numeric > 0
+    if transform == "log2":
+        result.loc[positive] = np.log2(numeric.loc[positive])
+        return result
+    result.loc[positive] = np.log10(numeric.loc[positive])
+    return result
+
+
+def prepare_snapshot_heatmap_inputs(
+    *,
+    df_in: pd.DataFrame | None,
+    fc_in: pd.DataFrame | None,
+    channel: str,
+    time: float,
+    time_tolerance: float,
+    value_transform: str | None,
+    fig: Mapping[str, Any] | None,
+    filename: str | None,
+    logger: logging.Logger | None = None,
+) -> SnapshotHeatmapInputs:
+    channel = str(channel)
+    wants_fc = channel.startswith("FC_") or channel.startswith("log2FC_")
+    fig_kwargs = dict(fig or {})
+    fig_kwargs.setdefault("time_tolerance", time_tolerance)
+
+    if wants_fc:
+        if fc_in is None:
+            raise ValueError(f"snapshot_heatmap: channel={channel!r} requires a fold_change.v1 input (reads.fc)")
+        target = channel.split("_", 1)[1]
+        use_col = "log2FC" if channel.startswith("log2FC_") else "FC"
+        table = fc_in.copy()
+        table = table[table["target"].astype(str) == target].copy()
+        if table.empty:
+            raise ValueError(f"snapshot_heatmap: no fold_change rows found for target={target!r}")
+        chosen_time = choose_nearest_time(
+            table["time"],
+            target_time=float(time),
+            tol=float(time_tolerance),
+            where="snapshot_heatmap",
+            logger=logger,
+        )
+        subset = table[pd.to_numeric(table["time"], errors="coerce") == chosen_time].copy()
+        subset = subset.rename(columns={use_col: "value"})
+        subset["channel"] = channel
+        keep = ["time", "channel", "value", "treatment", "design_id", "design_id_alias"]
+        df = subset[[column for column in keep if column in subset.columns]].copy()
+        resolved_filename = filename or f"snapshot_heatmap__{channel}__t{chosen_time:g}h"
+        fig_kwargs.setdefault("cbar_label", auto_cbar_label(channel, None))
+        return SnapshotHeatmapInputs(frame=df, filename=resolved_filename, fig_kwargs=fig_kwargs)
+
+    if df_in is None:
+        raise ValueError("snapshot_heatmap: tidy df is required when channel is not FC_/log2FC_-prefixed")
+
+    df = df_in.copy()
+    if value_transform and str(value_transform).lower() in {"log2", "log10"}:
+        mask = df["channel"].astype(str) == channel
+        df.loc[mask, "value"] = _transform_positive_values(
+            df.loc[mask, "value"],
+            transform=str(value_transform).lower(),
+        )
+    fig_kwargs.setdefault("cbar_label", auto_cbar_label(channel, value_transform))
+    return SnapshotHeatmapInputs(frame=df, filename=filename, fig_kwargs=fig_kwargs)
+
+
+__all__ = ["SnapshotHeatmapInputs", "auto_cbar_label", "prepare_snapshot_heatmap_inputs"]