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/CLAUDE.md b/CLAUDE.md new file mode 100644 index 0000000..911bcc1 --- /dev/null +++ b/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`. diff --git a/docs/contract/METHODS_CONFORMANCE.md b/docs/contract/METHODS_CONFORMANCE.md new file mode 100644 index 0000000..5f76437 --- /dev/null +++ b/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. diff --git a/docs/contract/METRICS_CONTRACT.md b/docs/contract/METRICS_CONTRACT.md new file mode 100644 index 0000000..7a613b8 --- /dev/null +++ b/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. diff --git a/docs/contract/MODEL_CONTRACT.md b/docs/contract/MODEL_CONTRACT.md new file mode 100644 index 0000000..43c3b26 --- /dev/null +++ b/docs/contract/MODEL_CONTRACT.md @@ -0,0 +1,191 @@ +# TCRI Model Contract (FROZEN) + +**The model this package implements is Supplementary Note 1** (`tcri_supplementary_methods_04_30_26.pdf`). +This document is the prose contract; `tcri/model/_model_contract.py` is its +machine-checkable form; `tests/test_model_contract_conformance.py` enforces it. + +Sibling of the API contract: `tcri/_contract.pyi` freezes the public *interface*, +this freezes the *mathematics*. `docs/contract/METHODS_CONFORMANCE.md` is the +eq-by-eq code map and deviation history. + +--- + +## THE RULE + +> **Changing the model's mathematics requires updating this contract *first*.** + +Concretely — adding/removing a stochastic site, changing a distribution family or +plate, altering the ELBO or the surrogate, or changing what a prior is scaled by: + +1. **Update the contract first** — this file *and* `_model_contract.py`, citing the + note equation and stating what changes in the joint distribution. +2. **Then change the code** so `test_model_contract_conformance` passes again. +3. **If the note itself is superseded**, say so explicitly here (with the new + reference). The note is the source of truth; the contract tracks it. + +**Never** make a conformance failure disappear by loosening the manifest to match +whatever the code now does. That silently rewrites the model the package claims to +implement, which is exactly what this guardrail exists to prevent. A failure means +*stop and decide*: is this an intended model change (update the contract) or a +regression (fix the code)? + +This applies to human and AI contributors alike. If you are an agent and a model +change seems necessary, surface the contract implication to the user rather than +editing the manifest to fit. + +--- + +## The generative model + +`p(Ω, Φ, z, x) = Π_c p(ω_c) · Π_m p(ϕ_m|ω_h(m)) · Π_i p(z_i)·p(z^ϕ_i|z_i,ϕ_g(i))·p(x_i|z_i)` + +| eq | site | distribution | plate | meaning | +|---|---|---|---|---| +| 1 | `p_c` | `MixtureDirichlet` | `clonotypes` (c=1..C) | `ω_c ~ (1/B_c) Σ_b Dir(α·ψ_b)` — clonotype-level phenotype distribution over archetypes `ψ_b` | +| 2 | `p_ct` | `Dirichlet` | `ct_plate` (m=1..M) | `ϕ_m \| ω_h(m) ~ Dir(β·ω_h(m))` — covariate-level, hierarchical under its clonotype | +| 3 | `latent` | `MixtureSameFamily` (VampPrior) | `data` (i=1..N) | `z_i ~ (1/B_z) Σ_k q(z\|u_k)` over learnable pseudo-inputs | +| 4 | *(surrogate)* | — | `data` | `ℓ_i = π·f_cls(z_i) + (1−π)·log ϕ_g(i)`; `z^ϕ_i ~ Cat(softmax(ℓ_i))` | +| 5 | `obs` | `ZeroInflatedNegativeBinomial` | `data` | `x_i ~ ZINB(g'_i, r_i, μ_i)` from the scVI decoder | + +**Scales are semantics, not tuning.** α (`global_scale`) scales eq 1's concentration +and β (`local_scale`) scales eq 2's. Dropping either changes the prior's *shape* — +with concentration entries < 1 a Dirichlet becomes U-shaped (mass at the simplex +corners), the opposite of a prior peaked at the archetype — and desynchronizes the +prior from the guide. Both are asserted by the conformance test. + +## The variational family (eq 6) + +`q(Ω,Φ,z|x) = Π_c Dir(ω_c|λ_c) · Π_m Dir(ϕ_m|λ'_m) · Π_i q(z_i|x_i;η_enc) · Π_i q(z^ϕ_i|z_i,ϕ;η_cls)` + +| site | distribution | learnable | +|---|---|---| +| `p_c` | `Dirichlet` | `q_p_c_raw` (λ_c), scaled by α | +| `p_ct` | `Dirichlet` | `q_p_ct_raw` (λ'_m), scaled by β | +| `latent` | `Normal(μ_i, diag(σ_i²))` | encoder `η_enc` | + +`z^ϕ` is **not** sampled — see the surrogate below. A categorical `q(z^ϕ)` site +reappearing in the guide changes the objective and is rejected by the test. + +## The objective + +The ELBO (eq 7) is `E_q[log p(x|z)] + E_q[log p(Ω,Φ,z,z^ϕ)] − E_q[log q]`, maximized +by SVI (Adam, reparameterized continuous latents). + +**The surrogate** ("Inference Details") replaces the discrete `z^ϕ` terms: + +> `L_new = L# + γ·Σ_i KL(probs_i ‖ ϕ_g(i))`, `probs_i = softmax(ℓ_i)` + +The KL is a **penalty** on misalignment (the note "penalizes misalignment"), so it is +*minimized*. Pyro's SVI **maximizes** the log-joint, therefore the factor is registered +with a **minus sign**: + +```python +pyro.factor("phenotype_alignment", -phenotype_kl_weight * kl) # −γ·KL ≤ 0 +``` + +A positive factor would push `probs` *away* from `ϕ`. The test asserts the factor's +log-value is ≤ 0 and non-zero. + +**This term is the only thing that trains `f_cls`.** Without it the classifier's +logits never enter the ELBO and it receives no gradient (recovery sits at chance). + +**The alignment target must use global cell indices.** `ϕ_g(i) = p_ct[ct_array[indices]]` +where `indices` are *global* cell ids threaded in via `_get_fn_args_from_batch`. The +pyro data-plate index is local (`0..batch_size−1`); using it scrambles each cell's +target across shuffled minibatches and collapses `f_cls` to a constant. + +## Gating (π) + +`gate_prob` = π ∈ (0,1), default **0.5** per the note. Endpoints are contract-tested: +π=1 ⇒ `predict()` is the pure classifier; π=0 ⇒ the pure clonotype prior; +π=`None` ⇒ the additive rule `f_cls + log ϕ`. + +--- + +## Sanctioned deviations + +Accepted departures from the note. **Anything not listed here that departs from the +note is a defect.** Keys match `SANCTIONED_DEVIATIONS` in `_model_contract.py` (the +test asserts they stay in sync). + +| key | departure | rationale | +|---|---|---| +| `E_reconstruction_loss_scale` | eq 7 weights `E[log p(x\|z)]` at 1; the `obs` site is scaled by `reconstruction_loss_scale` (default **`1e-2`**) | β-VAE-style reweighting, **re-measured and recalibrated** — see below. | + +### On `reconstruction_loss_scale` (deviation [E], resolved) + +Re-measured after the phantom optimizer was removed. Posterior-predictive library +ratio (simulated ÷ observed; 1.00 is calibrated): + +| scale | real yost (2259×1000) | synthetic (3000×60) | +|---|---|---| +| `1e-3` (old default) | **1.40** | 1.00 | +| **`1e-2` (new default)** | **0.99** | 1.00 | +| `1e-1` | 1.00 | 1.00 | +| `1.0` (eq-7 full weight) | — | **0.91** (over-corrects) | + +Dropout fraction matches observed at every setting (0.870 vs ~0.873 on real data). +Classifier recovery (1.000) and latent separation (7.19) are **unchanged** across +`1e-3`→`1e-1`, so the recalibration costs nothing. + +The originally-reported **~6× over-generation was mostly the phantom second +optimizer** shrinking the decoder (see `optimizer_weight_decay`); removing it took the +ratio 6× → 1.40, and this default closes the remainder. Note the synthetic data could +not detect this — only the real 1000-gene, 87%-dropout data discriminates. + +Three inconsistent defaults (`_model.train`=1e-3, `_module`=1e-3, `_training`=1e-2) +were unified to `1e-2`. +| `kl_warmup_z_only` | `kl_weight` anneals only the `latent` KL; the Dirichlet KLs are unscaled | Standard annealing; training-only, not part of eq 7. | +| `num_particles_enumeration_only` | `num_particles` applies only on the `TraceEnum_ELBO` path | Default `Trace_ELBO` uses 1 MC particle. | +| `F_perturbation_not_implemented` | in-silico perturbation (eqs 8–12) absent | Additive feature; explicitly out of scope for this release. | +| `optimizer_weight_decay` | the SVI optimizer applies Adam weight decay (default `1e-4`) to the network parameters | The note fixes the *objective* (eq 7 + the surrogate), not the optimizer. Applied inside Pyro's optimizer so it acts on the ELBO gradients. See the note below. | + +### On the optimizer (history worth keeping) + +Weight decay was previously applied by a **second** `torch.optim.Adam` over every +module parameter, installed by overriding `UnifiedTrainingPlan.configure_optimizers`. +That override replaced scvi's *deliberate no-op shim* — scvi returns +`Adam([self._dummy_param])` purely to advance Lightning's step counter — and it ran +**after** `SVI.step()` had already stepped and **zeroed the gradients**. + +Stepping Adam on zero gradients is not a no-op. The weight-decay term becomes the +entire gradient (`g = wd·p`), and Adam's normalization `g/√(g²)` then strips its +magnitude, so the update degenerates to **≈ `lr·sign(p)`** — a *scale-free* shrink of +roughly `lr` per step rather than proportional L2. Measured in isolation: a weight of +0.1 is driven to 5e-5 within 1000 steps, where true L2 (`(1−lr·wd)^n`) would leave it +at 0.9998. In the fitted model SVI pushed back, but the equilibrium sat at **~2.4× +smaller network weights** (encoder 0.107 vs 0.240). + +It also meant `train(lr=...)` never reached the optimizer that fits the model — Pyro +always used scvi's hard-coded `1e-3`, and `lr` only set the shrink rate. + +The override is removed; `lr`/`weight_decay`/`betas`/`eps` now go to Pyro's optimizer +via `optim_kwargs`, which is where the original intent belongs. `lr` is consequently a +**live** knob for the first time (verified: it moves recovery, weight scale, and the +final ELBO), so fits are not comparable across this change. + +## What the conformance test checks + +`tests/test_model_contract_conformance.py` traces the live `model()`/`guide()`: + +- every declared site exists with the right distribution family, plate, event-dim, + observed-flag — **and no undeclared site exists** (an extra site changes the joint); +- the guide's variational family + learnable params (λ_c, λ'_m), and that `z^ϕ` is not sampled; +- α scales eq 1's concentration; +- **eq 2 is hierarchical**: `p_ct`'s concentration is asserted *elementwise* to equal + `clamp(β·(ω_c[ct_to_c] + eps))` against the sampled `p_c` **from the same trace** — + pinning the scale, the source tensor, and the index map h(m) together; +- the surrogate factor is a *negative*, non-zero KL; +- **the alignment target is verified behaviorally**: on a minibatch whose global + indices differ from the local plate positions, the traced factor must equal the + surrogate recomputed under the *global* map and must differ from the local one; +- π endpoints reduce `predict()` to classifier / prior; +- every sanctioned deviation is documented here. + +**Assertions are behavioral, not textual.** Two earlier drafts of this test were +defeated in an adversarial audit: a scalar "concentration totals ≈ β" check passed +even with the clonotype→covariate hierarchy severed (every simplex row totals 1, so +any tensor under any permutation satisfies it), and a source-grep for +`ct_array[indices]` was defeated by routing the same wrong lookup through +`index_select`. Prefer assertions computed from a live trace over ones that read +source text or scalar summaries. diff --git a/docs/contract/REFACTOR_AGENDA.md b/docs/contract/REFACTOR_AGENDA.md index 09825c5..a59897e 100644 --- a/docs/contract/REFACTOR_AGENDA.md +++ b/docs/contract/REFACTOR_AGENDA.md @@ -33,31 +33,31 @@ 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 | -| 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 | -| 8 | `diag/` seeding | ☐ | low-med | 4,5 | PPC columns | -| 9 | PGM→docs; utils finalize | ☐ | low | 1,8 | import green sans daft | -| 10 | Notebook rewrite (fresh) | ☐ | low | 4–8 | nbmake tutorial | -| 11 | Public API + scverse CI | ☐ | low-med | all | ecosystem checklist | +| 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 | +| 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 | +| 8 | `diag/` seeding | ✅ | low-med | 4,5 | PPC columns | +| 9 | PGM→docs; utils finalize | ✅ | low | 1,8 | import green sans daft | +| 10 | Public API + scverse CI | ☐ | low-med | all | ecosystem checklist | ## Removal Ledger (the hard bar — every one MUST end deleted) Tick only when the symbol is gone from source AND `__all__`/imports AND `import tcri` is green. -**Phase 2 (dead / out-of-scope):** -- [ ] `pp.get_latent_embedding` · [ ] `pp.group_small_clones` · [ ] `pp.register_probability_columns` -- [ ] `pp.remove_meaningless_genes` · [ ] `pp.gene_entropy` · [ ] `pp.classify_phenotypes` -- [ ] `pl.polar_plot` · [ ] `pl.probability_distribution` · [ ] `pl.bayesian_mutual_information` -- [ ] `metrics._ent` · [ ] `tl.clone_fraction` · [ ] `metrics.dkl` (→ `_distance.kl_divergence`) -- [ ] `ut.probabilities` (+ its `_plotting.py:18` import, same PR) · [ ] `SankeyNode.hex_to_rgb` +**Phase 2 (dead / out-of-scope):** ✅ ALL DELETED (PR2) — 14 symbols, 384 lines, `import tcri` green. +- [x] `pp.get_latent_embedding` · [x] `pp.group_small_clones` · [x] `pp.register_probability_columns` +- [x] `pp.remove_meaningless_genes` · [x] `pp.gene_entropy` · [x] `pp.classify_phenotypes` +- [x] `pl.polar_plot` · [x] `pl.probability_distribution` · [x] `pl.bayesian_mutual_information` +- [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) +- [x] `tcri_clone_key` / `tcri_phenotype_key` / obsm `X_tcri_phenotypes` — **DONE**. The deferral said "still read by not-yet-refactored `metrics`/`plotting`", but those were rewritten in PR6/PR7; the last reader was a raw literal in `pp.clone_size`, now migrated to `uns[METADATA][CLONE_COL]` with a clear error. `to_anndata` no longer writes the shims and the three `LEGACY_*` constants are deleted (`LEGACY_MANAGER` stays — `save_tcri_session` still pops it defensively). **Phase 5/6 (consolidated away — delete WITH replacement, never before):** - [ ] `pp.joint_distribution_posterior` (→ unified `joint_distribution`) · [ ] `metrics._mi_from_joint` (→ `_mutual_information`) @@ -75,7 +75,7 @@ Tick only when the symbol is gone from source AND `__all__`/imports AND `import - [ ] `ut.build_nested_tcri_pgm` (→ `docs/`) · [ ] `ut.draw_tcri_pgm_nested` (→ `docs/`) · [ ] `daft` runtime dep **Phase 3/9 (model/utils cleanup):** -- [ ] `_ascii_hist` (+ all `graph=`/ASCII paths) · [ ] `ml.plot_loss` (→ `diag.loss`) · [ ] `ml.plot_archetypes` (→ `diag.archetypes`) +- [x] `_ascii_hist` (dead: zero callers) · [x] `ml.plot_loss` (→ `diag.loss`) · [x] `ml.plot_archetypes` (→ `diag.archetypes`) --- @@ -106,18 +106,106 @@ Template per PR: **Goal · Status · What happened · Issues & fixes · Added - **Usability:** internal only this step. - **`K.*` migration (done):** replaced **85** canonical key literals with `K.*` across preprocessing/metrics/plotting/utils via a verified script (model=0; only legacy keys there); added `test_no_canonical_key_literals` guard — none remain. `dkl` reassigned (dead `metrics.dkl`→Phase 2; `flux` inner→Phase 6). Legacy keys left as literals until their removal phases. **PR1 COMPLETE — 35 passed / 1 skipped.** -## PR 2 — Safe deletions · ☐ todo -_(diary to be filled — this is a REMOVAL PR; the ledger Phase-2 block must be fully ticked)_ - -## PR 3 — Model module split · ☐ todo -## PR 4 — Model→AnnData streamline · ☐ todo -## PR 5 — Engine consolidation · ☐ todo -## PR 6 — Metric-API consolidation · ☐ todo -## PR 7 — Plotting split + pl twins · ☐ todo -## PR 8 — diag/ seeding · ☐ todo -## PR 9 — PGM→docs; utils finalize · ☐ todo -## PR 10 — Notebook rewrite · ☐ todo -## PR 11 — Public API + scverse CI · ☐ todo +## PR 2 — Safe deletions · ✅ done (branch `refactor/pr2-safe-deletions`) +- **Goal:** delete the Phase-2 dead / out-of-scope symbols outright (they go to the trash, not to examples). Zero behavior change to the kept surface; `import tcri` stays green. +- **What happened:** verified **every** target has zero in-package call-sites (grep for calls/imports, plus a precise `NAME(` call-site check for the three ambiguous ones — `_ent`, module-level `dkl`, `probabilities` — all 0), then deleted via AST-span (top-level `FunctionDef` line ranges; the `SankeyNode.hex_to_rgb` method inside its `ClassDef`; the dead `from ..utils._utils import probabilities` line). **14 symbols, 384 lines removed** across metrics/plotting/sankey/preprocessing/utils. `import tcri` green; **full suite 35 passed / 1 skipped** (unchanged — nothing referenced them). +- **Why safe (not deferred):** the "dkl" refs that remain are the **string** distance-metric name in `flux` + the `_distance` registry, not the deleted module-level `dkl` function (`flux` keeps its own inner `dkl_func` → Phase 6). The `polar_plot`/`probability_distribution` "callers" were self-referential (own docstring examples / the self-recursion bug). `probabilities` was imported into plotting but never called. +- **Added:** n/a (pure removal PR). +- **Removed (hard bar):** ✅ all 14 Phase-2 ledger items ticked — `pp.get_latent_embedding`, `pp.group_small_clones`, `pp.register_probability_columns`, `pp.remove_meaningless_genes`, `pp.gene_entropy`, `pp.classify_phenotypes`, `pl.polar_plot`, `pl.probability_distribution`, `pl.bayesian_mutual_information`, `metrics._ent`, `tl.clone_fraction`, `metrics.dkl`, `ut.probabilities` (+ import), `SankeyNode.hex_to_rgb`. Confirmed gone: no `def` remains, no explicit re-export/`__all__` names them. +- **Test opportunities:** none new (removal); the existing suite is the regression gate and stayed green. +- **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 · ✅ 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 · ✅ 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 · ✅ COMPLETE (`tests/test_model_knobs.py`, 32 tests) + +**Two layers, and the split is the point.** + +- **WIRING** — does the value actually reach the object it configures? This layer was + added after `lr` sat marked "hooked up" for months while never reaching Pyro's + optimizer. The model still converged, so *every behavioral test passed*. A silently + ignored knob is invisible to convergence testing; assert the plumbing directly. +- **BEHAVIOR** — the mathematically-correct input→output assertion. + +| Knob | Wiring test | Behavior test | Status | +|---|---|---|---| +| `n_latent` | `module.n_latent` | latent width == n_latent | ✅ | +| `n_pseudo_obs` | `vamp_prior.pseudo_inputs.shape[0]` | — | ✅ | +| `K` | `mixture_concentration.shape[0]` | `centers.shape[0] == K` | ✅ | +| `n_hidden`,`n_layers` | encoder param count scales | — | ✅ | +| `global_scale` (α) | `module.global_scale` | scales the eq-1 prior concentration ([G] fix) | ✅ | +| `local_scale` (β) | `module.local_scale` | p_ct draw variance == p(1−p)/(β+1) | ✅ | +| `prior_temperature` | `module.prior_temperature` | T>1 raises `clone_phen_prior` row-entropy | ✅ | +| `guide_temperature` | `module.guide_temperature` | T<1 lowers `get_p_ct()` row-entropy | ✅ | +| `gate_prob` (π) | `module.gate_prob` | π=0 ⇒ softmax(log φ); π=1 ⇒ softmax(f_cls) | ✅ | +| `classifier_temperature` | `module.classifier_temperature` | logits(T=2) == logits(T=1)/2 | ✅ | +| `classifier_dropout` | `classifier.mlp[2].p` | — | ✅ (was ⛔ not-plumbed) | +| `classifier_hidden`,`classifier_n_layers` | layer widths / Linear count | — | ✅ (was ⛔ untrained) | +| `kl_weight_max` | `module.kl_weight_max` | ramp ceiling | ✅ | +| `guide_init_scale` | `module.guide_init_scale` | — | ✅ | +| `phenotype_kl_weight` (γ) | `module.phenotype_kl_weight` | classifier recovery (model contract) | ✅ | +| `lr`,`weight_decay`,`betas`,`eps` | **`plan.optim.pt_optim_args`** | recovery/ELBO move with lr | ✅ **(was DEAD)** | +| `n_steps_kl_warmup` | `plan.n_steps_kl_warmup` | `kl_weight` ramps 0→max, monotonic | ✅ | +| `reconstruction_loss_scale` | `module.reconstruction_loss_scale` | — (deviation [E]) | ✅ wiring | +| `max_epochs`,`patience` | `trainer.max_epochs`, EarlyStopping `.patience` | — | ✅ | +| `batch_size` | dataloader batch shape | `predict` invariant (float32 tol) | ✅ | +| `use_enumeration` | selects `TraceEnum_ELBO` vs `Trace_ELBO` | — | ✅ | +| `phenotype_weights` | — | — | **REMOVED** (was dead; deleted in Phase 1a) | + +**Findings from the run:** no new dead knobs. The two initial failures were both +*test* bugs, not code bugs — scvi installs `LoudEarlyStopping` (not `EarlyStopping`), +and `predict` differs by ~1.2e-07 across batch sizes (float32 kernel paths, not a +logic error). Previously-⛔ classifier knobs are unblocked now that the classifier +trains. `n_steps_kl_warmup`'s step-vs-epoch semantics remain open as **DUX-2**. + +## PR 5 — Engine consolidation · ✅ done (branch `refactor/pr5-engine`) +- **Goal:** build the unified `joint_distribution` engine (`tools/` + `_compute/`) that every metric will consume — the substrate. **Additive** this PR: the old `joint_distribution`/`joint_distribution_posterior` stay until Phase 6 migrates the metrics onto the new engine. HIGH risk (math-heavy). +- **What happened:** created `_compute/_xp.py` (torch-first device seam — CPU / torch-CUDA, lazy GPU import, `asnumpy` boundary; grafiti parity but torch-first since the draws are torch), `_compute/_joint.py::_joint_draws` (the `[S, n_clones, P]` core), `_compute/_reduce.py` (batched entropy/MI, **bits/log2 default** per the user, float64 accumulators), `tools/_joint.py` (the DataFrame wrapper), re-exported `tcri.joint_distribution`. Onboarded `tl.joint_distribution` into the contract (added `device`). +- **Engine invariants implemented (§7.1):** temper the base **once** (`T==1` is the *exact* identity — no eps round-trip, so both closed-form identities are exact/near-exact); **draw over all ct rows once then slice per covariate** (shared-draw invariant → draw-count == `n_samples` regardless of #covariates); `n_samples=0` deterministic; `use_logits=True` folds per-cell logits with `log(base)` **gate-aware** and scatter-adds per clone (matching `predict`); `weighted` scales clone rows by the **ct-keyed** cell count (fixes the old clone-indexed `Counter` bug); torch-seeded Dirichlet (`random_state`). +- **Verified (the gate — `tests/test_tools/test_joint.py`, all green):** `use_logits=False,n=0,T=1 == uns[P_CT]` restricted (**dev 0.0**); `use_logits=True,n=0,T=1 == predict` per-clone aggregation (**dev 4e-8**; compared to the *frozen* `X_PROBABILITIES` so it's store-independent); `n=0` bit-identical; `n>0` seeded-reproducible + Dirichlet mean→base; weighting == clone cell counts; `covariate=None` slice == per-covariate call (shared draw); provenance JSON-serializable. Plus `test_reduce` (entropy of uniform == `log2(P)` bits; MI independent==0, coupled==1 bit), gate-aware combine (direct), T≠1 temper, and the subset/`local_scale` guards. Suite **66 passed**. +- **Added:** ✅ `_compute/{_xp,_joint,_reduce}.py` ✅ `tools/{__init__,_joint}.py` ✅ `tcri.joint_distribution` re-export ✅ `tests/test_tools/{test_joint,test_reduce}.py`. +- **Removed (hard bar):** n/a — additive PR; the old engines are deleted in Phase 6 (with the metric migration), per the ledger. +- **Deferred (deliberate scope boundary, logged):** **`groupby`** — raises `NotImplementedError` for now. Its correct semantics (full-space cell/clone restriction, the clone-determined guard, and sharing the draw across groups) are substantial and land with the Phase-6 metric consumers + `_metric_boxplot` (Phase 7) that actually exercise it. `_compute/_reduce` is created but not yet consumed (Phase-6 metrics use it). The `+1e-8`-vs-clamp draw bug is **fixed** in the new engine (`clamp(local_scale·base, 1e-3)`); the old engines keep the bug until they're deleted. +- **Decision baked in (user):** all entropy/MI default to **bits (log2)**, consistent with `_distance`. +## PR 6–9 — Metrics · Plotting · diag · PGM · ✅ done (branch `refactor/pr6-9`, one PR at the end) +_(committed together on a single branch per the `/goal` directive — one multi-commit PR at the end.)_ +- **PR 6 (metric-API consolidation):** the 4 metrics + `compare_groups` rewritten onto the PR5 engine in `tools/` (bits/log2; `normalize_mode='min'` default; support-only/NaN fixes). `groupby` done at the metric level via the engine's `clones=` restriction (**with a clone-disjointness guard** — added after the audit — that raises if a clone spans groups, instead of silently contaminating). `tl` repointed `metrics`→`tools`; 5 metrics onboarded into the contract. **Removed:** `mi_compare`, `delta_*`, `flux_table`, `clonotypic_entropy_base`, `clonality`, `_mi_from_joint` (old), the old `metrics/` package, and the old `preprocessing` `joint_distribution`/`joint_distribution_posterior` engines. +- **PR 7 (plotting):** split the 1437-line `_plotting.py` into `_base`/`_colors`/`_entropy`/`_mutual_information`/`_flux` (explicit `__all__`); the 4 tl↔pl twins are **cache renderers** over the tidy `tl` results (no slice-and-call). `resolve_palette` mutates in place. Dropped the non-core plots + `tcri_boxplot`/`set_color_palette`. +- **PR 8 (diag):** `diagnostics/_ppc.py` (`joint_distribution_ppc`, `phenotype_calibration`, `reconstruction_ppc`, `permutation_null` — all → DataFrame) + `_training.py` (`loss`, `archetypes`, relocated off `TCRIModel`; `plot_loss`/`plot_archetypes` deleted). 6 onboarded into the contract. +- **PR 9 (PGM→docs):** moved `build_nested_tcri_pgm`/`draw_tcri_pgm_nested` out to `docs/model_pgm.py`; dropped `daft` from runtime deps (`import tcri` no longer imports daft). +- **Gate:** full suite **82 passed** (default pytest config) in the pinned venv; contract conformance green (all metrics + diag enforced). +- **Deferred (logged):** engine `groupby` param (metrics use the metric-level path); `_provenance` sidecar + GPU guardrails (PR5); the §7.2–§7.6 api-doc code blocks still lag the frozen `.pyi` (reconcile in the docs pass); the classifier-training fix (its own PR); flux Sankey (twin renders a box for now). +## Model PR — classifier training + methods conformance · ✅ done (branch `model/classifier-fix`) +_(the deferred "classifier-training fix, its own PR" above — plus a full audit of the model vs the Supplementary Methods note.)_ +- **The classifier now trains.** Two coupled bugs: (1) `cls_logits` never entered the ELBO → added `pyro.factor("phenotype_alignment", −γ·KL(probs‖φ))` in `model()` (the note's "Inference Details" surrogate, `γ=phenotype_kl_weight`); (2) the alignment target `φ=p_ct[ct_idx]` was indexed by the **local** pyro plate index → scrambled labels across shuffled minibatches → `f_cls` collapsed to a constant. Fixed by threading **global** cell `indices` through `_get_fn_args_from_batch → model()/guide()`; the `indices=None` path now `assert`s rather than silently falling back. Pure-classifier (gate=1.0) recovery on the perfect dataset: **0.200 (chance) → 1.000**. +- **Note conformance (new `docs/contract/METHODS_CONFORMANCE.md`).** Eq-by-eq code↔note map + deviation table. Fixed alongside: `gate_prob` default `None→0.5` (π), `classifier_dropout` plumbed into `PhenotypeClassifier`, dead `class_weights`/`phenotype_weights` removed (3 signatures), dead `encoder(x)` forward in `model()` removed. +- **Tests:** new `tests/test_model_classifier.py` (perfect-recovery guard at gate 1.0 + 0.5, asserts f_cls weights actually move — with a module-local param-store isolation fixture); round-trip now guards `phenotype_kl_weight`/`gate_prob`/`classifier_dropout`. Full suite **89 passed**. +- **[G] fixed (author-approved):** α (`global_scale`) now applied to the eq-1 clonotype prior (`expanded_conc = global_scale * centroids`), removing the prior/guide scale mismatch; classifier recovery unchanged (1.000), suite green. +- **Deferred:** **[E]** `reconstruction_loss_scale=1e-3` vs eq-7 full weight (over-generation symptom) — author deferred; may be an intentional β-VAE reweighting, and raising it needs a retrain + R/NR revalidation. Tracked as a follow-up investigation. **[F]** in-silico perturbation (eqs 8–12) not implemented (additive). +## PR 10 — Public API + scverse CI · ☐ todo - **Logged test (from grafiti parity):** once `pl.__all__` exists, add a conformance assertion `set(pl.__all__) == {pl entries in _contract.pyi}` — catches *extra/missing* plot functions (whole-surface), not just signature drift on onboarded ones. (tcri's namespaced `.pyi` checks drift incrementally via `IMPLEMENTED`; this closes the whole-surface gap grafiti gets from its markdown+`__all__` channel.) --- @@ -125,6 +213,16 @@ _(diary to be filled — this is a REMOVAL PR; the ledger Phase-2 block must be # AUDIT LOG _(dated entries; what was audited, findings, actions)_ +- **(GOAL RUN — metrics contract · knob test · [E] · legacy-key removal; self-audited):** five ordered items, all landed on `model/contract`. **1) Entropy verification → HALTED as instructed.** The code differs from Supplementary Note 1 eqs 3–4, and the *note* is wrong: both equations weight by the **marginal** while taking the log of the **conditional** (a cross-entropy), and eq 4's left-hand side is mislabelled `H(p(c))` while its right side sums over φ. Proof the code is right: MI 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 while the literal formula gives **−0.345883** — a negative MI. Author confirmed: keep the code, record the erratum. **2) Metrics contract (separate, by my call).** `_metrics_contract.py` + `METRICS_CONTRACT.md` + 12 identity tests; separate from the model contract because that one is verified by *tracing* `model()`/`guide()` while metrics are pure functions pinned by *numeric identities*. Keystone identity is the decomposition above. Fixed a real wart found en route: `np.where(p>0, p*log2(p), 0)` still evaluates `log2(0)` (numpy computes both branches) — masked before the log. **3) Knob test completed with a new WIRING layer** — 32 tests. The layer exists because `lr` sat marked "hooked up" for months while dead; convergence tests cannot see a silently-ignored knob. **No new dead knobs found**; both initial failures were *test* bugs (scvi installs `LoudEarlyStopping`, not `EarlyStopping`; `predict` differs ~1.2e-07 across batch sizes = float32 kernel paths, not logic). **4) [E] re-measured** now that the phantom optimizer is gone: real-yost library ratio **1.40 at 1e-3 → 0.99 at 1e-2**; recovery (1.000) and latent separation (7.19) unchanged, so the default was raised. The synthetic reads 1.00 everywhere and **could not** have detected this — only the real 1000-gene/87%-dropout data discriminates. Three inconsistent defaults (1e-3/1e-3/1e-2) unified. **5) Deferred Phase-4 legacy keys removed** — the deferral ("still read by metrics/plotting") was stale after PR6/PR7; the last reader was a raw literal in `pp.clone_size`, now on `uns[METADATA][CLONE_COL]` with a real error message. `_ascii_hist` deleted (zero callers). Ledger ticked. **SELF-AUDIT (mutation test of the new contract):** 5 mutations — note's-literal weighting, bits→nats, zero-mass→0, epsilon-clip, MI min→average — **all caught**. Note the first attempt at the note's-literal mutation was a **no-op** (`col[supp].sum() == col.sum()`, so it changed nothing) and appeared to "escape"; re-done with the true marginal it fails correctly. Mutations must be verified to actually change behavior before their result is trusted. Suite **152 passed**; all three conformance tests green (27 / 12 / 12). + +- **(MODEL PR — methods-conformance audit vs Supplementary Note 1 — WORKFLOW, 6 lenses × 2 adversarial verifiers, 53 agents):** **22 findings survived adversarial verification, 1 refuted.** Verdict: the classifier fix (ELBO factor + global-index alignment target) is **correct and faithful to the note's surrogate**. Caught a real **HIGH the local run missed**: the new `test_model_classifier.py` leaked the process-global Pyro param store → the *full* suite was RED (1 failed / 88 passed) though the file passed in isolation — fixed with a **module-local** autouse `clear_param_store` fixture (a conftest autouse would wipe the session-scoped `trained_model`). **Fixed now:** dead `class_weights`/`phenotype_weights` removed ([D]); `indices=None` silent fallback → `assert` (re-hardens A2); dead `encoder(x)` forward in `model()` removed; `gate_prob: Optional[float]`; surrogate-KL **sign** clarified in the doc (code's `−γ·KL` realizes the note's `+γ` *penalty* intent under SVI-maximization); KL-warmup z-only scope + `num_particles` enumeration-only scope documented; round-trip now guards the new scalars. Suite **89 passed**. **Confirmed CONFORMANT:** ZINB (eq 5), β (eq 2), VampPrior (eq 3), gated ℓ rule in `predict`. **Deferred (author sign-off, change fitted results):** [E] `reconstruction_loss_scale=1e-3` vs eq-7 full weight; [G] α not applied to the eq-1 clonotype prior (prior/guide scale mismatch). [F] perturbation (eqs 8–12) additive/not implemented. - **(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). +- **(PR7–9 + live-plan audit — WORKFLOW, 2 lenses × adversarial verify, 10 agents):** **8 confirmed (all MED/LOW).** Fixed: **the `pl.*` surface is now contract-enforced** (reconciled `_contract.pyi` `resolve_palette` + onboarded all 5 pl.* into `IMPLEMENTED`); **`reconstruction_ppc` now mirrors the module's generative distribution exactly** (`module.eps` + `nb_logits` clamp `[-10,10]`) so the PPC samples from what the model defines; api-doc §9.1 param names reconciled; the live plan now states its **underpowered-subset caveat** and uses an explicit `pre→post` flux direction. **Deferred (LOW):** `joint_distribution_ppc` per-covariate aggregate lives in `df.attrs`; a few surface pieces (`save/load` round-trip, `group_singletons`) not hit by the live plan (covered by CI tests). Suite **87 passed**. +- **(LIVE R/NR TEST — real yost data, `dev/live_test_rnr.py`):** **16/16 steps OK** on a 2259-cell / 76-clone / 10-patient subset (R+NR, pre+post, 6 clusters). Train → `to_anndata` → all 6 `diag` → all 4 metrics + `compare_groups` (R vs NR) → all 4 `pl` twins (figures render). Only failure was a **test-harness** pandas bug (`reset_index`), fixed. Diagnostics surfaced two genuine **model-quality** findings (ZINB decoder over-generates counts ~6×; classifier untrained → calibration reflects the prior) — both already tracked as model debt, **not refactor defects**. Detail in `REFACTOR_NOTES`. +- **(PR6 audit — WORKFLOW, 2 lenses × adversarial verify, 12 agents):** **7 confirmed (2 MED correctness/infra, 5 doc/schema).** Fixed: the metric `groupby` now **validates clone-disjointness** and raises on a clone spanning groups (was silent cross-group contamination) + `synthetic_adata` clones made patient-disjoint; test `__init__.py` added (prepend-mode basename collision); `groupby` on a precomputed joint raises the §7.9 `ValueError`; `compare_groups` unified column schema (keeps `p_lt`); added `n_clones_ref` to `clonotypic_entropy` (+ contract). Remaining doc-only: the §7.2–§7.6 code blocks lag the frozen `.pyi` (reconcile in the docs pass). PR7–9 audit runs next (combined). +- **(PR5 audit — WORKFLOW, 3 lenses × adversarial verify, 14 agents; correctness lens math-focused):** **11 findings confirmed, 0 refuted — 2 MED, 9 LOW.** Both MED were real safety gaps, **fixed here:** (C4) restored the subset/filtered-AnnData length guard the old engine had — a sliced AnnData now errors instead of silently misaligning full-space `uns` vs subset `obsm`; (C8) `n_samples>0` now **raises** on a missing `uns[LOCAL_SCALE]` instead of silently defaulting to 1.0. Also fixed: made `clones=` ordering consistent across the single- and MultiIndex paths, and added the missing tests — the **gate-aware combine** (direct `_joint_draws` unit test; the fixture is gate=None), the T≠1 temper, and both new guards (66 passed). Reconciled stale api-doc §4.1/§4.2 (`get_xp`→`torch_device`; the real `_joint_draws(decomposed args)->(blocks,n_draws)`). **Deferred (LOW, logged):** the `_provenance` sidecar column (§7.7, with the Phase-6 cache) and GPU guardrails 5/7/8 (§7.4, with the GPU path). The correctness lens confirmed the core math (exact `P_CT` identity, `predict` parity dev 4e-8, clamped-Dirichlet draw, ct-keyed weighting, shared-draw invariant, bits/log2). +- **(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..a4333aa 100644 --- a/docs/contract/tcri_api_and_responsibilities.md +++ b/docs/contract/tcri_api_and_responsibilities.md @@ -139,8 +139,8 @@ tcri/ _stats.py # stars, AUROC+permutation, bootstrap, MWU, prob_direction, hdi, summarize _distance.py # kl_divergence, l1_distance, js_divergence, phenotype_distance dispatcher _compute/ # NEW private numeric+device seam (grafiti-mirrored) - _xp.py # resolve_device, get_xp, asnumpy (torch-first, cupy optional, CPU default) - _joint.py # _joint_draws(...) -> ndarray[n_samples, n_clones, P] (scatter-add core) + _xp.py # resolve_device, torch_device, asnumpy (torch-first, cupy optional later, CPU default) + _joint.py # _joint_draws(p_ct, ct_to_cov, ct_to_c, ct_array, cov_array, *, ...) -> (blocks, n_draws) _reduce.py # batched entropy / mutual-information / distance reductions over the stack model/ # ml _model.py # TCRIModel @@ -289,12 +289,12 @@ The engine's numeric core is written **once** as a batched, device-routable func | Signature | Responsibility | |---|---| | `resolve_device(device)` | `None`/`"cpu"`→`"cpu"`; `"mps"`→`"cpu"`; `"cuda"`/`"gpu"`/`"auto"`→GPU **iff** the backend imports AND a device is present (`getDeviceCount()>0`), else CPU. Explicit `"cuda"` warns on fallback; `"auto"`/`"gpu"` silent; unknown warns. | -| `get_xp(device)` | Return the array module — torch(-cuda) preferred (already a hard dep → zero new deps), cupy optional, numpy default. GPU libs imported **lazily inside** the function. | +| `torch_device(device)` | Return the resolved `torch.device` (`resolve_device` maps the ladder to cpu/cuda). torch-first (already a hard dep → zero new deps); cupy optional later. GPU libs imported **lazily inside** the function. | | `asnumpy(x)` | Host-boundary shim: `cupy.asnumpy(x)` / `x.cpu().numpy()` / `np.asarray(x)`. Every accelerated function returns a plain numpy array. | ### 4.2 `_joint.py` / `_reduce.py` — the batched core -- **`_joint_draws(adata, *, covariate, clones, n_samples, use_logits, temperature, gate_prob, random_state, device) -> np.ndarray`** — returns the `[max(n_samples,1), n_clones, P]` joint stack. Precomputes clone integer codes **once**; draws all `n_samples` Dirichlet samples in one batched kernel from `clamp(s·m̃, 1e-3)`; softmaxes the (optionally gated) per-cell combination batched on the leading axis; reduces per clone with a **constant-index scatter-add** (`np.add.at` / `torch.index_add_` / `cupy.bincount`) instead of a per-draw `pandas.groupby` — the dominant win. Validates finiteness / nonnegativity / per-row sum $\approx1$ **on device** before returning; `float64` accumulators for CPU/GPU parity; `asnumpy` at the boundary; chunked over cells/draws to bound device memory. +- **`_joint_draws(p_ct, ct_to_cov, ct_to_c, ct_array, cov_array, *, local_scale, n_samples, temperature, use_logits, covariate_idx, logits, gate_prob, weighted, random_state, device) -> (blocks, n_draws)`** — the adata-unpacking lives in the `tools/_joint` wrapper; this core takes **decomposed uns arrays** and returns a **list of per-covariate `(cov_idx, clone_idx, J[S, n_rows, P])` blocks** plus the draw count (`covariate=None` stacks variable-length per-covariate blocks). Draws all `n_samples` Dirichlet samples over **all ct rows in one batched kernel** from `clamp(s·m̃, 1e-3)` (the shared-draw invariant), then slices per covariate; softmaxes the (optionally gated) per-cell combination batched on the leading axis; reduces per clone with a **scatter-add** (`torch.index_add_`) instead of a per-draw `pandas.groupby` — the dominant win. `float64` accumulators for CPU/GPU parity; `asnumpy` at the boundary. **Phase-6 (with the GPU path / `_reduce` wiring):** on-device per-row-sum $\approx1$ validation + chunking over cells/draws (§7.4 guardrails 5/7/8). - **`_reduce.py`** — batched `entropy`, `mutual_information`, `distance` as `xlogx`/outer-product reductions over the whole stack (no per-draw scipy call, no per-clone `.loc`), plus the `summarize`/`hdi` reduction over the sample axis. ### 4.3 GPU guardrails (replicated uniformly from grafiti) @@ -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. @@ -392,9 +392,10 @@ Called only by `TCRIModel.to_anndata`; folds in the old `register_phenotype_key` joint_distribution( adata, *, covariate=None, # None → ALL covariate values in one pass (shared draw) - groupby=None, + groupby=None, # NOTE: deferred to Phase 6 in code (raises NotImplementedError until then) n_samples=0, use_logits=True, # was posterior=; alias cell_informed=; classifier-mixing switch + weighted=False, # per-clonotype (False) vs cell-weighted (True); ct-keyed clones=None, temperature=1.0, random_state=None, @@ -716,8 +717,8 @@ Read-only checks on the finalized model. PPCs return `DataFrame`s; the two reloc |---|---| | `joint_distribution_ppc(adata, *, covariate=None, distance_metric="l1", temperature=1.0) -> pandas.DataFrame` | **(fixed `compare_joint_distribution`)** Model vs empirical per-clone phenotype frequencies. $P_{\text{model}}(\phi\mid c,m)=\texttt{joint\_distribution}(adata, covariate=m)[c]$; $P_{\text{emp}}(\phi\mid c,m)=\frac{\#\{i\in c,m:\text{pheno}_i=\phi\}}{\#\{i\in c,m\}}$; per-clone $\delta_c=\text{L1}$ or $\text{KL}(P_{\text{emp}}\Vert P_{\text{model}})$, plus per-covariate aggregate. **Model-free (adata only).** **Bug fix:** reads `clonotype_col`/`phenotype_col` from `uns[K.METADATA]` instead of the undefined global `model` (repairs the `NameError`). | | `phenotype_calibration(adata, *, n_bins=10) -> pandas.DataFrame` | Reliability of `predict()` probabilities: bin cells by predicted max-prob; per bin compare mean predicted prob to empirical accuracy; $\text{ECE}=\sum_b\frac{n_b}{N}|\text{acc}_b-\text{conf}_b|$. **adata only.** Returns `(bin, mean_pred, emp_freq, count)` + scalar `ECE`. | -| `reconstruction_ppc(model, adata=None, *, n_samples=100, seed=0) -> pandas.DataFrame` | ZINB reconstruction PPC: simulate from the fitted decoder ($\mu,\theta,\pi_{\text{dropout}}$), compare library size / per-gene dropout / mean–variance vs observed. **`model` REQUIRED** (live decoder lives on the module, not in `adata`). Returns statistic × {observed, simulated, discrepancy}. | -| `permutation_null(adata, *, metric="mutual_information", covariate=None, groupby=None, n_permutations=1000, seed=0) -> pandas.DataFrame` | Permute phenotype labels within each covariate $R$ times, recompute the metric to form a null; $p=\text{mean}(\text{null}\ge\text{obs})$, $z=\frac{\text{obs}-\overline{\text{null}}}{\text{sd(null)}}$. **adata only.** One shared draw stack (§7.8). Returns per stratum: `observed, null_mean, null_sd, z, p`. | +| `reconstruction_ppc(model, adata=None, *, n_sims=100, random_state=0) -> pandas.DataFrame` | ZINB reconstruction PPC: simulate from the fitted decoder ($\mu,\theta,\pi_{\text{dropout}}$), compare library size / per-gene dropout / mean–variance vs observed. **`model` REQUIRED** (live decoder lives on the module, not in `adata`). Returns statistic × {observed, simulated, discrepancy}. | +| `permutation_null(adata, *, metric="mutual_information", covariate=None, groupby=None, n_perm=1000, random_state=None) -> pandas.DataFrame` | Permute phenotype labels within each covariate $R$ times, recompute the metric to form a null; $p=\text{mean}(\text{null}\ge\text{obs})$, $z=\frac{\text{obs}-\overline{\text{null}}}{\text{sd(null)}}$. **adata only.** One shared draw stack (§7.8). Returns per stratum: `observed, null_mean, null_sd, z, p`. | `__all__ = ["joint_distribution_ppc", "phenotype_calibration", "reconstruction_ppc", "permutation_null"]` diff --git a/docs/contract/tcri_implementation_plan.md b/docs/contract/tcri_implementation_plan.md index cfc7747..12cdc12 100644 --- a/docs/contract/tcri_implementation_plan.md +++ b/docs/contract/tcri_implementation_plan.md @@ -117,7 +117,8 @@ Freeze the map **before** Phase 5 (the first breaking PR). Renames are breaking; | `mi_compare`, `delta_entropy_table`, `flux_table`, `delta_clonotypic_entropy`, `phenotypic_entropy_delta` | **deleted** — expressed via `groupby=` + `tl.compare_groups` | Phase 6 | | `tl.phenotypic_entropies` / `tl.clonotypic_entropies` (plural batch forms) | **deleted** — subsumed by `groupby=` on the singular metric | Phase 6 | | `get_cell_phenotype_probs` | `predict` (scvi/CellAssign idiom; order‑preserving loader, indexed by `obs_names`) | Phase 4 | -| `register_model` (+ `classify_phenotypes`, `register_*_key`) | `model.to_anndata` (thin) | Phase 4 | +| `register_model` (+ `register_*_key`) | `model.to_anndata` (thin) | Phase 4 | +| `classify_phenotypes` | — (DROP; superseded by `REDO_LIST.md`) | Phase 2 ✅ deleted | | `register_clonotype_key` / `register_phenotype_key` | folded (private) into `to_anndata` | Phase 4 | | `pl.clonotypic_entropy_by_phenotype` | `pl.clonotypic_entropy` | Phase 7 | | `plot_pheno_sankey` | `pl.phenotypic_flux` (sankey) | Phase 7 | diff --git a/docs/model_pgm.py b/docs/model_pgm.py new file mode 100644 index 0000000..493ac37 --- /dev/null +++ b/docs/model_pgm.py @@ -0,0 +1,176 @@ +"""Model plate-diagram (PGM) generators — moved OUT of the tcri package (PR9, §9). +Requires `daft` (a docs-only extra, NOT a tcri runtime dependency).""" +import daft +import matplotlib.pyplot as plt +import numpy as np + + +def build_nested_tcri_pgm(): + """ + A fully explicit TCRI PGM matching the implementation in _model.py + with improved layout to minimize edge crossings + """ + # Define colors + red, yellow, green, gray, blue = "#cd442a", "#f0bd00", "#7e9437", "#eee", "#009de1" + + # Create a PGM canvas + pgm = daft.PGM( + shape=[8, 8], # width x height + origin=[0, 0], + grid_unit=1.6, + node_unit=1.5 + ) + + # ------------------------------------------------------------------ + # 1) Global hyperparameters for Dirichlet priors + # ------------------------------------------------------------------ + pgm.add_node( + "global_scale", + r"$\mathrm{global\_scale}$", + 5.1, # x + 6.6, # y + fixed=True, + plot_params={"fc": "#DDD"} + ) + + # ------------------------------------------------------------------ + # 2) Plate: batch (b) - outermost plate + # ------------------------------------------------------------------ + + # Batch-level variables - aligned vertically + # ------------------------------------------------------------------ + # 3) Plate: clonotypes (c) - middle plate + # ------------------------------------------------------------------ + pgm.add_plate( + [1.0, 1.0, 6.0, 6.3], # [x, y, width, height] + label=r"clonotypes $(c)$", + shift=-0.1 + ) + + # p_c (Dirichlet) + pgm.add_node( + "p_c", + r"$p_c$", + 3.8, # x + 6.6, # y + observed=False, + plot_params={"fc": blue} + ) + + # Edge: global_scale -> p_c + pgm.add_edge("global_scale", "p_c") + + # ------------------------------------------------------------------ + # 4) Plate: clone-covariate (ct) - inner plate + # ------------------------------------------------------------------ + pgm.add_plate( + [1.5, 1.5, 5.0, 4.5], # [x, y, width, height] + label=r"clone-covariate $(ct)$", + shift=-0.1 + ) + + # local_scale moved inside clone-covariate plate + pgm.add_node( + "local_scale", + r"$\mathrm{local\_scale}$", + 5.2, # x + 5.3, # y + fixed=True, + plot_params={"fc": "#DDD"} + ) + + # p_ct (Dirichlet) + pgm.add_node( + "p_ct", + r"$p_{ct}$", + 3.8, # x + 5.3, # y + observed=False, + plot_params={"fc": yellow} + ) + + # Edges: p_c -> p_ct, local_scale -> p_ct + pgm.add_edge("p_c", "p_ct") + pgm.add_edge("local_scale", "p_ct") + + # ------------------------------------------------------------------ + # 5) Plate: data (i) - innermost plate + # ------------------------------------------------------------------ + pgm.add_plate( + [1.9, 2., 4.0, 2.7], # [x, y, width, height] + label=r"data $(i)$", + shift=-0.1 + ) + + # Grid layout for data-level variables - aligned vertically + # Column 1: Observed variables + pgm.add_node( + "obs", + r"$X_{i}$", + 5, # x + 4.0, # y + observed=True, + plot_params={"fc": gray} + ) + + pgm.add_node( + "obs_label", + r"$Pheno_{i}$", + 2.5, # x + 4.0, # y + observed=True, + plot_params={"fc": gray} + ) + + # Column 2: Latent variables + pgm.add_node( + "latent", + r"$z_i$", + 3.8, # x + 2.8, # y + observed=False, + plot_params={"fc": green} + ) + + pgm.add_node( + "z_i_phen", + r"$z_{i,\mathrm{phen}}$", + 3.8, # x + 4.0, # y + observed=False, + plot_params={"fc": red} + ) + + # Column 3: Decoder inputs + pgm.add_node( + "px_r", + r"$ZINB(X_{i})$", + 5, # x + 2.8, # y + observed=False, + plot_params={"fc": "#DDD"} + ) + + # Edges - now mostly vertical and horizontal + # Data-level edges + pgm.add_edge("p_ct", "z_i_phen") + pgm.add_edge("latent", "z_i_phen") + pgm.add_edge("z_i_phen", "obs_label") + + # Direct connections to z_i (previously through decoder) + pgm.add_edge("latent", "obs") + pgm.add_edge("px_r", "obs") + + # ------------------------------------------------------------------ + # Text / Title + # ------------------------------------------------------------------ + pgm.add_text(3.1,7.5, "TCRi Model", fontsize=14) + + return pgm + + +def draw_tcri_pgm_nested(): + pgm = build_nested_tcri_pgm() + pgm.render() + pgm.figure.savefig("tcri_model_fully_explicit.pdf", dpi=300) + plt.show() diff --git a/pyproject.toml b/pyproject.toml index e7fda8d..54c60dd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,7 +45,6 @@ dependencies = [ "seaborn>=0.13.2", "mpltern>=1.0.4", "gseapy>=1.1.4", - "daft", "tqdm>=4.66.5", ] 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/__init__.py b/tcri/__init__.py index d42c7a8..1131424 100644 --- a/tcri/__init__.py +++ b/tcri/__init__.py @@ -5,12 +5,16 @@ except PackageNotFoundError: # running from a source tree without an install __version__ = "0.0.0+unknown" -from . import metrics as tl +from . import tools as tl # PR6: tl repointed metrics -> tools (engine-backed metrics) from . import preprocessing as pp from . import plotting as pl from . import utils as ut from . import model as ml +from . import diagnostics as diag # PR8 + +# The unified engine, re-exported top-level for prominence. +from .tools import joint_distribution import sys -sys.modules.update({f'{__name__}.{m}': globals()[m] for m in ['tl', 'pp', 'pl', 'ut', 'ml']}) \ No newline at end of file +sys.modules.update({f'{__name__}.{m}': globals()[m] for m in ['tl', 'pp', 'pl', 'ut', 'ml', 'diag']}) \ No newline at end of file diff --git a/tcri/_compute/__init__.py b/tcri/_compute/__init__.py new file mode 100644 index 0000000..3fd57d7 --- /dev/null +++ b/tcri/_compute/__init__.py @@ -0,0 +1,8 @@ +"""Private numeric + device seam for the engine (grafiti ``_compute`` parity). + +- :mod:`._xp` — device dispatch (torch-first; CPU / torch-CUDA), ``asnumpy`` boundary. +- :mod:`._joint` — ``_joint_draws``: the batched ``[S, n_clones, P]`` engine core. +- :mod:`._reduce` — batched entropy / mutual-information reductions (**bits / log2** default). + +Nothing here is public API; GPU libs are imported lazily inside functions. +""" diff --git a/tcri/_compute/_joint.py b/tcri/_compute/_joint.py new file mode 100644 index 0000000..be1c663 --- /dev/null +++ b/tcri/_compute/_joint.py @@ -0,0 +1,165 @@ +"""Batched engine core — the clone×phenotype joint as a ``[S, n_clones, P]`` stack, +device-routed via :mod:`._xp` (§7.2). ``tools/joint_distribution`` is a thin +DataFrame wrapper over :func:`_joint_draws`. + +Design invariants (§7.1/§7.8): +- Temperature-temper the base **once**: ``T==1`` is the exact identity (``base = p_ct``, + no ``eps`` round-trip), else ``base = softmax(log(p_ct+1e-8)/T)``. +- Draws are made over **all ct rows at once** (a single seeded ``Dirichlet.sample((N,))``), + then sliced per covariate — so the draw count == ``n_samples`` regardless of how many + covariates/groups are requested (the draw-once invariant). +- ``n_samples=0`` is deterministic (base only; no Monte-Carlo, ``random_state`` ignored). +- ``use_logits=True`` folds per-cell logits with ``log(base)`` (gate-aware) exactly like + ``predict()`` and scatter-adds per clone; ``use_logits=False`` returns the ct-level base. +- ``weighted`` scales each clone row by its **ct-keyed** cell count (the fix for the old + clone-indexed ``Counter`` bug); ``weighted=False`` row-normalizes to a per-clone simplex. +""" +from __future__ import annotations + +import contextlib + +import numpy as np + +from . import _xp + +_EPS = 1e-8 + + +@contextlib.contextmanager +def _torch_seed(random_state): + """Seed the global torch RNG from ``random_state`` for the duration, then restore + it (so the draw is reproducible without a permanent global side effect). ``None`` + leaves the RNG untouched (non-reproducible).""" + import torch + + if random_state is None: + yield + return + if isinstance(random_state, torch.Generator): + seed = int(random_state.initial_seed()) + elif isinstance(random_state, (int, np.integer)): + seed = int(random_state) + elif hasattr(random_state, "integers"): # numpy Generator + seed = int(random_state.integers(0, 2**31 - 1)) + else: + seed = int(np.random.default_rng(random_state).integers(0, 2**31 - 1)) + prev = torch.random.get_rng_state() + torch.manual_seed(seed) + try: + yield + finally: + torch.random.set_rng_state(prev) + + +def _joint_draws( + p_ct, + ct_to_cov, + ct_to_c, + ct_array, + cov_array, + *, + local_scale, + n_samples=0, + temperature=1.0, + use_logits=True, + covariate_idx=None, + logits=None, + gate_prob=None, + weighted=False, + random_state=None, + device=None, +): + """Compute the joint stack per covariate. + + Returns ``(blocks, n_draws)`` where ``blocks`` is a list of + ``(cov_idx, clone_idx[n_rows], J[S, n_rows, P])`` (host numpy, float64) — one block + per covariate processed (all covariates when ``covariate_idx is None``) — and + ``n_draws`` is the number of Dirichlet draws performed (``== n_samples``, or ``0``). + """ + import torch + + dev = _xp.torch_device(device) + p_ct_t = torch.as_tensor(np.asarray(p_ct), dtype=torch.float64, device=dev) # [n_ct, P] + ct_to_cov_t = torch.as_tensor(np.asarray(ct_to_cov), dtype=torch.long, device=dev) + ct_to_c_t = torch.as_tensor(np.asarray(ct_to_c), dtype=torch.long, device=dev) + ct_arr_t = torch.as_tensor(np.asarray(ct_array), dtype=torch.long, device=dev) + cov_arr_t = torch.as_tensor(np.asarray(cov_array), dtype=torch.long, device=dev) + n_ct, P = p_ct_t.shape + + # sanitize any non-finite p_ct rows to uniform (matches the legacy guard) + bad = ~torch.isfinite(p_ct_t) + if bad.any(): + p_ct_t = torch.where(bad, torch.full_like(p_ct_t, 1.0 / P), p_ct_t) + + # temper the base once (T==1 == exact identity, no eps round-trip) + if float(temperature) == 1.0: + base = p_ct_t + else: + base = torch.softmax(torch.log(p_ct_t + _EPS) / float(temperature), dim=-1) + + # draw over ALL ct rows once (shared-draw invariant), seeded + if n_samples and int(n_samples) > 0: + conc = torch.clamp(float(local_scale) * base, min=1e-3) + with _torch_seed(random_state): + bases = torch.distributions.Dirichlet(conc).sample((int(n_samples),)) # [N, n_ct, P] + n_draws = int(n_samples) + S = int(n_samples) + else: + bases = base.unsqueeze(0) # [1, n_ct, P] + n_draws = 0 + S = 1 + + if use_logits: + logits_t = torch.as_tensor(np.asarray(logits), dtype=torch.float64, device=dev) # [n_obs, P] + use_gate = gate_prob is not None and not (isinstance(gate_prob, float) and np.isnan(gate_prob)) + g = float(gate_prob) if use_gate else None + + covs = ([int(covariate_idx)] if covariate_idx is not None + else sorted(int(c) for c in torch.unique(ct_to_cov_t).tolist())) + + blocks = [] + for m in covs: + ct_rows = torch.nonzero(ct_to_cov_t == m, as_tuple=True)[0] # ct indices at cov m + clones_m = ct_to_c_t[ct_rows] # clone id per ct row + # ct-keyed cell counts at this covariate (the weighting fix) + counts_per_ct = torch.bincount( + ct_arr_t[cov_arr_t == m], minlength=n_ct + ).to(torch.float64)[ct_rows] # [n_ct_m] + + if not use_logits: + J = bases[:, ct_rows, :].clone() # [S, n_ct_m, P] — each ct row is a clone + if weighted: + J = J * counts_per_ct[None, :, None] + else: + cells_m = torch.nonzero(cov_arr_t == m, as_tuple=True)[0] + ct_of_cell = ct_arr_t[cells_m] # global ct per cell + lut = torch.full((n_ct,), -1, dtype=torch.long, device=dev) + lut[ct_rows] = torch.arange(ct_rows.shape[0], device=dev) + local_ct = lut[ct_of_cell] # [n_cells_m] local clone-row index + ell = logits_t[cells_m].unsqueeze(0) # [1, n_cells_m, P] + J = torch.zeros((S, ct_rows.shape[0], P), dtype=torch.float64, device=dev) + # Chunk over draws. The per-cell chain below needs four live + # [S, n_cells_m, P] float64 tensors at once (~2.9 GB each at S=500, + # 60k cells, P=12 — 11 GB measured). Draws are independent, so chunking + # is bit-identical and caps peak memory at a few hundred MB. + n_cells_m = cells_m.shape[0] + per_draw = max(1, n_cells_m * P) + chunk = max(1, min(S, int(8_000_000 // per_draw))) # ~64 MB/temp in float64 + for s0 in range(0, S, chunk): + s1 = min(s0 + chunk, S) + b_cell = bases[s0:s1][:, ct_of_cell, :] # [c, n_cells_m, P] + log_b = torch.log(b_cell + _EPS) + combine = (g * ell + (1.0 - g) * log_b) if use_gate else (ell + log_b) + # §7.1: the base is already tempered (b) and the combine is divided by T here; + # at T==1 both are the identity, so this reproduces predict() bit-for-bit. T!=1 is + # a deliberate analysis-time temper (the two T's do not cancel — documented). + p_cell = torch.softmax(combine / float(temperature), dim=-1) + Jc = torch.zeros((s1 - s0, ct_rows.shape[0], P), dtype=torch.float64, device=dev) + Jc.index_add_(1, local_ct, p_cell) # sum P over cells per clone (row sum == count) + J[s0:s1] = Jc + if not weighted: + J = J / J.sum(-1, keepdim=True).clamp_min(_EPS) + + blocks.append((m, _xp.asnumpy(clones_m), _xp.asnumpy(J))) + + return blocks, n_draws diff --git a/tcri/_compute/_reduce.py b/tcri/_compute/_reduce.py new file mode 100644 index 0000000..16bd6e6 --- /dev/null +++ b/tcri/_compute/_reduce.py @@ -0,0 +1,56 @@ +"""Batched information-theoretic reductions over the joint stack ``[S, n_clones, P]``. + +**All quantities default to BITS (log2)** — the tcri convention (matches +:mod:`tcri._distance`); pass ``base=None`` for nats or ``base=b`` for an arbitrary +base. Reductions run in **float64** accumulators (§7.4 guardrail 6) so a future GPU +path matches CPU, and use the ``0·log0 := 0`` convention. Grafiti `joint.py::_entropy/_mi` +parity, generalized to batch over the leading sample/covariate/group axes. +""" +from __future__ import annotations + +import numpy as np + +_EPS = 1e-12 + + +def _logb(x, base): + if base == 2: + return np.log2(x) + if base is None: # nats + return np.log(x) + return np.log(x) / np.log(base) + + +def entropy(p, *, axis=-1, base=2): + """Shannon entropy of a distribution along ``axis`` (bits by default), batched + over every other axis. ``0·log0 := 0``; ``p`` is used as-is (assumed normalized).""" + p = np.asarray(p, dtype=np.float64) + with np.errstate(divide="ignore", invalid="ignore"): + terms = np.where(p > 0, p * _logb(p, base), 0.0) + return -np.sum(terms, axis=axis) + + +def mutual_information(joint, *, base=2): + """MI(X;Y) (bits by default) of a joint ``P[..., X, Y]`` over the last two axes, + batched over the leading axes. ``joint`` is renormalized per batch element, so + an un-normalized (e.g. cell-weighted count) table is accepted.""" + P = np.asarray(joint, dtype=np.float64) + Z = P.sum(axis=(-2, -1), keepdims=True) + P = np.where(Z > 0, P / Z, 0.0) + Px = P.sum(axis=-1, keepdims=True) # [..., X, 1] + Py = P.sum(axis=-2, keepdims=True) # [..., 1, Y] + indep = Px * Py + with np.errstate(divide="ignore", invalid="ignore"): + terms = np.where(P > 0, P * _logb((P + _EPS) / (indep + _EPS), base), 0.0) + return np.sum(terms, axis=(-2, -1)) + + +def joint_from_conditional(cond, weights): + """Assemble a normalized joint ``P[..., C, Y]`` from per-clone conditionals + ``cond`` = P(φ|c) ``[..., C, Y]`` and clone ``weights`` ``[..., C]`` (any scale; + normalized here). ``weighted=False`` passes uniform weights, ``weighted=True`` + passes clone cell counts — the clone-mass choice lives with the caller.""" + cond = np.asarray(cond, dtype=np.float64) + w = np.asarray(weights, dtype=np.float64) + w = w / np.clip(w.sum(axis=-1, keepdims=True), _EPS, None) + return cond * w[..., None] diff --git a/tcri/_compute/_xp.py b/tcri/_compute/_xp.py new file mode 100644 index 0000000..a795cb9 --- /dev/null +++ b/tcri/_compute/_xp.py @@ -0,0 +1,62 @@ +"""Device seam for the batched engine — **torch-first** (grafiti `_compute/_xp` +parity, adapted because tcri's numeric core is torch: the Dirichlet draws and the +softmax are `torch`, and torch≥2.4 is already a hard dep so torch.cuda is zero new +deps). `cupy` may be added later as a second backend; for now the ladder is +CPU / torch-CUDA. Every accelerated function returns a host numpy array via +:func:`asnumpy` (§7.4 guardrail 4). GPU libs are imported lazily inside functions +(guardrail 1) — importing this module never touches CUDA. +""" +from __future__ import annotations + +import warnings + +import numpy as np + + +def resolve_device(device: str | None) -> str: + """Resolve ``device`` to ``"cpu"`` or ``"cuda"`` (§7.1 ladder). + + ``None``/``"cpu"`` → cpu; ``"mps"`` → cpu (first pass: the Dirichlet core has no + Metal backend); ``"auto"``/``"gpu"``/``"cuda"`` → cuda only if torch reports a + device, else cpu (an explicit ``"cuda"`` warns on fallback; ``auto``/``gpu`` are + silent). + """ + if device in (None, "cpu"): + return "cpu" + if device == "mps": + return "cpu" + if device in ("cuda", "gpu", "auto"): + try: + import torch + + if torch.cuda.is_available() and torch.cuda.device_count() > 0: + return "cuda" + except Exception: + pass + if device == "cuda": + warnings.warn( + "device='cuda' requested but torch CUDA is unavailable; using CPU.", + stacklevel=2, + ) + return "cpu" + warnings.warn(f"device={device!r} not recognized; using CPU.", stacklevel=2) + return "cpu" + + +def torch_device(device: str | None): + """The :class:`torch.device` for the resolved device.""" + import torch + + return torch.device("cuda" if resolve_device(device) == "cuda" else "cpu") + + +def asnumpy(x) -> np.ndarray: + """Bring an array (possibly a torch tensor, possibly on GPU) back to host numpy.""" + try: + import torch + + if isinstance(x, torch.Tensor): + return x.detach().cpu().numpy() + except Exception: + pass + return np.asarray(x) diff --git a/tcri/_contract.pyi b/tcri/_contract.pyi index 1ac327f..d50fe37 100644 --- a/tcri/_contract.pyi +++ b/tcri/_contract.pyi @@ -5,6 +5,10 @@ The conformance test (``tests/test_contract_conformance.py``) checks each but-absent functions are the refactor worklist. Prose spec: ``docs/contract/tcri_api_and_responsibilities.md``. +This freezes the public *interface*. Its sibling — ``tcri/model/_model_contract.py`` +(prose: ``docs/contract/MODEL_CONTRACT.md``) — freezes the model's *mathematics* +(Supplementary Note 1). + RULES this file encodes: - Only the KEPT surface is declared. A symbol NOT in this file must NOT be public after the refactor (see the Removal Ledger in ``docs/contract/REFACTOR_AGENDA.md``). @@ -26,19 +30,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) ─────────────────────────────────────────────────────── @@ -56,11 +64,13 @@ class tl: adata: AnnData, *, covariate: Optional[str] = ..., groupby: Optional[str] = ..., n_samples: int = ..., use_logits: bool = ..., weighted: bool = ..., clones: Any = ..., temperature: float = ..., random_state: Any = ..., + device: Any = ..., ) -> pd.DataFrame: ... def clonotypic_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 = ..., + clones: Any = ..., weighted: bool = ..., normalized: bool = ..., + n_clones_ref: Any = ..., random_state: Any = ..., ) -> Any: ... def phenotypic_entropy( adata_or_jd: Any, *, covariate: Optional[str] = ..., groupby: Optional[str] = ..., @@ -114,7 +124,7 @@ class pl: distance_metric: str = ..., palette: Any = ..., ax: Any = ..., figsize: Any = ..., save: Any = ..., show: Any = ..., return_axes: bool = ..., ) -> Any: ... - def resolve_palette(adata: AnnData, columns: Any) -> Any: ... + def resolve_palette(adata: AnnData, columns: Any, *, palette: Any = ...) -> Any: ... # ── diagnostics (diag) — returns DataFrames ────────────────────────────────── diff --git a/tcri/_keys.py b/tcri/_keys.py index 5bf5940..894bd59 100644 --- a/tcri/_keys.py +++ b/tcri/_keys.py @@ -6,9 +6,9 @@ conformance test forbids it. Migrating a reader/writer means swapping the literal for the constant here. -The two legacy shadow keys (``tcri_clone_key`` / ``tcri_phenotype_key``) and the -legacy ``X_tcri_phenotypes`` obsm slot are listed only so the removal step can -find and delete them; new code uses ``METADATA`` + ``X_PROBABILITIES``. +The legacy shadow keys (``tcri_clone_key`` / ``tcri_phenotype_key``) and the old +``X_tcri_phenotypes`` obsm slot have been removed — use ``METADATA`` and +``X_PROBABILITIES``. """ # ── uns: metadata + learned priors ─────────────────────────────────────────── @@ -45,8 +45,10 @@ PHENOTYPE_COL = "phenotype_col" BATCH_COL = "batch_col" -# ── legacy — declared ONLY so the removal step can find + delete them ───────── -LEGACY_MANAGER = "tcri_manager" # non-picklable AnnDataManager stash → drop -LEGACY_CLONE_KEY = "tcri_clone_key" # shadow of METADATA[CLONE_COL] → drop -LEGACY_PHENOTYPE_KEY = "tcri_phenotype_key" # shadow of METADATA[PHENOTYPE_COL] → drop -LEGACY_X_PHENOTYPES = "X_tcri_phenotypes" # old prob slot → X_PROBABILITIES +# ── legacy ─────────────────────────────────────────────────────────────────── +# The shadow keys `tcri_clone_key` / `tcri_phenotype_key` and the old +# `X_tcri_phenotypes` obsm slot are GONE: `to_anndata` no longer writes them and +# nothing reads them (`pp.clone_size`, the last reader, now uses METADATA). +# `LEGACY_MANAGER` stays because it still does defensive work — `save_tcri_session` +# pops it so a stray non-picklable AnnDataManager can never be serialized. +LEGACY_MANAGER = "tcri_manager" # popped defensively before save diff --git a/tcri/_stats.py b/tcri/_stats.py index 6468298..1908ddd 100644 --- a/tcri/_stats.py +++ b/tcri/_stats.py @@ -11,7 +11,7 @@ import math import numpy as np -from scipy.stats import mannwhitneyu +from scipy.stats import mannwhitneyu, rankdata from sklearn.metrics import roc_auc_score @@ -35,26 +35,51 @@ def mann_whitney(a, b, *, alternative: str = "two-sided"): def auc_and_label_permutation(scores, labels, pos_label=None, n_perm=200_000, seed=42, max_exact=200_000): - """Observed AUROC + a label-permutation p-value (exact when feasible).""" + """Observed AUROC + a label-permutation p-value (exact when feasible). + + Under label permutation the *scores* never change, so the ranks are computed + ONCE and each permuted AUROC is a rank-sum over the permuted positive set via + the Mann–Whitney identity + + AUC = (Σ ranks[pos] − n_pos(n_pos+1)/2) / (n_pos·n_neg) + + which is exact (midranks reproduce ``roc_auc_score``'s tie handling) and turns + each draw from an O(n log n) re-sort into an O(n_pos) sum. Measured 137× on the + Monte-Carlo path (191 s → 1.4 s at the default ``n_perm``). + """ scores = np.asarray(scores, dtype=float) labels = np.asarray(labels) if pos_label is None: pos_label = sorted(set(labels))[-1] y = (labels == pos_label).astype(int) obs_auc = roc_auc_score(y, scores) + n = len(y) n_pos = int(y.sum()) - n_exact = math.comb(len(y), n_pos) + n_neg = n - n_pos + + if n_pos == 0 or n_neg == 0: # AUROC undefined; keep the old failure mode + perm_stats = np.array([]) + return obs_auc, float("nan"), perm_stats, "degenerate" + + # midranks: ties get the average rank, matching roc_auc_score exactly + ranks = rankdata(scores) + denom = float(n_pos) * float(n_neg) + offset = n_pos * (n_pos + 1) / 2.0 + + n_exact = math.comb(n, n_pos) if n_exact <= max_exact: - perm_stats = np.array([ - roc_auc_score(np.isin(np.arange(len(y)), idx).astype(int), scores) - for idx in itertools.combinations(range(len(y)), n_pos) - ]) + perm_stats = np.fromiter( + ((ranks[list(idx)].sum() - offset) / denom + for idx in itertools.combinations(range(n), n_pos)), + dtype=float, count=n_exact, + ) perm_mode = "exact" else: rng = np.random.default_rng(seed) - perm_stats = np.array([ - roc_auc_score(rng.permutation(y), scores) for _ in range(n_perm) - ]) + perm_stats = np.empty(n_perm, dtype=float) + for i in range(n_perm): + # a random size-n_pos subset of ranks == a random label permutation + perm_stats[i] = (rng.permutation(ranks)[:n_pos].sum() - offset) / denom perm_mode = "mc" p_perm = np.mean(np.abs(perm_stats - 0.5) >= np.abs(obs_auc - 0.5)) return obs_auc, p_perm, perm_stats, perm_mode diff --git a/tcri/diagnostics/__init__.py b/tcri/diagnostics/__init__.py new file mode 100644 index 0000000..7ab9e95 --- /dev/null +++ b/tcri/diagnostics/__init__.py @@ -0,0 +1,19 @@ +"""``tcri.diag`` — diagnostics: posterior-predictive checks (all return DataFrames) + the +two relocated training plots. Read-only checks on a finalized model. +""" +from ._ppc import ( + joint_distribution_ppc, + phenotype_calibration, + reconstruction_ppc, + permutation_null, +) +from ._training import loss, archetypes + +__all__ = [ + "joint_distribution_ppc", + "phenotype_calibration", + "reconstruction_ppc", + "permutation_null", + "loss", + "archetypes", +] diff --git a/tcri/diagnostics/_ppc.py b/tcri/diagnostics/_ppc.py new file mode 100644 index 0000000..9d1ec9f --- /dev/null +++ b/tcri/diagnostics/_ppc.py @@ -0,0 +1,170 @@ +"""``tcri.diag`` posterior-predictive checks (§9.1). All return DataFrames. ``model`` is +required only for ``reconstruction_ppc`` (the live ZINB decoder); the rest are adata-only. +""" +from __future__ import annotations + +import numpy as np +import pandas as pd + +from .. import _keys as K + +__all__ = ["joint_distribution_ppc", "phenotype_calibration", "reconstruction_ppc", "permutation_null"] + + +def joint_distribution_ppc(adata, *, covariate=None, distance_metric="l1", temperature=1.0, + clones=None, random_state=None): + """Model vs empirical per-clone phenotype frequencies (the fixed + ``compare_joint_distribution``). Per-clone distance + per-covariate aggregate. adata-only.""" + from .._distance import phenotype_distance + from ..tools import joint_distribution + + dist_fn = phenotype_distance(distance_metric) + meta = adata.uns[K.METADATA] + clone_col, pheno_col, cov_col = meta["clone_col"], meta["phenotype_col"], meta["covariate_col"] + phenos = list(adata.uns[K.PHENOTYPE_CATEGORIES]) + covs = [covariate] if covariate is not None else list(adata.uns[K.COVARIATE_CATEGORIES]) + + rows = [] + for m in covs: + Jm = joint_distribution(adata, covariate=m, use_logits=True, n_samples=0, + temperature=temperature, clones=clones) + cmask = adata.obs[cov_col].astype(str) == str(m) + sub = adata.obs.loc[cmask, [clone_col, pheno_col]] + for c in Jm.index: + emp = (sub.loc[sub[clone_col] == c, pheno_col].value_counts() + .reindex(phenos).fillna(0.0).to_numpy(dtype=float)) + if emp.sum() <= 0: + continue + emp = emp / emp.sum() + model_p = np.asarray(Jm.loc[c].reindex(phenos).to_numpy(), dtype=float) + rows.append({"covariate": m, "clonotype": c, + "distance": float(dist_fn(emp, model_p))}) + df = pd.DataFrame(rows) + if len(df): + agg = df.groupby("covariate")["distance"].mean().rename("mean_distance").reset_index() + df.attrs["per_covariate"] = agg + return df + + +def phenotype_calibration(adata, *, n_bins=10): + """Reliability of predict() probabilities: bin by predicted max-prob, compare mean + predicted prob to empirical accuracy per bin; scalar ECE in ``df.attrs['ECE']``. adata-only.""" + probs = np.asarray(adata.obsm[K.X_PROBABILITIES], dtype=float) + phenos = list(adata.uns[K.PHENOTYPE_CATEGORIES]) + conf = probs.max(axis=1) + pred = np.asarray(phenos)[probs.argmax(axis=1)] + true = adata.obs[adata.uns[K.METADATA]["phenotype_col"]].astype(str).to_numpy() + correct = (pred == true).astype(float) + + edges = np.linspace(0.0, 1.0, n_bins + 1) + rows = [] + ece = 0.0 + N = max(len(conf), 1) + for b in range(n_bins): + lo, hi = edges[b], edges[b + 1] + m = (conf >= lo) & (conf <= hi if b == n_bins - 1 else conf < hi) + n = int(m.sum()) + if n == 0: + rows.append({"bin": b, "mean_pred": np.nan, "emp_freq": np.nan, "count": 0}) + continue + mp, acc = float(conf[m].mean()), float(correct[m].mean()) + rows.append({"bin": b, "mean_pred": mp, "emp_freq": acc, "count": n}) + ece += (n / N) * abs(acc - mp) + df = pd.DataFrame(rows) + df.attrs["ECE"] = float(ece) + return df + + +def reconstruction_ppc(model, adata=None, *, n_sims=100, random_state=0): + """ZINB reconstruction PPC: simulate counts from the fitted decoder and compare library + size / dropout / mean / variance vs observed. ``model`` REQUIRED (live decoder).""" + import torch + import pyro.distributions as dist + from scvi import REGISTRY_KEYS + + module = model.module + module.eval() + adata = model._validate_anndata(adata) + device = next(module.parameters()).device + torch.manual_seed(int(random_state)) + loader = model._make_data_loader(adata=adata, batch_size=256) + + obs_lib, sim_lib, obs_all, sim_all = [], [], [], [] + with torch.no_grad(): + for tensors in loader: + x = tensors[REGISTRY_KEYS.X_KEY].to(device) + b = tensors[REGISTRY_KEYS.BATCH_KEY].long().to(device) + log_lib = torch.log(x.sum(1, keepdim=True) + 1e-6) + z_loc, _, _ = module.encoder(x, b) + _px_scale, _px_r, px_rate, px_dropout = module.decoder("gene", z_loc, log_lib, b) + gate = torch.sigmoid(px_dropout).clamp(1e-3, 1 - 1e-3) + # mirror the module's generative distribution exactly (same eps + [-10,10] clamp, + # tcri/model/_module.py) so the PPC samples from what the model actually defines + nb_logits = (px_rate + module.eps).log() - (module.px_r.exp() + module.eps).log() + nb_logits = torch.clamp(nb_logits, min=-10.0, max=10.0) + total = module.px_r.exp().clamp(max=1e4) + xd = dist.ZeroInflatedNegativeBinomial(gate=gate, total_count=total, + logits=nb_logits, validate_args=False) + sim = xd.sample() + obs_lib.append(x.sum(1).cpu().numpy()); sim_lib.append(sim.sum(1).cpu().numpy()) + obs_all.append(x.cpu().numpy()); sim_all.append(sim.cpu().numpy()) + + ol, sl = np.concatenate(obs_lib), np.concatenate(sim_lib) + oa, sa = np.concatenate(obs_all), np.concatenate(sim_all) + stats = { + "mean_library_size": (float(ol.mean()), float(sl.mean())), + "median_library_size": (float(np.median(ol)), float(np.median(sl))), + "dropout_fraction": (float((oa == 0).mean()), float((sa == 0).mean())), + "mean_expression": (float(oa.mean()), float(sa.mean())), + "var_expression": (float(oa.var()), float(sa.var())), + } + return pd.DataFrame([ + {"statistic": k, "observed": o, "simulated": s, "discrepancy": abs(o - s)} + for k, (o, s) in stats.items() + ]) + + +def permutation_null(adata, *, metric="mutual_information", covariate=None, groupby=None, + n_perm=1000, random_state=None): + """Permutation null for a clone↔phenotype metric: permute phenotype labels within each + covariate, recompute the metric on the **empirical** clone×phenotype joint to form a null; + report observed, null mean/sd, z, p. adata-only, model-free.""" + from ..tools._mutual_information import _mi_from_joint + + if metric != "mutual_information": + raise ValueError("permutation_null currently supports metric='mutual_information'.") + rng = np.random.default_rng(random_state) + meta = adata.uns[K.METADATA] + clone_col, pheno_col, cov_col = meta["clone_col"], meta["phenotype_col"], meta["covariate_col"] + phenos = list(adata.uns[K.PHENOTYPE_CATEGORIES]) + pheno_index = {p: i for i, p in enumerate(phenos)} + covs = [covariate] if covariate is not None else list(adata.uns[K.COVARIATE_CATEGORIES]) + + n_phenos = len(phenos) + + def _empirical_mi(clones_codes, pheno_codes, n_clones): + # bincount on the flattened (clone, phenotype) key rather than np.add.at — + # identical counts, ~3.7x faster (np.add.at is the unbuffered ufunc path). + flat = np.bincount( + clones_codes * n_phenos + pheno_codes, minlength=n_clones * n_phenos + ) + J = flat.astype(float).reshape(n_clones, n_phenos) + return _mi_from_joint(J, normalized=True, mode="min") + + rows = [] + for m in covs: + cmask = (adata.obs[cov_col].astype(str) == str(m)).to_numpy() + clones = adata.obs.loc[cmask, clone_col].astype(str).to_numpy() + uc = {c: i for i, c in enumerate(sorted(set(clones)))} + cc = np.array([uc[c] for c in clones]) + pc = adata.obs.loc[cmask, pheno_col].astype(str).map(pheno_index).to_numpy() + if len(cc) == 0: + continue + obs_mi = _empirical_mi(cc, pc, len(uc)) + null = np.array([_empirical_mi(cc, rng.permutation(pc), len(uc)) for _ in range(int(n_perm))]) + mu, sd = float(null.mean()), float(null.std(ddof=1) if null.size > 1 else 0.0) + z = (obs_mi - mu) / sd if sd > 0 else np.nan + p = float(np.mean(null >= obs_mi)) + rows.append({"covariate": m, "observed": obs_mi, "null_mean": mu, "null_sd": sd, + "z": z, "p": p}) + return pd.DataFrame(rows) diff --git a/tcri/diagnostics/_training.py b/tcri/diagnostics/_training.py new file mode 100644 index 0000000..745927b --- /dev/null +++ b/tcri/diagnostics/_training.py @@ -0,0 +1,78 @@ +"""``tcri.diag`` training diagnostics (§9.2) — the two model plots relocated off +``TCRIModel`` (``plot_loss`` → :func:`loss`, ``plot_archetypes`` → :func:`archetypes`).""" +from __future__ import annotations + +import numpy as np + +__all__ = ["loss", "archetypes"] + + +def _series(hist, key): + """Pull a 1-D value array from an scvi ``history_`` entry (DataFrame/Series/list).""" + v = hist.get(key) if hasattr(hist, "get") else None + if v is None: + return [] + if hasattr(v, "values"): + arr = np.asarray(v.values).ravel() + return arr.tolist() + return list(v) + + +def loss(model, *, log_scale=False, ax=None, save=None): + """Plot training/validation ELBO and prior-KL from ``model.history_``.""" + import matplotlib.pyplot as plt + + hist = getattr(model, "history_", {}) or {} + elbo_train = _series(hist, "elbo_train") + elbo_val = _series(hist, "elbo_validation") + kl_train = _series(hist, "kl_divergence_with_prior_train_epoch") + kl_val = _series(hist, "kl_divergence_with_prior_val") + + if ax is None: + fig, axes = plt.subplots(2, 1, figsize=(10, 10)) + else: + fig = ax.figure + axes = [ax, ax.figure.add_subplot(212)] if len(ax.figure.axes) < 2 else ax.figure.axes[:2] + + axes[0].plot(elbo_train, label="train ELBO") + if elbo_val: + axes[0].plot(elbo_val, label="val ELBO") + axes[0].set_xlabel("epoch"); axes[0].set_ylabel("ELBO"); axes[0].set_title("ELBO"); axes[0].legend() + if kl_train or kl_val: + if kl_train: + axes[1].plot(kl_train, label="train KL(prior)") + if kl_val: + axes[1].plot(kl_val, label="val KL(prior)") + axes[1].set_xlabel("epoch"); axes[1].set_ylabel("KL"); axes[1].set_title("prior KL"); axes[1].legend() + if log_scale: + for a in axes: + a.set_yscale("log") + fig.tight_layout() + if save: + fig.savefig(save, bbox_inches="tight", dpi=150) + return axes[0] + + +def archetypes(model, *, ax=None, save=None): + """Cluster-ordered clone×phenotype heatmap + archetype centroids, ordered by the + ``build_archetypes`` labels retained on the model.""" + import matplotlib.pyplot as plt + + labels = np.asarray(model.labels) + prior = np.asarray(model.clone_phenotype_prior) + centers = np.asarray(model.centers) + order = np.argsort(labels) + + fig, axes = plt.subplots(1, 2, figsize=(12, 5)) if ax is None else (ax.figure, [ax, ax.figure.add_subplot(122)]) + im0 = axes[0].imshow(prior[order, :], aspect="auto", cmap="viridis") + axes[0].set_title("clone phenotype prior (cluster-ordered)") + axes[0].set_xlabel("phenotype"); axes[0].set_ylabel("clone (by cluster)") + fig.colorbar(im0, ax=axes[0]) + im1 = axes[1].imshow(centers, aspect="auto", cmap="viridis") + axes[1].set_title("archetype centroids") + axes[1].set_xlabel("phenotype"); axes[1].set_ylabel("archetype") + fig.colorbar(im1, ax=axes[1]) + fig.tight_layout() + if save: + fig.savefig(save, bbox_inches="tight", dpi=150) + return axes[0] diff --git a/tcri/metrics/__init__.py b/tcri/metrics/__init__.py deleted file mode 100644 index 6195be4..0000000 --- a/tcri/metrics/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from ..metrics._metrics import * diff --git a/tcri/metrics/_metrics.py b/tcri/metrics/_metrics.py deleted file mode 100644 index 6ff3d51..0000000 --- a/tcri/metrics/_metrics.py +++ /dev/null @@ -1,1000 +0,0 @@ -# Standard library imports -import warnings -from .. import _keys as K -import numpy as np -import pandas as pd -import torch -import torch.nn.functional as F -import scanpy as sc -from typing import Optional, List, Union - -# Third-party imports -from scipy.stats import entropy -from scipy.spatial import distance -import gseapy as gp -import numpy as np, pandas as pd, torch, umap -from tqdm.auto import tqdm -from scvi import REGISTRY_KEYS -import itertools -import collections -# Local imports -from ..preprocessing._preprocessing import joint_distribution,joint_distribution_posterior - -warnings.filterwarnings('ignore') - -import numpy as np, pandas as pd -from typing import Optional, List -from tqdm.auto import tqdm # nice progress bar in notebooks -# ╭──────────────────────────────────────────────────────────────╮ -# │ Δ-E N T R O P Y T A B L E B U I L D E R (v2) │ -# ╰──────────────────────────────────────────────────────────────╯ -import numpy as np, pandas as pd -from typing import Optional, List -from tqdm.auto import tqdm - - -# ------------ simple ANSI helpers ------------ # -RESET = "\x1b[0m" -BOLD = "\x1b[1m" -DIM = "\x1b[2m" -GREEN = "\x1b[32m" -CYAN = "\x1b[36m" -MAGENT = "\x1b[35m" - -# ╭─ colour / pretty-print helpers ─────────────────────────────────────────╮ -RESET = "\x1b[0m"; BOLD = "\x1b[1m"; DIM = "\x1b[2m" -GRN = "\x1b[32m"; CYN = "\x1b[36m"; MAG = "\x1b[35m"; YLW = "\x1b[33m"; RED = "\x1b[31m" - -from .._console import _ok, _info, _warn, _fin -# ╰──────────────────────────────────────────────────────────────────────────╯ - - -# ╭─ tiny ASCII histogram (handy in notebooks/SSH) ──────────────────────────╮ -def _ascii_hist(samples, bins=25, width=40) -> str: - hist, edges = np.histogram(samples, bins=bins) - top = hist.max() - lines=[] - for h,e0,e1 in zip(hist, edges[:-1], edges[1:]): - bar = "█"*int(width*h/top) if top else "" - lines.append(f"{e0:7.3f}-{e1:7.3f} | {bar}") - return "\n".join(lines) -# ╰───────────────── - - -# ╭─ MI helper (single source of truth) ─────────────────────────────────────╮ -def _mi_from_joint(pxy: np.ndarray, normalised: bool, mode: str="average") -> float: - """ - Mutual information (optionally normalised) from an *already normalised* - joint table pxy (shape C×P). - """ - eps = 1e-15 - px = pxy.sum(1, keepdims=True) - py = pxy.sum(0, keepdims=True) - mi = np.sum(pxy * np.log2((pxy+eps) / (px @ py + eps))) - - if not normalised: - return mi - h_c = -np.sum(px * np.log2(px+eps)) - h_p = -np.sum(py * np.log2(py+eps)) - denom = 0.5*(h_c+h_p) if mode == "average" else min(h_c, h_p) - return mi/denom if denom > 0 else 0.0 -# ╰──────────────────────────────────────────────────────────────────────────╯ - - -def mi_compare(adata, groupby, groups=None, treatment=None, n_samples=50, - patient_col=None, clone_col=None, covariate_col=None, - verbose=True, **mi_kwargs): - meta = adata.uns[K.METADATA] - patient_col = patient_col or meta["batch_col"] - clone_col = clone_col or meta["clone_col"] - covariate_col = covariate_col or meta["covariate_col"] - - covariates = treatment - if covariates is None: - covariates = adata.obs[covariate_col].cat.categories.tolist() - elif isinstance(covariates, str): - covariates = [covariates] - - # resolve groups into list of 2-tuples - unique_groups = adata.obs[groupby].dropna().unique().tolist() - if groups is None: - pairs = list(itertools.combinations(sorted(unique_groups), 2)) - elif isinstance(groups[0], (list, tuple)): - pairs = [tuple(g) for g in groups] - else: - pairs = list(itertools.combinations(groups, 2)) - - # compute patient-level MI samples - keep_groups = set(g for pair in pairs for g in pair) - records = [] - patients = adata.obs[patient_col].unique() - for p in tqdm(patients, disable=not verbose, desc="MI per patient"): - pmask = adata.obs[patient_col] == p - group_val = adata.obs.loc[pmask, groupby].iloc[0] - if group_val not in keep_groups: - continue - clones = adata.obs.loc[pmask, clone_col].unique().tolist() - for cov in covariates: - try: - samples = mutual_information( - adata, cov, n_samples=n_samples, - clones=clones, verbose=False, **mi_kwargs - ) - except Exception as e: - if verbose: - print(f" Skip {p}/{cov}: {e}") - continue - for s in np.atleast_1d(samples): - records.append({"patient": p, "group": group_val, - "covariate": cov, "MI": float(s)}) - - result = pd.DataFrame(records) - summary = ( - result.groupby(["patient", "group", "covariate"])["MI"] - .agg(mean="mean", median="median", - lo=lambda x: np.quantile(x, 0.025), - hi=lambda x: np.quantile(x, 0.975), - sd="std", n="size") - .reset_index() - ) - - return { - "samples": result, - "summary": summary, - "pairs": pairs, - "params": {"groupby": groupby, "covariates": covariates, - "n_samples": n_samples, "patient_col": patient_col}, - } - - -def dkl(p, q): - epsilon = 1e-10 - p = np.clip(p, epsilon, 1) - q = np.clip(q, epsilon, 1) - return entropy(p, q) - -import numpy as np, pandas as pd -from scipy.stats import entropy # Shannon entropy -from typing import Optional, List, Union -import numpy as np, pandas as pd -from typing import Optional, List, Union -from scipy.stats import entropy # Shannon H - -# ---------- ANSI helpers ------------------------------------------- -RESET="\x1b[0m"; BOLD="\x1b[1m"; DIM="\x1b[2m" -GRN="\x1b[32m"; CYN="\x1b[36m"; MAG="\x1b[35m"; YLW="\x1b[33m" - -# small helper ---------------------------------------------- -def _ent(p, base=2): - p = np.asarray(p, dtype=float) - eps = 1e-15 - p = p.clip(eps) / p.sum() - return entropy(p, base=base) - -def clonotypic_entropy_base( - adata, - covariate_label: str, - phenotype: str, - *, - base: int = 2, - normalised: bool = True, - temperature: float = 1.0, - clones: Optional[List[str]] = None, - weighted: bool = False, - posterior: bool = True, - combine_with_logits: bool = True, -) -> float: - if posterior: - jd = joint_distribution_posterior( - adata, - covariate_label = covariate_label, - temperature = temperature, - clones = clones, - weighted = weighted, - combine_with_logits = combine_with_logits, - silent = True) - else: - jd = joint_distribution( - adata, - covariate_label = covariate_label, - temperature = temperature, - n_samples = 0, - clones = clones, - weighted = weighted) - - if jd is None or jd.empty or phenotype not in jd.columns: - return 0.0 - - vec = jd[phenotype].to_numpy(dtype=float) - eps = 1e-15 - vec = np.clip(vec, eps, None) - vec = vec / vec.sum() # ensure ∑=1 - - H = entropy(vec, base=base) - if normalised and len(vec) > 1: - H /= np.log(len(vec)) / np.log(base) # # max-entropy normalisation - return H - -def clonotypic_entropy( - adata, - covariate: str, - *, - point_estimate: bool = True, - n_samples: int = 200, - temperature: float = 1.0, - combine_with_logits: bool = True, - _clones: Optional[List[str]] = None, -) -> Union[pd.Series, np.ndarray]: - r"""Clonotypic entropy of each phenotype at one covariate value. - - For a phenotype :math:`\phi`, this is the normalized Shannon entropy of the - distribution over clonotypes carrying that phenotype, - :math:`H\!\left[P(c \mid \phi,\, m)\right]`, estimated from posterior draws of - the clone–phenotype joint distribution. High values mean the phenotype is - spread across many clones; low values mean a few clones dominate it. - - Parameters - ---------- - adata : AnnData - Registered object (see - :func:`~tcri.preprocessing._preprocessing.register_model`). - covariate : str - Covariate value :math:`m` to condition on (a category of the registered - covariate column). - point_estimate : bool, default True - If True, return the posterior-mean entropy per phenotype; if False, return - the full matrix of per-draw entropies. - n_samples : int, default 200 - Number of posterior draws to average over. Must be ``>= 1``. - temperature : float, default 1.0 - Sharpen (``<1``) or flatten (``>1``) the per-cell distribution before - aggregating. - combine_with_logits : bool, default True - Combine the sampled prior :math:`p_{ct}` with the per-cell classifier - logits (the full posterior) rather than the prior alone. - - Returns - ------- - pandas.Series or numpy.ndarray - If ``point_estimate`` is True, a Series indexed by phenotype name whose - values are the mean entropy in bits, normalized to :math:`[0, 1]` by - :math:`\log_2 n_\text{clones}`, over ``n_samples`` draws. Otherwise an - array of shape ``(n_samples, n_phenotypes)``. - - Raises - ------ - ValueError - If ``n_samples < 1``. - - See Also - -------- - phenotypic_entropy : the per-clone analogue, :math:`H[P(\phi \mid c, m)]`. - mutual_information : clone–phenotype coupling at a covariate. - - Examples - -------- - >>> covariate = adata.uns["tcri_covariate_categories"][0] - >>> ce = clonotypic_entropy(adata, covariate, n_samples=50) - >>> ce.sort_values(ascending=False).head() - """ - if n_samples < 1: - raise ValueError("n_samples must be >= 1") - - phenotypes = list(adata.uns[K.PHENOTYPE_CATEGORIES]) - n_pheno = len(phenotypes) - samples = np.empty((n_samples, n_pheno), dtype=float) - - for i in range(n_samples): - jd = joint_distribution_posterior( - adata, - covariate_label = covariate, - temperature = temperature, - clones = _clones, - combine_with_logits = combine_with_logits, - silent = True, - ) - if jd is None or jd.empty: - samples[i, :] = np.nan - continue - for j, p in enumerate(phenotypes): - if p not in jd.columns: - samples[i, j] = np.nan - continue - vec = jd[p].to_numpy(dtype=float) - vec = np.clip(vec, 1e-15, None) - vec = vec / vec.sum() - H = entropy(vec, base=2) - if len(vec) > 1: - H /= np.log2(len(vec)) - samples[i, j] = H - - if point_estimate: - return pd.Series(np.nanmean(samples, axis=0), index=phenotypes, name=covariate) - return samples - -def delta_clonotypic_entropy( - adata, - phenotype: str, - *, - cov_pre: str = "Pre-treatment", - cov_post: str = "Post-treatment", - n_samples: int = 1_000, - temperature: float = 1.0, - clones: Optional[List[str]] = None, - weighted: bool = False, - normalised: bool = True, - base: int = 2, - posterior: bool = True, - combine_with_logits: bool = True, - verbose: bool = True, - graph: bool = False, # ASCII plot of Δ posterior - seed: Optional[int] = None, -) -> np.ndarray: - """ - Sample Δ-entropy = H_post – H_pre for one phenotype. - - Returns - ------- - delta_samples : ndarray (shape = (n_samples,)) - Positive values ⇒ entropy increased from pre → post. - """ - if seed is not None: - np.random.seed(seed) - - if verbose: - print(f"{BOLD}{MAG}Δ-Entropy {phenotype}: " - f"'{cov_pre}' ⟶ '{cov_post}'{RESET}") - _info("# samples", n_samples) - _info("posterior", posterior) - _info("weighted", weighted) - - Δ = np.empty(n_samples, dtype=float) - - # --- Monte Carlo loop ----------------------------------------- - for i in range(n_samples): - H_pre = clonotypic_entropy_base( - adata, cov_pre, phenotype, - base=base, normalised=normalised, - temperature=temperature, - clones=clones, weighted=weighted, - posterior=posterior, - combine_with_logits=combine_with_logits) - - H_post = clonotypic_entropy_base( - adata, cov_post, phenotype, - base=base, normalised=normalised, - temperature=temperature, - clones=clones, weighted=weighted, - posterior=posterior, - combine_with_logits=combine_with_logits) - Δ[i] = H_post - H_pre - - if verbose: - _ok("sampling complete") - _info("mean ± sd", f"{Δ.mean():.4f} ± {Δ.std():.4f}") - lo,hi = np.percentile(Δ,[2.5,97.5]) - _info("95 % CI", f"[{lo:.4f}, {hi:.4f}]") - if graph: - print(f"{DIM}\nASCII histogram of Δ:\n{_ascii_hist(Δ)}{RESET}") - - return Δ - - -def delta_entropy_table( - adata, - *, - cov_pre : str = "Pre-treatment", - cov_post: str = "Post-treatment", - splitby : str = "response", - n_samples : int = 1_000, - temperature: float = 1.0, - weighted : bool = False, - normalised : bool = True, - base : int = 2, - posterior : bool = True, - combine_with_logits : bool = True, - seed : Optional[int] = 42, - show_progress: bool = True -) -> pd.DataFrame: - """ - Build a tidy Δ-clonotypic-entropy table (post – pre). - - Each row ⇢ one phenotype × one `splitby` group. - The `delta_samples` column keeps the full NumPy vector so you can - re-plot KDEs or run further stats without re-sampling. - """ - # reproducibility - if seed is not None: - np.random.seed(seed) - - meta = adata.uns[K.METADATA] - clone_col = meta["clone_col"] - phen_col = meta["phenotype_col"] - - groups = sorted(adata.obs[splitby].dropna().unique().tolist()) - phenotypes = adata.obs[phen_col].astype("category").cat.categories.tolist() - - records = [] - iterator = tqdm(groups, desc="Δ-entropy groups") if show_progress else groups - - for g in iterator: - # — restrict ONLY the clone list, keep full AnnData for index integrity - mask_g = adata.obs[splitby] == g - clones_g = adata.obs.loc[mask_g, clone_col].unique().tolist() - - for ph in phenotypes: - delta = delta_clonotypic_entropy( - adata, ph, - cov_pre = cov_pre, - cov_post = cov_post, - n_samples = n_samples, - temperature = temperature, - clones = clones_g, # ⬅ scoped clones - weighted = weighted, - normalised = normalised, - base = base, - posterior = posterior, - combine_with_logits = combine_with_logits, - verbose = False - ) - - d_mean = delta.mean() - d_sd = delta.std() - hdi_lo, hdi_hi = np.percentile(delta, [2.5, 97.5]) - p_gt = (delta > 0).mean() - p_lt = (delta < 0).mean() - - records.append(dict( - **{splitby: g, "phenotype": ph}, - delta_samples = delta, - delta_mean = d_mean, - delta_sd = d_sd, - hdi_low = hdi_lo, - hdi_high = hdi_hi, - p_greater = p_gt, - p_less = p_lt - )) - - return pd.DataFrame.from_records(records) - - -def phenotypic_entropy( - adata, - covariate: str, - *, - point_estimate: bool = True, - n_samples: int = 200, - temperature: float = 1.0, - combine_with_logits: bool = True, -) -> Union[pd.Series, np.ndarray]: - r"""Phenotypic entropy of each clonotype at one covariate value. - - For a clonotype :math:`c`, this is the normalized Shannon entropy of its - distribution over phenotypes, :math:`H\!\left[P(\phi \mid c,\, m)\right]`, - estimated from posterior draws. High values mean the clone is phenotypically - plastic; low values mean it is committed to one phenotype. - - Parameters - ---------- - adata : AnnData - Registered object (see - :func:`~tcri.preprocessing._preprocessing.register_model`). - covariate : str - Covariate value :math:`m` to condition on. - point_estimate : bool, default True - If True, return the posterior-mean entropy per clone; if False, the full - per-draw matrix. - n_samples : int, default 200 - Number of posterior draws to average over. Must be ``>= 1``. - temperature : float, default 1.0 - Sharpen (``<1``) or flatten (``>1``) the per-cell distribution. - combine_with_logits : bool, default True - Combine the sampled prior with the per-cell classifier logits. - - Returns - ------- - pandas.Series or numpy.ndarray - If ``point_estimate`` is True, a Series indexed by clonotype, in bits and - normalized to :math:`[0, 1]` by :math:`\log_2 n_\text{phenotypes}`. - Otherwise an array of shape ``(n_samples, n_clones)`` over the clones - present at ``covariate``. - - Raises - ------ - ValueError - If ``n_samples < 1``. - - See Also - -------- - clonotypic_entropy : the per-phenotype analogue. - flux : change in a clone's phenotype distribution between two covariates. - - Examples - -------- - >>> covariate = adata.uns["tcri_covariate_categories"][0] - >>> pe = phenotypic_entropy(adata, covariate, n_samples=50) - >>> pe.mean() # average phenotypic plasticity across clones - """ - if n_samples < 1: - raise ValueError("n_samples must be >= 1") - - meta = adata.uns[K.METADATA] - clone_col = meta["clone_col"] - covariate_col = meta["covariate_col"] - - clones_list = ( - adata.obs.loc[adata.obs[covariate_col] == covariate, clone_col] - .unique() - .tolist() - ) - - if len(clones_list) == 0: - if point_estimate: - return pd.Series(dtype=float, name=covariate) - return np.empty((n_samples, 0), dtype=float) - - n_clones = len(clones_list) - samples = np.empty((n_samples, n_clones), dtype=float) - - for i in range(n_samples): - jd = joint_distribution_posterior( - adata, - covariate_label = covariate, - temperature = temperature, - clones = clones_list, - combine_with_logits = combine_with_logits, - silent = True, - ) - if jd is None or jd.empty: - samples[i, :] = np.nan - continue - n_phen = jd.shape[1] - norm = np.log2(n_phen) if n_phen > 1 else 1.0 - for j, cl in enumerate(clones_list): - if cl not in jd.index: - samples[i, j] = np.nan - continue - p = jd.loc[cl].to_numpy(dtype=float) - p = np.clip(p, 1e-15, None) - p = p / p.sum() - samples[i, j] = entropy(p, base=2) / norm - - if point_estimate: - return pd.Series(np.nanmean(samples, axis=0), index=clones_list, name=covariate) - return samples - - -def clonality(adata): - r"""Phenotype clonality: how clonally concentrated each phenotype is. - - For each phenotype, clonality is :math:`1 - H / \log_2 K`, where :math:`H` is - the Shannon entropy (bits) of the clone-size distribution among cells of that - phenotype and :math:`K` is the number of distinct clones. It is the normalized - complement of clonotypic entropy: ``1`` means a single clone dominates the - phenotype, ``0`` means all clones are equally represented. - - Parameters - ---------- - adata : AnnData - Registered object. Uses the observed (hard) clone and phenotype labels - rather than the posterior, so it only needs the registered category keys. - - Returns - ------- - dict of {str: float} - Maps each phenotype to its clonality in :math:`[0, 1]`. - - See Also - -------- - clonotypic_entropy : the soft, posterior per-phenotype entropy. - - Examples - -------- - >>> clonality(adata) - {'A': 0.31, 'B': 0.07, 'C': 0.52} - """ - phenotypes = adata.obs[adata.uns["tcri_phenotype_key"]].tolist() - unique_phenotypes = np.unique(phenotypes) - entropys = dict() - df = adata.obs.copy() - for phenotype in unique_phenotypes: - pheno_df = df[df[adata.uns['tcri_phenotype_key']]==phenotype] - clonotypes = pheno_df[adata.uns["tcri_clone_key"]].tolist() - unique_clonotypes = np.unique(clonotypes) - nums = [] - for tcr in unique_clonotypes: - tcrs = pheno_df[pheno_df[adata.uns["tcri_clone_key"]]==tcr] - nums.append(len(tcrs)) - clonality = 1 - entropy(np.array(nums),base=2) / np.log2(len(nums)) - entropys[phenotype] = np.nan_to_num(clonality) - return entropys - -def clone_fraction(adata, groupby): - frequencies = dict() - for group in set(adata.obs[groupby]): - frequencies[group] = dict() - sdata = adata[adata.obs[groupby] == group] - total_cells = len(sdata.obs.index.tolist()) - clones = sdata.obs[sdata.uns["tcri_clone_key"]].tolist() - for c in set(clones): - frequencies[group][c] = clones.count(c) / total_cells - return frequencies - - -def mutual_information( - adata, - covariate: str, - *, - temperature: float = 1.0, - n_samples: int = 0, - clones: Optional[List[str]] = None, - normalised: bool = True, - normalise_mode: str = "average", - posterior: bool = True, - combine_with_logits: bool = True, - verbose: bool = True, - graph: bool = False, -) -> Union[float, np.ndarray]: - r"""Mutual information between clonotype and phenotype at one covariate value. - - :math:`I(c; \phi \mid m)` quantifies how much knowing a cell's clonotype tells - you about its phenotype at covariate :math:`m` — the strength of - clone–phenotype coupling. Zero means clonotype and phenotype are independent; - larger values mean clones are phenotypically structured. - - Parameters - ---------- - adata : AnnData - Registered object (see - :func:`~tcri.preprocessing._preprocessing.register_model`). - covariate : str - Covariate value :math:`m` to condition on. - temperature : float, default 1.0 - Sharpen (``<1``) or flatten (``>1``) the distributions before computing MI. - n_samples : int, default 0 - ``0`` returns a single point estimate; ``> 0`` returns one MI value per - posterior draw. - clones : list of str, optional - Restrict to these clonotypes; default uses all. - normalised : bool, default True - Normalize the MI (see ``normalise_mode``) to :math:`[0, 1]`. - normalise_mode : str, default "average" - Normalization denominator when ``normalised`` is True (e.g. the average of - the two marginal entropies). - posterior : bool, default True - Use a Dirichlet draw of :math:`p_{ct}` (optionally combined with the - classifier logits). If False, use the prior-only joint distribution. - combine_with_logits : bool, default True - Combine the sampled prior with the per-cell logits (only when ``posterior``). - verbose : bool, default True - Print a short progress summary. - graph : bool, default False - Print an ASCII histogram of the posterior MI draws. - - Returns - ------- - float or numpy.ndarray - A single (normalized) MI in bits if ``n_samples == 0``; otherwise an array - of shape ``(n_samples,)`` of per-draw values. - - Notes - ----- - .. math:: - - I(c; \phi) = \sum_{c,\, \phi} p(c, \phi)\, - \log_2 \frac{p(c, \phi)}{p(c)\, p(\phi)} - - See Also - -------- - clonotypic_entropy, phenotypic_entropy : the marginal entropies MI builds on. - - Examples - -------- - >>> covariate = adata.uns["tcri_covariate_categories"][0] - >>> mutual_information(adata, covariate, verbose=False) - 0.42 - """ - - if verbose: - print(f"{BOLD}{MAGENT}📊 MI for '{covariate}'{RESET}") - _info("posterior", posterior) - _info("n_samples", n_samples) - - # helper to obtain *one* joint table - def _get_df(silent_flag: bool): - if posterior: - return joint_distribution_posterior( - adata, - covariate_label = covariate, - temperature = temperature, - clones = clones, - combine_with_logits = combine_with_logits, - silent = silent_flag, - ) - else: - # fall back to prior-only version (not shown here) - raise NotImplementedError("prior-only joint_distribution not included") - - # ---------- single draw ---------------------------------------- - if n_samples == 0: - df = _get_df(silent_flag=not verbose) - pxy = df.to_numpy().astype(float) - pxy /= pxy.sum() - mi = _mi_from_joint(pxy, normalised, normalise_mode) - _ok("computed MI", quiet=not verbose) - _info("value", f"{mi:.4f}", quiet=not verbose) - return mi - - # ---------- multiple draws ------------------------------------- - mi_samples = np.empty(n_samples, dtype=float) - iter_bar = tqdm(range(n_samples), disable=not verbose, leave=False, - desc=f"sampling MI [{covariate}]") - for i in iter_bar: - df = _get_df(silent_flag=True) - pxy = df.to_numpy().astype(float) - pxy /= pxy.sum() - mi_samples[i] = _mi_from_joint(pxy, normalised, normalise_mode) - - _ok("computed MI for all samples", quiet=not verbose) - _info("mean ± sd", - f"{mi_samples.mean():.4f} ± {mi_samples.std():.4f}", - quiet=not verbose) - _info("95% CI", - f"[{np.percentile(mi_samples,2.5):.4f}, " - f"{np.percentile(mi_samples,97.5):.4f}]", - quiet=not verbose) - - if graph and n_samples > 1 and verbose: - print(f"{DIM}\nASCII histogram of MI posterior:\n" - f"{_ascii_hist(mi_samples)}{RESET}") - - return mi_samples - - - -# ──────────────────────────────────────────────────────────────── -def flux_table( - adata, - *, - cov_pre: str = "Pre-treatment", - cov_post: str = "Post-treatment", - splitby: str = "response", - n_samples: int = 0, - temperature: float = 1.0, - weighted: bool = False, - posterior: bool = True, - combine_with_logits: bool = True, - distance_metric: str = "l1", - seed: Optional[int] = 42, - show_progress: bool = True -) -> pd.DataFrame: - """ - Build a tidy table with per-clone flux distances and clone size. - - Nested progress bars: - • outer: groups in `splitby` - • inner: clones within that group - """ - - if seed is not None: - np.random.seed(seed) - - meta = adata.uns[K.METADATA] - clone_col = meta["clone_col"] - - groups = sorted(adata.obs[splitby].dropna().unique().tolist()) - records = [] - - outer_it = groups if not show_progress else tqdm(groups, desc="groups") - - for g in outer_it: - # -------- subset AnnData to *this* group ----------------- - mask_g = adata.obs[splitby] == g - adata_g = adata[mask_g] - clones_g = adata_g.obs[clone_col].unique().tolist() - - # -------- clone sizes (within this group) ---------------- - c_sizes = adata_g.obs[clone_col].value_counts().to_dict() - - # -------- flux distances (vector, index = clone_id) ----- - dist = flux( - adata, - from_this = cov_pre, - to_that = cov_post, - clones = clones_g, - temperature = temperature, - distance_metric = distance_metric, - n_samples = n_samples, - weighted = weighted, - posterior = posterior, - combine_with_logits = combine_with_logits, - graph = False # keep helper quiet - ) - - # nested bar over clones ---------------------------------- - inner_it = clones_g if not show_progress else tqdm( - clones_g, desc=f"{g}: clones", leave=False) - - for cl in inner_it: - if n_samples == 0: - val = float(dist.get(cl, np.nan)) - sd = 0.0 - vec = np.array([val]) - else: - # `dist` is (n_samples, n_clones); retrieve column - idx = clones_g.index(cl) - vec = dist[:, idx] - val = float(vec.mean()) - sd = float(vec.std(ddof=1)) - - records.append(dict( - **{splitby: g, "clone_id": cl}, - flux_samples = vec, - flux_mean = val, - flux_sd = sd, - clone_size = c_sizes.get(cl, 0) - )) - - return pd.DataFrame.from_records(records) - - -# ──────────────────────────────────────────────────────────────── -def flux( - adata, - *, - from_this : str, - to_that : str, - clones : Optional[Union[str, List[str]]] = None, - temperature : float = 1.0, - distance_metric : Union[str, callable] = "l1", # MODIFIED: Type hint updated - n_samples : int = 0, # posterior draws - weighted : bool = False, - posterior : bool = True, - combine_with_logits : bool = True, - graph : bool = False, # ASCII histogram - seed : Optional[int] = 42 -) -> Union[pd.Series, np.ndarray]: - r"""Phenotypic flux of each clonotype between two covariate values. - - For each clonotype, the distance between its phenotype distribution at - ``from_this`` and at ``to_that`` — how much the clone's phenotype mix shifts - across the two covariates. Useful for tracking phenotypic movement over time or - treatment. - - Parameters - ---------- - adata : AnnData - Registered object (see - :func:`~tcri.preprocessing._preprocessing.register_model`). - from_this, to_that : str - The two covariate values to compare (e.g. ``"Pre-treatment"`` and - ``"Post-treatment"``). - clones : str or list of str, optional - Restrict to these clonotypes; default uses all clones present. - temperature : float, default 1.0 - Sharpen (``<1``) or flatten (``>1``) the distributions. - distance_metric : str or callable, default "l1" - Distance between the two phenotype distributions — e.g. ``"l1"`` or - ``"dkl"``, or a callable ``f(p, q) -> float``. - n_samples : int, default 0 - ``0`` returns a point estimate per clone; ``> 0`` returns per-draw values. - weighted : bool, default False - Weight clones by size when building the joint distribution. - posterior : bool, default True - Use posterior draws of :math:`p_{ct}` (vs the prior-only joint). - combine_with_logits : bool, default True - Combine the sampled prior with the per-cell logits (only when ``posterior``). - graph : bool, default False - Print an ASCII histogram of the flux distribution. - seed : int, optional - Seed for the posterior sampling (default 42). - - Returns - ------- - pandas.Series or numpy.ndarray - A Series indexed by clonotype if ``n_samples == 0``; otherwise an array of - shape ``(n_samples, n_clones)`` whose rows are posterior draws. - - Raises - ------ - ValueError - If the requested clones do not overlap the data at both covariates. - - See Also - -------- - phenotypic_entropy : per-clone phenotypic spread at a single covariate. - - Examples - -------- - >>> flux(adata, from_this="T1", to_that="T2").sort_values(ascending=False).head() - """ - if seed is not None: - np.random.seed(seed) - - # ---------- which clones? ---------------------------------- - clone_col = adata.uns[K.METADATA]["clone_col"] - if clones is None: - clones = adata.obs[clone_col].unique().tolist() - elif isinstance(clones, str): - clones = [clones] - - # ---------- get joint tables ------------------------------- - get = joint_distribution_posterior if posterior else joint_distribution - - jd_from = get( - adata, - covariate_label = from_this, - temperature = temperature, - clones = clones, - weighted = weighted, - combine_with_logits = combine_with_logits if posterior else None, - silent = True - ) - jd_to = get( - adata, - covariate_label = to_that, - temperature = temperature, - clones = clones, - weighted = weighted, - combine_with_logits = combine_with_logits if posterior else None, - silent = True - ) - - if jd_from.empty or jd_to.empty: - raise ValueError("No overlap between requested clones and data.") - - common = jd_from.index.intersection(jd_to.index) - - # --- MINIMAL CHANGE BLOCK 1 --- - # Define _dkl once if needed - _dkl = None - if isinstance(distance_metric, str) and distance_metric.lower() == "dkl": - eps = 1e-15 - def dkl_func(p, q): - p = p.clip(eps)/p.sum(); q = q.clip(eps)/q.sum() - return float(np.sum(p*np.log(p/q))) - _dkl = dkl_func - - if isinstance(distance_metric, collections.abc.Callable): - dist = pd.Series( - {cl: distance_metric(jd_from.loc[cl], jd_to.loc[cl]) for cl in common} - ) - elif isinstance(distance_metric, str) and distance_metric.lower() == "l1": - dist = (jd_from.loc[common] - jd_to.loc[common]).abs().sum(axis=1) - elif _dkl is not None: - dist = pd.Series( - {cl: _dkl(jd_from.loc[cl], jd_to.loc[cl]) for cl in common} - ) - else: - raise ValueError("distance_metric must be 'l1', 'dkl', or a callable function.") - # --- END OF CHANGE --- - - # ---------- single / multi-sample handling ----------------- - if n_samples == 0: - return dist - - # posterior draws (one extra call per sample) - samples = np.empty((n_samples, len(common)), dtype=float) - for i in range(n_samples): - jd_from_s = get(adata, covariate_label = from_this, temperature=temperature, - clones=clones, weighted=weighted, - combine_with_logits=combine_with_logits if posterior else None, - silent=True) - jd_to_s = get(adata, covariate_label = to_that, temperature=temperature, - clones=clones, weighted=weighted, - combine_with_logits=combine_with_logits if posterior else None, - silent=True) - - # --- MINIMAL CHANGE BLOCK 2 --- - if isinstance(distance_metric, collections.abc.Callable): - samples[i] = [distance_metric(jd_from_s.loc[c], jd_to_s.loc[c]) for c in common] - elif isinstance(distance_metric, str) and distance_metric.lower() == "l1": - samples[i] = (jd_from_s.loc[common] - jd_to_s.loc[common]).abs().sum(axis=1) - else: # Handles the 'dkl' case using the function defined above - samples[i] = [ _dkl(jd_from_s.loc[c], jd_to_s.loc[c]) for c in common ] - # --- END OF CHANGE --- - - if graph: - print(_ascii_hist(samples.ravel())) - - return samples \ No newline at end of file diff --git a/tcri/model/_classifier.py b/tcri/model/_classifier.py new file mode 100644 index 0000000..8c41e43 --- /dev/null +++ b/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 diff --git a/tcri/model/_model.py b/tcri/model/_model.py index 8188252..ecfe3fe 100644 --- a/tcri/model/_model.py +++ b/tcri/model/_model.py @@ -1,33 +1,40 @@ +"""User-facing scvi-tools model class :class:`TCRIModel`. + +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`, +`predict`, `to_anndata`, `get_p_ct`, ...): + +- :mod:`._module` -- Pyro model/guide (:class:`TCRIModule`) +- :mod:`._priors` -- :class:`MixtureDirichlet`, :class:`VampPrior` +- :mod:`._classifier` -- :class:`PhenotypeClassifier` +- :mod:`._training` -- :class:`UnifiedTrainingPlan`, :func:`build_archetypes` +""" import logging +import os +import warnings + import numpy as np import pandas as pd +import pyro import torch -import matplotlib.pyplot as plt import torch.nn.functional as F -import pyro -import pyro.distributions as dist -import pyro.poutine as poutine -import torch.nn as nn - +import matplotlib.pyplot as plt -from typing import Dict, Optional +from typing import Optional from anndata import AnnData -import os -from scvi.data.fields import CategoricalObsField, LayerField +from scvi import REGISTRY_KEYS from scvi.data import AnnDataManager +from scvi.data.fields import CategoricalObsField, LayerField from scvi.model.base import BaseModelClass -from scvi.train import PyroTrainingPlan, TrainRunner -from scvi import REGISTRY_KEYS -from scvi.nn import Encoder, DecoderSCVI -from scvi.module.base import PyroBaseModuleClass, auto_move_data -from scvi.utils import setup_anndata_dsp +from scvi.train import TrainRunner from scvi.dataloaders import DataSplitter -from torch.nn.functional import cosine_similarity -from pyro.infer import TraceEnum_ELBO, Trace_ELBO -from torch.distributions import Categorical, Dirichlet, MixtureSameFamily -from sklearn.cluster import KMeans -import warnings + +from ._module import TCRIModule +from ._training import UnifiedTrainingPlan, build_archetypes + +__all__ = ["TCRIModel"] warnings.filterwarnings("ignore", category=UserWarning, message="Found auxiliary vars") warnings.filterwarnings( @@ -41,637 +48,27 @@ logger = logging.getLogger(__name__) - -def build_archetypes(c2p_mat, K=4): - kmeans = KMeans(n_clusters=K, random_state=42) - labels = kmeans.fit_predict(c2p_mat) - centers = kmeans.cluster_centers_ - centers = np.clip(centers, 1e-8, None) - centers = centers / centers.sum(axis=1, keepdims=True) - return centers, labels - - -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 - - -class VampPrior(torch.nn.Module): - def __init__(self, pseudo_inputs, encoder): - """ - Args: - pseudo_inputs (torch.Tensor): Initial pseudo-inputs of shape (K, input_dim). - encoder (torch.nn.Module): Encoder that takes an input and returns (mean, log_var) - for the approximate posterior q(z|x). - """ - super(VampPrior, self).__init__() - # Learnable pseudo-inputs; these are optimized during training. - self.pseudo_inputs = torch.nn.Parameter(pseudo_inputs) - self.encoder = encoder - - def get_mixture(self): - """ - Constructs the VampPrior as a uniform mixture of q(z|u_k) for each pseudo-input u_k. - """ - # Compute the approximate posterior parameters for each pseudo-input. - # Expected output shapes: means and log_vars: (K, latent_dim) - K = self.pseudo_inputs.size(0) - # Create a dummy categorical argument; unsqueeze to have shape (K, 1) - dummy_batch = torch.zeros(K, dtype=torch.long, device=self.pseudo_inputs.device).unsqueeze(1) - means, log_vars, _ = self.encoder(self.pseudo_inputs, dummy_batch) - scales = torch.sqrt(torch.exp(log_vars)) - component_dist = dist.Independent(dist.Normal(means, scales), 1) - mixture_weights = torch.ones(K, device=self.pseudo_inputs.device) / K - mixture = dist.MixtureSameFamily( - dist.Categorical(mixture_weights), - component_dist - ) - return mixture - - def log_prob(self, z): - """ - Computes log p(z) under the VampPrior. - """ - return self.get_mixture().log_prob(z) - - def sample(self, sample_shape=torch.Size()): - """ - Draws samples from the VampPrior. - """ - return self.get_mixture().sample(sample_shape) - - -############################################################################### -# 0) Mixture of Dirichlet Distributions - TODO: Refactor -############################################################################### -class MixtureDirichlet(dist.TorchDistribution): - """ - Mixture of Dirichlet distributions parameterized by mixture weights and concentration parameters." - """ - arg_constraints = { - "mixture_weights": dist.constraints.simplex, # shape: batch_shape + (B,) - "concentration": dist.constraints.positive, # shape: batch_shape + (B, K) - } - support = dist.constraints.simplex # each sample is a simplex over K categories - has_rsample = False - - def __init__( - self, - mixture_weights: torch.Tensor, - concentration: torch.Tensor, - validate_args=None, - ): - """ - mixture_weights: Tensor of shape batch_shape + (B,), with each row summing to 1. - concentration: Tensor of shape batch_shape + (B, K), where K is the number of categories. - """ - self.mixture_weights = mixture_weights - # Clamp concentrations to ensure positivity. - self.concentration = torch.clamp(concentration, min=1e-3) - # Determine batch shape, B, and K. - batch_shape = self.mixture_weights.shape[:-1] - self.B = self.mixture_weights.size(-1) - self.K = self.concentration.size(-1) - event_shape = (self.K,) - super(MixtureDirichlet, self).__init__( - batch_shape, event_shape, validate_args=validate_args - ) - - def sample(self, sample_shape=torch.Size()): - """ - Returns a sample of shape: sample_shape + batch_shape + (K,). - For each batch element, first sample a mixture component, then sample from the corresponding Dirichlet. - """ - # Create categorical for mixture weights. - cat = dist.Categorical(self.mixture_weights) - # Sample mixture indices; shape: sample_shape + batch_shape. - mixture_idx = cat.sample(sample_shape) - full_shape = mixture_idx.shape # sample_shape + batch_shape - - # Expand concentration to shape: sample_shape + batch_shape + (B, K). - target_shape = sample_shape + self.concentration.shape - expanded_concentration = self.concentration.expand(target_shape) - - # Flatten the sample and batch dimensions. - flat_shape = (-1, self.B, self.K) - flat_concentration = expanded_concentration.reshape(flat_shape) - flat_idx = mixture_idx.reshape(-1) # shape: (num_samples,) - - # Select the concentration parameters corresponding to the sampled mixture index. - selected_concentration = flat_concentration[ - torch.arange(flat_idx.size(0)), flat_idx - ] - - # Sample from the Dirichlet for each sample. - flat_samples = dist.Dirichlet(selected_concentration).sample() - # Reshape to sample_shape + batch_shape + (K,). - return flat_samples.reshape(full_shape + (self.K,)) - - def log_prob(self, value): - device = value.device # get the device from input tensor - - # Move tensors explicitly to the same device - value_expanded = value.unsqueeze(-2).to(device) - expanded_concentration = self.concentration.expand( - value.shape[:-1] + (self.B, self.K) - ).to(device) - - d = dist.Dirichlet(expanded_concentration) - - component_log_probs = d.log_prob( - value_expanded.expand(expanded_concentration.shape) - ) - - expanded_weights = self.mixture_weights.expand(value.shape[:-1] + (self.B,)).to(device) - mixture_log = torch.log(expanded_weights) - - return torch.logsumexp(mixture_log + component_log_probs, dim=-1) - - def score_parts(self, value): - # Compute log probability. - lp = self.log_prob(value) - # Return dummy zeros for the score function and entropy terms. - zeros = torch.zeros_like(lp) - return lp, zeros, zeros - - def __call__(self, *args, **kwargs): - return self.sample(*args, **kwargs) - - -############################################################################### -# 1) Pyro Module with CVAE + Hierarchical Priors -############################################################################### -class TCRIModule(PyroBaseModuleClass): - """ - Two-level model that incorporates hierarchical priors (clonotype-level) - and a CVAE structure that explicitly conditions gene expression on the - observed cell-level phenotype. - """ - - def __init__( - self, - n_input: int, - n_latent: int, - P: int, - n_batch: int, - global_scale: float = 10.0, - local_scale: float = 5.0, - prior_temperature: float = 1.0, - guide_temperature: float = 1.0, - gate_prob: float = 0.5, - mixture_concentration: torch.Tensor = None, - n_pseudo_obs: int = 10, - use_enumeration: bool = False, - classifier_hidden: int = 128, - classifier_dropout: float = 0.1, - classifier_n_layers: int = 3, - n_hidden: int = 128, - n_layers: int = 3, - class_weights: torch.Tensor = None, - kl_weight_max: float = 1.0, - guide_init_scale: float = 10.0, - classifier_temperature: float = 1.0, - ): - super().__init__() - self.n_input = n_input - self.n_latent = n_latent - self.P = P - self.n_hidden = n_hidden - self.n_layers = n_layers - self.global_scale = global_scale - self.local_scale = local_scale - self.prior_temperature = prior_temperature - self.guide_temperature = guide_temperature - self.mixture_concentration = mixture_concentration - self.n_pseudo_obs = n_pseudo_obs - self.gate_prob = gate_prob - # Assert that it is not None - assert ( - self.mixture_concentration is not None - ), "mixture_concentration must be provided" - self.use_enumeration = use_enumeration - self.eps = 1e-6 - self.classifier_hidden = classifier_hidden - self.classifier_dropout = classifier_dropout - self.kl_weight_max = kl_weight_max - self.classifier_n_layers = classifier_n_layers - self.guide_init_scale = guide_init_scale - self.classifier_temperature = classifier_temperature - - # Defaults so model()/guide() work before train() sets them - self.kl_weight = 1e-6 - self.reconstruction_loss_scale = 1e-3 - - self.encoder = Encoder( - n_input=n_input, - n_output=n_latent, - n_layers=n_layers, - n_hidden=n_hidden, - n_cat_list=[n_batch], - use_layer_norm=True, - ) - - # VampPrior - pseudo_inputs = torch.randn(self.n_pseudo_obs, self.n_input) - self.vamp_prior = VampPrior(pseudo_inputs, self.encoder) - - self.decoder_input_dim = self.n_latent - self.decoder = DecoderSCVI( - self.decoder_input_dim, - n_input, - n_layers=n_layers, - n_hidden=n_hidden, - n_cat_list=[n_batch], - scale_activation="softplus", - use_layer_norm=True, - ) - - self.px_r = torch.nn.Parameter(torch.ones(n_input)) - - self.classifier = PhenotypeClassifier( - n_latent=self.n_latent, - classifier_hidden=self.classifier_hidden, - P=self.P, - num_layers=self.classifier_n_layers, - temperature=self.classifier_temperature, - ) - - self.register_buffer("clone_phen_prior", torch.empty(0)) - self.register_buffer("ct_to_c", torch.empty(0, dtype=torch.long)) - self.register_buffer("c_array", torch.empty(0, dtype=torch.long)) - self.register_buffer("ct_array", torch.empty(0, dtype=torch.long)) - self.register_buffer("ct_to_cov", torch.empty(0, dtype=torch.long)) - self.c_count = 0 - self.ct_count = 0 - self.n_cells = 0 - - self.register_buffer("_target_phenotypes", torch.empty(0, dtype=torch.long)) - - # Store or compute log of class weights if provided - if class_weights is not None: - # Expect a tensor of shape (P,) - self.register_buffer("log_class_weights", torch.log(class_weights)) - else: - self.log_class_weights = None - - def prepare_two_level_params( - self, - c_count: int, - ct_count: int, - clone_phen_prior_mat: torch.Tensor, - ct_to_c_array: torch.Tensor, - c_array_for_cells: torch.Tensor, - ct_array_for_cells: torch.Tensor, - target_phenotypes: torch.Tensor, - ct_to_cov_array: torch.Tensor = None, - ): - self.c_count = c_count - self.ct_count = ct_count - self.n_cells = c_array_for_cells.shape[0] - - prior_mat = clone_phen_prior_mat + self.eps - prior_mat = prior_mat / prior_mat.sum(dim=1, keepdim=True) - - if self.prior_temperature != 1.0: - prior_mat = prior_mat ** (1.0 / self.prior_temperature) - prior_mat = prior_mat / prior_mat.sum(dim=1, keepdim=True) - - self.register_buffer("clone_phen_prior", prior_mat) - self.register_buffer("ct_to_c", ct_to_c_array) - self.register_buffer("c_array", c_array_for_cells) - self.register_buffer("ct_array", ct_array_for_cells) - self.register_buffer("_target_phenotypes", target_phenotypes) - - if ct_to_cov_array is not None: - self.register_buffer("ct_to_cov", ct_to_cov_array) - - @property - def use_gate(self) -> bool: - return self.gate_prob is not None - - @staticmethod - def _get_fn_args_from_batch(tensor_dict: Dict[str, torch.Tensor]): - x = tensor_dict[REGISTRY_KEYS.X_KEY] - batch_idx = tensor_dict[REGISTRY_KEYS.BATCH_KEY].long() - log_library = torch.log(torch.sum(x, dim=1, keepdim=True) + 1e-6) - return (x, batch_idx, log_library), {} - - @auto_move_data - def model( - self, x: torch.Tensor, batch_idx: torch.Tensor, log_library: torch.Tensor - ): - pyro.module("scvi", self) - - kl_weight = self.kl_weight - batch_size = x.shape[0] - - with pyro.plate("clonotypes", self.c_count): - B = self.mixture_concentration.shape[0] - mixture_weights = torch.ones(B, device=x.device) / B - # Expand mixture parameters to add a leading dimension for clonotypes. - # expanded_conc will have shape (self.c_count, B, K) - expanded_conc = self.mixture_concentration.unsqueeze(0).expand( - self.c_count, -1, -1 - ) - # expanded_weights will have shape (self.c_count, B) - expanded_weights = mixture_weights.unsqueeze(0).expand(self.c_count, -1) - mixture_dist = MixtureDirichlet(expanded_weights, expanded_conc) - p_c = pyro.sample("p_c", mixture_dist) - # print("p_c shape:", p_c.shape) - - with pyro.plate("ct_plate", self.ct_count): - base_p = p_c[self.ct_to_c] + self.eps - conc_ct = torch.clamp(self.local_scale * base_p, min=1e-3) - p_ct = pyro.sample("p_ct", dist.Dirichlet(conc_ct)) - - # Encoder - z_loc, z_scale, _ = self.encoder(x, batch_idx) - z_scale = torch.clamp(z_scale, min=1e-3, max=10.0) - - with pyro.plate("data", batch_size) as idx: - - with poutine.scale(scale=kl_weight): - # vamp_mixture = self.vamp_prior.get_mixture().to_event(1) - vamp_mixture = self.vamp_prior.get_mixture() - z = pyro.sample("latent", vamp_mixture) - - ct_idx = self.ct_array[idx] - prior_log = torch.log(p_ct[ct_idx] + 1e-8) # log of local p_ct - cls_logits = self.classifier(z)# + self.phenotype_decoder(z) - - px_scale, px_r_out, px_rate, px_dropout = self.decoder( - "gene", z, log_library, batch_idx - ) - - zi_gate_probs = torch.sigmoid(px_dropout).clamp(min=1e-3, max=1.0 - 1e-3) - nb_logits = (px_rate + self.eps).log() - (self.px_r.exp() + self.eps).log() - nb_logits = torch.clamp(nb_logits, min=-10.0, max=10.0) - total_count = self.px_r.exp().clamp(max=1e4) - - x_dist = dist.ZeroInflatedNegativeBinomial( - gate=zi_gate_probs, - total_count=total_count, - logits=nb_logits, - validate_args=False, - ) - scale_val = torch.tensor(self.reconstruction_loss_scale, device=x.device) - with poutine.scale(scale=scale_val): - pyro.sample("obs", x_dist.to_event(1), obs=x) - - @auto_move_data - def guide( - self, x: torch.Tensor, batch_idx: torch.Tensor, log_library: torch.Tensor - ): - pyro.module("scvi", self) - batch_size = x.shape[0] - - with pyro.plate("clonotypes", self.c_count): - # Start from a scaled version of the prior. - init_mat_c = self.clone_phen_prior * self.guide_init_scale + 1e-3 - init_mat_c = init_mat_c.to(x.device) - - # Learnable raw parameters for q(p_c) - if "q_p_c_raw" not in pyro.get_param_store(): - q_p_c_raw = pyro.param( - "q_p_c_raw", - init_mat_c.clone().detach(), - constraint=dist.constraints.positive - ) - else: - q_p_c_raw = pyro.param("q_p_c_raw") - - bad_c = ~torch.isfinite(q_p_c_raw) - if bad_c.any(): - q_p_c_raw = torch.where(bad_c, init_mat_c.to(q_p_c_raw.device), q_p_c_raw) - - # Apply a sharpening transformation controlled by guide_temperature. - q_p_c_sharp = q_p_c_raw ** (1.0 / self.guide_temperature) - q_p_c_sharp = torch.clamp(q_p_c_sharp, min=1e-8) # ← add this - q_p_c_sharp = q_p_c_sharp / q_p_c_sharp.sum(dim=1, keepdim=True) - conc_c_guide = torch.clamp(self.global_scale * q_p_c_sharp, min=1e-3) - - # Sample p_c from a single learned Dirichlet per clonotype. - pyro.sample("p_c", dist.Dirichlet(conc_c_guide)) - - with pyro.plate("ct_plate", self.ct_count): - init_mat = self.clone_phen_prior[self.ct_to_c, :] - init_mat = init_mat * self.guide_init_scale + 1e-3 - init_mat = init_mat.to(x.device) - if "q_p_ct_raw" not in pyro.get_param_store(): - q_p_ct_raw = pyro.param( - "q_p_ct_raw", - init_mat.clone().detach(), # Make sure it's not a leaf - constraint=dist.constraints.positive, - ) - else: - q_p_ct_raw = pyro.param("q_p_ct_raw") - - bad_ct = ~torch.isfinite(q_p_ct_raw) - if bad_ct.any(): - q_p_ct_raw = torch.where(bad_ct, init_mat.to(q_p_ct_raw.device), q_p_ct_raw) - - q_p_ct_sharp = q_p_ct_raw ** (1.0 / self.guide_temperature) - q_p_ct_sharp = torch.clamp(q_p_ct_sharp, min=1e-8) - q_p_ct_sharp = q_p_ct_sharp / q_p_ct_sharp.sum(dim=1, keepdim=True) - conc_ct_guide = torch.clamp(self.local_scale * q_p_ct_sharp, min=1e-3) - pyro.sample("p_ct", dist.Dirichlet(conc_ct_guide)) - - z_loc, z_scale, _ = self.encoder(x, batch_idx) - z_scale = torch.clamp(z_scale, min=1e-3, max=10.0) - - with pyro.plate("data", batch_size) as idx: - latent_posterior = dist.Normal(z_loc, z_scale) - with pyro.poutine.scale(scale=self.kl_weight): - z = pyro.sample("latent", latent_posterior.to_event(1)) - - @auto_move_data - def get_latent(self, tensor_dict: Dict[str, torch.Tensor]): - x = tensor_dict[REGISTRY_KEYS.X_KEY] - batch_idx = tensor_dict[REGISTRY_KEYS.BATCH_KEY].long() - z_loc, _, _ = self.encoder(x, batch_idx) - if z_loc.ndim == 3: - z_loc = z_loc.mean(dim=1) - return z_loc.cpu() - - @torch.no_grad() - def get_p_ct(self): - from pyro import get_param_store - - param_store = get_param_store() - q_p_ct_raw = param_store["q_p_ct_raw"] - bad = ~torch.isfinite(q_p_ct_raw) - if bad.any(): - n_phen = q_p_ct_raw.shape[1] - q_p_ct_raw = torch.where(bad, torch.ones_like(q_p_ct_raw) / n_phen, q_p_ct_raw) - if self.guide_temperature != 1.0: - q_p_ct_sharp = q_p_ct_raw ** (1.0 / self.guide_temperature) - q_p_ct_sharp = q_p_ct_sharp / q_p_ct_sharp.sum(dim=1, keepdim=True) - else: - q_p_ct_sharp = q_p_ct_raw / q_p_ct_raw.sum(dim=1, keepdim=True) - return q_p_ct_sharp - - -############################################################################### -# 2) Unified Training Plan with Validation Step for scvi Early Stopping -############################################################################### -class UnifiedTrainingPlan(PyroTrainingPlan): - """ - Training plan that includes classification, reconstruction losses, - KL warmup, plus a validation_step that logs 'elbo_validation' so scvi's - early stopping can monitor it. - """ - - def __init__( - self, - module: TCRIModule, - n_steps_kl_warmup: int = 1000, - reconstruction_loss_scale: float = 1e-2, - num_particles: int = 5, - optimizer_config: dict = None, - class_weights: torch.Tensor = None, - **kwargs, - ): - self.num_particles = num_particles - if module.use_enumeration: - print("Using Enumeration") - self._loss_fn = TraceEnum_ELBO( - max_plate_nesting=3, num_particles=self.num_particles - ) - else: - self._loss_fn = Trace_ELBO() - - super().__init__(module, n_steps_kl_warmup=n_steps_kl_warmup, **kwargs) - - self.n_steps_kl_warmup = n_steps_kl_warmup - self.reconstruction_loss_scale = reconstruction_loss_scale - self._my_global_step = 0 - self.class_weights = class_weights - self.optimizer_config = optimizer_config - - if optimizer_config is None: - optimizer_config = {"lr":1e-3,"betas":(0.9,0.999),"eps":1e-5,"weight_decay":1e-4} - self.optimizer_config = optimizer_config - - @property - def loss(self): - return self._loss_fn - - def configure_optimizers(self): - optimizer = torch.optim.Adam( - self.module.parameters(), - lr=self.optimizer_config["lr"], - betas=self.optimizer_config["betas"], - eps=self.optimizer_config["eps"], - weight_decay=self.optimizer_config["weight_decay"], - ) - return {"optimizer": optimizer} - - def training_step(self, batch, batch_idx): - # ── KL warmup ──────────────────────────────────────────── - if self.n_steps_kl_warmup > 0 and self._my_global_step < self.n_steps_kl_warmup: - kl_weight = max(1e-6, self.module.kl_weight_max * (self._my_global_step / self.n_steps_kl_warmup)) - else: - kl_weight = self.module.kl_weight_max - self.module.kl_weight = kl_weight - - # ── Pyro ELBO step ─────────────────────────────────────── - loss_dict = super().training_step(batch, batch_idx) - device = next(self.module.parameters()).device - - if not isinstance(loss_dict["loss"], torch.Tensor): - loss_dict["loss"] = torch.tensor(loss_dict["loss"], device=device, requires_grad=True) - else: - loss_dict["loss"] = loss_dict["loss"].to(device) - - # ── Diagnostics (no gradient contribution) ─────────────── - with torch.no_grad(): - z_diag = self.module.get_latent(batch).to(device) - idx_diag = batch["indices"].long().view(-1).to(device) - cls_logits = self.module.classifier(z_diag) - ct_idx = self.module.ct_array[idx_diag] - p_ct_prior = self.module.get_p_ct()[ct_idx].to(device) - prior_log = torch.log(p_ct_prior + 1e-8) - - 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 - - probs = F.softmax(local_logits, dim=-1) - kl_div = F.kl_div(probs.log(), p_ct_prior, reduction='batchmean') - entropy = -torch.sum(probs * torch.log(probs + 1e-8), dim=-1).mean() - confidence = (probs**2).sum(dim=-1).mean() - - self.log("kl_divergence_with_prior_train", kl_div, prog_bar=False, on_epoch=True) - self.log("entropy_train", entropy, prog_bar=False, on_epoch=True) - self.log("confidence_train", confidence, prog_bar=False, on_epoch=True) - - self._my_global_step += 1 - return loss_dict - - def validation_step(self, batch, batch_idx): - with torch.no_grad(): - self.module.eval() - val_dict = super().training_step(batch, batch_idx) - self.module.train() - - device = next(self.module.parameters()).device - - if not isinstance(val_dict["loss"], torch.Tensor): - val_dict["loss"] = torch.tensor(val_dict["loss"], device=device) - else: - val_dict["loss"] = val_dict["loss"].to(device) - - # ── Diagnostic only ────────────────────────────────────── - z_batch = self.module.get_latent(batch).to(device) - idx = batch["indices"].long().view(-1).to(device) - cls_logits = self.module.classifier(z_batch) - ct_idx = self.module.ct_array[idx] - p_ct_prior = self.module.get_p_ct()[ct_idx].to(device) - prior_log = torch.log(p_ct_prior + 1e-8) - - 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 - - probs = F.softmax(local_logits, dim=-1) - kl_divergence = F.kl_div(probs.log(), p_ct_prior, reduction='batchmean') - self.log("kl_divergence_with_prior_val", kl_divergence, prog_bar=False, on_epoch=True) - - self.log("elbo_validation", val_dict["loss"], prog_bar=True, on_epoch=True) - return val_dict - - -############################################################################### -# 3) High-Level scVI Model with scvi Early Stopping -############################################################################### class TCRIModel(BaseModelClass): @classmethod 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!") @@ -694,12 +91,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, @@ -718,14 +113,33 @@ def __init__( classifier_dropout: float = 0.1, n_pseudo_obs: int = 10, K: int = 10, - phenotype_weights: Optional[Dict[str, float]] = None, - gate_prob: Optional[float] = None, + gate_prob: Optional[float] = 0.5, # π (gating weight, methods §Generative Model); None = additive kl_weight_max: float = 1.0, guide_init_scale: float = 10.0, classifier_temperature: float = 1.0, + phenotype_kl_weight: float = 1.0, **kwargs, ): super().__init__(adata) + + # Pyro's param store is PROCESS-GLOBAL: a second TCRIModel in the same session + # silently inherits the first model's fitted q_p_c_raw/q_p_ct_raw and network + # weights, so it starts from the previous fit instead of from scratch. We warn + # rather than clearing, because clearing here would destroy the params of a + # model loaded earlier in the session (load_tcri_session restores the store + # after construction). Proper per-instance namespacing is a design change. + _tcri_params = [k for k in pyro.get_param_store().keys() if k.startswith(("q_p_c", "q_p_ct", "scvi$$$"))] + if _tcri_params: + warnings.warn( + "The global Pyro param store already holds TCRI parameters " + f"({len(_tcri_params)} entries). This model will CONTINUE that fit " + "rather than start fresh. Call `pyro.clear_param_store()` before " + "constructing a new model (note this invalidates any model already " + "loaded in this session).", + UserWarning, + stacklevel=2, + ) + n_vars = self.summary_stats["n_vars"] clonotype_col = self.adata_manager.registry["clonotype_col"] phenotype_col = self.adata_manager.registry["phenotype_col"] @@ -742,13 +156,25 @@ def __init__( c_count = len(cvals.cat.categories) c_array_np = cvals.cat.codes.values pvals_np = ph_series.cat.codes.values - c2p_mat = np.zeros((c_count, P), dtype=np.float32) + clone_phenotype_prior = np.zeros((c_count, P), dtype=np.float32) for i in range(len(c_array_np)): - c2p_mat[c_array_np[i], pvals_np[i]] += 1 - c2p_mat += 1e-6 - c2p_mat = c2p_mat / c2p_mat.sum(axis=1, keepdims=True) - self.c2p_mat = c2p_mat - self.centers, self.labels = build_archetypes(self.c2p_mat, K=K) + clone_phenotype_prior[c_array_np[i], pvals_np[i]] += 1 + clone_phenotype_prior += 1e-6 + clone_phenotype_prior = clone_phenotype_prior / clone_phenotype_prior.sum(axis=1, keepdims=True) + self.clone_phenotype_prior = clone_phenotype_prior + # K archetypes are k-means centroids over clonotypes, so K > n_clonotypes is + # not satisfiable (sklearn raises "n_samples < n_clusters"). Clamp instead of + # crashing — a dataset with few clones is legitimate (and is exactly what the + # synthetic examples use). + if K > c_count: + warnings.warn( + f"K={K} archetypes requested but the data has only {c_count} " + f"clonotype(s); using K={c_count}.", + UserWarning, + stacklevel=2, + ) + K = c_count + self.centers, self.labels = build_archetypes(self.clone_phenotype_prior, K=K) cov_series = self.adata.obs[covariate_col].astype("category") cov_array_np = cov_series.cat.codes.values df_ct = pd.DataFrame({"c": c_array_np, "t": cov_array_np}) @@ -769,24 +195,6 @@ def __init__( batch_series = self.adata.obs[batch_col].astype("category") n_batch = len(batch_series.cat.categories) - if phenotype_weights is None: - # Automatically compute inverse-frequency weights for each phenotype - freq_count = ph_series.value_counts(sort=False) - class_weights_arr = [] - for cat_name in ph_series.cat.categories: - c = freq_count[cat_name] - # inverse-frequency weight - weight = 1.0 / c - class_weights_arr.append(weight) - class_weights = torch.tensor(class_weights_arr, dtype=torch.float32) - else: - class_weights_arr = [] - for cat_name in ph_series.cat.categories: - weight = phenotype_weights.get(cat_name, 1.0) - class_weights_arr.append(weight) - class_weights = torch.tensor(class_weights_arr, dtype=torch.float32) - - self.class_weights = class_weights self.module = TCRIModule( n_input=n_vars, n_latent=n_latent, @@ -803,15 +211,15 @@ def __init__( use_enumeration=use_enumeration, classifier_hidden=classifier_hidden, classifier_dropout=classifier_dropout, - class_weights=self.class_weights, gate_prob=gate_prob, kl_weight_max=kl_weight_max, n_pseudo_obs=n_pseudo_obs, guide_init_scale=guide_init_scale, classifier_temperature=classifier_temperature, + phenotype_kl_weight=phenotype_kl_weight, ) self.init_params_ = self._get_init_params(locals()) - c2p_torch = torch.tensor(c2p_mat, dtype=torch.float32) + c2p_torch = torch.tensor(clone_phenotype_prior, dtype=torch.float32) c_array_torch = torch.tensor(c_array_np, dtype=torch.long) ct_array_torch = torch.tensor(ct_array_np, dtype=torch.long) ct_to_c_torch = torch.tensor(ct_to_c_list, dtype=torch.long) @@ -838,7 +246,7 @@ def train( max_epochs: int = 1000, batch_size: int = 1000, lr: float = 1e-3, - reconstruction_loss_scale: float = 1e-3, + reconstruction_loss_scale: float = 1e-2, n_steps_kl_warmup: int = 2000, **kwargs, ): @@ -849,6 +257,21 @@ def train( """ # Create a train/val split self.module.reconstruction_loss_scale = reconstruction_loss_scale + + # batch_size >= n_obs means ONE optimizer step per epoch, so the fixed + # per-epoch overhead is paid per gradient update — the pathology behind the + # "9-hour" synthetic run (1000 cells, batch_size=20000, max_epochs=1e6). + n_obs = self.adata.n_obs + if batch_size >= n_obs: + warnings.warn( + f"batch_size={batch_size} >= n_obs={n_obs}: each epoch is a SINGLE " + "optimizer step, so per-epoch overhead dominates and `max_epochs` " + "becomes the number of gradient updates. Use a smaller batch_size " + "(e.g. 256-1024) for a comparable number of updates in far less time.", + UserWarning, + stacklevel=2, + ) + splitter = DataSplitter( self.adata_manager, train_size=0.9, @@ -860,7 +283,6 @@ def train( module=self.module, n_steps_kl_warmup=n_steps_kl_warmup, reconstruction_loss_scale=reconstruction_loss_scale, - class_weights=self.class_weights, optimizer_config={ "lr": lr, "betas": (0.9, 0.999), @@ -869,18 +291,22 @@ def train( }, ) + # Defaults the caller can override: setdefault (not hard-coded keywords) so + # passing e.g. early_stopping_patience=10 or accelerator="gpu" through + # train(**kwargs) works instead of raising "got multiple values for keyword". + kwargs.setdefault("early_stopping", True) + kwargs.setdefault("early_stopping_monitor", "elbo_validation") + kwargs.setdefault("early_stopping_mode", "min") + kwargs.setdefault("early_stopping_patience", self.patience) + kwargs.setdefault("check_val_every_n_epoch", 5) + kwargs.setdefault("accelerator", "auto") + kwargs.setdefault("devices", "auto") + runner = TrainRunner( self, training_plan=plan, data_splitter=splitter, max_epochs=max_epochs, - early_stopping=True, - early_stopping_monitor="elbo_validation", - early_stopping_mode="min", - early_stopping_patience=self.patience, - check_val_every_n_epoch=5, - accelerator="auto", - devices="auto", **kwargs, ) @@ -888,7 +314,7 @@ def train( return @torch.no_grad() - def get_latent_representation(self, adata=None, indices=None, batch_size=None): + def get_latent_representation(self, adata=None, indices=None, batch_size=4096): adata = self._validate_anndata(adata) scdl = self._make_data_loader( adata=adata, indices=indices, batch_size=batch_size @@ -905,71 +331,132 @@ 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 = 4096, 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 = 4096, 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) + ) + + 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, @@ -986,11 +473,11 @@ def _ok(m): print(f"{GRN}✅ {m}{RST}") raise ValueError(f"phenotype '{phenotype_name}' not found. Choices: {list(cats)}") p_idx = list(cats).index(phenotype_name) - # 2) clone-level prior (numpy array stored in model.c2p_mat) - mat = self.c2p_mat.copy() + # 2) clone-level prior (numpy array stored in model.clone_phenotype_prior) + mat = self.clone_phenotype_prior.copy() mat[:, p_idx] *= boost_factor mat /= mat.sum(axis=1, keepdims=True) - self.c2p_mat = mat # keep external copy + self.clone_phenotype_prior = mat # keep external copy with torch.no_grad(): new_clone_prior = torch.tensor(mat, dtype=torch.float32, @@ -1014,61 +501,4 @@ def _ok(m): print(f"{GRN}✅ {m}{RST}") _ok("Read to train.") - def plot_archetypes(self): - order = np.argsort(self.labels) - ordered_mat = self.c2p_mat[order, :] - - # Plot heatmap of the clone phenotype distributions - plt.figure(figsize=(10, 6)) - plt.imshow(ordered_mat, aspect='auto', cmap='viridis') - plt.colorbar(label='Phenotype Distribution') - plt.title('Heatmap of Clone Phenotype Distributions (Ordered by Cluster)') - plt.xlabel('Phenotype') - plt.ylabel('Clone (ordered by cluster)') - plt.show() - - # Plot heatmap of the archetype centroids - plt.figure(figsize=(6, 4)) - plt.imshow(self.centers, aspect='auto', cmap='viridis') - plt.colorbar(label='Centroid Value') - plt.title('Heatmap of Archetype Centroids') - plt.xlabel('Phenotype') - plt.ylabel('Archetype') - plt.show() - - def plot_loss(self, log_scale=False): - # Retrieve loss and accuracy history - loss_history = self.history_.get("elbo_train", []) - loss_validation = self.history_.get("elbo_validation", []) - train_accuracy = self.history_.get("kl_divergence_with_prior_train_epoch", []) - val_accuracy = self.history_.get("kl_divergence_with_prior_val", []) - - # Create subplots - fig, axes = plt.subplots(2, 1, figsize=(10, 12)) - - # Plot ELBO loss - axes[0].plot(loss_history, label="Training ELBO Loss") - axes[0].plot(loss_validation, label="Validation ELBO Loss") - axes[0].set_xlabel("Epoch") - axes[0].set_ylabel("ELBO Loss") - axes[0].set_title("ELBO Loss Over Epochs") - axes[0].legend() - - # Plot Accuracy - if len(train_accuracy) > 0 or len(val_accuracy) > 0: - if len(train_accuracy) > 0: - axes[1].plot(train_accuracy, label="Training Accuracy") - if len(val_accuracy) > 0: - axes[1].plot(val_accuracy, label="Validation Accuracy") - axes[1].set_xlabel("Epoch") - axes[1].set_ylabel("dKL") - axes[1].set_title("DKL Over Epochs") - axes[1].legend() - - # Apply log scale if requested - if log_scale: - for ax in axes: - ax.set_yscale("log") - plt.tight_layout() - plt.show() diff --git a/tcri/model/_model_contract.py b/tcri/model/_model_contract.py new file mode 100644 index 0000000..0d48268 --- /dev/null +++ b/tcri/model/_model_contract.py @@ -0,0 +1,209 @@ +"""TCRI **model contract** (frozen) — the machine-checkable probabilistic structure. + +The API contract (``tcri/_contract.pyi``) freezes the public *interface*; this file +freezes the *model* — the generative program and variational family of +**Supplementary Note 1** (``docs/contract/MODEL_CONTRACT.md``, prose; the PDF is the +source of truth). The conformance test +(``tests/test_model_contract_conformance.py``) traces ``TCRIModule.model``/``.guide`` +and asserts the live program matches this manifest **exactly** — no missing sites and +**no extra sites**. + +RULES this file encodes: +- Every stochastic site in ``model()``/``guide()`` is declared here, with the + equation of the note it realizes. A site NOT declared here must NOT exist. +- Changing the model's mathematics **requires updating this contract first** + (and ``docs/contract/MODEL_CONTRACT.md``). The conformance test is the forcing + function: it fails until the contract and the code agree. Do not "fix" a failure + by loosening the manifest to match new code — update the contract deliberately, + with the note reference, so the change is reviewed as a model change. +- Sanctioned departures from the note live in ``SANCTIONED_DEVIATIONS`` with a + rationale. An unlisted departure is a defect, not a feature. + +Symbols (note -> code): ω_c->``p_c``, ϕ_m->``p_ct``, z_i->``latent``, x_i->``obs``, +f_cls->``classifier``, π->``gate_prob``, α->``global_scale``, β->``local_scale``, +γ->``phenotype_kl_weight``, g(i)->``ct_array``, h(m)->``ct_to_c``. +""" +from __future__ import annotations + +from typing import Optional + +__all__ = [ + "GENERATIVE_SITES", + "GUIDE_SITES", + "GUIDE_PARAMS", + "FORBIDDEN_GUIDE_SITES", + "PLATES", + "SEMANTIC_INVARIANTS", + "SANCTIONED_DEVIATIONS", + "SiteSpec", +] + + +class SiteSpec: + """One declared stochastic site: what distribution, in which plate, per the note. + + ``dist`` is the *unwrapped* distribution class name — the checker strips Pyro's + ``Independent``/``ExpandedDistribution`` wrappers (``.to_event()``/``.expand()`` + are implementation details, not model semantics). ``event_dim`` pins how many + trailing dims are dependent, which IS semantics. + """ + + def __init__( + self, + name: str, + dist: str, + plate: Optional[str], + eq: str, + observed: bool = False, + event_dim: Optional[int] = None, + note: str = "", + ): + self.name = name + self.dist = dist + self.plate = plate + self.eq = eq + self.observed = observed + self.event_dim = event_dim + self.note = note + + def __repr__(self) -> str: # pragma: no cover - debug aid + return f"SiteSpec({self.name!r}, dist={self.dist!r}, plate={self.plate!r}, eq={self.eq!r})" + + +# ── plates ─────────────────────────────────────────────────────────────────── +# name -> what one index of the plate is (the note's index set). +PLATES = { + "clonotypes": "c = 1..C (clonotype)", + "ct_plate": "m = 1..M (clonotype x covariate group)", + "data": "i = 1..N (cell; subsampled minibatch)", +} + + +# ── the generative program: p(Ω, Φ, z, x) ──────────────────────────────────── +GENERATIVE_SITES = [ + SiteSpec( + "p_c", "MixtureDirichlet", "clonotypes", eq="1", event_dim=1, + note="ω_c ~ (1/B_c) Σ_b Dir(α·ψ_b); ψ_b = archetype vectors, α = global_scale.", + ), + SiteSpec( + "p_ct", "Dirichlet", "ct_plate", eq="2", event_dim=1, + note="ϕ_m | ω_h(m) ~ Dir(β·ω_h(m)); β = local_scale, h(m) = ct_to_c.", + ), + SiteSpec( + "latent", "MixtureSameFamily", "data", eq="3", event_dim=1, + note="z_i ~ (1/B_z) Σ_k q(z|u_k) — VampPrior over learnable pseudo-inputs.", + ), + SiteSpec( + "phenotype_alignment", "Unit", "data", eq="Inference Details", observed=True, + note=( + "The surrogate replacing the discrete z^ϕ (eq 4): a pyro.factor carrying " + "−γ·KL(probs_i ‖ ϕ_g(i)), probs_i = softmax(ℓ_i), " + "ℓ_i = π·f_cls(z_i) + (1−π)·log ϕ_g(i). This term is what trains f_cls." + ), + ), + SiteSpec( + "obs", "ZeroInflatedNegativeBinomial", "data", eq="5", observed=True, event_dim=1, + note="x_i ~ ZINB(g'_i, r_i, μ_i) from the scVI decoder (+ library size).", + ), +] + + +# ── the variational family q(Ω, Φ, z | x) (eq 6) ──────────────────────────── +GUIDE_SITES = [ + SiteSpec("p_c", "Dirichlet", "clonotypes", eq="6", event_dim=1, + note="q(ω_c) = Dir(λ_c); concentration = α · normalized q_p_c_raw."), + SiteSpec("p_ct", "Dirichlet", "ct_plate", eq="6", event_dim=1, + note="q(ϕ_m) = Dir(λ'_m); concentration = β · normalized q_p_ct_raw."), + SiteSpec("latent", "Normal", "data", eq="6", event_dim=1, + note="q(z_i|x_i) = N(μ_i, diag(σ_i²)) from the encoder."), +] + +# Learnable variational parameters the guide must register (λ_c, λ'_m). +GUIDE_PARAMS = ["q_p_c_raw", "q_p_ct_raw"] + +# The discrete phenotype latent z^ϕ is NOT sampled: the note's "Inference Details" +# replaces it with the surrogate KL penalty. A q(z^ϕ) categorical site reappearing +# in the guide means someone re-introduced the enumerated path without updating the +# contract — which changes the objective. +FORBIDDEN_GUIDE_SITES = ["z_phi", "z_phenotype", "phenotype", "phenotype_alignment"] + + +# ── semantic invariants (behavior the structure alone can't pin) ───────────── +# Each is asserted by the conformance test; the string is the failure explanation. +SEMANTIC_INVARIANTS = { + "alpha_scales_clonotype_prior": ( + "eq 1: the p_c prior concentration must scale with α (global_scale). " + "Without it the prior is Dir(ψ_b) (sums to ~1, U-shaped: mass at the simplex " + "corners) and is scaled inconsistently with the guide, which does apply α." + ), + "beta_scales_covariate_prior": ( + "eq 2: the p_ct prior concentration must scale with β (local_scale), i.e. " + "conc = β·ω_h(m), so the total concentration is ~β." + ), + "hierarchy_ct_depends_on_c": ( + "eq 2 is HIERARCHICAL: p_ct's concentration must be built from the SAMPLED " + "ω_c (site 'p_c') indexed by ct_to_c = h(m) — not from the static empirical " + "`clone_phen_prior`, and not under any other index map. Severing this makes " + "ϕ_m independent of its clonotype (p_c becomes a dangling latent contributing " + "only its own KL) and no clonotype-level information reaches the covariate " + "level, which is the entire point of the two-level model." + ), + "factor_is_negative_kl": ( + "Inference Details: the phenotype_alignment factor must be −γ·KL(probs‖ϕ) " + "(≤ 0). SVI MAXIMIZES the log-joint and the note's γ·KL is a PENALTY, so the " + "sign must be negative; a positive factor would push probs AWAY from ϕ." + ), + "gate_mixes_classifier_and_prior": ( + "eq 4: ℓ_i = π·f_cls(z_i) + (1−π)·log ϕ_g(i). π=1 ⇒ pure classifier; " + "π=0 ⇒ pure clonotype prior; π=None ⇒ the additive rule f_cls + log ϕ." + ), + "alignment_target_uses_global_indices": ( + "The per-cell target ϕ_g(i) must be indexed by GLOBAL cell indices, never the " + "local pyro plate index (0..batch_size−1), which scrambles targets across " + "shuffled minibatches and collapses f_cls to a constant." + ), +} + + +# ── departures from the note that are accepted, with rationale ─────────────── +# Anything NOT listed here that departs from the note is a defect. +SANCTIONED_DEVIATIONS = { + "E_reconstruction_loss_scale": ( + "eq 7 weights E[log p(x|z)] at 1; the code scales the `obs` site by " + "`reconstruction_loss_scale`, a beta-VAE-style reweighting. Default RAISED " + "1e-3 -> 1e-2 after re-measurement: on real data (yost subset, 2259 cells x " + "1000 genes) the posterior-predictive library ratio was 1.40 at 1e-3 and " + "0.99 at 1e-2, with dropout matching throughout. Recovery (1.000) and latent " + "separation (7.19) are unchanged, so the calibration is free. NOTE the " + "original ~6x over-generation was mostly the phantom second optimizer " + "shrinking the decoder (see `optimizer_weight_decay`); removing it took the " + "ratio 6x -> 1.40, and this default closes the rest. Full weight (1.0) " + "over-corrects: ratio 0.91 on synthetic." + ), + "kl_warmup_z_only": ( + "Training-only: `kl_weight` anneals the `latent` (z) KL over " + "`n_steps_kl_warmup`; the two Dirichlet KLs are unscaled. Standard annealing, " + "not part of eq 7." + ), + "num_particles_enumeration_only": ( + "`num_particles` is honored only on the TraceEnum_ELBO path; the default " + "Trace_ELBO uses a single MC particle." + ), + "F_perturbation_not_implemented": ( + "In-silico perturbation (eqs 8-12) is not implemented. Additive feature; " + "explicitly out of scope for this release." + ), + "optimizer_weight_decay": ( + "Training-only: the SVI optimizer applies Adam weight decay (default 1e-4) " + "to the network parameters. The note specifies the objective (eq 7 + the " + "surrogate) but not the optimizer, so this is an L2 regularizer on top of " + "it. Applied inside Pyro's optimizer, where it acts on the ELBO gradients. " + "NOTE (history): this used to be applied by a second torch Adam installed " + "over `configure_optimizers`, which overrode scvi's deliberate no-op shim " + "and stepped AFTER `SVI.step()` had already zeroed the gradients. With zero " + "gradients Adam's normalization turns weight decay into a scale-free " + "shrink of ~lr per step (not proportional L2), which held the networks at " + "~2.4x smaller weights, and `train(lr=)` never reached the real optimizer. " + "Removed; see MODEL_CONTRACT.md." + ), +} diff --git a/tcri/model/_module.py b/tcri/model/_module.py new file mode 100644 index 0000000..78024e0 --- /dev/null +++ b/tcri/model/_module.py @@ -0,0 +1,366 @@ +"""The TCRI Pyro module: a CVAE (encoder/decoder over gene expression) coupled to +two-level hierarchical Dirichlet priors (clonotype -> clonotype x covariate) and a +phenotype classifier head.""" +from typing import Dict, Optional + +import torch +import pyro +import pyro.distributions as dist +import pyro.poutine as poutine + +from scvi import REGISTRY_KEYS +from scvi.nn import Encoder, DecoderSCVI +from scvi.module.base import PyroBaseModuleClass, auto_move_data + +from ._classifier import PhenotypeClassifier +from ._priors import VampPrior, MixtureDirichlet + +__all__ = ["TCRIModule"] + + +class TCRIModule(PyroBaseModuleClass): + """ + Two-level model that incorporates hierarchical priors (clonotype-level) + and a CVAE structure that explicitly conditions gene expression on the + observed cell-level phenotype. + """ + + def __init__( + self, + n_input: int, + n_latent: int, + P: int, + n_batch: int, + global_scale: float = 10.0, + local_scale: float = 5.0, + prior_temperature: float = 1.0, + guide_temperature: float = 1.0, + gate_prob: Optional[float] = 0.5, # None = additive (no gating) + mixture_concentration: torch.Tensor = None, + n_pseudo_obs: int = 10, + use_enumeration: bool = False, + classifier_hidden: int = 128, + classifier_dropout: float = 0.1, + classifier_n_layers: int = 3, + n_hidden: int = 128, + n_layers: int = 3, + kl_weight_max: float = 1.0, + guide_init_scale: float = 10.0, + classifier_temperature: float = 1.0, + phenotype_kl_weight: float = 1.0, + ): + super().__init__() + self.n_input = n_input + self.n_latent = n_latent + self.P = P + self.n_hidden = n_hidden + self.n_layers = n_layers + self.global_scale = global_scale + self.local_scale = local_scale + self.prior_temperature = prior_temperature + self.guide_temperature = guide_temperature + # register_buffer (not a plain attribute) so module.to(device) moves it with + # the rest of the module — as a bare Tensor it stays on CPU and the eq-1 prior + # then needs an ad-hoc .to() at every use, or fails outright on GPU. + if mixture_concentration is not None and not torch.is_tensor(mixture_concentration): + mixture_concentration = torch.as_tensor(mixture_concentration) + self.register_buffer("mixture_concentration", mixture_concentration) + self.n_pseudo_obs = n_pseudo_obs + self.gate_prob = gate_prob + # Assert that it is not None + assert ( + self.mixture_concentration is not None + ), "mixture_concentration must be provided" + self.use_enumeration = use_enumeration + self.eps = 1e-6 + self.classifier_hidden = classifier_hidden + self.classifier_dropout = classifier_dropout + self.kl_weight_max = kl_weight_max + self.classifier_n_layers = classifier_n_layers + self.guide_init_scale = guide_init_scale + self.classifier_temperature = classifier_temperature + self.phenotype_kl_weight = phenotype_kl_weight # γ (methods §Inference Details) + + # Defaults so model()/guide() work before train() sets them + self.kl_weight = 1e-6 + self.reconstruction_loss_scale = 1e-2 + + self.encoder = Encoder( + n_input=n_input, + n_output=n_latent, + n_layers=n_layers, + n_hidden=n_hidden, + n_cat_list=[n_batch], + use_layer_norm=True, + ) + + # VampPrior + pseudo_inputs = torch.randn(self.n_pseudo_obs, self.n_input) + self.vamp_prior = VampPrior(pseudo_inputs, self.encoder) + + self.decoder_input_dim = self.n_latent + self.decoder = DecoderSCVI( + self.decoder_input_dim, + n_input, + n_layers=n_layers, + n_hidden=n_hidden, + n_cat_list=[n_batch], + scale_activation="softplus", + use_layer_norm=True, + ) + + self.px_r = torch.nn.Parameter(torch.ones(n_input)) + + self.classifier = PhenotypeClassifier( + n_latent=self.n_latent, + classifier_hidden=self.classifier_hidden, + P=self.P, + num_layers=self.classifier_n_layers, + dropout_rate=self.classifier_dropout, + temperature=self.classifier_temperature, + ) + + self.register_buffer("clone_phen_prior", torch.empty(0)) + self.register_buffer("ct_to_c", torch.empty(0, dtype=torch.long)) + self.register_buffer("c_array", torch.empty(0, dtype=torch.long)) + self.register_buffer("ct_array", torch.empty(0, dtype=torch.long)) + self.register_buffer("ct_to_cov", torch.empty(0, dtype=torch.long)) + self.c_count = 0 + self.ct_count = 0 + self.n_cells = 0 + + self.register_buffer("_target_phenotypes", torch.empty(0, dtype=torch.long)) + + def prepare_two_level_params( + self, + c_count: int, + ct_count: int, + clone_phen_prior_mat: torch.Tensor, + ct_to_c_array: torch.Tensor, + c_array_for_cells: torch.Tensor, + ct_array_for_cells: torch.Tensor, + target_phenotypes: torch.Tensor, + ct_to_cov_array: torch.Tensor = None, + ): + self.c_count = c_count + self.ct_count = ct_count + self.n_cells = c_array_for_cells.shape[0] + + prior_mat = clone_phen_prior_mat + self.eps + prior_mat = prior_mat / prior_mat.sum(dim=1, keepdim=True) + + if self.prior_temperature != 1.0: + prior_mat = prior_mat ** (1.0 / self.prior_temperature) + prior_mat = prior_mat / prior_mat.sum(dim=1, keepdim=True) + + self.register_buffer("clone_phen_prior", prior_mat) + self.register_buffer("ct_to_c", ct_to_c_array) + self.register_buffer("c_array", c_array_for_cells) + self.register_buffer("ct_array", ct_array_for_cells) + self.register_buffer("_target_phenotypes", target_phenotypes) + + if ct_to_cov_array is not None: + self.register_buffer("ct_to_cov", ct_to_cov_array) + + @property + def use_gate(self) -> bool: + return self.gate_prob is not None + + @staticmethod + def _get_fn_args_from_batch(tensor_dict: Dict[str, torch.Tensor]): + x = tensor_dict[REGISTRY_KEYS.X_KEY] + batch_idx = tensor_dict[REGISTRY_KEYS.BATCH_KEY].long() + log_library = torch.log(torch.sum(x, dim=1, keepdim=True) + 1e-6) + # Global cell indices for this minibatch. Needed so model()/guide() can map + # each cell to its (clonotype x covariate) group via ``ct_array``; the pyro + # data-plate index is LOCAL (0..batch_size-1) and must NOT be used for this. + indices = tensor_dict["indices"].long().view(-1) + return (x, batch_idx, log_library, indices), {} + + @auto_move_data + def model( + self, + x: torch.Tensor, + batch_idx: torch.Tensor, + log_library: torch.Tensor, + indices: torch.Tensor = None, + ): + pyro.module("scvi", self) + + kl_weight = self.kl_weight + batch_size = x.shape[0] + + with pyro.plate("clonotypes", self.c_count): + B = self.mixture_concentration.shape[0] + mixture_weights = torch.ones(B, device=x.device) / B + # Expand mixture parameters to add a leading dimension for clonotypes. + # expanded_conc will have shape (self.c_count, B, K). + # eq 1: ω_c ~ (1/B) Σ_b Dir(α·ψ_b) — scale the archetype vectors ψ_b by + # α (global_scale), mirroring eq 2's β on the covariate prior. Without α + # the concentration sums to ~1 (U-shaped, mass at the simplex corners) and + # is scaled inconsistently with the guide q(ω_c), which does apply α. + expanded_conc = self.global_scale * self.mixture_concentration.unsqueeze( + 0 + ).expand(self.c_count, -1, -1) + # expanded_weights will have shape (self.c_count, B) + expanded_weights = mixture_weights.unsqueeze(0).expand(self.c_count, -1) + mixture_dist = MixtureDirichlet(expanded_weights, expanded_conc) + p_c = pyro.sample("p_c", mixture_dist) + # print("p_c shape:", p_c.shape) + + with pyro.plate("ct_plate", self.ct_count): + base_p = p_c[self.ct_to_c] + self.eps + conc_ct = torch.clamp(self.local_scale * base_p, min=1e-3) + p_ct = pyro.sample("p_ct", dist.Dirichlet(conc_ct)) + + with pyro.plate("data", batch_size) as idx: + + with poutine.scale(scale=kl_weight): + # vamp_mixture = self.vamp_prior.get_mixture().to_event(1) + vamp_mixture = self.vamp_prior.get_mixture() + z = pyro.sample("latent", vamp_mixture) + + # Map each cell to its (clonotype x covariate) group using GLOBAL indices, + # not the local plate index `idx` (which is 0..batch_size-1 and would + # scramble the alignment target across shuffled minibatches). Fail loudly + # rather than silently falling back to the (wrong) local index. + assert indices is not None, ( + "model() requires global cell indices (supplied by " + "_get_fn_args_from_batch); the local plate index must not be used " + "for the clonotype x covariate lookup." + ) + ct_idx = self.ct_array[indices] + cls_logits = self.classifier(z) # l_i = f_cls(z_i) (eq. 4) + + # Phenotype-alignment surrogate (Supplementary Note, "Inference Details"): + # ℓ_i = π·f_cls(z_i) + (1-π)·log φ_{g(i)}, probs_i = softmax(ℓ_i); + # add -γ·KL(probs_i ‖ φ_{g(i)}) to the log-joint so the ELBO trains the + # classifier f_cls (η_cls). φ (the covariate-level distribution) is the + # DETACHED alignment target. Without this factor cls_logits never enters + # the ELBO and f_cls receives no gradient. + phi = p_ct[ct_idx].detach() + log_phi = torch.log(phi + 1e-8) + if self.gate_prob is not None: + ell = self.gate_prob * cls_logits + (1.0 - self.gate_prob) * log_phi + else: + ell = cls_logits + log_phi + probs = torch.softmax(ell, dim=-1) + pheno_kl = (probs * (torch.log(probs + 1e-8) - log_phi)).sum(dim=-1) + pyro.factor("phenotype_alignment", -self.phenotype_kl_weight * pheno_kl) + + px_scale, px_r_out, px_rate, px_dropout = self.decoder( + "gene", z, log_library, batch_idx + ) + + zi_gate_probs = torch.sigmoid(px_dropout).clamp(min=1e-3, max=1.0 - 1e-3) + nb_logits = (px_rate + self.eps).log() - (self.px_r.exp() + self.eps).log() + nb_logits = torch.clamp(nb_logits, min=-10.0, max=10.0) + total_count = self.px_r.exp().clamp(max=1e4) + + x_dist = dist.ZeroInflatedNegativeBinomial( + gate=zi_gate_probs, + total_count=total_count, + logits=nb_logits, + validate_args=False, + ) + # plain float, not torch.tensor(..., device=x.device): poutine.scale takes a + # scalar, and materializing one on the device is a host->device copy (and a + # sync point) on every SVI step for a value that never changes mid-epoch. + with poutine.scale(scale=float(self.reconstruction_loss_scale)): + pyro.sample("obs", x_dist.to_event(1), obs=x) + + @auto_move_data + def guide( + self, + x: torch.Tensor, + batch_idx: torch.Tensor, + log_library: torch.Tensor, + indices: torch.Tensor = None, + ): + pyro.module("scvi", self) + batch_size = x.shape[0] + + with pyro.plate("clonotypes", self.c_count): + # Start from a scaled version of the prior. + init_mat_c = self.clone_phen_prior * self.guide_init_scale + 1e-3 + init_mat_c = init_mat_c.to(x.device) + + # Learnable raw parameters for q(p_c) + if "q_p_c_raw" not in pyro.get_param_store(): + q_p_c_raw = pyro.param( + "q_p_c_raw", + init_mat_c.clone().detach(), + constraint=dist.constraints.positive + ) + else: + q_p_c_raw = pyro.param("q_p_c_raw") + + bad_c = ~torch.isfinite(q_p_c_raw) + if bad_c.any(): + q_p_c_raw = torch.where(bad_c, init_mat_c.to(q_p_c_raw.device), q_p_c_raw) + + # Apply a sharpening transformation controlled by guide_temperature. + q_p_c_sharp = q_p_c_raw ** (1.0 / self.guide_temperature) + q_p_c_sharp = torch.clamp(q_p_c_sharp, min=1e-8) # ← add this + q_p_c_sharp = q_p_c_sharp / q_p_c_sharp.sum(dim=1, keepdim=True) + conc_c_guide = torch.clamp(self.global_scale * q_p_c_sharp, min=1e-3) + + # Sample p_c from a single learned Dirichlet per clonotype. + pyro.sample("p_c", dist.Dirichlet(conc_c_guide)) + + with pyro.plate("ct_plate", self.ct_count): + init_mat = self.clone_phen_prior[self.ct_to_c, :] + init_mat = init_mat * self.guide_init_scale + 1e-3 + init_mat = init_mat.to(x.device) + if "q_p_ct_raw" not in pyro.get_param_store(): + q_p_ct_raw = pyro.param( + "q_p_ct_raw", + init_mat.clone().detach(), # Make sure it's not a leaf + constraint=dist.constraints.positive, + ) + else: + q_p_ct_raw = pyro.param("q_p_ct_raw") + + bad_ct = ~torch.isfinite(q_p_ct_raw) + if bad_ct.any(): + q_p_ct_raw = torch.where(bad_ct, init_mat.to(q_p_ct_raw.device), q_p_ct_raw) + + q_p_ct_sharp = q_p_ct_raw ** (1.0 / self.guide_temperature) + q_p_ct_sharp = torch.clamp(q_p_ct_sharp, min=1e-8) + q_p_ct_sharp = q_p_ct_sharp / q_p_ct_sharp.sum(dim=1, keepdim=True) + conc_ct_guide = torch.clamp(self.local_scale * q_p_ct_sharp, min=1e-3) + pyro.sample("p_ct", dist.Dirichlet(conc_ct_guide)) + + z_loc, z_scale, _ = self.encoder(x, batch_idx) + z_scale = torch.clamp(z_scale, min=1e-3, max=10.0) + + with pyro.plate("data", batch_size) as idx: + latent_posterior = dist.Normal(z_loc, z_scale) + with pyro.poutine.scale(scale=self.kl_weight): + z = pyro.sample("latent", latent_posterior.to_event(1)) + + @auto_move_data + def get_latent(self, tensor_dict: Dict[str, torch.Tensor]): + x = tensor_dict[REGISTRY_KEYS.X_KEY] + batch_idx = tensor_dict[REGISTRY_KEYS.BATCH_KEY].long() + z_loc, _, _ = self.encoder(x, batch_idx) + if z_loc.ndim == 3: + z_loc = z_loc.mean(dim=1) + return z_loc.cpu() + + @torch.no_grad() + def get_p_ct(self): + from pyro import get_param_store + + param_store = get_param_store() + q_p_ct_raw = param_store["q_p_ct_raw"] + bad = ~torch.isfinite(q_p_ct_raw) + if bad.any(): + n_phen = q_p_ct_raw.shape[1] + q_p_ct_raw = torch.where(bad, torch.ones_like(q_p_ct_raw) / n_phen, q_p_ct_raw) + if self.guide_temperature != 1.0: + q_p_ct_sharp = q_p_ct_raw ** (1.0 / self.guide_temperature) + q_p_ct_sharp = q_p_ct_sharp / q_p_ct_sharp.sum(dim=1, keepdim=True) + else: + q_p_ct_sharp = q_p_ct_raw / q_p_ct_raw.sum(dim=1, keepdim=True) + return q_p_ct_sharp diff --git a/tcri/model/_priors.py b/tcri/model/_priors.py new file mode 100644 index 0000000..efa05ef --- /dev/null +++ b/tcri/model/_priors.py @@ -0,0 +1,149 @@ +"""Prior distributions for the TCRI generative model. + +`VampPrior` is a learnable mixture-of-encoders prior over the latent space; +`MixtureDirichlet` is the per-clonotype mixture-of-Dirichlets used as the +clonotype-level phenotype prior. +""" +import torch +import pyro.distributions as dist + +__all__ = ["MixtureDirichlet", "VampPrior"] + + +class VampPrior(torch.nn.Module): + def __init__(self, pseudo_inputs, encoder): + """ + Args: + pseudo_inputs (torch.Tensor): Initial pseudo-inputs of shape (K, input_dim). + encoder (torch.nn.Module): Encoder that takes an input and returns (mean, log_var) + for the approximate posterior q(z|x). + """ + super(VampPrior, self).__init__() + # Learnable pseudo-inputs; these are optimized during training. + self.pseudo_inputs = torch.nn.Parameter(pseudo_inputs) + self.encoder = encoder + + def get_mixture(self): + """ + Constructs the VampPrior as a uniform mixture of q(z|u_k) for each pseudo-input u_k. + """ + # Compute the approximate posterior parameters for each pseudo-input. + # Expected output shapes: means and log_vars: (K, latent_dim) + K = self.pseudo_inputs.size(0) + # Create a dummy categorical argument; unsqueeze to have shape (K, 1) + dummy_batch = torch.zeros(K, dtype=torch.long, device=self.pseudo_inputs.device).unsqueeze(1) + means, log_vars, _ = self.encoder(self.pseudo_inputs, dummy_batch) + scales = torch.sqrt(torch.exp(log_vars)) + component_dist = dist.Independent(dist.Normal(means, scales), 1) + mixture_weights = torch.ones(K, device=self.pseudo_inputs.device) / K + mixture = dist.MixtureSameFamily( + dist.Categorical(mixture_weights), + component_dist + ) + return mixture + + def log_prob(self, z): + """ + Computes log p(z) under the VampPrior. + """ + return self.get_mixture().log_prob(z) + + def sample(self, sample_shape=torch.Size()): + """ + Draws samples from the VampPrior. + """ + return self.get_mixture().sample(sample_shape) + + +class MixtureDirichlet(dist.TorchDistribution): + """ + Mixture of Dirichlet distributions parameterized by mixture weights and concentration parameters." + """ + arg_constraints = { + "mixture_weights": dist.constraints.simplex, # shape: batch_shape + (B,) + "concentration": dist.constraints.positive, # shape: batch_shape + (B, K) + } + support = dist.constraints.simplex # each sample is a simplex over K categories + has_rsample = False + + def __init__( + self, + mixture_weights: torch.Tensor, + concentration: torch.Tensor, + validate_args=None, + ): + """ + mixture_weights: Tensor of shape batch_shape + (B,), with each row summing to 1. + concentration: Tensor of shape batch_shape + (B, K), where K is the number of categories. + """ + self.mixture_weights = mixture_weights + # Clamp concentrations to ensure positivity. + self.concentration = torch.clamp(concentration, min=1e-3) + # Determine batch shape, B, and K. + batch_shape = self.mixture_weights.shape[:-1] + self.B = self.mixture_weights.size(-1) + self.K = self.concentration.size(-1) + event_shape = (self.K,) + super(MixtureDirichlet, self).__init__( + batch_shape, event_shape, validate_args=validate_args + ) + + def sample(self, sample_shape=torch.Size()): + """ + Returns a sample of shape: sample_shape + batch_shape + (K,). + For each batch element, first sample a mixture component, then sample from the corresponding Dirichlet. + """ + # Create categorical for mixture weights. + cat = dist.Categorical(self.mixture_weights) + # Sample mixture indices; shape: sample_shape + batch_shape. + mixture_idx = cat.sample(sample_shape) + full_shape = mixture_idx.shape # sample_shape + batch_shape + + # Expand concentration to shape: sample_shape + batch_shape + (B, K). + target_shape = sample_shape + self.concentration.shape + expanded_concentration = self.concentration.expand(target_shape) + + # Flatten the sample and batch dimensions. + flat_shape = (-1, self.B, self.K) + flat_concentration = expanded_concentration.reshape(flat_shape) + flat_idx = mixture_idx.reshape(-1) # shape: (num_samples,) + + # Select the concentration parameters corresponding to the sampled mixture index. + selected_concentration = flat_concentration[ + torch.arange(flat_idx.size(0)), flat_idx + ] + + # Sample from the Dirichlet for each sample. + flat_samples = dist.Dirichlet(selected_concentration).sample() + # Reshape to sample_shape + batch_shape + (K,). + return flat_samples.reshape(full_shape + (self.K,)) + + def log_prob(self, value): + device = value.device # get the device from input tensor + + # Move tensors explicitly to the same device + value_expanded = value.unsqueeze(-2).to(device) + expanded_concentration = self.concentration.expand( + value.shape[:-1] + (self.B, self.K) + ).to(device) + + d = dist.Dirichlet(expanded_concentration) + + component_log_probs = d.log_prob( + value_expanded.expand(expanded_concentration.shape) + ) + + expanded_weights = self.mixture_weights.expand(value.shape[:-1] + (self.B,)).to(device) + mixture_log = torch.log(expanded_weights) + + return torch.logsumexp(mixture_log + component_log_probs, dim=-1) + + def score_parts(self, value): + # Compute log probability. + lp = self.log_prob(value) + # Return dummy zeros for the score function and entropy terms. + zeros = torch.zeros_like(lp) + return lp, zeros, zeros + + def __call__(self, *args, **kwargs): + return self.sample(*args, **kwargs) diff --git a/tcri/model/_training.py b/tcri/model/_training.py new file mode 100644 index 0000000..85b0746 --- /dev/null +++ b/tcri/model/_training.py @@ -0,0 +1,169 @@ +"""Training plan for the TCRI model plus the archetype initializer. + +`UnifiedTrainingPlan` adds classification/reconstruction diagnostics and a +`validation_step` (logs ``elbo_validation`` for scvi early stopping) on top of +Pyro's ELBO step. `build_archetypes` K-means-clusters the clone x phenotype +matrix to seed the Dirichlet mixture (returns centers AND labels). +""" +import numpy as np +import torch +import torch.nn.functional as F + +from scvi.train import PyroTrainingPlan +from pyro.infer import TraceEnum_ELBO, Trace_ELBO +from sklearn.cluster import KMeans + +from ._module import TCRIModule + +__all__ = ["UnifiedTrainingPlan", "build_archetypes"] + + +def build_archetypes(clone_phenotype_prior, K=4): + kmeans = KMeans(n_clusters=K, random_state=42) + labels = kmeans.fit_predict(clone_phenotype_prior) + centers = kmeans.cluster_centers_ + centers = np.clip(centers, 1e-8, None) + centers = centers / centers.sum(axis=1, keepdims=True) + return centers, labels + + +class UnifiedTrainingPlan(PyroTrainingPlan): + """ + Training plan that includes classification, reconstruction losses, + KL warmup, plus a validation_step that logs 'elbo_validation' so scvi's + early stopping can monitor it. + """ + + def __init__( + self, + module: TCRIModule, + n_steps_kl_warmup: int = 1000, + reconstruction_loss_scale: float = 1e-2, + num_particles: int = 5, + optimizer_config: dict = None, + **kwargs, + ): + self.num_particles = num_particles + if module.use_enumeration: + print("Using Enumeration") + self._loss_fn = TraceEnum_ELBO( + max_plate_nesting=3, num_particles=self.num_particles + ) + else: + self._loss_fn = Trace_ELBO() + + if optimizer_config is None: + optimizer_config = {"lr": 1e-3, "betas": (0.9, 0.999), "eps": 1e-5, + "weight_decay": 1e-4} + self.optimizer_config = optimizer_config + + # Hand the optimizer settings to PYRO's optimizer — the one inside SVI that + # actually descends the ELBO gradients. Previously this class overrode + # `configure_optimizers` with a real torch Adam over every module parameter; + # that override replaced scvi's deliberate no-op shim ("a shim optimizer ... + # to keep Lightning happy") and ran AFTER SVI.step() had already stepped and + # ZEROED the gradients. Stepping Adam on zero gradients is not a no-op: the + # weight-decay term becomes the whole gradient, and Adam's own normalization + # (g/sqrt(g^2)) strips its magnitude, so the update degenerates to ~lr*sign(p) + # — a scale-free shrink, not proportional L2. Measured effect: network weights + # held at ~2.4x smaller than without it. It also meant `lr` never reached the + # real optimizer (Pyro always used scvi's hard-coded 1e-3). + super().__init__( + module, + n_steps_kl_warmup=n_steps_kl_warmup, + optim_kwargs={ + "lr": optimizer_config["lr"], + "betas": optimizer_config["betas"], + "eps": optimizer_config["eps"], + "weight_decay": optimizer_config["weight_decay"], + }, + **kwargs, + ) + + self.n_steps_kl_warmup = n_steps_kl_warmup + self.reconstruction_loss_scale = reconstruction_loss_scale + self._my_global_step = 0 + + @property + def loss(self): + return self._loss_fn + + # NOTE: configure_optimizers is deliberately NOT overridden — scvi's base class + # returns a shim over a single dummy parameter purely to advance Lightning's step + # counter. All real optimization happens in Pyro's SVI (see __init__). + + def training_step(self, batch, batch_idx): + # ── KL warmup ──────────────────────────────────────────── + if self.n_steps_kl_warmup > 0 and self._my_global_step < self.n_steps_kl_warmup: + kl_weight = max(1e-6, self.module.kl_weight_max * (self._my_global_step / self.n_steps_kl_warmup)) + else: + kl_weight = self.module.kl_weight_max + self.module.kl_weight = kl_weight + + # ── Pyro ELBO step ─────────────────────────────────────── + loss_dict = super().training_step(batch, batch_idx) + device = next(self.module.parameters()).device + + if not isinstance(loss_dict["loss"], torch.Tensor): + loss_dict["loss"] = torch.tensor(loss_dict["loss"], device=device, requires_grad=True) + else: + loss_dict["loss"] = loss_dict["loss"].to(device) + + # ── Diagnostics (no gradient contribution) ─────────────── + with torch.no_grad(): + z_diag = self.module.get_latent(batch).to(device) + idx_diag = batch["indices"].long().view(-1).to(device) + cls_logits = self.module.classifier(z_diag) + ct_idx = self.module.ct_array[idx_diag] + p_ct_prior = self.module.get_p_ct()[ct_idx].to(device) + prior_log = torch.log(p_ct_prior + 1e-8) + + 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 + + probs = F.softmax(local_logits, dim=-1) + kl_div = F.kl_div(probs.log(), p_ct_prior, reduction='batchmean') + entropy = -torch.sum(probs * torch.log(probs + 1e-8), dim=-1).mean() + confidence = (probs**2).sum(dim=-1).mean() + + self.log("kl_divergence_with_prior_train", kl_div, prog_bar=False, on_epoch=True) + self.log("entropy_train", entropy, prog_bar=False, on_epoch=True) + self.log("confidence_train", confidence, prog_bar=False, on_epoch=True) + + self._my_global_step += 1 + return loss_dict + + def validation_step(self, batch, batch_idx): + with torch.no_grad(): + self.module.eval() + val_dict = super().training_step(batch, batch_idx) + self.module.train() + + device = next(self.module.parameters()).device + + if not isinstance(val_dict["loss"], torch.Tensor): + val_dict["loss"] = torch.tensor(val_dict["loss"], device=device) + else: + val_dict["loss"] = val_dict["loss"].to(device) + + # ── Diagnostic only ────────────────────────────────────── + z_batch = self.module.get_latent(batch).to(device) + idx = batch["indices"].long().view(-1).to(device) + cls_logits = self.module.classifier(z_batch) + ct_idx = self.module.ct_array[idx] + p_ct_prior = self.module.get_p_ct()[ct_idx].to(device) + prior_log = torch.log(p_ct_prior + 1e-8) + + 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 + + probs = F.softmax(local_logits, dim=-1) + kl_divergence = F.kl_div(probs.log(), p_ct_prior, reduction='batchmean') + self.log("kl_divergence_with_prior_val", kl_divergence, prog_bar=False, on_epoch=True) + + self.log("elbo_validation", val_dict["loss"], prog_bar=True, on_epoch=True) + return val_dict diff --git a/tcri/plotting/__init__.py b/tcri/plotting/__init__.py index 7fed5f3..a718e95 100644 --- a/tcri/plotting/__init__.py +++ b/tcri/plotting/__init__.py @@ -1 +1,16 @@ -from ..plotting._plotting import * +"""``tcri.pl`` — plotting. The four ``tl``↔``pl`` metric twins (cache renderers over the +engine-backed metrics) plus the shared colors. Explicit ``__all__``; no ``import *``. +""" +from ._entropy import clonotypic_entropy, phenotypic_entropy +from ._mutual_information import mutual_information +from ._flux import phenotypic_flux +from ._colors import tcri_colors, resolve_palette + +__all__ = [ + "clonotypic_entropy", + "phenotypic_entropy", + "mutual_information", + "phenotypic_flux", + "tcri_colors", + "resolve_palette", +] diff --git a/tcri/plotting/_base.py b/tcri/plotting/_base.py new file mode 100644 index 0000000..531ca1d --- /dev/null +++ b/tcri/plotting/_base.py @@ -0,0 +1,78 @@ +"""``pl`` private plotting engine (§8.5) — the generic tidy-DataFrame box/strip renderer +shared by the metric twins, and the save/show finisher. + +The twins are **cache renderers**: they call the ``tl`` metric (which already does its own +``groupby`` via full-space ``clones=`` restriction), get a tidy DataFrame, and hand it here. +Nothing here slices the AnnData or re-computes metric math (fixes the old ``tcri_boxplot`` +slice-and-call pattern that tripped the engine alignment guard). +""" +from __future__ import annotations + +import matplotlib.pyplot as plt + +__all__: list[str] = [] # private module + + +def _finish(fig, ax, *, save=None, show=None): + if save: + fig.savefig(save, bbox_inches="tight", dpi=150) + if show: + plt.show() + return ax + + +def _metric_boxplot(df, *, x, y, hue=None, order=None, hue_order=None, palette=None, + ax=None, figsize=(8, 4), s=20, ylabel=None, rotation=90): + """Box + strip of a tidy metric DataFrame: box over the ``x`` categories, optional + ``hue`` split, individual units as strip dots. Returns ``(fig, ax)``.""" + import seaborn as sns + + if ax is None: + fig, ax = plt.subplots(1, 1, figsize=figsize) + else: + fig = ax.figure + + if df is None or len(df) == 0 or y not in getattr(df, "columns", []): + ax.text(0.5, 0.5, f"no data for {y}", ha="center", va="center", transform=ax.transAxes) + ax.set_ylabel(ylabel if ylabel is not None else y) + return fig, ax + d = df.dropna(subset=[y]) + if len(d) == 0: + ax.text(0.5, 0.5, f"no finite {y}", ha="center", va="center", transform=ax.transAxes) + ax.set_ylabel(ylabel if ylabel is not None else y) + return fig, ax + if order is None and x in d: + order = (d.groupby(x)[y].median().sort_values(ascending=False).index.tolist()) + + common = dict(data=d, x=x, y=y, order=order, ax=ax) + if hue is not None and hue in d.columns: + common.update(hue=hue, hue_order=hue_order) + sns.boxplot(**common, palette=palette, showfliers=False, + boxprops=dict(alpha=0.5)) + sns.stripplot(**common, palette=palette, dodge=hue is not None, size=s / 3, + edgecolor="black", linewidth=0.3) + ax.set_ylabel(ylabel if ylabel is not None else y) + ax.set_xlabel(x) + for label in ax.get_xticklabels(): + label.set_rotation(rotation) + if hue is not None and ax.get_legend() is not None: + ax.legend(bbox_to_anchor=(1.02, 1.0), loc="upper left", frameon=False, fontsize=8) + return fig, ax + + +def _barplot(series, *, ylabel, palette=None, ax=None, figsize=(8, 4), rotation=90): + """Fallback for a no-groupby metric (a Series over phenotypes/clones): a simple bar.""" + if ax is None: + fig, ax = plt.subplots(1, 1, figsize=figsize) + else: + fig = ax.figure + s = series.dropna() + ax.bar(range(len(s)), s.values, + color=(palette if isinstance(palette, str) else None) or tcri_bar_color) + ax.set_xticks(range(len(s))) + ax.set_xticklabels(list(s.index), rotation=rotation) + ax.set_ylabel(ylabel) + return fig, ax + + +tcri_bar_color = "#66D9EF" diff --git a/tcri/plotting/_colors.py b/tcri/plotting/_colors.py new file mode 100644 index 0000000..4c4da56 --- /dev/null +++ b/tcri/plotting/_colors.py @@ -0,0 +1,28 @@ +"""``pl`` colors — the shared categorical palette and `resolve_palette` (§8.6).""" +from __future__ import annotations + +__all__ = ["tcri_colors", "resolve_palette"] + +tcri_colors = [ + "#AE81FF", "#FD971F", "#66D9EF", "#A6E22E", "#F92672", "#E6DB74", "#75715E", + "#D65F0E", "#004d47", "#D291BC", "#3A506B", "#5D8A5E", "#A6A1E2", "#E97451", + "#6C8D67", "#832232", "#1E1E1E", "#F92659", "#272822", "#8B4513", +] + + +def resolve_palette(adata, columns, *, palette=None): + """Assign `tcri_colors` to each `obs` column's categories, store in + ``uns["_colors"]``, and return ``{col: {category: color}}``. + + Mutates ``adata`` **in place** (fixes the old ``set_color_palette`` bug of writing + onto a throwaway copy). ``palette`` overrides the default color cycle. + """ + cols = [columns] if isinstance(columns, str) else list(columns) + cycle = list(palette) if palette is not None else tcri_colors + out = {} + for col in cols: + cats = adata.obs[col].astype("category").cat.categories.tolist() + mapping = {c: cycle[i % len(cycle)] for i, c in enumerate(cats)} + adata.uns[f"{col}_colors"] = [mapping[c] for c in cats] + out[col] = mapping + return out diff --git a/tcri/plotting/_entropy.py b/tcri/plotting/_entropy.py new file mode 100644 index 0000000..51915e3 --- /dev/null +++ b/tcri/plotting/_entropy.py @@ -0,0 +1,54 @@ +"""``pl.clonotypic_entropy`` / ``pl.phenotypic_entropy`` (§8.1) — cache renderers that call +the ``tl`` twin and box/strip the tidy result. No metric math here.""" +from __future__ import annotations + +from ._base import _barplot, _finish, _metric_boxplot + +__all__ = ["clonotypic_entropy", "phenotypic_entropy"] + + +def clonotypic_entropy(adata, *, covariate=None, groupby=None, splitby=None, n_samples=0, + temperature=1.0, clones=None, weighted=False, normalized=True, + order=None, hue_order=None, palette=None, ax=None, figsize=(8, 4), + save=None, show=None, return_df=False): + """Per-phenotype clonotypic entropy (bits), boxed over ``groupby`` units and split by + ``splitby`` (x = phenotype).""" + from .. import tools as tl + + res = tl.clonotypic_entropy( + adata, covariate=covariate, groupby=groupby, splitby=splitby, n_samples=n_samples, + temperature=temperature, clones=clones, weighted=weighted, normalized=normalized, + ) + if return_df: + return res + if groupby is None: + fig, ax = _barplot(res, ylabel="clonotypic entropy (bits)", palette=palette, ax=ax, figsize=figsize) + else: + fig, ax = _metric_boxplot(res, x="phenotype", y="clonotypic_entropy", hue=splitby, + order=order, hue_order=hue_order, palette=palette, ax=ax, + figsize=figsize, ylabel="clonotypic entropy (bits)") + return _finish(fig, ax, save=save, show=show) + + +def phenotypic_entropy(adata, *, covariate=None, groupby=None, splitby=None, n_samples=0, + temperature=1.0, clones=None, weighted=False, normalized=True, + order=None, hue_order=None, palette=None, ax=None, figsize=(8, 4), + save=None, show=None, return_df=False): + """Distribution of per-clone phenotypic entropy (plasticity), boxed by ``splitby`` (or + ``groupby``) with each clone a dot.""" + from .. import tools as tl + + res = tl.phenotypic_entropy( + adata, covariate=covariate, groupby=groupby, splitby=splitby, n_samples=n_samples, + temperature=temperature, clones=clones, weighted=weighted, normalized=normalized, + ) + if return_df: + return res + if groupby is None: + fig, ax = _barplot(res, ylabel="phenotypic entropy (bits)", palette=palette, ax=ax, figsize=figsize) + else: + x = splitby if (splitby is not None and splitby in res.columns) else groupby + fig, ax = _metric_boxplot(res, x=x, y="phenotypic_entropy", hue=None, order=order, + palette=palette, ax=ax, figsize=figsize, + ylabel="phenotypic entropy (bits)") + return _finish(fig, ax, save=save, show=show) diff --git a/tcri/plotting/_flux.py b/tcri/plotting/_flux.py new file mode 100644 index 0000000..8a6734d --- /dev/null +++ b/tcri/plotting/_flux.py @@ -0,0 +1,37 @@ +"""``pl.phenotypic_flux`` (§8.3) — cache renderer for per-clonotype phenotype-distribution +flux across a covariate ``order``. Renders a box of the per-clone flux magnitudes split by +``splitby``. (A phenotype-flow Sankey over ``order`` is a deferred enhancement, tracked in +``REFACTOR_NOTES``; the tidy flux values are available via ``return_axes=True``.) +""" +from __future__ import annotations + +from ._base import _finish, _metric_boxplot + +__all__ = ["phenotypic_flux"] + + +def phenotypic_flux(adata, *, order, groupby=None, splitby=None, n_samples=0, temperature=1.0, + clones=None, weighted=False, distance_metric="l1", palette=None, ax=None, + figsize=(8, 4), save=None, show=None, return_axes=False): + """Per-clone phenotype-distribution flux from ``order[0]`` to ``order[-1]``, boxed by + cohort. ``order`` is the covariate sequence (>= 2 values).""" + from .. import tools as tl + from .. import _keys as K + + order = list(order) + if len(order) < 2: + raise ValueError("phenotypic_flux needs `order` with >= 2 covariate values.") + cov_from, cov_to = order[0], order[-1] + gb = groupby if groupby is not None else adata.uns[K.METADATA]["batch_col"] + + res = tl.phenotypic_flux( + adata, cov_from=cov_from, cov_to=cov_to, groupby=gb, splitby=splitby, + n_samples=n_samples, temperature=temperature, clones=clones, weighted=weighted, + distance_metric=distance_metric, + ) + if return_axes: + return res + x = splitby if (splitby is not None and splitby in res.columns) else gb + fig, ax = _metric_boxplot(res, x=x, y="phenotypic_flux", hue=None, palette=palette, ax=ax, + figsize=figsize, ylabel=f"phenotypic flux ({distance_metric})") + return _finish(fig, ax, save=save, show=show) diff --git a/tcri/plotting/_mutual_information.py b/tcri/plotting/_mutual_information.py new file mode 100644 index 0000000..bc19c16 --- /dev/null +++ b/tcri/plotting/_mutual_information.py @@ -0,0 +1,31 @@ +"""``pl.mutual_information`` (§8.2) — cache renderer: per-unit MI boxed by ``splitby``.""" +from __future__ import annotations + +from ._base import _finish, _metric_boxplot + +__all__ = ["mutual_information"] + + +def mutual_information(adata, *, covariate=None, groupby=None, splitby=None, n_samples=0, + temperature=1.0, clones=None, weighted=False, normalized=True, + normalize_mode="min", order=None, hue_order=None, palette=None, ax=None, + figsize=(8, 4), save=None, show=None, return_df=False): + """Clone↔phenotype MI (bits, normalized). With ``groupby`` (units, e.g. patient) and + ``splitby`` (cohort, e.g. response): one MI per unit, boxed by cohort.""" + from .. import tools as tl + from .. import _keys as K + + # a boxplot needs per-unit MI; default the aggregation unit to the batch (patient) column + gb = groupby if groupby is not None else adata.uns[K.METADATA]["batch_col"] + + res = tl.mutual_information( + adata, covariate=covariate, groupby=gb, splitby=splitby, n_samples=n_samples, + temperature=temperature, clones=clones, weighted=weighted, normalized=normalized, + normalize_mode=normalize_mode, + ) + if return_df: + return res + x = splitby if (splitby is not None and splitby in res.columns) else gb + fig, ax = _metric_boxplot(res, x=x, y="MI", hue=None, order=order, palette=palette, + ax=ax, figsize=figsize, ylabel="mutual information (bits)") + return _finish(fig, ax, save=save, show=show) diff --git a/tcri/plotting/_plotting.py b/tcri/plotting/_plotting.py deleted file mode 100644 index 1a17053..0000000 --- a/tcri/plotting/_plotting.py +++ /dev/null @@ -1,1429 +0,0 @@ -import numpy as np -from .. import _keys as K -import pandas as pd -import seaborn as sns -import matplotlib.pyplot as plt -import matplotlib.patches as mpatches -from scipy.cluster.hierarchy import dendrogram, linkage -import scanpy as sc -from gseapy import dotplot -import numpy as np, pandas as pd, torch, umap -from tqdm.auto import tqdm -from scvi import REGISTRY_KEYS -from scipy.stats import mannwhitneyu - -import collections -import operator -import itertools - -from ..utils._utils import probabilities -from ._sankey import SankeyNode -from ..preprocessing._preprocessing import clone_size, joint_distribution, joint_distribution_posterior -from ..metrics._metrics import clonotypic_entropy as centropy -from ..metrics._metrics import phenotypic_entropy as pentropy -from ..metrics._metrics import clonality as clonality_tl -from ..metrics._metrics import flux as flux_tl -from ..metrics._metrics import mutual_information as mutual_information_tl - - -import warnings -warnings.filterwarnings('ignore') - -sc._settings.settings._vector_friendly=True - -# ╭─ colour / pretty-print helpers ─────────────────────────────────────────╮ -RESET = "\x1b[0m"; BOLD = "\x1b[1m"; DIM = "\x1b[2m" -GRN = "\x1b[32m"; CYN = "\x1b[36m"; MAG = "\x1b[35m"; YLW = "\x1b[33m"; RED = "\x1b[31m" - -from .._console import _ok, _info, _warn, _fin -# ╰──────────────────────────────────────────────────────────────────────────╯ - -red = "#cd442a" -yellow = "#f0bd00" -green = "#7e9437" - -tcri_colors = [ - red, - yellow, - green, - "#004d47", # Darker Teal - "#AE81FF", # Purple - "#FD971F", # Orange - "#E6DB74", # Yellow - "#A6E22E", # Green - "#F28E7F", # Salmon - "#75715E", # Brown - "#F92659", # Pink - "#D65F0E", # Abricos - "#66D9EF", # Blue - "#F92672", # Red - "#1E1E1E", # Black - "#272822", # Background - "#D291BC", # Soft Pink - "#3A506B", # Dark Slate Blue - "#5D8A5E", # Sage Green - "#A6A1E2", # Dull Lavender - "#E97451", # Burnt Sienna - "#6C8D67", # Muted Lime Green - "#832232", # Dim Maroon - "#669999", # Desaturated Cyan - "#C08497", # Dusty Rose - "#587B7F", # Ocean Blue - "#9A8C98", # Muted Purple - "#F3B61F", # Goldenrod - "#FFD8B1", # Light Peach - "#88AB75", # Moss Green - "#C38D94", # Muted Rose - "#6D6A75", # Purple Gray -] - -sns.set_palette(sns.color_palette(tcri_colors)) - -from ..metrics._metrics import mi_compare as mi_compare_tl -from .._stats import auc_and_label_permutation, bootstrap_auc - - -def mi_compare(adata, groupby, groups=None, treatment=None, n_samples=50, - point="median", palette=None, patient_col=None, - clone_col=None, covariate_col=None, - ax=None, save=None, seed=42, verbose=True, **mi_kwargs): - - result = mi_compare_tl( - adata, groupby=groupby, groups=groups, treatment=treatment, - n_samples=n_samples, patient_col=patient_col, - clone_col=clone_col, covariate_col=covariate_col, - verbose=verbose, **mi_kwargs - ) - - summary = result["summary"] - pairs = result["pairs"] - covariates = result["params"]["covariates"] - - n_pairs = len(pairs) - n_covs = len(covariates) - fig_needed = ax is None - if fig_needed: - fig, axes = plt.subplots(n_pairs, n_covs, - figsize=(6 * n_covs, 5 * n_pairs), - squeeze=False) - else: - axes = np.atleast_2d(ax) - - all_stats = {} - rng = np.random.default_rng(seed) - - for i, (g0, g1) in enumerate(pairs): - if palette is None: - pal = {g0: tcri_colors[0], g1: tcri_colors[1]} - else: - pal = palette - - for j, cov in enumerate(covariates): - cur_ax = axes[i, j] - pat = summary[ - (summary["covariate"] == cov) & - (summary["group"].isin([g0, g1])) - ].copy() - pat["point"] = pat[point] - - g0_vals = pat.loc[pat["group"] == g0, "point"].values - g1_vals = pat.loc[pat["group"] == g1, "point"].values - - _, p_mwu = mannwhitneyu(g0_vals, g1_vals, alternative="two-sided") - auc, p_perm, _, perm_mode = auc_and_label_permutation( - pat["point"].values, pat["group"].values, pos_label=g1) - auc_lo, auc_hi = bootstrap_auc( - pat["point"].values, pat["group"].values, pos_label=g1) - - sns.boxplot(data=pat, x="group", y="point", order=[g0, g1], - width=0.45, showcaps=True, showfliers=False, - boxprops=dict(facecolor="white", alpha=0.9), - medianprops=dict(color="k", lw=2), ax=cur_ax) - - for x_pos, grp in enumerate([g0, g1]): - sub = pat[pat["group"] == grp].sort_values("point").reset_index(drop=True) - jitter = rng.uniform(-0.18, 0.18, len(sub)) - for k, row in sub.iterrows(): - xj = x_pos + jitter[k] - cur_ax.plot([xj, xj], [row["lo"], row["hi"]], - color=pal[grp], lw=1.2, alpha=0.45) - cur_ax.scatter(xj, row["point"], color=pal[grp], s=65, zorder=3) - cur_ax.text(xj, row["hi"] + 0.004, row["patient"], fontsize=6, - ha="center", va="bottom", color=pal[grp]) - - cur_ax.set_xlabel("") - cur_ax.set_ylabel(f"{cov} TCRi NMI") - cur_ax.set_title( - f"{cov} NMI: {g0} vs {g1}\n" - f"MWU p={p_mwu:.3g} | AUROC={auc:.2f} [{auc_lo:.2f}, {auc_hi:.2f}]" - f" | label-perm p={p_perm:.3g} ({perm_mode})", - fontweight="bold", fontsize=9) - sns.despine(ax=cur_ax) - - all_stats[(g0, g1, cov)] = { - "p_mwu": p_mwu, "auc": auc, - "auc_ci": (auc_lo, auc_hi), - "p_perm": p_perm, "perm_mode": perm_mode, - } - - if fig_needed: - fig.tight_layout() - if save is not None: - fig.savefig(save, dpi=200, bbox_inches="tight") - - return result, all_stats - -def compare_phenotypes(adata, variable1, variable2): - df = adata.obs[[variable1,variable2]] - df=pd.crosstab(df[variable1],df[variable2],normalize='index') - return sns.heatmap(df) - -def compare_joint_distribution(adata, temperature=1): - # ----------------------------- - # 1. Get Model-Inferred Distributions - # ----------------------------- - # Create a dictionary mapping each tissue (treatment group) to its clone phenotype DataFrame - covariate_col = adata.uns[K.METADATA]["covariate_col"] - model_dists = dict() - for tissue in set(adata.obs[covariate_col]): - # Use your function to get the inferred p_ct distribution (with temperature scaling) - df_tissue = joint_distribution(adata,tissue, temperature=temperature) - df_tissue[covariate_col] = tissue - # Set clonotype_id as index for easier merging/comparison later - #df_tissue.set_index("clonotype_id", inplace=True) - model_dists[tissue] = df_tissue - - # Concatenate the inferred distributions from all tissues into a single DataFrame - df_model = pd.concat(model_dists.values(), axis=0) - - empirical_dists = dict() - clonotype_col = model.adata_manager.registry["clonotype_col"] - phenotype_col = model.adata_manager.registry["phenotype_col"] - - for tissue in set(adata.obs[covariate_col]): - # Filter for cells in the given tissue - adata_tissue = adata[adata.obs[covariate_col] == tissue].copy() - # Group by clonotype and compute normalized counts of each phenotype - emp = ( - adata_tissue.obs.groupby(clonotype_col)[phenotype_col] - .value_counts(normalize=True) - .unstack(fill_value=0) - ) - # Ensure that the DataFrame uses the actual phenotype category names as columns. - # If some phenotype categories are missing in a tissue, add them with 0. - phenotype_categories = list(adata.obs[phenotype_col].astype("category").cat.categories) - for ph in phenotype_categories: - if ph not in emp.columns: - emp[ph] = 0.0 - # Reorder columns - emp = emp[phenotype_categories] - emp[covariate_col] = tissue - # Use clonotype ID (from the index) as a column if needed - emp.index.name = "clonotype_id" - empirical_dists[tissue] = emp - - # Concatenate the empirical distributions from all tissues into one DataFrame. - df_empirical = pd.concat(empirical_dists.values(), axis=0) - - # ----------------------------- - # 3. Compare Distributions: Plotting Side-by-Side - # ----------------------------- - # We'll loop over the unique tissues and for each, plot the inferred (model) and empirical distributions. - unique_tissues = df_model[covariate_col].unique() - n_tissues = len(unique_tissues) - - fig, axes = plt.subplots(n_tissues, 4, figsize=(20, 4 * n_tissues), - gridspec_kw={'width_ratios': [1, 4, 1, 4]}) - - # In case there's only one tissue, ensure axes is 2D. - if n_tissues == 1: - axes = np.expand_dims(axes, axis=0) - - for i, tissue in enumerate(unique_tissues): - # Select the rows for the current tissue for both distributions - model_data = df_model[df_model[covariate_col] == tissue] - empirical_data = df_empirical[df_empirical[covariate_col] == tissue] - - # Determine the phenotype columns (assumed to be common to both) - phenotype_cols = [col for col in model_data.columns if col not in ["clonotype_index", covariate_col]] - - # --- Model-Inferred Distribution --- - # Compute hierarchical clustering for the model distribution. - Z_model = linkage(model_data[phenotype_cols], method='average') - dendro_model = dendrogram(Z_model, orientation='left', ax=axes[i, 0], no_labels=True) - # Order the data - ordered_model = model_data.iloc[dendro_model['leaves']] - sns.heatmap(ordered_model[phenotype_cols], ax=axes[i, 1], cmap="viridis", cbar=True) - axes[i, 1].set_title(f"Model-Inferred: {tissue}") - axes[i, 0].set_title("Dendrogram") - axes[i, 0].set_xticks([]) - axes[i, 0].set_yticks([]) - axes[i, 1].set_yticklabels([]) - - # --- Empirical Distribution --- - # Compute hierarchical clustering for the empirical distribution. - Z_emp = linkage(empirical_data[phenotype_cols], method='average') - dendro_emp = dendrogram(Z_emp, orientation='left', ax=axes[i, 2], no_labels=True) - ordered_emp = empirical_data.iloc[dendro_emp['leaves']] - sns.heatmap(ordered_emp[phenotype_cols], ax=axes[i, 3], cmap="viridis", cbar=True) - axes[i, 3].set_title(f"Empirical: {tissue}") - axes[i, 2].set_title("Dendrogram") - axes[i, 2].set_xticks([]) - axes[i, 2].set_yticks([]) - axes[i, 3].set_yticklabels([]) - - plt.tight_layout() - plt.show() - -def _phenotype_mass_per_clone(adata, covariate, clones, normalize): - """Per-clone phenotype mass at one covariate. - - Returns dict {clone_id -> np.ndarray(n_phen)} aligned to - adata.uns['tcri_phenotype_categories']. Rows of joint_distribution - sum to 1 (per clone-celltype). Multiple ct rows for the same clone - are summed. With normalize=True, each clone-ct row contributes its - probability vector unweighted; with normalize=False, each row is - weighted by the clone's cell count at this covariate. - """ - phenotypes = list(adata.uns[K.PHENOTYPE_CATEGORIES]) - jd = joint_distribution(adata, covariate_label=covariate, clones=clones) - if jd is None or jd.empty: - return {} - - if not normalize: - meta = adata.uns[K.METADATA] - clone_col = meta["clone_col"] - cov_col = meta["covariate_col"] - counts_at_cov = ( - adata.obs.loc[adata.obs[cov_col] == covariate, clone_col] - .value_counts() - .to_dict() - ) - else: - counts_at_cov = None - - per_clone = {} - for clone_id, row in jd[phenotypes].iterrows(): - vec = row.to_numpy(dtype=float) - if counts_at_cov is not None: - vec = vec * counts_at_cov.get(clone_id, 0) - per_clone.setdefault(clone_id, np.zeros(len(phenotypes))) - per_clone[clone_id] = per_clone[clone_id] + vec - return per_clone - - -def plot_pheno_sankey( - adata, - *, - covariate_order, - clones=None, - phenotype_colors=None, - times=None, - time_rescale=1.0, - normalize=True, - ax=None, - figsize=(9, 5), - xlim=None, - ylim=None, - xlabel=None, - ylabel=None, - title=None, - show_legend=True, - fontsize=12, - return_axes=False, -): - """Sankey of phenotype-distribution flow across covariate values. - - Built directly from adata via joint_distribution; no CellRepertoire - construction. Flow geometry preserves the per-clone outer-product - semantics of the original implementation. - """ - phenotypes = list(adata.uns[K.PHENOTYPE_CATEGORIES]) - n_phen = len(phenotypes) - n_reps = len(covariate_order) - if n_reps == 0: - raise ValueError("covariate_order must contain at least one covariate") - - if phenotype_colors is None: - phenotype_colors = {p: tcri_colors[i % len(tcri_colors)] for i, p in enumerate(phenotypes)} - - if times is not None: - times = np.array(times) * time_rescale - dx = (max(times) - min(times)) / 500 if max(times) > min(times) else 0.2 - else: - times = list(range(n_reps)) - dx = 0.2 - - per_cov = [ - _phenotype_mass_per_clone(adata, cov, clones, normalize) - for cov in covariate_order - ] - - origin_nodes = [{} for _ in range(n_reps - 1)] - destination_nodes = [{} for _ in range(n_reps - 1)] - plot_nodes = [{} for _ in range(n_reps)] - - if n_reps == 1: - c_origin_node_vals = np.zeros(n_phen) - clone_keys = clones if clones is not None else list(per_cov[0].keys()) - for clone in clone_keys: - c_origin_node_vals += per_cov[0].get(clone, np.zeros(n_phen)) - origin_main_node_ys = np.array([0] + list(np.cumsum(c_origin_node_vals[:-1]))) - for j, p in enumerate(phenotypes): - plot_nodes[0][p] = SankeyNode( - times[0], origin_main_node_ys[j], c_origin_node_vals[j], - dx=dx, color=phenotype_colors[p], - ) - else: - for i in range(n_reps - 1): - origin = per_cov[i] - dest = per_cov[i + 1] - - c_origin_flow = np.zeros((n_phen, n_phen)) - c_dest_flow = np.zeros((n_phen, n_phen)) - c_origin_node_vals = np.zeros(n_phen) - c_dest_node_vals = np.zeros(n_phen) - - if clones is not None: - clone_keys = clones - else: - clone_keys = list(set(origin.keys()) | set(dest.keys())) - - for clone in clone_keys: - o = origin.get(clone, np.zeros(n_phen)) - d = dest.get(clone, np.zeros(n_phen)) - c_origin_node_vals += o - c_dest_node_vals += d - o_sum = o.sum() - d_sum = d.sum() - if o_sum > 0 and d_sum > 0: - c_origin_flow += np.outer(o, d / d_sum) - c_dest_flow += np.outer(o / o_sum, d) - - origin_main_node_ys = np.array([0] + list(np.cumsum(c_origin_node_vals[:-1]))) - dest_main_node_ys = np.array([0] + list(np.cumsum(c_dest_node_vals[:-1]))) - running_origin_node_ys = origin_main_node_ys.copy() - running_dest_node_ys = np.cumsum(c_dest_node_vals) - np.sum(c_dest_flow, axis=0) - - for j, op in enumerate(phenotypes): - plot_nodes[i][op] = SankeyNode( - times[i], origin_main_node_ys[j], c_origin_node_vals[j], - dx=dx, color=phenotype_colors[op], - ) - for k, dp in enumerate(phenotypes): - origin_nodes[i][(op, dp)] = SankeyNode( - times[i], running_origin_node_ys[j], c_origin_flow[j, k], - dx=dx, color=phenotype_colors[op], - ) - running_origin_node_ys[j] += c_origin_flow[j, k] - destination_nodes[i][(op, dp)] = SankeyNode( - times[i + 1], running_dest_node_ys[k], c_dest_flow[j, k], - dx=dx, color=phenotype_colors[dp], - ) - running_dest_node_ys[k] += c_dest_flow[j, k] - - if i == n_reps - 2: - for j, p in enumerate(phenotypes): - plot_nodes[i + 1][p] = SankeyNode( - times[i + 1], dest_main_node_ys[j], c_dest_node_vals[j], - dx=dx, color=phenotype_colors[p], - ) - - fig = None - if ax is None: - fig, ax = plt.subplots(figsize=figsize) - - for i in range(n_reps): - for p in phenotypes: - plot_nodes[i][p].plot(ax=ax) - if i < n_reps - 1: - for op in phenotypes: - for dp in phenotypes: - origin_nodes[i][(op, dp)].plot_node_connection( - destination_nodes[i][(op, dp)], ax=ax, alpha=0.5, - ) - - data_ymax = 0.0 - for i in range(n_reps): - for p in phenotypes: - node = plot_nodes[i][p] - if node.max_y > data_ymax: - data_ymax = node.max_y - - if xlim is not None: - ax.set_xlim(xlim) - if ylim is not None: - ax.set_ylim(ylim) - else: - ax.set_ylim([0, min(data_ymax * 1.02, 1) if normalize else data_ymax * 1.02]) - - if show_legend: - ax.legend( - [plot_nodes[0][p].patch for p in phenotypes], - list(phenotypes), - frameon=True, fontsize=fontsize, - bbox_to_anchor=(1.05, 1), loc='upper left', - ) - - if xlabel is not None: - ax.set_xlabel(xlabel, fontsize=fontsize) - if ylabel is not None: - ax.set_ylabel(ylabel, fontsize=fontsize) - else: - ax.set_ylabel('Fraction' if normalize else 'Cell Counts', fontsize=fontsize) - if title is not None: - ax.set_title(title, fontsize=fontsize) - - if fig is not None: - fig.tight_layout() - - if return_axes: - return fig, ax - return None - - -def phenotypic_flux( - adata, splitby, order, clones=None, normalize=True, - phenotype_colors=None, save=None, figsize=(6, 3), - show_legend=True, title=None, -): - """Sankey of phenotype flow across `order` values of `splitby`. - - Thin wrapper around plot_pheno_sankey. `splitby` is forwarded as the - x-axis label and must match adata.uns['tcri_metadata']['covariate_col']. - """ - times = list(range(len(order))) - if phenotype_colors is None: - phenotype_colors = dict(zip(adata.uns[K.PHENOTYPE_CATEGORIES], tcri_colors)) - fig, ax = plot_pheno_sankey( - adata, - covariate_order=order, - clones=clones, - phenotype_colors=phenotype_colors, - times=times, - xlim=(min(times), max(times)) if len(times) > 1 else None, - normalize=normalize, - xlabel=splitby, - title=title, - show_legend=show_legend, - figsize=figsize, - return_axes=True, - ) - ax.set_xticks(times) - ax.set_xticklabels(order) - if save is not None: - fig.savefig(save) - -def probability_distribution(adata, phenotype_order=None, color="#000000", rotation=90, splitby=None, order=None, figsize=(7,5), save=None): - columns = [] - if splitby != None: - ncols = len(set(adata.obs[splitby])) - else: - ncols = 1 - fig, ax = plt.subplots(1,ncols,figsize=figsize) - - if order == None: - order = list(sorted(adata.obs[splitby])) - for i, o in enumerate(order): - zdata = adata[adata.obs[splitby] == o] - pdist = probability_distribution(zdata) - sns.barplot(data=pdist,ax=ax[i],order=phenotype_order,color=color) - ax[i].set_xticklabels(ax[i].get_xticklabels(), rotation=rotation) - ax[i].set_title(o) - - fig.tight_layout() - if save != None: - fig.savefig(save) - - -def top_clone_umap(adata, reduction="umap", top_n=10, fg_alpha=0.9, fg_size=25, bg_size=0.1, bg_alpha=0.6, figsize=(12,5), return_df=False,save=None): - df = adata.obs - seq_column = adata.uns["tcri_clone_key"] - plt.figure(figsize = figsize) - clonotype_counts = collections.defaultdict(int) - for clonotype in df[seq_column]: - clonotype_counts[clonotype] += 1 - top_clonotypes = sorted(clonotype_counts.items(), key=operator.itemgetter(1),reverse=True) - top_clonotypes = [x[0] for x in top_clonotypes[:top_n]] - ax1 = plt.subplot(1,1,1) - x = [x[0] for x in adata.obsm["X_{}".format(reduction)]] - y = [x[1] for x in adata.obsm["X_{}".format(reduction)]] - sns.scatterplot(x=x,y=y, color="#75715E", alpha=bg_alpha, ax=ax1, s=bg_size, linewidth=0.0) - xonly = [] - yonly = [] - clonotype_labels = [] - size = [] - for clonotype,x1,y1 in zip(df[seq_column],x,y): - clonotype = str(clonotype) - if clonotype not in top_clonotypes or clonotype == "None" or clonotype == "nan": - continue - else: - xonly.append(x1) - yonly.append(y1) - size.append(1/clonotype_counts[clonotype]) - clonotype_labels.append(str(clonotype) + " {}".format(clonotype_counts[clonotype])) - dftop = pd.DataFrame.from_dict({"TCR Sequence":clonotype_labels,"Cells":size, "UMAP1":xonly,"UMAP2":yonly}) - colors = tcri_colors + tcri_colors + tcri_colors - order = [] - for c in set(clonotype_labels): - if c != "_Other": - order.append(c) - colors = colors[:len(set(clonotype_labels))] - sns.scatterplot(data=dftop, x="UMAP1", y="UMAP2", hue="TCR Sequence", hue_order=order, ax=ax1, alpha=fg_alpha,s=fg_size, linewidth=0.0,palette=colors) - ax1.set_xlabel('UMAP-1') - ax1.set_ylabel('UMAP-2') - ax1.xaxis.set_ticklabels([]) - ax1.yaxis.set_ticklabels([]) - ax1.xaxis.set_ticks([]) - ax1.yaxis.set_ticks([]) - ax1.set_title("Top 10 TCR Clone by Size") - h,l = ax1.get_legend_handles_labels() - ax1.legend(h[:top_n-1], l[:top_n-1], borderaxespad=2.,fontsize='9',bbox_to_anchor=(0, 1), loc='best') - plt.tight_layout() - if return_df: - return dftop - elif save != None: - plt.savefig(save) - - -def tcri_boxplot(adata, function, groupby=None,ylabel="", splitby=None,figsize=(8,4),s=20,order=None, palette=None): - if palette == None: - palette = tcri_colors - if groupby == None and splitby == None: - data = function(adata) - df = pd.DataFrame(list(data.items()), columns=['Phenotype', 'Clonotypic Entropy']) - df.replace([np.inf, -np.inf], np.nan, inplace=True) - df.dropna(inplace=True) - fig,ax=plt.subplots(1,1,figsize=figsize) - sns.stripplot(data=df,x="Phenotype",y=ylabel,s=s,ax=ax, palette=palette) - ax.set_ylim(0,max(df[ylabel] + 0.1)) - ax.set_ylabel(ylabel) - ax.set_title(ylabel) - fig.tight_layout() - elif groupby != None and splitby == None: - groups = adata.obs[groupby].unique() - dfs = [] - for group in groups: - data = function(adata[adata.obs[groupby]==group]) - df = pd.DataFrame(list(data.items()), columns=['Phenotype',ylabel]) - df[groupby] = group - dfs.append(df) - df = pd.concat(dfs) - 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() - 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) - ax.set_ylim(0,max(df[ylabel] + 0.1)) - ax.set_title(ylabel) - ax.set_ylabel(ylabel) - sns.move_legend(ax, "upper left", bbox_to_anchor=(1, 1)) - fig.tight_layout() - elif groupby != None and splitby != None: - groups = adata.obs[groupby].unique() - dfs = [] - for group in groups: - sub = adata[adata.obs[groupby]==group] - splits = sub.obs[splitby].unique() - for split in splits: - data = function(sub[sub.obs[splitby]==split]) - df = pd.DataFrame(list(data.items()), columns=['Phenotype', ylabel]) - df[groupby] = group - df[splitby] = split - dfs.append(df) - df = pd.concat(dfs) - df.replace([np.inf, -np.inf], np.nan, inplace=True) - 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() - 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) - sns.move_legend(ax, "upper left", bbox_to_anchor=(1, 1)) - ax.set_ylabel(ylabel) - fig.tight_layout() - else: - raise ValueError("'groupby' must be set to use 'splitby'.") - return ax - -def clonality(adata, groupby = None, splitby=None, s=10, order=None, figsize=(12,5), palette=None): - return tcri_boxplot(adata,clonality_tl, ylabel="Clonality", groupby=groupby, splitby=splitby, s=s, figsize=figsize, order=order, palette=palette) - -def clonotypic_entropy_by_phenotype( - adata, - *, - temperature = 1.0, - n_samples = 200, - combine_with_logits = True, - bayesian = True, - bayes_samples = 1_000, - palette = None, - group_colors = None, # {covariate: colour} - hue_order = None, - legend_fontsize = 6, - bbox_to_anchor = (1.15, 1.), - figsize = (6, 3), - rotation = 90, - save = None, - return_df = False, - progress = True, -): - """Box-and-dot plot of clonotypic entropy per phenotype / covariate.""" - - # ---- meta columns --------------------------------------------- # - meta = adata.uns[K.METADATA] - cov_col = meta["covariate_col"] - clone_col = meta["clone_col"] - phen_col = meta["phenotype_col"] - batch_col = meta["batch_col"] - - covariates = adata.obs[cov_col].astype("category").cat.categories.tolist() - phenotypes = adata.obs[phen_col].astype("category").cat.categories.tolist() - batches = adata.obs[batch_col].astype("category").cat.categories.tolist() - - if hue_order is None: - hue_order = covariates - palette=tcri_colors - cov2col = dict(zip(hue_order, palette)) - - # ---- compute entropy values ----------------------------------- # - records = [] - iterator = list(itertools.product(batches, hue_order)) - if progress: - iterator = tqdm.tqdm(iterator, desc="clonotypic entropy") - - for patient, cov_value in iterator: - mask_pc = ( - (adata.obs[batch_col] == patient) & - (adata.obs[cov_col] == cov_value) - ) - if not mask_pc.any(): - continue - clones = list(set(adata.obs.loc[mask_pc, clone_col])) - if not clones: - continue - ent_series = centropy( - adata, - covariate = cov_value, - point_estimate = True, - n_samples = n_samples, - temperature = temperature, - combine_with_logits= combine_with_logits, - _clones = clones, - ) - for phen in phenotypes: - if not (mask_pc & (adata.obs[phen_col] == phen)).any(): - continue - if phen not in ent_series.index: - continue - records.append( - {cov_col: cov_value, batch_col: patient, - phen_col: phen, "entropy": float(ent_series[phen])} - ) - - df = pd.DataFrame.from_records(records) - if df.empty: - _warn("no data – nothing to plot") - return None - - # ---- x-position jitter per hue -------------------------------- # - x_levels = df[phen_col].unique().tolist() - n_hue = len(hue_order) - w_tot = .8 - step = w_tot / n_hue - offsets = np.linspace(-w_tot/2 + step/2, w_tot/2 - step/2, n_hue) - cov2off = dict(zip(hue_order, offsets)) - df["_x"] = df[phen_col].map(lambda x: x_levels.index(x)) + df[cov_col].map(cov2off) - - # ---- plotting ------------------------------------------------- # - fig, ax = plt.subplots(figsize=figsize) - - sns.boxplot( - data = df, - x = phen_col, - y = "entropy", - hue = cov_col, - palette = palette, - hue_order = hue_order, - width = w_tot, - fliersize = 0, - ax = ax, - zorder = 1, - ) - ax.legend_.remove() - - # ---- Bayesian / MW stats inside each phenotype ---------------- # - y_max = df["entropy"].max(); y_min = df["entropy"].min() - h_bar = (y_max - y_min) * .05 - v_gap = h_bar * 1.25 - - for i, ph in enumerate(x_levels): - for j, (g1, g2) in enumerate(itertools.combinations(hue_order, 2)): - d1 = df[(df[phen_col]==ph)&(df[cov_col]==g1)]["entropy"].values - d2 = df[(df[phen_col]==ph)&(df[cov_col]==g2)]["entropy"].values - if len(d1)==0 or len(d2)==0: - continue - - if bayesian: - samp1 = np.random.choice(d1, (bayes_samples, len(d1)), replace=True).mean(1) - samp2 = np.random.choice(d2, (bayes_samples, len(d2)), replace=True).mean(1) - delta = samp2 - samp1 - lbl = f"Δ={delta.mean():.2g}" - else: - from scipy.stats import mannwhitneyu - _, p = mannwhitneyu(d1,d2,alternative="two-sided") - lbl = "ns" if p>=.05 else ("*" if p<.05 else "**" if p<.01 else "***" if p<.001 else "****") - - x1 = i+cov2off[g1]; x2 = i+cov2off[g2]; y = y_max + j*v_gap - ax.plot([x1,x1,x2,x2], [y,y+h_bar,y+h_bar,y], color="k", lw=1.4) - ax.text((x1+x2)/2, y+h_bar, lbl, ha="center", va="bottom", fontsize=legend_fontsize) - - # patient dots - pt_pal = sns.color_palette("tab20b", len(batches)) - for pt,c in zip(batches, pt_pal): - sub = df[df[batch_col]==pt] - ax.scatter(sub["_x"], sub["entropy"], color=c, s=80, alpha=.85, zorder=3, label=pt) - ax.legend(title="Patient", fontsize=legend_fontsize, loc="upper right", bbox_to_anchor=bbox_to_anchor) - - ax.set_xticklabels(x_levels, rotation=rotation) - ax.set_xlabel("Phenotype"); ax.set_ylabel("Clonotypic entropy") - ax.set_title("Clonotypic entropy per phenotype / covariate") - fig.tight_layout() - - if save: - fig.savefig(save, dpi=150); _ok(f"figure saved → {save}") - else: - plt.show() - - if return_df: - return df - -def plot_phenotype_probabilities(adata, phenotype_prob_slot="X_tcri_phenotypes", add_outline=False, save=None,ncols=2,cmap="magma"): - phenotypes = adata.uns[K.PHENOTYPE_CATEGORIES] - prob_labels = [] - adata = adata.copy() - for y,x in zip(phenotypes, adata.obsm[phenotype_prob_slot].T): - adata.obs['{}_probability'.format(y)] = x - prob_labels.append('{}_probability'.format(y)) - sc.pl.umap(adata,color=prob_labels, cmap=cmap, s=30, ncols=ncols, show=False, add_outline=add_outline) - if save != None: - plt.savefig(save) - -def clone_size_umap(adata, reduction="umap",figsize=(10,8),size=1,alpha=0.7,palette="coolwarm",save=None): - clone_size(adata) - df = adata.obs - reduction="umap" - sizes = np.log10(adata.obs[K.CLONE_SIZE].to_numpy()) - df["UMAP1"] = [x[0] for x in adata.obsm["X_{}".format(reduction)]] - df["UMAP2"] = [x[1] for x in adata.obsm["X_{}".format(reduction)]] - df["log(Clone Size)"] = sizes - fig,ax=plt.subplots(1,1,figsize=figsize) - sns.scatterplot(data=df,x="UMAP1", y="UMAP2", hue="log(Clone Size)",s=size,palette=palette, ax=ax, alpha=alpha,linewidth=0.) - ax.set_xlabel('UMAP-1') - ax.set_ylabel('UMAP-2') - ax.xaxis.set_ticklabels([]) - ax.yaxis.set_ticklabels([]) - ax.xaxis.set_ticks([]) - ax.yaxis.set_ticks([]) - fig.tight_layout() - if save != None: - fig.savefig(save) - return ax - - -def ridge_delta_entropy( - df_delta : pd.DataFrame, - *, - splitby : str = "complete_response", - order_group : list = None, - order_phen : list = None, - palette : dict = None, - bw_adjust : float = 0.8, - jitter : float = .15, - density_scale : float = .9, - significance : bool = True, - sig_test : str = "mannwhitney", # or "bayesian" - bayes_iters : int = 5_000, - bracket_pad : float = 0.15, # ↑ lift for bracket - star_size : int = 16, # ★ marker size - figsize : tuple = (10, 6), - ax = None -): - """ - Ridge plot of Δ-entropy posteriors per phenotype. - - For each phenotype the *first two* groups in ``order_group`` - are compared and annotated:: - - ┌──────────┐ - CR │ │ NR - ★ p-value / stars - - Bracket anchors are placed at the group means. - """ - # ───── tidy → long Δ samples ───────────────────────────────── - long = (df_delta - .explode("delta_samples") - .rename(columns={"delta_samples": "delta"})) - - if order_group is None: - order_group = sorted(long[splitby].unique()) - if order_phen is None: - order_phen = sorted(long["phenotype"].unique()) - - # ───── palette ─────────────────────────────────────────────── - if palette is None: - tab = cm.tab10.colors - palette = {g: tab[i % 10] for i, g in enumerate(order_group)} - - # ───── basic geometry ─────────────────────────────────────── - phen2y = {p: i for i, p in enumerate(order_phen)} - n_groups = len(order_group) - - x_all = long["delta"].astype(float).to_numpy() - x_min,x_max = np.percentile(x_all, [0.5, 99.5]) - pad = 0.06*(x_max-x_min) - xs = np.linspace(x_min-pad, x_max+pad, 500) - - if ax is None: - fig, ax = plt.subplots(figsize=figsize) - else: - fig = ax.figure - - # ───── draw ridges ─────────────────────────────────────────── - for ph in order_phen: - base_y = phen2y[ph] - - # plot each group’s ridge - means = {} # keep per-group mean (for bracket) - for g_idx, g in enumerate(order_group): - data = long[(long["phenotype"]==ph) & (long[splitby]==g)]["delta"].astype(float) - if data.empty: - continue - kde = st.gaussian_kde(data, bw_method=bw_adjust) - ys = kde(xs) - ys = ys/ys.max()*density_scale - - shift = (g_idx - (n_groups-1)/2)*2*jitter - y_line = base_y + shift - - ax.fill_between(xs, y_line, y_line+ys, - color=palette[g], alpha=.85, lw=0) - ax.plot(xs, y_line+ys, color=palette[g], lw=.8) - - means[g] = data.mean() - - # ── bracket + significance for first two groups ───────── - if significance and len(order_group) >= 2 and all(k in means for k in order_group[:2]): - g1, g2 = order_group[:2] - d1 = long[(long["phenotype"]==ph)&(long[splitby]==g1)]["delta"].to_numpy(float) - d2 = long[(long["phenotype"]==ph)&(long[splitby]==g2)]["delta"].to_numpy(float) - - # statistical label - if sig_test == "mannwhitney": - _, pval = st.mannwhitneyu(d1,d2,alternative="two-sided") - label = ("ns" if pval>=.05 else - "*" if pval<.05 else - "**" if pval<.01 else - "***" if pval<.001 else "****") - else: # Bayesian Δ > 0 - idx = np.random.randint(0, min(len(d1),len(d2)), size=bayes_iters) - p_gt = ((d2[idx]-d1[idx]) > 0).mean() - label = f"P={p_gt:.2f}" - - # bracket coordinates - x1, x2 = means[g1], means[g2] - y_brk = base_y + density_scale + bracket_pad - ax.plot([x1, x1, x2, x2], [y_brk, y_brk+bracket_pad, - y_brk+bracket_pad, y_brk], - color="k", lw=1.2) - - # star / label on top - x_star = (x1+x2)/2 - ax.text(x_star, y_brk+bracket_pad*1.05, label, - ha="center", va="bottom", fontsize=star_size, - color="k") - - # ───── axis cosmetics ─────────────────────────────────────── - ax.set_yticks(list(phen2y.values()), list(order_phen)) - ax.axvline(0, color="k", ls="--", lw=.8) - ax.set_xlabel("Δ clonotypic entropy (post – pre)") - ax.set_xlim(x_min-pad, x_max+pad) - ax.set_title("Δ Clonotypic Entropy") - - # legend - handles = [plt.Line2D([0],[0],lw=8,color=palette[g],label=g) for g in order_group] - ax.legend(handles=handles, title=splitby, frameon=False) - - plt.tight_layout() - return fig, ax - - - -def phenotypic_entropy(adata, splitby=None, temperature=1, n_samples=0, normalized=True, palette=None, save=None, legend_fontsize=6, bbox_to_anchor=(1.15,1.), figsize=(8,4), rotation=90): - if palette == None: - palette=tcri_colors - cov_col = adata.uns[K.METADATA]["covariate_col"] - clone_col = adata.uns[K.METADATA]["clone_col"] - phenotype_col = adata.uns[K.METADATA]["phenotype_col"] - batch_col = adata.uns[K.METADATA]["batch_col"] - - covs = adata.obs[cov_col].astype("category").cat.categories.tolist() - clones = adata.obs[clone_col].astype("category").cat.categories.tolist() - phenotypes = adata.obs[phenotype_col].astype("category").cat.categories.tolist() - batches = adata.obs[batch_col].astype("category").cat.categories.tolist() - - mi = [] - ps = [] - ts = [] - rs = [] - cl = [] - phs = [] - for p in tqdm.tqdm(batches): - sub = adata[adata.obs[batch_col] == p].copy() - for t in covs: - subt = sub[sub.obs[cov_col] == t] - vclones = list(set(subt.obs[clone_col])) - if len(vclones) == 0: continue - for ph in phenotypes: - if splitby == None: - vclones = list(set(subt.obs[clone_col])) - mi.append(pentropy(subt,t,ph, temperature=temperature, clones=vclones,n_samples=n_samples, normalized=normalized)) - ps.append(p) - ts.append(t) - phs.append(ph) - else: - for s in set(subt.obs[splitby]): - subts = subt[subt.obs[cov_col] == t] - vclones = list(set(subts.obs[clone_col])) - mi.append(pentropy(subts,t,ph, temperature=temperature, clones=vclones,n_samples=n_samples, normalized=normalized)) - ps.append(p) - ts.append(t) - cl.append(s) - phs.append(ph) - fig,ax = plt.subplots(1,1,figsize=figsize) - if splitby != None: - df = pd.DataFrame.from_dict({cov_col: ts, batch_col:ps, "Phenotypic Entropy":mi, splitby:cl, clone_col:phs}) - sns.boxplot(data=df,x=splitby,y="Phenotypic Entropy",hue=cov_col,palette=palette) - palette_black = {level: "black" for level in df[cov_col].unique()} - sns.stripplot(data=df,x=splitby,y="Phenotypic Entropy",hue=cov_col, dodge=True, palette=palette_black) - else: - df = pd.DataFrame.from_dict({cov_col: ts, batch_col:ps, "Phenotypic Entropy":mi,clone_col:phs}) - sns.boxplot(data=df,x=cov_col,y="Phenotypic Entropy",color="#999999") - sns.stripplot(data=df,x=cov_col,y="Phenotypic Entropy",palette=palette,dodge=False) - plt.xticks(rotation=rotation) - leg = ax.legend(loc='upper right', bbox_to_anchor=bbox_to_anchor, fontsize=legend_fontsize) - fig.tight_layout() - if save: - fig.savefig(save) - -def set_color_palette(adata, columns): - i = 0 - main_color_map = dict() - adata = adata.copy() - colors = tcri_colors.copy() + tcri_colors.copy() + tcri_colors.copy() - for x in columns: - ct = [] - for i, val in enumerate(set(adata.obs[x].tolist())): - c = colors.pop(i) - ct.append(c) - main_color_map[val] = c - adata.uns["{}_colors".format(x)] = ct - return main_color_map - -def flux(adata, key, order, groupby, paint_dict=None, method="probabilistic", paint=None, distance_metric="l1", figsize=(12,5), paint_order=None, palette=None): - dfs = [] - if paint != None: - palette = [] - legend_handles = [] - paint_categories = adata.obs[paint].unique() - if paint_dict != None: - pcolors = paint_dict - else: - pcolors = dict(zip(paint_categories, tcri_colors)) - for category in paint_categories: - handle = mpatches.Patch(color=pcolors[category], label=category) - legend_handles.append(handle) - else: - if palette == None: - if "{}_colors".format(paint) in adata.uns: - palette = adata.uns["{}_colors".format(paint)] - else: - palette = tcri_colors - for x in tqdm.tqdm(list(set(adata.obs[groupby]))): - sdata = adata[adata.obs[groupby]==x] - hue_order = [] - for i in range(len(order)-1): - l1_distances = flux_tl(sdata,key=key,from_this=order[i],to_that=order[i+1],distance_metric=distance_metric) - df = pd.DataFrame(list(l1_distances.items()), columns=['Clone', distance_metric]) - df[groupby] = x - if paint!=None: - pcat = sdata.obs[paint].unique().tolist()[0] - palette.append(pcolors[pcat]) - df["Comparison"] = pcat - dfs.append(df) - print(palette) - df = pd.concat(dfs) - df.replace([np.inf, -np.inf], np.nan, inplace=True) - df.dropna(inplace=True) - order = df.groupby(groupby).median(distance_metric).sort_values(distance_metric).index.tolist() - fig,ax=plt.subplots(1,1,figsize=figsize) - sns.boxplot(data=df,x=groupby,y=distance_metric,hue="Comparison",order=order,palette=pcolors,ax=ax) - fig.tight_layout() - return ax - -def mutual_information(adata, splitby=None, temperature=1.0, n_samples=0, normalized=True, palette=None, save=None, legend_fontsize=6, bbox_to_anchor=(1.15,1.), figsize=(8,4), rotation=90, weighted=True, return_plot=True): - """ - Compute and plot mutual information between clonotypes and phenotypes. - - This function calculates mutual information between TCR clonotypes and cell phenotypes, - which quantifies how much information one variable provides about the other. The function - can optionally split the calculation by a specified categorical variable (e.g., timepoint, - condition). Results are displayed as a box plot with individual data points. - - Parameters - ---------- - adata : AnnData - AnnData object containing the data with TCR and phenotype information - splitby : str, optional - Column name to split the data by. If None, uses the covariate column stored in adata.uns['tcri_metadata'] - temperature : float, default=1.0 - Temperature parameter for softening/sharpening distributions in the joint distribution calculation - n_samples : int, default=0 - Number of samples to use for Monte Carlo estimation of mutual information - normalized : bool, default=True - Whether to normalize the mutual information values to [0,1] range - palette : list, optional - Color palette for the plot. If None, uses tcri_colors - save : str, optional - Path to save the plot figure - legend_fontsize : int, default=6 - Font size for the plot legend - bbox_to_anchor : tuple, default=(1.15,1.) - Position of the legend box - figsize : tuple, default=(8,4) - Size of the figure in inches (width, height) - rotation : int, default=90 - Rotation angle for x-axis labels - weighted : bool, default=True - Whether to weight the mutual information by clone size - return_plot : bool, default=True - Whether to return the plot axis. If False, returns only the DataFrame with MI values - - Returns - ------- - Union[matplotlib.axes.Axes, pd.DataFrame] - If return_plot is True, returns the plot axis. Otherwise returns a DataFrame with MI values. - - Examples - -------- - >>> import tcri - >>> # Calculate and plot mutual information - >>> ax = tcri.pl.mutual_information(adata, splitby="timepoint", temperature=1.0) - >>> - >>> # Get the mutual information values as a DataFrame without plotting - >>> mi_df = tcri.pl.mutual_information(adata, splitby="condition", return_plot=False) - """ - if palette is None: - palette = tcri_colors - - # Retrieve metadata from adata - cov_col = adata.uns[K.METADATA]["covariate_col"] - clone_col = adata.uns[K.METADATA]["clone_col"] - phenotype_col = adata.uns[K.METADATA]["phenotype_col"] - batch_col = adata.uns[K.METADATA]["batch_col"] - - covs = adata.obs[cov_col].astype("category").cat.categories.tolist() - batches = adata.obs[batch_col].astype("category").cat.categories.tolist() - - mi_vals = [] - ps = [] - ts = [] - cl = [] - - for p in tqdm.tqdm(batches, desc="Computing mutual information"): - sub = adata[adata.obs[batch_col] == p].copy() - for t in covs: - subt = sub[sub.obs[cov_col] == t] - # figure out which clones are actually present - vclones = list(set(subt.obs[clone_col])) - - if splitby is None: - # Compute MI for all clones in cov t, batch p - val = mutual_information_tl( - subt, t, - temperature=temperature, - clones=vclones, - n_samples=n_samples, - weighted=weighted - ) - mi_vals.append(val) - ps.append(p) - ts.append(t) - - else: - # If you want to split further by some obs column - for s in sorted(subt.obs[splitby].unique()): - subts = subt[subt.obs[splitby] == s] - vclones2 = list(set(subts.obs[clone_col])) - val = mutual_information_tl( - subts, t, - temperature=temperature, - clones=vclones2, - n_samples=n_samples, - weighted=weighted - ) - mi_vals.append(val) - ps.append(p) - ts.append(t) - cl.append(s) - - # Build a DataFrame for plotting - if splitby is None: - df = pd.DataFrame({ - cov_col: ts, - batch_col: ps, - "Mutual Information": mi_vals - }) - else: - df = pd.DataFrame({ - cov_col: ts, - batch_col: ps, - "Mutual Information": mi_vals, - splitby: cl - }) - - if not return_plot: - return df - - # Create the plot - fig, ax = plt.subplots(1,1, figsize=figsize) - if splitby is None: - sns.boxplot(data=df, x=cov_col, y="Mutual Information", color="#999999", ax=ax) - sns.stripplot(data=df, x=cov_col, y="Mutual Information", palette=palette, dodge=False, ax=ax) - else: - sns.boxplot(data=df, x=splitby, y="Mutual Information", hue=cov_col, color="#999999", ax=ax) - sns.stripplot(data=df, x=splitby, y="Mutual Information", hue=cov_col, palette=palette, dodge=True, ax=ax) - ax.legend(loc='upper right', bbox_to_anchor=bbox_to_anchor, fontsize=legend_fontsize) - - plt.xticks(rotation=rotation) - plt.title("Mutual Information" + (" (Weighted)" if weighted else "")) - fig.tight_layout() - - if save: - plt.savefig(save, dpi=150) - plt.show() - - return ax - - -def bayesian_mutual_information( - adata, - *, - group1, - group2, - splitby, - n_samples = 200, - temperature = 1.0, - normalised = True, - normalise_mode = "average", - weighted = False, - posterior = True, - combine_with_logits=True, - seed = 42, - palette = None, -): - np.random.seed(seed) - - meta = adata.uns[K.METADATA] - cov_col = meta["covariate_col"] - clone_col = meta["clone_col"] - - groups = sorted(adata.obs[splitby].dropna().unique().tolist()) - if palette == None: - palette = dict() - for i, g in enumerate(groups): - palette[g] = tcri_colors[i] - print(f"{BOLD}{MAG}──────── MI summary ({group1} → {group2}) ────────{RESET}") - _info("split column", splitby) - _info("# groups", len(groups)) - _info("Δ samples / group", n_samples) - _info("weighted", weighted) - print(f"{MAG}────────────────────────────────────────────────────{RESET}") - - results = {} - - # ---------- iterate over strata -------------------------------- - for g in groups: - mask_g = adata.obs[splitby] == g - clones = adata.obs.loc[mask_g, clone_col].unique().tolist() - - mi_pre = []; mi_post = [] - bar = tqdm(range(n_samples), desc=f"Δ-MI samples ({g})") - for _ in bar: - mi_pre.append( mutual_information_tl( - adata, group1, temperature=temperature, n_samples=1, - clones=clones, weighted=weighted, normalised=normalised, - normalise_mode=normalise_mode, posterior=posterior, - combine_with_logits=combine_with_logits, verbose=False)) - mi_post.append( mutual_information_tl( - adata, group2, temperature=temperature, n_samples=1, - clones=clones, weighted=weighted, normalised=normalised, - normalise_mode=normalise_mode, posterior=posterior, - combine_with_logits=combine_with_logits, verbose=False)) - - mi_pre = np.array(mi_pre) - mi_post = np.array(mi_post) - delta = mi_post - mi_pre - - d_mean, d_std = delta.mean(), delta.std() - hdi_low, hdi_hi = np.percentile(delta, [2.5,97.5]) - p_gt = (delta>0).mean(); p_lt = 1-p_gt - cohens_d = d_mean/d_std if d_std>0 else 0.0 - - print(f"\n{BOLD}{g}{RESET} ΔMI = {d_mean:.4f} ± {d_std:.4f} " - f"95 % HDI [{hdi_low:.4f}, {hdi_hi:.4f}] " - f"P(>0)={p_gt:.3f}") - - results[g] = dict(delta_samples=delta, mi_pre_samples=mi_pre, - mi_post_samples=mi_post, delta_mean=d_mean, - delta_std=d_std, cohens_d=cohens_d, - p_greater=p_gt, p_less=p_lt, - hdi=(hdi_low, hdi_hi)) - - fig, ax = plt.subplots(1, 3, figsize=(15, 4), - gridspec_kw=dict(width_ratios=[2, 2, 1]), - constrained_layout=True) - - # A ─── Δ-MI KDEs - for i, g in enumerate(groups): - sns.kdeplot(results[g]["delta_samples"], - fill=True, ax=ax[0], - palette=[palette[g]], - alpha=.9, linewidth=1.2, label=g) - ax[0].axvline(0, color="k", ls="--") - ax[0].set(title="Δ MI (post – pre)", - xlabel="Δ normalised MI", ylabel="density") - ax[0].legend(title=splitby) - - # B ─── pre / post KDEs - for i, g in enumerate(groups): - sns.kdeplot(results[g]["mi_pre_samples"], - ax=ax[1], palette=[palette[g]], - ls="-", label=f"{g} – {group1}") - sns.kdeplot(results[g]["mi_post_samples"], - ax=ax[1], palette=[palette[g]], - ls="--", label=f"{g} – {group2}") - ax[1].set(title="MI posterior per condition", - xlabel="normalised MI") - ax[1].legend() - - # C ─── bar summary of Δ - means = [results[g]["delta_mean"] for g in groups] - errs = [results[g]["delta_std"] for g in groups] - ax[2].bar(groups, means, yerr=errs, capsize=5, - color=[palette[g] for g in groups]) - ax[2].axhline(0, color="k", ls="--") - ax[2].set(title="Δ MI summary", ylabel="Δ normalised MI") - - fig.suptitle("Bayesian MI Analysis of clonotype ⇄ phenotype coupling", - fontsize=14, weight="bold") - return results - - - -def polar_plot(adata, phenotypes=None, statistic="distribution", method="joint_distribution", splitby=None, color_dict=None, temperature=1.0): - """ - Create a polar plot showing phenotype distributions or entropies. - - This function creates a radar/polar chart that visualizes either the distribution of phenotypes - or entropy values across different conditions. It's useful for comparing phenotype proportions - or entropy patterns across experimental groups. - - Parameters - ---------- - adata : AnnData - AnnData object containing the data with TCR and phenotype information - phenotypes : list, optional - List of phenotype names to include in the plot. If None, uses all phenotypes - defined in adata.uns['tcri_metadata']["phenotype_col"] - statistic : str, default="distribution" - Type of statistic to plot, one of "distribution" or "entropy" - method : str, default="joint_distribution" - Method to compute phenotype distributions, either "joint_distribution" (model-based) - or "empirical" (raw cell counts) - splitby : str, optional - Column name to split the data by. If None, uses the covariate column stored - in adata.uns['tcri_metadata']["covariate_col"] - color_dict : dict, optional - Dictionary mapping split categories to colors - temperature : float, default=1.0 - Temperature parameter for softening/sharpening distributions - - Returns - ------- - matplotlib.axes.Axes - The polar plot axis - - Examples - -------- - >>> import tcri - >>> # Basic phenotype distribution polar plot - >>> ax = tcri.pl.polar_plot(adata, statistic="distribution") - >>> - >>> # Entropy polar plot with custom colors - >>> color_dict = {"Day0": "#FF5733", "Day7": "#33FF57", "Day14": "#3357FF"} - >>> ax = tcri.pl.polar_plot(adata, statistic="entropy", color_dict=color_dict) - """ - if phenotypes is None: - phenotypes = adata.uns[K.METADATA]["phenotype_col"] - - if splitby is None: - splitby = adata.uns[K.METADATA]["covariate_col"] - - # Get unique splits - splits = adata.obs[splitby].unique() - - # Create figure and axis - fig, ax = plt.subplots(figsize=(8, 8), subplot_kw={'projection': 'polar'}) - - # Calculate angles for each phenotype - angles = np.linspace(0, 2*np.pi, len(phenotypes), endpoint=False) - - # Plot for each split - for i, split in enumerate(splits): - if statistic == "distribution": - if method == "joint_distribution": - jd = joint_distribution(adata, split, temperature=temperature) - values = jd.mean().values - else: - subset = adata[adata.obs[splitby] == split] - values = np.zeros(len(phenotypes)) - for j, pheno in enumerate(phenotypes): - mask = subset.obs[adata.uns[K.METADATA]["phenotype_col"]] == pheno - values[j] = np.sum(mask) / len(subset) - else: # entropy - values = np.zeros(len(phenotypes)) - for j, pheno in enumerate(phenotypes): - values[j] = clonotypic_entropy(adata, split, pheno, temperature=temperature) - - # Normalize values - values = values / np.sum(values) - - # Plot values - color = color_dict[split] if color_dict else None - ax.plot(angles, values, 'o-', linewidth=2, label=split, color=color) - ax.fill(angles, values, alpha=0.25, color=color) - - # Set the labels - ax.set_xticks(angles) - ax.set_xticklabels(phenotypes) - - # Add legend - ax.legend(loc='upper right', bbox_to_anchor=(0.1, 0.1)) - - plt.tight_layout() - return ax \ No newline at end of file diff --git a/tcri/plotting/_sankey.py b/tcri/plotting/_sankey.py index 33f317f..6c6d03c 100644 --- a/tcri/plotting/_sankey.py +++ b/tcri/plotting/_sankey.py @@ -25,9 +25,6 @@ def plot(self, ax): ax.add_patch(self.patch) - def hex_to_rgb(self, hex_color): - hex_color = hex_color.lstrip('#') - return tuple(int(hex_color[i:i+2], 16) / 255.0 for i in (0, 2, 4)) def plot_node_connection(self, destination_node, ax, **kwargs): num_segments = 500 diff --git a/tcri/preprocessing/_preprocessing.py b/tcri/preprocessing/_preprocessing.py index 9cfc0eb..d48c8a5 100644 --- a/tcri/preprocessing/_preprocessing.py +++ b/tcri/preprocessing/_preprocessing.py @@ -17,7 +17,6 @@ import torch import torch.nn.functional as F from typing import Optional -from sklearn.metrics.pairwise import cosine_similarity import umap import numpy as np, pandas as pd, torch, umap from tqdm.auto import tqdm @@ -50,27 +49,6 @@ from .._console import _ok, _info, _warn, _fin -def _ascii_hist(samples, bins=25, width=40) -> str: - hist, edges = np.histogram(samples, bins=bins) - top = hist.max() - lines=[] - for h,e0,e1 in zip(hist, edges[:-1], edges[1:]): - bar = "█"*int(width*h/top) if top else "" - lines.append(f"{e0:7.3f}-{e1:7.3f} | {bar}") - 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) clone_counts = adata.obs["trb_candidate"].value_counts() @@ -83,461 +61,31 @@ def collapse_singleton(row): adata.obs[target_col] = adata.obs.apply(collapse_singleton, axis=1) -def classify_phenotypes(adata, phenotype_prob_slot="X_tcri_phenotypes", phenotype_assignment_obs=K.PHENOTYPE): - print("\t...classifying phenotypes...\n") - phenotype_col = adata.uns[K.METADATA]["phenotype_col"] - ct_array = adata.uns[K.CT_ARRAY] - unique_cts = np.unique(ct_array) - phenotype_probs_posterior = adata.uns[K.P_CT] - phenotypes = adata.uns[K.PHENOTYPE_CATEGORIES] - latent_z = adata.obsm[K.X_TCRI] - - # Pre-compute phenotype archetype embeddings - archetype_matrix = np.vstack([ - latent_z[adata.obs[phenotype_col].values == phenotype].mean(axis=0) - for phenotype in phenotypes - ]) - - all_probs = np.zeros((adata.n_obs, len(phenotypes))) - - # Iterate over each unique ct - for ct in unique_cts: - ct_indices = np.where(ct_array == ct)[0] - ct_embeddings = latent_z[ct_indices] - - # Cosine similarity (cells x phenotypes) - similarity = cosine_similarity(ct_embeddings, archetype_matrix) - similarity = (similarity + 1) / 2 # Normalize cosine similarity to [0,1] - - # Adjust similarity scores by posterior phenotype probabilities - adjusted_scores = similarity * phenotype_probs_posterior[ct] - - # Normalize to probabilities per cell - probs_normalized = adjusted_scores / adjusted_scores.sum(axis=1, keepdims=True) - all_probs[ct_indices] = probs_normalized - - # Store the normalized probabilities in AnnData - adata.obsm[phenotype_prob_slot] = all_probs - - # Assign phenotype with highest probability - assignments = all_probs.argmax(axis=1) - adata.obs[phenotype_assignment_obs] = pd.Categorical.from_codes( - assignments, categories=phenotypes - ) # ------------ 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( - adata, covariate_label, *, temperature=1.0, clones=None, - weighted=False, combine_with_logits=True, precision=3, silent=False): - - meta = adata.uns[K.METADATA]; cov_col = meta["covariate_col"] - clone_col = meta["clone_col"]; ph_cats = adata.uns[K.PHENOTYPE_CATEGORIES] - cov_idx = adata.uns[K.COVARIATE_CATEGORIES].index(covariate_label) - - ct_per_cell = adata.uns[K.CT_ARRAY] - cov_per_cell = adata.uns[K.COV_ARRAY] - clone_labels = adata.obs[clone_col].values - - # Guard against filtered AnnData (view or subset copy). The per-cell arrays in - # .uns are stored in the original full-cell space and are NOT subset when adata - # is sliced, whereas .obs/.obsm ARE subset. Indexing one with positions derived - # from the other then silently misaligns cells (Notion #4). Fail loudly instead - # of returning wrong numbers. - n_obs = adata.n_obs - if len(ct_per_cell) != n_obs or len(cov_per_cell) != n_obs: - raise ValueError( - "joint_distribution_posterior received an AnnData whose per-cell " - f"registration arrays (len {len(ct_per_cell)}) do not match adata.n_obs " - 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 " - "filtered AnnData, or pass the full object and filter with `clones=`." - ) - idx_cov = np.nonzero(cov_per_cell == cov_idx)[0] - if clones is not None: - idx_cov = idx_cov[np.isin(clone_labels[idx_cov], clones)] - _ok(f"selected {len(idx_cov):,} cells", silent) - p_ct_mean = torch.tensor(adata.uns[K.P_CT]) - local_scale = adata.uns.get(K.LOCAL_SCALE, 1.0) - bad = ~torch.isfinite(p_ct_mean) - if bad.any(): - n_phen = p_ct_mean.shape[1] - p_ct_mean = torch.where(bad, torch.ones_like(p_ct_mean) / n_phen, p_ct_mean) - p_ct_sample = Dirichlet(local_scale * p_ct_mean + 1e-8).sample().numpy() - _ok("sampled one draw from posterior p_ct", silent) - if combine_with_logits: - if K.X_LOGITS not in adata.obsm: - raise RuntimeError("X_tcri_logits missing in adata.") - logits = adata.obsm[K.X_LOGITS][idx_cov] - ct_idx_sel = ct_per_cell[idx_cov] - log_prior = np.log(p_ct_sample[ct_idx_sel] + 1e-8) - probs_cell = softmax((logits + log_prior)/temperature, axis=1) - _ok("combined logits with sampled prior", silent) - else: - probs_cell = p_ct_sample[ct_per_cell[idx_cov]] - _ok("using sampled p_ct only", silent) - df = (pd.DataFrame(probs_cell, columns=ph_cats, - index=clone_labels[idx_cov]) - .groupby(level=0).sum().astype(float)) - if not weighted: - df = df.div(df.sum(1), axis=0).fillna(0.0) - if clones is not None: - df = df.reindex(clones).fillna(0.0) - _info("resulting DataFrame", df.shape, silent); _fin(silent) - return df.round(precision) -def remove_meaningless_genes(adata, include_mt=True, include_rp=True, include_mtrn=True, include_hsp=True, include_tcr=True): - genes = [x for x in adata.var.index.tolist() if "RIK" not in x.upper()] - genes = [x for x in genes if "GM" not in x] - genes = [x for x in genes if "-" not in x or "HLA" in x] - genes = [x for x in genes if "." not in x or "HLA" in x] - genes = [x for x in genes if "LINC" not in x.upper()] - if include_mtrn: - genes = [x for x in adata.var.index.tolist() if "MTRN" not in x] - if include_hsp: - genes = [x for x in adata.var.index.tolist() if "HSP" not in x] - if include_mt: - genes = [x for x in genes if "MT-" not in x.upper()] - if include_rp: - genes = [x for x in genes if "RP" not in x.upper()] - if include_tcr: - genes = [x for x in genes if "TRAV" not in x] - genes = [x for x in genes if "TRAJ" not in x] - genes = [x for x in genes if "TRAD" not in x] - genes = [x for x in genes if "TRBV" not in x] - genes = [x for x in genes if "TRBJ" not in x] - genes = [x for x in genes if "TRBD" not in x] - - genes = [x for x in genes if "TRGV" not in x] - genes = [x for x in genes if "TRGJ" not in x] - genes = [x for x in genes if "TRGD" not in x] - - genes = [x for x in genes if "TRDV" not in x] - genes = [x for x in genes if "TRDJ" not in x] - genes = [x for x in genes if "TRDD" not in x] - adata = adata[:,genes] - return adata.copy() - -def joint_distribution( - adata, - covariate_label: str, - temperature: float = 1.0, - n_samples: int = 0, - clones=None, - weighted: bool = False, -) -> pd.DataFrame: - - p_ct = torch.tensor(adata.uns[K.P_CT]) - ct_to_cov = torch.tensor(adata.uns[K.CT_TO_COV]) - ct_to_c = torch.tensor(adata.uns[K.CT_TO_C]) - - covariate_categories = adata.uns[K.COVARIATE_CATEGORIES] - phenotype_categories = adata.uns[K.PHENOTYPE_CATEGORIES] - clonotype_categories = adata.uns[K.CLONOTYPE_CATEGORIES] - - metadata = adata.uns[K.METADATA] - covariate_col = metadata["covariate_col"] - - # Convert covariate_label to index - try: - cov_value = covariate_categories.index(covariate_label) - except ValueError: - raise ValueError(f"Covariate label '{covariate_label}' not found among: {covariate_categories}") - - # Get data specific to this covariate - chosen_mask = (ct_to_cov == cov_value) - chosen_idx = chosen_mask.nonzero(as_tuple=True)[0] - p_ct_for_cov = p_ct[chosen_mask] - - # Apply temperature scaling - eps = 1e-8 - p_ct_for_cov = F.softmax(torch.log(p_ct_for_cov + eps) / temperature, dim=-1) - - # Get clonotype indices for each chosen ct - clone_indices = ct_to_c[chosen_idx].numpy() - - # Get cell counts for each clonotype-covariate pair (for weighting) - ct_array_for_cells = adata.uns[K.CT_ARRAY] - cov_array_for_cells = adata.uns[K.COV_ARRAY] - - from collections import Counter - cell_mask = (cov_array_for_cells == cov_value) - cts_in_cov = ct_array_for_cells[cell_mask] - ct_counts_dict = Counter(cts_in_cov.tolist()) - - p_ct_arr = p_ct_for_cov.numpy() - - if n_samples == 0: - # Build dataframe with point estimates (no sampling) - df = pd.DataFrame(p_ct_arr, columns=phenotype_categories) - df["clonotype_index"] = clone_indices - df["clonotype_id"] = [clonotype_categories[i] for i in clone_indices] - - # Filter to requested clones - if clones is not None: - df = df[df["clonotype_id"].isin(clones)] - - # Apply clone size weighting if requested - if weighted: - counts = [] - for i, row in df.iterrows(): - ct_i = row["clonotype_index"] - c_count = ct_counts_dict.get(ct_i, 0) - counts.append(c_count) - - counts = np.array(counts, dtype=float) - df.loc[:, phenotype_categories] = df[phenotype_categories].values * counts[:, None] - - total_mass = df[phenotype_categories].sum().sum() - if total_mass > 0: - df.loc[:, phenotype_categories] = df[phenotype_categories] / total_mass - - # Set the index and clean up columns - df.index = df["clonotype_id"] - df = df[[col for col in df.columns if "clonotype" not in col]] - return df - - else: - # Sample from Dirichlet distribution - local_scale = adata.uns.get(K.LOCAL_SCALE, 1.0) - conc = local_scale * p_ct_for_cov - - samples = Dirichlet(conc).sample((n_samples,)) - samples_np = samples.cpu().numpy() - - # Reshape samples for DataFrame creation - num_chosen, num_pheno = p_ct_arr.shape - samples_expanded = samples_np.transpose(1, 0, 2).reshape(-1, num_pheno) - - # Create arrays for sample tracking - clonotype_indices_expanded = np.repeat(clone_indices, n_samples) - clonotype_ids_expanded = [clonotype_categories[i] for i in clonotype_indices_expanded] - sample_ids = np.tile(np.arange(n_samples), num_chosen) - - # Build dataframe - df_samples = pd.DataFrame(samples_expanded, columns=phenotype_categories) - df_samples["clonotype_index"] = clonotype_indices_expanded - df_samples["clonotype_id"] = clonotype_ids_expanded - df_samples["sample_id"] = sample_ids - - # Filter to requested clones - if clones is not None: - df_samples = df_samples[df_samples["clonotype_id"].isin(clones)] - - # Apply clone size weighting if requested - if weighted: - counts = [] - for i, row in df_samples.iterrows(): - ct_i = row["clonotype_index"] - c_count = ct_counts_dict.get(ct_i, 0) - counts.append(c_count) - counts = np.array(counts, dtype=float) - df_samples.loc[:, phenotype_categories] = ( - df_samples[phenotype_categories].values * counts[:, None] - ) - total_mass = df_samples[phenotype_categories].sum().sum() - if total_mass > 0: - df_samples.loc[:, phenotype_categories] /= total_mass - - # Set the index and clean up columns - df_samples.index = [ - f"{cid}_{sid}" for cid, sid in zip(df_samples["clonotype_id"], df_samples["sample_id"]) - ] - df_samples = df_samples[[col for col in df_samples.columns if col not in ["clonotype_id","clonotype_index","sample_id"]]] - return df_samples - -def get_latent_embedding( - adata, - latent_slot: str = K.X_TCRI, - n_samples: int = 0, - posterior_scale: float = 1.0 -) -> "np.ndarray": - mean_z = adata.obsm[latent_slot] - n_cells, latent_dim = mean_z.shape - samples = np.random.normal( - loc=mean_z, - scale=posterior_scale, - size=(n_samples, n_cells, latent_dim) - ) - return samples - -def group_small_clones(adata, patient_key=""): - ct = [] - for x, s, p in zip(adata.obs["trb"], adata.obs[K.CLONE_SIZE], adata.obs[patient_key]): - if s < 4: - ct.append("Singleton_{}".format(p)) - else: - ct.append("{}_{}".format(x,p)) - adata.obs["trb_unique"] = ct - -def register_probability_columns(adata, probability_columns): - adata.uns["probability_columns"] = probability_columns - -def gene_entropy(adata, key_added="entropy", batch_key=None, agg_function=None): - import tqdm - if batch_key == None: - X = adata.X.todense() - X = np.array(X.T) - gene_to_row = list(zip(adata.var.index.tolist(), X)) - entropies = [] - for _, exp in tqdm.tqdm(gene_to_row): - counts = np.unique(exp, return_counts = True) - entropies.append(entropy(counts[1][1:])) - adata.var[key_added] = entropies - else: - if agg_function == None: - agg_function = np.mean - entropies = collections.defaultdict(list) - for x in tqdm.tqdm(list(set(adata.obs[batch_key]))): - sdata = adata[adata.obs[batch_key]==x] - X = sdata.X.todense() - X = np.array(X.T) - gene_to_row = list(zip(sdata.var.index.tolist(), X)) - for symbol, exp in gene_to_row: - counts = np.unique(exp, return_counts = True) - entropies[symbol].append(entropy(counts[1][1:])) - aggregated_entropies = [] - for g in adata.var.index.tolist(): - ent = agg_function(entropies[g]) - aggregated_entropies.append(ent) - adata.var[key_added] = aggregated_entropies def clone_size(adata, key_added=K.CLONE_SIZE, return_counts=False): - tcr_key = adata.uns["tcri_clone_key"] + # Canonical source is uns[METADATA]['clone_col'] (written by to_anndata). This + # used to read the legacy uns['tcri_clone_key'] shadow key — the last reader of + # it, which is why the shim outlived Phase 4. + meta = adata.uns.get(K.METADATA) + if not meta or K.CLONE_COL not in meta: + raise KeyError( + f"adata.uns[{K.METADATA!r}][{K.CLONE_COL!r}] is missing — run " + "model.to_anndata(adata) first (or load a session) so the clonotype " + "column is registered." + ) + tcr_key = meta[K.CLONE_COL] res = np.unique(adata.obs[tcr_key].tolist(), return_counts=True) clone_sizes = dict(zip(res[0],res[1])) sizes = [] diff --git a/tcri/tools/__init__.py b/tcri/tools/__init__.py new file mode 100644 index 0000000..db9b016 --- /dev/null +++ b/tcri/tools/__init__.py @@ -0,0 +1,19 @@ +"""``tcri.tools`` (``tl``) — the joint-distribution engine + the four metric twins + +``compare_groups``, all engine-backed (PR6). ``tl`` points here (repointed from the old +``tcri.metrics`` in PR6); the metrics default to **bits (log2)** and read the unified +``joint_distribution``. +""" +from ._joint import joint_distribution +from ._entropy import clonotypic_entropy, phenotypic_entropy +from ._mutual_information import mutual_information +from ._flux import phenotypic_flux +from ._compare import compare_groups + +__all__ = [ + "joint_distribution", + "clonotypic_entropy", + "phenotypic_entropy", + "mutual_information", + "phenotypic_flux", + "compare_groups", +] diff --git a/tcri/tools/_common.py b/tcri/tools/_common.py new file mode 100644 index 0000000..fc727fd --- /dev/null +++ b/tcri/tools/_common.py @@ -0,0 +1,164 @@ +"""Shared reduction helpers for the ``tl`` metrics. + +Each metric pulls the clone×phenotype joint from the engine (:func:`joint_distribution`, +``use_logits=True``), reduces per draw, and — for ``n_samples>0`` — summarizes the draw +distribution (mean / sd / HDI). ``groupby`` is implemented here by **restricting clones +per group** (§7.1: full-space clone masks + ``clones=``, never slicing the AnnData), which +relies on clones being disjoint across groups (a TCR clone never spans two patients). +""" +from __future__ import annotations + +import numpy as np +import pandas as pd + +from .._stats import hdi +from ._joint import joint_distribution + + +def is_precomputed_joint(x) -> bool: + """A precomputed joint (fast path, §7.9) is a plain DataFrame, not an AnnData.""" + return isinstance(x, pd.DataFrame) + + +def joint_draws(adata, covariate, *, n_samples, weighted, temperature, clones, random_state, + use_logits=True): + """Return ``(draws, phenotype_cols)`` where ``draws`` is a list of ``(clone_ids, [C, P])`` + per posterior draw (length 1 for ``n_samples=0``). + + Consumes the engine's raw blocks directly. Going through + :func:`~tcri.tools._joint.joint_distribution` would flatten ``[S, n_rows, P]`` + into a MultiIndex DataFrame only for this function to ``groupby('sample_id')`` + and unpack it straight back to arrays — measured at ~2x the engine core itself. + Ordering matches the DataFrame path exactly (blocks in covariate order, clones in + block order, then the ``clones=`` filter applied as a stable reorder). + """ + from ._joint import _engine_blocks + + blocks, _n_draws, clonotype_cats, _cov_cats, cols = _engine_blocks( + adata, + covariate=covariate, + groupby=None, + n_samples=n_samples, + use_logits=use_logits, + weighted=weighted, + temperature=temperature, + random_state=random_state, + device=None, + ) + + # Row labels in DataFrame-concat order (per covariate block, per clone). When + # covariate is None the DataFrame carries a leading `covariate` index level, so + # its labels are (covariate, clonotype) TUPLES — reproduce that exactly, since + # callers key on whatever this returns. + all_cov = covariate is None + clone_names, ids = [], [] + for m, clone_idx, _J in blocks: + for i in clone_idx: + c = clonotype_cats[i] + clone_names.append(c) + ids.append((_cov_cats[m], c) if all_cov else c) + + keep = None + if clones is not None: + clones = list(clones) + rank = {c: i for i, c in enumerate(clones)} + sel = [j for j, c in enumerate(clone_names) if c in rank] + # stable sort by requested order — mirrors the MultiIndex argsort(kind="stable"), + # which ranks on the clonotype level only + keep = sorted(sel, key=lambda j: rank[clone_names[j]]) + ids = [ids[j] for j in keep] + + n_draws_out = blocks[0][2].shape[0] if blocks else 1 + draws = [] + for s in range(n_draws_out): + arr = np.concatenate([J[s] for _m, _ci, J in blocks], axis=0) if blocks else np.empty((0, len(cols))) + arr = arr.astype(float, copy=False) + if keep is not None: + arr = arr[keep] + draws.append((list(ids), arr)) + return draws, list(cols) + + +def summarize(values, *, hdi_prob=0.94) -> dict: + """Summarize a 1-D array of per-draw metric values → mean / sd / hdi_low / hdi_high.""" + v = np.asarray(values, dtype=float) + v = v[np.isfinite(v)] + if v.size == 0: + return {"mean": np.nan, "sd": np.nan, "hdi_low": np.nan, "hdi_high": np.nan} + if v.size == 1: + return {"mean": float(v[0]), "sd": 0.0, "hdi_low": float(v[0]), "hdi_high": float(v[0])} + lo, hi = hdi(v, prob=hdi_prob) + return {"mean": float(v.mean()), "sd": float(v.std(ddof=1)), "hdi_low": lo, "hdi_high": hi} + + +def clone_col(adata): + from .. import _keys as K + return adata.uns[K.METADATA]["clone_col"] + + +def _validate_group_clones(obs, groupby, cc): + """The metric ``groupby`` restricts the engine by clone id (``clones=``), which is only + correct when clones are **disjoint across groups** (a clone's cells all live in one group). + Raise loudly if a clone id spans groups — otherwise a group's estimate would silently + absorb that clone's cells from other groups (§7.1 groupby↔covariate semantics).""" + seen = {} + for g in obs[groupby].dropna().unique().tolist(): + for c in obs.loc[obs[groupby] == g, cc].dropna().unique(): + if c in seen and seen[c] != g: + raise ValueError( + f"groupby={groupby!r}: clonotype {c!r} spans groups {seen[c]!r} and {g!r}. " + f"The metric groupby restricts by clone id (clones=), which requires clones " + f"to be disjoint across groups (e.g. patient-specific `trb_unique`). Use a " + f"clone-disjoint groupby, or pre-filter with `clones=`." + ) + seen[c] = g + + +def grouped_scalar(adata, *, groupby, splitby, value, compute, hdi_prob=0.94): + """Loop over ``adata.obs[groupby]`` values, restrict to each group's clones, compute a + scalar-valued metric per group, and tidy into a DataFrame with the ``splitby`` label. + + ``compute(clones) -> (point, draws_or_None)``: ``point`` is the ``n_samples=0`` scalar (or + the draw mean), ``draws`` is the per-draw vector (or ``None`` at ``n_samples=0``). + """ + cc = clone_col(adata) + obs = adata.obs + _validate_group_clones(obs, groupby, cc) + rows = [] + for g in obs[groupby].dropna().unique().tolist(): + gmask = obs[groupby] == g + clones_g = obs.loc[gmask, cc].dropna().unique().tolist() + point, draws = compute(clones_g) + row = {groupby: g, value: point} + if splitby is not None and splitby in obs.columns: + row[splitby] = obs.loc[gmask, splitby].iloc[0] + if draws is not None: + row.update(summarize(draws, hdi_prob=hdi_prob)) + rows.append(row) + return pd.DataFrame(rows) + + +def grouped_series(adata, *, groupby, splitby, item_name, value, compute, hdi_prob=0.94): + """Like :func:`grouped_scalar` but the metric is a per-item (phenotype/clone) map. + + ``compute(clones) -> (point, draws_or_None)``: ``point`` is ``{item: value}``; + ``draws`` is ``{item: [per-draw values]}`` (or ``None`` at ``n_samples=0``). Tidies to + one row per (group, item). + """ + cc = clone_col(adata) + obs = adata.obs + _validate_group_clones(obs, groupby, cc) + rows = [] + for g in obs[groupby].dropna().unique().tolist(): + gmask = obs[groupby] == g + clones_g = obs.loc[gmask, cc].dropna().unique().tolist() + point, draws = compute(clones_g) + split_val = obs.loc[gmask, splitby].iloc[0] if (splitby and splitby in obs.columns) else None + for item, val in point.items(): + row = {groupby: g, item_name: item, value: val} + if split_val is not None: + row[splitby] = split_val + if draws is not None and item in draws: + row.update(summarize(draws[item], hdi_prob=hdi_prob)) + rows.append(row) + return pd.DataFrame(rows) diff --git a/tcri/tools/_compare.py b/tcri/tools/_compare.py new file mode 100644 index 0000000..fa9fd47 --- /dev/null +++ b/tcri/tools/_compare.py @@ -0,0 +1,81 @@ +"""``tl.compare_groups`` (§7.6) — the public group-comparison orchestrator that replaces +the deleted ``mi_compare`` / ``delta_entropy_table`` / ``flux_table``. + +Turns a tidy ``groupby`` result (per-unit point estimates, e.g. per patient) into group +contrasts. Unpaired: Mann–Whitney U + two-sided p + Δ of means. Paired (``paired=True``): +per-unit posterior-draw vectors aligned by ``pair_on`` → signed Δ draws, HDI, and the +direction probability ``p_gt`` via ``prob_direction`` (the only place a direction prob is +emitted). +""" +from __future__ import annotations + +import itertools + +import numpy as np +import pandas as pd + +from .._stats import hdi, mann_whitney, prob_direction, stars + +__all__ = ["compare_groups"] + + +def compare_groups(df, *, value, splitby, reference=None, paired=False, pair_on=None, + hdi_prob=0.94, alternative="two-sided"): + """Contrast ``value`` across ``df[splitby]`` groups. Returns a tidy DataFrame, one row + per contrast (``reference`` vs each other level, or all pairs when ``reference`` is None).""" + levels = [g for g in df[splitby].dropna().unique().tolist()] + if reference is not None: + if reference not in levels: + raise ValueError(f"reference {reference!r} not in {splitby} levels {levels}") + pairs = [(reference, g) for g in levels if g != reference] + else: + pairs = list(itertools.combinations(sorted(levels, key=str), 2)) + + rows = [] + for a, b in pairs: + da = df[df[splitby] == a] + db = df[df[splitby] == b] + if paired: + if pair_on is None: + raise ValueError("paired=True requires pair_on (the per-unit id to align draws on).") + sa = da.set_index(pair_on)[value] + sb = db.set_index(pair_on)[value] + diffs = [] + for u in sa.index.intersection(sb.index): + va = np.asarray(sa[u], dtype=float).ravel() + vb = np.asarray(sb[u], dtype=float).ravel() + n = min(va.size, vb.size) + if n: + diffs.append(vb[:n] - va[:n]) + delta = np.concatenate(diffs) if diffs else np.array([]) + if delta.size == 0: + continue + p_gt, p_lt = prob_direction(delta) + lo, hi = hdi(delta, prob=hdi_prob) if delta.size > 1 else (float(delta[0]), float(delta[0])) + rows.append(_row(a, b, delta=float(np.mean(delta)), p_gt=float(p_gt), + p_lt=float(p_lt), hdi_low=lo, hdi_high=hi)) + else: + va = pd.to_numeric(da[value], errors="coerce").dropna().to_numpy() + vb = pd.to_numeric(db[value], errors="coerce").dropna().to_numpy() + if va.size == 0 or vb.size == 0: + continue + try: + U, p = mann_whitney(va, vb, alternative=alternative) + except ValueError: # e.g. all-identical values + U, p = np.nan, np.nan + rows.append(_row(a, b, mean_a=float(np.mean(va)), mean_b=float(np.mean(vb)), + delta=float(np.mean(vb) - np.mean(va)), U=float(U), p=float(p), + stars=stars(p))) + return pd.DataFrame(rows, columns=_COLUMNS) + + +_COLUMNS = ["group_a", "group_b", "mean_a", "mean_b", "delta", "U", "p", "stars", + "p_gt", "p_lt", "hdi_low", "hdi_high"] + + +def _row(a, b, **vals): + """A contrast row with the unified §7.6 schema; inapplicable stats are NaN.""" + row = {c: np.nan for c in _COLUMNS} + row["group_a"], row["group_b"] = a, b + row.update(vals) + return row diff --git a/tcri/tools/_entropy.py b/tcri/tools/_entropy.py new file mode 100644 index 0000000..531b2f9 --- /dev/null +++ b/tcri/tools/_entropy.py @@ -0,0 +1,123 @@ +"""``tl.clonotypic_entropy`` and ``tl.phenotypic_entropy`` — normalized Shannon entropies +in **bits** (§7.2/§7.3), engine-backed. + +- ``clonotypic_entropy``: per phenotype φ, H[P(c|φ)] over the **supported** clones (spread of + a phenotype across clones). Absent/zero-support clones are excluded before normalizing (no + ε-clip fabricating uniform mass). Normalizer = log2(#supported clones). +- ``phenotypic_entropy``: per clone c, H[P(φ|c)] (plasticity vs commitment). A clone with zero + posterior mass returns **NaN** (not reindexed-to-zeros → spurious H=1). Normalizer = log2(P). + +``n_samples=0`` is the deterministic plug-in; ``n_samples>0`` returns the posterior-mean + +sd + HDI of the entropy (plug-in ≥ posterior-mean for entropy — documented as distinct). +""" +from __future__ import annotations + +import numpy as np +import pandas as pd + +from ._common import grouped_series, is_precomputed_joint, joint_draws, summarize + +__all__ = ["clonotypic_entropy", "phenotypic_entropy"] + + +def _clonotypic_one(J, cols, *, normalized, n_clones_ref=None): + """{phenotype: H[P(c|φ)] bits} over supported clones. ``n_clones_ref`` fixes the + normalizer log2(C) for cross-group comparability (else the #supported clones).""" + out = {} + for j, ph in enumerate(cols): + col = np.asarray(J[:, j], dtype=np.float64) + supp = col > 0 + s = col[supp].sum() + if supp.sum() == 0 or s <= 0: + out[ph] = np.nan + continue + v = col[supp] / s + H = float(-np.sum(v * np.log2(v))) + if normalized: + C = int(n_clones_ref) if n_clones_ref else int(supp.sum()) + if C > 1: + H /= np.log2(C) + out[ph] = H + return out + + +def _phenotypic_one(clone_ids, J, cols, *, normalized): + """{clone: H[P(φ|c)] bits}; zero-mass clone → NaN.""" + P = len(cols) + out = {} + for i, c in enumerate(clone_ids): + row = np.asarray(J[i], dtype=np.float64) + s = row.sum() + if s <= 0: + out[c] = np.nan + continue + p = row / s + # 0*log0 := 0. Mask first rather than np.where(p>0, p*log2(p), 0): numpy + # evaluates BOTH branches, so log2(0) still fires divide-by-zero warnings. + nz = p[p > 0] + H = float(-np.sum(nz * np.log2(nz))) + if normalized and P > 1: + H /= np.log2(P) + out[c] = H + return out + + +def _entropy_metric(adata_or_jd, *, kind, covariate, groupby, splitby, n_samples, temperature, + clones, weighted, normalized, random_state, n_clones_ref=None): + item_name = "phenotype" if kind == "clonotypic" else "clonotype" + value = f"{kind}_entropy" + + def _one(clone_ids, J, cols): + if kind == "clonotypic": + return _clonotypic_one(J, cols, normalized=normalized, n_clones_ref=n_clones_ref) + return _phenotypic_one(clone_ids, J, cols, normalized=normalized) + + if groupby is not None: + if is_precomputed_joint(adata_or_jd): + 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, + temperature=temperature, clones=cl, random_state=random_state) + per = [_one(ids, J, cols) for ids, J in draws] + keys = list(per[0].keys()) + point = {k: float(np.nanmean([p[k] for p in per])) for k in keys} + drawsd = {k: [p[k] for p in per] for k in keys} if (n_samples and int(n_samples) > 0) else None + return point, drawsd + return grouped_series(adata_or_jd, groupby=groupby, splitby=splitby, + item_name=item_name, value=value, compute=_compute) + + if is_precomputed_joint(adata_or_jd): + if n_samples and int(n_samples) > 0: + raise ValueError("precomputed-joint fast path is valid only at n_samples=0 (§7.9).") + 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, + temperature=temperature, clones=clones, random_state=random_state) + per = [_one(ids, J, cols) for ids, J in draws] + keys = list(per[0].keys()) + if n_samples and int(n_samples) > 0: + return pd.DataFrame({k: summarize([p[k] for p in per]) for k in keys}).T + return pd.Series(per[0], name=value) + + +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): + """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) + + +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): + """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) diff --git a/tcri/tools/_flux.py b/tcri/tools/_flux.py new file mode 100644 index 0000000..587b09f --- /dev/null +++ b/tcri/tools/_flux.py @@ -0,0 +1,66 @@ +"""``tl.phenotypic_flux`` (renamed from ``flux``) — per-clonotype phenotype-distribution +shift between two covariate values (§7.5), engine-backed. + +For each clone in the ``cov_from`` ∩ ``cov_to`` intersection, the distance between its +phenotype distribution at ``cov_from`` and ``cov_to``, via ``_distance`` (l1 / kl / jsd, +KL & JSD in **bits**). ``n_samples=0`` is the plug-in (a clone with no real shift reads +exactly 0); ``n_samples>0`` redraws both sides coherently (same seed) and summarizes. +""" +from __future__ import annotations + +import numpy as np +import pandas as pd + +from .._distance import phenotype_distance +from ._common import grouped_series, joint_draws, summarize + +__all__ = ["phenotypic_flux"] + + +def _flux_once(adata, *, cov_from, cov_to, n_samples, weighted, temperature, clones, + distance_metric, random_state): + dist_fn = phenotype_distance(distance_metric) + draws_from, _ = joint_draws(adata, cov_from, n_samples=n_samples, weighted=weighted, + temperature=temperature, clones=clones, random_state=random_state) + draws_to, _ = joint_draws(adata, cov_to, n_samples=n_samples, weighted=weighted, + temperature=temperature, clones=clones, random_state=random_state) + per = [] + for (ids_f, Jf), (ids_t, Jt) in zip(draws_from, draws_to): + idx_t = {c: i for i, c in enumerate(ids_t)} + d = {} + for i, c in enumerate(ids_f): + if c not in idx_t: + continue + p = np.asarray(Jf[i], dtype=np.float64); q = np.asarray(Jt[idx_t[c]], dtype=np.float64) + ps, qs = p.sum(), q.sum() + if ps <= 0 or qs <= 0: + d[c] = np.nan + continue + d[c] = float(dist_fn(p / ps, q / qs)) + per.append(d) + keys = sorted(set().union(*[set(p) for p in per])) if per else [] + point = {c: float(np.nanmean([p.get(c, np.nan) for p in per])) for c in keys} + drawsd = ({c: [p.get(c, np.nan) for p in per] for c in keys} + if (n_samples and int(n_samples) > 0) else None) + return point, drawsd + + +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): + """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) + 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) + 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/_joint.py b/tcri/tools/_joint.py new file mode 100644 index 0000000..6fe1b80 --- /dev/null +++ b/tcri/tools/_joint.py @@ -0,0 +1,220 @@ +"""The joint-distribution engine (``tl``) — a thin DataFrame wrapper over +:func:`tcri._compute._joint._joint_draws`. Re-exported top-level as +``tcri.joint_distribution``; unifies the old ``joint_distribution`` + +``joint_distribution_posterior``. + +See §7.1 of ``docs/contract/tcri_api_and_responsibilities.md`` for the math. This +is the substrate every metric consumes (Phase 6 migrates them onto it). +""" +from __future__ import annotations + +import numpy as np +import pandas as pd + +from .. import _keys as K +from .._compute._joint import _joint_draws + +__all__ = ["joint_distribution"] + + +def _engine_blocks( + adata, + *, + covariate, + groupby, + n_samples, + use_logits, + weighted, + temperature, + random_state, + device, +): + """Validate, then run the numeric core; return the raw per-covariate blocks. + + Shared by :func:`joint_distribution` (which formats them into a labelled + DataFrame) and the metric fast path (which consumes the arrays directly), so + both see identical validation, the same engine call, and one shared draw. + + Returns ``(blocks, n_draws, clonotype_cats, covariate_cats, phenotype_cats)`` + where ``blocks`` is a list of ``(covariate_index, clone_idx, J)`` and ``J`` has + shape ``[S, n_rows, P]``. + """ + if groupby is not None: + raise NotImplementedError( + "joint_distribution(groupby=...) is not implemented in the engine yet; it " + "lands with the metric migration (Phase 6) so it can share the draw across " + "groups and enforce the clone-determined guard. Restrict with `clones=` for now." + ) + + phenotype_cats = list(adata.uns[K.PHENOTYPE_CATEGORIES]) + clonotype_cats = list(adata.uns[K.CLONOTYPE_CATEGORIES]) + covariate_cats = list(adata.uns[K.COVARIATE_CATEGORIES]) + + logits = None + if use_logits: + if K.X_LOGITS not in adata.obsm: + raise RuntimeError( + f"obsm[{K.X_LOGITS!r}] missing — run model.to_anndata(...) or pass use_logits=False." + ) + logits = adata.obsm[K.X_LOGITS] + + cov_idx = None + if covariate is not None: + try: + cov_idx = covariate_cats.index(covariate) + except ValueError: + raise ValueError(f"covariate {covariate!r} not found among {covariate_cats}") + + # subset/filtered-AnnData guard: the per-cell uns arrays live in full-cell space and + # are NOT sliced when adata is subset, whereas obsm/obs ARE — so a slice silently + # misaligns cells. Fail loudly (mirrors the legacy joint_distribution_posterior guard). + n_obs = adata.n_obs + n_reg = len(np.asarray(adata.uns[K.CT_ARRAY])) + if n_reg != n_obs or len(np.asarray(adata.uns[K.COV_ARRAY])) != n_obs: + raise ValueError( + f"joint_distribution received an AnnData whose per-cell registration arrays " + f"(uns[{K.CT_ARRAY!r}], len {n_reg}) do not match adata.n_obs ({n_obs}). This " + f"happens on a filtered/sliced AnnData: the full-space uns arrays misalign against " + f"the subset obsm/obs. Re-run model.to_anndata(...) on the filtered object, or pass " + f"the full object and filter with `clones=`." + ) + + # local_scale is required for the Dirichlet draw; refuse to silently fall back at n>0. + local_scale = adata.uns.get(K.LOCAL_SCALE, None) + if n_samples and int(n_samples) > 0 and local_scale is None: + raise RuntimeError( + f"n_samples>0 needs uns[{K.LOCAL_SCALE!r}] for the clamped-Dirichlet draw, but it " + f"is missing; run model.to_anndata(...)." + ) + local_scale = float(local_scale) if local_scale is not None else 1.0 + + blocks, n_draws = _joint_draws( + adata.uns[K.P_CT], + adata.uns[K.CT_TO_COV], + adata.uns[K.CT_TO_C], + adata.uns[K.CT_ARRAY], + adata.uns[K.COV_ARRAY], + local_scale=local_scale, + n_samples=n_samples, + temperature=temperature, + use_logits=use_logits, + covariate_idx=cov_idx, + logits=logits, + gate_prob=adata.uns.get(K.GATE_PROB, None), + weighted=weighted, + random_state=random_state, + device=device, + ) + return blocks, n_draws, clonotype_cats, covariate_cats, phenotype_cats + + +def joint_distribution( + adata, + *, + covariate=None, + groupby=None, + n_samples=0, + use_logits=True, + weighted=False, + clones=None, + temperature=1.0, + random_state=None, + device=None, +) -> pd.DataFrame: + """Clone×phenotype distribution from the learned posterior of ``p_ct``. + + Parameters + ---------- + covariate : str | None + A covariate value; ``None`` computes all covariate values in one shared-draw + pass (adds a leading ``covariate`` index level). + groupby : str | None + Not yet implemented in the engine (lands with the metric migration, Phase 6); + restrict with ``clones=`` for now. + n_samples : int + ``0`` → deterministic posterior-mean table; ``N`` → ``N`` clamped-Dirichlet + draws (adds a ``sample_id`` index level). Only place ``local_scale`` enters. + use_logits : bool + ``True`` → fold per-cell classifier logits with ``log(base)`` (gate-aware) and + aggregate per clone, matching :meth:`~tcri.model._model.TCRIModel.predict`; + ``False`` → the ct-level base table. Neither touches the generative prior. + weighted : bool + ``False`` → each clone is one unit (per-clone simplex). ``True`` → each clone + row is scaled by its (ct-keyed) cell count (cell-weighted). + temperature : float + Tempers the base once; ``T=1`` is the identity (and reproduces ``predict()`` + on the ``use_logits=True`` path). + random_state : int | numpy.Generator | torch.Generator | None + Seeds the torch Dirichlet generator for ``n_samples>0``; ignored at ``0``. + device : str | None + Routes the numeric core through ``_compute/_xp`` (CPU / torch-CUDA). Result is + always a host DataFrame. + + Returns + ------- + pandas.DataFrame + Columns = phenotype categories. Index = clonotype (``+ sample_id`` for + ``n_samples>0``, ``+ covariate`` leading level for ``covariate=None``). + Provenance in ``df.attrs["params"]``. + """ + blocks, n_draws, clonotype_cats, covariate_cats, phenotype_cats = _engine_blocks( + adata, + covariate=covariate, + groupby=groupby, + n_samples=n_samples, + use_logits=use_logits, + weighted=weighted, + temperature=temperature, + random_state=random_state, + device=device, + ) + + sampling = bool(n_samples and int(n_samples) > 0) + all_cov = covariate is None + frames = [] + for m, clone_idx, J in blocks: # J: [S, n_rows, P] + clone_ids = [clonotype_cats[i] for i in clone_idx] + S = J.shape[0] + if sampling: + arr = J.transpose(1, 0, 2).reshape(-1, J.shape[2]) # [n_rows*S, P] + idx_clone = np.repeat(clone_ids, S) + idx_samp = np.tile(np.arange(S), len(clone_ids)) + cols = [idx_clone, idx_samp] + names = ["clonotype", "sample_id"] + else: + arr = J[0] # [n_rows, P] + cols = [clone_ids] + names = ["clonotype"] + if all_cov: + cols = [[covariate_cats[m]] * arr.shape[0], *cols] + names = ["covariate", *names] + index = (pd.MultiIndex.from_arrays(cols, names=names) + if len(cols) > 1 else pd.Index(cols[0], name=names[0])) + frames.append(pd.DataFrame(arr, columns=phenotype_cats, index=index)) + + df = pd.concat(frames) if len(frames) > 1 else frames[0] + + if clones is not None: + clones = list(clones) + if isinstance(df.index, pd.MultiIndex): + # filter to the listed clones (absent dropped, not all-zero) then order by the + # requested list — stable within the sample_id/covariate levels (matches the + # single-index reindex; §7.1 "reindex to the exact list"). + keep = df.index.get_level_values("clonotype").isin(clones) + df = df[keep] + rank = pd.Index(df.index.get_level_values("clonotype")).map({c: i for i, c in enumerate(clones)}) + df = df.iloc[np.argsort(np.asarray(rank), kind="stable")] + else: + df = df.reindex([c for c in clones if c in df.index]) + + df.attrs["params"] = { + "covariate": covariate, + "groupby": groupby, + "n_samples": int(n_samples), + "use_logits": bool(use_logits), + "weighted": bool(weighted), + "temperature": float(temperature), + "n_draws": int(n_draws), + "clones": None if clones is None else [str(c) for c in clones], + } + return df diff --git a/tcri/tools/_metrics_contract.py b/tcri/tools/_metrics_contract.py new file mode 100644 index 0000000..d3d2ada --- /dev/null +++ b/tcri/tools/_metrics_contract.py @@ -0,0 +1,185 @@ +"""Frozen definitions of the information-theoretic metrics (the *metrics contract*). + +Companion to the **model** contract (:mod:`tcri.model._model_contract`), which freezes +the generative mathematics. This module freezes what the **metrics** compute: the +entropies and mutual information over a clone x phenotype joint. + +The two are separate on purpose — they are checked by different means. The model +contract is verified by *tracing* ``model()``/``guide()`` for sites, plates and +distribution families. Metrics are pure functions of a joint table, so they are +verified by *numeric identities* (uniform -> log2(k), independent -> MI 0, and the +entropy/MI decomposition) in +``tests/test_metrics_contract_conformance.py``. + +Source of truth: **Supplementary Note 1**, "Entropy" section (eqs 2-4) and the mutual +information it defines — with the erratum recorded in ``SOURCE_ERRATA`` below. + +Changing any definition here means changing what the published numbers mean. +**Update this manifest and ``docs/contract/METRICS_CONTRACT.md`` FIRST, then the +code** — never loosen a definition to make a failing conformance test pass. +""" +from __future__ import annotations + +__all__ = [ + "LOG_BASE", + "METRIC_SPECS", + "IDENTITIES", + "SOURCE_ERRATA", + "SANCTIONED_EXTENSIONS", + "MetricSpec", +] + +#: All entropies/MI are reported in **bits**. The note writes an unspecified ``log``; +#: tcri fixes base 2 throughout so entropies read as bits and normalizers are log2(k). +LOG_BASE = 2 + + +class MetricSpec: + """One frozen metric definition.""" + + def __init__(self, name, formula, per, support, normalizer, empty, note_eq): + self.name = name + self.formula = formula # the exact quantity computed + self.per = per # what one output value corresponds to + self.support = support # how zero/absent mass is handled + self.normalizer = normalizer + self.empty = empty # value when there is no mass + self.note_eq = note_eq + + def __repr__(self): # pragma: no cover - debug aid + return f"MetricSpec({self.name!r}, {self.formula!r})" + + +METRIC_SPECS = { + "clonotypic_entropy": MetricSpec( + name="clonotypic_entropy", + formula="H[P(c|phi)] = -sum_c P(c|phi) log2 P(c|phi)", + per="one value per PHENOTYPE (how spread that phenotype is across clones)", + support=( + "SUPPORT-ONLY: clones with zero mass in the column are dropped BEFORE " + "renormalizing. No epsilon clip — fabricating uniform mass on absent " + "clones would inflate H toward 1." + ), + normalizer="log2(#supported clones), or log2(n_clones_ref) when given", + empty="NaN when the phenotype column has no positive mass", + note_eq="eq 3 (see SOURCE_ERRATA['eq3_weights_marginal'])", + ), + "phenotypic_entropy": MetricSpec( + name="phenotypic_entropy", + formula="H[P(phi|c)] = -sum_phi P(phi|c) log2 P(phi|c)", + per="one value per CLONE (plasticity vs commitment of that clone)", + support=( + "All P phenotypes are in the sum; 0*log0 is taken as 0. A clone with zero " + "total mass yields NaN — it is NOT reindexed to zeros, which would report " + "a spurious H=1 for a clone that was never observed." + ), + normalizer="log2(P), P = number of phenotype categories", + empty="NaN when the clone row has no positive mass", + note_eq="eq 4 (see SOURCE_ERRATA['eq4_label_and_weights'])", + ), + "mutual_information": MetricSpec( + name="mutual_information", + formula="I(c;phi) = sum_{c,phi} P(c,phi) log2( P(c,phi) / (P(c) P(phi)) )", + per="one value per (covariate) joint table", + support=( + "The table is renormalized to a joint; a numerical epsilon (1e-15) guards " + "log(0). Rows/columns with zero mass contribute zero." + ), + normalizer=( + "normalize_mode='min' (DEFAULT): I / min(H(c), H(phi)) — the coefficient of " + "constraint. 'average': I / (0.5*(H(c)+H(phi))). 'min' is the default " + "because the 'average' denominator scales with log2(C) and is therefore NOT " + "comparable across groups with different clone counts." + ), + empty="NaN when the table has no positive mass", + note_eq="the MI defined over the joint in the note's Entropy section", + ), +} + + +#: Identities the conformance test enforces. These are what make a redefinition +#: detectable: any change to the metrics that breaks one of these is a contract change. +IDENTITIES = { + "entropy_uniform_is_log2_k": ( + "A uniform distribution over k supported outcomes has H = log2(k) bits, and " + "normalized H = 1.0." + ), + "entropy_degenerate_is_zero": ( + "All mass on one outcome gives H = 0 (normalized and unnormalized)." + ), + "entropy_zero_mass_is_nan": ( + "A clone/phenotype with no posterior mass yields NaN — never 0 and never a " + "spurious 1 from reindexing absent entries to zeros." + ), + "mi_independent_is_zero": ( + "For an independent joint P(c,phi) = P(c)P(phi), I(c;phi) = 0." + ), + "mi_is_symmetric": "I(c;phi) == I(phi;c) (transposing the table is a no-op).", + "mi_is_nonnegative": "I(c;phi) >= 0 for every joint.", + "mi_perfect_coupling_is_one": ( + "For a permutation-like joint (each clone in exactly one phenotype and vice " + "versa), normalized MI with mode='min' is 1.0." + ), + "mi_entropy_decomposition": ( + "THE cross-metric identity: I(c;phi) = H(c) - sum_phi P(phi) * H[P(c|phi)], " + "where H[P(c|phi)] is the UNNORMALIZED clonotypic entropy. This ties the two " + "metric families together — it is the identity that proved the note's eqs 3-4 " + "are mistranscribed (weighting by the marginal instead of the conditional " + "makes this yield a NEGATIVE mutual information)." + ), +} + + +#: Errors in the source document, kept explicit so nobody "fixes" the code to match a +#: typo. The code is correct; the note's transcription is not. +SOURCE_ERRATA = { + "eq3_weights_marginal": ( + "Note eq 3 reads H(p(c|phi)) = -sum_c p(c) log p(c|phi) — it weights by the " + "MARGINAL p(c) while taking the log of the CONDITIONAL. That is a " + "cross-entropy, not an entropy. The intended (and implemented) quantity is " + "-sum_c p(c|phi) log p(c|phi)." + ), + "eq4_label_and_weights": ( + "Note eq 4 is labelled H(p(c)) but its right-hand side sums over phi and uses " + "p(phi|c), so the label is wrong; it also weights by the marginal p(phi) " + "rather than the conditional. The intended (and implemented) quantity is " + "-sum_phi p(phi|c) log p(phi|c)." + ), + "prose_says_marginal": ( + "The prose introduces eqs 3-4 as 'the entropy of the marginal distributions', " + "but both equations are conditionals. A marginal entropy would be " + "-sum_c p(c) log p(c)." + ), + "why_the_code_is_right": ( + "Decisive check: mutual information must satisfy " + "I(c;phi) = H(c) - E_phi[H(c|phi)]. On a test joint with true MI 0.288703, the " + "implemented conditional entropy reproduces it exactly, while the note's " + "literal formula gives -0.345883 — a negative mutual information, which is " + "impossible. The literal equations are inconsistent with the note's own MI." + ), +} + + +#: Deliberate additions beyond the note. Not deviations from its mathematics — the +#: note simply does not specify them. +SANCTIONED_EXTENSIONS = { + "bits_log2": ( + "The note writes an unspecified `log`; tcri fixes base 2 so all entropies are " + "in bits and normalizers are log2(k)." + ), + "normalization": ( + "`normalized=True` divides by the maximum-entropy value so results land in " + "[0,1] and compare across groups of different size. The note defines only the " + "raw entropies." + ), + "n_clones_ref": ( + "clonotypic_entropy accepts `n_clones_ref` to FIX the normalizer across groups; " + "without it each group is normalized by its own supported-clone count, which is " + "not comparable between groups." + ), + "posterior_summaries": ( + "`n_samples>0` returns mean/sd/HDI of the metric over posterior draws instead of " + "a single plug-in value. The plug-in entropy is >= the posterior mean (Jensen), " + "so the two are reported as distinct quantities." + ), +} diff --git a/tcri/tools/_mutual_information.py b/tcri/tools/_mutual_information.py new file mode 100644 index 0000000..e836e87 --- /dev/null +++ b/tcri/tools/_mutual_information.py @@ -0,0 +1,69 @@ +"""``tl.mutual_information`` — clone↔phenotype coupling I(c;φ|m) in **bits** (§7.4). + +Engine-backed rewrite of the old ``metrics.mutual_information``. Default +``normalize_mode="min"`` (coefficient of constraint I/min(H_c,H_p)) — the ``"average"`` +denominator throttles normalized MI by ~1/log2(C) and is non-comparable across groups with +different clone counts (the blocking fix). ``n_samples=0`` is the deterministic plug-in. +""" +from __future__ import annotations + +import numpy as np + +from ._common import grouped_scalar, is_precomputed_joint, joint_draws, summarize + +__all__ = ["mutual_information"] + +_EPS = 1e-15 + + +def _mi_from_joint(J: np.ndarray, *, normalized: bool = True, mode: str = "min") -> float: + """MI (bits) of a clone×phenotype table ``J`` (any scale — renormalized here).""" + J = np.asarray(J, dtype=np.float64) + total = J.sum() + if total <= 0: + return np.nan + pxy = J / total + px = pxy.sum(1, keepdims=True) # P(clone) + py = pxy.sum(0, keepdims=True) # P(phenotype) + mi = float(np.sum(pxy * np.log2((pxy + _EPS) / (px @ py + _EPS)))) + if not normalized: + return mi + h_c = float(-np.sum(px * np.log2(px + _EPS))) + h_p = float(-np.sum(py * np.log2(py + _EPS))) + denom = min(h_c, h_p) if mode == "min" else 0.5 * (h_c + h_p) + return mi / denom if denom > 0 else 0.0 + + +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, +): + """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``).""" + if groupby is not None: + if is_precomputed_joint(adata_or_jd): + 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, + temperature=temperature, clones=cl, random_state=random_state, + ) + vals = [_mi_from_joint(J, normalized=normalized, mode=normalize_mode) for _, J in draws] + return (float(np.nanmean(vals)), vals if (n_samples and int(n_samples) > 0) else None) + return grouped_scalar(adata_or_jd, groupby=groupby, splitby=splitby, value="MI", compute=_compute) + + if is_precomputed_joint(adata_or_jd): + if n_samples and int(n_samples) > 0: + raise ValueError("precomputed-joint fast path is valid only at n_samples=0 (§7.9).") + 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, + temperature=temperature, clones=clones, random_state=random_state, + ) + vals = [_mi_from_joint(J, normalized=normalized, mode=normalize_mode) for _, J in draws] + if n_samples and int(n_samples) > 0: + return summarize(vals) + return vals[0] diff --git a/tcri/utils/_utils.py b/tcri/utils/_utils.py index dd1c338..7ef3401 100644 --- a/tcri/utils/_utils.py +++ b/tcri/utils/_utils.py @@ -236,15 +236,6 @@ def load_tcri_session( -def probabilities(adata): - matrix = adata.obs[adata.uns["probability_columns"]] - barcodes = matrix.index.tolist() - cells = np.nan_to_num(matrix.to_numpy()) - index = adata.uns["joint_distribution"].index - probabs = dict() - for bc, cell in zip(barcodes, cells): - probabs[bc] = dict(zip(index, cell)) - return probabs tcri_colors = [ @@ -280,178 +271,10 @@ def probabilities(adata): "#6D6A75", # Purple Gray ] -import daft import matplotlib.pyplot as plt -def build_nested_tcri_pgm(): - """ - A fully explicit TCRI PGM matching the implementation in _model.py - with improved layout to minimize edge crossings - """ - # Define colors - red, yellow, green, gray, blue = "#cd442a", "#f0bd00", "#7e9437", "#eee", "#009de1" - - # Create a PGM canvas - pgm = daft.PGM( - shape=[8, 8], # width x height - origin=[0, 0], - grid_unit=1.6, - node_unit=1.5 - ) - - # ------------------------------------------------------------------ - # 1) Global hyperparameters for Dirichlet priors - # ------------------------------------------------------------------ - pgm.add_node( - "global_scale", - r"$\mathrm{global\_scale}$", - 5.1, # x - 6.6, # y - fixed=True, - plot_params={"fc": "#DDD"} - ) - - # ------------------------------------------------------------------ - # 2) Plate: batch (b) - outermost plate - # ------------------------------------------------------------------ - - # Batch-level variables - aligned vertically - # ------------------------------------------------------------------ - # 3) Plate: clonotypes (c) - middle plate - # ------------------------------------------------------------------ - pgm.add_plate( - [1.0, 1.0, 6.0, 6.3], # [x, y, width, height] - label=r"clonotypes $(c)$", - shift=-0.1 - ) - - # p_c (Dirichlet) - pgm.add_node( - "p_c", - r"$p_c$", - 3.8, # x - 6.6, # y - observed=False, - plot_params={"fc": blue} - ) - # Edge: global_scale -> p_c - pgm.add_edge("global_scale", "p_c") - # ------------------------------------------------------------------ - # 4) Plate: clone-covariate (ct) - inner plate - # ------------------------------------------------------------------ - pgm.add_plate( - [1.5, 1.5, 5.0, 4.5], # [x, y, width, height] - label=r"clone-covariate $(ct)$", - shift=-0.1 - ) - - # local_scale moved inside clone-covariate plate - pgm.add_node( - "local_scale", - r"$\mathrm{local\_scale}$", - 5.2, # x - 5.3, # y - fixed=True, - plot_params={"fc": "#DDD"} - ) - - # p_ct (Dirichlet) - pgm.add_node( - "p_ct", - r"$p_{ct}$", - 3.8, # x - 5.3, # y - observed=False, - plot_params={"fc": yellow} - ) - - # Edges: p_c -> p_ct, local_scale -> p_ct - pgm.add_edge("p_c", "p_ct") - pgm.add_edge("local_scale", "p_ct") - - # ------------------------------------------------------------------ - # 5) Plate: data (i) - innermost plate - # ------------------------------------------------------------------ - pgm.add_plate( - [1.9, 2., 4.0, 2.7], # [x, y, width, height] - label=r"data $(i)$", - shift=-0.1 - ) - - # Grid layout for data-level variables - aligned vertically - # Column 1: Observed variables - pgm.add_node( - "obs", - r"$X_{i}$", - 5, # x - 4.0, # y - observed=True, - plot_params={"fc": gray} - ) - - pgm.add_node( - "obs_label", - r"$Pheno_{i}$", - 2.5, # x - 4.0, # y - observed=True, - plot_params={"fc": gray} - ) - - # Column 2: Latent variables - pgm.add_node( - "latent", - r"$z_i$", - 3.8, # x - 2.8, # y - observed=False, - plot_params={"fc": green} - ) - - pgm.add_node( - "z_i_phen", - r"$z_{i,\mathrm{phen}}$", - 3.8, # x - 4.0, # y - observed=False, - plot_params={"fc": red} - ) - - # Column 3: Decoder inputs - pgm.add_node( - "px_r", - r"$ZINB(X_{i})$", - 5, # x - 2.8, # y - observed=False, - plot_params={"fc": "#DDD"} - ) - - # Edges - now mostly vertical and horizontal - # Data-level edges - pgm.add_edge("p_ct", "z_i_phen") - pgm.add_edge("latent", "z_i_phen") - pgm.add_edge("z_i_phen", "obs_label") - - # Direct connections to z_i (previously through decoder) - pgm.add_edge("latent", "obs") - pgm.add_edge("px_r", "obs") - - # ------------------------------------------------------------------ - # Text / Title - # ------------------------------------------------------------------ - pgm.add_text(3.1,7.5, "TCRi Model", fontsize=14) - - return pgm - - -def draw_tcri_pgm_nested(): - pgm = build_nested_tcri_pgm() - pgm.render() - pgm.figure.savefig("tcri_model_fully_explicit.pdf", dpi=300) - plt.show() @@ -493,22 +316,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] = {} @@ -590,10 +398,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..bed7668 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -102,13 +102,16 @@ def synthetic_adata(): X = rng.poisson(lam=1.5, size=(n_cells, n_genes)).astype(np.float32) + patient = rng.choice(batches, size=n_cells) + base_clone = rng.choice([f"clone_{i}" for i in range(n_clones)], size=n_cells) + obs = pd.DataFrame({ - "unique_clone_id": rng.choice( - [f"clone_{i}" for i in range(n_clones)], size=n_cells - ), + # patient-specific clone ids (disjoint across patients), as real `trb_unique` is — + # so metric groupby='patient' is valid (clones don't span groups). + "unique_clone_id": [f"{c}_{p}" for c, p in zip(base_clone, patient)], "phenotype_col": rng.choice(phenotypes, size=n_cells), "timepoint": rng.choice(covariates, size=n_cells), - "patient": rng.choice(batches, size=n_cells), + "patient": patient, }) for col in obs.columns: obs[col] = obs[col].astype("category") @@ -118,11 +121,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 +156,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..642483c 100644 --- a/tests/test_contract_conformance.py +++ b/tests/test_contract_conformance.py @@ -20,9 +20,33 @@ 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]] = { + "tl.joint_distribution": ("tcri.tools._joint", "joint_distribution"), + "tl.clonotypic_entropy": ("tcri.tools._entropy", "clonotypic_entropy"), + "tl.phenotypic_entropy": ("tcri.tools._entropy", "phenotypic_entropy"), + "tl.mutual_information": ("tcri.tools._mutual_information", "mutual_information"), + "tl.phenotypic_flux": ("tcri.tools._flux", "phenotypic_flux"), + "tl.compare_groups": ("tcri.tools._compare", "compare_groups"), + "diag.joint_distribution_ppc": ("tcri.diagnostics._ppc", "joint_distribution_ppc"), + "diag.phenotype_calibration": ("tcri.diagnostics._ppc", "phenotype_calibration"), + "diag.reconstruction_ppc": ("tcri.diagnostics._ppc", "reconstruction_ppc"), + "diag.permutation_null": ("tcri.diagnostics._ppc", "permutation_null"), + "diag.loss": ("tcri.diagnostics._training", "loss"), + "diag.archetypes": ("tcri.diagnostics._training", "archetypes"), + "pl.clonotypic_entropy": ("tcri.plotting._entropy", "clonotypic_entropy"), + "pl.phenotypic_entropy": ("tcri.plotting._entropy", "phenotypic_entropy"), + "pl.mutual_information": ("tcri.plotting._mutual_information", "mutual_information"), + "pl.phenotypic_flux": ("tcri.plotting._flux", "phenotypic_flux"), + "pl.resolve_palette": ("tcri.plotting._colors", "resolve_palette"), + "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_diag/__init__.py b/tests/test_diag/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_diag/test_diag.py b/tests/test_diag/test_diag.py new file mode 100644 index 0000000..f86e4da --- /dev/null +++ b/tests/test_diag/test_diag.py @@ -0,0 +1,43 @@ +"""tcri.diag smoke — each PPC returns a DataFrame; the two plots return axes.""" +import matplotlib + +matplotlib.use("Agg") + +import pandas as pd + +import tcri + + +def test_joint_distribution_ppc(trained_model): + _, adata = trained_model + df = tcri.diag.joint_distribution_ppc(adata, distance_metric="l1") + assert isinstance(df, pd.DataFrame) + if len(df): + assert {"covariate", "clonotype", "distance"}.issubset(df.columns) + + +def test_phenotype_calibration(trained_model): + _, adata = trained_model + df = tcri.diag.phenotype_calibration(adata, n_bins=5) + assert {"bin", "mean_pred", "emp_freq", "count"}.issubset(df.columns) + assert "ECE" in df.attrs and df.attrs["ECE"] >= 0 + + +def test_reconstruction_ppc(trained_model): + model, adata = trained_model + df = tcri.diag.reconstruction_ppc(model, adata, n_sims=1, random_state=0) + assert {"statistic", "observed", "simulated", "discrepancy"}.issubset(df.columns) + assert (df["discrepancy"] >= 0).all() + + +def test_permutation_null(trained_model): + _, adata = trained_model + df = tcri.diag.permutation_null(adata, n_perm=30, random_state=0) + assert {"covariate", "observed", "null_mean", "null_sd", "z", "p"}.issubset(df.columns) + assert ((df["p"] >= 0) & (df["p"] <= 1)).all() + + +def test_loss_and_archetypes(trained_model): + model, _ = trained_model + assert tcri.diag.loss(model) is not None + assert tcri.diag.archetypes(model) is not None diff --git a/tests/test_helpers.py b/tests/test_helpers.py index 137d37d..20f716e 100644 --- a/tests/test_helpers.py +++ b/tests/test_helpers.py @@ -11,8 +11,12 @@ def test_keys_constants(): assert K.X_LOGITS == "X_tcri_logits" assert K.METADATA == "tcri_metadata" assert K.X_PROBABILITIES == "X_tcri_probabilities" - # legacy keys are present (so the removal step can target them) but distinct - assert K.LEGACY_CLONE_KEY == "tcri_clone_key" != K.CLONE_COL + # the legacy shadow keys are GONE (Phase 4 removal completed); only the + # defensively-popped manager stash name remains + assert not hasattr(K, "LEGACY_CLONE_KEY") + assert not hasattr(K, "LEGACY_PHENOTYPE_KEY") + assert not hasattr(K, "LEGACY_X_PHENOTYPES") + assert K.LEGACY_MANAGER == "tcri_manager" def test_console_aliases_and_callables(): @@ -55,6 +59,43 @@ def test_auc_helpers_on_perfect_separation(): assert 0.0 <= lo <= hi <= 1.0 +def test_auc_permutation_matches_sklearn_exactly_with_ties(): + """The Mann–Whitney rank-sum identity must reproduce ``roc_auc_score`` exactly. + + Guards the O(n log n) -> O(n_pos) optimization of the permutation loop. Ties are + the failure mode that matters: the identity is only equivalent under *midranks*, + so scores are rounded here to force repeated values. + """ + import itertools + + from sklearn.metrics import roc_auc_score + + rng = np.random.default_rng(0) + for _ in range(40): + n = int(rng.integers(4, 11)) + scores = np.round(rng.normal(size=n), int(rng.integers(0, 2))) # forces ties + labels = rng.integers(0, 2, n) + if labels.sum() in (0, n): + continue + auc, _, perm, mode = S.auc_and_label_permutation(scores, labels) + assert mode == "exact" + y = (labels == 1).astype(int) + assert auc == pytest.approx(roc_auc_score(y, scores), abs=1e-12) + ref = sorted( + roc_auc_score(np.isin(np.arange(n), idx).astype(int), scores) + for idx in itertools.combinations(range(n), int(y.sum())) + ) + np.testing.assert_allclose(sorted(perm), ref, atol=1e-12) + + +def test_auc_permutation_degenerate_single_class(): + """A single-class label vector has no defined AUROC — report, don't crash.""" + auc, p, perm, mode = S.auc_and_label_permutation( + np.array([0.1, 0.2, 0.3]), np.array([1, 1, 1]) + ) + assert mode == "degenerate" and np.isnan(p) and perm.size == 0 + + def test_distance_kernels(): assert abs(D.kl_divergence([1, 0], [1, 0])) < 1e-9 # KL(p‖p)=0 assert D.kl_divergence([0.9, 0.1], [0.1, 0.9]) > 0 # asymmetric, positive diff --git a/tests/test_metrics/test_metrics.py b/tests/test_metrics/test_metrics.py deleted file mode 100644 index 03792f6..0000000 --- a/tests/test_metrics/test_metrics.py +++ /dev/null @@ -1,88 +0,0 @@ -import numpy as np -import pandas as pd -import pytest - -from tcri.metrics._metrics import ( - clonotypic_entropy, - phenotypic_entropy, - mutual_information, - clonality, -) - - -def test_clonotypic_entropy(trained_model): - _, adata = trained_model - covariate = adata.uns["tcri_covariate_categories"][0] - phenotypes = adata.uns["tcri_phenotype_categories"] - - s = clonotypic_entropy(adata, covariate, n_samples=20) - assert isinstance(s, pd.Series) - assert list(s.index) == list(phenotypes) - vals = s.to_numpy() - assert np.all(np.isfinite(vals)) - assert np.all(vals >= 0) - - arr = clonotypic_entropy( - adata, covariate, point_estimate=False, n_samples=5 - ) - assert isinstance(arr, np.ndarray) - assert arr.shape == (5, len(phenotypes)) - assert np.all(np.isfinite(arr)) - assert np.all(arr >= 0) - - with pytest.raises(ValueError): - clonotypic_entropy(adata, covariate, n_samples=0) - - -def test_phenotypic_entropy(trained_model): - _, adata = trained_model - covariate = adata.uns["tcri_covariate_categories"][0] - - meta = adata.uns["tcri_metadata"] - expected_clones = ( - adata.obs.loc[adata.obs[meta["covariate_col"]] == covariate, meta["clone_col"]] - .unique() - .tolist() - ) - - s = phenotypic_entropy(adata, covariate, n_samples=20) - assert isinstance(s, pd.Series) - assert set(s.index) == set(expected_clones) - vals = s.to_numpy() - assert np.all(np.isfinite(vals)) - assert np.all(vals >= 0) - - arr = phenotypic_entropy( - adata, covariate, point_estimate=False, n_samples=5 - ) - assert isinstance(arr, np.ndarray) - assert arr.shape == (5, len(expected_clones)) - assert np.all(np.isfinite(arr)) - assert np.all(arr >= 0) - - with pytest.raises(ValueError): - phenotypic_entropy(adata, covariate, n_samples=0) - - -def test_mutual_information(trained_model): - _, adata = trained_model - covariate = adata.uns["tcri_covariate_categories"][0] - - mi = mutual_information(adata, covariate=covariate, temperature=1.0, verbose=False) - assert np.isfinite(mi) - assert mi >= 0 - - mi_alt = mutual_information(adata, covariate=covariate, temperature=0.5, verbose=False) - assert np.isfinite(mi_alt) - assert mi_alt >= 0 - assert mi_alt != mi - - -def test_clonality(mock_adata): - """Test clonality function.""" - clonality_dict = clonality(mock_adata) - - for val in clonality_dict.values(): - assert 0 <= val <= 1 - - assert set(clonality_dict.keys()) == set(mock_adata.uns["tcri_phenotype_categories"]) diff --git a/tests/test_metrics_contract_conformance.py b/tests/test_metrics_contract_conformance.py new file mode 100644 index 0000000..52d2f38 --- /dev/null +++ b/tests/test_metrics_contract_conformance.py @@ -0,0 +1,178 @@ +"""Conformance test for the METRICS contract (``tcri/tools/_metrics_contract.py``). + +Companion to ``test_model_contract_conformance.py``. Where the model contract is +verified by tracing ``model()``/``guide()``, metrics are pure functions of a joint +table, so they are pinned by **numeric identities**: uniform -> log2(k), +independent -> MI 0, and the entropy/MI decomposition. + +A failure here means the *meaning* of a published number changed. Update the manifest +and ``docs/contract/METRICS_CONTRACT.md`` first, deliberately — never relax an +identity to make this pass. +""" +from __future__ import annotations + +import numpy as np +import pytest + +from tcri.tools import _metrics_contract as MC +from tcri.tools._entropy import _clonotypic_one, _phenotypic_one +from tcri.tools._mutual_information import _mi_from_joint + + +# ── manifest hygiene ──────────────────────────────────────────────────────── +def test_manifest_is_complete(): + """Every public metric is specified, and every spec field is filled in.""" + assert set(MC.METRIC_SPECS) == { + "clonotypic_entropy", "phenotypic_entropy", "mutual_information" + } + for name, spec in MC.METRIC_SPECS.items(): + for field in ("formula", "per", "support", "normalizer", "empty", "note_eq"): + assert getattr(spec, field), f"{name}.{field} is empty" + assert MC.LOG_BASE == 2 + + +def test_source_errata_are_documented(): + """The note's eq 3/4 errors must stay recorded, with the justification.""" + for key in ("eq3_weights_marginal", "eq4_label_and_weights", "why_the_code_is_right"): + assert key in MC.SOURCE_ERRATA and len(MC.SOURCE_ERRATA[key]) > 40 + + +# ── entropy identities ────────────────────────────────────────────────────── +def test_entropy_uniform_is_log2_k(): + """IDENTITIES['entropy_uniform_is_log2_k'].""" + for k in (2, 4, 8): + J = np.ones((k, k)) + # clonotypic: each phenotype column is uniform over k clones + raw = _clonotypic_one(J, [f"p{j}" for j in range(k)], normalized=False) + for v in raw.values(): + assert v == pytest.approx(np.log2(k)) + nrm = _clonotypic_one(J, [f"p{j}" for j in range(k)], normalized=True) + for v in nrm.values(): + assert v == pytest.approx(1.0) + # phenotypic: each clone row is uniform over k phenotypes + rawp = _phenotypic_one([f"c{i}" for i in range(k)], J, + [f"p{j}" for j in range(k)], normalized=False) + for v in rawp.values(): + assert v == pytest.approx(np.log2(k)) + + +def test_entropy_degenerate_is_zero(): + """IDENTITIES['entropy_degenerate_is_zero'].""" + J = np.array([[5.0, 0.0], [0.0, 3.0]]) + c = _clonotypic_one(J, ["a", "b"], normalized=False) + assert c["a"] == pytest.approx(0.0) and c["b"] == pytest.approx(0.0) + p = _phenotypic_one(["c0", "c1"], J, ["a", "b"], normalized=False) + assert p["c0"] == pytest.approx(0.0) and p["c1"] == pytest.approx(0.0) + + +def test_entropy_zero_mass_is_nan_not_zero_or_one(): + """IDENTITIES['entropy_zero_mass_is_nan'] — the spurious-H=1 regression.""" + J = np.array([[1.0, 0.0], [1.0, 0.0]]) # phenotype 'b' has no mass + c = _clonotypic_one(J, ["a", "b"], normalized=True) + assert np.isnan(c["b"]), "empty phenotype column must be NaN" + J2 = np.array([[1.0, 1.0], [0.0, 0.0]]) # clone 'c1' has no mass + p = _phenotypic_one(["c0", "c1"], J2, ["a", "b"], normalized=True) + assert np.isnan(p["c1"]), "zero-mass clone must be NaN, never a spurious 1.0" + + +def test_clonotypic_entropy_is_support_only(): + """Absent clones are dropped BEFORE normalizing (no epsilon-clip inflation). + + Two supported clones out of five: H must be log2(2), normalized 1.0 — not the + log2(5)-normalized value an epsilon clip would produce. + """ + J = np.zeros((5, 1)) + J[0, 0] = J[1, 0] = 1.0 + raw = _clonotypic_one(J, ["a"], normalized=False)["a"] + nrm = _clonotypic_one(J, ["a"], normalized=True)["a"] + assert raw == pytest.approx(np.log2(2)) + assert nrm == pytest.approx(1.0) + # n_clones_ref fixes the normalizer instead, for cross-group comparability + fixed = _clonotypic_one(J, ["a"], normalized=True, n_clones_ref=5)["a"] + assert fixed == pytest.approx(np.log2(2) / np.log2(5)) + + +# ── mutual information identities ─────────────────────────────────────────── +def test_mi_independent_is_zero(): + """IDENTITIES['mi_independent_is_zero'].""" + px = np.array([0.2, 0.3, 0.5])[:, None] + py = np.array([0.4, 0.6])[None, :] + J = px @ py # exactly independent + assert _mi_from_joint(J, normalized=False) == pytest.approx(0.0, abs=1e-9) + + +def test_mi_is_symmetric_and_nonnegative(): + """IDENTITIES['mi_is_symmetric'] + ['mi_is_nonnegative'].""" + rng = np.random.default_rng(0) + for _ in range(25): + J = rng.random((rng.integers(2, 7), rng.integers(2, 7))) + mi = _mi_from_joint(J, normalized=False) + assert mi >= -1e-12 + assert mi == pytest.approx(_mi_from_joint(J.T, normalized=False)) + + +def test_mi_perfect_coupling_is_one(): + """IDENTITIES['mi_perfect_coupling_is_one'] (normalize_mode='min').""" + J = np.eye(4) * 7.0 + assert _mi_from_joint(J, normalized=True, mode="min") == pytest.approx(1.0) + + +def test_mi_normalize_modes_differ_as_specified(): + """'min' vs 'average' denominators — 'min' is the default for comparability.""" + J = np.array([[8.0, 1.0, 1.0], [1.0, 5.0, 1.0]]) + mi = _mi_from_joint(J, normalized=False) + pxy = J / J.sum() + px, py = pxy.sum(1), pxy.sum(0) + h_c = -np.sum(px * np.log2(px)) + h_p = -np.sum(py * np.log2(py)) + assert _mi_from_joint(J, normalized=True, mode="min") == pytest.approx(mi / min(h_c, h_p), rel=1e-6) + assert _mi_from_joint(J, normalized=True, mode="average") == pytest.approx( + mi / (0.5 * (h_c + h_p)), rel=1e-6) + + +# ── the cross-metric identity that caught the note's erratum ──────────────── +def test_mi_equals_marginal_minus_expected_conditional_entropy(): + """IDENTITIES['mi_entropy_decomposition'] — I(c;phi) = H(c) - E_phi[H(c|phi)]. + + This is the identity that proves the implemented conditional entropy is the right + one: weighting by the marginal (as the note's eqs 3-4 literally read) makes this + yield a NEGATIVE mutual information. It ties the entropy and MI families together, + so redefining either alone breaks it. + """ + rng = np.random.default_rng(7) + for _ in range(20): + n_c, n_p = int(rng.integers(2, 8)), int(rng.integers(2, 6)) + J = rng.random((n_c, n_p)) + 0.05 # strictly positive: full support + cols = [f"p{j}" for j in range(n_p)] + + pxy = J / J.sum() + p_c, p_ph = pxy.sum(1), pxy.sum(0) + H_c = -np.sum(p_c * np.log2(p_c)) + + # UNNORMALIZED clonotypic entropy per phenotype, weighted by P(phi) + H_cond = _clonotypic_one(J, cols, normalized=False) + expected_cond = sum(p_ph[j] * H_cond[cols[j]] for j in range(n_p)) + + mi = _mi_from_joint(J, normalized=False) + assert mi == pytest.approx(H_c - expected_cond, abs=1e-9) + + +def test_note_literal_formula_would_break_the_decomposition(): + """Guards the erratum itself: the note's literal eq 3 gives a NEGATIVE MI. + + If someone 'fixes' the code to match the mistranscribed equation, this test + documents exactly why that is wrong. + """ + J = np.array([[4.0, 1.0], [1.0, 1.0], [1.0, 6.0]]) + pxy = J / J.sum() + p_c, p_ph = pxy.sum(1), pxy.sum(0) + H_c = -np.sum(p_c * np.log2(p_c)) + mi = _mi_from_joint(J, normalized=False) + + # the note as literally written: weight by the MARGINAL p(c) + literal = sum( + p_ph[j] * (-np.sum(p_c * np.log2(J[:, j] / J[:, j].sum()))) + for j in range(J.shape[1]) + ) + assert H_c - literal < 0, "expected the literal formula to give a negative MI" + assert mi > 0 diff --git a/tests/test_model_classifier.py b/tests/test_model_classifier.py new file mode 100644 index 0000000..cd9fe25 --- /dev/null +++ b/tests/test_model_classifier.py @@ -0,0 +1,130 @@ +"""Phenotype-classifier recovery test. + +Locks in the two coupled fixes that make ``f_cls`` (the phenotype classifier head, +methods eq. 4) actually train: + + 1. The classifier enters the ELBO through ``pyro.factor("phenotype_alignment", ...)`` + in ``TCRIModule.model()`` (the surrogate KL objective, methods "Inference + Details"). Without it ``cls_logits`` never touches the log-joint and f_cls + gets no gradient (weight change == 0). + 2. The per-cell alignment target ``phi = p_ct[ct_idx]`` is indexed with the + GLOBAL cell indices, not the local pyro data-plate index. Indexing with the + local index silently scrambles each cell's target across shuffled minibatches, + which trains the classifier on the wrong labels and collapses it to a constant + (recovery == chance). + +The dataset is "perfect": each clonotype expresses one unique marker gene and maps +to one phenotype, so a correctly-trained classifier recovers the phenotype from gene +expression alone (gate_prob=1.0, the pure-classifier path). +""" +import contextlib +import io + +import numpy as np +import pandas as pd +import pyro +import pytest +import torch +from anndata import AnnData + +from tcri.model._model import TCRIModel + + +@pytest.fixture(autouse=True) +def _isolate_param_store(): + """Own the process-global Pyro param store for this module's tests. + + Each test here trains a small model, leaving shape-specific params + (``q_p_c_raw`` etc.) in the global store; without clearing on teardown the + next model test in the suite reuses a stale-shaped param and crashes. Scoped + to THIS module (not a conftest autouse) so it never wipes the session-scoped + ``trained_model`` fixture's params mid-suite. + """ + pyro.clear_param_store() + yield + pyro.clear_param_store() + + +def _perfect_adata(n_clones=5, n_per=60, n_genes=6, seed=0): + """One marker gene per clone, one phenotype per clone (n_conditions=1).""" + rng = np.random.default_rng(seed) + rows, clone, phen = [], [], [] + for c in range(n_clones): + for _ in range(n_per): + v = rng.poisson(0.2, size=n_genes).astype("float32") + v[c % n_genes] = 100.0 # unique hot gene marks the clone + rows.append(v) + clone.append(f"clone_{c}") + phen.append(f"phen_{c}") + obs = pd.DataFrame( + { + "clone_id": clone, + "true_phenotype": phen, + "covariate": "c0", + "patient": "P1", + } + ) + ad = AnnData( + X=np.asarray(rows), + obs=obs, + var=pd.DataFrame(index=[f"g{g}" for g in range(n_genes)]), + ) + ad.layers["counts"] = ad.X.copy() + return ad + + +@pytest.mark.parametrize("gate_prob", [1.0, 0.5]) +def test_classifier_perfect_recovery(gate_prob): + """f_cls recovers the phenotype on a linearly-separable dataset. + + gate_prob=1.0 is the strict test (phenotype comes from the classifier alone); + gate_prob=0.5 is the methods default (classifier + clonotype prior). + """ + np.random.seed(0) + torch.manual_seed(0) + pyro.set_rng_seed(0) # param store cleared by the autouse fixture + import scvi + + scvi.settings.seed = 0 + + ad = _perfect_adata() + truth = ad.obs["true_phenotype"].to_numpy() + + TCRIModel.setup_anndata( + ad, + layer="counts", + clonotype_key="clone_id", + phenotype_key="true_phenotype", + covariate_key="covariate", + batch_key="patient", + ) + model = TCRIModel( + ad, + n_latent=8, + n_hidden=16, + n_layers=1, + classifier_n_layers=1, + classifier_hidden=16, + K=5, + n_pseudo_obs=3, + gate_prob=gate_prob, + ) + + # snapshot classifier weights to prove the head actually trains + w0 = {k: v.detach().clone() for k, v in model.module.classifier.state_dict().items()} + + with contextlib.redirect_stdout(io.StringIO()): + model.train( + max_epochs=200, + batch_size=128, + enable_progress_bar=False, + enable_model_summary=False, + ) + + w1 = model.module.classifier.state_dict() + dw = sum((w1[k] - w0[k]).pow(2).sum().item() for k in w0) ** 0.5 + assert dw > 1e-3, f"classifier head did not train (ΔL2={dw:.2e}); is it in the ELBO?" + + pred = model.predict(ad) + recovery = (pred.columns[pred.values.argmax(1)] == truth).mean() + assert recovery >= 0.9, f"phenotype recovery {recovery:.3f} (gate={gate_prob}); chance=0.2" diff --git a/tests/test_model_contract_conformance.py b/tests/test_model_contract_conformance.py new file mode 100644 index 0000000..d53bf36 --- /dev/null +++ b/tests/test_model_contract_conformance.py @@ -0,0 +1,403 @@ +"""Model-contract conformance — the guardrail for the model's *mathematics*. + +``tcri/model/_model_contract.py`` freezes the probabilistic structure of +Supplementary Note 1 (prose: ``docs/contract/MODEL_CONTRACT.md``). This test traces +the live ``TCRIModule.model``/``.guide`` and asserts they match that manifest +exactly — declared sites present with the right distribution family / plate / +event-dim, **no undeclared sites**, the guide's variational family intact, and the +semantic invariants (α and β scaling, the surrogate's sign, the gate rule, global +alignment indices) holding. + +If you changed the model and this test fails: **update the contract first** +(``_model_contract.py`` + ``MODEL_CONTRACT.md``, citing the note), then make the +code agree. Do not loosen the manifest to match new code — that silently rewrites +the model the package claims to implement. + +Sibling of ``test_contract_conformance.py`` (which does the same for the public API). +""" +import contextlib +import io + +import numpy as np +import pandas as pd +import pyro +import pyro.poutine as poutine +import pytest +import torch +from anndata import AnnData + +from tcri.model._model import TCRIModel +from tcri.model import _model_contract as MC + +# pyro wrappers that carry no model semantics — unwrapped before comparing. +_WRAPPERS = ("Independent", "ExpandedDistribution", "MaskedDistribution") + + +@pytest.fixture(scope="module") +def traced(): + """Trace model() and guide() once on a tiny fixture (no training needed).""" + pyro.clear_param_store() + np.random.seed(0) + torch.manual_seed(0) + pyro.set_rng_seed(0) + + rng = np.random.default_rng(0) + n_cells, n_genes = 60, 6 + X = rng.poisson(1.0, size=(n_cells, n_genes)).astype("float32") + obs = pd.DataFrame( + { + "clone": [f"c{i % 5}" for i in range(n_cells)], + "phen": [f"p{i % 3}" for i in range(n_cells)], + "cov": ["a", "b"] * (n_cells // 2), + "pt": ["P1"] * n_cells, + } + ) + ad = AnnData(X=X, obs=obs, var=pd.DataFrame(index=[f"g{i}" for i in range(n_genes)])) + ad.layers["counts"] = ad.X.copy() + + TCRIModel.setup_anndata( + ad, layer="counts", clonotype_key="clone", phenotype_key="phen", + covariate_key="cov", batch_key="pt", + ) + model = TCRIModel( + ad, n_latent=4, n_hidden=8, n_layers=1, classifier_n_layers=1, + classifier_hidden=8, K=3, n_pseudo_obs=2, + ) + + loader = model._make_data_loader(adata=model.adata, batch_size=32, shuffle=False) + batch = next(iter(loader)) + args, kwargs = model.module._get_fn_args_from_batch(batch) + + with contextlib.redirect_stdout(io.StringIO()): + m_trace = poutine.trace(model.module.model).get_trace(*args, **kwargs) + g_trace = poutine.trace(model.module.guide).get_trace(*args, **kwargs) + + yield model, m_trace, g_trace, args, kwargs + pyro.clear_param_store() + + +def _unwrap(d): + """Strip Independent/Expanded wrappers to the semantic distribution class.""" + seen = 0 + while type(d).__name__ in _WRAPPERS and seen < 8: + d = getattr(d, "base_dist", None) or getattr(d, "base_distribution", None) + if d is None: + return None + seen += 1 + return d + + +def _stochastic(trace): + """{name: node} for real stochastic sites (drops plate _Subsample bookkeeping).""" + out = {} + for name, node in trace.nodes.items(): + if node["type"] != "sample": + continue + if type(node["fn"]).__name__ == "_Subsample": + continue + out[name] = node + return out + + +def _plates(node): + return {f.name for f in node.get("cond_indep_stack", ())} + + +# ── structure: the generative program ──────────────────────────────────────── + +def test_generative_sites_match_contract(traced): + _, m_trace, _, _, _ = traced + live = _stochastic(m_trace) + + for spec in MC.GENERATIVE_SITES: + assert spec.name in live, ( + f"model() is missing the declared site '{spec.name}' (note eq {spec.eq}). " + "If you removed it, update tcri/model/_model_contract.py + " + "docs/contract/MODEL_CONTRACT.md first." + ) + node = live[spec.name] + base = _unwrap(node["fn"]) + assert base is not None and type(base).__name__ == spec.dist, ( + f"site '{spec.name}' (eq {spec.eq}) is " + f"{type(base).__name__ if base is not None else None}, contract says " + f"{spec.dist}. {spec.note}" + ) + if spec.plate is not None: + assert spec.plate in _plates(node), ( + f"site '{spec.name}' must live in plate '{spec.plate}' " + f"({MC.PLATES.get(spec.plate, '')}); found {_plates(node)}." + ) + assert bool(node.get("is_observed")) == spec.observed, ( + f"site '{spec.name}' observed={bool(node.get('is_observed'))}, " + f"contract says observed={spec.observed}." + ) + if spec.event_dim is not None: + assert len(node["fn"].event_shape) == spec.event_dim, ( + f"site '{spec.name}' event_dim={len(node['fn'].event_shape)}, " + f"contract says {spec.event_dim}." + ) + + +def test_no_undeclared_generative_sites(traced): + """An EXTRA stochastic site changes the joint distribution — the contract must say so.""" + _, m_trace, _, _, _ = traced + live = set(_stochastic(m_trace)) + declared = {s.name for s in MC.GENERATIVE_SITES} + extra = live - declared + assert not extra, ( + f"model() has undeclared stochastic site(s): {sorted(extra)}. Every site " + "changes the joint p(Ω,Φ,z,x). Declare it in tcri/model/_model_contract.py " + "(with its note equation) and document it in docs/contract/MODEL_CONTRACT.md." + ) + + +# ── structure: the variational family ──────────────────────────────────────── + +def test_guide_family_matches_contract(traced): + _, _, g_trace, _, _ = traced + live = _stochastic(g_trace) + + for spec in MC.GUIDE_SITES: + assert spec.name in live, ( + f"guide() is missing q({spec.name}) (eq 6). {spec.note}" + ) + base = _unwrap(live[spec.name]["fn"]) + assert base is not None and type(base).__name__ == spec.dist, ( + f"q({spec.name}) is {type(base).__name__ if base is not None else None}, " + f"contract says {spec.dist}. {spec.note}" + ) + + declared = {s.name for s in MC.GUIDE_SITES} + extra = set(live) - declared + assert not extra, ( + f"guide() has undeclared site(s): {sorted(extra)}. This changes the " + "variational family (eq 6) and therefore the ELBO." + ) + + +def test_guide_registers_variational_params(traced): + _, _, g_trace, _, _ = traced + params = {n for n, nd in g_trace.nodes.items() if nd["type"] == "param"} + for p in MC.GUIDE_PARAMS: + assert p in params, ( + f"guide() must register the learnable variational parameter '{p}' " + "(λ_c / λ'_m of eq 6); without it the Dirichlet posteriors are not learned." + ) + + +def test_discrete_phenotype_latent_is_not_sampled(traced): + """z^ϕ is replaced by the surrogate; a q(z^ϕ) site would change the objective.""" + _, _, g_trace, _, _ = traced + live = set(_stochastic(g_trace)) + for forbidden in MC.FORBIDDEN_GUIDE_SITES: + assert forbidden not in live, ( + f"guide() samples '{forbidden}', but the note's Inference Details replace " + "the discrete z^ϕ with the phenotype_alignment surrogate. Re-introducing " + "it changes the optimized objective — update the contract first." + ) + + +# ── semantics: invariants the structure alone cannot pin ───────────────────── + +def test_alpha_scales_the_clonotype_prior(traced): + """eq 1: p_c prior concentration must scale with α (global_scale).""" + model, m_trace, _, args, kwargs = traced + conc = _unwrap(m_trace.nodes["p_c"]["fn"]).concentration + alpha = float(model.module.global_scale) + archetypes = model.module.mixture_concentration # rows sum to ~1 + expected_total = alpha * float(archetypes.sum(-1).mean()) + live_total = float(conc.sum(-1).mean()) + assert live_total == pytest.approx(expected_total, rel=1e-4), ( + f"p_c concentration totals {live_total:.4f}, expected ≈{expected_total:.4f} " + f"(α={alpha} × archetype). {MC.SEMANTIC_INVARIANTS['alpha_scales_clonotype_prior']}" + ) + + +def test_beta_scales_the_covariate_prior(traced): + """eq 2: p_ct concentration must be β·ω_h(m) — the SAMPLED ω under ct_to_c. + + Asserted as an elementwise identity against the same trace, which pins three + things at once: the scale (β), the source tensor (the sampled ``p_c``, not the + static empirical prior), and the index map (``ct_to_c`` = h(m)). A scalar + "totals ≈ β" check cannot do this — every simplex row totals 1, so any tensor + under any permutation would satisfy it while the hierarchy is severed. + """ + model, m_trace, _, _, _ = traced + mod = model.module + conc = _unwrap(m_trace.nodes["p_ct"]["fn"]).concentration + omega = m_trace.nodes["p_c"]["value"] # the sampled ω_c from THIS trace + beta = float(mod.local_scale) + + expected = torch.clamp(beta * (omega[mod.ct_to_c] + mod.eps), min=1e-3) + assert torch.allclose(conc, expected, rtol=1e-5, atol=1e-6), ( + "p_ct concentration is not β·ω_h(m) built from the sampled p_c under " + f"ct_to_c (max|diff|={float((conc - expected).abs().max()):.3e}). " + f"{MC.SEMANTIC_INVARIANTS['beta_scales_covariate_prior']} " + f"{MC.SEMANTIC_INVARIANTS['hierarchy_ct_depends_on_c']}" + ) + + +def test_alignment_factor_is_a_negative_kl(traced): + """Inference Details: the surrogate must ENTER as −γ·KL (a penalty), never +γ·KL.""" + _, m_trace, _, _, _ = traced + node = m_trace.nodes["phenotype_alignment"] + val = node["fn"].log_factor if hasattr(node["fn"], "log_factor") else node["value"] + val = torch.as_tensor(val) + assert torch.all(val <= 1e-6), ( + f"phenotype_alignment carries a positive log-factor (max={float(val.max()):.4e}). " + f"{MC.SEMANTIC_INVARIANTS['factor_is_negative_kl']}" + ) + # and it must be non-trivial (a zero factor trains nothing) + assert float(val.abs().sum()) > 0, ( + "phenotype_alignment is identically zero — f_cls would receive no gradient " + "(this is exactly the bug the surrogate exists to fix)." + ) + + +def test_alignment_target_uses_global_indices(traced): + """The ϕ target must be indexed by GLOBAL cell indices, not the local plate index. + + Checked *behaviorally*: trace a minibatch whose global indices differ from the + local plate positions (0..B−1), then recompute the surrogate from the global + map and compare to the traced factor. Under the local-index bug the two differ. + A source-text assertion cannot do this — it is defeated by any rename or by + routing the same wrong lookup through ``index_select``. + """ + model, _, _, _, _ = traced + mod = model.module + + # a batch whose global indices are NOT 0..B-1 (so local != global) + loader = model._make_data_loader(adata=model.adata, batch_size=16, shuffle=False) + batches = list(loader) + assert len(batches) >= 2, "need >1 batch for local-vs-global to differ" + args, kwargs = mod._get_fn_args_from_batch(batches[1]) + global_idx = args[3] + assert not torch.equal( + global_idx, torch.arange(global_idx.numel(), device=global_idx.device) + ), "fixture batch must have global indices != local positions" + + # eval mode: classifier dropout is stochastic in train mode, which would make + # the recomputation below irreproducible. Restored afterwards. + was_training = mod.training + mod.eval() + try: + with contextlib.redirect_stdout(io.StringIO()): + tr = poutine.trace(mod.model).get_trace(*args, **kwargs) + finally: + if was_training: + mod.train() + + z = tr.nodes["latent"]["value"] + p_ct = tr.nodes["p_ct"]["value"] + live = torch.as_tensor(tr.nodes["phenotype_alignment"]["fn"].log_factor).detach() + + def _surrogate(ct_index): + phi = p_ct[ct_index].detach() + log_phi = torch.log(phi + 1e-8) + with torch.no_grad(): + mod.eval() + logits = mod.classifier(z) + if was_training: + mod.train() + ell = ( + mod.gate_prob * logits + (1.0 - mod.gate_prob) * log_phi + if mod.gate_prob is not None + else logits + log_phi + ) + probs = torch.softmax(ell, dim=-1) + kl = (probs * (torch.log(probs + 1e-8) - log_phi)).sum(-1) + return (-mod.phenotype_kl_weight * kl).detach() + + expected_global = _surrogate(mod.ct_array[global_idx]) + local_idx = torch.arange(global_idx.numel(), device=global_idx.device) + expected_local = _surrogate(mod.ct_array[local_idx]) + + # the local-index variant must be a genuinely different target, or this + # fixture cannot discriminate and the test would be vacuous + assert not torch.allclose(expected_global, expected_local, rtol=1e-4, atol=1e-6), ( + "fixture cannot distinguish global from local indexing — strengthen it." + ) + assert torch.allclose(live, expected_global, rtol=1e-4, atol=1e-5), ( + "the phenotype_alignment target does not match the GLOBAL-index mapping " + f"(max|diff| vs global={float((live - expected_global).abs().max()):.3e}, " + f"vs local={float((live - expected_local).abs().max()):.3e}). " + f"{MC.SEMANTIC_INVARIANTS['alignment_target_uses_global_indices']}" + ) + + +@pytest.mark.parametrize( + "gate,expect", [(1.0, "classifier"), (0.0, "prior")] +) +def test_gate_rule_endpoints(gate, expect): + """eq 4: π=1 ⇒ predict is the pure classifier; π=0 ⇒ the pure clonotype prior.""" + import torch.nn.functional as F + + pyro.clear_param_store() + np.random.seed(0) + torch.manual_seed(0) + pyro.set_rng_seed(0) + + rng = np.random.default_rng(0) + n_cells, n_genes = 40, 5 + X = rng.poisson(1.0, size=(n_cells, n_genes)).astype("float32") + obs = pd.DataFrame( + { + "clone": [f"c{i % 4}" for i in range(n_cells)], + "phen": [f"p{i % 3}" for i in range(n_cells)], + "cov": ["a"] * n_cells, + "pt": ["P1"] * n_cells, + } + ) + ad = AnnData(X=X, obs=obs, var=pd.DataFrame(index=[f"g{i}" for i in range(n_genes)])) + ad.layers["counts"] = ad.X.copy() + TCRIModel.setup_anndata( + ad, layer="counts", clonotype_key="clone", phenotype_key="phen", + covariate_key="cov", batch_key="pt", + ) + m = TCRIModel( + ad, n_latent=4, n_hidden=8, n_layers=1, classifier_n_layers=1, + classifier_hidden=8, K=3, n_pseudo_obs=2, gate_prob=gate, + ) + # run the guide once so q_p_ct_raw exists in the global param store (get_p_ct + # reads it); no training needed — this is a formula identity, not a fit. + loader = m._make_data_loader(adata=m.adata, batch_size=32, shuffle=False) + g_args, g_kwargs = m.module._get_fn_args_from_batch(next(iter(loader))) + with contextlib.redirect_stdout(io.StringIO()): + poutine.trace(m.module.guide).get_trace(*g_args, **g_kwargs) + m.module.eval() + + probs = m.predict(ad).values + x = torch.tensor(ad.layers["counts"]) + b = torch.zeros(x.shape[0], 1) + with torch.no_grad(): + z_loc, _, _ = m.module.encoder(x, b) + cls = F.softmax(m.module.classifier(z_loc), dim=-1).numpy() + p_ct = m.module.get_p_ct() + prior = F.softmax( + torch.log(p_ct[m.module.ct_array] + 1e-8), dim=-1 + ).numpy() + + target = cls if expect == "classifier" else prior + np.testing.assert_allclose(probs, target, atol=1e-4, err_msg=( + f"gate_prob={gate} must reduce predict() to the pure {expect}. " + f"{MC.SEMANTIC_INVARIANTS['gate_mixes_classifier_and_prior']}" + )) + pyro.clear_param_store() + + +# ── the contract must stay in sync with its prose ──────────────────────────── + +def test_sanctioned_deviations_are_documented(): + """Every accepted departure from the note must be spelled out in the prose contract.""" + from pathlib import Path + + import tcri + + doc = Path(tcri.__file__).parent.parent / "docs" / "contract" / "MODEL_CONTRACT.md" + assert doc.exists(), f"missing prose model contract: {doc}" + text = doc.read_text() + for key in MC.SANCTIONED_DEVIATIONS: + assert key in text, ( + f"sanctioned deviation '{key}' is declared in _model_contract.py but not " + f"documented in {doc.name}. Keep the machine contract and the prose in sync." + ) diff --git a/tests/test_model_guardrails.py b/tests/test_model_guardrails.py new file mode 100644 index 0000000..12d290c --- /dev/null +++ b/tests/test_model_guardrails.py @@ -0,0 +1,141 @@ +"""Guardrails on TCRIModel construction/training defaults. + +These pin fixes for footguns found in the performance audit: a hard crash on small +datasets, silent cross-model contamination via Pyro's process-global param store, +a pathological batch size, and Trainer knobs the caller could not override. +""" +import contextlib +import io + +import numpy as np +import pandas as pd +import pyro +import pytest +from anndata import AnnData + +from tcri.model._model import TCRIModel + + +@pytest.fixture(autouse=True) +def _isolate_param_store(): + """Own the process-global Pyro param store for this module. + + These tests train small models, leaving shape-specific params (``q_p_c_raw`` + etc.) behind; without clearing on teardown the next model test in the suite + reuses a stale-shaped param and fails. Module-local (not a conftest autouse) so + it never wipes the session-scoped ``trained_model`` fixture mid-suite. + """ + pyro.clear_param_store() + yield + pyro.clear_param_store() + + +@pytest.fixture +def tiny_adata(): + """5 clonotypes — fewer than the default K=10.""" + n_clones, n_per, n_genes = 5, 40, 5 + rows, clone, phen, cov = [], [], [], [] + for c in range(n_clones): + for j in range(2): + for _ in range(n_per): + v = np.zeros(n_genes, dtype="float32") + v[c % n_genes] = 100.0 + rows.append(v) + clone.append(f"clone_{c}") + phen.append(f"phen_{c}") + cov.append(f"cond_{j}") + ad = AnnData( + X=np.asarray(rows), + obs=pd.DataFrame({"clone_id": clone, "true_phenotype": phen, + "covariate": cov, "patient": "P1"}), + var=pd.DataFrame(index=[f"g{g}" for g in range(n_genes)]), + ) + ad.layers["counts"] = ad.X.copy() + TCRIModel.setup_anndata( + ad, layer="counts", clonotype_key="clone_id", phenotype_key="true_phenotype", + covariate_key="covariate", batch_key="patient", + ) + return ad + + +def _model(adata, **kw): + kw.setdefault("n_latent", 8) + kw.setdefault("n_hidden", 16) + kw.setdefault("n_layers", 1) + kw.setdefault("classifier_n_layers", 1) + kw.setdefault("classifier_hidden", 16) + return TCRIModel(adata, **kw) + + +def test_K_clamped_to_n_clonotypes(tiny_adata): + """K > n_clonotypes used to raise from sklearn KMeans; it now clamps + warns.""" + pyro.clear_param_store() + with pytest.warns(UserWarning, match="archetype"): + model = _model(tiny_adata) # K=10 default, only 5 clones + assert model.centers.shape[0] == 5 + + +def test_second_model_warns_about_shared_param_store(tiny_adata): + """Pyro's param store is process-global — a 2nd model silently continues the 1st fit.""" + pyro.clear_param_store() + m1 = _model(tiny_adata, K=5) + with contextlib.redirect_stdout(io.StringIO()): + m1.train(max_epochs=5, batch_size=256, + enable_progress_bar=False, enable_model_summary=False) + with pytest.warns(UserWarning, match="param store"): + _model(tiny_adata, K=5) + + +def test_batch_size_at_or_above_n_obs_warns(tiny_adata): + """batch_size >= n_obs => 1 optimizer step/epoch (the 9-hour-notebook pathology).""" + pyro.clear_param_store() + model = _model(tiny_adata, K=5) + with pytest.warns(UserWarning, match="SINGLE"): + with contextlib.redirect_stdout(io.StringIO()): + model.train(max_epochs=2, batch_size=10_000, + enable_progress_bar=False, enable_model_summary=False) + + +def test_lr_and_weight_decay_reach_pyros_optimizer(tiny_adata): + """The optimizer settings must configure SVI, not a side optimizer. + + ``UnifiedTrainingPlan`` used to override ``configure_optimizers`` with a real + torch Adam over every module parameter. That replaced scvi's deliberate no-op + shim and ran *after* ``SVI.step()`` had already zeroed the gradients, so it + only ever applied a scale-free ~lr*sign(p) shrink — and ``lr`` never reached + the optimizer that actually descends the ELBO (Pyro stayed at scvi's 1e-3). + """ + from tcri.model._training import UnifiedTrainingPlan + + model = _model(tiny_adata, K=5) + plan = UnifiedTrainingPlan( + module=model.module, n_steps_kl_warmup=10, reconstruction_loss_scale=1e-3, + optimizer_config={"lr": 0.05, "betas": (0.9, 0.999), "eps": 1e-5, + "weight_decay": 1e-4}, + ) + args = plan.optim.pt_optim_args + assert args["lr"] == 0.05, f"lr did not reach Pyro's SVI optimizer: {args}" + assert args["weight_decay"] == 1e-4, f"weight_decay did not reach Pyro: {args}" + + # and the Lightning-facing optimizer must be scvi's dummy shim, not the module + opt = plan.configure_optimizers() + opt = opt["optimizer"] if isinstance(opt, dict) else opt + n_opt = sum(p.numel() for g in opt.param_groups for p in g["params"]) + n_module = sum(p.numel() for p in model.module.parameters()) + assert n_opt == 1 < n_module, ( + f"configure_optimizers covers {n_opt} params (module has {n_module}); it must " + "stay scvi's single-dummy-param shim so nothing steps on zeroed gradients." + ) + + +def test_trainer_knobs_are_overridable(tiny_adata): + """These were hard-coded keywords: passing them raised 'got multiple values'.""" + pyro.clear_param_store() + model = _model(tiny_adata, K=5) + with contextlib.redirect_stdout(io.StringIO()): + model.train( # would previously raise TypeError + max_epochs=6, batch_size=256, + early_stopping_patience=1, check_val_every_n_epoch=1, + enable_progress_bar=False, enable_model_summary=False, + ) + assert model.trainer.current_epoch <= 6 diff --git a/tests/test_model_knobs.py b/tests/test_model_knobs.py new file mode 100644 index 0000000..2b7224c --- /dev/null +++ b/tests/test_model_knobs.py @@ -0,0 +1,327 @@ +"""Knob-test matrix — every constructor/train knob gets a correctness test. + +Two layers, and the split matters: + +**WIRING** — does the value actually reach the object it claims to configure? +This layer exists because ``lr`` was marked "hooked up" in the matrix for months +while never reaching Pyro's optimizer: the model still converged, so every +behavioral test passed. A knob that is silently ignored is invisible to +convergence tests. Assert the plumbing directly. + +**BEHAVIOR** — the mathematically-correct input->output assertion (draw variance +scales as 1/(scale+1), temperature sharpens, gate endpoints reduce to closed +forms, batch size is an invariance). +""" +from __future__ import annotations + +import contextlib +import io + +import numpy as np +import pandas as pd +import pyro +import pytest +import torch +from anndata import AnnData + +from tcri.model._model import TCRIModel +from tcri.model._training import UnifiedTrainingPlan + + +@pytest.fixture(autouse=True) +def _isolate_param_store(): + pyro.clear_param_store() + yield + pyro.clear_param_store() + + +@pytest.fixture +def adata(): + n_clones, n_per, n_genes = 6, 25, 8 + rng = np.random.default_rng(0) + rows, clone, phen, cov = [], [], [], [] + for c in range(n_clones): + for j in range(2): + for _ in range(n_per): + v = rng.poisson(0.3, size=n_genes).astype("float32") + v[c % n_genes] = 60.0 + rows.append(v) + clone.append(f"clone_{c}") + phen.append(f"phen_{c % 4}") + cov.append(f"cond_{j}") + ad = AnnData( + X=np.asarray(rows), + obs=pd.DataFrame({"clone_id": clone, "true_phenotype": phen, + "covariate": cov, "patient": "P1"}), + var=pd.DataFrame(index=[f"g{g}" for g in range(n_genes)]), + ) + ad.layers["counts"] = ad.X.copy() + TCRIModel.setup_anndata( + ad, layer="counts", clonotype_key="clone_id", phenotype_key="true_phenotype", + covariate_key="covariate", batch_key="patient", + ) + return ad + + +def _model(ad, **kw): + kw.setdefault("n_latent", 8) + kw.setdefault("n_hidden", 16) + kw.setdefault("n_layers", 1) + kw.setdefault("classifier_n_layers", 1) + kw.setdefault("classifier_hidden", 16) + kw.setdefault("K", 3) + return TCRIModel(ad, **kw) + + +def _train(model, **kw): + kw.setdefault("max_epochs", 3) + kw.setdefault("batch_size", 128) + with contextlib.redirect_stdout(io.StringIO()): + model.train(enable_progress_bar=False, enable_model_summary=False, **kw) + + +# ══════════════════════════ WIRING ══════════════════════════════════════════ +# "does the value reach the thing it configures?" + +@pytest.mark.parametrize("knob,value,reader", [ + ("n_latent", 6, lambda m: m.module.n_latent), + ("n_pseudo_obs", 4, lambda m: m.module.vamp_prior.pseudo_inputs.shape[0]), + ("K", 3, lambda m: m.module.mixture_concentration.shape[0]), + ("global_scale", 7.5, lambda m: m.module.global_scale), + ("local_scale", 2.5, lambda m: m.module.local_scale), + ("prior_temperature", 1.5, lambda m: m.module.prior_temperature), + ("guide_temperature", 0.5, lambda m: m.module.guide_temperature), + ("gate_prob", 0.25, lambda m: m.module.gate_prob), + ("classifier_temperature", 2.0, lambda m: m.module.classifier_temperature), + ("classifier_dropout", 0.3, lambda m: m.module.classifier.mlp[2].p), + ("classifier_hidden", 12, lambda m: m.module.classifier.mlp[0].out_features), + ("kl_weight_max", 0.7, lambda m: m.module.kl_weight_max), + ("guide_init_scale", 4.0, lambda m: m.module.guide_init_scale), + ("phenotype_kl_weight", 3.0, lambda m: m.module.phenotype_kl_weight), +]) +def test_constructor_knob_is_wired(adata, knob, value, reader): + """Each constructor knob must be readable back off the module.""" + m = _model(adata, **{knob: value}) + assert reader(m) == pytest.approx(value), f"{knob}={value} did not reach the module" + + +def test_n_latent_determines_latent_width(adata): + m = _model(adata, n_latent=6) + _train(m) + assert m.get_latent_representation().shape[1] == 6 + + +def test_K_determines_archetype_count(adata): + m = _model(adata, K=3) + assert m.centers.shape[0] == 3 + assert m.module.mixture_concentration.shape[0] == 3 + + +def test_classifier_depth_is_wired(adata): + """classifier_n_layers controls the number of Linear blocks.""" + shallow = _model(adata, classifier_n_layers=1) + deep = _model(adata, classifier_n_layers=3) + n_lin = lambda m: sum(isinstance(x, torch.nn.Linear) for x in m.module.classifier.mlp) + assert n_lin(deep) == n_lin(shallow) + 2 + + +def test_train_knobs_reach_the_optimizer_and_plan(adata): + """lr/weight_decay must reach PYRO's optimizer — the regression that started this. + + Marked "hooked up" in the knob matrix while Pyro silently used scvi's hard-coded + 1e-3; the model still converged, so no behavioral test noticed. + """ + m = _model(adata) + plan = UnifiedTrainingPlan( + module=m.module, n_steps_kl_warmup=123, reconstruction_loss_scale=5e-3, + optimizer_config={"lr": 0.07, "betas": (0.8, 0.99), "eps": 1e-6, + "weight_decay": 2e-4}, + ) + args = plan.optim.pt_optim_args + assert args["lr"] == 0.07 + assert args["weight_decay"] == 2e-4 + assert args["betas"] == (0.8, 0.99) + assert args["eps"] == 1e-6 + assert plan.n_steps_kl_warmup == 123 + assert plan.reconstruction_loss_scale == 5e-3 + + +def test_reconstruction_loss_scale_reaches_the_module(adata): + m = _model(adata) + _train(m, reconstruction_loss_scale=7e-3) + assert m.module.reconstruction_loss_scale == pytest.approx(7e-3) + + +def test_max_epochs_and_patience_reach_the_trainer(adata): + m = _model(adata, patience=4) + _train(m, max_epochs=5) + assert m.trainer.max_epochs == 5 + # scvi installs its own EarlyStopping subclass (LoudEarlyStopping), so match on + # the attribute rather than the class name. + es = [cb for cb in m.trainer.callbacks if hasattr(cb, "patience")] + assert es, f"no early-stopping callback installed: {[type(c).__name__ for c in m.trainer.callbacks]}" + assert es[0].patience == 4, "patience did not reach the early-stopping callback" + + +def test_network_geometry_knobs_are_wired(adata): + """n_hidden/n_layers must actually size the encoder.""" + small = _model(adata, n_hidden=16, n_layers=1) + big = _model(adata, n_hidden=64, n_layers=3) + n = lambda m: sum(p.numel() for p in m.module.encoder.parameters()) + assert n(big) > 10 * n(small), f"encoder did not grow: {n(small)} -> {n(big)}" + + +def test_batch_size_reaches_the_dataloader(adata): + m = _model(adata) + for bs in (32, 200): + loader = m._make_data_loader(adata=m.adata, batch_size=bs, shuffle=False) + assert next(iter(loader))["X"].shape[0] == bs + + +def test_n_steps_kl_warmup_ramps_the_kl_weight(adata): + """The warmup must actually anneal module.kl_weight from ~0 up to kl_weight_max. + + NOTE the warmup is counted in optimizer STEPS while max_epochs is in epochs; with + batch_size >= n_obs that is one step per epoch (tracked as deviation DUX-2). + """ + m = _model(adata) + seen = [] + orig = UnifiedTrainingPlan.training_step + + def spy(self, batch, batch_idx): + out = orig(self, batch, batch_idx) + seen.append(float(self.module.kl_weight)) + return out + + UnifiedTrainingPlan.training_step = spy + try: + _train(m, max_epochs=12, batch_size=64, n_steps_kl_warmup=20) + finally: + UnifiedTrainingPlan.training_step = orig + warm = seen[:20] # the ramp itself; it plateaus afterwards + assert warm[0] < warm[len(warm) // 2] < warm[-1], f"kl_weight did not ramp: {warm[:6]}" + assert all(b >= a - 1e-12 for a, b in zip(seen, seen[1:])), "ramp must be monotonic" + assert max(seen) == pytest.approx(m.module.kl_weight_max), "ramp must reach the ceiling" + assert seen[0] < 1e-3, "ramp must start near zero" + + +def test_use_enumeration_selects_the_elbo(adata): + """use_enumeration picks TraceEnum_ELBO vs Trace_ELBO.""" + from pyro.infer import TraceEnum_ELBO, Trace_ELBO + + plain = UnifiedTrainingPlan(module=_model(adata, use_enumeration=False).module, + n_steps_kl_warmup=1, reconstruction_loss_scale=1e-3) + assert isinstance(plain.loss, Trace_ELBO) + pyro.clear_param_store() + enum = UnifiedTrainingPlan(module=_model(adata, use_enumeration=True).module, + n_steps_kl_warmup=1, reconstruction_loss_scale=1e-3) + assert isinstance(enum.loss, TraceEnum_ELBO) + + +# ══════════════════════════ BEHAVIOR ════════════════════════════════════════ + +def test_local_scale_controls_p_ct_draw_variance(adata): + """Dirichlet(β·p): Var = p(1−p)/(β+1) — variance strictly decreases in β.""" + m = _model(adata) + _train(m) + base = torch.as_tensor(m.get_p_ct(), dtype=torch.float64) + vs = [] + for beta in (1.0, 10.0, 100.0): + conc = torch.clamp(beta * base, min=1e-3) + d = torch.distributions.Dirichlet(conc) + vs.append(float(d.sample((4000,)).var(0).mean())) + assert vs[0] > vs[1] > vs[2], f"p_ct draw variance must fall with local_scale: {vs}" + # exact identity at β=10 on the first row + beta = 10.0 + p = base[0] / base[0].sum() + expected = (p * (1 - p) / (beta + 1)).mean() + got = torch.distributions.Dirichlet(torch.clamp(beta * base[0], min=1e-3)).sample((20000,)).var(0).mean() + assert float(got) == pytest.approx(float(expected), rel=0.15) + + +def test_global_scale_enters_the_eq1_prior(adata): + """α scales the clonotype-prior concentration (deviation [G] fix).""" + lo = _model(adata, global_scale=1.0) + hi = _model(adata, global_scale=50.0) + conc = lambda m: float((m.module.global_scale * m.module.mixture_concentration).sum(-1).mean()) + assert conc(hi) == pytest.approx(50.0 * conc(lo) / 1.0, rel=1e-5) + + +def test_prior_temperature_raises_clone_prior_entropy(adata): + """clone_phen_prior = normalize(prior**(1/T)); T>1 flattens ⇒ higher entropy.""" + ent = lambda M: float(np.mean([-np.sum(r[r > 0] * np.log2(r[r > 0])) for r in M])) + t1 = _model(adata, prior_temperature=1.0).module.clone_phen_prior.numpy() + t3 = _model(adata, prior_temperature=3.0).module.clone_phen_prior.numpy() + assert ent(t3) > ent(t1), "T>1 must raise the row-entropy of clone_phen_prior" + + +def test_guide_temperature_sharpens_get_p_ct(adata): + """get_p_ct sharpens q**(1/T); T<1 ⇒ lower row-entropy, same q_p_ct_raw.""" + m = _model(adata) + _train(m) + ent = lambda M: float(np.mean([-np.sum(r[r > 0] * np.log2(r[r > 0])) for r in M])) + m.module.guide_temperature = 1.0 + e1 = ent(m.get_p_ct()) + m.module.guide_temperature = 0.25 # same params, sharper read-out + e_sharp = ent(m.get_p_ct()) + assert e_sharp < e1, "T<1 must lower the row-entropy of get_p_ct()" + + +def test_classifier_temperature_divides_logits(adata): + """forward() divides by T ⇒ logits(T=2) == logits(T=1)/2 for fixed weights.""" + m = _model(adata, classifier_temperature=1.0) + z = torch.randn(5, m.module.n_latent) + m.module.eval() + with torch.no_grad(): + a = m.module.classifier(z) + m.module.classifier.temperature = 2.0 + b = m.module.classifier(z) + torch.testing.assert_close(b, a / 2.0) + + +@pytest.mark.parametrize("gate,expect", [(0.0, "prior"), (1.0, "classifier")]) +def test_gate_prob_endpoints_reduce_to_closed_forms(adata, gate, expect): + """ℓ = π·f_cls + (1−π)·log φ: π=0 ⇒ softmax(log φ); π=1 ⇒ softmax(f_cls).""" + import torch.nn.functional as F + from scvi import REGISTRY_KEYS + + m = _model(adata, gate_prob=gate) + _train(m) + got = m.predict(adata).values + + mod = m.module + mod.eval() + p_ct = torch.as_tensor(m.get_p_ct()) + with torch.no_grad(): + loader = m._make_data_loader(adata=m.adata, batch_size=4096) + chunks, start = [], 0 + for tensors in loader: + x = tensors[REGISTRY_KEYS.X_KEY] + b = tensors[REGISTRY_KEYS.BATCH_KEY].long() + n = x.shape[0] + prior = p_ct[mod.ct_array[start:start + n]] + z_loc, _, _ = mod.encoder(x, b) + if expect == "prior": + logits = torch.log(prior + 1e-8) + else: + logits = mod.classifier(z_loc) + chunks.append(F.softmax(logits, dim=-1)) + start += n + want = torch.cat(chunks).numpy() + np.testing.assert_allclose(got, want, atol=1e-5) + + +def test_predict_is_invariant_to_batch_size(adata): + """batch_size is a chunking detail, not a modelling one. + + Tolerance is float32-scale, not exact: different batch shapes take different BLAS + kernel paths, so the encoder's accumulations differ in the last bits (~1e-7 here). + Anything materially larger would mean batch_size is affecting the computation. + """ + m = _model(adata) + _train(m) + a = m.predict(adata, batch_size=64).values + b = m.predict(adata, batch_size=4096).values + 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) 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 new file mode 100644 index 0000000..9171b87 --- /dev/null +++ b/tests/test_model_smoke.py @@ -0,0 +1,65 @@ +"""Runtime smoke for the model: construct -> train (2 epochs) -> latent / p_ct / predict. + +Exercises the pieces split into sibling modules in PR3 (``TCRIModule`` model/guide, +``UnifiedTrainingPlan``, ``build_archetypes``, ``MixtureDirichlet``, ``VampPrior``) +end-to-end. Before this the suite only covered ``setup_anndata``; this locks in that +the split preserves the full construct/train/query path. +""" +import contextlib +import io + +import numpy as np + +from tcri.model._model import TCRIModel, build_archetypes + + +def test_model_construct_train_predict(synthetic_adata): + 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, + ) + + # PR3 rename: the clone x phenotype prior attribute is clone_phenotype_prior (was c2p_mat). + assert hasattr(model, "clone_phenotype_prior") + assert not hasattr(model, "c2p_mat") + n_clones, P = model.clone_phenotype_prior.shape + + # build_archetypes returns centers AND labels (M5). + centers, labels = build_archetypes(model.clone_phenotype_prior, K=3) + assert centers.shape == (3, P) + assert labels.shape == (n_clones,) + + buf = io.StringIO() + with contextlib.redirect_stdout(buf): + model.train( + max_epochs=2, + batch_size=64, + enable_progress_bar=False, + enable_model_summary=False, + ) + + n_cells = adata.n_obs + z = model.get_latent_representation() + assert z.shape == (n_cells, 8) + + p_ct = model.get_p_ct() + assert p_ct.ndim == 2 and p_ct.shape[1] == P + + probs = model.predict() + assert probs.shape == (n_cells, P) + 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_plotting/__init__.py b/tests/test_plotting/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_plotting/test_plotting.py b/tests/test_plotting/test_plotting.py deleted file mode 100644 index 8dc0ea6..0000000 --- a/tests/test_plotting/test_plotting.py +++ /dev/null @@ -1,16 +0,0 @@ -import matplotlib.pyplot as plt - -from tcri.plotting._plotting import clonality - - -def test_clonality_plot(mock_adata): - """Test clonality plotting function.""" - # Test basic functionality - ax = clonality(mock_adata, groupby="timepoint") - assert isinstance(ax, plt.Axes) - plt.close() - - # Test with splitby parameter - ax = clonality(mock_adata, groupby="timepoint", splitby="patient") - assert isinstance(ax, plt.Axes) - plt.close() diff --git a/tests/test_plotting/test_sankey.py b/tests/test_plotting/test_sankey.py deleted file mode 100644 index 149b9bc..0000000 --- a/tests/test_plotting/test_sankey.py +++ /dev/null @@ -1,60 +0,0 @@ -import matplotlib - -matplotlib.use("Agg") -import matplotlib.pyplot as plt - -from tcri.plotting._plotting import plot_pheno_sankey - - -def test_sankey_smoke_two_covariates(trained_model): - """Two covariates, return_axes=True: figure has axes with patches.""" - _model, adata = trained_model - covariates = list(adata.uns["tcri_covariate_categories"]) - assert len(covariates) >= 2 - - fig, ax = plot_pheno_sankey( - adata, - covariate_order=covariates[:2], - return_axes=True, - ) - assert fig is not None - assert ax in fig.axes - assert len(ax.patches) > 0 - plt.close(fig) - - -def test_sankey_clones_filtering(trained_model): - """Subset of clones runs without error and produces fewer connections.""" - _model, adata = trained_model - covariates = list(adata.uns["tcri_covariate_categories"])[:2] - all_clones = list(adata.uns["tcri_clonotype_categories"]) - - fig_full, _ = plot_pheno_sankey( - adata, covariate_order=covariates, return_axes=True, - ) - fig_subset, ax_subset = plot_pheno_sankey( - adata, covariate_order=covariates, - clones=all_clones[: max(1, len(all_clones) // 4)], - return_axes=True, - ) - assert fig_subset is not None - assert ax_subset in fig_subset.axes - plt.close(fig_full) - plt.close(fig_subset) - - -def test_sankey_normalize_false(trained_model): - """normalize=False (raw cell counts) produces an axis.""" - _model, adata = trained_model - covariates = list(adata.uns["tcri_covariate_categories"])[:2] - - fig, ax = plot_pheno_sankey( - adata, - covariate_order=covariates, - normalize=False, - return_axes=True, - ) - assert fig is not None - assert ax in fig.axes - assert len(ax.patches) > 0 - plt.close(fig) diff --git a/tests/test_plotting/test_twins.py b/tests/test_plotting/test_twins.py new file mode 100644 index 0000000..8efbaf7 --- /dev/null +++ b/tests/test_plotting/test_twins.py @@ -0,0 +1,45 @@ +"""The four pl↔tl twins (cache renderers) + resolve_palette (PR7). Smoke: each runs and +returns a Matplotlib axes (or the tidy df via return_df).""" +import matplotlib + +matplotlib.use("Agg") + +import tcri +from tcri import _keys as K + + +def _cov(adata): + return list(adata.uns[K.COVARIATE_CATEGORIES])[0] + + +def test_pl_mutual_information(trained_model): + _, adata = trained_model + ax = tcri.pl.mutual_information(adata, covariate=_cov(adata)) + assert ax is not None + df = tcri.pl.mutual_information(adata, covariate=_cov(adata), return_df=True) + assert "MI" in df.columns + + +def test_pl_entropy_twins(trained_model): + _, adata = trained_model + cov = _cov(adata) + assert tcri.pl.clonotypic_entropy(adata, covariate=cov, groupby="patient") is not None + assert tcri.pl.phenotypic_entropy(adata, covariate=cov, groupby="patient") is not None + # no-groupby falls back to a bar of the Series + assert tcri.pl.clonotypic_entropy(adata, covariate=cov) is not None + + +def test_pl_phenotypic_flux(trained_model): + _, adata = trained_model + covs = list(adata.uns[K.COVARIATE_CATEGORIES]) + ax = tcri.pl.phenotypic_flux(adata, order=covs) + assert ax is not None + df = tcri.pl.phenotypic_flux(adata, order=covs, return_axes=True) + assert df is not None + + +def test_resolve_palette_in_place(trained_model): + _, adata = trained_model + mapping = tcri.pl.resolve_palette(adata, "patient") + assert "patient_colors" in adata.uns + assert isinstance(mapping, dict) and "patient" in mapping diff --git a/tests/test_preprocessing/test_joint_distribution_posterior.py b/tests/test_preprocessing/test_joint_distribution_posterior.py deleted file mode 100644 index 41ffcd2..0000000 --- a/tests/test_preprocessing/test_joint_distribution_posterior.py +++ /dev/null @@ -1,61 +0,0 @@ -"""Tests for joint_distribution_posterior view/subset handling (Notion #4 / T7). - -Filtered AnnData (views or subset copies) keep their per-cell `.uns` arrays in the -original full-cell space while `.obs`/`.obsm` are subset. The function used to index -one with positions derived from the other, silently returning misaligned results. -It must now raise instead. -""" - -import numpy as np -import pandas as pd -import pytest - -from tcri.preprocessing._preprocessing import joint_distribution_posterior - - -def _covariate(adata): - return adata.uns["tcri_covariate_categories"][0] - - -def test_jd_posterior_full_adata_ok(trained_model): - """Baseline: the full registered AnnData returns a valid distribution.""" - _, adata = trained_model - df = joint_distribution_posterior(adata, _covariate(adata), silent=True) - assert isinstance(df, pd.DataFrame) - assert df.shape[1] == len(adata.uns["tcri_phenotype_categories"]) - assert np.all(np.isfinite(df.to_numpy())) - - -def test_jd_posterior_rejects_filtered_view(trained_model): - """A cell-filtered view must raise, not silently misalign (Notion #4).""" - _, adata = trained_model - cov_col = adata.uns["tcri_metadata"]["covariate_col"] - mask = np.asarray(adata.obs[cov_col] == _covariate(adata)) - view = adata[mask] - assert view.n_obs < adata.n_obs - - with pytest.raises(ValueError, match=r"filtered|subset|register_model"): - joint_distribution_posterior(view, _covariate(adata), silent=True) - - -def test_jd_posterior_rejects_filtered_copy(trained_model): - """A cell-filtered copy is also misaligned: .uns stays full-length.""" - _, adata = trained_model - cov_col = adata.uns["tcri_metadata"]["covariate_col"] - mask = np.asarray(adata.obs[cov_col] == _covariate(adata)) - sub = adata[mask].copy() - assert sub.n_obs < adata.n_obs - - with pytest.raises(ValueError, match=r"filtered|subset|register_model"): - joint_distribution_posterior(sub, _covariate(adata), silent=True) - - -def test_jd_posterior_allows_gene_subset(trained_model): - """Var-only subsetting keeps n_obs intact and must NOT trip the guard.""" - _, adata = trained_model - sub = adata[:, : adata.n_vars // 2] - assert sub.n_obs == adata.n_obs - - df = joint_distribution_posterior(sub, _covariate(adata), silent=True) - assert isinstance(df, pd.DataFrame) - assert df.shape[1] == len(adata.uns["tcri_phenotype_categories"]) diff --git a/tests/test_session_round_trip.py b/tests/test_session_round_trip.py index 1f1941d..f1b5a6c 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,35 @@ 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 + ) + + # 4) the model scalars carried only in init_params_ (not in the AnnData) survive + # the save/load — guards the classifier knobs added in the model PR. + assert loaded_model.module.phenotype_kl_weight == model.module.phenotype_kl_weight + assert loaded_model.module.gate_prob == model.module.gate_prob + assert loaded_model.module.classifier_dropout == model.module.classifier_dropout diff --git a/tests/test_tools/__init__.py b/tests/test_tools/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_tools/test_joint.py b/tests/test_tools/test_joint.py new file mode 100644 index 0000000..a0d5d3e --- /dev/null +++ b/tests/test_tools/test_joint.py @@ -0,0 +1,201 @@ +"""The Phase-5 engine gate — ``tcri.joint_distribution`` identities (§7.1). + +Comparisons use the *frozen* canonical keys written by ``to_anndata`` (``uns[P_CT]``, +``obsm[X_LOGITS]``, ``obsm[X_PROBABILITIES]``) rather than a live ``predict()`` call, +so the identities do not depend on the process-global pyro param store (§5.2). +""" +import json + +import numpy as np +import pandas as pd + +import tcri +from tcri import _keys as K + + +def _first_covariate(adata): + return list(adata.uns[K.COVARIATE_CATEGORIES])[0] + + +def test_top_level_reexport(): + assert hasattr(tcri, "joint_distribution") + assert tcri.joint_distribution is tcri.tools.joint_distribution + + +def test_identity_ct_table_equals_p_ct(trained_model): + """use_logits=False, n_samples=0, T=1 == uns[P_CT] restricted to the covariate.""" + _, adata = trained_model + cov = _first_covariate(adata) + df = tcri.joint_distribution(adata, covariate=cov, use_logits=False, n_samples=0, temperature=1.0) + + cov_i = list(adata.uns[K.COVARIATE_CATEGORIES]).index(cov) + rows = np.where(np.asarray(adata.uns[K.CT_TO_COV]) == cov_i)[0] + expected = np.asarray(adata.uns[K.P_CT])[rows] # one ct row per clone + np.testing.assert_allclose(df.values, expected, atol=1e-6) + + +def test_identity_use_logits_equals_predict_aggregation(trained_model): + """use_logits=True, n_samples=0, T=1 == per-clone mean of predict() (== frozen X_PROBABILITIES).""" + _, adata = trained_model + cov = _first_covariate(adata) + df = tcri.joint_distribution(adata, covariate=cov, use_logits=True, n_samples=0, temperature=1.0) + + cov_i = list(adata.uns[K.COVARIATE_CATEGORIES]).index(cov) + mask = np.asarray(adata.uns[K.COV_ARRAY]) == cov_i + clone_col = adata.uns[K.METADATA]["clone_col"] + probs = adata.obsm[K.X_PROBABILITIES][mask] + agg = pd.DataFrame(probs, columns=list(adata.uns[K.PHENOTYPE_CATEGORIES])) + agg["c"] = adata.obs[clone_col].values[mask] + agg = agg.groupby("c")[list(adata.uns[K.PHENOTYPE_CATEGORIES])].mean().reindex(df.index) + np.testing.assert_allclose(df.values, agg.values, atol=1e-4) + + +def test_point_estimate_is_deterministic(trained_model): + """n_samples=0 is deterministic and bit-identical across repeated calls.""" + _, adata = trained_model + cov = _first_covariate(adata) + a = tcri.joint_distribution(adata, covariate=cov, n_samples=0) + b = tcri.joint_distribution(adata, covariate=cov, n_samples=0) + assert np.array_equal(a.values, b.values) + + +def test_sampling_is_seeded_reproducible(trained_model): + """n_samples>0 is reproducible under a fixed random_state, varies otherwise, and + carries a (clonotype, sample_id) MultiIndex.""" + _, adata = trained_model + cov = _first_covariate(adata) + s1 = tcri.joint_distribution(adata, covariate=cov, n_samples=8, random_state=0) + s2 = tcri.joint_distribution(adata, covariate=cov, n_samples=8, random_state=0) + s3 = tcri.joint_distribution(adata, covariate=cov, n_samples=8, random_state=1) + np.testing.assert_allclose(s1.values, s2.values) + assert not np.allclose(s1.values, s3.values) + assert list(s1.index.names) == ["clonotype", "sample_id"] + # draws are valid simplices + np.testing.assert_allclose(s1.values.sum(axis=1), 1.0, atol=1e-5) + + +def test_sampling_mean_approaches_tempered_base(trained_model): + """Many draws average to the Dirichlet mean == the (tempered) base — validates the + draw is Dirichlet(clamp(local_scale·base)).""" + _, adata = trained_model + cov = _first_covariate(adata) + base = tcri.joint_distribution(adata, covariate=cov, use_logits=False, n_samples=0) + draws = tcri.joint_distribution(adata, covariate=cov, use_logits=False, n_samples=400, random_state=0) + mean = draws.groupby(level="clonotype", sort=False).mean().reindex(base.index) + np.testing.assert_allclose(mean.values, base.values, atol=0.05) + + +def test_weighted_scales_rows_by_ct_cell_count(trained_model): + """weighted=True scales each clone row by its (ct-keyed) cell count; weighted=False + is a per-clone simplex.""" + _, adata = trained_model + cov = _first_covariate(adata) + w0 = tcri.joint_distribution(adata, covariate=cov, use_logits=False, n_samples=0, weighted=False) + w1 = tcri.joint_distribution(adata, covariate=cov, use_logits=False, n_samples=0, weighted=True) + np.testing.assert_allclose(w0.values.sum(axis=1), 1.0, atol=1e-6) # unweighted rows sum to 1 + assert not np.allclose(w0.values, w1.values) + # weighted row sum == that clone's cell count at this covariate + cov_i = list(adata.uns[K.COVARIATE_CATEGORIES]).index(cov) + mask = np.asarray(adata.uns[K.COV_ARRAY]) == cov_i + clone_col = adata.uns[K.METADATA]["clone_col"] + counts = adata.obs[clone_col].values[mask] + expected = pd.Series(counts).value_counts().reindex(w1.index).astype(float) + np.testing.assert_allclose(w1.values.sum(axis=1), expected.values, atol=1e-6) + + +def test_shared_draw_invariant_across_covariates(trained_model): + """covariate=None draws once and shares it: the slice for a covariate equals the + per-covariate call at the same random_state (draw-count == n_samples, not ×#covariates).""" + _, adata = trained_model + cov = _first_covariate(adata) + allcov = tcri.joint_distribution(adata, covariate=None, n_samples=6, random_state=7) + percov = tcri.joint_distribution(adata, covariate=cov, n_samples=6, random_state=7) + assert "covariate" in allcov.index.names + sl = allcov.xs(cov, level="covariate") + np.testing.assert_allclose(sl.values, percov.values, atol=1e-6) + assert allcov.attrs["params"]["n_draws"] == 6 # one draw block, not ×#covariates + + +def test_provenance_is_json_serializable(trained_model): + _, adata = trained_model + cov = _first_covariate(adata) + df = tcri.joint_distribution(adata, covariate=cov, n_samples=4, random_state=0) + json.dumps(df.attrs["params"]) # must not raise + assert df.attrs["params"]["use_logits"] is True + assert df.attrs["params"]["n_draws"] == 4 + + +def test_groupby_deferred(trained_model): + _, adata = trained_model + cov = _first_covariate(adata) + import pytest + with pytest.raises(NotImplementedError, match="groupby"): + tcri.joint_distribution(adata, covariate=cov, groupby="patient") + + +def test_gate_aware_combine_direct(): + """use_logits=True with a numeric gate_prob uses g·logits + (1-g)·log(base), then + mean per clone — exercised directly (the trained_model fixture has gate_prob=None).""" + from tcri._compute._joint import _joint_draws + + p_ct = np.array([[0.7, 0.3], [0.2, 0.8]], dtype=float) # 2 ct rows == 2 clones + ct_to_cov = np.array([0, 0]); ct_to_c = np.array([0, 1]) + ct_array = np.array([0, 0, 1, 1]); cov_array = np.array([0, 0, 0, 0]) + logits = np.array([[2.0, -1.0], [1.0, 0.0], [-1.0, 2.0], [0.0, 1.0]], dtype=float) + g = 0.5 + blocks, _ = _joint_draws( + p_ct, ct_to_cov, ct_to_c, ct_array, cov_array, local_scale=1.0, n_samples=0, + temperature=1.0, use_logits=True, covariate_idx=0, logits=logits, gate_prob=g, + ) + _, _, J = blocks[0] # [1, 2, 2] + + eps = 1e-8 + def _sm(x): + e = np.exp(x - x.max()); return e / e.sum() + exp = np.zeros((2, 2)) + for c in range(2): + cells = [i for i in range(4) if ct_to_c[ct_array[i]] == c] + exp[c] = np.mean([_sm(g * logits[i] + (1 - g) * np.log(p_ct[ct_array[i]] + eps)) for i in cells], 0) + np.testing.assert_allclose(J[0], exp, atol=1e-10) + + +def test_temperature_tempers_base_on_ct_path(trained_model): + """use_logits=False, T!=1 == softmax(log(P_CT+1e-8)/T) restricted to the covariate.""" + from scipy.special import softmax + _, adata = trained_model + cov = _first_covariate(adata) + df = tcri.joint_distribution(adata, covariate=cov, use_logits=False, n_samples=0, temperature=2.0) + cov_i = list(adata.uns[K.COVARIATE_CATEGORIES]).index(cov) + rows = np.where(np.asarray(adata.uns[K.CT_TO_COV]) == cov_i)[0] + expected = softmax(np.log(np.asarray(adata.uns[K.P_CT])[rows] + 1e-8) / 2.0, axis=1) + np.testing.assert_allclose(df.values, expected, atol=1e-5) + + +def test_temperature_changes_use_logits_joint(trained_model): + """T!=1 tempers the cell-informed joint too (behavior lock, per §7.1).""" + _, adata = trained_model + cov = _first_covariate(adata) + a = tcri.joint_distribution(adata, covariate=cov, use_logits=True, n_samples=0, temperature=1.0) + b = tcri.joint_distribution(adata, covariate=cov, use_logits=True, n_samples=0, temperature=2.0) + assert not np.allclose(a.values, b.values) + + +def test_subset_anndata_raises(trained_model): + """A filtered/sliced AnnData (full-space uns vs subset obsm) is refused, not silently + misaligned.""" + import pytest + _, adata = trained_model + sub = adata[: adata.n_obs // 2].copy() # uns per-cell arrays stay full-space + with pytest.raises(ValueError, match="n_obs|filtered|subset"): + tcri.joint_distribution(sub, covariate=_first_covariate(adata)) + + +def test_missing_local_scale_raises_only_when_sampling(trained_model): + """n_samples>0 without uns[LOCAL_SCALE] errors (no silent 1.0 fallback); n_samples=0 is fine.""" + import pytest + _, adata = trained_model + a2 = adata.copy() + del a2.uns[K.LOCAL_SCALE] + with pytest.raises(RuntimeError, match="LOCAL_SCALE|local_scale"): + tcri.joint_distribution(a2, covariate=_first_covariate(adata), n_samples=4) + tcri.joint_distribution(a2, covariate=_first_covariate(adata), n_samples=0) # ok diff --git a/tests/test_tools/test_metrics.py b/tests/test_tools/test_metrics.py new file mode 100644 index 0000000..89e1457 --- /dev/null +++ b/tests/test_tools/test_metrics.py @@ -0,0 +1,100 @@ +"""The four engine-backed metric twins + compare_groups (PR6). Shapes, bits/normalization, +and the group-comparison math.""" +import numpy as np +import pandas as pd +import pytest + +import tcri +from tcri import _keys as K + + +def _cov(adata): + return list(adata.uns[K.COVARIATE_CATEGORIES])[0] + + +def test_clonotypic_entropy_shapes_and_range(trained_model): + _, adata = trained_model + cov = _cov(adata) + s = tcri.tl.clonotypic_entropy(adata, covariate=cov, n_samples=0) + assert isinstance(s, pd.Series) + assert list(s.index) == list(adata.uns[K.PHENOTYPE_CATEGORIES]) + finite = s.dropna().to_numpy() + assert (finite >= -1e-9).all() and (finite <= 1 + 1e-9).all() # normalized bits in [0,1] + df = tcri.tl.clonotypic_entropy(adata, covariate=cov, n_samples=8, random_state=0) + assert set(["mean", "sd", "hdi_low", "hdi_high"]).issubset(df.columns) + + +def test_phenotypic_entropy_shapes_and_range(trained_model): + _, adata = trained_model + cov = _cov(adata) + s = tcri.tl.phenotypic_entropy(adata, covariate=cov, n_samples=0) + assert isinstance(s, pd.Series) + finite = s.dropna().to_numpy() + assert (finite >= -1e-9).all() and (finite <= 1 + 1e-9).all() + + +def test_mutual_information_scalar_and_fastpath(trained_model): + _, adata = trained_model + cov = _cov(adata) + mi = tcri.tl.mutual_information(adata, covariate=cov, n_samples=0) + assert isinstance(mi, float) and -1e-9 <= mi <= 1 + 1e-9 # normalized (min) in [0,1] + jd = tcri.joint_distribution(adata, covariate=cov, n_samples=0) + assert np.isclose(tcri.tl.mutual_information(jd, n_samples=0), mi) # fast path + summ = tcri.tl.mutual_information(adata, covariate=cov, n_samples=8, random_state=0) + assert set(["mean", "sd", "hdi_low", "hdi_high"]) == set(summ) + + +def test_mutual_information_unnormalized_is_bits(trained_model): + _, adata = trained_model + cov = _cov(adata) + mi_bits = tcri.tl.mutual_information(adata, covariate=cov, n_samples=0, normalized=False) + assert mi_bits >= -1e-9 # raw MI in bits, non-negative + + +def test_metric_groupby_tidy(trained_model): + _, adata = trained_model + cov = _cov(adata) + df = tcri.tl.mutual_information(adata, covariate=cov, groupby="patient", n_samples=0) + assert "patient" in df.columns and "MI" in df.columns + ce = tcri.tl.clonotypic_entropy(adata, covariate=cov, groupby="patient", n_samples=0) + assert set(["patient", "phenotype", "clonotypic_entropy"]).issubset(ce.columns) + + +def test_phenotypic_flux_over_common_clones(trained_model): + _, adata = trained_model + covs = list(adata.uns[K.COVARIATE_CATEGORIES]) + if len(covs) < 2: + pytest.skip("needs >=2 covariates") + fx = tcri.tl.phenotypic_flux(adata, cov_from=covs[0], cov_to=covs[1], n_samples=0, distance_metric="l1") + assert isinstance(fx, pd.Series) + v = fx.dropna().to_numpy() + assert (v >= -1e-9).all() and (v <= 2 + 1e-9).all() # l1 on the simplex is bounded [0,2] + + +def test_compare_groups_unpaired(): + """Mann–Whitney contrast on a tidy per-unit frame (R vs NR).""" + df = pd.DataFrame({ + "patient": [f"p{i}" for i in range(8)], + "response": ["R"] * 4 + ["NR"] * 4, + "MI": [0.8, 0.75, 0.82, 0.79, 0.4, 0.35, 0.45, 0.5], + }) + out = tcri.tl.compare_groups(df, value="MI", splitby="response", reference="NR") + assert len(out) == 1 + row = out.iloc[0] + assert {row["group_a"], row["group_b"]} == {"R", "NR"} + assert row["mean_b"] > row["mean_a"] or row["mean_a"] > row["mean_b"] + assert "p" in out.columns and "stars" in out.columns and "delta" in out.columns + + +def test_compare_groups_paired_direction(): + """Paired posterior-draw contrast emits p_gt + HDI via prob_direction.""" + rng = np.random.default_rng(0) + rows = [] + for u in range(5): + rows.append({"unit": u, "arm": "A", "v": rng.normal(0.0, 0.1, size=200)}) + rows.append({"unit": u, "arm": "B", "v": rng.normal(0.5, 0.1, size=200)}) # B > A + df = pd.DataFrame(rows) + out = tcri.tl.compare_groups(df, value="v", splitby="arm", reference="A", paired=True, pair_on="unit") + assert len(out) == 1 + assert out.iloc[0]["delta"] > 0 and out.iloc[0]["p_gt"] > 0.9 # B-A positive, high direction prob + assert "hdi_low" in out.columns diff --git a/tests/test_tools/test_reduce.py b/tests/test_tools/test_reduce.py new file mode 100644 index 0000000..88ee7ff --- /dev/null +++ b/tests/test_tools/test_reduce.py @@ -0,0 +1,43 @@ +"""`_compute/_reduce` — the batched information-theoretic reductions default to +**bits (log2)**, the tcri convention.""" +import numpy as np + +from tcri._compute import _reduce + + +def test_entropy_uniform_is_log2_P_bits(): + """Entropy of the uniform distribution over P outcomes = log2(P) bits.""" + for P in (2, 3, 8): + u = np.full(P, 1.0 / P) + np.testing.assert_allclose(_reduce.entropy(u), np.log2(P), atol=1e-12) + # a point mass has zero entropy + assert _reduce.entropy(np.array([1.0, 0.0, 0.0])) == 0.0 + + +def test_entropy_base_override_nats(): + u = np.full(4, 0.25) + np.testing.assert_allclose(_reduce.entropy(u, base=None), np.log(4), atol=1e-12) # nats + + +def test_entropy_batched_over_leading_axes(): + p = np.stack([np.full(4, 0.25), np.array([1.0, 0, 0, 0])]) # [2, 4] + out = _reduce.entropy(p) # -> [2] + np.testing.assert_allclose(out, [2.0, 0.0], atol=1e-12) + + +def test_mutual_information_independent_is_zero(): + Px = np.array([0.5, 0.5]); Py = np.array([0.25, 0.75]) + joint = np.outer(Px, Py) # independent + np.testing.assert_allclose(_reduce.mutual_information(joint), 0.0, atol=1e-12) + + +def test_mutual_information_perfectly_coupled_equals_marginal_entropy(): + """A diagonal joint (X determines Y) has MI == H(X) == H(Y) (bits).""" + joint = np.diag([0.5, 0.5]) # perfectly coupled 2x2 + np.testing.assert_allclose(_reduce.mutual_information(joint), 1.0, atol=1e-9) # 1 bit + + +def test_mutual_information_renormalizes_unnormalized_joint(): + """An un-normalized (cell-weighted count) joint is accepted (renormalized).""" + joint = np.diag([5.0, 5.0]) # counts, not probabilities + np.testing.assert_allclose(_reduce.mutual_information(joint), 1.0, atol=1e-9)