diff --git a/docs/contract/REFACTOR_AGENDA.md b/docs/contract/REFACTOR_AGENDA.md index a59897e..f891cd5 100644 --- a/docs/contract/REFACTOR_AGENDA.md +++ b/docs/contract/REFACTOR_AGENDA.md @@ -60,19 +60,19 @@ Tick only when the symbol is gone from source AND `__all__`/imports AND `import - [x] `tcri_clone_key` / `tcri_phenotype_key` / obsm `X_tcri_phenotypes` — **DONE**. The deferral said "still read by not-yet-refactored `metrics`/`plotting`", but those were rewritten in PR6/PR7; the last reader was a raw literal in `pp.clone_size`, now migrated to `uns[METADATA][CLONE_COL]` with a clear error. `to_anndata` no longer writes the shims and the three `LEGACY_*` constants are deleted (`LEGACY_MANAGER` stays — `save_tcri_session` still pops it defensively). **Phase 5/6 (consolidated away — delete WITH replacement, never before):** -- [ ] `pp.joint_distribution_posterior` (→ unified `joint_distribution`) · [ ] `metrics._mi_from_joint` (→ `_mutual_information`) -- [ ] `tl.mi_compare` (→ `compare_groups`) · [ ] `tl.delta_clonotypic_entropy` · [ ] `tl.delta_entropy_table` · [ ] `tl.flux_table` -- [ ] `tl.clonotypic_entropy_base` · [ ] `tl.clonality` · [ ] `tl.dkl` local `dkl_func` -- [ ] plural `*_entropies` shims · [ ] `metrics/` package (after migration to `tools/`) +- [x] `pp.joint_distribution_posterior` · [x] `metrics._mi_from_joint` (verified: no `def` remains in `tcri/`) +- [x] `tl.mi_compare` · [x] `tl.delta_clonotypic_entropy` · [x] `tl.delta_entropy_table` · [x] `tl.flux_table` +- [x] `tl.clonotypic_entropy_base` · [x] `tl.clonality` · [x] `tl.dkl` local `dkl_func` +- [x] plural `*_entropies` shims · [x] `metrics/` package — the `.py` files went in PR6 but an **empty dir with a `.DS_Store` survived**; removed now **Phase 7 (non-core plots — DROP, not to examples):** -- [ ] `pl.probability_ternary` · [ ] `pl.top_clone_umap` · [ ] `pl.clone_size_umap` · [ ] `pl.plot_phenotype_probabilities` -- [ ] `pl.compare_phenotypes` · [ ] `pl.ridge_delta_entropy` · [ ] `pl.flux` boxplot · [ ] `pl.clonality` plot -- [ ] `pl.tcri_boxplot` (→ private `_metric_boxplot`) · [ ] `pl.set_color_palette` (→ `resolve_palette`) -- [ ] `pl.plot_pheno_sankey` (→ private `_sankey`) · [ ] leaked aliases `centropy`/`pentropy`/`*_tl` +- [x] `pl.probability_ternary` · [x] `pl.top_clone_umap` · [x] `pl.clone_size_umap` · [x] `pl.plot_phenotype_probabilities` +- [x] `pl.compare_phenotypes` · [x] `pl.ridge_delta_entropy` · [x] `pl.flux` boxplot · [x] `pl.clonality` plot +- [x] `pl.tcri_boxplot` (→ `_metric_boxplot`) · [x] `pl.set_color_palette` (→ `resolve_palette`) +- [x] `pl.plot_pheno_sankey` (→ `_sankey`) · [x] leaked aliases `centropy`/`pentropy`/`*_tl` **Phase 9 (out of the package):** -- [ ] `ut.build_nested_tcri_pgm` (→ `docs/`) · [ ] `ut.draw_tcri_pgm_nested` (→ `docs/`) · [ ] `daft` runtime dep +- [x] `ut.build_nested_tcri_pgm` (→ `docs/`) · [x] `ut.draw_tcri_pgm_nested` (→ `docs/`) · [x] `daft` runtime dep **Phase 3/9 (model/utils cleanup):** - [x] `_ascii_hist` (dead: zero callers) · [x] `ml.plot_loss` (→ `diag.loss`) · [x] `ml.plot_archetypes` (→ `diag.archetypes`) diff --git a/pyproject.toml b/pyproject.toml index 54c60dd..c915782 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,8 +43,6 @@ dependencies = [ "umap-learn>=0.5.6", "matplotlib>=3.9.2", "seaborn>=0.13.2", - "mpltern>=1.0.4", - "gseapy>=1.1.4", "tqdm>=4.66.5", ] diff --git a/requirements.txt b/requirements.txt index 154b8ad..9ed9954 100644 --- a/requirements.txt +++ b/requirements.txt @@ -15,8 +15,6 @@ scvi-tools==1.5.0.post1 umap-learn==0.5.12 matplotlib==3.10.9 seaborn==0.13.2 -mpltern==1.0.5 -gseapy==1.3.0 daft==0.6.14 tqdm==4.68.4 pytest==9.1.1 diff --git a/tcri/__init__.py b/tcri/__init__.py index 1131424..d673ee8 100644 --- a/tcri/__init__.py +++ b/tcri/__init__.py @@ -11,6 +11,7 @@ from . import utils as ut from . import model as ml from . import diagnostics as diag # PR8 +from . import datasets # synthetic data with a known-MI oracle # The unified engine, re-exported top-level for prominence. from .tools import joint_distribution diff --git a/tcri/datasets/__init__.py b/tcri/datasets/__init__.py new file mode 100644 index 0000000..9f1fd2f --- /dev/null +++ b/tcri/datasets/__init__.py @@ -0,0 +1,8 @@ +"""Synthetic datasets with known ground truth (``tcri.datasets``). + +:func:`simulate_tcri` generates a TCR+RNA dataset whose **mutual information is +known in closed form**, which is what makes statistical recovery testing possible. +""" +from ._simulate import mi_from_joint_oracle, simulate_tcri + +__all__ = ["simulate_tcri", "mi_from_joint_oracle"] diff --git a/tcri/datasets/_simulate.py b/tcri/datasets/_simulate.py new file mode 100644 index 0000000..af879c7 --- /dev/null +++ b/tcri/datasets/_simulate.py @@ -0,0 +1,232 @@ +"""Self-contained synthetic generator with an **exact mutual-information oracle**. + +Implements the semi-synthetic generative story from Supplementary Note 1 +("Generative Model for Semi-Synthetic Simulations"):: + + pi ~ Dirichlet(alpha_pi) # clone abundance P(c) + omega_c ~ Dirichlet(alpha_omega) # P(phi | c) + z_i ~ Categorical(pi) # clone of cell i + phi_i|z_i ~ Categorical(omega[z_i]) # phenotype of cell i + x_i|phi_i ~ Poisson(U_i @ V), U_i ~ Gamma(program of phi_i) + +Because ``pi`` and ``omega`` are known, the **population** mutual information +I(c;phi) is available in closed form — which is what makes statistical *recovery* +testing possible at all. Nothing else in the test suite has an oracle: the contract +tests check structure and identities, never accuracy. + +Two oracles are reported, and the distinction matters for writing honest tests: + +``true_mi_*`` + the **population** value implied by ``(pi, omega)`` — what an estimator should + approach as ``n_cells -> inf``. +``empirical_mi_*`` + the value implied by the **realized** ``(clone, phenotype)`` counts in this + particular sample. A perfect estimator applied to *this* dataset returns this, + not the population value. They differ by finite-sample noise plus the plug-in + estimator's upward bias, roughly ``(C-1)(P-1) / (2 N ln2)`` bits. + +Both are given under both normalizations (``min`` and ``average``): tcri defaults to +``normalize_mode="min"`` while the note's benchmark used the mean denominator, so a +like-for-like comparison has to pick deliberately. + +Unlike the original ``sc_simulator``, this needs no real dataset to fit — the gene +programs are generated directly — so it is importable, seeded, and fast. +""" +from __future__ import annotations + +import numpy as np +import pandas as pd +from anndata import AnnData + +__all__ = ["simulate_tcri", "mi_from_joint_oracle"] + + +def mi_from_joint_oracle(joint: np.ndarray) -> dict: + """Exact MI / entropies (bits) of a clone x phenotype **probability** table. + + ``joint`` must sum to 1. Returns ``mi``, ``h_clone``, ``h_phenotype`` and both + normalized variants. This is the oracle — deliberately a small, independent + implementation so it cannot drift with the package's own metric code (a test + that computes the expected value with the code under test proves nothing). + """ + P = np.asarray(joint, dtype=np.float64) + total = P.sum() + if total <= 0: + raise ValueError("joint has no mass") + P = P / total + p_c = P.sum(1) + p_ph = P.sum(0) + + nz = P > 0 + outer = np.outer(p_c, p_ph) + mi = float(np.sum(P[nz] * np.log2(P[nz] / outer[nz]))) + + h_c = float(-np.sum(p_c[p_c > 0] * np.log2(p_c[p_c > 0]))) + h_p = float(-np.sum(p_ph[p_ph > 0] * np.log2(p_ph[p_ph > 0]))) + + denom_min = min(h_c, h_p) + denom_avg = 0.5 * (h_c + h_p) + return { + "mi": mi, + "h_clone": h_c, + "h_phenotype": h_p, + "nmi_min": mi / denom_min if denom_min > 0 else 0.0, + "nmi_average": mi / denom_avg if denom_avg > 0 else 0.0, + } + + +def _phenotype_programs(rng, n_phenotypes, n_factors, fuzziness): + """Per-phenotype Gamma params over latent factors, blended by ``fuzziness``. + + ``fuzziness=0`` leaves the phenotypes fully distinct; ``fuzziness=1`` collapses + them to a common program so phenotype is **unidentifiable from expression** while + the clone->phenotype coupling (and hence the true MI) is untouched. That + separation is the point: it varies estimation difficulty at fixed ground truth. + + Interpolation is in natural-parameter space ``theta = [alpha-1, -beta]``, matching + the original ``interpolate_gamma_params``. + """ + alpha = rng.uniform(1.5, 6.0, size=(n_phenotypes, n_factors)) + beta = rng.uniform(1.0, 3.0, size=(n_phenotypes, n_factors)) + + theta = np.concatenate([alpha - 1.0, -beta], axis=1) + theta = (1.0 - fuzziness) * theta + fuzziness * theta.mean(0, keepdims=True) + + alpha_f = theta[:, :n_factors] + 1.0 + beta_f = -theta[:, n_factors:] + # keep the Gamma valid after blending + return np.clip(alpha_f, 1e-3, None), np.clip(beta_f, 1e-3, None) + + +def simulate_tcri( + *, + n_clones: int = 30, + n_phenotypes: int = 4, + n_genes: int = 60, + n_cells: int = 2000, + n_covariates: int = 1, + n_factors: int = 8, + omega_concentration: float = 0.35, + pi_concentration: float = 2.0, + fuzziness: float = 0.0, + label_error_rate: float = 0.0, + seed: int = 0, +) -> AnnData: + """Simulate a TCR+RNA dataset whose true mutual information is known exactly. + + Parameters + ---------- + n_clones, n_phenotypes, n_genes, n_cells + Problem size. ``n_cells`` drives how close the realized sample sits to the + population oracle. + n_covariates + Cells are split across covariates. ``omega`` is shared, so each covariate's + population MI equals the global one — which lets a per-covariate estimate be + compared against the same oracle. + omega_concentration + Dirichlet concentration for ``P(phi|c)`` and **the knob that sets the true + MI**: small (<1) gives near-one-hot rows and MI near ``H(phi)``; large gives + near-uniform rows and MI near 0. + pi_concentration + Dirichlet concentration for clone abundance. + fuzziness + In ``[0,1]``. Blends the per-phenotype expression programs, making phenotype + harder to read off expression **without changing the true MI**. + label_error_rate + Fraction of cells whose recorded phenotype is replaced by a uniform draw. + This DOES change the realized coupling, so ``empirical_*`` is recomputed from + the corrupted labels while ``true_*`` keeps the uncorrupted population value. + seed + Seeds everything. + + Returns + ------- + AnnData + ``X`` / ``layers['counts']`` integer counts; ``obs`` with ``clone_id``, + ``phenotype`` (possibly corrupted), ``true_phenotype``, ``covariate``, + ``batch``; ``uns['tcri_truth']`` holding ``omega``, ``pi``, the population + oracle (``true_mi``, ``true_nmi_min``, ``true_nmi_average``, entropies) and + the realized-sample oracle (``empirical_*``), plus the generating settings. + """ + if not 0.0 <= fuzziness <= 1.0: + raise ValueError("fuzziness must be in [0, 1]") + if not 0.0 <= label_error_rate <= 1.0: + raise ValueError("label_error_rate must be in [0, 1]") + if n_covariates < 1: + raise ValueError("n_covariates must be >= 1") + + rng = np.random.default_rng(seed) + + # ── ground truth: clone abundance and the clone -> phenotype coupling ──── + pi = rng.dirichlet(np.full(n_clones, pi_concentration)) + omega = rng.dirichlet(np.full(n_phenotypes, omega_concentration), size=n_clones) + + # ── population oracle (closed form, independent of any sampling) ───────── + truth = mi_from_joint_oracle(pi[:, None] * omega) + + # ── sample cells ──────────────────────────────────────────────────────── + z = rng.choice(n_clones, size=n_cells, p=pi) + phi_true = np.array([rng.choice(n_phenotypes, p=omega[c]) for c in z]) + + phi = phi_true.copy() + if label_error_rate > 0: + flip = rng.random(n_cells) < label_error_rate + phi[flip] = rng.integers(0, n_phenotypes, size=int(flip.sum())) + + # ── expression: x_i ~ Poisson(U_i @ V), U_i ~ Gamma(program of phi_i) ──── + alpha, beta = _phenotype_programs(rng, n_phenotypes, n_factors, fuzziness) + V = rng.gamma(2.0, 1.0, size=(n_factors, n_genes)) + U = rng.gamma(alpha[phi_true], 1.0 / beta[phi_true]) # (n_cells, n_factors) + X = rng.poisson(U @ V).astype("float32") + + # ── realized-sample oracle, from the labels actually recorded ─────────── + counts = np.zeros((n_clones, n_phenotypes), dtype=np.float64) + np.add.at(counts, (z, phi), 1.0) + empirical = mi_from_joint_oracle(counts) if counts.sum() > 0 else dict.fromkeys(truth, np.nan) + + obs = pd.DataFrame( + { + "clone_id": pd.Categorical([f"clone_{i}" for i in z]), + "phenotype": pd.Categorical([f"phen_{p}" for p in phi]), + "true_phenotype": pd.Categorical([f"phen_{p}" for p in phi_true]), + "covariate": pd.Categorical( + [f"cov_{i}" for i in rng.integers(0, n_covariates, size=n_cells)] + ), + "batch": pd.Categorical(["batch_0"] * n_cells), + }, + index=[f"cell_{i}" for i in range(n_cells)], + ) + + adata = AnnData( + X=X, + obs=obs, + var=pd.DataFrame(index=[f"gene_{g}" for g in range(n_genes)]), + ) + adata.layers["counts"] = adata.X.copy() + + adata.uns["tcri_truth"] = { + "omega": omega, + "pi": pi, + "true_mi": truth["mi"], + "true_nmi_min": truth["nmi_min"], + "true_nmi_average": truth["nmi_average"], + "true_h_clone": truth["h_clone"], + "true_h_phenotype": truth["h_phenotype"], + "empirical_mi": empirical["mi"], + "empirical_nmi_min": empirical["nmi_min"], + "empirical_nmi_average": empirical["nmi_average"], + "settings": { + "n_clones": n_clones, + "n_phenotypes": n_phenotypes, + "n_genes": n_genes, + "n_cells": n_cells, + "n_covariates": n_covariates, + "n_factors": n_factors, + "omega_concentration": omega_concentration, + "pi_concentration": pi_concentration, + "fuzziness": fuzziness, + "label_error_rate": label_error_rate, + "seed": seed, + }, + } + return adata diff --git a/tcri/preprocessing/_preprocessing.py b/tcri/preprocessing/_preprocessing.py index d48c8a5..e0c1430 100644 --- a/tcri/preprocessing/_preprocessing.py +++ b/tcri/preprocessing/_preprocessing.py @@ -1,53 +1,15 @@ -from scipy.stats import entropy -from .. import _keys as K -import numpy as np -import tqdm -import pandas as pd -import collections -import warnings -import torch -import torch.nn.functional as F -import datetime -import pyro.distributions as dist -from pyro.distributions import Dirichlet -import pyro -import scanpy as sc -import numpy as np -import pandas as pd -import torch -import torch.nn.functional as F -from typing import Optional -import umap -import numpy as np, pandas as pd, torch, umap -from tqdm.auto import tqdm -from scvi import REGISTRY_KEYS +"""Preprocessing helpers: clonotype grouping and clone sizes. +Deliberately light on imports — this module is loaded by ``import tcri``, so an +eager ``import umap`` here cost ~2.9 s of every import (umap → pynndescent → +numba/llvmlite) for a dependency this file never used. +""" import numpy as np -import pandas as pd -import torch -import torch.nn.functional as F -from torch.distributions import Dirichlet -import numpy as np -import pandas as pd -from scipy.special import softmax -from torch.distributions import Dirichlet -import torch -import warnings - -warnings.filterwarnings('ignore') - -# ------------ simple ANSI helpers ------------ # -RESET = "\x1b[0m" -BOLD = "\x1b[1m" -DIM = "\x1b[2m" -GREEN = "\x1b[32m" -CYAN = "\x1b[36m" -MAGENT = "\x1b[35m" -# ╭─ colour / pretty-print helpers ─────────────────────────────────────────╮ -RESET = "\x1b[0m"; BOLD = "\x1b[1m"; DIM = "\x1b[2m" -GRN = "\x1b[32m"; CYN = "\x1b[36m"; MAG = "\x1b[35m"; YLW = "\x1b[33m"; RED = "\x1b[31m" - -from .._console import _ok, _info, _warn, _fin + +from .. import _keys as K + +__all__ = ["group_singletons", "clone_size"] + def group_singletons(adata,clonotype_key="trb",groupby="patient", target_col="trb_unique", min_clone_size=10): adata.obs["trb_candidate"] = adata.obs[clonotype_key].astype(str) + "_" + adata.obs[groupby].astype(str) diff --git a/tcri/utils/_utils.py b/tcri/utils/_utils.py index 7ef3401..f020c3e 100644 --- a/tcri/utils/_utils.py +++ b/tcri/utils/_utils.py @@ -8,7 +8,6 @@ import matplotlib as mpl import matplotlib.patches as mpatches from matplotlib.collections import LineCollection -import mpltern import numpy as np from scipy.stats import fisher_exact#, binom_test from matplotlib.colors import LinearSegmentedColormap diff --git a/tests/conftest.py b/tests/conftest.py index bed7668..048d826 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -10,10 +10,37 @@ from scipy import sparse +def pytest_addoption(parser): + parser.addoption( + "--runslow", action="store_true", default=False, + help="run slow statistical-recovery tests (model fits over many configs)", + ) + + def pytest_configure(config): """Deterministic RNG for all tests (Notion T13).""" np.random.seed(42) torch.manual_seed(42) + config.addinivalue_line( + "markers", + "slow: statistical-recovery test that fits models over several configs; " + "skipped unless --runslow is passed", + ) + + +def pytest_collection_modifyitems(config, items): + """Skip ``@pytest.mark.slow`` unless ``--runslow``. + + Recovery tests fit real models across a grid, so they are minutes-scale and do + not belong in the per-commit suite — but they are the only tests with an + accuracy oracle, so they must stay runnable (nightly / pre-release). + """ + if config.getoption("--runslow"): + return + skip_slow = pytest.mark.skip(reason="needs --runslow") + for item in items: + if "slow" in item.keywords: + item.add_marker(skip_slow) def _seed_all(seed: int = 0) -> None: diff --git a/tests/test_helpers.py b/tests/test_helpers.py index 20f716e..a6aa7ee 100644 --- a/tests/test_helpers.py +++ b/tests/test_helpers.py @@ -19,6 +19,47 @@ def test_keys_constants(): assert K.LEGACY_MANAGER == "tcri_manager" +def test_import_tcri_does_not_hijack_global_warning_filters(): + """A library must not silence the application's warnings. + + ``_preprocessing`` used to call a blanket ``warnings.filterwarnings('ignore')`` + at module scope, which runs on ``import tcri`` and silenced EVERY warning in the + user's session — including this package's own guardrails (the K clamp, the + param-store-reuse notice, the batch_size warning). Narrow message-specific + filters are tolerated; a catch-all is not. + """ + import subprocess + import sys + + code = ( + "import warnings, io, contextlib; import tcri; " + "buf = io.StringIO()\n" + "with contextlib.redirect_stderr(buf): warnings.warn('probe', UserWarning)\n" + "print('VISIBLE' if 'probe' in buf.getvalue() else 'SILENCED')" + ) + out = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True) + assert "VISIBLE" in out.stdout, ( + "importing tcri silenced a user warning — something re-added a blanket " + f"filterwarnings('ignore'). stdout={out.stdout!r}" + ) + + +def test_preprocessing_import_is_light(): + """``import tcri`` must not drag in umap (~2.9 s via pynndescent/numba). + + umap is only needed for the opt-in ``to_anndata(compute_umap=True)`` path, which + imports it locally. An eager module-level import cost every user ~2.9 s. + """ + import subprocess + import sys + + out = subprocess.run( + [sys.executable, "-c", "import tcri, sys; print('umap' in sys.modules)"], + capture_output=True, text=True, + ) + assert "False" in out.stdout, f"import tcri eagerly loaded umap: {out.stdout!r}" + + def test_console_aliases_and_callables(): assert _console.MAG == _console.MAGENT assert _console.GRN == _console.GREEN diff --git a/tests/test_recovery.py b/tests/test_recovery.py new file mode 100644 index 0000000..b456456 --- /dev/null +++ b/tests/test_recovery.py @@ -0,0 +1,294 @@ +"""Statistical **recovery** tests — the only tests in this suite with an accuracy oracle. + +Everything else here checks *structure* (contracts, identities, wiring). Those cannot +catch an estimator that is well-formed but wrong. These use +:func:`tcri.datasets.simulate_tcri`, whose mutual information is known in closed form, +and ask: **does the number come out right?** + +Tiering (see ``conftest.py``): + +* unmarked — oracle self-consistency, exactness of the metric against an independent + MI implementation, and metamorphic invariances. Model-free, so fast; runs every commit. +* ``@pytest.mark.slow`` — needs model fits and/or replication over seeds. Skipped unless + ``--runslow``. Run nightly / before a release. + +Two honest caveats baked into the assertions below: + +1. **Finite-sample bias.** The plug-in MI estimator is biased *upward* by roughly + ``(C-1)(P-1)/(2N ln2)`` bits, so a realized sample's MI exceeds the population value. + Tests compare against the *realized* oracle where that matters, and use + bias-aware tolerances where they compare against the population value. +2. **Single-seed monotonicity is flaky.** Sampling noise can make a larger N land + *further* from the truth on one seed (observed: gap +0.003 at N=5000 vs +0.012 at + N=20000 on the same seed). Convergence is therefore asserted on a **mean over + seeds**, never on one draw. +""" +from __future__ import annotations + +import contextlib +import io + +import numpy as np +import pandas as pd +import pyro +import pytest + +import tcri +from tcri.datasets import mi_from_joint_oracle, simulate_tcri + + +# ── helpers ───────────────────────────────────────────────────────────────── +def _empirical_joint(adata, phenotype_col="phenotype"): + """Realized clone x phenotype count table as a DataFrame (a precomputed joint).""" + return pd.crosstab(adata.obs["clone_id"], adata.obs[phenotype_col]) + + +def _truth(adata): + return adata.uns["tcri_truth"] + + +# ══════════════════════ TIER A — the oracle itself ══════════════════════════ +def test_oracle_respects_information_bounds(): + """MI <= min(H(c), H(phi)), and both normalizations land in [0, 1].""" + for conc in (0.05, 0.5, 5.0): + t = _truth(simulate_tcri(omega_concentration=conc, n_cells=500, seed=0)) + assert t["true_mi"] >= -1e-12 + assert t["true_mi"] <= min(t["true_h_clone"], t["true_h_phenotype"]) + 1e-9 + assert 0.0 - 1e-12 <= t["true_nmi_min"] <= 1.0 + 1e-9 + assert 0.0 - 1e-12 <= t["true_nmi_average"] <= 1.0 + 1e-9 + + +def test_omega_concentration_controls_the_true_mi(): + """The difficulty knob must actually move the ground truth, monotonically. + + Small Dirichlet concentration => near-one-hot P(phi|c) => MI approaches H(phi). + Large => near-uniform rows => MI approaches 0. Averaged over seeds so the + ordering reflects the parameter, not one draw. + """ + means = [] + for conc in (0.05, 0.5, 5.0, 50.0): + vals = [ + _truth(simulate_tcri(omega_concentration=conc, n_cells=300, seed=s))["true_mi"] + for s in range(4) + ] + means.append(float(np.mean(vals))) + assert means == sorted(means, reverse=True), f"MI must fall as concentration rises: {means}" + assert means[0] > 10 * means[-1], f"knob has too little range: {means}" + + +def test_fuzziness_changes_difficulty_not_truth(): + """fuzziness blends the expression programs only — the true MI must be identical. + + This separation is what lets a benchmark vary estimation difficulty while holding + the estimand fixed; if it leaked into the truth, MAE-vs-fuzziness would be + uninterpretable. + """ + ref = _truth(simulate_tcri(fuzziness=0.0, n_cells=400, seed=3))["true_mi"] + for f in (0.25, 0.5, 1.0): + t = _truth(simulate_tcri(fuzziness=f, n_cells=400, seed=3)) + assert t["true_mi"] == pytest.approx(ref, abs=1e-12), f"fuzziness={f} moved the truth" + + +def test_simulation_is_deterministic_given_a_seed(): + a = simulate_tcri(n_cells=250, seed=7) + b = simulate_tcri(n_cells=250, seed=7) + np.testing.assert_array_equal(a.X, b.X) + assert list(a.obs["clone_id"]) == list(b.obs["clone_id"]) + assert _truth(a)["true_mi"] == _truth(b)["true_mi"] + c = simulate_tcri(n_cells=250, seed=8) + assert not np.array_equal(a.X, c.X), "different seeds must differ" + + +def test_label_error_degrades_the_realized_coupling_only(): + """Corrupting labels must lower the *realized* MI while the population truth stands.""" + clean = _truth(simulate_tcri(label_error_rate=0.0, n_cells=3000, seed=4)) + noisy = _truth(simulate_tcri(label_error_rate=0.6, n_cells=3000, seed=4)) + assert noisy["empirical_mi"] < clean["empirical_mi"], "label noise must reduce realized MI" + assert noisy["true_mi"] == pytest.approx(clean["true_mi"], abs=1e-12) + + +# ══════════════ TIER B — does tcri's metric equal an independent oracle? ════ +def test_tcri_mi_matches_an_independent_implementation(): + """``tl.mutual_information`` on a realized joint == the oracle, both normalizations. + + The strongest fast test available: the oracle in ``tcri.datasets`` is a separate, + deliberately independent implementation, so agreement is real evidence rather + than a tautology. Catches sign errors, wrong log base, and denominator swaps. + """ + for seed in range(3): + adata = simulate_tcri(n_cells=1500, n_clones=15, n_phenotypes=4, seed=seed) + jd = _empirical_joint(adata) + oracle = mi_from_joint_oracle(jd.values) + + raw = tcri.tl.mutual_information(jd, normalized=False) + assert raw == pytest.approx(oracle["mi"], rel=1e-9, abs=1e-12) + + nmi_min = tcri.tl.mutual_information(jd, normalized=True, normalize_mode="min") + assert nmi_min == pytest.approx(oracle["nmi_min"], rel=1e-9, abs=1e-12) + + nmi_avg = tcri.tl.mutual_information(jd, normalized=True, normalize_mode="average") + assert nmi_avg == pytest.approx(oracle["nmi_average"], rel=1e-9, abs=1e-12) + + +def test_the_two_normalizations_are_not_interchangeable(): + """Guards the benchmark trap: tcri defaults to 'min', the note's grid used the mean. + + ``min <= average`` denominators means nmi_min >= nmi_average, so comparing a + 'min'-normalized estimate to a mean-normalized ground truth silently inflates the + estimate. A benchmark must pick deliberately. + """ + adata = simulate_tcri(n_cells=1200, n_clones=25, n_phenotypes=4, seed=1) + jd = _empirical_joint(adata) + nmi_min = tcri.tl.mutual_information(jd, normalize_mode="min") + nmi_avg = tcri.tl.mutual_information(jd, normalize_mode="average") + assert nmi_min > nmi_avg + t = _truth(adata) + assert t["empirical_nmi_min"] > t["empirical_nmi_average"] + + +def test_empirical_mi_is_biased_upward_at_small_n(): + """The plug-in estimator over-reports on small samples — assert the known sign. + + Documents why a recovery test cannot demand ``estimate == truth`` at small N. + """ + gaps = [ + _truth(simulate_tcri(n_cells=150, n_clones=25, n_phenotypes=5, seed=s))["empirical_mi"] + - _truth(simulate_tcri(n_cells=150, n_clones=25, n_phenotypes=5, seed=s))["true_mi"] + for s in range(8) + ] + assert float(np.mean(gaps)) > 0, f"expected upward plug-in bias, got {np.mean(gaps):+.4f}" + + +# ══════════════════ TIER D — metamorphic invariances ════════════════════════ +def test_mi_is_invariant_to_relabeling(): + """Renaming clones or phenotypes cannot change an information quantity.""" + adata = simulate_tcri(n_cells=900, n_clones=12, n_phenotypes=4, seed=2) + jd = _empirical_joint(adata) + base = tcri.tl.mutual_information(jd, normalized=False) + + rng = np.random.default_rng(0) + shuffled = jd.iloc[rng.permutation(jd.shape[0]), rng.permutation(jd.shape[1])] + assert tcri.tl.mutual_information(shuffled, normalized=False) == pytest.approx(base, rel=1e-9) + + +def test_mi_is_invariant_to_uniform_replication(): + """Doubling every count is the same distribution — MI must not move.""" + adata = simulate_tcri(n_cells=800, n_clones=10, n_phenotypes=3, seed=5) + jd = _empirical_joint(adata) + base = tcri.tl.mutual_information(jd, normalized=False) + assert tcri.tl.mutual_information(jd * 7, normalized=False) == pytest.approx(base, rel=1e-9) + + +# ══════════════════════ TIER C/E — slow, model-based ════════════════════════ +@pytest.mark.slow +def test_empirical_mi_converges_to_the_population_value(): + """|empirical - true| must shrink with N **on average over seeds**. + + Deliberately not a single-seed monotonicity check: one draw can land further away + at larger N (observed +0.003 at N=5000 vs +0.012 at N=20000, same seed). + """ + n_seeds = 6 + mae = {} + for n in (250, 1000, 4000): + errs = [ + abs(_truth(simulate_tcri(n_cells=n, n_clones=20, n_phenotypes=4, seed=s))["empirical_mi"] + - _truth(simulate_tcri(n_cells=n, n_clones=20, n_phenotypes=4, seed=s))["true_mi"]) + for s in range(n_seeds) + ] + mae[n] = float(np.mean(errs)) + assert mae[250] > mae[1000] > mae[4000], f"MAE must fall with N: {mae}" + + +@pytest.mark.slow +def test_model_mi_tracks_the_true_mi_across_difficulty(): + """A *fitted* model's MI must respond to the ground truth, not sit at a constant. + + The weakest defensible claim about the full pipeline: three datasets with very + different true MI must come back ordered correctly. Not an equality test — the + model estimator is posterior-based and differs from the plug-in value. + """ + from tcri.model._model import TCRIModel + + got = [] + for conc in (0.05, 1.0, 20.0): + pyro.clear_param_store() + adata = simulate_tcri( + omega_concentration=conc, n_cells=1200, n_clones=20, + n_phenotypes=4, n_genes=40, seed=11, + ) + TCRIModel.setup_anndata( + adata, layer="counts", clonotype_key="clone_id", + phenotype_key="phenotype", covariate_key="covariate", batch_key="batch", + ) + model = TCRIModel( + adata, n_latent=16, n_hidden=32, n_layers=1, + classifier_n_layers=1, classifier_hidden=32, K=4, + ) + with contextlib.redirect_stdout(io.StringIO()): + model.train(max_epochs=60, batch_size=256, + enable_progress_bar=False, enable_model_summary=False) + model.to_anndata(adata) + est = tcri.tl.mutual_information( + adata, covariate="cov_0", weighted=True, normalize_mode="average", + ) + got.append((_truth(adata)["true_nmi_average"], float(est))) + + truths = [g[0] for g in got] + ests = [g[1] for g in got] + assert truths == sorted(truths, reverse=True), f"fixture truths not ordered: {truths}" + assert ests == sorted(ests, reverse=True), ( + f"model MI did not track the truth: truths={truths} estimates={ests}" + ) + + +@pytest.mark.slow +def test_posterior_hdi_covers_the_truth(): + """Calibration: the 94% HDI from ``n_samples>0`` should contain the realized MI. + + Coverage testing (the lightweight form of simulation-based calibration). Validates + the whole Dirichlet-draw path — nothing else in the suite asserts that the + posterior interval *means* anything. + + Measured at design time over 8 independent replicates: **8/8 covered** (mean HDI + width 0.103), with posterior means tracking the realized value closely (e.g. + truth 0.2868 vs mean 0.2825). With only 8 replicates that is consistent with true + coverage anywhere from roughly 70% to 100%, so this is evidence of *no + miscalibration* rather than proof of calibration — establishing 94% would need + ~50+ replicates and belongs in ``benchmarks/``, not a test. + + The bar (>=6/8) is set to catch gross miscalibration or a degenerate/inverted + interval while tolerating sampling noise. It is deliberately not >=2/5, which + would pass even at 40% coverage. + """ + from tcri.model._model import TCRIModel + + covered = 0 + n_rep = 8 + for seed in range(n_rep): + pyro.clear_param_store() + adata = simulate_tcri( + n_cells=1200, n_clones=18, n_phenotypes=4, n_genes=40, + omega_concentration=0.4, seed=100 + seed, + ) + TCRIModel.setup_anndata( + adata, layer="counts", clonotype_key="clone_id", + phenotype_key="phenotype", covariate_key="covariate", batch_key="batch", + ) + model = TCRIModel( + adata, n_latent=16, n_hidden=32, n_layers=1, + classifier_n_layers=1, classifier_hidden=32, K=4, + ) + with contextlib.redirect_stdout(io.StringIO()): + model.train(max_epochs=60, batch_size=256, + enable_progress_bar=False, enable_model_summary=False) + model.to_anndata(adata) + summary = tcri.tl.mutual_information( + adata, covariate="cov_0", n_samples=100, weighted=True, + normalize_mode="average", random_state=seed, + ) + lo, hi = summary["hdi_low"], summary["hdi_high"] + assert hi > lo and np.isfinite(lo) and np.isfinite(hi), f"degenerate HDI: {summary}" + assert hi - lo < 0.5, f"HDI too wide to be informative: [{lo}, {hi}]" + if lo <= _truth(adata)["empirical_nmi_average"] <= hi: + covered += 1 + assert covered >= 6, f"94% HDI covered only {covered}/{n_rep} replicates"