diff --git a/benchmarks/run_grid.py b/benchmarks/run_grid.py new file mode 100644 index 0000000..a4fb8e8 --- /dev/null +++ b/benchmarks/run_grid.py @@ -0,0 +1,287 @@ +#!/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") + +# 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]), + "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))), + # 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))), +} + + +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, temperature=1.0, fit_params=None, profile=False, + local_scale=None): + """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() + 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() + TCRIModel.setup_anndata( + adata, layer="counts", clonotype_key="clone_id", phenotype_key="phenotype", + covariate_key="covariate", batch_key="batch", + ) + # 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() + 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) + + # 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, 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, + 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, + ) + 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") + 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'.") + 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") + 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] + 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 " + "(the synthetic omega cannot reproduce the published anchors)") + combos = list(itertools.product(grid["fuzziness"], grid["n_cells"], + 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, 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, + 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} " + 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() 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. 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/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..d8a87e4 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,134 @@ 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) + + # 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 + + 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 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)