diffsvg 0.1.0__py3-none-any.whl

diffsvg/pl/plot_results.py

@@ -0,0 +1,484 @@
+"""Plots for the output of `diffsvg.tl.fit_model`: the statistics table and the per-condition curves."""
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+
+# ============================================================
+# Label placement helpers (used by plot_volcano)
+# ============================================================
+# A dependency-free stand-in for `adjustText`: try a ring of offsets around each point and take the first
+# label box that overlaps nothing placed so far, else the least-overlapping one. Best-effort.
+
+
+def _get_renderer(fig):
+    try:
+        fig.canvas.draw()
+        return fig.canvas.get_renderer()
+    except Exception:  # noqa: BLE001 - collision avoidance is a nice-to-have, never worth crashing the plot
+        return None
+
+
+def _bbox_tuple(bbox):
+    return (bbox.x0, bbox.y0, bbox.x1, bbox.y1)
+
+
+def _boxes_overlap(a, b, pad=2.0):
+    ax0, ay0, ax1, ay1 = a
+    bx0, by0, bx1, by1 = b
+    return not (ax1 + pad < bx0 or bx1 + pad < ax0 or ay1 + pad < by0 or by1 + pad < ay0)
+
+
+def _dot_boxes(ax, xs, ys, marker_size, pad_px=2.0):
+    """Small display-coordinate boxes around each `(x, y)` scatter point, for label-collision checks."""
+    if len(xs) == 0:
+        return []
+    pts = ax.transData.transform(np.column_stack([xs, ys]))
+    r = np.sqrt(marker_size / np.pi) * ax.get_figure().dpi / 72.0 + pad_px
+    return [(px - r, py - r, px + r, py + r) for px, py in pts]
+
+
+def _place_label(
+    ax,
+    renderer,
+    x,
+    y,
+    text,
+    occupied,
+    fontsize=7,
+    color="black",
+    fontweight="normal",
+    preferred_offset=None,
+    base_radius=10.0,
+    n_angles=10,
+    n_rings=4,
+    arrow=False,
+    bbox_kwargs=None,
+    zorder=12,
+):
+    """Place `text` near data point `(x, y)`, avoiding the display-coordinate boxes in `occupied`.
+
+    Tries `preferred_offset` (data units) first, then rings of pixel offsets; takes the first label box
+    that overlaps nothing in `occupied`, else the least-overlapping one, and appends it to `occupied`.
+    Without a `renderer` the label goes at `preferred_offset` (or the point) unchecked.
+    Returns the `matplotlib.text.Annotation`.
+    """
+    point_disp = ax.transData.transform((x, y))
+    style = {
+        "textcoords": "data",
+        "fontsize": fontsize,
+        "color": color,
+        "fontweight": fontweight,
+        "ha": "center",
+        "va": "center",
+        "zorder": zorder,
+        "bbox": bbox_kwargs,
+    }
+
+    def _measure(dx, dy):
+        # Without the arrow: its window extent reaches back to the point and would count as occupied.
+        cand = ax.transData.inverted().transform((point_disp[0] + dx, point_disp[1] + dy))
+        return cand, ax.annotate(text, xy=(x, y), xytext=cand, **style)
+
+    def _finalize(cand, ann):
+        if not arrow:
+            return ann
+        ann.remove()
+        arrowprops = {"arrowstyle": "->", "lw": 1.0, "color": color}
+        return ax.annotate(text, xy=(x, y), xytext=cand, arrowprops=arrowprops, **style)
+
+    if renderer is None:
+        dx, dy = (
+            (0.0, 0.0)
+            if preferred_offset is None
+            else tuple(ax.transData.transform((x + preferred_offset[0], y + preferred_offset[1])) - point_disp)
+        )
+        cand, ann = _measure(dx, dy)
+        return _finalize(cand, ann)
+
+    candidates = []
+    if preferred_offset is not None:
+        pdx, pdy = preferred_offset
+        candidates.append(tuple(ax.transData.transform((x + pdx, y + pdy)) - point_disp))
+    for ring in range(n_rings):
+        radius = base_radius * (1 + ring)
+        for i in range(n_angles):
+            angle = 2 * np.pi * i / n_angles + ring * 0.35
+            candidates.append((np.cos(angle) * radius, np.sin(angle) * radius))
+
+    best, best_score, best_box, best_cand = None, None, None, None
+    for dx, dy in candidates:
+        cand, ann = _measure(dx, dy)
+        box = _bbox_tuple(ann.get_window_extent(renderer=renderer))
+        score = sum(1 for other in occupied if _boxes_overlap(box, other))
+        if score == 0:
+            if best is not None:
+                best.remove()
+            occupied.append(box)
+            return _finalize(cand, ann)
+        # On ties prefer the later candidate: labels sharing a preferred spot would otherwise stack.
+        if best_score is None or score <= best_score:
+            if best is not None:
+                best.remove()
+            best, best_score, best_box, best_cand = ann, score, box, cand
+        else:
+            ann.remove()
+
+    occupied.append(best_box)
+    return _finalize(best_cand, best)
+
+
+def _is_parallel_table(stats_df):
+    """Whether `stats_df` comes from ``fit_model(..., test="parallel")`` rather than the Wald test."""
+    return "parallel_pval" in stats_df.columns
+
+
+_AXIS_LABELS = {
+    # `integral` is the area under f1 - f0 (condition 1 minus condition 0).
+    "integral": r"$\Delta_{f_1 - f_0}$",
+    "slope_diff": r"$\Delta$ slope ($f_1 - f_0$)",
+    "level_diff": r"$\Delta$ level ($f_1 - f_0$)",
+    "shape_diff": "shape difference",
+}
+
+
+def plot_volcano(
+    stats_df,
+    x=None,
+    pvalue=None,
+    pvalue_adj=None,
+    alpha=0.05,
+    show_dots=True,
+    n_label=20,
+    highlight_genes=None,
+    highlight_offset=(0.025, 1.25),
+    ax=None,
+    figsize=(4, 6),
+    save_path=None,
+):
+    """Volcano plot of per-gene curve difference vs. significance.
+
+    Parameters
+    ----------
+    stats_df
+        A statistics table indexed by gene, e.g. ``adata.varm["dsvg_stats"]``.
+    x
+        Column used for the x-axis. Defaults to ``"integral"`` (area between
+        the two condition curves) for the Wald test and ``"slope_diff"``
+        (linear trend of ``f1 - f0`` over the anchor range) for
+        ``test="parallel"``; ``"integral"`` there shows level against shape.
+    pvalue
+        Column used for the y-axis, plotted as ``-log10(pvalue)``. Every
+        point above the ``-log10(alpha)`` reference line is colored by the
+        sign of `x` (blue = positive, red = negative); points at or below it
+        are black. Defaults to ``"wald_pval"`` for the Wald test and
+        ``"parallel_pval"`` for ``test="parallel"``. p-values of exactly 0
+        (underflow) are drawn at the smallest non-zero one; genes with a
+        missing p-value or `x` (e.g. not fitted) are left out.
+    pvalue_adj
+        Column used to rank genes for labeling (the `n_label` genes with the
+        smallest `pvalue_adj` among those above the line get a text label).
+        Defaults to the adjusted version of the default `pvalue`.
+    alpha
+        Significance threshold on `pvalue`, drawn as a horizontal reference
+        line at ``-log10(alpha)`` and used for the color cutoff described
+        under `pvalue`.
+    show_dots
+        If `True` (default), every gene is drawn as a colored dot. If
+        `False`, every gene's name is drawn instead, in the same color, in
+        place of a dot - there is no separate background scatter in that
+        mode. Best suited to modest gene panels (a few hundred genes);
+        subset `stats_df` first for larger ones.
+    n_label
+        Label at most this many of the genes above the line (the ones with
+        the smallest `pvalue_adj`) with their gene name. Use `None` to label
+        every gene above the line, or `0` to label none. Labels use the
+        same collision-avoidance as `highlight_genes` (see below).
+    highlight_genes
+        A gene name, or list of gene names, to draw with a larger red point
+        and a bold, boxed, arrow-annotated label on top of the rest. Raises
+        `KeyError` if any given gene is not in `stats_df.index`.
+    highlight_offset
+        ``(dx, dy)`` offset (in data units) tried first for each
+        `highlight_genes` label; if that spot is already occupied (by a
+        dot, another gene's name, or another label), a nearby spot is used
+        instead - so with several highlighted genes this is a preference,
+        not a guarantee.
+    ax
+        Existing axes to draw into; a new figure is created if omitted.
+    figsize
+        Passed to `plt.subplots` when `ax` is omitted.
+    save_path
+        If given, save the figure here (``dpi=300, bbox_inches="tight"``).
+
+    Returns
+    -------
+    matplotlib.axes.Axes
+    """
+    if _is_parallel_table(stats_df):
+        defaults = ("slope_diff", "parallel_pval", "parallel_pval_adj")
+    elif "slope_pval" in stats_df.columns:
+        defaults = ("slope_diff", "slope_pval", "slope_pval_adj")
+    elif "level_pval" in stats_df.columns:
+        defaults = ("level_diff", "level_pval", "level_pval_adj")
+    else:
+        defaults = ("integral", "wald_pval", "wald_pval_adj")
+    x, pvalue, pvalue_adj = (
+        given if given is not None else default
+        for given, default in zip((x, pvalue, pvalue_adj), defaults, strict=True)
+    )
+
+    if ax is None:
+        _, ax = plt.subplots(figsize=figsize)
+    fig = ax.get_figure()
+
+    if highlight_genes is None:
+        highlight_genes = []
+    elif isinstance(highlight_genes, str):
+        highlight_genes = [highlight_genes]
+    else:
+        highlight_genes = list(highlight_genes)
+    missing = [g for g in highlight_genes if g not in stats_df.index]
+    if missing:
+        raise KeyError(f"highlight_genes not found in stats_df.index: {missing}")
+
+    # Missing p-values (genes that were not fitted) would otherwise be drawn as the most significant.
+    stats_df = stats_df[
+        np.isfinite(stats_df[pvalue].to_numpy(dtype=float)) & np.isfinite(stats_df[x].to_numpy(dtype=float))
+    ]
+    missing = [g for g in highlight_genes if g not in stats_df.index]
+    if missing:
+        raise KeyError(f"highlight_genes have no p-value or `x`: {missing}")
+    x_vals = stats_df[x]
+    pvals = stats_df[pvalue].to_numpy(dtype=float)
+    positive = pvals[pvals > 0]
+    floor = positive.min() if positive.size else np.finfo(float).tiny
+    minus_log10_p = pd.Series(-np.log10(np.where(pvals > 0, pvals, floor)), index=stats_df.index)
+    threshold = -np.log10(alpha)
+    sig = minus_log10_p > threshold
+    pos_sign = x_vals > 0
+    neg_sign = x_vals < 0
+
+    ax.axvline(0, color="black", linestyle="--", lw=1)
+    ax.axhline(threshold, color="gray", linestyle="--", lw=1)
+
+    renderer = _get_renderer(fig) if (n_label != 0 or highlight_genes or not show_dots) else None
+    occupied = []
+
+    if show_dots:
+        ax.scatter(x_vals[~sig], minus_log10_p[~sig], color="black", s=20, alpha=0.7, linewidth=0)
+        ax.scatter(x_vals[sig & pos_sign], minus_log10_p[sig & pos_sign], color="blue", s=20, alpha=0.7, linewidth=0)
+        ax.scatter(x_vals[sig & neg_sign], minus_log10_p[sig & neg_sign], color="red", s=20, alpha=0.7, linewidth=0)
+        if renderer is not None:
+            occupied.extend(_dot_boxes(ax, x_vals.to_numpy(), minus_log10_p.to_numpy(), marker_size=20))
+    else:
+        for gene in stats_df.index:
+            color = "black" if not sig[gene] else ("blue" if pos_sign[gene] else "red")
+            txt = ax.text(
+                x_vals[gene], minus_log10_p[gene], gene, fontsize=6, color=color, alpha=0.85, ha="center", va="center"
+            )
+            if renderer is not None:
+                occupied.append(_bbox_tuple(txt.get_window_extent(renderer=renderer)))
+
+    label_genes = stats_df.index[sig]
+    if n_label is not None and len(label_genes) > n_label:
+        label_genes = stats_df.loc[label_genes, pvalue_adj].sort_values().index[:n_label]
+
+    for gene in label_genes:
+        if gene in highlight_genes:
+            continue  # gets the larger highlight treatment below instead
+        _place_label(ax, renderer, x_vals[gene], minus_log10_p[gene], gene, occupied, fontsize=7, base_radius=8.0)
+
+    for gene in highlight_genes:
+        hx, hy = x_vals[gene], minus_log10_p[gene]
+        ax.scatter(hx, hy, s=80, color="red", edgecolor="black", linewidth=1.5, zorder=10)
+        _place_label(
+            ax,
+            renderer,
+            hx,
+            hy,
+            gene,
+            occupied,
+            fontsize=13,
+            fontweight="bold",
+            preferred_offset=highlight_offset,
+            base_radius=45.0,
+            n_rings=6,
+            arrow=True,
+            bbox_kwargs={"boxstyle": "round,pad=0.3", "fc": "white", "ec": "black", "alpha": 0.9},
+            zorder=13,
+        )
+
+    ax.spines["top"].set_visible(False)
+    ax.spines["right"].set_visible(False)
+    ax.set_xlabel(_AXIS_LABELS.get(x, x), fontsize=12)
+    ax.set_ylabel(r"$-\log_{10}(p)$", fontsize=12)
+    ax.tick_params(labelsize=10)
+
+    fig.tight_layout()
+    if save_path is not None:
+        fig.savefig(save_path, dpi=300, bbox_inches="tight")
+
+    return ax
+
+
+def _panel_grid(n_panels, ncols, figsize):
+    """A figure with `n_panels` axes in rows of `ncols`, each panel `figsize` inches; returns ``(fig, axes)``, flat."""
+    ncols = max(1, min(ncols, n_panels))
+    nrows = (n_panels + ncols - 1) // ncols
+    width, height = figsize
+    fig, axes = plt.subplots(nrows, ncols, figsize=(ncols * width, nrows * height), squeeze=False)
+    axes = axes.ravel()
+    for spare in axes[n_panels:]:
+        spare.set_visible(False)
+    return fig, axes[:n_panels]
+
+
+def _draw_prediction_curve(ax, adata, gene, curves_key, condition_labels, colors):
+    """One gene's two condition curves (mean + credible band) on `ax`."""
+    for cat in (0, 1):
+        curves = adata.varm[f"{curves_key}_cond{cat}_mean"]
+        lower = adata.varm[f"{curves_key}_cond{cat}_ci_low"].loc[gene]
+        upper = adata.varm[f"{curves_key}_cond{cat}_ci_high"].loc[gene]
+        # Taken from the columns, so the bin grid is read off the table itself.
+        x = curves.columns.to_numpy(dtype=float)
+        ax.fill_between(x, lower.to_numpy(dtype=float), upper.to_numpy(dtype=float), color=colors[cat], alpha=0.2, lw=0)
+        ax.plot(x, curves.loc[gene].to_numpy(dtype=float), color=colors[cat], label=condition_labels[str(cat)])
+    ax.spines["top"].set_visible(False)
+    ax.spines["right"].set_visible(False)
+    ax.set_xlabel("Normalized distance")
+    ax.set_ylabel("smoothed expression")
+    ax.set_title(gene)
+
+
+def plot_prediction_curve(
+    adata,
+    gene,
+    curves_key="dsvg_curves",
+    condition_labels=None,
+    colors=None,
+    ncols=4,
+    figsize=(4, 3.2),
+    ax=None,
+    dpi=300,
+    save_path=None,
+):
+    """Plot the predicted expression curve (posterior mean +/- credible band) per gene, for both conditions.
+
+    Reads the six ``adata.varm[f"{curves_key}_cond{0,1}_{mean,ci_low,ci_high}"]``
+    tables written by `diffsvg.tl.fit_model`. Several genes are drawn as a
+    grid of panels, one gene each, each with its own axis ranges because genes
+    differ in scale.
+
+    Parameters
+    ----------
+    adata
+        Annotated data matrix that `diffsvg.tl.fit_model` wrote into.
+    gene
+        A gene name, or a list of them (all must be in ``adata.var_names``).
+    curves_key
+        Prefix used by `diffsvg.tl.fit_model` for the per-condition curve
+        `adata.varm` entries.
+    condition_labels
+        ``{"0": label, "1": label}`` for the legend. Defaults to
+        ``adata.uns[f"{curves_key}_condition_categories"]`` if present,
+        otherwise generic ``"condition 0"``/``"condition 1"`` labels.
+    colors
+        ``{0: color, 1: color}`` for the two curves.
+    ncols
+        Panels per row when several genes are given. Fewer genes than this
+        give a single, narrower row.
+    figsize
+        ``(width, height)`` of a *single* panel in inches.
+    ax
+        Existing axes to draw into, for a single gene; a new figure is created
+        if omitted.
+    dpi
+        Resolution used when saving to `save_path`.
+    save_path
+        If given, save the figure here (``bbox_inches="tight"``).
+
+    Returns
+    -------
+    A `matplotlib.axes.Axes` when one gene is given, otherwise an array of
+    them, one per gene in the order passed.
+
+    Examples
+    --------
+    >>> dvg.pl.plot_prediction_curve(adata, gene=["PAX7", "IRX3"])
+    """
+    genes = [gene] if isinstance(gene, str) else list(gene)
+    if not genes:
+        raise ValueError("`gene` is empty - pass at least one gene name.")
+    missing = [g for g in genes if g not in adata.var_names]
+    if missing:
+        raise KeyError(f"genes not found in adata.var_names: {missing}")
+    if ax is not None and len(genes) > 1:
+        raise ValueError("`ax` holds a single panel - leave it out to plot several genes.")
+
+    if condition_labels is None:
+        condition_labels = adata.uns.get(f"{curves_key}_condition_categories", {"0": "condition 0", "1": "condition 1"})
+    if colors is None:
+        colors = {0: "steelblue", 1: "darkorange"}
+
+    fig, axes = (ax.get_figure(), np.array([ax])) if ax is not None else _panel_grid(len(genes), ncols, figsize)
+    for panel, name in zip(axes, genes, strict=True):
+        _draw_prediction_curve(panel, adata, name, curves_key, condition_labels, colors)
+    # One shared legend: repeating it in every panel is just noise.
+    axes[0].legend(frameon=False, fontsize=9)
+    if ax is None:
+        fig.tight_layout()
+
+    if save_path is not None:
+        fig.savefig(save_path, dpi=dpi, bbox_inches="tight")
+    return axes[0] if len(genes) == 1 else axes
+
+
+def plot_pvalue_histogram(
+    stats_df,
+    pvalue=None,
+    bins=20,
+    ax=None,
+    figsize=(4, 4),
+    save_path=None,
+):
+    """QC histogram of a per-gene p-value column; should look roughly uniform under the null.
+
+    Parameters
+    ----------
+    stats_df
+        A statistics table indexed by gene, e.g. ``adata.varm["dsvg_stats"]``.
+    pvalue
+        Column to histogram. Defaults to ``"wald_pval"`` for the Wald
+        test and ``"parallel_pval"`` for ``test="parallel"``.
+    bins
+        Passed to `matplotlib.axes.Axes.hist`.
+    ax
+        Existing axes to draw into; a new figure is created if omitted.
+    figsize
+        Passed to `plt.subplots` when `ax` is omitted.
+    save_path
+        If given, save the figure here (``dpi=300, bbox_inches="tight"``).
+
+    Returns
+    -------
+    matplotlib.axes.Axes
+    """
+    if pvalue is None:
+        pvalue = next(c for c in ("parallel_pval", "slope_pval", "level_pval", "wald_pval") if c in stats_df.columns)
+    if ax is None:
+        _, ax = plt.subplots(figsize=figsize)
+
+    ax.hist(stats_df[pvalue].dropna(), bins=bins, edgecolor="black", linewidth=1.0)
+
+    ax.spines["top"].set_visible(False)
+    ax.spines["right"].set_visible(False)
+    ax.set_xlabel(pvalue, fontsize=12)
+    ax.set_ylabel("count", fontsize=12)
+
+    fig = ax.get_figure()
+    fig.tight_layout()
+    if save_path is not None:
+        fig.savefig(save_path, dpi=300, bbox_inches="tight")
+
+    return ax

