diff --git a/.gitignore b/.gitignore index ccdf6e3..d6c0452 100644 --- a/.gitignore +++ b/.gitignore @@ -144,3 +144,6 @@ example/*.pdf .DS_Store + +# local-only dev harnesses (real-data tests; NOT CI) +dev/ diff --git a/docs/contract/REFACTOR_AGENDA.md b/docs/contract/REFACTOR_AGENDA.md index b7c4a2d..ce8c74f 100644 --- a/docs/contract/REFACTOR_AGENDA.md +++ b/docs/contract/REFACTOR_AGENDA.md @@ -35,7 +35,7 @@ tracker + running diary for the whole refactor. The detailed spec lives in `tcri | 1 | Shared helpers + `_keys` | ✅ | low | 0 | existing tests green | | 2 | Safe deletions | ✅ | very low | 1 | import-graph clean | | 3 | Model module split | ✅ | low | 1 | model/pyro tests green | -| 4 | Model→AnnData streamline | ☐ | HIGH | 1,3 | session round-trip | +| 4 | Model→AnnData streamline | ✅ | HIGH | 1,3 | session round-trip | | 5 | Engine consolidation | ☐ | HIGH | 4 | joint identities | | 6 | Metric-API consolidation | ☐ | HIGH | 5 | metric tests | | 7 | Plotting split + pl twins | ☐ | medium | 6,1 | twins render | @@ -54,10 +54,11 @@ Tick only when the symbol is gone from source AND `__all__`/imports AND `import - [x] `metrics._ent` · [x] `tl.clone_fraction` · [x] `metrics.dkl` (→ `_distance.kl_divergence`) - [x] `ut.probabilities` (+ its `_plotting.py` import, same PR) · [x] `SankeyNode.hex_to_rgb` -**Phase 4 (folded into `to_anndata` / session):** -- [ ] `pp.register_model` (→ `model.to_anndata`) · [ ] `pp.register_phenotype_key` · [ ] `pp.register_clonotype_key` -- [ ] `pp._compute_logits_and_prior` · [ ] `ut.write_adata_safely` · [ ] `ut._pop_nonserializables` -- [ ] uns keys `tcri_manager`, `tcri_clone_key`, `tcri_phenotype_key`, obsm `X_tcri_phenotypes` +**Phase 4 (folded into `to_anndata` / session):** ✅ functions folded + manager stash retired (PR4). +- [x] `pp.register_model` (→ `model.to_anndata`) · [x] `pp.register_phenotype_key` · [x] `pp.register_clonotype_key` +- [x] `pp._compute_logits_and_prior` · [x] `ut.write_adata_safely` · [x] `ut._pop_nonserializables` +- [x] uns key `tcri_manager` (retired at `setup_anndata`; `test_model_setup` asserts it's gone) +- [ ] `tcri_clone_key` / `tcri_phenotype_key` / obsm `X_tcri_phenotypes` **→ DEFERRED to Phase 6/7**: still read by not-yet-refactored `metrics`/`plotting`; `to_anndata` writes the two `tcri_*_key` shims until their readers move. (Logged in `REFACTOR_NOTES`.) **Phase 5/6 (consolidated away — delete WITH replacement, never before):** - [ ] `pp.joint_distribution_posterior` (→ unified `joint_distribution`) · [ ] `metrics._mi_from_joint` (→ `_mutual_information`) @@ -127,7 +128,48 @@ Template per PR: **Goal · Status · What happened · Issues & fixes · Added - **Usability:** each file now has a docstring stating its role; the model file reads as a clean `BaseModelClass` API surface. - **Deferred (logged):** **M5** (`build_archetypes` default `K=4` vs `TCRIModel` `K=10`) — behavior-neutral today (the model always passes `K=10` explicitly), reconciled with persisted `labels` when `diag.archetypes` lands (Phase 8). Not touched here to keep the split purely mechanical. Also deferred (auditor's own recommendation): the stale `c2p_mat` descriptors in the contract **generator** (`build_tcri_contract.py:81,267`) + regenerating the contract HTML — bundled with the Phase-8 `diag.archetypes`/M5 pass (they describe that future function). The inventory rename-table row (`c2p_mat → clone_phenotype_prior`) is correct and stays. - **Audit (workflow — 3 lenses × adversarial verify, 8 agents):** 2 lenses PASS, plan-contract FIX. 5 findings, **all confirmed, all LOW/MED** — no behavior/correctness defect (behavior lens confirmed byte-identical class bodies + zero F821 undefined-names + suite green). Fixed here: the **MED** — explicit `__all__` per module (plan §Phase 3) was omitted — now added to all 5 files, which also resolves the two LOW "surface not byte-for-byte" findings (surface is now the explicit `{TCRIModel}`; diary wording corrected; the 3 `# noqa: F401` re-exports removed as no longer needed). Suite 36 passed / 1 skipped. -## PR 4 — Model→AnnData streamline · ☐ todo +## PR 4 — Model→AnnData streamline · ✅ done (branch `refactor/pr4-model-anndata`) +- **Goal:** kill the `uns['tcri_manager']` hack; fold `register_model → model.to_anndata` (writing the full canonical set incl. new `GATE_PROB`/`CLASSIFIER_TEMPERATURE`); rename `get_cell_phenotype_probs → predict` (labelled DataFrame); rewrite the round-trip gate. Behavior change. +- **Env (prereq):** built a fresh py3.12 venv on the latest scverse stack (anndata 0.13.1, scanpy 1.12.2, **scvi-tools 1.5.0**, torch 2.13, numpy 2.4, **pandas 3.0.3**), pinned in `requirements.txt`; all runs use `.venv`. One pandas-3.0 compat fix (legacy `tcri_boxplot` `groupby.median()` positional `numeric_only`). Suite green in the new env. +- **What happened:** `setup_anndata` → keyword-only, returns `None`, no manager stash (registration only). `predict` → labelled `DataFrame` (obs_names × phenotypes), `eval()` → deterministic. `to_anndata` → metadata + categories, `P_CT`/`CT_TO_COV`/`CT_TO_C`/per-cell `CT_ARRAY`/`COV_ARRAY`, `LOCAL_SCALE`, **`GATE_PROB`**, **`CLASSIFIER_TEMPERATURE`**, `X_TCRI`/`X_LOGITS`/`X_LOGPOSTERIOR`, `X_PROBABILITIES` (from `predict`) + argmax labels. Deleted the register cluster + `write_adata_safely`/`_pop_nonserializables` (inlined h5ad write). Onboarded the 6 `TCRIModel` methods into the contract (`IMPLEMENTED`) — conformance now enforces the model surface. +- **Issues & fixes:** the rewritten round-trip test passed in isolation but failed in the full suite — the **process-global pyro param store** (§5.2) is clobbered by other model-training tests, so the session-scoped `trained_model`'s store isn't its own at save-time. Fixed with a function-scoped `fresh_trained_model` fixture that trains inside the test and owns the store. +- **Correction to the PR3 note:** my "train path was uncovered" claim was wrong — the fixture is `trained_model` (I grepped the wrong name `fitted_model`), and `test_session_round_trip` consumes it, so train WAS covered. The PR3 smoke is still a faster/targeted guard; the framing was off. +- **Added:** ✅ rewritten `test_session_round_trip.py` (canonical write-set · setup-obs invariant · reloaded model reproduces p_ct/latent/predict) ✅ `dev/real_data_to_anndata.py` — **LOCAL-only, gitignored, not CI** — builds the 50-largest patient-specific clones of yost (`trb_unique = trb+patient`; 7682 cells / 10 patients) and runs setup→train→to_anndata (canonical keys OK; prior-driven recovery 0.90). +- **Removed (hard bar):** ✅ `register_model`, `register_phenotype_key`, `register_clonotype_key`, `_compute_logits_and_prior`, `write_adata_safely`, `_pop_nonserializables`; `uns['tcri_manager']` stash. Legacy `tcri_clone_key`/`tcri_phenotype_key`/`X_tcri_phenotypes` deferred to Phase 6/7 (live readers). +- **Test opportunities:** the **Model knob-test matrix** below (new deliverable). **Two dead knobs surfaced** — see "Model correctness debt." +- **Streamline / Usability:** `predict` returns a labelled DataFrame; `to_anndata` is one call replacing the heavy `register_model`; save/load is plain h5ad. + +--- + +# MODEL KNOB-TEST MATRIX *(PR4 deliverable — each knob gets ≥1 correctness test by end of refactor, or a justification)* + +**Legend for "Test":** the mathematically-correct input→output assertion. **Status:** ✅ tested · ◐ partial · ☐ planned (target PR) · ⛔ blocked (see debt). + +| Knob (default) | Category | Hooked up? | Mathematically-correct input→output test | Status | +|---|---|---|---|---| +| `n_latent` (128) | structural | yes | `to_anndata` ⇒ `obsm[X_tcri].shape[1] == n_latent`; `get_latent_representation` width | ☐ PR8-diag / a model-unit test | +| `n_pseudo_obs` (10) | structural | yes | `module.vamp_prior.pseudo_inputs.shape[0] == n_pseudo_obs` | ☐ model-unit | +| `K` (10) | structural | yes | `model.centers.shape[0] == K` **and** `module.mixture_concentration.shape[0] == K`; `build_archetypes` returns `centers,(labels)` with `K` clusters | ◐ smoke asserts `build_archetypes` K; add centers/mixture shape | +| `n_hidden`,`n_layers` (128,3) | structural | yes | encoder/decoder layer count/width from `module` submodules | ☐ model-unit (low value) | +| `local_scale` (3.0) | serialized + distributional | yes | (a) `to_anndata` ⇒ `uns[LOCAL_SCALE] == module.local_scale` ✅ (value equality, not just finite); (b) Dirichlet total-concentration of the `p_ct` draw ⇒ **draw variance = pᵢ(1−pᵢ)/(local_scale+1)** (exact; or assert strictly-decreasing in `local_scale`) | ◐ (a) ✅ round-trip; (b) ☐ **model-unit**, co-located with `global_scale` | +| `gate_prob` (None) | serialized + predict-mix | yes | (a) `uns[GATE_PROB]` written == configured ✅; (b) **formula identities, testable NOW** (independent of training): `gate=0` ⇒ `predict == softmax(log prior)`, `gate=1` ⇒ `predict == softmax(cls_logits)`; (c) `gate=1` phenotype **recovery** ⛔ (needs a trained classifier) | ◐ (a) ✅; (b) ☐ model-unit; (c) ⛔ | +| `classifier_temperature` (1.0) | serialized + functional | yes (division only) | (a) `uns[CLASSIFIER_TEMPERATURE]` written == configured ✅; (b) `PhenotypeClassifier.forward` divides logits by `T` ⇒ `logits(T=2)==logits(T=1)/2` for fixed weights | ◐ (a) ✅; (b) ☐ classifier-unit (division is live even if training isn't) | +| `prior_temperature` (1.0) | distributional | yes | `prepare_two_level_params`: `clone_phen_prior = normalize(prior**(1/T))` ⇒ `T>1` **raises** the row-entropy of `module.clone_phen_prior` vs `T=1` | ☐ model-unit / Phase-8 | +| `guide_temperature` (1.0) | distributional | yes | `get_p_ct` sharpens `q**(1/T)` ⇒ `T<1` **lowers** row-entropy of `get_p_ct()`. **Post-train** test (or inject a synthetic `q_p_ct_raw`): `get_p_ct` reads the *global* Pyro store, populated only after the guide runs — hold `q_p_ct_raw` fixed and flip `T` | ☐ **post-train** model-unit / Phase-8 | +| `global_scale` (5.0) | distributional | yes | Dirichlet total-concentration of the `p_c` guide draw ⇒ **draw variance = pᵢ(1−pᵢ)/(global_scale+1)** (exact; same test-kind as `local_scale`) | ☐ **model-unit**, co-located with `local_scale` | +| `batch_size` (train/predict) | invariance | yes | order/size invariance: `predict(batch_size=a).values ≈ predict(batch_size=b).values` | ☐ model-unit (cheap; add next) | +| `max_epochs`,`lr`,`n_steps_kl_warmup`,`patience`,`guide_init_scale` | training dynamics | yes | **Justification:** no closed-form per-call output; correctness is convergence — covered by the perfect-recovery test (post-classifier-fix) + the smoke/round-trip trains | ☐ justification accepted | +| `kl_weight_max`,`reconstruction_loss_scale` | training dynamics (z/recon path) | yes | **Justification (weaker):** act on the ELBO **z/reconstruction** path, which is **decoupled** from the prior-driven phenotype recovery — so covered only by "trains-without-error" (smoke/round-trip), NOT by the recovery test | ☐ justification (weaker) | +| `use_enumeration` (False) | training path | yes (near-inert) | Selects `TraceEnum_ELBO` vs `Trace_ELBO`, but the model has **no discrete latent to enumerate**, so the paths are near-equivalent. Only the default (`Trace_ELBO`) path is smoke-tested today | ☐ add a parametrized `use_enumeration=True` smoke, or document as inert | +| `classifier_hidden`,`classifier_n_layers` | classifier | **NO (untrained)** | ⛔ **BLOCKED** — the classifier never receives gradient (debt below); functional tests land **with** the classification-loss fix | ⛔ | +| `classifier_dropout` (0.1) | classifier | **NO (not plumbed)** | ⛔ **BLOCKED + wiring bug**: never forwarded to `PhenotypeClassifier` (`_module.py` constructs it without `dropout_rate=self.classifier_dropout`), so the classifier uses its own default 0.1 regardless. One-line plumbing fix (fold into the classifier-fix PR); also moot at inference (`eval()` disables dropout) | ⛔ | +| `phenotype_weights` (None) | class-imbalance | **NO (dead)** | ⛔ **BLOCKED** — `log_class_weights` registered but never read; belongs to the missing classification loss | ⛔ | + +### Model correctness debt *(surfaced in PR4; fix scheduled per user — "add it to the agenda for the correct PR")* +- **The phenotype classifier is never trained.** `TCRIModule.model()` computes `cls_logits = self.classifier(z)` but never uses it in any `pyro.sample`/ELBO factor; the training plan only evaluates the classifier under `torch.no_grad()`. Verified: classifier weight **Δ = 0.0** over 150 epochs; isolated (`gate=1.0`) recovery = **chance**. End-to-end `predict` perfect-data recovery (1.0) is entirely **prior-driven** (`p_ct`). +- **`phenotype_weights`/`class_weights` is dead** for the same root cause (would weight the absent classification loss). +- **`classifier_dropout` is not plumbed** (separate one-line wiring bug): `TCRIModule.__init__` constructs `PhenotypeClassifier(...)` without `dropout_rate=self.classifier_dropout`, so the knob is inert (the classifier keeps its own default 0.1). Fold the fix into the classifier PR. +- **Fix (deferred to its own PR, before Phase 6 metrics depend on `predict`):** wire a supervised cross-entropy of `cls_logits` vs the observed phenotype into training so the classifier learns; then add — the perfect-recovery **CI** test (user-provided `create_perfect_synthetic_anndata`), the `gate=1.0` isolated-classifier test, and the ⛔ knob tests above. Model-behavior change (shifts learned metrics + the round-trip fixture values) → isolated PR + re-validation. ## PR 5 — Engine consolidation · ☐ todo ## PR 6 — Metric-API consolidation · ☐ todo ## PR 7 — Plotting split + pl twins · ☐ todo @@ -146,4 +188,5 @@ _(dated entries; what was audited, findings, actions)_ - **(PR1 ◐):** shared-helper foundation created (`_keys`/`_console`/`_stats`/`_distance`) + 8 unit tests. Caught & fixed an `hdi` off-by-one before it shipped. **Adoption pending** (dedup, stats-move, `K.*` migration) — no ledger items ticked yet; foundation is additive, suite green. Logged: key-literal test (PR1), `pl.__all__` whole-surface test (PR11). - **(PR0+PR1 multi-agent audit — 3 lenses):** verdict FIX. Caught a real regression — the `K.*` find/replace over-reached into **10** display/warning/docstring strings (`register_model`/`load_tcri_session` printed `"K.X_LOGITS"` etc.). **Fixed:** restored readable key text in all 10 (AST-span, delimiter-safe); made the key-literal guard **AST-based** (checks real subscripts/`.get`, ignores prose); removed 3 dead `utils` imports the audit flagged. Suite 35 passed. Two non-blocking items deferred to `REFACTOR_NOTES` (contract↔api-doc reconciliation; helper-name canonicalization) — noted in the PR body. - **(PR2 multi-agent audit — 3 lenses):** PASS on all three (doc↔code · deletion safety · plan/contract). Independently re-derived: all 14 deletions have zero call-sites, all on-plan Phase-2/DROP, none in `_contract.pyi`. 3 LOW items fixed before push (orphaned `cosine_similarity` import; a −129→−127 count; a stale plan line calling `classify_phenotypes` a Phase-4 fold). +- **(PR4 audit — WORKFLOW, 4 lenses × adversarial verify, 23 agents):** the standard three (doc↔code · correctness · plan/contract) **plus a dedicated knob-test-plan lens**. **19 findings confirmed, 0 refuted — all LOW/MED, zero HIGH.** The correctness lens independently **confirmed the streamline is behavior-preserving** (44 passed) and reproduced both dead-knob findings (classifier untrained; `class_weights` dead). **Fixed here:** requirements count (36→44), a value-equality assertion for `LOCAL_SCALE`/`GATE_PROB`/`CLASSIFIER_TEMPERATURE` in the round-trip, `pyro.clear_param_store()` in the `trained_model` fixture, api-doc `to_anndata` signature reconciled to the contract, and **7 knob-matrix corrections** — the Dirichlet variance is `∝ 1/(scale+1)` (not `1/scale`); `classifier_dropout` split out as a *separate* not-plumbed knob; gate=1 *formula* is testable now (only recovery blocked); `guide_temperature` needs a post-train test; scale-variance tests co-located; `kl_weight_max`/`reconstruction_loss_scale` and `use_enumeration` justifications sharpened. **Deferred (LOW/MED, logged in `REFACTOR_NOTES`):** the `group_singletons`-ordering guard, `predict` order hardening, round-trip exclusivity (Phase 6/7), and the stale `build_tcri_depgraph.py`. - **(PR3 audit — WORKFLOW, 3 lenses × adversarial verify, 8 agents):** behavior + doc-code lenses PASS, plan-contract FIX. **5 findings, all confirmed, all LOW/MED — zero behavior/correctness defect.** Behavior lens verified class bodies are byte-identical to the pre-split monolith (modulo the sanctioned rename), zero F821 undefined-names across all 5 files (every import header complete), and suite/smoke green. **Fixed:** the MED — explicit `__all__` per module (plan §Phase 3, line 279) was omitted → added to all 5 files (also resolves the two LOW surface-wording findings; surface now pinned to `{TCRIModel}`). **Deferred (auditor-recommended):** stale `c2p_mat` in the contract *generator* → Phase 8 with `diag.archetypes`. Suite 36 passed / 1 skipped. diff --git a/docs/contract/tcri_api_and_responsibilities.md b/docs/contract/tcri_api_and_responsibilities.md index 1146bab..7670ca5 100644 --- a/docs/contract/tcri_api_and_responsibilities.md +++ b/docs/contract/tcri_api_and_responsibilities.md @@ -317,7 +317,7 @@ Lazy GPU imports (never at module top — `import tcri` never touches a GPU lib; | `get_latent_representation(self, adata=None, *, indices=None, batch_size=None) -> np.ndarray` | Batched encode to the `(n_cells, n_latent)` posterior-mean latent. | | `predict(self, adata=None, *, batch_size=256, eps=1e-8) -> pd.DataFrame` | **(renamed from `get_cell_phenotype_probs`)** Per-cell phenotype-probability `DataFrame` (index = `adata.obs_names`, columns = phenotypes). Combines classifier logits with $\log p_{ct}$ (gate or additive), matching training (scvi/CellAssign idiom). **Reference the `use_logits=True` joint must reproduce at $T=1$** (§0.9, §7.1). Uses an **order-preserving loader** (shuffle=False / sequential sampler) and the registered `indices` field so ct-lookup and barcode labels cannot drift. | | `get_p_ct(self, *, guide_temperature=1.0) -> np.ndarray` | Return the learned `(ct_count, P)` posterior mean $m=\text{normalize}(q\_p\_ct\_raw)$. At the default `guide_temperature=1.0` this equals `uns[K.P_CT]` exactly. | -| `to_anndata(self, adata=None, *, latent_key="X_tcri", logits_key="X_tcri_logits", predictions_key="X_tcri_probabilities", label_key="tcri_phenotype") -> AnnData` | **(replaces the heavy `register_model`)** Thin writer of the **canonical minimum**: metadata + categories (from registry); `X_tcri` latent; **`obsm[K.X_LOGITS]` per-cell logits** (restored — the `use_logits=True` engine path hard-requires them); `predict()` probs + argmax hard labels; `p_ct` (+ `ct_to_cov`, `ct_to_c`, per-cell ct/cov arrays); **`local_scale`**, **`gate_prob`**, **`classifier_temperature`**. No manager stash; no other writes. | +| `to_anndata(self, adata=None, *, batch_size=256, compute_umap=False) -> AnnData` | **(replaces the heavy `register_model`; signature matches the frozen `_contract.pyi` as of PR4 — the canonical key names come from `_keys`, NOT per-call arguments)** Thin writer of the **canonical minimum**: metadata + categories (from registry); `X_tcri` latent; **`obsm[K.X_LOGITS]` per-cell logits** (restored — the `use_logits=True` engine path hard-requires them); `predict()` probs + argmax hard labels; `p_ct` (+ `ct_to_cov`, `ct_to_c`, per-cell ct/cov arrays); **`local_scale`**, **`gate_prob`**, **`classifier_temperature`**. No manager stash; no other writes. | > Relocated off the model: `plot_archetypes`→`diag.archetypes`; `plot_loss`→`diag.loss`. `boost_phenotype_prior`, `use_gate` remain internal. diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..154b8ad --- /dev/null +++ b/requirements.txt @@ -0,0 +1,23 @@ +# Verified-green dependency lock for tcri (scverse ecosystem). +# Generated + tested in a fresh Python 3.12 venv; `pytest` = 44 passed / 0 skipped. +# pyproject.toml holds the compatible ranges; this file pins the tested-recent set. +# Regenerate by creating a fresh venv and `pip install -e ".[test]"`. + +numpy==2.4.6 +pandas==3.0.3 +scipy==1.18.0 +scikit-learn==1.9.0 +anndata==0.13.1 +scanpy==1.12.2 +torch==2.13.0 +pyro-ppl==1.9.1 +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 +pytest-cov==7.1.0 diff --git a/tcri/_contract.pyi b/tcri/_contract.pyi index 1ac327f..f1435c1 100644 --- a/tcri/_contract.pyi +++ b/tcri/_contract.pyi @@ -26,19 +26,23 @@ class TCRIModel: def setup_anndata( cls, adata: AnnData, *, layer: Optional[str] = ..., clonotype_key: str = ..., phenotype_key: str = ..., - covariate_key: str = ..., batch_key: str = ..., + covariate_key: str = ..., batch_key: str = ..., **kwargs: Any, ) -> None: ... def train( self, max_epochs: int = ..., batch_size: int = ..., lr: float = ..., - reconstruction_loss_scale: float = ..., n_steps_kl_warmup: int = ..., + reconstruction_loss_scale: float = ..., n_steps_kl_warmup: int = ..., **kwargs: Any, ) -> None: ... def get_latent_representation( self, adata: Optional[AnnData] = ..., indices: Any = ..., batch_size: Optional[int] = ..., ) -> Any: ... - def predict(self, adata: Optional[AnnData] = ..., *, batch_size: int = ...) -> pd.DataFrame: ... + def predict( + self, adata: Optional[AnnData] = ..., *, batch_size: int = ..., eps: float = ..., + ) -> pd.DataFrame: ... def get_p_ct(self) -> Any: ... - def to_anndata(self, adata: AnnData, *, batch_size: int = ..., compute_umap: bool = ...) -> AnnData: ... + def to_anndata( + self, adata: Optional[AnnData] = ..., *, batch_size: int = ..., compute_umap: bool = ..., + ) -> AnnData: ... # ── preprocessing (pp) ─────────────────────────────────────────────────────── diff --git a/tcri/metrics/_metrics.py b/tcri/metrics/_metrics.py index 2f28b26..3a89fec 100644 --- a/tcri/metrics/_metrics.py +++ b/tcri/metrics/_metrics.py @@ -227,7 +227,7 @@ def clonotypic_entropy( ---------- adata : AnnData Registered object (see - :func:`~tcri.preprocessing._preprocessing.register_model`). + :meth:`~tcri.model._model.TCRIModel.to_anndata`). covariate : str Covariate value :math:`m` to condition on (a category of the registered covariate column). @@ -468,7 +468,7 @@ def phenotypic_entropy( ---------- adata : AnnData Registered object (see - :func:`~tcri.preprocessing._preprocessing.register_model`). + :meth:`~tcri.model._model.TCRIModel.to_anndata`). covariate : str Covariate value :math:`m` to condition on. point_estimate : bool, default True @@ -626,7 +626,7 @@ def mutual_information( ---------- adata : AnnData Registered object (see - :func:`~tcri.preprocessing._preprocessing.register_model`). + :meth:`~tcri.model._model.TCRIModel.to_anndata`). covariate : str Covariate value :math:`m` to condition on. temperature : float, default 1.0 @@ -845,7 +845,7 @@ def flux( ---------- adata : AnnData Registered object (see - :func:`~tcri.preprocessing._preprocessing.register_model`). + :meth:`~tcri.model._model.TCRIModel.to_anndata`). from_this, to_that : str The two covariate values to compare (e.g. ``"Pre-treatment"`` and ``"Post-treatment"``). diff --git a/tcri/model/_model.py b/tcri/model/_model.py index f25507a..8234694 100644 --- a/tcri/model/_model.py +++ b/tcri/model/_model.py @@ -3,7 +3,7 @@ The generative model, priors, classifier, and training plan live in sibling modules; this file holds only the high-level `BaseModelClass` API (`setup_anndata`, `__init__`, `train`, `get_latent_representation`, -`get_cell_phenotype_probs`, `get_p_ct`, ...): +`predict`, `to_anndata`, `get_p_ct`, ...): - :mod:`._module` -- Pyro model/guide (:class:`TCRIModule`) - :mod:`._priors` -- :class:`MixtureDirichlet`, :class:`VampPrior` @@ -52,13 +52,22 @@ class TCRIModel(BaseModelClass): def setup_anndata( cls, adata: AnnData, + *, layer: Optional[str] = None, clonotype_key: str = "unique_clone_id", phenotype_key: str = "phenotype_col", covariate_key: str = "timepoint", batch_key: str = "patient", **kwargs, - ): + ) -> None: + """Register clonotype/phenotype/covariate/batch/count fields with scvi. + + Registration only. Writes ``obs['indices']`` (scvi glue that the + training/validation steps consume via ``batch['indices']``) and records the + layer, but performs **no** analysis/label ``obs`` mutation and does **not** + stash the ``AnnDataManager`` in ``uns`` (the retired ``tcri_manager`` hack) — + learned outputs are written solely by :meth:`to_anndata`. + """ for col in [clonotype_key, phenotype_key, covariate_key, batch_key]: if col not in adata.obs: raise ValueError(f"{col} not in adata.obs!") @@ -81,12 +90,10 @@ def setup_anndata( adata_manager.registry["covariate_col"] = covariate_key adata_manager.registry["batch_col"] = batch_key cls.register_manager(adata_manager) - adata.uns["tcri_manager"] = adata_manager if layer is None: adata.uns.pop("tcri_layer", None) else: adata.uns["tcri_layer"] = layer - return adata def __init__( self, @@ -292,71 +299,136 @@ def get_p_ct(self): return self.module.get_p_ct().cpu().numpy() @torch.no_grad() - def get_cell_phenotype_probs( - self, adata=None, batch_size: int = 256, eps: float = 1e-8 - ) -> np.ndarray: - """ - Computes the cell-level phenotype probabilities in the same way as training. - - If ``self.module.gate_prob`` is set (i.e. ``use_gate`` is True), uses: - local_logits = gate_prob * cls_logits + (1 - gate_prob) * log(prior) - Otherwise uses the additive (Bayesian product) rule: - local_logits = cls_logits + log(prior) - - Parameters - ---------- - adata - If None, defaults to the AnnData used in training. - batch_size : int - Mini-batch size for data loader. - eps : float - Small epsilon for numerical stability in logs. - - Returns - ------- - probs : np.ndarray - Array of shape (n_cells, P) of phenotype probabilities. + def predict(self, adata=None, *, batch_size: int = 256, eps: float = 1e-8) -> pd.DataFrame: + """Per-cell phenotype-probability ``DataFrame`` (index ``adata.obs_names``, + columns = phenotypes) — the single source of phenotype probabilities. + + Combines classifier logits with ``log p_ct`` exactly as in training: if + ``self.module.gate_prob`` is set (``use_gate``), + ``gate_prob * cls_logits + (1 - gate_prob) * log(prior)``; otherwise the + additive rule ``cls_logits + log(prior)``. Renamed from + ``get_cell_phenotype_probs`` (which returned a bare ``ndarray``). The module + is put in ``eval`` mode so classifier dropout is off and the result is + deterministic; the sequential loader keeps the row order aligned to + ``obs_names``. """ adata = self._validate_anndata(adata) + self.module.eval() device = next(self.module.parameters()).device scdl = self._make_data_loader(adata=adata, batch_size=batch_size) - # The learned posterior p_ct -> shape (ct_count, P) p_ct = self.module.get_p_ct().to(device) - # Map each cell to its clonotype-covariate index -> shape (n_cells,) ct_array = self.module.ct_array.to(device) all_probs = [] current_idx = 0 - for tensors in scdl: x = tensors[REGISTRY_KEYS.X_KEY].to(device) b = tensors[REGISTRY_KEYS.BATCH_KEY].long().to(device) - this_batch_size = x.shape[0] - - # Which (clonotype, covariate) does each cell belong to? - ct_indices = ct_array[current_idx : current_idx + this_batch_size] - clone_cov_posterior = p_ct[ct_indices] # (batch_size, P) - - # Encode to get latent z + n = x.shape[0] + clone_cov_posterior = p_ct[ct_array[current_idx : current_idx + n]] z_loc, _, _ = self.module.encoder(x, b) - - # ----- 1) Compute classifier logits (same as training) ----- cls_logits = self.module.classifier(z_loc) - prior_log = torch.log(clone_cov_posterior + eps) if self.module.use_gate: local_logits = self.module.gate_prob * cls_logits + (1.0 - self.module.gate_prob) * prior_log else: local_logits = cls_logits + prior_log + all_probs.append(F.softmax(local_logits, dim=-1).cpu()) + current_idx += n - probs = F.softmax(local_logits, dim=-1) + probs = torch.cat(all_probs, dim=0).numpy() + phenotype_col = self.adata_manager.registry["phenotype_col"] + pheno_cats = self.adata.obs[phenotype_col].astype("category").cat.categories.tolist() + return pd.DataFrame(probs, index=adata.obs_names, columns=pheno_cats) - all_probs.append(probs.cpu()) - current_idx += this_batch_size + @torch.no_grad() + def to_anndata(self, adata=None, *, batch_size: int = 256, compute_umap: bool = False) -> AnnData: + """Write the model's learned state onto ``adata`` under the canonical + ``tcri_*`` keys (from :mod:`tcri._keys`) and return it. Replaces the old + ``preprocessing.register_model``; writes no manager stash. + + Writes — ``uns``: ``METADATA`` + covariate/clonotype/phenotype categories, + ``P_CT`` (posterior-mean ``p_ct``), ``CT_TO_COV``/``CT_TO_C``, per-cell + ``CT_ARRAY``/``COV_ARRAY``, ``LOCAL_SCALE``, ``GATE_PROB``, + ``CLASSIFIER_TEMPERATURE``; ``obsm``: ``X_TCRI`` latent, ``X_LOGITS``, + ``X_LOGPOSTERIOR``, ``X_PROBABILITIES`` (from :meth:`predict`); ``obs``: + ``PHENOTYPE`` argmax hard label. + """ + from .. import _keys as K - # Concatenate into final array of shape (n_cells, P) - return torch.cat(all_probs, dim=0).numpy() + adata = self._validate_anndata(adata) + self.module.eval() + device = next(self.module.parameters()).device + reg = self.adata_manager.registry + + # 1) metadata + category orders (order = training) -------------------- + meta = { + K.COVARIATE_COL: reg["covariate_col"], + K.CLONE_COL: reg["clonotype_col"], + K.PHENOTYPE_COL: reg["phenotype_col"], + K.BATCH_COL: reg["batch_col"], + } + adata.uns[K.METADATA] = meta + for col_key, cat_key in ( + (K.COVARIATE_COL, K.COVARIATE_CATEGORIES), + (K.CLONE_COL, K.CLONOTYPE_CATEGORIES), + (K.PHENOTYPE_COL, K.PHENOTYPE_CATEGORIES), + ): + adata.uns[cat_key] = adata.obs[meta[col_key]].astype("category").cat.categories.tolist() + + # 2) learned priors + per-cell index arrays -------------------------- + ct_arr = self.module.ct_array.cpu().numpy() + adata.uns[K.P_CT] = self.module.get_p_ct().cpu().numpy() + adata.uns[K.CT_TO_COV] = self.module.ct_to_cov.cpu().numpy() + adata.uns[K.CT_TO_C] = self.module.ct_to_c.cpu().numpy() + adata.uns[K.CT_ARRAY] = ct_arr + adata.uns[K.COV_ARRAY] = self.module.ct_to_cov.cpu().numpy()[ct_arr] + adata.uns[K.LOCAL_SCALE] = float(self.module.local_scale) + gp = self.module.gate_prob + adata.uns[K.GATE_PROB] = float(gp) if gp is not None else float("nan") + adata.uns[K.CLASSIFIER_TEMPERATURE] = float(self.module.classifier_temperature) + + # 3) latent mean ----------------------------------------------------- + adata.obsm[K.X_TCRI] = self.get_latent_representation( + adata=adata, batch_size=batch_size + ).astype("float32") + + # 4) per-cell logits + additive log-posterior (folds _compute_logits_and_prior) + loader = self._make_data_loader(adata=adata, batch_size=batch_size) + p_ct_t = self.module.get_p_ct().to(device) + ct_arr_t = self.module.ct_array.to(device) + logits_buf, prior_buf = [], [] + start = 0 + for tensors in loader: + x = tensors[REGISTRY_KEYS.X_KEY].to(device) + b = tensors[REGISTRY_KEYS.BATCH_KEY].long().to(device) + n = x.shape[0] + z_loc, _, _ = self.module.encoder(x, b) + logits_buf.append(self.module.classifier(z_loc).cpu()) + prior_buf.append(torch.log(p_ct_t[ct_arr_t[start : start + n]] + 1e-8).cpu()) + start += n + cls_logits = torch.cat(logits_buf).numpy().astype("float32") + prior_log = torch.cat(prior_buf).numpy().astype("float32") + adata.obsm[K.X_LOGITS] = cls_logits + adata.obsm[K.X_LOGPOSTERIOR] = cls_logits + prior_log + + # 5) probabilities (gate-aware, canonical) + argmax hard labels ------ + probs_df = self.predict(adata, batch_size=batch_size) + adata.obsm[K.X_PROBABILITIES] = probs_df.values.astype("float32") + adata.obs[K.PHENOTYPE] = pd.Categorical.from_codes( + probs_df.values.argmax(1), categories=list(probs_df.columns) + ) + + # 6) legacy compat keys — retired in Phase 6/7 with their readers ----- + adata.uns[K.LEGACY_CLONE_KEY] = meta[K.CLONE_COL] + adata.uns[K.LEGACY_PHENOTYPE_KEY] = K.PHENOTYPE + + if compute_umap: + import umap + adata.obsm[K.X_UMAP] = umap.UMAP(random_state=42).fit_transform(adata.obsm[K.X_TCRI]) + + return adata def boost_phenotype_prior( self, diff --git a/tcri/plotting/_plotting.py b/tcri/plotting/_plotting.py index cd0fce6..940a859 100644 --- a/tcri/plotting/_plotting.py +++ b/tcri/plotting/_plotting.py @@ -591,7 +591,7 @@ def tcri_boxplot(adata, function, groupby=None,ylabel="", splitby=None,figsize=( df.replace([np.inf, -np.inf], np.nan, inplace=True) df.dropna(inplace=True) if order == None: - order = df.groupby(["Phenotype"]).median(ylabel).sort_values(ylabel).index.tolist() + order = df.groupby("Phenotype")[ylabel].median().sort_values().index.tolist() fig,ax=plt.subplots(1,1,figsize=figsize) sns.stripplot(data=df,x="Phenotype",y=ylabel,s=s,hue=groupby,ax=ax,order=order, palette=palette) sns.boxplot(data=df,x="Phenotype",y=ylabel,ax=ax, color="#999999",order=order) @@ -617,7 +617,7 @@ def tcri_boxplot(adata, function, groupby=None,ylabel="", splitby=None,figsize=( df.dropna(inplace=True) fig,ax=plt.subplots(1,1,figsize=figsize) if order == None: - order = df.groupby(["Phenotype"]).median(ylabel).sort_values(ylabel).index.tolist() + order = df.groupby("Phenotype")[ylabel].median().sort_values().index.tolist() sns.boxplot(data=df,x="Phenotype",y=ylabel,ax=ax, hue=splitby,order=order,palette=palette) ax.set_ylim(0,max(df[ylabel] + 0.1)) ax.set_title(ylabel) diff --git a/tcri/preprocessing/_preprocessing.py b/tcri/preprocessing/_preprocessing.py index a45fac0..cafeb87 100644 --- a/tcri/preprocessing/_preprocessing.py +++ b/tcri/preprocessing/_preprocessing.py @@ -59,16 +59,7 @@ def _ascii_hist(samples, bins=25, width=40) -> str: return "\n".join(lines) -def register_phenotype_key(adata, phenotype_key, order=None): - assert phenotype_key in adata.obs, "Key {} not found.".format(phenotype_key) - if order==None: - adata.uns["tcri_unique_phenotypes"] = np.unique(adata.obs[phenotype_key].tolist()) - adata.uns["tcri_phenotype_key"] = phenotype_key -def register_clonotype_key(adata, tcr_key): - assert tcr_key in adata.obs, "Key {} not found.".format(tcr_key) - adata.uns["tcri_clone_key"] = tcr_key - adata.uns["tcri_unique_clonotypes"] = np.unique(adata.obs[tcr_key].tolist()) 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) @@ -84,130 +75,8 @@ def collapse_singleton(row): # ------------ helper to extract logits -------- # -@torch.no_grad() -def _compute_logits_and_prior(model, adata, batch_size=256, eps=1e-8): - device = next(model.module.parameters()).device - loader = model._make_data_loader(adata=adata, batch_size=batch_size) - ct_arr = model.module.ct_array.to(device) - p_ct = model.module.get_p_ct().to(device) - - logits_buf, prior_buf = [], [] - start = 0 - for tensors in tqdm(loader, desc="extracting logits", leave=False): - x = tensors[REGISTRY_KEYS.X_KEY].to(device) - b = tensors[REGISTRY_KEYS.BATCH_KEY].long().to(device) - n = x.size(0) - - z_loc, _, _ = model.module.encoder(x, b) - logits = model.module.classifier(z_loc) - prior_log = torch.log(p_ct[ct_arr[start:start+n]] + eps) - - logits_buf.append(logits.cpu()) - prior_buf.append(prior_log.cpu()) - start += n - - return (torch.cat(logits_buf).numpy().astype("float32"), - torch.cat(prior_buf).numpy().astype("float32")) # ------------ main routine -------------------- # -@torch.no_grad() -def register_model( - adata, model, - phenotype_prob_slot=K.X_PROBABILITIES, - phenotype_assignment_obs=K.PHENOTYPE, - latent_slot=K.X_TCRI, - batch_size=256, - store_logits=True, - store_logposterior=True, - compute_umap=False, - umap_n_neighbors=50, - umap_min_dist=1e-3, - umap_metric="euclidean", - umap_random_state=42, - umap_output_metric="euclidean", - clonotype_key="trb_unique", -): - print(f"{BOLD}{MAGENT}🔗 Registering TCRi model outputs …{RESET}") - - # 1) priors & arrays ------------------------------------------------- - adata.uns[K.P_CT] = model.module.get_p_ct().cpu().numpy() - adata.uns[K.CT_TO_COV] = model.module.ct_to_cov.cpu().numpy() - adata.uns[K.CT_TO_C] = model.module.ct_to_c.cpu().numpy() - adata.uns[K.LOCAL_SCALE] = model.module.local_scale - _ok("stored hierarchical priors") - for k in (K.P_CT,K.CT_TO_COV,K.CT_TO_C): - _info(f"uns['{k}']", np.shape(adata.uns[k])) - - # 2) metadata -------------------------------------------------------- - meta = { - "covariate_col": model.adata_manager.registry["covariate_col"], - "clone_col": model.adata_manager.registry["clonotype_col"], - "phenotype_col": model.adata_manager.registry["phenotype_col"], - "batch_col": model.adata_manager.registry["batch_col"], - } - adata.uns[K.METADATA] = meta - _ok("stored metadata dictionary") - - # categories - for key, col in (("covariate","covariate_col"), - ("clonotype","clone_col"), - ("phenotype","phenotype_col")): - cats = adata.obs[meta[col]].astype("category").cat.categories.tolist() - adata.uns[f"tcri_{key}_categories"] = cats - _info(f"uns['tcri_{key}_categories']", len(cats)) - - # per-cell ct / cov arrays - ct_arr = model.module.ct_array.cpu().numpy() - adata.uns[K.CT_ARRAY] = ct_arr - cov_arr = model.module.ct_to_cov.cpu().numpy()[ct_arr] - adata.uns[K.COV_ARRAY] = cov_arr - _ok("stored per-cell ct / cov indices") - - # 3) latent means ---------------------------------------------------- - z = model.get_latent_representation(batch_size=batch_size).astype("float32") - adata.obsm[latent_slot] = z - _ok("stored latent means") - _info(f"obsm['{latent_slot}']", z.shape) - - # 4) logits & log-posterior ----------------------------------------- - cls_logits, prior_log = _compute_logits_and_prior(model, adata, batch_size) - if store_logits: - adata.obsm[K.X_LOGITS] = cls_logits - _info("obsm['X_tcri_logits']", cls_logits.shape) - if store_logposterior: - adata.obsm[K.X_LOGPOSTERIOR] = cls_logits + prior_log - _info("obsm['X_tcri_logposterior']", cls_logits.shape) - _ok("computed logits & additive log-posterior") - - # 5) probabilities & hard labels ------------------------------------ - if phenotype_prob_slot not in adata.obsm: - from scipy.special import softmax - probs = softmax(cls_logits + prior_log, axis=1).astype("float32") - adata.obsm[phenotype_prob_slot] = probs - _info(f"obsm['{phenotype_prob_slot}']", probs.shape) - - adata.obs[phenotype_assignment_obs] = pd.Categorical.from_codes( - adata.obsm[phenotype_prob_slot].argmax(1), - categories=adata.uns[K.PHENOTYPE_CATEGORIES], - ) - _ok("stored probabilities and hard labels") - - # 6) optional UMAP --------------------------------------------------- - if compute_umap: - print(f"{CYAN}🗺️ computing UMAP …{RESET}") - reducer = umap.UMAP( - n_neighbors=umap_n_neighbors, min_dist=umap_min_dist, - metric=umap_metric, random_state=umap_random_state, - output_metric=umap_output_metric, - ) - adata.obsm["X_umap"] = reducer.fit_transform(z) - _info("obsm['X_umap']", adata.obsm["X_umap"].shape) - - register_phenotype_key(adata,phenotype_assignment_obs) - register_clonotype_key(adata,clonotype_key) - - print(f"{MAGENT}✨ All TCRi artefacts registered!{RESET}") - return adata def joint_distribution_posterior( @@ -235,7 +104,7 @@ def joint_distribution_posterior( f"({n_obs}). This happens when the function is called on a filtered " "AnnData view or subset: the 'tcri_*_array_for_cells' arrays in .uns " "remain in the original full-cell space while .obs/.obsm are subset, so " - "cell indices silently misalign. Re-run register_model(...) on the " + "cell indices silently misalign. Re-run model.to_anndata(...) on the " "filtered AnnData, or pass the full object and filter with `clones=`." ) diff --git a/tcri/utils/_utils.py b/tcri/utils/_utils.py index 2a6096b..d7d14d3 100644 --- a/tcri/utils/_utils.py +++ b/tcri/utils/_utils.py @@ -484,22 +484,7 @@ def _to_jsonable(x: Any) -> Any: pass return str(x) -def _pop_nonserializables(adata: "_ad.AnnData") -> Dict[str, Any]: - sidecar = {} - if "tcri_manager" in adata.uns: - sidecar["tcri_manager"] = "dropped (AnnDataManager is not serializable)" - adata.uns.pop("tcri_manager") - return sidecar - -def write_adata_safely(adata: "_ad.AnnData", path: str, *, compression: str = "gzip") -> None: - removed = {} - try: - removed = _pop_nonserializables(adata) - adata.write_h5ad(path, compression=compression) - finally: - # We intentionally do not restore removed manager objects back into .uns - # because they are session-bound and will be reconstructed on load. - pass + def _collect_setup_from_adata_or_model(adata: "_ad.AnnData", model: Any) -> Dict[str, Any]: setup: Dict[str, Any] = {} @@ -581,10 +566,12 @@ def save_tcri_session( _json.dump(setup, f, indent=2) paths["setup"] = _os.path.join(out_dir, SETUP_FILE) - # 4) Save sanitized AnnData (without tcri_manager) + # 4) Save the AnnData (plain h5ad; the manager stash is retired at setup_anndata) if save_adata: - write_adata_safely(adata, _os.path.join(out_dir, AD_FILE), compression=compression) - paths["adata"] = _os.path.join(out_dir, AD_FILE) + ad_path = _os.path.join(out_dir, AD_FILE) + adata.uns.pop(K.LEGACY_MANAGER, None) # defensive: never serialize a stray AnnDataManager + adata.write_h5ad(ad_path, compression=compression) + paths["adata"] = ad_path # 5) Meta / versions meta = { diff --git a/tests/conftest.py b/tests/conftest.py index 01c43ef..4b59322 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -118,11 +118,13 @@ def synthetic_adata(): @pytest.fixture(scope="session") def trained_model(synthetic_adata): - """TCRIModel fit for 50 epochs on synthetic_adata, with register_model applied.""" + """TCRIModel fit for 50 epochs on synthetic_adata, with to_anndata applied.""" _seed_all(0) + import pyro + pyro.clear_param_store() # own the process-global store (§5.2 cross-test contamination) + from tcri.model._model import TCRIModel - from tcri.preprocessing._preprocessing import register_model adata = synthetic_adata.copy() TCRIModel.setup_anndata( @@ -151,5 +153,5 @@ def trained_model(synthetic_adata): enable_progress_bar=False, enable_model_summary=False, ) - register_model(adata, model, clonotype_key="unique_clone_id") + model.to_anndata(adata) return model, adata diff --git a/tests/test_contract_conformance.py b/tests/test_contract_conformance.py index 6cfa858..458c474 100644 --- a/tests/test_contract_conformance.py +++ b/tests/test_contract_conformance.py @@ -20,9 +20,16 @@ PYI = Path(tcri.__file__).parent / "_contract.pyi" # contract key ("Namespace.func" / "TCRIModel.method") -> (module, dotted attr). -# EMPTY at PR0 — nothing has been migrated to the new surface yet. Each PR adds -# its landed functions here; the signature test then enforces live == contract. -IMPLEMENTED: dict[str, tuple[str, str]] = {} +# Each PR onboards its landed functions here; the signature test then enforces +# live == contract. PR4 landed the model→AnnData surface. +IMPLEMENTED: dict[str, tuple[str, str]] = { + "TCRIModel.setup_anndata": ("tcri.model._model", "TCRIModel.setup_anndata"), + "TCRIModel.train": ("tcri.model._model", "TCRIModel.train"), + "TCRIModel.get_latent_representation": ("tcri.model._model", "TCRIModel.get_latent_representation"), + "TCRIModel.predict": ("tcri.model._model", "TCRIModel.predict"), + "TCRIModel.get_p_ct": ("tcri.model._model", "TCRIModel.get_p_ct"), + "TCRIModel.to_anndata": ("tcri.model._model", "TCRIModel.to_anndata"), +} def _params_from_ast(a: ast.arguments): diff --git a/tests/test_model_setup.py b/tests/test_model_setup.py index 0275954..73053bb 100644 --- a/tests/test_model_setup.py +++ b/tests/test_model_setup.py @@ -26,9 +26,11 @@ def test_setup_anndata_defaults_to_x_matrix(): TCRIModel.setup_anndata(adata) - manager = adata.uns["tcri_manager"] + manager = TCRIModel._get_most_recent_anndata_manager(adata) assert manager.registry["setup_args"]["layer"] is None assert "tcri_layer" not in adata.uns + # Phase 4: the AnnDataManager is no longer stashed in uns (tcri_manager retired). + assert "tcri_manager" not in adata.uns np.testing.assert_array_equal( manager.get_from_registry(REGISTRY_KEYS.X_KEY), adata.X, @@ -42,7 +44,7 @@ def test_setup_anndata_explicit_none_uses_x_matrix_with_layers_present(): TCRIModel.setup_anndata(adata, layer=None) - manager = adata.uns["tcri_manager"] + manager = TCRIModel._get_most_recent_anndata_manager(adata) assert manager.registry["setup_args"]["layer"] is None assert "tcri_layer" not in adata.uns np.testing.assert_array_equal( @@ -57,7 +59,7 @@ def test_setup_anndata_default_uses_x_matrix_with_one_layer_present(): TCRIModel.setup_anndata(adata) - manager = adata.uns["tcri_manager"] + manager = TCRIModel._get_most_recent_anndata_manager(adata) assert manager.registry["setup_args"]["layer"] is None assert "tcri_layer" not in adata.uns np.testing.assert_array_equal( @@ -73,7 +75,7 @@ def test_setup_anndata_default_uses_x_matrix_with_layers_present(): TCRIModel.setup_anndata(adata) - manager = adata.uns["tcri_manager"] + manager = TCRIModel._get_most_recent_anndata_manager(adata) assert manager.registry["setup_args"]["layer"] is None assert "tcri_layer" not in adata.uns np.testing.assert_array_equal( @@ -88,7 +90,7 @@ def test_setup_anndata_accepts_single_explicit_layer(): TCRIModel.setup_anndata(adata, layer="counts") - manager = adata.uns["tcri_manager"] + manager = TCRIModel._get_most_recent_anndata_manager(adata) assert manager.registry["setup_args"]["layer"] == "counts" assert adata.uns["tcri_layer"] == "counts" np.testing.assert_array_equal( @@ -113,7 +115,7 @@ def test_setup_anndata_accepts_explicit_layer_when_multiple_layers_present( TCRIModel.setup_anndata(adata, layer=layer) - manager = adata.uns["tcri_manager"] + manager = TCRIModel._get_most_recent_anndata_manager(adata) assert manager.registry["setup_args"]["layer"] == layer assert adata.uns["tcri_layer"] == layer np.testing.assert_array_equal( diff --git a/tests/test_model_smoke.py b/tests/test_model_smoke.py index 1e968c1..9171b87 100644 --- a/tests/test_model_smoke.py +++ b/tests/test_model_smoke.py @@ -59,6 +59,7 @@ def test_model_construct_train_predict(synthetic_adata): p_ct = model.get_p_ct() assert p_ct.ndim == 2 and p_ct.shape[1] == P - probs = model.get_cell_phenotype_probs() + probs = model.predict() assert probs.shape == (n_cells, P) - np.testing.assert_allclose(probs.sum(axis=1), 1.0, atol=1e-4) + assert list(probs.index) == list(adata.obs_names) # order-preserving, labelled + np.testing.assert_allclose(probs.values.sum(axis=1), 1.0, atol=1e-4) diff --git a/tests/test_session_round_trip.py b/tests/test_session_round_trip.py index 1f1941d..633ab90 100644 --- a/tests/test_session_round_trip.py +++ b/tests/test_session_round_trip.py @@ -1,20 +1,100 @@ -"""Round-trip save/load test (Notion T4 + T11). +"""Session save/load round-trip + the Phase-4 model→AnnData streamline contract. -Guards against the PyTorch 2.6+ weights_only=True regression that silently -broke pyro param store loading. If load fails silently, _ensure_pyro_posterior_params -re-initializes q_p_ct_raw to a uniform 1/P matrix; we detect that here. +Locks three things: + 1. ``to_anndata`` writes exactly the canonical key set, including the Phase-4 + additions (logits / gate / classifier-temperature / local-scale), and no + ``tcri_manager`` stash. + 2. ``setup_anndata`` is registration-only — no analysis/label ``obs`` mutation. + 3. save → load reproduces ``p_ct`` + latent + ``predict`` from the *reloaded model* + (this also guards the PyTorch ``weights_only=True`` regression that silently + broke pyro param-store loading; a failed load re-inits ``q_p_ct_raw`` to a + uniform 1/P matrix, which we detect via row variance). """ import contextlib import io import numpy as np import pyro +import pytest +from tcri import _keys as K from tcri.utils._utils import load_tcri_session, save_tcri_session -def test_session_round_trip(trained_model, tmp_path): +@pytest.fixture +def fresh_trained_model(synthetic_adata): + """A model trained fresh inside this test so it OWNS the process-global pyro + param store for the save (§5.2 — loading/training a second model in one process + clobbers ``q_p_ct_raw``, so the shared session-scoped ``trained_model`` store is + unsafe for a round-trip that recomputes from the reloaded model).""" + from tcri.model._model import TCRIModel + + pyro.clear_param_store() + adata = synthetic_adata.copy() + TCRIModel.setup_anndata( + adata, clonotype_key="unique_clone_id", phenotype_key="phenotype_col", + covariate_key="timepoint", batch_key="patient", + ) + model = TCRIModel( + adata, n_latent=8, n_hidden=16, n_layers=1, classifier_n_layers=1, + classifier_hidden=16, K=3, n_pseudo_obs=3, + ) + with contextlib.redirect_stdout(io.StringIO()): + model.train(max_epochs=50, batch_size=64, enable_progress_bar=False, + enable_model_summary=False) + model.to_anndata(adata) + return model, adata + +CANONICAL_UNS = [ + K.METADATA, K.P_CT, K.CT_TO_COV, K.CT_TO_C, K.CT_ARRAY, K.COV_ARRAY, + K.LOCAL_SCALE, K.GATE_PROB, K.CLASSIFIER_TEMPERATURE, + K.COVARIATE_CATEGORIES, K.CLONOTYPE_CATEGORIES, K.PHENOTYPE_CATEGORIES, +] +CANONICAL_OBSM = [K.X_TCRI, K.X_LOGITS, K.X_LOGPOSTERIOR, K.X_PROBABILITIES] + + +def test_to_anndata_writes_canonical_set(trained_model): + """to_anndata writes the full canonical key set incl. the Phase-4 additions, + with the scalar knobs equal to the model's configured values; no manager stash.""" model, adata = trained_model + for k in CANONICAL_UNS: + assert k in adata.uns, f"to_anndata did not write uns[{k}]" + for k in CANONICAL_OBSM: + assert k in adata.obsm, f"to_anndata did not write obsm[{k}]" + assert K.PHENOTYPE in adata.obs, "to_anndata did not write the hard-label obs" + assert "tcri_manager" not in adata.uns, "manager stash was not retired" + # Phase-4 scalar knobs: written == the model's configured value (not just finite) + assert adata.uns[K.LOCAL_SCALE] == float(model.module.local_scale) + assert adata.uns[K.CLASSIFIER_TEMPERATURE] == float(model.module.classifier_temperature) + gp = model.module.gate_prob + if gp is None: + assert np.isnan(adata.uns[K.GATE_PROB]) # no gate -> nan sentinel (serializable) + else: + assert adata.uns[K.GATE_PROB] == float(gp) + + +def test_setup_anndata_leaves_analysis_obs_untouched(synthetic_adata): + """setup_anndata is registration-only: it may add the 'indices' glue column but + writes no analysis/label obs (that is exclusively to_anndata's job).""" + from tcri.model._model import TCRIModel + + adata = synthetic_adata.copy() + obs_before = set(adata.obs.columns) + TCRIModel.setup_anndata( + adata, clonotype_key="unique_clone_id", phenotype_key="phenotype_col", + covariate_key="timepoint", batch_key="patient", + ) + new_cols = set(adata.obs.columns) - obs_before + # registration glue is allowed ('indices' + scvi's internal '_scvi_*' columns); + # any OTHER new obs column would be analysis/label leakage. + non_glue = {c for c in new_cols if c != "indices" and not c.startswith("_scvi")} + assert not non_glue, f"setup_anndata mutated analysis obs: {non_glue}" + assert K.PHENOTYPE not in adata.obs, "setup_anndata must not write hard labels" + assert "tcri_manager" not in adata.uns + + +def test_session_round_trip(fresh_trained_model, tmp_path): + model, adata = fresh_trained_model out_dir = tmp_path / "session" save_tcri_session(model, adata, str(out_dir)) @@ -23,23 +103,29 @@ def test_session_round_trip(trained_model, tmp_path): buf = io.StringIO() with contextlib.redirect_stdout(buf): - _, loaded = load_tcri_session(str(out_dir)) + loaded_model, loaded = load_tcri_session(str(out_dir)) - np.testing.assert_allclose( - adata.obsm["X_tcri"], loaded.obsm["X_tcri"], atol=1e-5 - ) + # 1) serialization: the saved AnnData survives the h5ad round-trip np.testing.assert_array_equal( - adata.obsm["X_tcri_probabilities"], loaded.obsm["X_tcri_probabilities"] - ) - np.testing.assert_allclose( - adata.uns["tcri_p_ct"], loaded.uns["tcri_p_ct"], atol=1e-5 + adata.obsm[K.X_PROBABILITIES], loaded.obsm[K.X_PROBABILITIES] ) + np.testing.assert_allclose(adata.obsm[K.X_TCRI], loaded.obsm[K.X_TCRI], atol=1e-6) + np.testing.assert_allclose(adata.uns[K.P_CT], loaded.uns[K.P_CT], atol=1e-6) + assert "tcri_manager" not in loaded.uns + # 2) pyro store restored, not silently re-initialized to a uniform prior store = pyro.get_param_store() assert "q_p_ct_raw" in store, "q_p_ct_raw missing from pyro store after load" q = store["q_p_ct_raw"].detach().cpu().numpy() - row_var = q.var(axis=-1) - assert row_var.mean() > 1e-6, ( - f"q_p_ct_raw rows are uniform (mean row var {row_var.mean():.2e}); " - "pyro load silently failed (PyTorch weights_only regression)" + assert q.var(axis=-1).mean() > 1e-6, ( + "q_p_ct_raw rows are uniform; pyro load silently failed (weights_only regression)" + ) + + # 3) the RELOADED model reproduces what to_anndata wrote before save + np.testing.assert_allclose(adata.uns[K.P_CT], loaded_model.get_p_ct(), atol=1e-4) + np.testing.assert_allclose( + adata.obsm[K.X_TCRI], loaded_model.get_latent_representation(loaded), atol=1e-4 + ) + np.testing.assert_allclose( + adata.obsm[K.X_PROBABILITIES], loaded_model.predict(loaded).values, atol=1e-4 )