Add volcano plot helper with tests

src/rnaseq_tutorial/__init__.py

@@ -8,6 +8,7 @@
 from rnaseq_tutorial.data import load_counts, simulate_counts
 from rnaseq_tutorial.features import filter_low_counts, normalize_cpm
 from rnaseq_tutorial.model import differential_expression
+from rnaseq_tutorial.plots import volcano_data
 
 __version__ = "0.1.0"
 
@@ -17,4 +18,5 @@
     "filter_low_counts",
     "normalize_cpm",
     "differential_expression",
+    "volcano_data",
 ]

src/rnaseq_tutorial/plots.py

@@ -0,0 +1,71 @@
+"""Plotting helpers.
+
+Design note: the *computation* (what to plot) is separated from the *drawing*
+(matplotlib). ``volcano_data`` is pure and testable; ``volcano_plot`` only draws.
+This keeps the testable logic free of a hard matplotlib dependency.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import pandas as pd
+
+
+def volcano_data(
+    results: pd.DataFrame,
+    lfc_threshold: float = 1.0,
+    alpha: float = 0.05,
+) -> pd.DataFrame:
+    """Compute volcano-plot coordinates from a differential-expression table.
+
+    A volcano plot shows effect size (x = log2 fold change) against significance
+    (y = -log10 p-value), so the most interesting genes — large change *and* highly
+    significant — sit in the top corners.
+
+    Expects the columns produced by ``model.differential_expression``:
+    ``log2_fold_change``, ``p_value``, ``p_adj``.
+
+    Returns the input frame plus ``neg_log10_p`` and a ``category`` label of
+    "up", "down", or "ns" (not significant).
+    """
+    out = results.copy()
+    out["neg_log10_p"] = -np.log10(out["p_value"].clip(lower=1e-300))
+
+    sig = out["p_adj"] < alpha
+    up = sig & (out["log2_fold_change"] >= lfc_threshold)
+    down = sig & (out["log2_fold_change"] <= -lfc_threshold)
+
+    out["category"] = "ns"
+    out.loc[up, "category"] = "up"
+    out.loc[down, "category"] = "down"
+    return out
+
+
+def volcano_plot(results: pd.DataFrame, lfc_threshold: float = 1.0, alpha: float = 0.05):
+    """Draw a volcano plot and return the matplotlib Axes.
+
+    matplotlib is imported lazily so importing this module (and running the unit
+    tests) doesn't require a plotting backend to be installed.
+    """
+    import matplotlib.pyplot as plt
+
+    data = volcano_data(results, lfc_threshold=lfc_threshold, alpha=alpha)
+    colors = {"up": "#d62728", "down": "#1f77b4", "ns": "#999999"}
+
+    fig, ax = plt.subplots(figsize=(6, 5))
+    for category, group in data.groupby("category"):
+        ax.scatter(
+            group["log2_fold_change"],
+            group["neg_log10_p"],
+            s=10,
+            c=colors[category],
+            label=category,
+            alpha=0.7,
+        )
+    ax.axvline(lfc_threshold, ls="--", c="grey", lw=0.8)
+    ax.axvline(-lfc_threshold, ls="--", c="grey", lw=0.8)
+    ax.set_xlabel("log2 fold change")
+    ax.set_ylabel("-log10 p-value")
+    ax.set_title("Volcano plot")
+    ax.legend(title="regulation")
+    return ax

tests/test_plots.py

@@ -0,0 +1,34 @@
+"""Tests for the volcano-plot data preparation.
+
+We test the *pure* function (volcano_data), not the drawing, because the logic
+worth protecting is the categorization — not matplotlib's pixels.
+"""
+
+import pandas as pd
+
+from rnaseq_tutorial.plots import volcano_data
+
+
+def _results():
+    # Three hand-crafted genes: clearly up, clearly down, and not significant.
+    return pd.DataFrame(
+        {
+            "log2_fold_change": [3.0, -2.5, 0.1],
+            "p_value": [1e-8, 1e-7, 0.6],
+            "p_adj": [1e-6, 1e-5, 0.8],
+        },
+        index=["UP", "DOWN", "NS"],
+    )
+
+
+def test_volcano_data_categorizes_genes():
+    out = volcano_data(_results(), lfc_threshold=1.0, alpha=0.05)
+    assert out.loc["UP", "category"] == "up"
+    assert out.loc["DOWN", "category"] == "down"
+    assert out.loc["NS", "category"] == "ns"
+
+
+def test_volcano_data_adds_neg_log10_p():
+    out = volcano_data(_results())
+    # -log10(1e-8) == 8
+    assert abs(out.loc["UP", "neg_log10_p"] - 8.0) < 1e-9