Add distributional diagnostics to model comparison notebook

vitkl · claude · vitkl · commit 40ec399a6a0e · 2026-02-22T19:56:48.000Z
Add overlaid ECDF, overlaid density histograms, quantile table (p1-p99),
extreme-value fractions (genes with theta below thresholds), and
distribution spread metrics (IQR, 90/98% range) to reveal tail behavior
across models.

Co-Authored-By: Claude Opus 4.6 &lt;noreply@anthropic.com&gt;
diff --git a/docs/notebooks/model_comparison_theta.ipynb b/docs/notebooks/model_comparison_theta.ipynb
@@ -16,8 +16,8 @@
     "- **regularizedvi_nb** (tutorial): NB mode, prior on $\\sqrt{\\theta}$ (pushes $\\theta$ small)\n",
     "\n",
     "The Exponential(rate=3) prior has:\n",
-    "- GammaPoisson direction: $E[1/\\sqrt{\\theta}] = 1/3$ → $\\theta$ centered ~9\n",
-    "- NB direction: $E[\\sqrt{\\theta}] = 1/3$ → $\\theta$ centered ~1/9\n",
+    "- GammaPoisson direction: $E[1/\\sqrt{\\theta}] = 1/3$ \u2192 $\\theta$ centered ~9\n",
+    "- NB direction: $E[\\sqrt{\\theta}] = 1/3$ \u2192 $\\theta$ centered ~1/9\n",
     "\n",
     "Key question: do the unconstrained scVI models learn theta values that are very different from both 1/9 and 9?"
    ]
@@ -64,14 +64,14 @@
     "    (\"scvi_custom\", \"scvi_custom\", \"NB, 512/128, custom train, no prior\"),\n",
     "    (\"scvi_custom_default_train\", \"scvi_custom_default_train\", \"NB, 512/128, default train, no prior\"),\n",
     "    (\"scvi_custom_zinb\", \"scvi_custom_zinb\", \"ZINB, 512/128, custom train, no prior\"),\n",
-    "    (\"regularizedvi_gamma_poisson\", \"regularizedvi_gamma_poisson\", \"GammaPoisson, 512/128, prior→Poisson\"),\n",
-    "    (\"regularizedvi_nb\", \"regularizedvi_nb\", \"NB, 512/128, prior→overdispersed\"),\n",
-    "    (\"regularizedvi_gp_small\", \"regularizedvi_gamma_poisson_small\", \"GammaPoisson, 128/10, prior→Poisson\"),\n",
-    "    (\"regularizedvi_gp_medium\", \"regularizedvi_gamma_poisson_medium\", \"GammaPoisson, 256/30, prior→Poisson\"),\n",
+    "    (\"regularizedvi_gamma_poisson\", \"regularizedvi_gamma_poisson\", \"GammaPoisson, 512/128, prior\u2192Poisson\"),\n",
+    "    (\"regularizedvi_nb\", \"regularizedvi_nb\", \"NB, 512/128, prior\u2192overdispersed\"),\n",
+    "    (\"regularizedvi_gp_small\", \"regularizedvi_gamma_poisson_small\", \"GammaPoisson, 128/10, prior\u2192Poisson\"),\n",
+    "    (\"regularizedvi_gp_medium\", \"regularizedvi_gamma_poisson_medium\", \"GammaPoisson, 256/30, prior\u2192Poisson\"),\n",
     "    (\n",
     "        \"regularizedvi_gp_small_default\",\n",
     "        \"regularizedvi_gamma_poisson_small_default_train\",\n",
-    "        \"GammaPoisson, 128/10, default train, prior→Poisson\",\n",
+    "        \"GammaPoisson, 128/10, default train, prior\u2192Poisson\",\n",
     "    ),\n",
     "]\n",
     "\n",
@@ -112,7 +112,49 @@
    "id": "hist-theta",
    "metadata": {},
    "outputs": [],
