From 72f13b779f30f2b452d609d5460582034066ea06 Mon Sep 17 00:00:00 2001 From: Nicholas Ceglia Date: Mon, 3 Aug 2026 16:16:52 -0400 Subject: [PATCH 01/10] feat(gpu): make the device seam reachable from every metric (CU-01) The `_compute/_xp` device seam has existed since PR5 (resolve_device / torch_device, CPU + torch-CUDA), and `joint_distribution` accepted `device=` -- but NO metric exposed it, and `joint_draws` hard-coded `device=None` when calling the engine. So the GPU path was documented and completely unreachable: every metric ran on CPU no matter what the caller asked for. Threaded `device=` from the four public metrics down to `_joint_draws`: tl.mutual_information / clonotypic_entropy / phenotypic_entropy / phenotypic_flux -> _entropy_metric / _flux_once -> joint_draws -> _engine_blocks -> _joint_draws Verified by spying on `_joint_draws`: all four now deliver the value (flux twice, once per covariate). Results are unchanged -- device is a placement detail, not a numerical one -- pinned by test_device_does_not_change_results. The API contract caught this correctly: adding a public parameter is frozen interface drift, and test_contract_conformance failed on all four signatures until `_contract.pyi` was updated deliberately. It also caught a parameter-ORDER inconsistency -- `device` had landed before `random_state` in flux but after it in the other three -- so the ordering is now uniform (`device` last everywhere). This unblocks GPU work: without it there was no code path to profile. Full suite 166 passed, 3 skipped. Co-Authored-By: Claude Opus 4.8 --- tcri/_contract.pyi | 6 ++-- tcri/tools/_common.py | 4 +-- tcri/tools/_entropy.py | 15 ++++---- tcri/tools/_flux.py | 14 ++++---- tcri/tools/_mutual_information.py | 6 ++-- tests/test_model_knobs.py | 57 +++++++++++++++++++++++++++++++ 6 files changed, 82 insertions(+), 20 deletions(-) diff --git a/tcri/_contract.pyi b/tcri/_contract.pyi index d50fe37..c76acae 100644 --- a/tcri/_contract.pyi +++ b/tcri/_contract.pyi @@ -70,23 +70,25 @@ class tl: adata_or_jd: Any, *, covariate: Optional[str] = ..., groupby: Optional[str] = ..., splitby: Optional[str] = ..., n_samples: int = ..., temperature: float = ..., clones: Any = ..., weighted: bool = ..., normalized: bool = ..., - n_clones_ref: Any = ..., random_state: Any = ..., + n_clones_ref: Any = ..., random_state: Any = ..., device: Any = ..., ) -> Any: ... def phenotypic_entropy( adata_or_jd: Any, *, covariate: Optional[str] = ..., groupby: Optional[str] = ..., splitby: Optional[str] = ..., n_samples: int = ..., temperature: float = ..., clones: Any = ..., weighted: bool = ..., normalized: bool = ..., random_state: Any = ..., + device: Any = ..., ) -> Any: ... def mutual_information( adata_or_jd: Any, *, covariate: Optional[str] = ..., groupby: Optional[str] = ..., splitby: Optional[str] = ..., n_samples: int = ..., temperature: float = ..., clones: Any = ..., weighted: bool = ..., normalized: bool = ..., - normalize_mode: str = ..., random_state: Any = ..., + normalize_mode: str = ..., random_state: Any = ..., device: Any = ..., ) -> Any: ... def phenotypic_flux( adata: AnnData, *, cov_from: str, cov_to: str, groupby: Optional[str] = ..., splitby: Optional[str] = ..., n_samples: int = ..., temperature: float = ..., clones: Any = ..., weighted: bool = ..., distance_metric: str = ..., random_state: Any = ..., + device: Any = ..., ) -> Any: ... def compare_groups( df: pd.DataFrame, *, value: str, splitby: str, reference: Optional[str] = ..., diff --git a/tcri/tools/_common.py b/tcri/tools/_common.py index fc727fd..4e4db4d 100644 --- a/tcri/tools/_common.py +++ b/tcri/tools/_common.py @@ -21,7 +21,7 @@ def is_precomputed_joint(x) -> bool: def joint_draws(adata, covariate, *, n_samples, weighted, temperature, clones, random_state, - use_logits=True): + use_logits=True, device=None): """Return ``(draws, phenotype_cols)`` where ``draws`` is a list of ``(clone_ids, [C, P])`` per posterior draw (length 1 for ``n_samples=0``). @@ -43,7 +43,7 @@ def joint_draws(adata, covariate, *, n_samples, weighted, temperature, clones, r weighted=weighted, temperature=temperature, random_state=random_state, - device=None, + device=device, ) # Row labels in DataFrame-concat order (per covariate block, per clone). When diff --git a/tcri/tools/_entropy.py b/tcri/tools/_entropy.py index 531b2f9..5ef4c23 100644 --- a/tcri/tools/_entropy.py +++ b/tcri/tools/_entropy.py @@ -63,7 +63,7 @@ def _phenotypic_one(clone_ids, J, cols, *, normalized): def _entropy_metric(adata_or_jd, *, kind, covariate, groupby, splitby, n_samples, temperature, - clones, weighted, normalized, random_state, n_clones_ref=None): + clones, weighted, normalized, random_state, n_clones_ref=None, device=None): item_name = "phenotype" if kind == "clonotypic" else "clonotype" value = f"{kind}_entropy" @@ -77,7 +77,7 @@ def _one(clone_ids, J, cols): raise ValueError("groupby requires an AnnData, not a precomputed joint (§7.9).") def _compute(cl): - draws, cols = joint_draws(adata_or_jd, covariate, n_samples=n_samples, weighted=weighted, + draws, cols = joint_draws(adata_or_jd, covariate, n_samples=n_samples, weighted=weighted, device=device, temperature=temperature, clones=cl, random_state=random_state) per = [_one(ids, J, cols) for ids, J in draws] keys = list(per[0].keys()) @@ -93,7 +93,7 @@ def _compute(cl): one = _one(list(adata_or_jd.index), adata_or_jd.values, list(adata_or_jd.columns)) return pd.Series(one, name=value) - draws, cols = joint_draws(adata_or_jd, covariate, n_samples=n_samples, weighted=weighted, + draws, cols = joint_draws(adata_or_jd, covariate, n_samples=n_samples, weighted=weighted, device=device, temperature=temperature, clones=clones, random_state=random_state) per = [_one(ids, J, cols) for ids, J in draws] keys = list(per[0].keys()) @@ -104,20 +104,21 @@ def _compute(cl): def clonotypic_entropy(adata_or_jd, *, covariate=None, groupby=None, splitby=None, n_samples=0, temperature=1.0, clones=None, weighted=False, normalized=True, - n_clones_ref=None, random_state=None): + n_clones_ref=None, random_state=None, device=None): """H[P(c|φ)] per phenotype (bits). ``n_clones_ref`` fixes the normalizer for cross-group comparability (else per-group #supported clones). See module docstring.""" return _entropy_metric(adata_or_jd, kind="clonotypic", covariate=covariate, groupby=groupby, splitby=splitby, n_samples=n_samples, temperature=temperature, clones=clones, weighted=weighted, normalized=normalized, - random_state=random_state, n_clones_ref=n_clones_ref) + random_state=random_state, n_clones_ref=n_clones_ref, + device=device) def phenotypic_entropy(adata_or_jd, *, covariate=None, groupby=None, splitby=None, n_samples=0, temperature=1.0, clones=None, weighted=False, normalized=True, - random_state=None): + random_state=None, device=None): """H[P(φ|c)] per clone (bits). See module docstring.""" return _entropy_metric(adata_or_jd, kind="phenotypic", covariate=covariate, groupby=groupby, splitby=splitby, n_samples=n_samples, temperature=temperature, clones=clones, weighted=weighted, normalized=normalized, - random_state=random_state) + random_state=random_state, device=device) diff --git a/tcri/tools/_flux.py b/tcri/tools/_flux.py index 587b09f..7f2b174 100644 --- a/tcri/tools/_flux.py +++ b/tcri/tools/_flux.py @@ -18,11 +18,11 @@ def _flux_once(adata, *, cov_from, cov_to, n_samples, weighted, temperature, clones, - distance_metric, random_state): + distance_metric, random_state, device=None): dist_fn = phenotype_distance(distance_metric) - draws_from, _ = joint_draws(adata, cov_from, n_samples=n_samples, weighted=weighted, + draws_from, _ = joint_draws(adata, cov_from, n_samples=n_samples, weighted=weighted, device=device, temperature=temperature, clones=clones, random_state=random_state) - draws_to, _ = joint_draws(adata, cov_to, n_samples=n_samples, weighted=weighted, + draws_to, _ = joint_draws(adata, cov_to, n_samples=n_samples, weighted=weighted, device=device, temperature=temperature, clones=clones, random_state=random_state) per = [] for (ids_f, Jf), (ids_t, Jt) in zip(draws_from, draws_to): @@ -47,20 +47,22 @@ def _flux_once(adata, *, cov_from, cov_to, n_samples, weighted, temperature, clo def phenotypic_flux(adata, *, cov_from, cov_to, groupby=None, splitby=None, n_samples=0, temperature=1.0, clones=None, weighted=False, distance_metric="l1", - random_state=None): + random_state=None, device=None): """Per-clone phenotype-distribution distance from ``cov_from`` to ``cov_to`` (bits for kl/jsd). ``groupby`` → tidy DataFrame (one row per group×clone).""" if groupby is not None: def _compute(cl): return _flux_once(adata, cov_from=cov_from, cov_to=cov_to, n_samples=n_samples, weighted=weighted, temperature=temperature, clones=cl, - distance_metric=distance_metric, random_state=random_state) + distance_metric=distance_metric, random_state=random_state, + device=device) return grouped_series(adata, groupby=groupby, splitby=splitby, item_name="clonotype", value="phenotypic_flux", compute=_compute) point, drawsd = _flux_once(adata, cov_from=cov_from, cov_to=cov_to, n_samples=n_samples, weighted=weighted, temperature=temperature, clones=clones, - distance_metric=distance_metric, random_state=random_state) + distance_metric=distance_metric, random_state=random_state, + device=device) if n_samples and int(n_samples) > 0: return pd.DataFrame({c: summarize(drawsd[c]) for c in point}).T return pd.Series(point, name="phenotypic_flux") diff --git a/tcri/tools/_mutual_information.py b/tcri/tools/_mutual_information.py index e836e87..00c0ca2 100644 --- a/tcri/tools/_mutual_information.py +++ b/tcri/tools/_mutual_information.py @@ -37,7 +37,7 @@ def _mi_from_joint(J: np.ndarray, *, normalized: bool = True, mode: str = "min") def mutual_information( adata_or_jd, *, covariate=None, groupby=None, splitby=None, n_samples=0, temperature=1.0, clones=None, weighted=False, normalized=True, - normalize_mode="min", random_state=None, + normalize_mode="min", random_state=None, device=None, ): """I(c;φ|covariate) in bits. ``groupby`` → tidy DataFrame (one row per group); otherwise a scalar (``n_samples=0``) or a mean/sd/hdi summary (``n_samples>0``).""" @@ -47,7 +47,7 @@ def mutual_information( def _compute(cl): draws, cols = joint_draws( - adata_or_jd, covariate, n_samples=n_samples, weighted=weighted, + adata_or_jd, covariate, n_samples=n_samples, weighted=weighted, device=device, temperature=temperature, clones=cl, random_state=random_state, ) vals = [_mi_from_joint(J, normalized=normalized, mode=normalize_mode) for _, J in draws] @@ -60,7 +60,7 @@ def _compute(cl): return _mi_from_joint(adata_or_jd.values, normalized=normalized, mode=normalize_mode) draws, cols = joint_draws( - adata_or_jd, covariate, n_samples=n_samples, weighted=weighted, + adata_or_jd, covariate, n_samples=n_samples, weighted=weighted, device=device, temperature=temperature, clones=clones, random_state=random_state, ) vals = [_mi_from_joint(J, normalized=normalized, mode=normalize_mode) for _, J in draws] diff --git a/tests/test_model_knobs.py b/tests/test_model_knobs.py index 2b7224c..ad43d98 100644 --- a/tests/test_model_knobs.py +++ b/tests/test_model_knobs.py @@ -325,3 +325,60 @@ def test_predict_is_invariant_to_batch_size(adata): np.testing.assert_allclose(a, b, atol=1e-6, rtol=1e-5) # rows must still be probability vectors regardless of chunking np.testing.assert_allclose(a.sum(1), 1.0, atol=1e-5) + + +# ══════════════════════ device seam (CU-01) ═════════════════════════════════ + +def test_device_reaches_the_engine_from_every_metric(adata): + """``device=`` must actually configure the numeric core. + + The seam existed in ``_compute/_xp`` from PR5 but no metric exposed it, so it was + unreachable — GPU was documented and dead. This asserts the value arrives at + ``_joint_draws`` for every public metric, which is the only thing that makes a + GPU run possible. + """ + import tcri + import tcri._compute._joint as CJ + import tcri.tools._joint as TJ + + m = _model(adata) + _train(m) + m.to_anndata(adata) + covs = list(adata.uns["tcri_covariate_categories"]) + + seen = [] + orig = CJ._joint_draws + + def spy(*a, **k): + seen.append(k.get("device")) + return orig(*a, **k) + + TJ._joint_draws = spy + try: + for call in ( + lambda: tcri.tl.mutual_information(adata, covariate=covs[0], device="cpu"), + lambda: tcri.tl.clonotypic_entropy(adata, covariate=covs[0], device="cpu"), + lambda: tcri.tl.phenotypic_entropy(adata, covariate=covs[0], device="cpu"), + lambda: tcri.tl.phenotypic_flux( + adata, cov_from=covs[0], cov_to=covs[1], device="cpu"), + ): + seen.clear() + call() + assert seen and all(d == "cpu" for d in seen), ( + f"device did not reach the engine: {seen}" + ) + finally: + TJ._joint_draws = orig + + +def test_device_does_not_change_results(adata): + """Routing through the device seam is a placement detail, not a numerical one.""" + import tcri + + m = _model(adata) + _train(m) + m.to_anndata(adata) + cov = list(adata.uns["tcri_covariate_categories"])[0] + a = tcri.tl.mutual_information(adata, covariate=cov, device=None) + b = tcri.tl.mutual_information(adata, covariate=cov, device="cpu") + assert a == pytest.approx(b, rel=1e-12) From b9476efd6331930e873cf7ad4acf752ca3f91884 Mon Sep 17 00:00:00 2001 From: Nicholas Ceglia Date: Mon, 3 Aug 2026 16:20:11 -0400 Subject: [PATCH 02/10] feat(benchmarks): synthetic MAE grid vs the known-MI oracle Grid runner reproducing Supplementary Note 1's benchmark design: sweep fuzziness x N x K-at-inference x seeds, compare TCRi's normalized MI against the closed-form truth from tcri.datasets.simulate_tcri, plus a KMeans/GMM baseline. Two traps the harness handles explicitly: - NORMALIZATION. eq 6 uses the MEAN denominator; tcri defaults to 'min'. Comparing a min-normalized estimate to a mean-normalized truth silently inflates it, so --normalize-mode is explicit and defaults to 'average' (eq 6). - WHICH ORACLE. true_* is the population value, empirical_* is what a perfect estimator returns on the realized sample; they differ by the plug-in bias at small N. MAE is reported against both. Presets: smoke (2 cells) / reduced (24) / full (1500, the note's grid). --profile attaches the torch profiler to one training run. Local smoke: TCRi AE 0.034-0.037 vs KMeans 0.116-0.189 -- the baseline-dominance claim reproduces. HDI covered the realized value 2/2. Co-Authored-By: Claude Opus 4.8 --- benchmarks/run_grid.py | 213 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 213 insertions(+) create mode 100644 benchmarks/run_grid.py diff --git a/benchmarks/run_grid.py b/benchmarks/run_grid.py new file mode 100644 index 0000000..275aaf2 --- /dev/null +++ b/benchmarks/run_grid.py @@ -0,0 +1,213 @@ +#!/usr/bin/env python3 +"""Synthetic benchmark grid — MAE of the NMI estimate against a known oracle. + +Reproduces the design of Supplementary Note 1's "Benchmarks" section: sweep the +difficulty (``fuzziness``), the sample size (``N``), and the number of phenotypes +supplied at inference (``K``), and compare TCRi's normalized MI against the closed-form +truth from :func:`tcri.datasets.simulate_tcri` — plus a GMM/KMeans baseline. + +Two things this harness is careful about, both of which are easy to get wrong: + +* **Normalization.** The note's eq 6 uses the MEAN denominator; tcri defaults to + ``min``. Comparing a ``min``-normalized estimate against a mean-normalized truth + silently inflates the estimate, so ``--normalize-mode`` is explicit and defaults to + ``average`` (i.e. eq 6). +* **Which oracle.** ``true_*`` is the population value; ``empirical_*`` is what a + perfect estimator returns on the realized finite sample. MAE is reported against + both, because at small N they differ by the plug-in bias. + +Usage:: + + python benchmarks/run_grid.py --preset reduced --device cuda --out results.csv + python benchmarks/run_grid.py --preset full --device cuda --profile +""" +from __future__ import annotations + +import argparse +import contextlib +import io +import itertools +import json +import time +import warnings + +import numpy as np +import pandas as pd + +warnings.filterwarnings("ignore") + +PRESETS = { + # shake-out grid: proves the pipeline and gives a real MAE-vs-fuzziness curve + "smoke": dict(fuzziness=[0.0, 0.9], n_cells=[1000], k_infer=[5], seeds=[0]), + "reduced": dict(fuzziness=[0.0, 0.3, 0.6, 0.9], n_cells=[500, 2000], + k_infer=[5], seeds=[0, 1, 2]), + # the note's grid (K supplied at inference varies around the true 5) + "full": dict(fuzziness=[round(0.1 * i, 1) for i in range(10)], + n_cells=[250, 500, 1000, 2000, 5000], + k_infer=[4, 5, 6], seeds=list(range(10))), +} + + +def _baseline_nmi(adata, k, seed, method="kmeans"): + """Cluster expression, then compute NMI from the (clone, cluster) table. + + The comparison point from the note: an estimator that ignores the hierarchical + model and just clusters cells, then measures clone/cluster coupling. + """ + from sklearn.cluster import KMeans + from sklearn.mixture import GaussianMixture + from sklearn.preprocessing import StandardScaler + + from tcri.datasets import mi_from_joint_oracle + + X = np.asarray(adata.layers["counts"], dtype=float) + X = X / np.clip(X.sum(1, keepdims=True), 1e-9, None) * 1e4 + X = StandardScaler().fit_transform(X) + if method == "gmm": + labels = GaussianMixture(n_components=k, random_state=seed, + covariance_type="diag").fit_predict(X) + else: + labels = KMeans(n_clusters=k, random_state=seed, n_init=10).fit_predict(X) + tab = pd.crosstab(adata.obs["clone_id"], pd.Series(labels, index=adata.obs_names)) + return mi_from_joint_oracle(tab.values) + + +def run_cell(fuzz, n_cells, k_infer, seed, *, device, n_samples, epochs, + normalize_mode, baseline, profile=False): + """One grid point -> a dict of results.""" + import pyro + + import tcri + from tcri.datasets import simulate_tcri + from tcri.model._model import TCRIModel + + t_all = time.time() + adata = simulate_tcri( + n_clones=40, n_phenotypes=5, n_genes=200, n_cells=n_cells, + omega_concentration=0.4, fuzziness=fuzz, seed=seed, + ) + truth = adata.uns["tcri_truth"] + + pyro.clear_param_store() + TCRIModel.setup_anndata( + adata, layer="counts", clonotype_key="clone_id", phenotype_key="phenotype", + covariate_key="covariate", batch_key="batch", + ) + model = TCRIModel(adata, n_latent=32, n_hidden=64, n_layers=2, + classifier_n_layers=1, classifier_hidden=64, K=k_infer) + + acc = "gpu" if device == "cuda" else "cpu" + t0 = time.time() + with contextlib.redirect_stdout(io.StringIO()): + model.train(max_epochs=epochs, batch_size=1024, accelerator=acc, + enable_progress_bar=False, enable_model_summary=False) + model.to_anndata(adata) + t_train = time.time() - t0 + + t0 = time.time() + est = tcri.tl.mutual_information( + adata, covariate="cov_0", n_samples=n_samples, weighted=True, + normalize_mode=normalize_mode, device=device, random_state=seed, + ) + t_metric = time.time() - t0 + est_mean = float(est["mean"]) if isinstance(est, dict) else float(est) + + key = "nmi_average" if normalize_mode == "average" else "nmi_min" + true_v, emp_v = truth[f"true_{key}"], truth[f"empirical_{key}"] + + row = dict( + fuzziness=fuzz, n_cells=n_cells, k_infer=k_infer, seed=seed, device=device, + true_nmi=true_v, empirical_nmi=emp_v, tcri_nmi=est_mean, + ae_vs_true=abs(est_mean - true_v), ae_vs_empirical=abs(est_mean - emp_v), + t_train=t_train, t_metric=t_metric, t_total=time.time() - t_all, + ) + if isinstance(est, dict): + row.update(hdi_low=est["hdi_low"], hdi_high=est["hdi_high"], + covers_empirical=bool(est["hdi_low"] <= emp_v <= est["hdi_high"])) + if baseline: + b = _baseline_nmi(adata, k_infer, seed, method=baseline) + row[f"{baseline}_nmi"] = b[key] + row[f"ae_{baseline}_vs_true"] = abs(b[key] - true_v) + return row + + +def main(): + ap = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--preset", choices=list(PRESETS), default="reduced") + ap.add_argument("--device", default=None, help="None|cpu|cuda (metrics engine)") + ap.add_argument("--epochs", type=int, default=60) + ap.add_argument("--n-samples", type=int, default=200) + ap.add_argument("--normalize-mode", choices=["average", "min"], default="average", + help="'average' == the note's eq 6 (default); 'min' == tcri's default") + ap.add_argument("--baseline", choices=["kmeans", "gmm", "none"], default="kmeans") + ap.add_argument("--out", default="benchmark_results.csv") + ap.add_argument("--profile", action="store_true", help="torch profiler on one cell") + args = ap.parse_args() + + grid = PRESETS[args.preset] + combos = list(itertools.product(grid["fuzziness"], grid["n_cells"], + grid["k_infer"], grid["seeds"])) + baseline = None if args.baseline == "none" else args.baseline + print(f"preset={args.preset} cells={len(combos)} device={args.device} " + f"normalize_mode={args.normalize_mode} baseline={baseline}", flush=True) + + rows = [] + t0 = time.time() + for i, (f, n, k, s) in enumerate(combos, 1): + r = run_cell(f, n, k, s, device=args.device, n_samples=args.n_samples, + epochs=args.epochs, normalize_mode=args.normalize_mode, + baseline=baseline) + rows.append(r) + print(f"[{i:>4}/{len(combos)}] f={f} N={n} K={k} s={s} | " + f"tcri={r['tcri_nmi']:.4f} true={r['true_nmi']:.4f} " + f"AE={r['ae_vs_true']:.4f} | {r['t_total']:.1f}s", flush=True) + + df = pd.DataFrame(rows) + df.to_csv(args.out, index=False) + print(f"\nwrote {args.out} ({len(df)} rows, {time.time()-t0:.0f}s total)") + + print("\n=== MAE vs fuzziness (mean over N, K, seeds) ===") + cols = ["ae_vs_true", "ae_vs_empirical"] + ( + [f"ae_{baseline}_vs_true"] if baseline else []) + print(df.groupby("fuzziness")[cols].mean().round(4).to_string()) + if "covers_empirical" in df: + print(f"\nHDI coverage of the realized value: " + f"{df['covers_empirical'].mean():.1%} ({int(df['covers_empirical'].sum())}/{len(df)})") + + if args.profile: + _profile_one(args) + + +def _profile_one(args): + """torch-profiler breakdown of a single training run + device-sync count.""" + import torch + from torch.profiler import ProfilerActivity, profile + + import pyro + + from tcri.datasets import simulate_tcri + from tcri.model._model import TCRIModel + + print("\n=== profile: one training run ===", flush=True) + adata = simulate_tcri(n_clones=40, n_phenotypes=5, n_genes=200, n_cells=4000, seed=0) + pyro.clear_param_store() + TCRIModel.setup_anndata(adata, layer="counts", clonotype_key="clone_id", + phenotype_key="phenotype", covariate_key="covariate", + batch_key="batch") + model = TCRIModel(adata, n_latent=32, n_hidden=64, n_layers=2, + classifier_n_layers=1, classifier_hidden=64, K=5) + acts = [ProfilerActivity.CPU] + if args.device == "cuda" and torch.cuda.is_available(): + acts.append(ProfilerActivity.CUDA) + acc = "gpu" if args.device == "cuda" else "cpu" + with profile(activities=acts, record_shapes=False) as prof: + with contextlib.redirect_stdout(io.StringIO()): + model.train(max_epochs=10, batch_size=1024, accelerator=acc, + enable_progress_bar=False, enable_model_summary=False) + sort_key = "cuda_time_total" if ProfilerActivity.CUDA in acts else "cpu_time_total" + print(prof.key_averages().table(sort_by=sort_key, row_limit=18)) + + +if __name__ == "__main__": + main() From 666ccbf4766e7ff4f4b99432f4a5ac7ec99302c9 Mon Sep 17 00:00:00 2001 From: Nicholas Ceglia Date: Mon, 3 Aug 2026 23:10:26 -0400 Subject: [PATCH 03/10] chore: remove stray handbook swept in by git add -A docs/contract/ss-small-team-handbook.md (1998 lines) was an untracked local file that my 'git add -A' picked up into the metrics-contract commit. It is unrelated to the package and was never meant to be committed. Removed from the tree and added to .gitignore so it cannot be swept up again. Co-Authored-By: Claude Opus 4.8 --- .gitignore | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitignore b/.gitignore index d6c0452..eb4c0d5 100644 --- a/.gitignore +++ b/.gitignore @@ -147,3 +147,4 @@ example/*.pdf # local-only dev harnesses (real-data tests; NOT CI) dev/ +docs/contract/ss-small-team-handbook.md From 59032d40dbc5c8474eb82a4cdfc8f40b2b6830df Mon Sep 17 00:00:00 2001 From: Nicholas Ceglia Date: Mon, 3 Aug 2026 23:25:32 -0400 Subject: [PATCH 04/10] feat(datasets): generate from the published benchmark's empirical fit simulate_from_fit_params() reproduces sc_simulator.simulate_dataset from a fitted (pi, omega, gamma_params, V), plus temperature_scale() for the omega^(1/T) axis the published benchmark sweeps. Needed because the synthetic generator CANNOT stand in for the empirical one. The benchmark's true-NMI anchors are 0.520/0.316/0.182 at T=0.1/0.5/1.0; the best symmetric-Dirichlet fit (omega_concentration=0.763) gives 0.633/0.320/0.150, and its sharpening ratio T=0.1->T=1.0 is 4.22x versus the true 2.86x. The mismatch is in the SHAPE of omega's response to temperature, so no reparameterization fixes it. The empirical fit matches all three exactly. Because gamma_params/V are in the pickle, the expression generator is reproducible too, so estimator difficulty matches -- not only the ground truth. Co-Authored-By: Claude Opus 4.8 --- tcri/datasets/__init__.py | 10 ++- tcri/datasets/_simulate.py | 131 ++++++++++++++++++++++++++++++++++++- 2 files changed, 138 insertions(+), 3 deletions(-) diff --git a/tcri/datasets/__init__.py b/tcri/datasets/__init__.py index 9f1fd2f..f84841a 100644 --- a/tcri/datasets/__init__.py +++ b/tcri/datasets/__init__.py @@ -3,6 +3,12 @@ :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 +from ._simulate import ( + mi_from_joint_oracle, + simulate_from_fit_params, + simulate_tcri, + temperature_scale, +) -__all__ = ["simulate_tcri", "mi_from_joint_oracle"] +__all__ = ["simulate_tcri", "mi_from_joint_oracle", "simulate_from_fit_params", + "temperature_scale"] diff --git a/tcri/datasets/_simulate.py b/tcri/datasets/_simulate.py index af879c7..5691195 100644 --- a/tcri/datasets/_simulate.py +++ b/tcri/datasets/_simulate.py @@ -38,7 +38,8 @@ import pandas as pd from anndata import AnnData -__all__ = ["simulate_tcri", "mi_from_joint_oracle"] +__all__ = ["simulate_tcri", "mi_from_joint_oracle", "simulate_from_fit_params", + "temperature_scale"] def mi_from_joint_oracle(joint: np.ndarray) -> dict: @@ -230,3 +231,131 @@ def simulate_tcri( }, } return adata + + +# ── benchmark reproduction: generate from an empirical fit ────────────────── + +def temperature_scale(P, T, eps=1e-12): + """Sharpen/flatten a row-stochastic matrix: ``P**(1/T)`` renormalized. + + Verbatim behaviour of ``sc_simulator.temperature_scale_conditional``. ``T<1`` + sharpens (raising I(c;phi)), ``T>1`` flattens. This is the axis the published + benchmark sweeps, and it changes the GROUND TRUTH, not just the difficulty. + """ + P = np.clip(np.asarray(P, dtype=float), eps, None) + Pp = P ** (1.0 / T) + return Pp / Pp.sum(axis=1, keepdims=True) + + +def simulate_from_fit_params( + params, + *, + n_cells: int = 1000, + temperature: float = 1.0, + fuzziness: float = 0.0, + label_error_rate: float = 0.0, + seed: int = 0, +) -> AnnData: + """Simulate from an **empirically fitted** ``(pi, omega, gamma_params, V)``. + + Reproduces ``sc_simulator.simulate_dataset``: ``z ~ Cat(pi)``, + ``phi|z ~ Cat(omega[z])``, ``U ~ Gamma(alpha_phi, 1/beta_phi)``, + ``x ~ Poisson(U @ V)``. + + Use this — rather than :func:`simulate_tcri` — whenever the point is to compare + against the published benchmark. A symmetric-Dirichlet ``omega`` cannot + reproduce the benchmark's true-NMI anchors: its response to temperature has the + wrong SHAPE (sharpening ratio 4.22x vs the true 2.86x), so no reparameterization + of the synthetic generator suffices. The empirical fit matches all three anchors + exactly (0.520 / 0.316 / 0.182 at T = 0.1 / 0.5 / 1.0). + + Parameters + ---------- + params + Path to a ``fit_params.pkl``, or the already-unpickled dict. Needs + ``pi``, ``omega``, ``gamma_params``, ``V``, ``L``. + temperature + Applied to ``omega`` BEFORE sampling, so it moves the ground truth. + fuzziness + Blends the per-phenotype Gamma programs toward their mean — difficulty only, + the truth is untouched. + """ + import pickle + + if not isinstance(params, dict): + with open(params, "rb") as fh: + params = pickle.load(fh) + for key in ("pi", "omega", "gamma_params", "V"): + if key not in params: + raise KeyError(f"fit params missing {key!r}; got {sorted(params)}") + + rng = np.random.default_rng(seed) + pi = np.asarray(params["pi"], dtype=float) + pi = pi / pi.sum() + omega = np.asarray(params["omega"], dtype=float) + if temperature != 1.0: + omega = temperature_scale(omega, temperature) + omega = omega / omega.sum(axis=1, keepdims=True) + + V = np.asarray(params["V"], dtype=float) # (L, D) + L, D = V.shape + n_clones, P = omega.shape + + truth = mi_from_joint_oracle(pi[:, None] * omega) + + z = rng.choice(n_clones, size=n_cells, p=pi) + phi_true = np.array([rng.choice(P, 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, P, size=int(flip.sum())) + + # per-phenotype Gamma programs, optionally blended toward the mean + gp = params["gamma_params"] + keys = sorted(gp.keys()) + alpha = np.stack([np.asarray(gp[k]["alpha"], dtype=float) for k in keys]) + beta = np.stack([np.asarray(gp[k]["beta"], dtype=float) for k in keys]) + if fuzziness > 0: + theta = np.concatenate([alpha - 1.0, -beta], axis=1) + theta = (1.0 - fuzziness) * theta + fuzziness * theta.mean(0, keepdims=True) + alpha = np.clip(theta[:, :L] + 1.0, 1e-3, None) + beta = np.clip(-theta[:, L:], 1e-3, None) + + U = rng.gamma(alpha[phi_true], 1.0 / beta[phi_true]) + X = rng.poisson(U @ V).astype("float32") + + counts = np.zeros((n_clones, P), 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) + + clone_levels = params.get("clone_levels") + clone_names = ([str(clone_levels[i]) for i in z] if clone_levels is not None + else [f"clone_{i}" for i in z]) + + obs = pd.DataFrame({ + "clone_id": pd.Categorical(clone_names), + "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(["cov_0"] * 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(D)])) + 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": {"source": "empirical fit", "n_cells": n_cells, + "temperature": temperature, "fuzziness": fuzziness, + "label_error_rate": label_error_rate, "seed": seed, + "n_clones": n_clones, "n_phenotypes": P, "n_genes": D, "L": L}, + } + return adata From 0604f8ed58ba9cb0f61ca241c4585e484d71140f Mon Sep 17 00:00:00 2001 From: Nicholas Ceglia Date: Mon, 3 Aug 2026 23:26:04 -0400 Subject: [PATCH 05/10] fix(datasets): V is stored (D, L) in the fit; transpose to (L, D) The reference sampler does np.dot(U_i, params['V'].T), i.e. V is (n_genes, n_factors). I had assumed (L, D), which raised a matmul dimension error on the real pickle (984 vs 5). Detect the orientation from L rather than assuming. Co-Authored-By: Claude Opus 4.8 --- tcri/datasets/_simulate.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tcri/datasets/_simulate.py b/tcri/datasets/_simulate.py index 5691195..d8a87e4 100644 --- a/tcri/datasets/_simulate.py +++ b/tcri/datasets/_simulate.py @@ -297,7 +297,10 @@ def simulate_from_fit_params( omega = temperature_scale(omega, temperature) omega = omega / omega.sum(axis=1, keepdims=True) - V = np.asarray(params["V"], dtype=float) # (L, D) + # NB the fit stores V as (D, L) and the reference sampler does U @ V.T + V = np.asarray(params["V"], dtype=float) + if V.shape[0] != int(params.get("L", V.shape[1])): + V = V.T # -> (L, D) L, D = V.shape n_clones, P = omega.shape From d32b3a4aabfcc7b837ab49ab806a5913776943fa Mon Sep 17 00:00:00 2001 From: Nicholas Ceglia Date: Mon, 3 Aug 2026 23:27:58 -0400 Subject: [PATCH 06/10] feat(benchmarks): reproduce the published grid from the empirical fit Adds --fit-params plus a temperature axis, and the 'published'/'published_quick' presets matching the figures' axes (K in {8,10,12} at inference, T in {0.1,0.5,1.0}, N in {250..5000}, fuzziness 0.0-0.9, 10 seeds). Guards the trap: a preset that sweeps temperature REQUIRES --fit-params, because the synthetic Dirichlet omega cannot reproduce the published true-NMI anchors (its sharpening response is 4.22x vs the true 2.86x). Verified end to end -- the generator reproduces 0.520/0.316/0.182 at T=0.1/0.5/1.0 exactly, with 47 clones/984 genes/K=10. Co-Authored-By: Claude Opus 4.8 --- benchmarks/run_grid.py | 43 +++++++++++++++++++++++++++++++++--------- 1 file changed, 34 insertions(+), 9 deletions(-) diff --git a/benchmarks/run_grid.py b/benchmarks/run_grid.py index 275aaf2..ed1dcd0 100644 --- a/benchmarks/run_grid.py +++ b/benchmarks/run_grid.py @@ -36,6 +36,12 @@ warnings.filterwarnings("ignore") +# the published benchmark's axes (Supplementary Note 1, Benchmarks) +PUBLISHED = dict(fuzziness=[round(0.1*i,1) for i in range(10)], + n_cells=[250,500,1000,2000,5000], + k_infer=[8,10,12], temperature=[0.1,0.5,1.0], + seeds=list(range(10))) + PRESETS = { # shake-out grid: proves the pipeline and gives a real MAE-vs-fuzziness curve "smoke": dict(fuzziness=[0.0, 0.9], n_cells=[1000], k_infer=[5], seeds=[0]), @@ -45,6 +51,11 @@ "full": dict(fuzziness=[round(0.1 * i, 1) for i in range(10)], n_cells=[250, 500, 1000, 2000, 5000], k_infer=[4, 5, 6], seeds=list(range(10))), + # reproduce the published figures: needs --fit-params + "published": PUBLISHED, + "published_quick": dict(fuzziness=[0.1], n_cells=[250,1000,5000], + k_infer=[10], temperature=[0.1,0.5,1.0], + seeds=list(range(3))), } @@ -73,7 +84,7 @@ def _baseline_nmi(adata, k, seed, method="kmeans"): def run_cell(fuzz, n_cells, k_infer, seed, *, device, n_samples, epochs, - normalize_mode, baseline, profile=False): + normalize_mode, baseline, temperature=1.0, fit_params=None, profile=False): """One grid point -> a dict of results.""" import pyro @@ -82,10 +93,17 @@ def run_cell(fuzz, n_cells, k_infer, seed, *, device, n_samples, epochs, from tcri.model._model import TCRIModel t_all = time.time() - adata = simulate_tcri( - n_clones=40, n_phenotypes=5, n_genes=200, n_cells=n_cells, - omega_concentration=0.4, fuzziness=fuzz, seed=seed, - ) + if fit_params is not None: + from tcri.datasets import simulate_from_fit_params + adata = simulate_from_fit_params( + fit_params, n_cells=n_cells, temperature=temperature, + fuzziness=fuzz, seed=seed, + ) + else: + adata = simulate_tcri( + n_clones=40, n_phenotypes=5, n_genes=200, n_cells=n_cells, + omega_concentration=0.4, fuzziness=fuzz, seed=seed, + ) truth = adata.uns["tcri_truth"] pyro.clear_param_store() @@ -117,6 +135,7 @@ def run_cell(fuzz, n_cells, k_infer, seed, *, device, n_samples, epochs, row = dict( fuzziness=fuzz, n_cells=n_cells, k_infer=k_infer, seed=seed, device=device, + temperature=temperature, true_nmi=true_v, empirical_nmi=emp_v, tcri_nmi=est_mean, ae_vs_true=abs(est_mean - true_v), ae_vs_empirical=abs(est_mean - emp_v), t_train=t_train, t_metric=t_metric, t_total=time.time() - t_all, @@ -143,23 +162,29 @@ def main(): ap.add_argument("--baseline", choices=["kmeans", "gmm", "none"], default="kmeans") ap.add_argument("--out", default="benchmark_results.csv") ap.add_argument("--profile", action="store_true", help="torch profiler on one cell") + ap.add_argument("--fit-params", default=None, + help="path to a fitted params.pkl -> reproduce the published benchmark") args = ap.parse_args() grid = PRESETS[args.preset] + temps = grid.get("temperature", [1.0]) + if args.fit_params is None and temps != [1.0]: + ap.error("--fit-params is required for a preset that sweeps temperature " + "(the synthetic omega cannot reproduce the published anchors)") combos = list(itertools.product(grid["fuzziness"], grid["n_cells"], - grid["k_infer"], grid["seeds"])) + grid["k_infer"], temps, grid["seeds"])) baseline = None if args.baseline == "none" else args.baseline print(f"preset={args.preset} cells={len(combos)} device={args.device} " f"normalize_mode={args.normalize_mode} baseline={baseline}", flush=True) rows = [] t0 = time.time() - for i, (f, n, k, s) in enumerate(combos, 1): + for i, (f, n, k, T, s) in enumerate(combos, 1): r = run_cell(f, n, k, s, device=args.device, n_samples=args.n_samples, epochs=args.epochs, normalize_mode=args.normalize_mode, - baseline=baseline) + baseline=baseline, temperature=T, fit_params=args.fit_params) rows.append(r) - print(f"[{i:>4}/{len(combos)}] f={f} N={n} K={k} s={s} | " + print(f"[{i:>4}/{len(combos)}] f={f} N={n} K={k} T={T} s={s} | " f"tcri={r['tcri_nmi']:.4f} true={r['true_nmi']:.4f} " f"AE={r['ae_vs_true']:.4f} | {r['t_total']:.1f}s", flush=True) From 74c32d6de6540d89f94980d08458af9d402e1c6c Mon Sep 17 00:00:00 2001 From: Nicholas Ceglia Date: Tue, 4 Aug 2026 01:43:59 -0400 Subject: [PATCH 07/10] feat(benchmarks): --k-infer override so fixture sweeps aren't confounded by K The presets pin k_infer=10. Sweeping different fitted omegas (K=8/10/12) without this would conflate 'different fixture' with 'wrong K at inference'. Co-Authored-By: Claude Opus 5 --- benchmarks/run_grid.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/benchmarks/run_grid.py b/benchmarks/run_grid.py index ed1dcd0..b1f8838 100644 --- a/benchmarks/run_grid.py +++ b/benchmarks/run_grid.py @@ -164,9 +164,15 @@ def main(): ap.add_argument("--profile", action="store_true", help="torch profiler on one cell") ap.add_argument("--fit-params", default=None, help="path to a fitted params.pkl -> reproduce the published benchmark") + ap.add_argument("--k-infer", type=int, default=None, + help="override the preset's k_infer. Needed when sweeping FIXTURES: the " + "presets pin k_infer=10, so running a K=8 or K=12 fit unchanged would " + "conflate 'different omega' with 'wrong K at inference'.") args = ap.parse_args() grid = PRESETS[args.preset] + if args.k_infer is not None: + grid = dict(grid, k_infer=[args.k_infer]) temps = grid.get("temperature", [1.0]) if args.fit_params is None and temps != [1.0]: ap.error("--fit-params is required for a preset that sweeps temperature " From 7223079c3984dc6d05532f0b54d99be4f5375ce1 Mon Sep 17 00:00:00 2001 From: Nicholas Ceglia Date: Tue, 4 Aug 2026 10:10:48 -0400 Subject: [PATCH 08/10] feat(benchmarks): carry NMI(E[J]) alongside E[NMI(J)], plus single-cell overrides The metric reports E_s[NMI(J_s)]. NMI is nonlinear in the joint, so that is not the NMI of the posterior. Read the same draws both ways and record the gap; it needs no ground truth, so it also works on real data. --temperature/--n-cells let a single grid cell be driven to convergence without paying for the whole sweep. Co-Authored-By: Claude Opus 5 --- benchmarks/run_grid.py | 31 ++++++++++++++++++++++++++++++- 1 file changed, 30 insertions(+), 1 deletion(-) diff --git a/benchmarks/run_grid.py b/benchmarks/run_grid.py index b1f8838..bea7206 100644 --- a/benchmarks/run_grid.py +++ b/benchmarks/run_grid.py @@ -130,13 +130,34 @@ def run_cell(fuzz, n_cells, k_infer, seed, *, device, n_samples, epochs, t_metric = time.time() - t0 est_mean = float(est["mean"]) if isinstance(est, dict) else float(est) + # The shipped metric reports E_s[NMI(J_s)] (_mutual_information.py:66-68). NMI is a + # nonlinear functional of the joint, so that is not NMI of the posterior — read the + # SAME draws the other way round, NMI(E_s[J_s]), and carry both. The gap between them + # is the Jensen term, measurable with no ground truth. + mean_joint_nmi = float("nan") + if n_samples and int(n_samples) > 0: + from tcri.tools._common import joint_draws + from tcri.tools._mutual_information import _mi_from_joint + draws, _cols = joint_draws( + adata, "cov_0", n_samples=n_samples, weighted=True, + temperature=1.0, # METRIC temperature, matching the call above -- + clones=None, # NOT the generator temperature + random_state=seed, device=device, + ) + mean_joint_nmi = float(_mi_from_joint( + np.mean([J for _ids, J in draws], axis=0), + normalized=True, mode=normalize_mode)) + key = "nmi_average" if normalize_mode == "average" else "nmi_min" true_v, emp_v = truth[f"true_{key}"], truth[f"empirical_{key}"] row = dict( fuzziness=fuzz, n_cells=n_cells, k_infer=k_infer, seed=seed, device=device, - temperature=temperature, + temperature=temperature, epochs=epochs, true_nmi=true_v, empirical_nmi=emp_v, tcri_nmi=est_mean, + tcri_nmi_meanjoint=mean_joint_nmi, + jensen_gap=est_mean - mean_joint_nmi, + ae_meanjoint_vs_true=abs(mean_joint_nmi - true_v), ae_vs_true=abs(est_mean - true_v), ae_vs_empirical=abs(est_mean - emp_v), t_train=t_train, t_metric=t_metric, t_total=time.time() - t_all, ) @@ -168,11 +189,19 @@ def main(): help="override the preset's k_infer. Needed when sweeping FIXTURES: the " "presets pin k_infer=10, so running a K=8 or K=12 fit unchanged would " "conflate 'different omega' with 'wrong K at inference'.") + ap.add_argument("--temperature", type=float, default=None, + help="override the preset's temperature sweep with a single value") + ap.add_argument("--n-cells", type=int, default=None, + help="override the preset's n_cells sweep with a single value") args = ap.parse_args() grid = PRESETS[args.preset] if args.k_infer is not None: grid = dict(grid, k_infer=[args.k_infer]) + if args.n_cells is not None: + grid = dict(grid, n_cells=[args.n_cells]) + if args.temperature is not None: + grid = dict(grid, temperature=[args.temperature]) temps = grid.get("temperature", [1.0]) if args.fit_params is None and temps != [1.0]: ap.error("--fit-params is required for a preset that sweeps temperature " From 0bf980b7d19b6080681d90d34dc47b576455e188 Mon Sep 17 00:00:00 2001 From: Nicholas Ceglia Date: Tue, 4 Aug 2026 16:13:40 -0400 Subject: [PATCH 09/10] docs: record the coupling-dependent NMI bias found by the benchmark Synthetic-only finding, mechanism identified and magnitude bounded, fix undecided. Written up so the open questions (posterior summary convention, eq-2 concentration) are on record before anyone acts on the benchmark numbers. Co-Authored-By: Claude Opus 5 --- docs/nmi_temperature_bias.md | 139 +++++++++++++++++++++++++++++++++++ 1 file changed, 139 insertions(+) create mode 100644 docs/nmi_temperature_bias.md diff --git a/docs/nmi_temperature_bias.md b/docs/nmi_temperature_bias.md new file mode 100644 index 0000000..34c9895 --- /dev/null +++ b/docs/nmi_temperature_bias.md @@ -0,0 +1,139 @@ +# A coupling-strength-dependent bias in the NMI estimate + +**Status:** open. Mechanism identified, magnitude bounded, fix not yet decided. +**Scope:** `tcri.tl.mutual_information` at `n_samples > 0`. Evidence is synthetic. + +--- + +## 1. What the testing set out to do + +The package reports normalized mutual information between clonotype and phenotype. That +number appears in published figures, so it needs a calibration: given data where the true +clone→phenotype coupling is known exactly, does the estimator recover it? + +The test bed generates cells from the *same fitted parameters* behind the published +benchmark — 47 clonotypes, 984 genes, K=10 phenotypes — so the estimate is comparable to +the published numbers rather than to a toy problem. Coupling strength is swept by +sharpening or flattening the fitted ω, and the generator reproduces the published +ground-truth anchors exactly (0.520 / 0.316 / 0.182). + +Each grid cell carries three reference points: + +- **truth** — the population NMI, computed in closed form from ω. +- **label oracle** — the plug-in NMI over the realized cells' true labels. What a method + with perfect label knowledge would report on exactly this sample. +- **GMM** — a clustering baseline, as an independent control. + +## 2. What the test revealed + +TCRi's estimate is biased **upward**, and the size of the bias depends on how strong the +true coupling is. At the flattest setting it reads 0.267 where the sample's own true +labels give 0.191 and the truth is 0.182. + +The direction never flips. The estimator over-reads at every coupling strength; it simply +over-reads more when there is less real structure to find. + +Two properties made this worth pursuing rather than dismissing as sampling noise: + +- The **label oracle already absorbs finite-sample bias**. At N=5000 the oracle sits + 0.010 above truth, so an estimate 0.076 above the oracle is claiming structure the + realized sample does not contain. +- **GMM is flat across the sweep** (error 0.005–0.033, no trend). The dependence is + specific to this estimator, not a property of the metric or the test bed. + +## 3. How it was confirmed + +**It is not one fixture.** Three independently fitted parameter sets, each run at its own +K, reproduce the pattern. Error at N=5000, 1000 epochs: + +| coupling | K8 | K10 | K12 | +|---|---|---|---| +| sharp (T=0.1) | 0.013 | 0.015 | 0.017 | +| flat (T=1.0) | 0.035 | 0.058 | 0.061 | +| ratio | 2.7× | 3.8× | 3.7× | + +**It is not under-training.** An epoch ladder run to plateau at the flat setting, N=5000: + +| epochs | estimate | error | +|---|---|---| +| 60 | 0.2673 | 0.086 | +| 1000 | 0.2397 | 0.058 | +| 2000 | 0.2311 | 0.050 | +| 4000 | 0.2270 | 0.045 | +| 8000 | 0.2272 | 0.046 | + +Between 4000 and 8000 epochs the estimate moves by 0.0002. There is a real floor that +training does not remove. The early rows fall steeply, which is why a short run looks like +a convergence problem — the default 60-epoch budget lands in exactly that misleading zone. + +**The mechanism is two opposing distortions, not one.** The reported value passes through +two steps that each bias it, in opposite directions: + +- The per-cell blend applies a square root to each clone's phenotype row, flattening it. + This pushes the estimate **down**. +- The posterior is summarized as the mean of the per-draw NMI, `E[NMI(J)]`, rather than + the NMI of the mean joint, `NMI(E[J])`. NMI is nonlinear in the joint, and each Dirichlet + draw is sharper than the posterior mean it came from, so this pushes the estimate **up**. + +The reported number is the residual of the two. That is why the sharp regime looked +accurate: there the two happened to nearly cancel. + +**The nonlinearity term is itself coupling-dependent**, measured at 4000 epochs: +0.017 at +sharp coupling versus +0.099 at flat coupling. The Dirichlet concentration is +`local_scale × p_ct`, so its total is fixed at `local_scale` regardless of the data. With +`local_scale = 3` over 10 phenotypes the per-entry concentration is well below 1, which +makes draws corner-seeking — and flatter rows produce more dispersed draws, hence a larger +gap. This is the first independent support for the proposed mechanism. + +## 4. Magnitude and consequence + +**The floor is roughly 0.02 to 0.05 NMI**, depending on coupling strength — about 0.045 at +flat coupling and 0.02 at sharp, once trained to plateau. + +**The coupling dependence is about 2×**, not the ~12× the default configuration suggests. +Most of that apparent 12× was the short training budget, and it disappears with a longer +run. + +Three consequences follow: + +- **Cross-regime comparisons are not supportable as printed.** A figure claiming a method + ranking that changes with coupling strength is reading a bias that changes with coupling + strength. The ordering of methods within a single regime is less affected. +- **The accuracy at sharp coupling was partly an artifact of the default budget.** Trained + longer, the error at sharp coupling *grows* (0.007 → 0.022), because the cancellation + that made it look accurate degrades. +- **The obvious fix does not work alone.** Switching to `NMI(E[J])` removes a term larger + than the error itself, converting a +0.045 over-read into a −0.054 under-read. Correcting + one distortion without the other makes accuracy worse. + +Separately, the benchmark's smallest cells are not interpretable at all: at N=250 with flat +coupling, a label-permutation null — an estimator with no information — scores 0.214 +against a truth of 0.182. Any method evaluated there is being scored above its own noise +floor. + +## 5. Next steps + +1. **Sweep `local_scale`.** It is the single knob predicted to move the nonlinearity term + and the floor together. If raising it collapses both, the mechanism is confirmed and the + fix is a calibration rather than a redesign. +2. **Decide the posterior summary convention.** `E[NMI(J)]` and `NMI(E[J])` are different + quantities and the package should state which it reports and why. This belongs in the + metrics contract before it goes in the code, and it affects every metric that accepts + `n_samples > 0`, not only mutual information. +3. **Resolve the concentration question against Supplementary Note 1.** Total posterior + concentration does not scale with a clone's cell count, so the reported uncertainty is + set by a prior rather than informed by data. This is β on eq 2 — a model-contract + decision, not a code change. +4. **Report the noise floor per benchmark cell.** A permutation null costs almost nothing + and would have flagged the small-N cells automatically. +5. **Check whether any of this is visible on real data.** Everything above is synthetic. + Two diagnostics work without ground truth and can run on a real dataset: the value the + pipeline reports on a table with the clone structure removed, and whether posterior + width narrows as clones get more cells. + +## What is not yet known + +- Whether the plateau holds beyond 8000 epochs (16000-epoch runs pending). +- Whether the effect appears under the package's default `normalize_mode="min"`; all of + the above uses `"average"`. +- Whether the estimator behaves this way on real repertoire data. From fde342f09d7dc417f2d0e9f55ba69755848d065b Mon Sep 17 00:00:00 2001 From: Nicholas Ceglia Date: Tue, 4 Aug 2026 16:16:23 -0400 Subject: [PATCH 10/10] feat(benchmarks): --local-scale, to test the corner-seeking hypothesis Per-entry Dirichlet concentration is local_scale/P. Below 1 the posterior draws are corner-seeking, which is the proposed source of the upward NMI bias. This exposes the knob so the prediction can be tested: raising it should collapse the Jensen gap and the error floor together. Moves both the guide's posterior and (via uns) the metric's draw, so a result here is not attributable to one side without a follow-up that pins the metric separately. Co-Authored-By: Claude Opus 5 --- benchmarks/run_grid.py | 22 ++++++++++++++++++---- 1 file changed, 18 insertions(+), 4 deletions(-) diff --git a/benchmarks/run_grid.py b/benchmarks/run_grid.py index bea7206..a4fb8e8 100644 --- a/benchmarks/run_grid.py +++ b/benchmarks/run_grid.py @@ -84,7 +84,8 @@ def _baseline_nmi(adata, k, seed, method="kmeans"): def run_cell(fuzz, n_cells, k_infer, seed, *, device, n_samples, epochs, - normalize_mode, baseline, temperature=1.0, fit_params=None, profile=False): + normalize_mode, baseline, temperature=1.0, fit_params=None, profile=False, + local_scale=None): """One grid point -> a dict of results.""" import pyro @@ -111,8 +112,16 @@ def run_cell(fuzz, n_cells, k_infer, seed, *, device, n_samples, epochs, adata, layer="counts", clonotype_key="clone_id", phenotype_key="phenotype", covariate_key="covariate", batch_key="batch", ) - model = TCRIModel(adata, n_latent=32, n_hidden=64, n_layers=2, - classifier_n_layers=1, classifier_hidden=64, K=k_infer) + # local_scale sets the TOTAL Dirichlet concentration on p_ct, so per-entry + # concentration is local_scale/P. Below 1 the draws are corner-seeking, which is the + # proposed source of the upward NMI bias. Note it moves BOTH sides at once: the guide's + # posterior and, via uns, the metric's draw -- so a change here is not attributable to + # one or the other without a follow-up that pins the metric side separately. + mk = dict(n_latent=32, n_hidden=64, n_layers=2, + classifier_n_layers=1, classifier_hidden=64, K=k_infer) + if local_scale is not None: + mk["local_scale"] = float(local_scale) + model = TCRIModel(adata, **mk) acc = "gpu" if device == "cuda" else "cpu" t0 = time.time() @@ -154,6 +163,7 @@ def run_cell(fuzz, n_cells, k_infer, seed, *, device, n_samples, epochs, row = dict( fuzziness=fuzz, n_cells=n_cells, k_infer=k_infer, seed=seed, device=device, temperature=temperature, epochs=epochs, + local_scale=(local_scale if local_scale is not None else float("nan")), true_nmi=true_v, empirical_nmi=emp_v, tcri_nmi=est_mean, tcri_nmi_meanjoint=mean_joint_nmi, jensen_gap=est_mean - mean_joint_nmi, @@ -193,6 +203,9 @@ def main(): help="override the preset's temperature sweep with a single value") ap.add_argument("--n-cells", type=int, default=None, help="override the preset's n_cells sweep with a single value") + ap.add_argument("--local-scale", type=float, default=None, + help="total Dirichlet concentration on p_ct (per-entry = local_scale/P). " + "Below 1 per entry the posterior draws are corner-seeking.") args = ap.parse_args() grid = PRESETS[args.preset] @@ -217,7 +230,8 @@ def main(): for i, (f, n, k, T, s) in enumerate(combos, 1): r = run_cell(f, n, k, s, device=args.device, n_samples=args.n_samples, epochs=args.epochs, normalize_mode=args.normalize_mode, - baseline=baseline, temperature=T, fit_params=args.fit_params) + baseline=baseline, temperature=T, fit_params=args.fit_params, + local_scale=args.local_scale) rows.append(r) print(f"[{i:>4}/{len(combos)}] f={f} N={n} K={k} T={T} s={s} | " f"tcri={r['tcri_nmi']:.4f} true={r['true_nmi']:.4f} "