refactor: land the completed refactor in main (PRs #31–#37)

.gitignore

@@ -144,3 +144,6 @@ example/*.pdf
 .DS_Store
 
 
+
+# local-only dev harnesses (real-data tests; NOT CI)
+dev/

CLAUDE.md

@@ -0,0 +1,64 @@
+# tcri — contributor rules
+
+Single-cell TCR+RNA information-theory metrics on scvi-tools / pyro / scanpy.
+
+## The three contracts
+
+This repo is governed by three frozen contracts. All are machine-checked; a failing
+conformance test means **stop and decide**, not "adjust the contract until it passes."
+
+| | freezes | manifest | prose | test |
+|---|---|---|---|---|
+| **API contract** | the public *interface* | `tcri/_contract.pyi` | `docs/contract/tcri_api_and_responsibilities.md` | `tests/test_contract_conformance.py` |
+| **Model contract** | the generative *mathematics* | `tcri/model/_model_contract.py` | `docs/contract/MODEL_CONTRACT.md` | `tests/test_model_contract_conformance.py` |
+| **Metrics contract** | what the *metrics compute* | `tcri/tools/_metrics_contract.py` | `docs/contract/METRICS_CONTRACT.md` | `tests/test_metrics_contract_conformance.py` |
+
+### Model integrity (read before touching `tcri/model/`)
+
+The model implements **Supplementary Note 1** (`tcri_supplementary_methods_04_30_26.pdf`)
+— the source of truth. Changing its mathematics means: adding/removing a stochastic
+site, changing a distribution family or plate, altering the ELBO or the phenotype
+surrogate, or changing what a prior is scaled by (α on eq 1, β on eq 2).
+
+**Update the model contract FIRST, then the code.** Cite the note equation and state
+what changes in the joint distribution. Then make the code agree.
+
+**Never loosen the manifest to make a conformance failure go away.** That silently
+rewrites the model the package claims to implement. A failure is either an intended
+model change (update the contract deliberately, as a reviewed model change) or a
+regression (fix the code).
+
+If you are an AI agent and a model change appears necessary, **surface the contract
+implication to the user** rather than editing the manifest to fit your change.
+
+Accepted departures from the note live in `SANCTIONED_DEVIATIONS`
+(`_model_contract.py`) with a rationale, mirrored in `MODEL_CONTRACT.md`. Anything
+not listed there that departs from the note is a defect.
+
+`docs/contract/METHODS_CONFORMANCE.md` is the eq-by-eq code map + deviation history.
+
+### Metric integrity (read before touching `tcri/tools/`)
+
+The entropies and mutual information are frozen by the **metrics contract**. Changing
+a definition means changing what every published number means. Same rule: update
+`_metrics_contract.py` + `METRICS_CONTRACT.md` first, then the code.
+
+The conformance test pins numeric identities, the keystone being
+`I(c;φ) = H(c) − Σ_φ P(φ)·H[P(c|φ)]` — it ties the entropy and MI families together so
+neither can be redefined alone.
+
+**Note 1's eqs 3–4 are mistranscribed** (they weight by the marginal, making them
+cross-entropies, and eq 4's label is wrong). The code is correct — proven by the
+decomposition above, which the literal equations violate by producing a *negative*
+mutual information. These are recorded in `SOURCE_ERRATA`; do **not** "fix" the code to
+match them.
+
+## Working agreement
+
+- **`docs/contract/REFACTOR_AGENDA.md` is the living tracker** — read it before
+  starting, write a diary entry after each PR, and run the Standing Audit in it.
+- **Removal is a hard bar.** Delete dead code rather than keeping it "just in case";
+  git has it. Tick the Removal Ledger.
+- **Never read the `example/` notebooks.** They are disposable *outputs* of the
+  refactor, never an input — no caller census, no "is-it-used" checks.
+- Run tests with the pinned venv: `MPLBACKEND=Agg .venv/bin/python -m pytest tests/ -q`.

docs/contract/METHODS_CONFORMANCE.md

@@ -0,0 +1,116 @@
+# Methods Conformance — code ↔ Supplementary Note 1
+
+> **This is the eq-by-eq map + deviation history.** The *enforced* contract is
+> `docs/contract/MODEL_CONTRACT.md` (prose) + `tcri/model/_model_contract.py`
+> (manifest), checked by `tests/test_model_contract_conformance.py`. Model math
+> changes require updating that contract **first**.
+
+Maps the TCRi generative model in **Supplementary Note 1: Methods for Information
+theoretic metrics for single cell RNA and T-cell receptor sequencing**
+(`tcri_supplementary_methods_04_30_26.pdf`) to the implementation, and records
+every known deviation. The PDF is the source of truth; this file is the living
+conformance record (update it whenever the model changes).
+
+Code: `tcri/model/_module.py` (`TCRIModule.model`/`.guide`), `tcri/model/_priors.py`
+(`VampPrior`, `MixtureDirichlet`), `tcri/model/_classifier.py`
+(`PhenotypeClassifier`), `tcri/model/_model.py` (`TCRIModel`, `.predict`).
+
+## Symbols
+
+| Note | Meaning | Code |
+|---|---|---|
+| `ω_c` | clonotype-level phenotype dist | `p_c` (sample site `"p_c"`) |
+| `ϕ_m` | covariate-level phenotype dist | `p_ct` (sample site `"p_ct"`); `get_p_ct()` |
+| `z_i` | continuous latent embedding | `z` (sample site `"latent"`) |
+| `z^ϕ_i` | **discrete phenotype latent** | not sampled — replaced by the surrogate (below) |
+| `x_i` | gene expression | `x` (sample site `"obs"`) |
+| `f_cls` | classifier `R^L → R^P` (η_cls) | `self.classifier` |
+| `π` | gating weight | `gate_prob` (default **0.5**) |
+| `α` | global Dirichlet scale | `global_scale` |
+| `β` | local Dirichlet scale | `local_scale` |
+| `γ` | surrogate KL weight | `phenotype_kl_weight` (default 1.0) |
+| `g(i)` | covariate-group of cell `i` | `ct_array[i]` (clone×covariate index) |
+| `h(m)` | clonotype of group `m` | `ct_to_c[m]` |
+
+## Generative model
+
+| Eq | Note | Code (`_module.py::model`) | Status |
+|---|---|---|---|
+| 1 | `ω_c ~ (1/B_c) Σ_b Dir(α ψ_b)` | plate `"clonotypes"` → `MixtureDirichlet(weights, global_scale * mixture_concentration)`, sampled `"p_c"`. `ψ_b` = archetype centroids (`build_archetypes`). | ✅ α = `global_scale` (**[G]** fixed) |
+| 2 | `ϕ_m \| ω_h(m) ~ Dir(β ω_h(m))` | plate `"ct_plate"` → `conc_ct = clamp(local_scale * p_c[ct_to_c])`, sampled `"p_ct"` | ✅ β = `local_scale` |
+| 3 | `z_i ~ (1/B_z) Σ_k q(z\|u_k)` | `VampPrior.get_mixture()` (mixture of encoder-posteriors at learnable pseudo-inputs), sampled `"latent"` | ✅ |
+| 4 | `l_i = f_cls(z_i)`, `ℓ_i = π l_i + (1-π) log ϕ_g(i)`; `z^ϕ_i ~ Cat(softmax(ℓ_i))` | `cls_logits = classifier(z)`; `ell = gate_prob*cls_logits + (1-gate_prob)*log_phi`; discrete `z^ϕ` **not** sampled — see surrogate | ◐ via surrogate (below) |
+| 5 | `x_i ~ ZINB(g'_i, r_i, μ_i)` | `DecoderSCVI` → `ZeroInflatedNegativeBinomial(gate, total_count, logits)`, sampled `"obs"` | ✅ (scaled — **[E]**) |
+
+## Variational family (eq 6) — `_module.py::guide`
+
+- `q(ω_c) = Dir(λ_c)` — `q_p_c_raw` param → `conc_c_guide = clamp(global_scale * q_p_c_sharp)`. α = `global_scale`.
+- `q(ϕ_m) = Dir(λ'_m)` — `q_p_ct_raw` param → `conc_ct_guide = clamp(local_scale * q_p_ct_sharp)`.
+- `q(z_i\|x_i) = N(μ_i, σ_i²)` — `encoder(x, batch)` → `Normal(z_loc, z_scale)`, sampled `"latent"`.
+- `q(z^ϕ_i\|z_i, ϕ) = Cat(softmax(ℓ_i))` — represented by the surrogate, not an explicit categorical sample.
+
+## ELBO (eq 7) and the surrogate ("Inference Details")
+
+Eq 7 is the standard SVI ELBO (`Trace_ELBO`; `TraceEnum_ELBO` when `use_enumeration`),
+`E[log p(x|z)] + E[log p(Ω,Φ,z,z^ϕ)] − E[log q]`, maximized by Adam.
+
+The note replaces the discrete `z^ϕ` terms with a surrogate:
+
+> `L_new = L#  +  γ Σ_i KL(probs_i ‖ ϕ_g(i))`,  `probs_i = softmax(ℓ_i)`
+
+where `L#` is eq 7 with the `z^ϕ` terms removed and `γ>0`. The KL is a **penalty**
+(the note "penalizes misalignment"), i.e. the objective is to *minimize* it. Pyro's SVI
+**maximizes** the ELBO / log-joint, so the penalty enters the factor with a **minus
+sign** — `−γ·KL`. (Reading the note's `+γ·ΣKL` as something to maximize would push
+`probs` *away* from `ϕ`; the sign below is the one that realizes the note's intent.)
+Implemented in `model()`'s `"data"` plate:
+
+```python
+phi     = p_ct[ct_idx].detach()          # ϕ_g(i), detached alignment target
+ell     = gate_prob*cls_logits + (1-gate_prob)*log_phi   # ℓ_i  (eq 4)
+probs   = softmax(ell)
+pheno_kl = (probs * (log(probs) - log_phi)).sum(-1)      # KL(probs ‖ ϕ)
+pyro.factor("phenotype_alignment", -phenotype_kl_weight * pheno_kl)
+```
+
+- `ct_idx = ct_array[indices]` uses **global** cell indices (threaded in via
+  `_get_fn_args_from_batch`), never the local pyro plate index — indexing with the
+  local index scrambles the per-cell target across shuffled minibatches.
+- Optimum of the surrogate is `f_cls → log ϕ + const` (distinct per clone), i.e. the
+  classifier learns to predict the clonotype-informed phenotype from expression.
+- `predict()` applies the same `ℓ_i` rule with `z_loc` (encoder mean, dropout off).
+
+## In-silico perturbation (eqs 8–12) — **[F] not implemented**
+
+`I_j = Σ_p |ϕ̄_p^(0) − ϕ̃_p^(j)|` (L1 shift after zeroing gene `j`). Additive feature;
+no code path yet.
+
+## Deviations
+
+| id | deviation | severity | status |
+|---|---|---|---|
+| A | classifier had no ELBO gradient (missing factor) | HIGH | **fixed** — `pyro.factor("phenotype_alignment", …)` |
+| A2 | surrogate target indexed by local plate idx → scrambled labels → f_cls collapse | HIGH | **fixed** — global `indices` threaded into `model()`/`guide()`; the `indices=None` path now `assert`s instead of silently falling back |
+| B | `gate_prob` default was `None`; note sets π=0.5 | LOW | **fixed** — default `0.5` (typed `Optional[float]`) |
+| C | `classifier_dropout` constructed but not passed to `PhenotypeClassifier` | LOW | **fixed** — plumbed |
+| D | `class_weights`/`log_class_weights` — not in the note; was dead (computed + plumbed through 3 signatures, never read) | LOW | **fixed** — removed (with `phenotype_weights`) from `_model`/`_module`/`_training` |
+| H | dead per-cell `encoder(x)` forward in `model()` (result discarded; the VampPrior carries its own encoder) | INFO | **fixed** — removed |
+| G | α (`global_scale`) not applied to the clonotype prior (eq 1) in `model()`; concentration = normalized archetype centroid (sum≈1, U-shaped), so the prior was far more diffuse than `Dir(α·ψ_b)` and scaled inconsistently with the guide `q(ω_c)` | MED | **fixed** — `expanded_conc = global_scale * centroids` (eq 1); classifier recovery unchanged (1.000), suite green |
+| E | `reconstruction_loss_scale` down-weights ZINB vs eq-7 full weight | MED | **resolved** — default raised 1e-3 → 1e-2; real-data library ratio 1.40 → 0.99 (recovery/latent unchanged). The original ~6× over-generation was mostly the phantom optimizer shrinking the decoder. |
+| F | in-silico perturbation (eqs 8–12) not implemented | — | deferred — additive feature |
+
+**Training-only deviations from eq 7 (intentional, documented here):**
+- **KL warmup + z-only scope.** `UnifiedTrainingPlan` ramps `kl_weight` over `n_steps_kl_warmup`, and it scales only the `latent` (z) KL — the two Dirichlet KLs (`p_c`, `p_ct`) are unscaled. A standard annealing schedule; symmetric (no correctness bug) but not part of eq 7's full-weight KL.
+- **`num_particles`** on `UnifiedTrainingPlan` is honored only on the enumeration path (`TraceEnum_ELBO`); the default `Trace_ELBO` uses 1 MC particle regardless.
+
+A/A2/B/C/D/H were fixed in the model PR that introduced this file. **G and E are now
+resolved too** — α is applied to the eq-1 prior, and `reconstruction_loss_scale` was
+re-measured and recalibrated to `1e-2` (real-data library ratio 1.40 → 0.99). Both
+change fitted results, so runs are not comparable across them. **F** (in-silico
+perturbation) remains out of scope for this release.
+
+A further training-only deviation was found and removed: a **second torch Adam over all
+module parameters**, installed by overriding scvi's deliberate no-op `configure_optimizers`
+shim. It stepped after `SVI.step()` had zeroed the gradients, so weight decay degenerated
+to a scale-free `~lr·sign(p)` shrink (networks held ~2.4× small), and `train(lr=)` never
+reached Pyro's optimizer. See `optimizer_weight_decay` in the model contract.

docs/contract/METRICS_CONTRACT.md

@@ -0,0 +1,106 @@
+# Metrics Contract — what the numbers mean
+
+Freezes the **information-theoretic metrics**: the two entropies and mutual
+information over a clone × phenotype joint.
+
+| | freezes | manifest | prose | test |
+|---|---|---|---|---|
+| API contract | the public *interface* | `tcri/_contract.pyi` | `tcri_api_and_responsibilities.md` | `test_contract_conformance.py` |
+| Model contract | the *generative mathematics* | `tcri/model/_model_contract.py` | `MODEL_CONTRACT.md` | `test_model_contract_conformance.py` |
+| **Metrics contract** | **what the metrics compute** | `tcri/tools/_metrics_contract.py` | this file | `test_metrics_contract_conformance.py` |
+
+**Why separate from the model contract.** The two are verified by different means. The
+model contract *traces* `model()`/`guide()` and inspects sample sites, plates and
+distribution families. Metrics are pure functions of a joint table, so they are pinned
+by **numeric identities** — uniform → log₂(k), independent → MI 0, and the
+entropy/MI decomposition. Folding them together would force one mechanism to do a job
+it is bad at.
+
+Source of truth: **Supplementary Note 1**, "Entropy" section (eqs 2–4) — with the
+errata below.
+
+**Governance: update this file and the manifest FIRST, then the code.** A failing
+conformance test means the *meaning of a published number* changed. Never relax an
+identity to make it pass.
+
+## Definitions (all in **bits**, log base 2)
+
+### `clonotypic_entropy` — one value per **phenotype**
+
+```
+H[P(c|φ)] = − Σ_c P(c|φ) log₂ P(c|φ)
+```
+
+How spread a phenotype is across clones. **Support-only**: clones with zero mass in
+that column are dropped *before* renormalizing — no epsilon clip, which would fabricate
+uniform mass on absent clones and inflate H toward 1. Normalizer `log₂(#supported
+clones)`, or `log₂(n_clones_ref)` when supplied. Empty column → **NaN**.
+
+### `phenotypic_entropy` — one value per **clone**
+
+```
+H[P(φ|c)] = − Σ_φ P(φ|c) log₂ P(φ|c)
+```
+
+Plasticity vs commitment of a clone. All P phenotypes are in the sum with `0·log0 := 0`.
+Normalizer `log₂(P)`. A clone with zero mass → **NaN**, never reindexed to zeros (which
+would report a spurious `H=1` for a clone that was never observed).
+
+### `mutual_information` — one value per joint
+
+```
+I(c;φ) = Σ_{c,φ} P(c,φ) log₂( P(c,φ) / (P(c)·P(φ)) )
+```
+
+Default `normalize_mode="min"` → `I / min(H(c), H(φ))`, the coefficient of constraint.
+`"average"` → `I / (½(H(c)+H(φ)))`. **`min` is the default because the `average`
+denominator scales with `log₂(C)` and is therefore not comparable across groups with
+different clone counts.**
+
+## Enforced identities
+
+| identity | what it catches |
+|---|---|
+| uniform over k → `log₂(k)`, normalized `1.0` | a wrong log base or normalizer |
+| all mass on one outcome → `0` | sign/normalization errors |
+| zero-mass clone/phenotype → **NaN** | the spurious-`H=1` reindexing regression |
+| support-only normalization | an epsilon clip creeping back in |
+| independent joint → `I = 0` | a broken MI |
+| `I(c;φ) = I(φ;c)`, `I ≥ 0` | transpose/sign errors |
+| permutation joint → normalized `I = 1` | a wrong denominator |
+| **`I(c;φ) = H(c) − Σ_φ P(φ)·H[P(c|φ)]`** | **redefining either family alone** |
+
+That last one is the keystone: it ties entropy and MI together, so you cannot change
+one without breaking it.
+
+## Errata in Supplementary Note 1 (the code is correct)
+
+The note's eqs 3–4, read literally, do **not** match the implementation — and the note
+is the one that is wrong. Recorded here so nobody "fixes" the code to match a typo.
+
+1. **Eq 3** reads `H(p(c|φ)) = − Σ_c p(c) log p(c|φ)` — it weights by the **marginal**
+   `p(c)` while taking the log of the **conditional**. That is a cross-entropy, not an
+   entropy.
+2. **Eq 4** is labelled `H(p(c))` but its right-hand side sums over φ and uses `p(φ|c)`,
+   so the label is wrong; it also weights by the marginal.
+3. The **prose** introduces both as "the entropy of the marginal distributions", but the
+   equations are conditionals.
+
+**Why the code is right.** Mutual information must satisfy
+`I(c;φ) = H(c) − E_φ[H(c|φ)]`. On a test joint with true MI **0.288703**:
+
+- the implemented conditional entropy reproduces it **exactly** (0.288703);
+- the note's literal formula yields **−0.345883** — a *negative* mutual information,
+  which is impossible.
+
+The literal equations are inconsistent with the note's own MI, so they cannot be what
+was intended. `test_note_literal_formula_would_break_the_decomposition` pins this.
+
+## Sanctioned extensions (the note does not specify these)
+
+- **bits / log₂** — the note writes an unspecified `log`.
+- **`normalized=True`** — divide by the maximum-entropy value so results land in [0,1].
+- **`n_clones_ref`** — fix the clonotypic normalizer across groups; without it each
+  group normalizes by its own supported-clone count and the values are not comparable.
+- **`n_samples>0`** — return mean/sd/HDI over posterior draws. The plug-in entropy is
+  ≥ the posterior mean (Jensen), so the two are reported as distinct quantities.