-   "source": "fig, axes = plt.subplots(1, len(models), figsize=(4 * len(models), 4), sharey=True)\nif len(models) == 1:\n    axes = [axes]\n\nbins = np.linspace(-3, 3, 80)  # log10 scale: 0.001 to 1000\n\nfor ax, (name, data) in zip(axes, models.items(), strict=False):\n    log_theta = np.log10(data[\"theta\"])\n    ax.hist(log_theta, bins=bins, alpha=0.7, edgecolor=\"black\", linewidth=0.3)\n    ax.axvline(np.log10(1 / 9), color=\"red\", linestyle=\"--\", alpha=0.7, label=\"θ=1/9 (NB prior)\")\n    ax.axvline(np.log10(9), color=\"blue\", linestyle=\"--\", alpha=0.7, label=\"θ=9 (GP prior)\")\n    ax.axvline(\n        np.log10(np.median(data[\"theta\"])),\n        color=\"green\",\n        linestyle=\"-\",\n        alpha=0.9,\n        label=f\"median={np.median(data['theta']):.1f}\",\n    )\n    ax.set_xlabel(\"log₁₀(θ)\")\n    ax.set_title(name, fontsize=9)\n    ax.legend(fontsize=7)\n\naxes[0].set_ylabel(\"Gene count\")\nfig.suptitle(\"Per-gene inverse dispersion θ = exp(px_r)\", fontsize=12)\nplt.tight_layout()\nplt.show()"
+   "source": "fig, axes = plt.subplots(1, len(models), figsize=(4 * len(models), 4), sharey=True)\nif len(models) == 1:\n    axes = [axes]\n\nbins = np.linspace(-3, 3, 80)  # log10 scale: 0.001 to 1000\n\nfor ax, (name, data) in zip(axes, models.items(), strict=False):\n    log_theta = np.log10(data[\"theta\"])\n    ax.hist(log_theta, bins=bins, alpha=0.7, edgecolor=\"black\", linewidth=0.3)\n    ax.axvline(np.log10(1 / 9), color=\"red\", linestyle=\"--\", alpha=0.7, label=\"\u03b8=1/9 (NB prior)\")\n    ax.axvline(np.log10(9), color=\"blue\", linestyle=\"--\", alpha=0.7, label=\"\u03b8=9 (GP prior)\")\n    ax.axvline(\n        np.log10(np.median(data[\"theta\"])),\n        color=\"green\",\n        linestyle=\"-\",\n        alpha=0.9,\n        label=f\"median={np.median(data['theta']):.1f}\",\n    )\n    ax.set_xlabel(\"log\u2081\u2080(\u03b8)\")\n    ax.set_title(name, fontsize=9)\n    ax.legend(fontsize=7)\n\naxes[0].set_ylabel(\"Gene count\")\nfig.suptitle(\"Per-gene inverse dispersion \u03b8 = exp(px_r)\", fontsize=12)\nplt.tight_layout()\nplt.show()"
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3d6kr3mrx8p",
+   "source": "## 2b. Overlaid ECDF \u2014 all models on same axes\n\nThe empirical CDF of $\\log_{10}(\\theta)$ reveals the full shape of each distribution and makes tail differences immediately visible.",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "id": "bbavnahmdh7",
+   "source": "fig, ax = plt.subplots(figsize=(10, 6))\n\nfor name, data in models.items():\n    log_theta = np.sort(np.log10(data[\"theta\"]))\n    ecdf = np.arange(1, len(log_theta) + 1) / len(log_theta)\n    ax.plot(log_theta, ecdf, label=f\"{name} ({data['desc']})\", linewidth=1.5)\n\nax.axvline(np.log10(1 / 9), color=\"red\", linestyle=\"--\", alpha=0.5, label=\"\u03b8=1/9\")\nax.axvline(np.log10(9), color=\"blue\", linestyle=\"--\", alpha=0.5, label=\"\u03b8=9\")\nax.set_xlabel(\"log\u2081\u2080(\u03b8)\")\nax.set_ylabel(\"Cumulative fraction of genes\")\nax.set_title(\"ECDF of per-gene \u03b8 across all models\")\nax.legend(fontsize=7, loc=\"upper left\")\nax.set_xlim(-3, 3)\nax.grid(True, alpha=0.3)\nplt.tight_layout()\nplt.show()",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "68q53gvczvm",
+   "source": "## 2c. Quantile table and extreme-value fractions\n\nPercentiles reveal how spread the distribution is. Extreme-value fractions show what fraction of genes have very low \u03b8 (extreme overdispersion) or very high \u03b8 (near Poisson).",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "id": "f8yks65q64w",
+   "source": "quantiles = [1, 5, 25, 50, 75, 95, 99]\nheader = f\"{'Model':<35}\" + \"\".join(f\"{'p' + str(q):>8}\" for q in quantiles)\nprint(\"=== Theta quantiles ===\")\nprint(header)\nprint(\"-\" * len(header))\nfor name, data in models.items():\n    vals = \"\".join(f\"{np.percentile(data['theta'], q):>8.2f}\" for q in quantiles)\n    print(f\"{name:<35}{vals}\")\n\nprint()\n\n# Extreme-value fractions\nthresholds = [0.01, 0.1, 1, 10, 100, 1000]\nheader2 = f\"{'Model':<35}\" + \"\".join(f\"{'<' + str(t):>8}\" for t in thresholds)\nprint(\"=== Fraction of genes with \u03b8 below threshold ===\")\nprint(header2)\nprint(\"-\" * len(header2))\nfor name, data in models.items():\n    theta = data[\"theta\"]\n    fracs = \"\".join(f\"{(theta < t).mean():>8.3f}\" for t in thresholds)\n    print(f\"{name:<35}{fracs}\")\n\nprint()\n\n# IQR and range as spread measure\nprint(\"=== Distribution spread (log10 scale) ===\")\nprint(f\"{'Model':<35}{'IQR':>8}{'90% range':>12}{'98% range':>12}{'full range':>12}\")\nprint(\"-\" * 79)\nfor name, data in models.items():\n    lt = np.log10(data[\"theta\"])\n    iqr = np.percentile(lt, 75) - np.percentile(lt, 25)\n    r90 = np.percentile(lt, 95) - np.percentile(lt, 5)\n    r98 = np.percentile(lt, 99) - np.percentile(lt, 1)\n    full = lt.max() - lt.min()\n    print(f\"{name:<35}{iqr:>8.2f}{r90:>12.2f}{r98:>12.2f}{full:>12.2f}\")",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "vzxfo0slei",
+   "source": "## 2d. Overlaid histograms (density)\n\nSame data as Section 2 but overlaid on a single plot for direct visual comparison of distribution shapes.",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "id": "e2rwc55svjq",
+   "source": "fig, ax = plt.subplots(figsize=(10, 5))\nbins = np.linspace(-3, 3, 80)\n\nfor name, data in models.items():\n    log_theta = np.log10(data[\"theta\"])\n    ax.hist(log_theta, bins=bins, alpha=0.3, density=True, label=name)\n\nax.axvline(np.log10(1 / 9), color=\"red\", linestyle=\"--\", alpha=0.7, label=\"\u03b8=1/9\")\nax.axvline(np.log10(9), color=\"blue\", linestyle=\"--\", alpha=0.7, label=\"\u03b8=9\")\nax.set_xlabel(\"log\u2081\u2080(\u03b8)\")\nax.set_ylabel(\"Density\")\nax.set_title(\"Overlaid density of per-gene \u03b8\")\nax.legend(fontsize=7)\nplt.tight_layout()\nplt.show()",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
   },
   {
    "cell_type": "markdown",
@@ -130,7 +172,7 @@
    "id": "scatter-theta",
    "metadata": {},
    "outputs": [],
-   "source": "# Compare each model against scvi_default as baseline (if available)\nbaseline = \"scvi_default\"\nif baseline in models:\n    others = [k for k in models if k != baseline]\n    n = len(others)\n    fig, axes = plt.subplots(1, n, figsize=(4 * n, 4))\n    if n == 1:\n        axes = [axes]\n\n    for ax, other_name in zip(axes, others, strict=False):\n        x = np.log10(models[baseline][\"theta\"])\n        y = np.log10(models[other_name][\"theta\"])\n        # Ensure same number of genes\n        n_genes = min(len(x), len(y))\n        ax.hist2d(x[:n_genes], y[:n_genes], bins=80, cmap=\"viridis\", range=[[-3, 3], [-3, 3]])\n        ax.plot([-3, 3], [-3, 3], \"r--\", alpha=0.5)\n        ax.set_xlabel(f\"log₁₀(θ) — {baseline}\")\n        ax.set_ylabel(f\"log₁₀(θ) — {other_name}\")\n        ax.set_title(other_name, fontsize=9)\n\n    fig.suptitle(f\"Per-gene θ: each model vs {baseline}\", fontsize=12)\n    plt.tight_layout()\n    plt.show()\nelse:\n    print(f\"Baseline model '{baseline}' not found\")"
+   "source": "# Compare each model against scvi_default as baseline (if available)\nbaseline = \"scvi_default\"\nif baseline in models:\n    others = [k for k in models if k != baseline]\n    n = len(others)\n    fig, axes = plt.subplots(1, n, figsize=(4 * n, 4))\n    if n == 1:\n        axes = [axes]\n\n    for ax, other_name in zip(axes, others, strict=False):\n        x = np.log10(models[baseline][\"theta\"])\n        y = np.log10(models[other_name][\"theta\"])\n        # Ensure same number of genes\n        n_genes = min(len(x), len(y))\n        ax.hist2d(x[:n_genes], y[:n_genes], bins=80, cmap=\"viridis\", range=[[-3, 3], [-3, 3]])\n        ax.plot([-3, 3], [-3, 3], \"r--\", alpha=0.5)\n        ax.set_xlabel(f\"log\u2081\u2080(\u03b8) \u2014 {baseline}\")\n        ax.set_ylabel(f\"log\u2081\u2080(\u03b8) \u2014 {other_name}\")\n        ax.set_title(other_name, fontsize=9)\n\n    fig.suptitle(f\"Per-gene \u03b8: each model vs {baseline}\", fontsize=12)\n    plt.tight_layout()\n    plt.show()\nelse:\n    print(f\"Baseline model '{baseline}' not found\")"
   },
   {
    "cell_type": "markdown",
@@ -151,8 +193,8 @@
    "source": [
     "mu_typical = 5.0  # typical gene expression\n",
     "\n",
-    "print(f\"Variance = μ + μ²/θ = {mu_typical} + {mu_typical**2}/θ\")\n",
-    "print(f\"{'Model':<35} {'Median θ':>10} {'Var(x)':>10} {'Var/μ':>10} {'Poisson Var':>12}\")\n",
+    "print(f\"Variance = \u03bc + \u03bc\u00b2/\u03b8 = {mu_typical} + {mu_typical**2}/\u03b8\")\n",
+    "print(f\"{'Model':<35} {'Median \u03b8':>10} {'Var(x)':>10} {'Var/\u03bc':>10} {'Poisson Var':>12}\")\n",
     "print(\"-\" * 80)\n",
     "for name, data in models.items():\n",
     "    med_theta = np.median(data[\"theta\"])\n",
@@ -185,9 +227,9 @@
     "        print(f\"  CV across batches: median={np.median(cv):.3f}, mean={cv.mean():.3f}, max={cv.max():.3f}\")\n",
     "\n",
     "        plt.hist(cv, bins=50, alpha=0.7, edgecolor=\"black\", linewidth=0.3)\n",
-    "        plt.xlabel(\"CV of θ across batches\")\n",
+    "        plt.xlabel(\"CV of \u03b8 across batches\")\n",
     "        plt.ylabel(\"Gene count\")\n",
-    "        plt.title(f\"{name}: per-gene variation of θ across {theta_full.shape[1]} batches\")\n",
+    "        plt.title(f\"{name}: per-gene variation of \u03b8 across {theta_full.shape[1]} batches\")\n",
     "        plt.show()"
    ]
   },