diffsvg/pl/plot_stability.py

@@ -0,0 +1,184 @@
+"""Plots for the output of `diffsvg.tl.check_stability`."""
+
+import matplotlib.pyplot as plt
+import numpy as np
+from scipy.cluster.hierarchy import dendrogram, linkage
+from scipy.spatial.distance import squareform
+
+
+def _clustered_order(similarity):
+    """Order runs by hierarchical clustering on ``1 - similarity``, returning ``(order, linkage)``.
+
+    The linkage is what the dendrogram is drawn from; clustering the
+    configurations is the point of the plot, since it shows which parameter
+    changes land in the same place and which one sits on its own branch.
+    """
+    distance = 1.0 - similarity.to_numpy()
+    distance = (distance + distance.T) / 2.0  # squareform insists on exact symmetry
+    np.fill_diagonal(distance, 0.0)
+    if distance.shape[0] < 3:
+        return list(range(distance.shape[0])), None
+    link = linkage(squareform(distance, checks=False), method="average")
+    return dendrogram(link, no_plot=True)["leaves"], link
+
+
+def _heatmap(result, matrix, label, floor, ax, cmap, annotate, title):
+    order, link = _clustered_order(matrix)
+    ordered = matrix.iloc[order, order]
+    values = ordered.to_numpy()
+    n = len(ordered)
+
+    fig = ax.get_figure()
+    tree_ax, colorbar_ax = None, None
+    if ax.get_subplotspec() is not None:
+        # Give the dendrogram and the colorbar their own cells rather than
+        # letting the colorbar steal width from the heatmap: the dendrogram
+        # only lines up with the columns if both share one grid column.
+        rows = [1, 4] if link is not None else [0.001, 4]
+        grid = ax.get_subplotspec().subgridspec(
+            2, 2, height_ratios=rows, width_ratios=[20, 1], hspace=0.04, wspace=0.04
+        )
+        ax.remove()
+        if link is not None:
+            tree_ax = fig.add_subplot(grid[0, 0])
+        ax = fig.add_subplot(grid[1, 0])
+        colorbar_ax = fig.add_subplot(grid[1, 1])
+
+    if tree_ax is not None:
+        dendrogram(link, ax=tree_ax, color_threshold=0, above_threshold_color="black", no_labels=True)
+        # scipy places leaf i at 10*i + 5; pinning the limits to (0, 10n)
+        # makes that the same fractional position as column i under imshow.
+        tree_ax.set_xlim(0, 10 * n)
+        tree_ax.set_axis_off()
+
+    # Scaled to the data rather than a fixed 0-1: agreement is usually
+    # bunched near the top, and a fixed floor would flatten the differences
+    # that matter. Negative correlations stay visible.
+    image = ax.imshow(values, cmap=cmap, vmin=float(values.min()), vmax=1.0, aspect="auto")
+    ax.set_xticks(range(n), ordered.columns, rotation=90, fontsize=7)
+    ax.set_yticks(range(n), ordered.index, fontsize=7)
+
+    if annotate:
+        midpoint = (float(values.min()) + 1.0) / 2
+        for i in range(n):
+            for j in range(n):
+                ax.text(
+                    j,
+                    i,
+                    f"{values[i, j]:.2f}",
+                    ha="center",
+                    va="center",
+                    fontsize=6,
+                    color="white" if values[i, j] < midpoint else "black",
+                )
+
+    colorbar = fig.colorbar(image, cax=colorbar_ax, ax=None if colorbar_ax else ax)
+    colorbar.set_label(label, fontsize=9)
+    if np.isfinite(floor) and float(values.min()) <= floor <= 1.0:
+        colorbar.ax.axhline(floor, color="red", lw=1.5)
+
+    # On the dendrogram axes when there is one, so it clears the tree.
+    (tree_ax or ax).set_title(title, fontsize=11)
+    return ax
+
+
+def _sensitivity(result, metric, ax, color):
+    data = result.sensitivity
+    if data.empty:
+        raise ValueError(
+            "`sensitivity` is empty - it needs configurations that differ from the baseline in exactly "
+            "one parameter, which requires at least one grid axis with more than one value."
+        )
+    labels = [f"{row.parameter}={row.value}" for row in data.itertuples()]
+    positions = np.arange(len(data))
+
+    ax.barh(positions, data[metric].to_numpy(), color=color, alpha=0.85)
+    ax.set_yticks(positions, labels, fontsize=8)
+    ax.invert_yaxis()
+    ax.set_xlim(0, 1)
+    ax.set_xlabel(f"{metric} vs. baseline", fontsize=10)
+
+    floor = result.noise_floor.get(metric, float("nan"))
+    if np.isfinite(floor):
+        ax.axvline(floor, color="red", linestyle="--", lw=1.5, label="noise floor")
+        ax.legend(frameon=False, fontsize=8, loc="lower right")
+    ax.spines["top"].set_visible(False)
+    ax.spines["right"].set_visible(False)
+    return ax
+
+
+def plot_stability(
+    result,
+    kind="heatmap",
+    metric="spearman",
+    ax=None,
+    figsize=None,
+    cmap="viridis",
+    annotate=True,
+    save_path=None,
+):
+    """Visualize a `diffsvg.tl.check_stability` result.
+
+    Parameters
+    ----------
+    result
+        A `diffsvg.tl.StabilityResult`.
+    kind
+        ``"heatmap"`` (default) draws the run-by-run agreement matrix with a
+        dendrogram above it, so configurations that rank genes the same way
+        cluster together. ``"sensitivity"`` draws the one-factor-at-a-time
+        view instead: how far each single parameter change moves the ranking
+        away from the baseline.
+    metric
+        ``"spearman"`` (rank correlation over all genes) or ``"top_overlap"``
+        (fraction of the top `result.top_k` genes shared). The first is the
+        global view, the second is the decision-relevant one.
+    ax
+        Existing axes to draw into; a new figure is created if omitted. For
+        ``kind="heatmap"`` the axes is split to make room for the
+        dendrogram, so pass one that can be subdivided (a normal subplot).
+    figsize
+        Defaults to something that scales with the number of runs.
+    cmap, annotate
+        Heatmap colormap, and whether to write each cell's value into it
+        (readable up to roughly a dozen runs).
+    save_path
+        If given, save the figure here (``dpi=300, bbox_inches="tight"``).
+
+    Returns
+    -------
+    matplotlib.axes.Axes
+        The axes the data was drawn into. Note that for ``kind="heatmap"``
+        this is *not* the `ax` that was passed in, since that one is split.
+    """
+    if metric not in ("spearman", "top_overlap"):
+        raise ValueError(f"Unknown metric {metric!r}; expected 'spearman' or 'top_overlap'.")
+
+    if kind == "heatmap":
+        matrix = result.similarity if metric == "spearman" else result.top_overlap
+        label = "Spearman rank correlation" if metric == "spearman" else f"Top-{result.top_k} overlap"
+        floor = result.noise_floor.get(metric, float("nan"))
+        title = f"Ranking agreement across runs (by {result.rank_by})"
+        if np.isfinite(floor):
+            title += f"\nnoise floor {floor:.3f} (red line on the colorbar)"
+        if figsize is None:
+            side = max(5.0, 0.55 * len(matrix) + 3.0)
+            figsize = (side, side)
+        if ax is None:
+            _, ax = plt.subplots(figsize=figsize)
+        drawn = _heatmap(result, matrix, label, floor, ax, cmap, annotate, title)
+    elif kind == "sensitivity":
+        if figsize is None:
+            figsize = (6.0, max(2.5, 0.4 * len(result.sensitivity) + 1.5))
+        if ax is None:
+            _, ax = plt.subplots(figsize=figsize)
+        drawn = _sensitivity(result, metric, ax, "steelblue")
+        drawn.set_title(f"Effect of one parameter change (baseline: {result.baseline})", fontsize=10)
+    else:
+        raise ValueError(f"Unknown kind {kind!r}; expected 'heatmap' or 'sensitivity'.")
+
+    fig = drawn.get_figure()
+    fig.tight_layout()
+    if save_path is not None:
+        fig.savefig(save_path, dpi=300, bbox_inches="tight")
+    return drawn

diffsvg/pp/__init__.py

@@ -0,0 +1,10 @@
+from .utils import BinnedData, bin_counts, bin_obs, build_spline_basis, compute_interior_knots, spline_basis_at
+
+__all__ = [
+    "BinnedData",
+    "bin_counts",
+    "bin_obs",
+    "build_spline_basis",
+    "compute_interior_knots",
+    "spline_basis_at",
+]