refactor(PR4): model→AnnData streamline (to_anndata, predict; retire tcri_manager)

[nceglia]

Stacked on #32 (PR3). Base = refactor/pr3-model-split.

Model→AnnData streamline (Phase 4) — behavior change. Retires the uns['tcri_manager'] hack, folds register_model into model.to_anndata, renames get_cell_phenotype_probs → predict (labelled DataFrame), and rewrites the round-trip gate. Full suite 44 passed in the pinned venv.

Environment (prerequisite, this round)

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 in legacy tcri_boxplot (groupby.median() positional numeric_only).

Model surface

• setup_anndata: keyword-only, returns None, registration-only — no uns['tcri_manager'] stash, no analysis/label obs mutation.
• predict (was get_cell_phenotype_probs): labelled DataFrame (obs_names × phenotypes), eval() for deterministic inference, gate-aware.
• to_anndata (replaces register_model): writes the canonical set incl. the new GATE_PROB + CLASSIFIER_TEMPERATURE (+ P_CT, ct maps, per-cell arrays, LOCAL_SCALE, X_TCRI/X_LOGITS/X_LOGPOSTERIOR/X_PROBABILITIES from predict, argmax labels).
• Deleted register_model / register_phenotype_key / register_clonotype_key / _compute_logits_and_prior; write_adata_safely / _pop_nonserializables (h5ad write inlined). The 6 TCRIModel methods onboarded into the frozen contract — conformance now enforces the model surface.

Round-trip gate

Rewritten test_session_round_trip: canonical write-set · setup-obs invariant · reloaded model reproduces p_ct/latent/predict (a function-scoped fixture owns the process-global pyro param store, §5.2).

⚠️ Model correctness finding (answering the "is the classifier working" ask)

The phenotype classifier is never trained — cls_logits is computed in model() but never enters any ELBO factor, and the training plan only touches it under no_grad (verified: weight Δ = 0.0, isolated gate=1.0 recovery = chance). End-to-end predict recovers perfectly on a perfect dataset, but entirely via the prior p_ct. phenotype_weights/class_weights is dead for the same reason. Per direction, this is tracked in the agenda (Model knob-test matrix + correctness-debt note) for a dedicated fix PR before Phase 6 — not touched here.

Tests

• Rewritten round-trip test (CI).
• dev/real_data_to_anndata.py — LOCAL-only (gitignored, not CI): builds the 50 largest patient-specific yost clones (trb_unique = trb+patient; 7682 cells / 10 patients) → setup→train→to_anndata.
• Model knob-test matrix added to the agenda: every constructor/train knob → its mathematically-correct input→output test or a justification (with target PRs).

Deferred (logged in REFACTOR_NOTES)

Legacy tcri_clone_key / tcri_phenotype_key / X_tcri_phenotypes retire in Phase 6/7 (live readers in not-yet-refactored metrics/plotting). Classifier-training fix → dedicated PR.

Audit

A workflow audit ran on this branch — 4 lenses (doc↔code · behavior/correctness · plan/contract · a dedicated knob-test-plan lens), each finding adversarially verified (23 agents). 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.

Fixed in this branch: value-equality assertions for LOCAL_SCALE/GATE_PROB/CLASSIFIER_TEMPERATURE; pyro.clear_param_store() in the trained_model fixture; requirements count; api-doc to_anndata reconciled to the frozen contract; and 7 knob-matrix corrections — the Dirichlet draw-variance is ∝ 1/(scale+1) (not 1/scale); classifier_dropout is a separate un-plumbed knob (never forwarded to the classifier); gate=1's formula is testable now (only recovery is blocked); guide_temperature needs a post-train test; scale-variance tests co-located; kl_weight_max/reconstruction_loss_scale/use_enumeration justifications sharpened.

Deferred (LOW/MED, logged in REFACTOR_NOTES): the group_singletons-ordering guard, predict order-preservation hardening, round-trip key-exclusivity (Phase 6/7), and a stale depgraph doc-generator.

Stacked so the diff shows only PR4's changes. Living agenda + audit log + knob-test matrix: docs/contract/REFACTOR_AGENDA.md.

🤖 Generated with Claude Code