@@ -198,14 +240,14 @@
    "source": [
     "## Summary\n",
     "\n",
-    "| Model | Prior | Median θ | Interpretation |\n",
+    "| Model | Prior | Median \u03b8 | Interpretation |\n",
     "|-------|-------|----------|----------------|\n",
     "| scvi_default | None | ~1.3 | Moderate overdispersion |\n",
     "| scvi_custom | None | ~2.4 | Moderate overdispersion |\n",
-    "| regularizedvi_nb | Exp on √θ (→small) | ? | θ→1/9 expected |\n",
-    "| regularizedvi_gamma_poisson | Exp on 1/√θ (→large) | ~53 | Near Poisson |\n",
+    "| regularizedvi_nb | Exp on \u221a\u03b8 (\u2192small) | ? | \u03b8\u21921/9 expected |\n",
+    "| regularizedvi_gamma_poisson | Exp on 1/\u221a\u03b8 (\u2192large) | ~53 | Near Poisson |\n",
     "\n",
-    "Both prior directions constrain θ via √ transform, but the unconstrained models land at θ~1-2, which is far from both 1/9 and 9. This suggests the prior's main effect is constraining θ to a definite range rather than the specific direction mattering."
+    "Both prior directions constrain \u03b8 via \u221a transform, but the unconstrained models land at \u03b8~1-2, which is far from both 1/9 and 9. This suggests the prior's main effect is constraining \u03b8 to a definite range rather than the specific direction mattering."
    ]
   }
  ],
