refactor(PR3): split model/_model.py into scvi-style modules + explicit __all__

docs/contract/REFACTOR_AGENDA.md

@@ -34,7 +34,7 @@ tracker + running diary for the whole refactor. The detailed spec lives in `tcri
 | 0 | Contract freeze + CI scaffolding | ✅ | none | — | conformance green |
 | 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 |
+| 3 | Model module split | ✅ | low | 1 | model/pyro tests green |
 | 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 |
@@ -116,7 +116,17 @@ Template per PR: **Goal · Status · What happened · Issues & fixes · Added 
 - **Streamline:** shrinks preprocessing (−127) and plotting (−225) meaningfully ahead of the Phase 3/4/7 splits. Removed the one import this PR orphaned (`cosine_similarity` — sole user was the deleted `classify_phenotypes`); other module-top imports left for the file-split phases (conservative; the PR1 audit already handled the utils ones).
 - **Usability:** removes broken/dead public entry points (`probability_distribution` self-recursion, `bayesian_mutual_information` bad kwarg, `polar_plot` undefined-name) from the surface so nobody trips on them.
 
-## PR 3 — Model module split  ·  ☐ todo
+## PR 3 — Model module split  ·  ✅ done (branch `refactor/pr3-model-split`)
+- **Goal:** split the 1074-line `model/_model.py` into scvi-style sibling files (`_model` + `_module` + `_priors` + `_classifier` + `_training`), rename `c2p_mat → clone_phenotype_prior`. Mechanical; **no behavior change**.
+- **What happened:** verified up front that **no code outside `model/` references any moved internal** (only `TCRIModel` is imported externally). Extracted the 7 top-level defs via `ast.get_source_segment` (formatting-preserving) into the target files along the clean dependency DAG `_classifier`/`_priors` (leaf) → `_module` → `_training` → `_model`. Each module declares an explicit `__all__` (a Phase-3 deliverable — plan §Phase 3), so `tcri.model.*` is now pinned to exactly the public API the frozen contract promises — `{TCRIModel}` — and the incidental third-party re-export leaks the old `import *` exposed (`pyro`/`dist`/`Encoder`/`KMeans`/… 17 names, none tcri-defined, none referenced anywhere) are no longer surfaced. Applied the `c2p_mat → clone_phenotype_prior` rename with a word-boundary regex (13 sites; left the unrelated `c2p_torch` local and the module buffer `clone_phen_prior` untouched). Dropped 3 provably-dead top-level imports surfaced by the per-file import rebuild (`setup_anndata_dsp`, `cosine_similarity`, the `torch.distributions` `Categorical/Dirichlet/MixtureSameFamily` trio — all uses were `dist.`-prefixed pyro). File sizes: `_model` 462, `_module` 326, `_training` 154, `_priors` 147, `_classifier` 21.
+- **Issues & fixes:** (1) my first smoke silently imported a **stale `site-packages/tcri`** copy (a script's dir, not the repo, leads `sys.path`) and failed in `setup_anndata` on an old-copy/scvi mismatch — a red herring; forcing the repo copy onto the path, the smoke passes. (The stale install is an env-hygiene note, not a code issue — pytest already uses the repo copy, which is why the suite validates the split.) (2) The train path was **entirely uncovered** by the suite (`trained_model` fixture defined but unused) — so the split was only import-verified. Fixed by adding a real end-to-end smoke test.
+- **Added:** ✅ explicit `__all__` in all 5 model modules (`_model`=`{TCRIModel}`; siblings export their own class(es)). ✅ `tests/test_model_smoke.py` — construct → train (2 epochs) → `get_latent_representation` / `get_p_ct` / `get_cell_phenotype_probs` (asserts shapes + prob normalization), plus asserts the `clone_phenotype_prior` rename landed and `build_archetypes` returns centers **and** labels. Runs in ~1s inside the suite.
+- **Removed (hard bar):** n/a — PR3 is a structural split, not a removal PR. `ml.plot_loss`/`ml.plot_archetypes` stay on `TCRIModel` until `diag/` exists (Phase 8); no Phase-2 style deletions here.
+- **Test opportunities:** ✅ closed the biggest gap (model construct/train/query now covered). The rewritten `test_session_round_trip` (Phase 4) will extend this to save/load.
+- **Streamline:** the split makes Phase 4 (model→AnnData) and Phase 5 (engine) tractable — the pyro module, priors, classifier, and training plan are now editable in isolation.
+- **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 5 — Engine consolidation  ·  ☐ todo
 ## PR 6 — Metric-API consolidation  ·  ☐ todo
@@ -135,3 +145,5 @@ _(dated entries; what was audited, findings, actions)_
 - **(PR0 ✅):** agenda + removal ledger established; standing-audit checklist defined. Contract frozen (27 fns) + conformance guardrail live. Full suite 26 passed / 1 skipped, zero regressions.
 - **(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).
+- **(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.

tcri/model/_classifier.py

@@ -0,0 +1,23 @@
+"""Phenotype classifier head for the TCRI Pyro module (:class:`~tcri.model._module.TCRIModule`)."""
+import torch.nn as nn
+
+__all__ = ["PhenotypeClassifier"]
+
+
+class PhenotypeClassifier(nn.Module):
+    def __init__(self, n_latent, classifier_hidden, P, num_layers=3, dropout_rate=0.1, temperature=1.0):
+        super(PhenotypeClassifier, self).__init__()
+        layers = []
+        input_dim = n_latent
+        for _ in range(num_layers):
+            layers.append(nn.Linear(input_dim, classifier_hidden))
+            layers.append(nn.ReLU())
+            layers.append(nn.Dropout(dropout_rate))
+            input_dim = classifier_hidden
+        layers.append(nn.Linear(classifier_hidden, P))
+        self.mlp = nn.Sequential(*layers)
+        self.temperature = temperature  # Add temperature parameter
+
+    def forward(self, x):
+        logits = self.mlp(x)
+        return logits / self.temperature