diff --git a/.gitignore b/.gitignore index ee9453c..ccdf6e3 100644 --- a/.gitignore +++ b/.gitignore @@ -9,6 +9,9 @@ __pycache__/ # Backup / scratch files *.bak +# Refactor scratchpad — running notes, intentionally not tracked +docs/contract/REFACTOR_NOTES.md + # Distribution / packaging .Python build/ diff --git a/docs/conf.py b/docs/conf.py index 64efafd..fc1895b 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -55,7 +55,10 @@ def _get_release(): # List of patterns, relative to source directory, that match files and # directories to ignore when looking for source files. -exclude_patterns = ['_build', 'Thumbs.db', '.DS_Store'] +# 'contract' holds internal refactor-planning docs (target API contract + +# dependency map), generated by docs/contract/build_*.py — not part of the +# published site, so keep them out of the Sphinx toctree. +exclude_patterns = ['_build', 'Thumbs.db', '.DS_Store', 'contract'] # The name of the Pygments (syntax highlighting) style to use. pygments_style = 'sphinx' diff --git a/docs/contract/README.md b/docs/contract/README.md new file mode 100644 index 0000000..4d03e01 --- /dev/null +++ b/docs/contract/README.md @@ -0,0 +1,38 @@ +# TCRI refactor contract + +Internal planning/design docs for the `tcri` API refactor (target: a standalone scverse-ecosystem +package, "Door A"). **Not** part of the published Sphinx site (excluded via `exclude_patterns` in +`docs/conf.py`). Structural reference is the `../grafiti` package. + +## ⚠ Read first — governance + +| File | What it is | +|---|---| +| **`REFACTOR_HISTORY.md`** | Single source of truth: the **Hard Rules**, the chronology, and all settled decisions (§2). Read before touching anything. | +| **`REDO_LIST.md`** | The corrected **disposition map** (keep / drop) + the clean re-implementation plan. | + +**Hard Rules (never violate):** the `example/` notebooks are **disposable** — never read them or use them +for any decision. **Non-core = DROP (delete); nothing moves to `examples/`.** Disposition is decided by one +test: *is it core?* + +## The authoritative pair + +| File | What it is | +|---|---| +| **`tcri_api_and_responsibilities.md`** | Final API surface + per-function math/stats spec + prior-vs-mean resolution. **Clean** (post-recovery). | +| **`tcri_implementation_plan.md`** | Ordered PR sequence + model→AnnData streamline + GPU architecture + testing/CI. **Clean** (post-recovery). | +| `tcri_consistency_sufficiency.md` | Argument consistency deltas + sufficiency-from-responsibility (clean re-derivation). | + +## Supporting artifacts + +| File | What it is | +|---|---| +| `tcri_function_inventory.md` (+ data, `build_tcri_inventory.py`) | 131-function labeled inventory | +| `tcri_api_contract.*`, `tcri_dependency_map.*` (+ build scripts) | early target API + dependency graph | + +## `_quarantine/` — do not use + +Earlier versions of the two authoritative docs + their audit data, contaminated by a **notebook-caller-census +leak** (a workflow prompt treated the disposable notebooks as an authority on what to keep, resurrecting +`gene_entropy`/`probability_ternary`). Their **disposition sections are wrong**; their math/design was clean and +was salvaged into the current authoritative docs. Kept only as a record. See `REFACTOR_HISTORY.md`. diff --git a/docs/contract/REDO_LIST.md b/docs/contract/REDO_LIST.md new file mode 100644 index 0000000..b983003 --- /dev/null +++ b/docs/contract/REDO_LIST.md @@ -0,0 +1,61 @@ +# Redo list — clean re-implementation after the notebook-census leak + +Ordered work to redo, from the leak point onward. Each item is re-implemented **clean** (Hard Rules in +`REFACTOR_HISTORY.md`: no notebooks, no `examples/`, disposition = "is it core"). Check each against +`REFACTOR_HISTORY.md` §2 as it lands. + +## What we did since the leak point (to re-implement) +1. **Corrected disposition map** (below) — the foundation the two docs got wrong. **Done in this file.** +2. **API + Function-Responsibility doc** — rebuild clean from the quarantined version: keep its (clean) + math/stats spec, prior-vs-mean resolution, per-arg math, `covariate=None`=all semantics; **replace §11 and every + disposition/`examples`/notebook reference** with the map below. +3. **Implementation Plan doc** — rebuild clean: keep the (clean) PR sequence, model→anndata streamline, GPU + architecture, testing/CI; **delete Phase 0's "caller census," every "move to examples," and every "N notebook + uses" justification**; notebook rewrite is a downstream chore, never a disposition driver. +4. **Fold the (clean) consistency findings** into the API doc: `random_state` naming; canonical arg order; + `distance_metric` (not `metric`); `palette` (not `phenotype_colors`); `normalized` (not `normalize`); + `order=`/`hue_order=` parity; `figsize` unify; keyword-only after `*`; `show=` on public `pl`; American spelling. +5. **Re-derive sufficiency from RESPONSIBILITY only** — a kept function gets an arg iff its own job needs it + (`random_state` on sampling fns; `clones=` on per-clone metrics; `order=` for plot category ordering). + **Drop every notebook-parity item** (`gene_entropy`, `probability_ternary`, `weighted=`, `minimum_clone_size=`, + `base=`, `covariate_key=`, `pair_on=`-for-notebooks, etc.) unless responsibility independently requires it. +6. **Necessary/sufficient audit** of the rebuilt docs (consistency matrix over the surface; sufficiency vs responsibility). +7. **Final audit** of everything against `REFACTOR_HISTORY.md`. + +--- + +## CORRECTED DISPOSITION MAP (authoritative; supersedes quarantined §11 / §3) + +Decided by "is it core," never by callers. `examples/` does not appear. Allowed non-drop destinations: a kept +namespace, `diag/`, `docs/` (PGM only), or a private helper module. + +### KEEP — core surface +- **ml:** `TCRIModel` = `setup_anndata`, `train`, `get_latent_representation`, `predict` (was `get_cell_phenotype_probs`), + `get_p_ct`; private internals `TCRIModule` / priors / classifier / training-plan / `build_archetypes`. +- **pp:** `register_model`→`model.to_anndata` (streamlined), `group_singletons`, `clone_size`. +- **tl:** `joint_distribution` (+ top-level `tcri.joint_distribution`), `clonotypic_entropy`, `phenotypic_entropy`, + `mutual_information`, `phenotypic_flux`, `compare_groups`. +- **pl:** `clonotypic_entropy`, `phenotypic_entropy`, `mutual_information`, `phenotypic_flux`; private + `_sankey` (`SankeyNode` + `_phenotype_mass_per_clone`), `_metric_boxplot` (was `tcri_boxplot`), `_colors`/`resolve_palette`. +- **diag:** `joint_distribution_ppc` (fixed `compare_joint_distribution`), `phenotype_calibration`, + `reconstruction_ppc`, `permutation_null`, `loss` (was `plot_loss`), `archetypes` (was `plot_archetypes`). +- **ut:** `save_tcri_session`, `load_tcri_session`. +- **shared:** `_keys`, `_console`, `_stats`, `_distance`. + +### TRANSFORM (kept, relocated — NOT examples) +- `compare_joint_distribution` → `diag.joint_distribution_ppc` (fixed). `plot_loss`→`diag.loss`; + `plot_archetypes`→`diag.archetypes`. `build_nested_tcri_pgm`/`draw_tcri_pgm_nested` → **`docs/`** figure script. +- `register_phenotype_key`/`register_clonotype_key`/`_compute_logits_and_prior` → folded into `to_anndata` (private). +- `dkl` → `_distance.kl_divergence`. `tcri_boxplot` → private `_metric_boxplot`. + +### DROP — deleted (NOT moved to examples) +`clonality`, `probability_distribution`, `bayesian_mutual_information`, `polar_plot`, `compare_phenotypes`, +**`probability_ternary`** (was wrongly kept), **`gene_entropy`** (was wrongly →examples), **`top_clone_umap`**, +**`clone_size_umap`**, **`plot_phenotype_probabilities`** (all three were wrongly →examples), `mi_compare`, +`delta_clonotypic_entropy`, `delta_entropy_table`, `flux_table`, `clonotypic_entropy_base`, `ridge_delta_entropy`, +`classify_phenotypes`, `get_latent_embedding`, `group_small_clones`, `register_probability_columns`, +`remove_meaningless_genes`, `clone_fraction`, `probabilities`, `_ent`, `_ascii_hist` (+ all `graph=`/ASCII paths), +`SankeyNode.hex_to_rgb`, `write_adata_safely`, `_pop_nonserializables`, the plural `*_entropies` shims, the `*_tl`/ +`centropy`/`pentropy` aliases. + +> The 6 items in **bold** are the ones the leak wrongly kept-or-moved-to-examples. They are DROPPED. diff --git a/docs/contract/REFACTOR_AGENDA.md b/docs/contract/REFACTOR_AGENDA.md new file mode 100644 index 0000000..09825c5 --- /dev/null +++ b/docs/contract/REFACTOR_AGENDA.md @@ -0,0 +1,130 @@ +# TCRI Refactor — Agenda, Tasklist & Diary (LIVING DOC) + +**Update this every working session. Refer back to it before starting anything.** It is the operational +tracker + running diary for the whole refactor. The detailed spec lives in `tcri_api_and_responsibilities.md` ++ `tcri_implementation_plan.md`; the rules/history in `REFACTOR_HISTORY.md`; the scratch/deferred pile in +`REFACTOR_NOTES.md` (not checked in). + +## How to use this doc +1. **Removal is a hard bar.** The default failure mode is keeping old code around and over-engineering to + preserve it. **Do the opposite.** Every function in the Removal Ledger must actually be deleted, its + `__all__`/import sites cleaned, and the checkbox ticked. If keeping something "just in case" feels tempting — + don't. Delete it; git has it. +2. **Never read the `example/` notebooks.** They are disposable and are an *output* of the refactor, never an + input. No caller census, no "is-it-used," no "sufficiency." +3. **After every PR, run the Standing Audit** (below) and write a diary entry. +4. **Frequent audits**: at minimum after each PR; ideally mid-PR when a component is touched. Log them in the + Audit Log. +5. **Usability is a first-class check** — every session ask "is this easier to use than before?" + +## Standing Audit (run after each PR — copy into the diary entry) +- [ ] **Removed everything slated?** (cross-check the Removal Ledger; `__all__` + import-sites clean; `import tcri` green) +- [ ] **Added everything wanted?** (the PR's deliverables all present) +- [ ] **Tests:** what components can now be unit-tested that weren't? Added? +- [ ] **Streamline:** any duplication / dead branch / needless complexity spotted? Removed or logged? +- [ ] **Usability:** simpler signatures / clearer errors / fewer steps than before? +- [ ] **Contract conformance green?** (`test_contract_conformance`) +- [ ] Diary entry written; ledger + statuses updated. + +## Status legend: ☐ todo · ◐ in progress · ✅ done · ⚠ blocked + +## PR Agenda +| # | PR | Status | Risk | Depends | Gate | +|---|---|---|---|---|---| +| 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 | + +## 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 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 5/6 (consolidated away — delete WITH replacement, never before):** +- [ ] `pp.joint_distribution_posterior` (→ unified `joint_distribution`) · [ ] `metrics._mi_from_joint` (→ `_mutual_information`) +- [ ] `tl.mi_compare` (→ `compare_groups`) · [ ] `tl.delta_clonotypic_entropy` · [ ] `tl.delta_entropy_table` · [ ] `tl.flux_table` +- [ ] `tl.clonotypic_entropy_base` · [ ] `tl.clonality` · [ ] `tl.dkl` local `dkl_func` +- [ ] plural `*_entropies` shims · [ ] `metrics/` package (after migration to `tools/`) + +**Phase 7 (non-core plots — DROP, not to examples):** +- [ ] `pl.probability_ternary` · [ ] `pl.top_clone_umap` · [ ] `pl.clone_size_umap` · [ ] `pl.plot_phenotype_probabilities` +- [ ] `pl.compare_phenotypes` · [ ] `pl.ridge_delta_entropy` · [ ] `pl.flux` boxplot · [ ] `pl.clonality` plot +- [ ] `pl.tcri_boxplot` (→ private `_metric_boxplot`) · [ ] `pl.set_color_palette` (→ `resolve_palette`) +- [ ] `pl.plot_pheno_sankey` (→ private `_sankey`) · [ ] leaked aliases `centropy`/`pentropy`/`*_tl` + +**Phase 9 (out of the package):** +- [ ] `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`) + +--- + +# DIARY + +Template per PR: **Goal · Status · What happened · Issues & fixes · Added ✓ · Removed ✓ (hard bar) · Test opportunities · Streamline opportunities · Usability.** + +## PR 0 — Contract freeze + CI scaffolding · ✅ done (branch `refactor/pr0-contract-freeze`) +- **Goal:** freeze the target public API as a `.pyi` contract; add signature-drift + import-smoke CI. Zero package behavior change. +- **What happened:** hand-authored `tcri/_contract.pyi` (27 target functions, namespaced `tl/pp/pl/diag/ut` + `TCRIModel`; locked decisions baked in — `n_samples=250`, `weighted=False`, `use_logits`, `normalize_mode`). Ported grafiti's AST signature logic into `tests/test_contract_conformance.py` (contract-parses · live-vs-contract for `IMPLEMENTED` · unimplemented worklist · import-smoke). `IMPLEMENTED={}` (nothing migrated yet). **Full suite: 26 passed, 1 skipped** (23 existing + PR0; additive, zero regressions). +- **Issues & fixes:** `tl`/`pl` share function names → used **namespace container classes** in the `.pyi` so both twins declare cleanly; parser keys as `Namespace.func`; `_strip_receiver` drops `self`/`cls` for the future `TCRIModel` method checks. +- **Added:** ✅ `tcri/_contract.pyi` ✅ `tests/test_contract_conformance.py` ✅ import-smoke +- **Removed (hard bar):** n/a (additive). The contract lists ONLY the kept surface → the 27 declared + the Removal Ledger are the two halves of "done." +- **Test opportunities:** conformance is now a live guardrail; each future PR onboards its functions into `IMPLEMENTED` and drift fails CI. TODO: add `--nbmake` (Phase 10) and multi-py import-smoke in CI yaml. +- **Streamline:** none this PR. +- **Usability:** `_contract.pyi` doubles as the one-screen human-readable target signature reference. +- **Standing Audit:** removed✅(n/a) · added✅ · tests✅(guardrail live) · streamline✅(none) · usability✅ · conformance✅ green · diary✅. +- **Committed** on branch `refactor/pr0-contract-freeze`. + +## PR 1 — Shared helpers + `_keys` · ✅ done +- **Goal:** create `_keys`/`_console`/`_stats`/`_distance`; then adopt `_keys` at every read/write site, dedup console, move stats out of utils. API unchanged. +- **What happened (foundation, done):** created the 4 helper modules — `_keys.py` (all uns/obsm/obs constants incl. NEW `GATE_PROB`/`CLASSIFIER_TEMPERATURE`; legacy keys listed for the removal step); `_console.py` (single `_ok/_info/_warn/_fin` + all ANSI aliases both spellings); `_stats.py` (`stars`/`auc_and_label_permutation`/`bootstrap_auc` + **NEW** true `hdi`, `eti`, `prob_direction`, `mann_whitney`); `_distance.py` (`kl_divergence`/`l1_distance`/`jensen_shannon` + `phenotype_distance` dispatch, single eps, bits). Added `tests/test_helpers.py` (8 tests). Suite green. +- **Issues & fixes:** **`hdi` off-by-one** — first version spanned `ceil(prob·n)+1` points, so on a right-skewed sample it returned the full range (HDI==ETI). Caught by a sanity assertion; fixed to the arviz `floor(prob·n)` window → now correctly hugs the low-mass region (`HDI=(0,0.2)` vs `ETI=(0,2.12)`). A concrete "validate the math up front" catch. +- **Added:** ✅ `_keys.py` ✅ `_console.py` ✅ `_stats.py` ✅ `_distance.py` ✅ `tests/test_helpers.py` +- **Removed / deduped (hard bar):** ✅ console dedup — deleted the **12** copied `_ok/_info/_warn/_fin` defs (metrics+preprocessing+plotting) → one `_console`. ✅ moved `stars`/`auc_and_label_permutation`/`bootstrap_auc` out of `utils` → `_stats` (plotting repointed; dead `stars` import dropped). Suite green (34 passed). +- **Test opportunities:** ✅ done — 8 unit tests for pure helpers that were previously embedded/untestable. +- **Streamline:** foundation enables deleting ~3 console dupes + 2 `dkl` copies + the `utils` stats block on adoption. +- **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 +- **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.) + +--- + +# AUDIT LOG +_(dated entries; what was audited, findings, actions)_ + +- **(PR0 ✅):** agenda + removal ledger established; standing-audit checklist defined. Contract frozen (27 fns) + conformance guardrail live. Full suite 26 passed / 1 skipped, zero regressions. +- **(PR1 ◐):** shared-helper foundation created (`_keys`/`_console`/`_stats`/`_distance`) + 8 unit tests. Caught & fixed an `hdi` off-by-one before it shipped. **Adoption pending** (dedup, stats-move, `K.*` migration) — no ledger items ticked yet; foundation is additive, suite green. Logged: key-literal test (PR1), `pl.__all__` whole-surface test (PR11). +- **(PR0+PR1 multi-agent audit — 3 lenses):** verdict FIX. Caught a real regression — the `K.*` find/replace over-reached into **10** display/warning/docstring strings (`register_model`/`load_tcri_session` printed `"K.X_LOGITS"` etc.). **Fixed:** restored readable key text in all 10 (AST-span, delimiter-safe); made the key-literal guard **AST-based** (checks real subscripts/`.get`, ignores prose); removed 3 dead `utils` imports the audit flagged. Suite 35 passed. Two non-blocking items deferred to `REFACTOR_NOTES` (contract↔api-doc reconciliation; helper-name canonicalization) — noted in the PR body. diff --git a/docs/contract/REFACTOR_HISTORY.md b/docs/contract/REFACTOR_HISTORY.md new file mode 100644 index 0000000..18bb396 --- /dev/null +++ b/docs/contract/REFACTOR_HISTORY.md @@ -0,0 +1,115 @@ +# TCRI Refactor — History, Decisions, and Hard Rules + +This is the single source of truth for **what we decided**, **the order we did things**, **what went +wrong**, and **the rules that must never be violated again**. Read this before touching any refactor doc. + +--- + +## 0. HARD RULES (never violate) + +1. **The `example/` notebooks are DISPOSABLE.** They call an old/divergent API and will be rewritten. + **Never read them, never rely on them, never use them as evidence for ANY decision** — not disposition, + not "sufficiency," not "is it used," not deletion-safety. A notebook using a function is **NOT** a reason + to keep it. +2. **Non-core = DROP (delete).** **Nothing moves to `examples/`.** `examples/` is not a destination and is + not consulted for anything. +3. **Disposition is decided by ONE question: "is it core?"** per the settled design (§2). Core = the model, + the joint-distribution engine, the four metrics + their plots, session I/O, PPC diagnostics, shared helpers. + Everything else is dropped. +4. **Sufficiency means:** does a *kept* function's arguments cover *its own responsibility*? Never "did a + notebook pass this arg." +5. **Never write a workflow prompt that runs a "caller census over notebooks."** That instruction is what + caused the leak. +6. **Settled decisions are not re-litigated.** They are in §2. + +--- + +## 1. Chronology of what we did (in order) + +1. **Full codebase read** → catalogued every function (metrics/preprocessing/plotting/model/utils). Found: + no `__all__` anywhere; ANSI/console helpers triplicated; `tl`/`pl` twin names; phenotype-prob computed 3×; + several broken/dead functions; the joint-distribution engine as the substrate. +2. **Target API design** (chat). Created `docs/contract/` with `build_tcri_contract.py` → `tcri_api_contract.*` + and `build_tcri_depgraph.py` → `tcri_dependency_map.*`. **(CLEAN — code-based only.)** +3. **scvi-tools / scverse investigation.** Decided **Door A** = standalone scverse-ecosystem package; + Door C (scvi-hub) later; Door B (in-tree `scvi.external`) a possible stretch. Structural reference = the + `../grafiti` package (one file per topic, private `_state`/`_compute`, explicit `__all__` NOT `import *`, + `diagnostics` returns DataFrames). +4. **Package validation.** `import tcri` OK; 23 tests pass; the notebooks call a **drifted old API** + (`phenotypic_entropies`/`clonotypic_entropies` plural, `probability_ternary`, `phenotypic_entropy_delta`, + `ml.JointProbabilityDistribution`). **Decision: notebooks are a mess to be REWRITTEN clean — never a + proving ground.** (This is exactly why Rule 1 exists.) +5. **131-function inventory** — 9-agent workflow → `tcri_function_inventory.md` (+ `tcri_inventory_data.json`, + `build_tcri_inventory.py`), labeled from CODE, critic-verified 0 missing. **(CLEAN input.)** + *Caveat: its plotting-triage used a "move to examples" disposition — now BANNED by Rule 2; those become DROP.* +6. **Design decisions locked** (see §2). +7. **← LEAK POINT. 13-agent draft→audit→synthesize workflow** → `tcri_api_and_responsibilities.md` + + `tcri_implementation_plan.md`. The math/stats resolution, prior-vs-mean answer, and GPU architecture it + produced are CLEAN. **But my audit prompts instructed a "caller census over `example/` + `docs/` notebooks," + which resurrected `gene_entropy` (→examples) and `probability_ternary` (→kept) *because notebooks use them*, + and baked a "deletion-safety gated on notebook callers" rule into §11.** +8. **Consistency/sufficiency** — 6-agent workflow → `tcri_arg_consistency_sufficiency.md`. **Amplified the + leak** (sufficiency measured against notebooks); the consistency findings themselves are clean. + +## 1b. What is CLEAN vs CONTAMINATED + +- **CLEAN, keep:** `tcri_api_contract.*`, `tcri_dependency_map.*`, `tcri_function_inventory.md` (+ data) + *(apply examples→DROP)*. And, *as salvageable content*: the math/stats resolution, prior-vs-mean answer, GPU + architecture, metric design, layout, rename map — all inside the quarantined docs but not themselves contaminated. +- **CONTAMINATED, quarantined:** `tcri_api_and_responsibilities.md`, `tcri_implementation_plan.md`, + `tcri_arg_consistency_sufficiency.md`, `tcri_refactor_audit_data.json`, `tcri_arg_audit_data.json`. + The contamination is confined to the **deletion/disposition** sections (API §11; Plan §3 + Phase 0 + risk rows) + and the **sufficiency** findings. + +--- + +## 2. Settled design decisions (canonical — do not re-litigate) + +**Target:** standalone scverse package (Door A). Layout mirrors grafiti; explicit `__all__` re-export, **no `import *`**; +`_keys.py` now, `@tl_result` uns-cache deferred (build toward it). + +**Engine (first-class FUNCTION, not an object):** `tcri.tl.joint_distribution(adata, *, covariate=None, groupby=None, +n_samples=0, use_logits=True, clones=None, temperature=1.0, random_state=None) -> DataFrame`. Re-exported top-level +as `tcri.joint_distribution`. `covariate=None` → all covariate values in one shared-draw pass. Unifies the two old engines. + +**Four metrics (tl↔pl twins; `groupby` = aggregation unit, `splitby` = comparison cohort — BOTH kept, distinct):** +`clonotypic_entropy`, `phenotypic_entropy`, `mutual_information`, `phenotypic_flux` (renamed from `flux`). +`compare_groups` is the single public stats helper (replaces `mi_compare`/`*_table`/`*_delta`). **No `*_delta`/`*_table`.** + +**Sampling:** `n_samples=0` = deterministic **posterior-MEAN** point estimate (`E_q[p_ct] = get_p_ct() = uns[tcri_p_ct]`); +`n_samples>0` = draws from the exact guide `Dirichlet(clamp(local_scale·m̃, 1e-3))`. Drop `point_estimate=`. `n_samples=0` +and `n_samples>0`-mean are **different estimators** (Jensen gap) — documented, never asserted equal. + +**Prior-vs-mean RESOLVED:** the point estimate is the closed-form posterior mean; the **prior is rejected** (argmax-hard-label +init → leakage). The old `posterior=` flag was never prior-vs-posterior — it's a logits-mixing switch → **renamed `use_logits`**. + +**Model (ml):** `TCRIModel.setup_anndata / train / get_latent_representation / predict (was get_cell_phenotype_probs) / +get_p_ct`. Streamline: `setup_anndata → TCRIModel → train → model.to_anndata` (thin; writes minimal canonical state; +no manager-in-uns; kills `write_adata_safely`). + +**pp** shrinks to `register/to_anndata`, `group_singletons` (separate), `clone_size`. + +**diagnostics (diag):** PPCs + model validation, returns DataFrames. Seeds: `joint_distribution_ppc` (fixed +`compare_joint_distribution`), calibration, reconstruction PPC; relocate `loss` (was plot_loss), `archetypes`. + +**Shared helpers:** `_keys.py`, `_console.py` (leveled colored logging via scanpy verbosity; drop raw ANSI prints; +**drop `_ascii_hist` and all `graph=`/ASCII paths**), `_stats.py`, `_distance.py`. Sankey primitives → private `plotting/_sankey.py`. + +**American spelling everywhere** (`normalized`, `normalize_mode`, `color`). Reproducibility via a seeded `torch.Generator` +named `random_state` (the old `np.random.seed` was a no-op). + +**PGM** (`build_nested_tcri_pgm`/`draw_tcri_pgm_nested`) → `docs/`, out of the package. + +**DROP (delete) — NOT to examples:** `clonality`, `probability_distribution`, `bayesian_mutual_information`, `polar_plot`, +`compare_phenotypes`, `probability_ternary`, `gene_entropy`, `top_clone_umap`, `clone_size_umap`, +`plot_phenotype_probabilities`, all `*_table`/`*_delta`, `mi_compare`, `clonotypic_entropy_base`, `ridge_delta_entropy`, +`classify_phenotypes`, `get_latent_embedding`, `group_small_clones`, `register_probability_columns`, +`remove_meaningless_genes`, `clone_fraction`, `dkl` (→ `_distance.kl_divergence`), `probabilities`, `_ent`, +`SankeyNode.hex_to_rgb`. (Deletion decided by "not core" — never by notebook callers.) + +--- + +## 3. Open (still parked — not decided by this recovery) + +- Whether to keep a `posterior=`/prior axis at all beyond `use_logits` (audit leaned: no). +- Adopt `@tl_result` uns-cache now vs later (recommend: `_keys.py` now, decorator later). diff --git a/docs/contract/_quarantine/README.md b/docs/contract/_quarantine/README.md new file mode 100644 index 0000000..fe4eb31 --- /dev/null +++ b/docs/contract/_quarantine/README.md @@ -0,0 +1,3 @@ +# QUARANTINE — contaminated by the notebook-caller-census leak + +These docs (from the 13-agent math-check workflow onward) baked a 'deletion gated on notebook callers' rule that resurrected dropped functions (gene_entropy, probability_ternary). Their DISPOSITION sections are WRONG. Their math/stats resolution, prior-vs-mean answer, GPU architecture, and metric design are CLEAN and may be salvaged during clean re-implementation. NOT authoritative. See ../REFACTOR_HISTORY.md. diff --git a/docs/contract/_quarantine/tcri_api_and_responsibilities.md b/docs/contract/_quarantine/tcri_api_and_responsibilities.md new file mode 100644 index 0000000..697cd0a --- /dev/null +++ b/docs/contract/_quarantine/tcri_api_and_responsibilities.md @@ -0,0 +1,803 @@ +# TCRI — API & Function Responsibilities (Final) + +*The definitive, publishable specification for the refactored, grafiti-mirrored `tcri` package (Door A: standalone scverse package). It folds the API-surface draft and the math/stats draft into a single artifact and applies every fix from the plan-correctness, math/stats, prior-vs-mean, missing-links, and GPU/optimization audits **inline** — no known-wrong math survives below. For each function you get: exact final signature and module home; responsibility; the precise math/stats it performs; how every argument dictates that math; and the exact return shape per argument combination.* + +--- + +## 0. Conventions, notation, and resolved decisions + +### 0.1 Layout principle + +Mirrors `grafiti`: one file per topic (never a monolith), private cross-cutting helper modules, explicit `__all__` re-export, **never `import *`**. Five view directories map to short handles exactly as grafiti does — `model→ml`, `tools→tl`, `preprocessing→pp`, `plotting→pl`, `diagnostics→diag` — plus `utils→ut` (tcri-specific session I/O) and a new private `_compute/` numeric+device seam. The `tl` view lives in `tcri/tools/` (grafiti `tools↔tl`), superseding the earlier working name `metrics/`. + +### 0.2 Notation + +| Symbol | Meaning | +|---|---| +| $P$ | number of phenotypes (columns of every joint) | +| $C$ | number of clonotypes (rows of a joint at one covariate) | +| $M$ | number of covariate values (e.g. timepoints) | +| $\mathrm{ct}$ | a $(\text{clonotype},\text{covariate})$ group; `ct_count` = number of them | +| $p_{ct}\in\Delta^{P}$ | learned per-`ct` phenotype distribution; `uns[K.P_CT]`, shape `(ct_count, P)` | +| $m$ | $=\text{normalize}(q\_p\_ct\_raw)=\mathbb{E}_q[p_{ct}]=$ `uns[K.P_CT]` | +| $\ell_i\in\mathbb{R}^P$ | per-cell classifier logits, `obsm[K.X_LOGITS]` (already scaled by classifier temperature) | +| $s$ | local scale, `uns[K.LOCAL_SCALE]` (Dirichlet total concentration; controls draw spread only) | +| $g$ | `gate_prob`, `uns[K.GATE_PROB]` (scalar $\in[0,1]$, or `None` → additive) | +| $\tau_{\text{cls}}$ | classifier temperature, `uns[K.CLASSIFIER_TEMPERATURE]` | +| $T$ | analysis-time `temperature` argument | +| $\varepsilon$ | numerical floor; values stated in situ | + +Index maps: `uns[K.CT_TO_COV]`, `uns[K.CT_TO_C]`, `uns[K.CT_ARRAY]` (per-cell ct index), `uns[K.COV_ARRAY]` (per-cell covariate index). + +### 0.3 The one substrate every metric reads + +Training learns a variational Dirichlet posterior per `ct` row: + +$$q(p_{ct})=\mathrm{Dirichlet}(\alpha),\qquad \alpha=\mathrm{clamp}(s\cdot m,\ \min=10^{-3}),\quad m=\text{normalize}(q\_p\_ct\_raw)\in\Delta^{P}.$$ + +Where the clamp is **inactive** (the common case), $\alpha=s\cdot m$, so $\sum_k\alpha_k=s$ and the mean is closed-form: + +$$\mathbb{E}_q[p_{ct}]=\frac{\alpha}{\sum_k\alpha_k}=m=\texttt{get\_p\_ct()}=\texttt{uns[K.P\_CT]}.$$ + +The local scale $s$ **cancels in the mean** and matters only for the spread of draws. + +> **Persisted-substrate decision (fixes the double-tempering bug).** `uns[K.P_CT]` stores the **raw** normalized posterior mean $m=\text{normalize}(q\_p\_ct\_raw)$ with **no** guide-temperature and **no** analysis-temperature baked in. `TCRIModel.get_p_ct()` returns exactly this at its default `guide_temperature=1.0`. The analysis-time `temperature` argument (§0.9) is therefore the **single** tempering knob; it is never composed on top of a pre-tempered vector. + +### 0.4 RESOLVED — the point estimate is the closed-form posterior MEAN (prior vs mean vs MAP vs draw) + +This is the audit's required decision. `n_samples=0` returns the **closed-form variational posterior mean** $\mathbb{E}_q[p_{ct}]=m=\texttt{uns[K.P\_CT]}=\texttt{get\_p\_ct()}$ (temperature-adjusted per §0.9). It is read directly and **never sampled**. + +Options analyzed and their disposition: + +| Option | Verdict | Reason | +|---|---|---| +| **(a) Closed-form posterior mean** $\mathbb{E}_q[p_{ct}]=m$ | **ADOPTED** | Exact, closed-form, deterministic, reproducible; already stored in `uns[K.P_CT]` (zero sampling cost). It is the Bayes point estimator under squared-error loss and lies in the simplex interior. Invariant to $s$ (which only sets spread). | +| (b) Generative prior `clone_phen_prior` / archetype `mixture_concentration` | **REJECTED** | It is guide **initialization** / generative anchor, not what training learned; built from argmax hard labels (leakage/circularity); indexed at clone/archetype level, not the `ct` level metrics need. No metric may read it. | +| (c) MAP / posterior mode $(\alpha-1)/(\sum\alpha-P)$ | **REJECTED** | $\alpha_k=s\,m_k$ is routinely $<1$ (small $s$, $m_k<1$), so the mode sits on the simplex boundary / is undefined — unstable, discontinuous in $s$. | +| (d) Mean of many Dirichlet draws | **REJECTED** | Converges to (a) only up to Monte-Carlo noise; a strictly Rao-Blackwell-dominated, non-reproducible estimator of a quantity available in closed form. | + +**Fate of the `posterior=` argument.** The *only* real difference between today's two engines was never prior-vs-posterior — both already used the posterior mean of $p_{ct}$; neither ever touched the generative prior. The real axis is **whether per-cell classifier logits $\ell_i$ are folded in**. Therefore: + +- **`posterior=` is DELETED from all four metrics and from `phenotypic_flux`.** They always use the learned posterior (mean at `n_samples=0`, draws at `n_samples>0`) and, given an `adata`, compute the joint with logits folded in. +- **On the engine `joint_distribution`, the flag survives but is REDEFINED and RENAMED to `use_logits`** (alias `cell_informed=`), replacing both `posterior=` and the old `combine_with_logits=`. It is a **classifier-mixing switch**, not a prior/posterior switch: `use_logits=True` folds per-cell logits into $\log(\text{base})$ exactly like `predict()`; `use_logits=False` returns the `ct`-level table directly. **Both branches use the posterior mean/draws of $p_{ct}$; neither ever touches the generative prior.** The dead `mutual_information(posterior=False) → NotImplementedError` branch and the "prior path" are removed, not implemented. + +### 0.5 Uniform sampling convention (`n_samples`) + +| `n_samples` | Operation | +|---|---| +| `0` | **Deterministic point estimate.** Use the posterior mean $m$ (temperature-adjusted); **no Dirichlet draw**; bit-reproducible on repeat calls. Fixes today's latent bug where `mutual_information`/`flux` at `n_samples=0` returned *one random draw*. | +| `N>0` | **$N$ i.i.d. posterior draws** from the **exact guide posterior** $p_{ct}^{(s)}\sim\mathrm{Dirichlet}\big(\mathrm{clamp}(s\cdot\tilde p_{ct},\ \min=10^{-3})\big)$, seeded (§0.11). Adds a sample axis; posterior mean/HDI of the functional fall out. | + +The old `point_estimate=` argument is **deleted**; `n_samples` is the only point-vs-draws knob. + +> **Clamp fix (blocking).** Draws use the guide's clamped concentration `clamp(local_scale·m̃, 1e-3)` — **not** the bare `local_scale·m̃` or `local_scale·m̃+1e-8` variants that appear in today's three inconsistent engines and summarize a distribution the model never learned. See §0.10 for the induced (documented, intentional) mean discrepancy on committed clones. + +### 0.6 Estimator honesty — plug-in vs posterior-mean (Jensen gap) + +Entropy, MI, and KL/L1 flux are **nonlinear** functionals of $p$, so $\text{metric}(\mathbb{E}_q[p])\neq\mathbb{E}_q[\text{metric}(p)]$; the difference is a **Jensen gap**, *not* Monte-Carlo noise. Consequently the two numbers below are **different estimators by design** and must be documented and tested as such: + +- **`n_samples=0`** computes the **plug-in-at-posterior-mean** estimator $\text{metric}(m̃)$. +- The **`mean` summary column of `n_samples>0`** estimates the **posterior mean of the functional** $\mathbb{E}_q[\text{metric}(p)]$. + +Directions of the gap (Shannon entropy concave; L1/KL flux convex): + +| Metric | Relationship | +|---|---| +| clonotypic / phenotypic entropy | plug-in $\ge$ posterior-mean (over-estimates) | +| phenotypic flux (L1, KL) | plug-in $\le$ posterior-mean (under-estimates); a clone with no real shift reads exactly $0$ at `n_samples=0` but strictly $>0$ in the `n_samples>0` mean | +| mutual information $=H(\phi)-H(\phi\mid c)$ | gaps partly cancel; **sign indeterminate**, magnitudes differ | + +Additionally, because draws use the **clamped** concentration while the `n_samples=0` base is the **unclamped** mean $m$, the two also differ on **committed clones** (where $s\,m_k<10^{-3}$) by a second, clamp-induced term. **No conformance test may assert `n_samples=0 == mean(n_samples>0)`.** Docstrings label the two estimators distinctly. + +### 0.7 Uniform return-shape rule + +| `groupby` | `n_samples` | Return | +|---|---|---| +| unset | `0` | scalar (MI, single-clone flux) or `Series` (entropies over phenotypes/clones) | +| unset | `N>0` | draw array with a sample axis **+** summary columns `mean, sd, hdi_low, hdi_high` | +| set | `0` | tidy `DataFrame`, one row per group [× phenotype / × clone] | +| set | `N>0` | tidy `DataFrame`, one row per group [× phenotype / × clone] + `mean, sd, hdi_low, hdi_high` | + +> **`p_gt` fix (blocking).** `p_gt`/`P(>0)` is **removed from every single-metric summary**. Entropy, MI, L1 and KL flux are all **$\ge 0$**, so `P(draw>0) ≈ 1` always and is vacuous. A signed-direction probability is emitted **only** by `tl.compare_groups` on a between-group **difference** $\Delta$ (§7.6), where it is meaningful. + +> **HDI fix.** The interval columns are a **true highest-density interval** (`hdi_low/hdi_high`), i.e. the narrowest interval containing `hdi_prob` mass — **not** the equal-tailed `np.percentile(x,[2.5,97.5])` mislabeled "HDI" in today's code. For the bounded, right-skewed entropy/flux posteriors (mass piled against the boundary for committed clones) the equal-tailed interval is materially wrong. HDIs from few hundred draws near a boundary are documented as unstable. + +> **Draw-coherence rule (correctness).** For `n_samples>0`, **all clones within one sample share the same $p_{ct}$ draw** (one coherent joint per sample). Metrics iterate the `sample_id` level and compute the full-joint metric per draw, then summarize — never independent per-clone draws. + +### 0.8 The `weighted` axis — removed (uniform-clonotype prior) + +`weighted` is **dropped from every public signature** (engine and all four metrics). Each clonotype is one unit on the simplex regardless of cell count — every entropy/MI/flux is a **repertoire-level (per-clonotype) statistic**, stated in each docstring. This removes (i) the current weight-lookup bug (a `ct`-indexed `Counter` keyed with clone indices) and (ii) the inconsistency where `joint_distribution` normalized the whole table to sum 1 while `joint_distribution_posterior` returned un-normalized counts. + +> **Behavior-change note (changelog).** `pl.mutual_information` currently defaults `weighted=True`; removing it flips displayed MI from cell-weighted to per-clonotype. `tl.mutual_information` never accepted `weighted`, so only the plotting default changes user-visible numbers. + +### 0.9 Temperature — single knob, one consistent placement + +$T$ power-tempers the base **once**, identically in the mean and draw paths: + +$$\tilde p_{ct}=\mathrm{softmax}\!\Big(\tfrac1T\log(m+\varepsilon)\Big)=\frac{m^{1/T}}{\sum_\phi m_\phi^{1/T}},\qquad \varepsilon=10^{-8}.$$ + +$T=1$ is the identity (renormalization only); $T<1$ sharpens, $T>1$ flattens. This fixes today's split where the two engines tempered at different stages. + +- For `use_logits=True`, the combined per-cell logit is divided by $T$ **once**: $P(\phi\mid i)=\mathrm{softmax}\big(\text{combine}(\ell_i,\log\tilde b)/T\big)$. At **$T=1$** this reproduces `predict()` **bit-for-bit** (classifier temperature $\tau_{\text{cls}}$ is already baked into $\ell_i$; no second division). $T\neq1$ is an analysis-time temper that intentionally diverges from `predict()`; documented. +- For `n_samples>0`, draws are centered on the **re-tempered** $\tilde p_{ct}$; docstrings state that $T\neq1$ makes the sampled distribution a re-tempered object, not the raw learned posterior. + +### 0.10 Reproducibility / seeding + +`random_state` (`int | numpy.Generator | torch.Generator | None`) is added to `joint_distribution`, all four metrics, `compare_groups`'s bootstrap, and the sampling `diag` functions. It seeds a **`torch.Generator`** (and, on the GPU path, the CUDA RNG) because draws are `torch` Dirichlet — fixing the standing no-op where `seed` only touched `np.random`. `n_samples=0` is deterministic regardless of `random_state`. + +--- + +## 1. Package tree + +``` +tcri/ + __init__.py # explicit re-export + sys.modules aliases (tl/pp/pl/ml/ut/diag); top-level joint_distribution; NO import * + _keys.py # single source of every uns/obsm/obs key string (constants only) + _console.py # leveled, silenceable logging over scanpy.logging (no raw ANSI, no _ascii_hist) + _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) + _reduce.py # batched entropy / mutual-information / distance reductions over the stack + model/ # ml + _model.py # TCRIModel + _module.py # TCRIModule (pyro model/guide, get_latent, get_p_ct) + _priors.py # MixtureDirichlet, VampPrior + _classifier.py # PhenotypeClassifier + _training.py # UnifiedTrainingPlan, build_archetypes + preprocessing/ # pp + _register.py # registration writers behind TCRIModel.to_anndata (all private) + _clones.py # group_singletons, clone_size + tools/ # tl (mirrors pl by filename) + _joint.py # joint_distribution (THE ENGINE) + _entropy.py # clonotypic_entropy, phenotypic_entropy + _mutual_information.py # mutual_information (+ private _mi_from_joint) + _flux.py # phenotypic_flux + _compare.py # compare_groups (public group-comparison orchestrator) + plotting/ # pl + _base.py # _metric_boxplot, _finish + _colors.py # tcri_colors, resolve_palette + _entropy.py # clonotypic_entropy, phenotypic_entropy + _mutual_information.py # mutual_information + _flux.py # phenotypic_flux (sankey) + _ternary.py # probability_ternary + _sankey.py # SankeyNode, _phenotype_mass_per_clone (private) + diagnostics/ # diag (NEW) + _ppc.py # joint_distribution_ppc, phenotype_calibration, reconstruction_ppc, permutation_null + _training.py # loss, archetypes + utils/ # ut + _session.py # save_tcri_session, load_tcri_session (+ private helpers) +``` + +`examples/` (outside the package): `top_clone_umap`, `clone_size_umap`, `plot_phenotype_probabilities`, `gene_entropy`, `polar_plot`, rewritten notebooks. `docs/`: the model PGM (`build_nested_tcri_pgm` / `draw_tcri_pgm_nested`). + +--- + +## 2. Top-level `__init__.py` and the `__all__` story + +`__all__` is declared at **both** levels (grafiti pattern). Every impl module declares its own `__all__`; every view `__init__` imports symbols by name and re-declares an aggregate `__all__`. The root imports the six view packages, aliases them into `sys.modules`, and re-exports `joint_distribution` for prominence. **No `import *` anywhere** — numpy/pandas/torch and every `_helper` stay unexported; GPU libs are never imported at module top (§4.3). + +```python +# tcri/__init__.py +from importlib.metadata import PackageNotFoundError, version as _version +try: + __version__ = _version("tcri") +except PackageNotFoundError: + __version__ = "0.0.0+unknown" + +from . import tools as tl +from . import preprocessing as pp +from . import plotting as pl +from . import model as ml +from . import diagnostics as diag +from . import utils as ut +from .tools import joint_distribution # tcri.joint_distribution + +import sys +sys.modules.update({f"{__name__}.{m}": globals()[m] + for m in ("tl", "pp", "pl", "ml", "diag", "ut")}) + +__all__ = ["tl", "pp", "pl", "ml", "diag", "ut", "joint_distribution", "__version__"] +``` + +| View `__init__` | `__all__` | +|---|---| +| `tools/__init__.py` (`tl`) | `joint_distribution`, `clonotypic_entropy`, `phenotypic_entropy`, `mutual_information`, `phenotypic_flux`, `compare_groups` | +| `preprocessing/__init__.py` (`pp`) | `group_singletons`, `clone_size` | +| `plotting/__init__.py` (`pl`) | `clonotypic_entropy`, `phenotypic_entropy`, `mutual_information`, `phenotypic_flux`, `probability_ternary`, `tcri_colors`, `resolve_palette` | +| `model/__init__.py` (`ml`) | `TCRIModel` | +| `diagnostics/__init__.py` (`diag`) | `joint_distribution_ppc`, `phenotype_calibration`, `reconstruction_ppc`, `permutation_null`, `loss`, `archetypes` | +| `utils/__init__.py` (`ut`) | `save_tcri_session`, `load_tcri_session` | + +Private modules (`_keys`, `_console`, `_stats`, `_distance`, `_compute/*`) and every private symbol (`_mi_from_joint`, `_metric_boxplot`, `_finish`, `SankeyNode`, `_phenotype_mass_per_clone`, all `_register.py` writers) are **absent** from every `__all__`. + +--- + +## 3. Shared private helper modules + +### 3.1 `tcri/_keys.py` — canonical AnnData key registry (constants only) + +Import as `from tcri import _keys as K`; no key literal lives anywhere else. + +| Constant | Value | Slot | Meaning | +|---|---|---|---| +| `K.METADATA` | `"tcri_metadata"` | `uns` | dict: `covariate_col`, `clone_col`, `phenotype_col`, `batch_col` (single scheme; replaces dual `tcri_clone_key`/`tcri_phenotype_key`) | +| `K.PHENOTYPE_CATEGORIES` | `"tcri_phenotype_categories"` | `uns` | ordered phenotype categories | +| `K.CLONOTYPE_CATEGORIES` | `"tcri_clonotype_categories"` | `uns` | ordered clonotype categories | +| `K.COVARIATE_CATEGORIES` | `"tcri_covariate_categories"` | `uns` | ordered covariate categories | +| `K.P_CT` | `"tcri_p_ct"` | `uns` | `(ct_count, P)` **raw** posterior mean $m=\mathbb{E}_q[p_{ct}]$ (no temper baked in) | +| `K.CT_TO_COV` | `"tcri_ct_to_cov"` | `uns` | `(ct_count,)` ct→covariate index | +| `K.CT_TO_C` | `"tcri_ct_to_c"` | `uns` | `(ct_count,)` ct→clonotype index | +| `K.CT_ARRAY` | `"tcri_ct_array_for_cells"` | `uns` | `(n_obs,)` per-cell ct index | +| `K.COV_ARRAY` | `"tcri_cov_array_for_cells"` | `uns` | `(n_obs,)` per-cell covariate index | +| `K.LOCAL_SCALE` | `"tcri_local_scale"` | `uns` | scalar Dirichlet total concentration $s$ (draws only) | +| **`K.GATE_PROB`** | `"tcri_gate_prob"` | `uns` | **NEW** — scalar gate $g\in[0,1]$ or `None`; required for gate-aware `use_logits=True` parity with `predict()` | +| **`K.CLASSIFIER_TEMPERATURE`** | `"tcri_classifier_temperature"` | `uns` | **NEW** — $\tau_{\text{cls}}$; persisted for reproducibility/auditing (already baked into stored logits) | +| `K.X_LATENT` | `"X_tcri"` | `obsm` | `(n_obs, n_latent)` encoder posterior-mean latent | +| **`K.X_LOGITS`** | `"X_tcri_logits"` | `obsm` | `(n_obs, P)` classifier logits — **part of the canonical write-set** (§5.1); the `use_logits=True` engine path requires it | +| `K.X_PROBABILITIES` | `"X_tcri_probabilities"` | `obsm` | `(n_obs, P)` per-cell phenotype probabilities (`predict()`) | +| `K.PHENOTYPE_LABEL` | `"tcri_phenotype"` | `obs` | per-cell argmax hard label | +| `K.CLONE_SIZE` | `"clone_size"` | `obs` | per-cell clone cell-count | +| `K.OBS_INDICES` | `"indices"` | `obs` | per-cell integer index — **registration glue** written by `setup_anndata` (§5.1); not analysis output | + +> The AnnDataManager is **no longer** stashed in `uns` (`tcri_manager` retired) — deleting the `write_adata_safely` / `_pop_nonserializables` hack. The stash lived in `setup_anndata`, so it is removed there (§5.1). + +### 3.2 `tcri/_console.py` — leveled, silenceable logging (private) + +Thin wrappers over `scanpy.logging`; respects scanpy verbosity. Raw ANSI prints and `_ascii_hist` (and every `graph=`/ASCII-histogram code path) are **deleted**. + +| Signature | Responsibility | +|---|---| +| `info(msg, *, deep=None)` | `scanpy.logging.info`; silenced by scanpy verbosity. | +| `warning(msg)` | `scanpy.logging.warning`. | +| `success(msg)` | `scanpy.logging.hint`. | +| `done(msg="done")` | terminal completion line. | + +### 3.3 `tcri/_stats.py` — significance + posterior-comparison statistics (private) + +| Signature | Responsibility / math | +|---|---| +| `stars(p)` | $p<10^{-4}\to$`****`; $<10^{-3}\to$`***`; $<10^{-2}\to$`**`; $<0.05\to$`*`; else `ns`. | +| `mann_whitney(a, b, *, alternative="two-sided")` | Mann–Whitney $U=\sum_{i,j}\mathbb1[a_i>b_j]+\tfrac12\mathbb1[a_i=b_j]$; two-sided $p$ from the rank-sum null (`scipy.stats.mannwhitneyu`). Returns `(U, p)`. | +| `prob_direction(delta)` | **Signed-contrast only.** Given a paired difference-draw vector $\Delta$: $p_{gt}=\frac1N\sum_s\mathbb1[\Delta^{(s)}>0]$, $p_{lt}=\frac1N\sum_s\mathbb1[\Delta^{(s)}<0]$. Returns `(p_gt, p_lt)`. | +| `hdi(samples, *, hdi_prob=0.94)` | **True** highest-density interval: over sorted samples, the **narrowest** window containing $\lceil hdi\_prob\cdot N\rceil$ points. Returns `(low, high)`. Documented unstable near a bounded posterior's boundary for small $N$. | +| `summarize(samples, *, hdi_prob=0.94)` | Reducer for a **raw metric** draw vector → `{mean, sd, hdi_low, hdi_high}`. **No `p_gt`** (vacuous for $\ge0$ metrics, §0.7). | +| `auc_and_label_permutation(scores, labels, *, pos_label=None, n_perm=200_000, seed=42, max_exact=200_000)` | Observed ROC-AUC + two-sided permutation $p$: exact enumeration when $\binom{n}{k}\le$`max_exact`, else Monte-Carlo; $p_{\text{perm}}=\text{mean}(|\mathrm{AUC}_{\text{perm}}-0.5|\ge|\mathrm{AUC}_{\text{obs}}-0.5|)$. Returns `(auc, p, perm_stats, mode)`. | +| `bootstrap_auc(scores, labels, *, pos_label=None, n_boot=5000, seed=42)` | Resample cells with replacement (reject draws missing a class), recompute AUROC, return the 2.5/97.5 quantiles. Returns `np.array([lo, hi])`. | + +### 3.4 `tcri/_distance.py` — phenotype-distribution distances (private) + +Dedupes the old module-level `dkl` and `flux.dkl_func`; **one base (bits, $\log_2$) and one $\varepsilon=10^{-12}$ library-wide**, matching entropy/MI. + +| Signature | Responsibility / math | +|---|---| +| `l1_distance(p, q)` | $\sum_i|p_i-q_i|$; symmetric; range $[0,2]$ on the simplex. Defensively renormalizes inputs. Returns float. | +| `kl_divergence(p, q, *, base=2.0, eps=1e-12)` | $\mathrm{KL}(p\Vert q)=\sum_i p_i\log_2\frac{p_i}{q_i}$; clip to $[\varepsilon,1]$ then renormalize. **Asymmetric**, range $[0,\infty)$. **Single base fixed to $\log_2$ / single $\varepsilon$** (repairs the natural-log-vs-$\log_2$ and $10^{-10}$-vs-$10^{-15}$ divergence between the two dead copies). Returns float. | +| `js_divergence(p, q, *, base=2.0, eps=1e-12)` | **NEW** Jensen–Shannon $=\tfrac12\mathrm{KL}(p\Vert r)+\tfrac12\mathrm{KL}(q\Vert r)$, $r=\tfrac12(p+q)$; **symmetric, bounded $[0,1]$ bit** — the recommended symmetric shift measure. Returns float. | +| `phenotype_distance(p, q, *, metric="l1")` | Dispatcher: `"l1"`→`l1_distance`; `"kl"`/`"dkl"`→`kl_divergence` (directional, unbounded, bits); `"jsd"`→`js_divergence`; a callable `f(p,q)`; else `ValueError`. | + +--- + +## 4. `tcri/_compute/` — numeric + device seam (NEW, private) + +The engine's numeric core is written **once** as a batched, device-routable function so the acceleration is additive and reversible, and so the four metrics share one joint-draw stack. + +### 4.1 `_xp.py` — the device seam (grafiti reference, copied 1:1) + +| 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. | +| `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. +- **`_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) + +Lazy GPU imports (never at module top — `import tcri` never touches a GPU lib; the old module-top `import umap` is moved inside its function); GPU deps never in `install_requires` (CPU path always fully functional); permissive device ladder with device-count verification; `asnumpy`/`output_type="numpy"` at every boundary; `try/except` degrade-to-CPU reporting which backend ran; `float64` where parity must hold; validate-before-compute on **per-row** invariants; chunked reductions to bound host+device memory on the large `[n_samples, n_cells, P]` tensor. `n_samples=0` performs **zero** draws (closed-form read of `uns[K.P_CT]`). cuML UMAP for the one-off latent embedding sits behind a `_use_gpu(device)` gate with the `umap-learn` CPU fallback; layouts differ (both valid), documented, not claimed bit-identical. + +--- + +## 5. `tcri.ml` — model (`model/`) + +### 5.1 `model/_model.py` + +**`class TCRIModel(BaseModelClass)`** — register → build → train → extract → write. + +| Method (signature) | Responsibility | +|---|---| +| `@classmethod setup_anndata(cls, adata, *, layer=None, clonotype_key="unique_clone_id", phenotype_key="phenotype_col", covariate_key="timepoint", batch_key="patient", **kwargs)` | **Registration only** — register clonotype/phenotype/covariate/batch/count fields with scvi and store the layer. **Writes `obs["indices"]=range(n)` and registers it** (`CategoricalObsField`) — this is registration glue that `training_step`/`validation_step` consume via `batch["indices"]`; it is **not** analysis output. Invariant is **"no analysis/label `obs` mutation"** (labels/probabilities are written only by `to_anndata`). **Removes the `uns["tcri_manager"]` stash** (was here), deleting the need for `write_adata_safely`. | +| `__init__(self, adata, *, n_latent=128, n_hidden=128, n_layers=3, classifier_n_layers=3, global_scale=5.0, local_scale=3.0, prior_temperature=1.0, guide_temperature=1.0, use_enumeration=False, patience=300, classifier_hidden=128, classifier_dropout=0.1, n_pseudo_obs=10, K=10, phenotype_weights=None, gate_prob=None, kl_weight_max=1.0, guide_init_scale=10.0, classifier_temperature=1.0, **kwargs)` | Build the empirical clone→phenotype prior + KMeans archetypes + clonotype/covariate index maps + class weights, then construct/prime `TCRIModule`. Note `gate_prob=None` default ⇒ ungated model; the gate-parity guarantee is only exercised when a gate is trained. | +| `train(self, *, max_epochs=1000, batch_size=1000, lr=1e-3, reconstruction_loss_scale=1e-3, n_steps_kl_warmup=2000, **kwargs)` | 0.9/0.1 split, `UnifiedTrainingPlan`, `TrainRunner` with `elbo_validation` early stopping. | +| `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. | + +> Relocated off the model: `plot_archetypes`→`diag.archetypes`; `plot_loss`→`diag.loss`. `boost_phenotype_prior`, `use_gate` remain internal. + +### 5.2 `model/_module.py` + +**`class TCRIModule(PyroBaseModuleClass)`** — Pyro CVAE with hierarchical clonotype→(clonotype×covariate) Dirichlet priors and a phenotype classifier. *(internal)* + +| Member (signature) | Responsibility | +|---|---| +| `__init__(self, n_input, n_latent, P, n_batch, *, global_scale=10.0, local_scale=5.0, prior_temperature=1.0, guide_temperature=1.0, gate_prob=0.5, mixture_concentration=None, n_pseudo_obs=10, use_enumeration=False, classifier_hidden=128, classifier_dropout=0.1, classifier_n_layers=3, n_hidden=128, n_layers=3, class_weights=None, kl_weight_max=1.0, guide_init_scale=10.0, classifier_temperature=1.0)` | Construct encoder/decoder/classifier/VampPrior, `px_r`; register empty two-level buffers + class weights (`mixture_concentration` required). | +| `prepare_two_level_params(self, clone_phen_prior_mat, ct_to_c, ct_to_cov, ct_array_for_cells, cov_array_for_cells, *, eps=1e-6)` | Normalize/temperature the clone-phenotype prior; register two-level index buffers. | +| `model(self, x, batch_idx, log_library, ...)` | Generative: sample $p_c$ (MixtureDirichlet), $p_{ct}$ (Dirichlet centered at $p_c$), latent $z$ (VampPrior), ZINB gene obs. | +| `guide(self, x, batch_idx, log_library, ...)` | Guide: learnable Dirichlet params $q(p_c)$, $q(p_{ct})$ with **`clamp(min=1e-3)`** on the concentration (the floor the draw path must reproduce, §0.5); Normal $q(z)$ from the encoder. | +| `get_latent(self, tensor_dict) -> torch.Tensor` | Encode a batch to posterior-mean latent $z_{\text{loc}}$. | +| `get_p_ct(self) -> torch.Tensor` | Read `q_p_ct_raw` from the (process-global) param store; return the row-normalized `(ct_count, P)` posterior mean. | +| `use_gate(self) -> bool` (property) | `True` when `gate_prob is not None`. | +| `@staticmethod _get_fn_args_from_batch(tensor_dict) -> tuple` | Extract `(x, batch_idx, log_library)` from a scvi batch dict. | + +> **Param-store caveat (documented).** `get_p_ct` reads the **process-global** Pyro param store (`q_p_ct_raw`); loading two sessions in one process clobbers it. `to_anndata` and every `diag` PPC must be called immediately after the intended model's params are set; `load_tcri_session` sets the store before any `get_p_ct`/`to_anndata` call. Single-model-per-process otherwise. + +### 5.3 `model/_priors.py` + +**`class MixtureDirichlet(dist.TorchDistribution)`** *(internal)* — clonotype prior $p_c$. Members: `__init__(self, mixture_weights, concentration, validate_args=None)`; `sample`; `log_prob` (log-sum-exp of component Dirichlet log-probs); `score_parts` (returns `(log_prob, 0, 0)` → reparam-free); `__call__` (alias for `sample`). + +**`class VampPrior(torch.nn.Module)`** *(internal)* — VampPrior over $z$. Members: `__init__(self, pseudo_inputs, encoder)`; `get_mixture` (uniform `MixtureSameFamily` of `Independent` Normals); `log_prob(self, z)`; `sample`. + +### 5.4 `model/_classifier.py` + +**`class PhenotypeClassifier(nn.Module)`** *(internal)* — `__init__(self, n_latent, classifier_hidden, P, *, num_layers=3, dropout_rate=0.1, temperature=1.0)`; `forward(self, x)` returns MLP logits divided by `temperature` ($\tau_{\text{cls}}$, baked into the stored logits). + +### 5.5 `model/_training.py` + +**`class UnifiedTrainingPlan(PyroTrainingPlan)`** *(internal)* — `__init__(self, module, *, n_steps_kl_warmup=1000, reconstruction_loss_scale=1e-2, num_particles=5, optimizer_config=None, class_weights=None, **kwargs)`; `loss` (property); `configure_optimizers`; `training_step`; `validation_step` (logs `elbo_validation`). + +**Module function:** `build_archetypes(c2p_mat, *, K=10) -> tuple[np.ndarray, np.ndarray]` — KMeans-cluster clone→phenotype rows into `K` normalized archetype centroids. **Returns `(centers, labels)`** — labels are retained so `diag.archetypes` can reproduce the cluster-ordered heatmap. `K` default is **10** (aligned to `TCRIModel`, repairing the former `K=4` default mismatch). + +--- + +## 6. `tcri.pp` — preprocessing (`preprocessing/`) + +### 6.1 `preprocessing/_clones.py` — public + +| Signature | Responsibility | +|---|---| +| `group_singletons(adata, *, clonotype_key="trb", groupby="patient", target_col="trb_unique", min_clone_size=10) -> AnnData` | Collapse clones smaller than `min_clone_size` (per `groupby`) into `"Singleton_{group}"` labels in `target_col`. **Ordering invariant (documented + enforced):** any clone relabeling must run **before** `setup_anndata`/`train`, else the learned clonotype categories and `p_ct`'s `ct_to_c` map desync from `obs`; `setup_anndata` refuses registration if a later relabel is detected. | +| `clone_size(adata, *, key_added="clone_size", return_counts=False)` | Per-clone cell counts, written per cell into `obs[key_added]`. **Reads `uns[K.METADATA]["clone_col"]`** (migrated off the retired `tcri_clone_key` in the same change that stops writing it). | + +**Private inner:** `group_singletons.collapse_singleton(row)`. + +### 6.2 `preprocessing/_register.py` — private (the `to_anndata` writers) + +Called only by `TCRIModel.to_anndata`; folds in the old `register_phenotype_key` / `register_clonotype_key` / `_compute_logits_and_prior`. + +| Signature | Responsibility | +|---|---| +| `_write_metadata(adata, model)` | `uns[K.METADATA]` (single scheme) + the three category lists from the registry. | +| `_register_clonotype_key(adata, clonotype_key, *, order=None)` | Register the clonotype `obs` column + ordered categories. | +| `_register_phenotype_key(adata, phenotype_key, *, order=None)` | Register the phenotype `obs` column + ordered categories. | +| `_write_latent(adata, model, *, latent_key="X_tcri", batch_size=256)` | Encoder posterior-mean latent → `obsm`. | +| `_write_logits(adata, model, *, logits_key="X_tcri_logits", batch_size=256)` | **Per-cell classifier logits → `obsm[K.X_LOGITS]`** (canonical; required by the default engine path). | +| `_write_predictions(adata, model, *, predictions_key="X_tcri_probabilities", label_key="tcri_phenotype", batch_size=256)` | `predict()` probs → `obsm`; argmax hard labels → `obs`. | +| `_write_p_ct(adata, model)` | `p_ct`, `ct_to_cov`, `ct_to_c`, per-cell ct/cov arrays, `local_scale`, **`gate_prob`**, **`classifier_temperature`** → `uns`. | +| `_compute_logits_and_prior(model, adata, *, batch_size=256, eps=1e-8) -> tuple[np.ndarray, np.ndarray]` | Run encoder+classifier to extract per-cell logits and $\log p_{ct}$ from `get_p_ct()`. | + +--- + +## 7. `tcri.tl` — tools / metrics (`tools/`) + +### 7.1 Engine — `tools/_joint.py` + +```python +joint_distribution( + adata, *, + covariate=None, # None → ALL covariate values in one pass (shared draw) + groupby=None, + n_samples=0, + use_logits=True, # was posterior=; alias cell_informed=; classifier-mixing switch + clones=None, + temperature=1.0, + random_state=None, + device=None, +) -> pandas.DataFrame +``` +Re-exported top-level as `tcri.joint_distribution`. Unifies today's `joint_distribution` + `joint_distribution_posterior`. + +**(a) Responsibility.** Produce the clone×phenotype distribution (the substrate every metric consumes) at a covariate value from the learned variational posterior of $p_{ct}$ — a deterministic posterior-mean table or posterior draws. Provenance in `df.attrs["params"]` (and, for cache-friendliness, in a flat `_provenance` sidecar column, §7.7). + +**(b) Math.** Select `ct` rows for covariate $m$ via `ct_to_cov`; each clone maps to exactly one `ct`, so rows index by clonotype. Temperature-temper the base **once** (§0.9): $\tilde p_{ct}=\mathrm{softmax}(\tfrac1T\log(m+10^{-8}))$. + +*Base per point/draw:* +- `n_samples=0`: base $b=\tilde p_{ct}$ (posterior mean; deterministic). +- `n_samples=N`: bases $b^{(s)}\sim\mathrm{Dirichlet}\big(\mathrm{clamp}(s\cdot\tilde p_{ct},\ \min=10^{-3})\big)$, $s=$`local_scale`, seeded (§0.10). One coherent joint per sample (§0.7). + +*`use_logits=False` (ct-level table):* row $c=b$. For `n_samples=0` this is exactly $\tilde p_{ct}$; at $T=1$ it equals `uns[K.P_CT]` restricted to the covariate — the clean closed-form identity used in tests. + +*`use_logits=True` (fold per-cell logits; matches `predict()`):* per cell $i$ with clonotype $c(i)$, group $\mathrm{ct}(i)$, base $b_{\mathrm{ct}(i)}$: +$$P(\phi\mid i)=\mathrm{softmax}\!\Big(\tfrac1T\,\mathrm{combine}\big(\ell_i,\ \log(b_{\mathrm{ct}(i)}+\varepsilon)\big)\Big),\quad +\mathrm{combine}=\begin{cases}g\,\ell_i+(1-g)\log b & g=\texttt{gate\_prob}\neq\text{None}\\ \ell_i+\log b & \text{additive}\end{cases}$$ +then $J[c,\phi]=\sum_{i\in c}P(\phi\mid i)$, row-normalize. At $T=1$, `n_samples=0`, this equals `predict()` aggregated per clone — **bit-for-bit**, gate-aware (fixes the standing disagreement where metrics used one Dirichlet *draw* and *never* applied the gate). + +*`covariate=None`:* compute the joint for **all** covariate values from a **single shared draw** per sample (the draw-once invariant, §7.8), stacking a covariate axis. + +*`groupby=g`:* restrict by cell/clone masks into the **full-space** `uns` arrays (never by slicing the AnnData — see the guard note below), computing per group value and stacking a group axis. + +**(c) Arguments → math.** + +| Argument | Effect | +|---|---| +| `covariate` | Selects `ct` rows via `ct_to_cov`. `None` → all covariates in one shared-draw pass. | +| `groupby` | Separate joint per group value (adds a group axis), implemented by **restriction over full adata**, not slicing. Requires the cell-informed path or a clone-constant key (see semantics note). | +| `n_samples` | `0` → posterior-mean table (deterministic); `N` → $N$ clamped-Dirichlet draws. Only place `local_scale` enters. | +| `use_logits` | `True` → fold logits with $\log b$ (gate-aware), aggregate per clone, row-normalize; `False` → `ct`-level $\tilde p_{ct}$ rows. Neither is the generative prior. | +| `clones` | Filters rows to the listed clonotypes; with `use_logits=True` also restricts aggregated cells; final reindex to the exact list (absent clones → dropped, **not** all-zero rows — see §7.2 fix). | +| `temperature` | $T$ tempers the base once (§0.9). $T=1$ identity; at $T=1$, `use_logits=True` reproduces `predict()`. | +| `random_state` | Seeds the torch (and CUDA) Dirichlet generator for `n_samples>0`. Ignored at `n_samples=0`. | +| `device` | Routes the numeric core through `_compute/_xp` (CPU / torch-CUDA / cupy); result is always host numpy. | + +**(d) Return shape** — `pandas.DataFrame`, columns = phenotype categories. + +| `covariate` | `groupby` | `n_samples` | Index / axes | +|---|---|---|---| +| set | unset | `0` | rows = clonotype id; `(C_m, P)` | +| set | unset | `N>0` | MultiIndex (clonotype, `sample_id`); `(C_m·N, P)` | +| set | set | `0` | MultiIndex (group, clonotype) | +| set | set | `N>0` | MultiIndex (group, clonotype, `sample_id`) | +| `None` | — | — | adds a leading covariate level to any of the above | + +> **groupby ↔ alignment guard (blocking fix).** `joint_distribution_posterior` hard-raises if per-cell `uns[...array_for_cells]` lengths $\neq$ `n_obs`. Passing a **sliced** AnnData (today's `tcri_boxplot` pattern) trips this. groupby is therefore implemented by **positional cell/clone masks into the full-space `uns` arrays** + `clones=`, never by handing a slice to the engine. `_metric_boxplot` (§8.5) is rewritten off the slice-and-call pattern. + +> **groupby ↔ covariate semantics.** `p_ct` is indexed by `ct=(clonotype, covariate)` only. A `groupby` key that is **not** functionally determined by clonotype-at-fixed-covariate (e.g. a tissue cross-cutting one clone) is unrepresentable in the `use_logits=False` table and is only well-defined on the cell-informed `use_logits=True` path. The engine **requires the cell-informed path for such keys**, or requires the key be clone-nested / constant within a clone×covariate; it errors/warns on `use_logits=False` + a non-clone-determined groupby. The whole per-group scheme assumes **clones are disjoint across groups** (a TCR clone never spans two patients) — stated explicitly. + +`__all__ = ["joint_distribution"]` + +### 7.2 `tools/_entropy.py` — `clonotypic_entropy` + +```python +clonotypic_entropy( + adata_or_jd, *, + covariate=None, groupby=None, + n_samples=0, temperature=1.0, + clones=None, normalized=True, n_clones_ref=None, + random_state=None, device=None, +) -> float | pandas.Series | pandas.DataFrame +``` + +**(a) Responsibility.** For each phenotype $\phi$ (at covariate $m$), the normalized Shannon entropy of the distribution over clonotypes carrying that phenotype, $H[P(c\mid\phi,m)]$ — spread of a phenotype across clones. **Repertoire-level (uniform-clonotype prior; §0.8).** + +**(b) Math.** From joint $J$, take column $\phi$ over the **supported** clones only (absent/all-zero clones excluded — see fix), normalize, take entropy in bits: +$$v_c=\frac{J[c,\phi]}{\sum_{c'\in\text{supp}}J[c',\phi]},\qquad H_\phi=-\sum_{c\in\text{supp}} v_c\log_2 v_c.$$ +If `normalized`: divide by $\log_2 C_{\text{den}}$ where $C_{\text{den}}$ = number of **supported** clones (default), or `n_clones_ref` if given (a fixed reference for cross-group comparability). No division when $C_{\text{den}}\le1$. Base fixed to 2. + +- **Estimator (§0.6):** `n_samples=0` = plug-in $H_\phi(m̃)$; `n_samples>0` `mean` = $\mathbb{E}_q[H_\phi]$ (plug-in $\ge$ posterior-mean for entropy). Documented as distinct. +- **Fixes:** deterministic `n_samples=0` (no forced draw); `normalized` exposed (was hard-normalized); base fixed to 2; **absent/zero-support clones excluded before normalizing** (no $\varepsilon$-clip fabricating uniform mass or inflating $C$). + +**(c) Arguments → math.** + +| Argument | Effect | +|---|---| +| `adata_or_jd` | AnnData → compute $J$ internally via §7.1 (`use_logits=True`; `covariate` required); precomputed joint DataFrame → skip to entropy (fast path — valid only at `n_samples=0`, `groupby=None`; `clones` just re-filters; else `ValueError`, §7.9). | +| `covariate` / `groupby` | Condition $m$; per-group entropy → tidy rows (group × phenotype). | +| `n_samples` | `0` → plug-in per phenotype; `N` → per-draw + summary. | +| `temperature` | Tempers $J$ before the column is read. | +| `clones` | Restricts the clone set → changes support and the default $\log_2 C_{\text{den}}$. | +| `normalized` / `n_clones_ref` | `True` → divide by $\log_2 C_{\text{den}}$ (range $[0,1]$); `n_clones_ref` fixes the denominator for comparability; `False` → raw bits. | +| `random_state` / `device` | Seeding / backend routing for `n_samples>0`. | + +**(d) Return shape.** `n_samples=0`, no `groupby` → `Series` over phenotypes; `n_samples>0`, no `groupby` → per-phenotype `mean, sd, hdi_low, hdi_high`; `groupby` → tidy DataFrame row per (group, phenotype) [+ summary]. **Absent phenotype → `NaN`, not 0.** + +> **Comparability note.** Because the default denominator is group-specific, normalized clonotypic entropy is **within-group** unless `n_clones_ref` (a common denominator) is supplied. The `pl` twin defaults cross-group plots to a common `n_clones_ref`. + +### 7.3 `tools/_entropy.py` — `phenotypic_entropy` + +```python +phenotypic_entropy( + adata_or_jd, *, + covariate=None, groupby=None, + n_samples=0, temperature=1.0, + clones=None, normalized=True, + random_state=None, device=None, +) -> float | pandas.Series | pandas.DataFrame +``` + +**(a) Responsibility.** For each clonotype $c$, the normalized Shannon entropy of its phenotype distribution $H[P(\phi\mid c,m)]$ — plasticity vs commitment. + +**(b) Math.** Row $c$ = $P(\phi\mid c)$; normalize over phenotypes, entropy in bits, divide by $\log_2 P$ if `normalized` and $P>1$: +$$p_\phi=\frac{J[c,\phi]}{\sum_{\phi'}J[c,\phi']},\quad H_c=-\sum_\phi p_\phi\log_2 p_\phi.$$ +Estimator convention as §0.6 (plug-in at `n_samples=0`). **Critical bug fix:** a clone present in `obs` but with **zero posterior mass** returns **`NaN`** and is excluded — it is **not** reindexed to zeros, $\varepsilon$-clipped, and normalized to uniform → spurious $H=1.0$. Base fixed to 2; `normalized` exposed. + +**(c) Arguments → math.** As §7.2, except the normalizer is $\log_2 P$ (depends on $P$, not clone count), so `clones` does not change the divisor; $P\le1\Rightarrow$ divisor 1. + +**(d) Return shape.** `n_samples=0`, no `groupby` → `Series` over clonotypes; `n_samples>0` → per-clone `mean, sd, hdi_low, hdi_high`; `groupby` → tidy DataFrame row per (group, clone) [+ summary]. + +### 7.4 `tools/_mutual_information.py` — `mutual_information` (+ kernel) + +```python +mutual_information( + adata_or_jd, *, + covariate=None, groupby=None, + n_samples=0, temperature=1.0, + clones=None, normalized=True, + normalize_mode="min", # CHANGED default (was "average") + random_state=None, device=None, +) -> float | numpy.ndarray | pandas.DataFrame +``` + +**(a) Responsibility.** $I(c;\phi\mid m)$ in bits — strength of clone–phenotype coupling. Optionally normalized to $[0,1]$. + +**(b) Math — kernel `_mi_from_joint(pxy, *, normalized, mode="min") -> float`.** Normalize the whole joint: $\text{pxy}=J/\sum J$; marginals $p_x=\sum_\phi\text{pxy}$, $p_y=\sum_c\text{pxy}$. +$$I=\sum_{c,\phi}\text{pxy}\,\log_2\frac{\text{pxy}+\varepsilon}{p_x p_y^\top+\varepsilon}\ \text{(bits)},\quad \varepsilon=10^{-15}.$$ +With $H_c=-\sum p_x\log_2 p_x$, $H_p=-\sum p_y\log_2 p_y$: +$$I_{\text{norm}}=\frac{I}{D},\quad D=\begin{cases}\min(H_c,H_p) & \texttt{mode="min"}\ \text{(coefficient of constraint; default)}\\[2pt]\tfrac12(H_c+H_p) & \texttt{mode="average"}\end{cases}$$ +returning 0 if $D\le0$. + +> **Default `normalize_mode="min"` (blocking fix).** Under the uniform-clonotype prior (§0.8), each row sums to 1 and the table is divided by its sum, so $p_x=P(c)=1/C$ **exactly** and $H_c=\log_2 C$ is **structural and maximal**. `"average"` denom $=\tfrac12(\log_2 C+H_p)$ throttles normalized MI by $\sim1/\log_2 C$ and shrinks its ceiling as $C$ grows — non-comparable across groups/covariates with different $C$, breaking the groupby-comparison workflow. `"min"` gives $I/H_p$ (reaches 1 when clone determines phenotype, $C$-independent) and is the default. Docstring states $H_c=\log_2 C$ is not a meaningful normalizer here. + +**Fixes:** `n_samples=0` = deterministic plug-in $I(m̃)$ (was one random draw); `posterior=False → NotImplementedError` deleted (§0.4). Estimator honesty per §0.6 (MI Jensen-gap sign indeterminate). + +**(c) Arguments → math.** As the shared table; additionally `normalize_mode` selects $D$. `clones` restricts rows; `normalized` toggles $I$ vs $I/D$. + +**(d) Return shape.** `n_samples=0`, no `groupby` → scalar `float`; `n_samples>0`, no `groupby` → `(N,)` array + `mean, sd, hdi_low, hdi_high`; `groupby` → tidy DataFrame, one row per group, column `MI` [+ summary]. Fast path (precomputed jd) valid only at `n_samples=0`, `groupby=None`. + +`__all__ = ["mutual_information"]` + +### 7.5 `tools/_flux.py` — `phenotypic_flux` (renamed from `flux`) + +```python +phenotypic_flux( + adata, *, + cov_from, cov_to, + groupby=None, + n_samples=0, temperature=1.0, + clones=None, + distance_metric="l1", + random_state=None, device=None, +) -> pandas.Series | numpy.ndarray | pandas.DataFrame +``` + +**(a) Responsibility.** Per-clonotype distance between a clone's phenotype distribution at `cov_from` vs `cov_to`, over the clone intersection. + +**(b) Math.** $J_{\text{from}}$, $J_{\text{to}}$ (rows $P(\phi\mid c)$); over common clones $c$, $p=J_{\text{from}}[c]$, $q=J_{\text{to}}[c]$: +$$d_c=\text{phenotype\_distance}(p,q,\ \text{metric}=\texttt{distance\_metric}),$$ +dispatched through `_distance` (§3.4): `"l1"` (default, bounded $[0,2]$), `"kl"`/`"dkl"` (directional, unbounded, **bits**), `"jsd"` (symmetric, bounded $[0,1]$ bit), or callable. + +**Fixes carried by the rewrite:** (1) the `posterior=False` dead branch is gone (no more `TypeError` from unsupported `silent=`/`combine_with_logits=` kwargs); (2) **reproducibility** — draws use a seeded **torch** generator (old `seed=` only touched NumPy → no-op); (3) **unit consistency** — KL is $\log_2$/bits, single $\varepsilon$; (4) the `flux_table` column-misalignment (`clones_g.index(cl)` vs `common`-ordered columns) is removed by returning a tidy per-(group,clone) frame keyed by clone id. **Estimator honesty (§0.6):** `n_samples=0` = plug-in $d_c(m̃)$ (convex → under-estimates $\mathbb{E}_q[d_c]$; a clone with no real shift reads exactly 0 at `n_samples=0` but $>0$ in the `n_samples>0` mean). + +**(c) Arguments → math.** + +| Argument | Effect | +|---|---| +| `cov_from`, `cov_to` | The two conditions compared (was `from_this`/`to_that`). | +| `groupby` | Per-group flux → tidy rows (group × clone) with a `clone_size` column (replaces `flux_table`), via full-space restriction. | +| `n_samples` | `0` → deterministic per-clone plug-in `Series`; `N` → $N$ redrawn distance vectors + summary. | +| `temperature` | Tempers both joints identically before differencing. | +| `clones` | Restricts both sides; distances over the intersection. | +| `distance_metric` | `"l1"` / `"kl"` / `"jsd"` / callable. | +| `random_state` / `device` | Seeding / backend. | + +**(d) Return shape.** `n_samples=0`, no `groupby` → `Series` over common clones; `n_samples>0` → per-clone `mean, sd, hdi_low, hdi_high`; `groupby` → tidy DataFrame row per (group, clone) + `clone_size` [+ summary]. + +`__all__ = ["phenotypic_flux"]` + +### 7.6 `tools/_compare.py` — `compare_groups` (public group-comparison orchestrator) + +```python +compare_groups( + df, *, + value, # column holding the per-unit metric value or draw vector + by, # grouping column (e.g. "response") + reference=None, # baseline level; None → all pairwise + paired=False, # True → paired posterior-draw contrast (uses prob_direction) + hdi_prob=0.94, + alternative="two-sided", +) -> pandas.DataFrame +``` + +**Responsibility.** The **public** replacement for the deleted `mi_compare` / `delta_entropy_table` / `flux_table`: turn a tidy `groupby` result (per-unit point estimates, e.g. per patient) or paired posterior-draw vectors into group contrasts. This closes the audit gap where "`groupby` + `_stats` subsumes `*_compare`/`*_delta`" was non-functional because `_stats` is private. + +**Math.** For each contrast (`reference` vs other, or all pairs): +- **Unpaired point estimates:** Mann–Whitney $U$ + two-sided $p$ (`_stats.mann_whitney`), group means, and $\Delta=\text{mean}_B-\text{mean}_A$. +- **Paired posterior draws** (`paired=True`, one draw vector per group per unit, aligned by `sample_id`): the signed difference $\Delta^{(s)}=\text{metric}_B^{(s)}-\text{metric}_A^{(s)}$, then `mean(Δ)`, `hdi(Δ)`, and **`p_gt`/`p_lt` via `prob_direction`** — the **only** place a direction probability is emitted (§0.7). + +**Return.** Tidy DataFrame, one row per contrast: `group_a, group_b, mean_a, mean_b, delta, U, p, p_gt, hdi_low, hdi_high, stars`. Recreates `mi_compare`'s per-pair output exactly. + +`tools/__init__.py __all__ = ["joint_distribution", "clonotypic_entropy", "phenotypic_entropy", "mutual_information", "phenotypic_flux", "compare_groups"]` + +### 7.7 h5ad-serializable return shapes (forward-compat with the deferred `@tl_result` uns-cache) + +Every `tl` return frame is constrained **now** to survive an h5ad round-trip so the deferred `@tl_result` cache is a one-line wrapper later: **flat columns only, no object-dtype "samples" columns**, and provenance in a serializable `_provenance` column (JSON string) **plus** `df.attrs["params"]` (attrs are convenience-only; the column is the durable copy). Per-draw values live in a **separate long frame** carrying an explicit `sample_id` level, never as numpy vectors embedded in object columns (which `AnnData.write` drops). Cache key = hash of `(covariate, groupby, n_samples, temperature, clones, normalized, normalize_mode, distance_metric, use_logits, random_state)`. + +### 7.8 Draw-once efficiency invariant + +For `n_samples>0`, the engine draws the `p_ct` table **once per sample** and **reuses that draw across all covariates, groups, and clones**; groups are formed by cell/clone masking, not re-drawing. `covariate=None`, the flux sankey's pairwise series, per-patient analyses, and `diag.permutation_null` all consume one shared draw stack. A test/counter asserts the number of Dirichlet draws equals `n_samples`, independent of `#groups` and `#covariates`. + +### 7.9 Precomputed-joint fast path constraints + +A bare precomputed joint carries no `p_ct`/`local_scale`/logits/cells, so it is valid **only for `n_samples=0` and `groupby=None`** (`clones=` merely re-filters rows). `n_samples>0` or `groupby` on a bare joint raises a clear `ValueError`; a jd that already contains a `sample_id` axis is accepted, with `n_samples` treated as validation, not resampling. The adata-path == precomputed-jd equivalence test is scoped to `n_samples=0`. (`phenotypic_flux` takes adata only — it needs two joints — so it is unaffected.) Metrics propagate the input jd's `_provenance` into their output. + +--- + +## 8. `tcri.pl` — plotting (`plotting/`) + +Twins mirror `tl` by filename and function name. Each renderer accepts its `tl` twin's metric arguments (computing the joint internally when needed) plus rendering args. Cross-group comparison is driven by **`groupby`** (dots) and **`splitby`** (box hue) — **both axes are retained** because `splitby` has 116 example call sites and most published figures carry two categorical axes simultaneously (e.g. dots = patient, boxes = response, x = phenotype). Statistics come from `_stats` / `compare_groups`. + +### 8.1 `plotting/_entropy.py` + +```python +clonotypic_entropy( + adata, *, + covariate=None, groupby=None, splitby=None, + n_samples=0, temperature=1.0, clones=None, normalized=True, n_clones_ref=None, + palette=None, hue_order=None, ax=None, figsize=(6, 3), + rotation=90, legend_fontsize=6, bbox_to_anchor=(1.15, 1.0), + random_state=None, save=None, return_df=False, +) +``` +**(renamed from `clonotypic_entropy_by_phenotype`)** Box-and-dot plot of clonotypic entropy per phenotype across covariate values, per-group dots, significance brackets. Cross-group plots default to a common `n_clones_ref` for comparability (§7.2). + +```python +phenotypic_entropy( + adata, *, + covariate=None, groupby=None, splitby=None, + n_samples=0, temperature=1.0, clones=None, normalized=True, + palette=None, ax=None, figsize=(8, 4), + rotation=90, legend_fontsize=6, bbox_to_anchor=(1.15, 1.0), + random_state=None, save=None, return_df=False, +) +``` +**[FIXED]** Box/strip plot of phenotypic entropy per covariate/group. + +`__all__ = ["clonotypic_entropy", "phenotypic_entropy"]` + +### 8.2 `plotting/_mutual_information.py` + +```python +mutual_information( + adata, *, + covariate=None, groupby=None, splitby=None, + n_samples=0, temperature=1.0, clones=None, + normalized=True, normalize_mode="min", + palette=None, ax=None, figsize=(8, 4), rotation=90, + legend_fontsize=6, bbox_to_anchor=(1.15, 1.0), + random_state=None, save=None, return_df=False, +) +``` +**[FIXED; subsumes `mi_compare`]** Box/strip plot of clone×phenotype MI per covariate; `groupby` (e.g. `"patient"`) supplies per-group points and drives AUROC/MWU/label-permutation stats. `weighted` removed (behavior-change note, §0.8); default `normalize_mode="min"` (§7.4). + +`__all__ = ["mutual_information"]` + +### 8.3 `plotting/_flux.py` + +```python +phenotypic_flux( + adata, *, + order, + groupby=None, clones=None, + normalize=True, temperature=1.0, + distance_metric="l1", + phenotype_colors=None, ax=None, figsize=(6, 3), + show_legend=True, title=None, random_state=None, + save=None, return_axes=False, +) +``` +The Sankey. Draws phenotype-distribution flow across the ordered `order` sequence of covariate values, calling `tl.phenotypic_flux` pairwise between consecutive values under one shared draw stack (§7.8). `order` replaces pairwise `cov_from`/`cov_to` because a Sankey spans the full ordered series. + +`__all__ = ["phenotypic_flux"]` + +### 8.4 `plotting/_sankey.py` — private drawing primitives + +**`class SankeyNode`** *(internal)* — `__init__(self, x, y, val, *, dx=0.2, color=None, **kwargs)`; `plot(self, ax)`; `plot_node_connection(self, destination_node, ax, **kwargs)` (curved, color-interpolated ribbon). `_phenotype_mass_per_clone(adata, covariate, clones, normalize) -> dict[str, np.ndarray]` — `{clone → phenotype-mass vector}` at one covariate. `SankeyNode.hex_to_rgb` is **deleted** (0 callers; ribbons use `mcolors.to_rgb`). + +### 8.5 `plotting/_base.py` — private plotting engine + +| Signature | Responsibility | +|---|---| +| `_metric_boxplot(adata, function, *, groupby=None, splitby=None, ylabel="", order=None, palette=None, s=20, ax=None, figsize=(8, 4)) -> (fig, ax)` | **(renamed from `tcri_boxplot`)** Generic per-phenotype metric box/strip engine across `groupby`/`splitby` strata. **Rewritten to compute each stratum by full-space restriction (`clones=`/masks), never `function(adata[mask])`** — so engine-backed metrics never trip the alignment guard (§7.1). | +| `_finish(fig, ax, *, save=None, show=None, return_axes=False)` | scanpy-style show/save/return finalizer. | + +### 8.6 `plotting/_colors.py` + +| Symbol | Responsibility | +|---|---| +| `tcri_colors` (`list[str]`) | Canonical categorical hex palette. | +| `resolve_palette(adata, columns, *, palette=None) -> dict` | **(renamed from `set_color_palette`)** Assign `tcri_colors` to each `obs` column's categories, store in `uns["_colors"]`, return the map. **Fixes the "writes on `adata.copy()`" bug (mutates in place).** | + +`__all__ = ["tcri_colors", "resolve_palette"]` + +### 8.7 `plotting/_ternary.py` — `probability_ternary` (dispositioned; 24 live callers) + +```python +probability_ternary( + adata, *, + phenotypes, # the 3 phenotype axes of the simplex + groupby=None, clones=None, + palette=None, ax=None, figsize=(5, 5), + save=None, return_axes=False, +) +``` +Ternary phenotype-simplex scatter of per-cell/per-clone phenotype probabilities. **Kept public** (heavily used in notebooks) and migrated onto `K.X_PROBABILITIES` and the single metadata scheme; `weighted` removed. `__all__ = ["probability_ternary"]`. + +--- + +## 9. `tcri.diag` — diagnostics (`diagnostics/`) — NEW + +Read-only checks on the finalized model. PPCs return `DataFrame`s; the two relocated model plots render figures. **`model` is required exactly where the live decoder/param store is needed, optional where `adata` suffices** (stated per function). + +### 9.1 `diagnostics/_ppc.py` + +| Signature | Responsibility / math | +|---|---| +| `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`. | + +`__all__ = ["joint_distribution_ppc", "phenotype_calibration", "reconstruction_ppc", "permutation_null"]` + +### 9.2 `diagnostics/_training.py` + +| Signature | Responsibility | +|---|---| +| `loss(model, *, log_scale=False, ax=None, save=None)` | **(relocated `plot_loss`)** Plot training/validation ELBO and prior-KL from `model.history_`. | +| `archetypes(model, *, ax=None, save=None)` | **(relocated `plot_archetypes`)** Cluster-ordered clone-phenotype heatmap + archetype centroids, ordered by the `labels` from `build_archetypes` (retained, §5.5). | + +`diagnostics/__init__.py __all__ = ["joint_distribution_ppc", "phenotype_calibration", "reconstruction_ppc", "permutation_null", "loss", "archetypes"]` + +> The model PGM (`build_nested_tcri_pgm` / `draw_tcri_pgm_nested`) is moved **out of the package** to `docs/`. + +--- + +## 10. `tcri.ut` — utilities (`utils/`) + +### 10.1 `utils/_session.py` — public + +| Signature | Responsibility | +|---|---| +| `save_tcri_session(model, adata, out_dir, *, save_adata=True, compression="gzip") -> dict` | Persist a trained session: scvi model (weights + registry, no embedded adata), Pyro param store, `setup.json`, the h5ad. | +| `load_tcri_session(run_dir, *, adata_path=None, map_location=None, layer=None) -> (TCRIModel, AnnData)` | Reconstruct `TCRIModel` + `AnnData`: read h5ad, restore setup/category order, re-run `setup_anndata`, load model + Pyro params. **Sets the global Pyro store before any `get_p_ct`/`to_anndata` call** (param-store caveat, §5.2). | + +### 10.2 `utils/_session.py` — private helpers (not re-exported) + +`_to_jsonable(x)`; `_collect_setup_from_adata_or_model(adata, model)`; `_restore_category_order(adata, setup)`; `_resolve_TCRIModel()`; `_disable_scvi_onload_train()`; `_ensure_pyro_posterior_params(model, adata)` (guarantees `q_p_ct_raw`; if missing, warn + re-init to uniform $1/P$); `_pyro_load(path, *, map_location=None)`; `_ensure_dir(path)`. + +> **Removed from utils:** `write_adata_safely`, `_pop_nonserializables` (manager stash retired at `setup_anndata`); `probabilities` (dead: read a never-written `uns` key — **and its module-top import in `_plotting.py` is removed in the same PR**, §11); `build_nested_tcri_pgm`/`draw_tcri_pgm_nested` (→ `docs/`); `stars`/`auc_and_label_permutation`/`bootstrap_auc` (→ `_stats.py`). + +--- + +## 11. Surface deltas (removed / renamed / moved) and deletion-safety census + +**Deletion-safety rule (applied).** Every "safe deletion" is gated on a caller census over **`example/` + `docs/` notebooks**, not package source alone (notebook execution is itself an acceptance gate). Symbols with live notebook callers are **moved with their notebook, or the notebook is rewritten in the same PR** — never hard-deleted on package-only evidence. Every "0-caller deletion" PR also greps **import-sites** (not just call-sites) first. + +- **Deleted (dead/broken/out-of-scope, 0 live callers after census):** `clonality` (tl + pl), `clonotypic_entropy_base`, `delta_clonotypic_entropy`, `delta_entropy_table`, `mi_compare` (tl + pl), `flux_table`, `bayesian_mutual_information`, `probability_distribution`, `classify_phenotypes`, `get_latent_embedding`, `group_small_clones`, `register_probability_columns`, `remove_meaningless_genes`, `clone_fraction`, `_ent`, `ridge_delta_entropy`, `dkl` (→ `_distance.kl_divergence`), `probabilities` (**and its `_plotting.py` import**), `SankeyNode.hex_to_rgb`, `_ascii_hist` (+ all `graph=`/ASCII paths), `write_adata_safely`, `_pop_nonserializables`, and the retired `uns` keys `tcri_manager`, `tcri_clone_key`, `tcri_phenotype_key`, the `X_tcri_phenotypes` obsm slot. +- **Renamed:** `flux`→`tl.phenotypic_flux`; `get_cell_phenotype_probs`→`TCRIModel.predict`; `register_model`→`TCRIModel.to_anndata`; `clonotypic_entropy_by_phenotype`→`pl.clonotypic_entropy`; `tcri_boxplot`→`_base._metric_boxplot`; `set_color_palette`→`resolve_palette`; params `from_this`/`to_that`→`cov_from`/`cov_to`; engine `posterior=`→`use_logits=` (alias `cell_informed=`); `point_estimate=`→removed (use `n_samples`); `weighted=`→removed (uniform-clonotype prior). +- **Made private (with a public migration path):** `register_clonotype_key` / `register_phenotype_key` → `preprocessing/_register._register_*_key` (34 notebook callers → notebooks rewritten to `setup_anndata`/`to_anndata`, or a documented public key-registration shim provided in the same PR). +- **Moved to `examples/` (with notebooks):** `top_clone_umap`, `clone_size_umap`, `plot_phenotype_probabilities` (updated to read `K.X_PROBABILITIES`), `gene_entropy` (5 callers), `polar_plot` (1 caller). **Moved to `docs/`:** `build_nested_tcri_pgm`, `draw_tcri_pgm_nested`. **Moved to `diag/`:** `compare_joint_distribution`(→`joint_distribution_ppc`), `plot_loss`(→`loss`), `plot_archetypes`(→`archetypes`). +- **Subsumed by `groupby` + `compare_groups` (removed, with migration recipe):** the plural batch wrappers `clonotypic_entropies` / `phenotypic_entropies`, `pl.phenotypic_entropy_delta`, and every `*_compare` / `*_delta` / `*_table` variant. + +--- + +## 12. Appendix — current → target math/stats deltas (what changed and why) + +| # | Site | Current | Target | Rationale | +|---|---|---|---|---| +| 1 | engine `n_samples=0` | `joint_distribution_posterior` always draws 1 Dirichlet sample | closed-form posterior **mean** $m$, no draw | reproducible, Rao-Blackwell (§0.4) | +| 2 | `mutual_information`/`flux` `n_samples=0` | returns one random draw | deterministic plug-in point estimate | latent bug (§0.5) | +| 3 | `posterior=` semantics | conflates draw-vs-mean **and** logit-folding; MI `posterior=False` raises | axis renamed `use_logits`, means *fold per-cell logits* only; both branches use the posterior, never the generative prior | §0.4 | +| 4 | metric ↔ model agreement | metrics use a Dirichlet **draw** and **never** apply the gate | `use_logits=True` at $T=1$ uses the same gate-aware, mean-prior rule as `predict()` (needs persisted `X_tcri_logits`, `gate_prob`, `classifier_temperature`) | removes silent disagreement (§0.9, §5.1) | +| 5 | **plug-in vs posterior-mean** | drafts equate `n_samples=0` with `mean(n_samples>0)` | documented as **different estimators** (Jensen gap: entropy plug-in $\ge$ mean; flux plug-in $\le$ mean; MI indeterminate); **no equality test** | §0.6 | +| 6 | **`p_gt` summary** | attached to every `n_samples>0` metric | **removed** from single-metric summaries; emitted only by `compare_groups` on a signed $\Delta$ | metrics are $\ge0$ ⇒ $P(>0)\approx1$ (§0.7) | +| 7 | **posterior draw concentration** | three inconsistent variants: `clamp(s·m,1e-3)` (guide) vs `s·p_ct` vs `s·p_ct+1e-8` | draw from the **exact guide** `Dirichlet(clamp(s·m̃, 1e-3))` | HDIs must summarize the learned posterior (§0.5) | +| 8 | **MI `normalize_mode` default** | `"average"` ⇒ denom $\tfrac12(\log_2C+H_p)$, $C$-dependent | **`"min"`** ⇒ $I/H_p$, $C$-independent; document $H_c=\log_2C$ structural | cross-group comparability (§7.4) | +| 9 | clonotypic-entropy denominator | $\log_2$ of raw reindexed row count (inflated by absent clones) | $\log_2$ of **supported** clones; optional fixed `n_clones_ref` | comparability (§7.2) | +| 10 | phenotypic-entropy zero clone | zero-mass clone → uniform → $H=1.0$ | zero-support clone → **`NaN`/excluded** | §7.3 | +| 11 | flux `seed` | seeds NumPy only; torch draws unaffected | seed a **torch (+CUDA) Generator**; `random_state` on engine/metrics | reproducibility (§0.10) | +| 12 | KL base/$\varepsilon$ | natural log in flux; $\log_2$ elsewhere; mixed $\varepsilon$ | one base ($\log_2$/bits), one $\varepsilon=10^{-12}$; add bounded symmetric `jsd` | unit consistency (§3.4) | +| 13 | "HDI" | equal-tailed percentiles labeled HDI | **true** highest-density interval `hdi_low/hdi_high` | correct for skewed bounded posteriors (§0.7) | +| 14 | temperature | applied at different stages in the two engines; double-tempered with `guide_temperature` | **single** power-temper of the base; `uns[K.P_CT]` stores the **raw** mean; $T=1$ reproduces `predict()` | §0.9 | +| 15 | joint-distribution PPC | references undefined global `model` → `NameError` | reads cols from `uns[K.METADATA]` | §9.1 | +| 16 | groupby via slicing | `function(adata[mask])` trips the full-space alignment guard | full-space cell/clone **restriction**; `_metric_boxplot` rewritten | §7.1, §8.5 | +| 17 | `local_scale` fallback | `uns.get("tcri_local_scale", 1.0)` silently corrupts draw variance if unwritten | `to_anndata` always writes `K.LOCAL_SCALE`; engine **raises** (no `1.0` default) when missing at `n_samples>0` | draw-variance integrity (§5.1) | +| 18 | partial posterior | intervals silently read as full predictive uncertainty | documented: `n_samples>0` captures **`p_ct` uncertainty only** (classifier logits fixed at their posterior-mean encoding) | §0.6 | + +--- + +*Source of truth cross-checked against `tcri/model/_model.py`, `tcri/model/_module.py`, `tcri/preprocessing/_preprocessing.py`, `tcri/metrics/_metrics.py`, `tcri/plotting/_plotting.py`, `tcri/plotting/_sankey.py`, `tcri/utils/_utils.py`, and the grafiti reference at `/Users/ceglian/Codebase/GitHub/grafiti/grafiti`. Intended document home: `/Users/ceglian/Codebase/GitHub/tcri/docs/contract/tcri_api_and_responsibilities.md`.* \ No newline at end of file diff --git a/docs/contract/_quarantine/tcri_arg_audit_data.json b/docs/contract/_quarantine/tcri_arg_audit_data.json new file mode 100644 index 0000000..4a679b8 --- /dev/null +++ b/docs/contract/_quarantine/tcri_arg_audit_data.json @@ -0,0 +1 @@ +{"finalDoc": "The audit data is fully grounded against the live contract doc. Here is the final publishable document.\n\n---\n\n# TCRI \u2014 Argument Consistency & Sufficiency (Final)\n\n**Scope.** The analysis/plot/diagnostic/preprocessing/stats/engine surface defined in `docs/contract/tcri_api_and_responsibilities.md` (\u00a70\u2013\u00a712). Pure internal helpers with no shared-vocabulary arguments (`_console`, priors/module/classifier/training internals, `_register._write_*`, session utils, `_compute` reducers, `resolve_device`/`get_xp`/`asnumpy`) are omitted as all-`\u2014` rows.\n\n**Method.** A full parameter matrix was extracted, every shared concept was checked against the four RULES (same concept \u2192 same name/order/default/type; kw-only except the primary positional; American spelling; `groupby`\u2260`splitby`; `covariate=None`=all; `n_samples=0`=point estimate), then an argument-sufficiency pass compared the surface to the six example notebooks. A verifier then adjudicated every finding against both contract docs and the notebooks.\n\n**This document folds in the verifier's verdicts:** refuted findings are dropped (listed once, for the record); revised findings are restated in their corrected form; the verifier's missed items are added. Each item carries its audit id and status: **[confirmed]**, **[revised]**, or **[missed]**.\n\n**Dropped as refuted (not actionable \u2014 recorded so they are not re-raised):**\n- **S3 `weighted=`** \u2014 deliberately removed and documented (\u00a70.8, appendix row-and-changelog). Behavior change, not an oversight.\n- **S8 `pp.gene_entropy(batch_key=)`** \u2014 relocated to `examples/` with its notebook (\u00a711), so `batch_key`/`agg_function` travel with it. Not lost.\n- **S12 `base=`** \u2014 single-base (`log2`/bits) is an intentional unit-consistency fix (\u00a73.4, \u00a77.2/7.3, appendix row 12). Re-adding `base=` would undo it. `decimals=` is cosmetic; drop.\n\n---\n\n## 1. PARAMETER MATRIX\n\nCells show `name=default` (keyword-only unless noted), `name(req)` = required keyword-only (no default), `POS` = positional, `\u2014` = absent, **DIVERGENT** = the concept exists under a different name/default, **\u2014 (missing)** = the concept is applicable but absent.\n\n### Table A \u2014 engine / tl / pl / diag / pp core vocabulary\n\n| Function (ns/module) | primary(pos) | covariate | cov_from/cov_to | order | groupby | splitby | n_samples | temperature | clones | use_logits | normalized | normalize_mode | distance_metric | n_clones_ref | random_state | device | palette | ax | figsize | save | return |\n|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|\n| tl.joint_distribution | adata | covariate=None | \u2014 | \u2014 | groupby=None | \u2014 | n_samples=0 | temperature=1.0 | clones=None | use_logits=True | \u2014 | \u2014 | \u2014 | \u2014 | random_state=None | device=None | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 |\n| tl.clonotypic_entropy | adata_or_jd | covariate=None | \u2014 | \u2014 | groupby=None | \u2014 | n_samples=0 | temperature=1.0 | clones=None | \u2014 | normalized=True | \u2014 | \u2014 | n_clones_ref=None | random_state=None | device=None | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 |\n| tl.phenotypic_entropy | adata_or_jd | covariate=None | \u2014 | \u2014 | groupby=None | \u2014 | n_samples=0 | temperature=1.0 | clones=None | \u2014 | normalized=True | \u2014 | \u2014 | \u2014 | random_state=None | device=None | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 |\n| tl.mutual_information | adata_or_jd | covariate=None | \u2014 | \u2014 | groupby=None | \u2014 | n_samples=0 | temperature=1.0 | clones=None | \u2014 | normalized=True | normalize_mode=\"min\" | \u2014 | \u2014 | random_state=None | device=None | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 |\n| tl.phenotypic_flux | adata | \u2014 | cov_from(req), cov_to(req) | \u2014 | groupby=None | \u2014 | n_samples=0 | temperature=1.0 | clones=None | \u2014 | \u2014 | \u2014 | distance_metric=\"l1\" | \u2014 | random_state=None | device=None | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 |\n| tl.compare_groups | df | \u2014 | \u2014 | \u2014 | **DIVERGENT: by(req)** | **DIVERGENT: by(req)** | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 |\n| pl.clonotypic_entropy | adata | covariate=None | \u2014 | **\u2014 (missing)** | groupby=None | splitby=None | n_samples=0 | temperature=1.0 | clones=None | \u2014 | normalized=True | \u2014 | \u2014 | n_clones_ref=None | random_state=None | **\u2014 (missing)** | palette=None | ax=None | figsize=(6,3) | save=None | return_df=False |\n| pl.phenotypic_entropy | adata | covariate=None | \u2014 | **\u2014 (missing)** | groupby=None | splitby=None | n_samples=0 | temperature=1.0 | clones=None | \u2014 | normalized=True | \u2014 | \u2014 | \u2014 | random_state=None | **\u2014 (missing)** | palette=None | ax=None | figsize=(8,4) | save=None | return_df=False |\n| pl.mutual_information | adata | covariate=None | \u2014 | **\u2014 (missing)** | groupby=None | splitby=None | n_samples=0 | temperature=1.0 | clones=None | \u2014 | normalized=True | normalize_mode=\"min\" | \u2014 | \u2014 | random_state=None | **\u2014 (missing)** | palette=None | ax=None | figsize=(8,4) | save=None | return_df=False |\n| pl.phenotypic_flux | adata | \u2014 | \u2014 | order(req) | groupby=None | **\u2014 (missing)** | **\u2014 (missing n_samples)** | temperature=1.0 | clones=None | \u2014 | **DIVERGENT: normalize=True** | \u2014 | distance_metric=\"l1\" | \u2014 | random_state=None | **\u2014 (missing)** | **DIVERGENT: phenotype_colors=None** | ax=None | figsize=(6,3) | save=None | return_axes=False |\n| pl.probability_ternary | adata | **\u2014 (missing)** | \u2014 | \u2014 | groupby=None | **\u2014 (missing)** | **\u2014 (missing)** | **\u2014 (missing)** | clones=None | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | **\u2014 (missing)** | \u2014 | palette=None | ax=None | figsize=(5,5) | save=None | return_axes=False |\n| pl._metric_boxplot (priv) | adata, function POS | \u2014 | \u2014 | order=None | groupby=None | splitby=None | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | palette=None | ax=None | figsize=(8,4) | \u2014 | \u2014 |\n| pl._finish (priv) | fig, ax POS | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | ax POS | \u2014 | save=None | return_axes=False; **show=None (only here)** |\n| pl.resolve_palette | adata, columns POS | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | palette=None | \u2014 | \u2014 | \u2014 | \u2014 |\n| diag.joint_distribution_ppc | adata | covariate=None | \u2014 | \u2014 | **\u2014 (missing)** | \u2014 | **\u2014 (missing)** | temperature=1.0 | \u2014 | \u2014 | \u2014 | \u2014 | distance_metric=\"l1\" | \u2014 | **\u2014 (missing)** | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 |\n| diag.phenotype_calibration | adata | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | (n_bins=10) |\n| diag.reconstruction_ppc | model, adata=None POS | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | **DIVERGENT: n_samples=100** | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | **DIVERGENT: seed=0** | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 |\n| diag.permutation_null | adata | covariate=None | \u2014 | \u2014 | groupby=None | **\u2014 (missing)** | **\u2014 (metric=\"mutual_information\"; no metric passthrough)** | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | **DIVERGENT: seed=0** | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | (n_permutations=1000) |\n| diag.loss | model | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | ax=None | \u2014 | save=None | (log_scale=False) |\n| diag.archetypes | model | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | ax=None | \u2014 | save=None | \u2014 |\n| pp.group_singletons | adata | \u2014 | \u2014 | \u2014 | **DIVERGENT default: groupby=\"patient\"** | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | (clonotype_key=\"trb\", target_col=\"trb_unique\", min_clone_size=10) |\n| pp.clone_size | adata | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | (key_added=\"clone_size\", return_counts=False) |\n\n### Table B \u2014 distance / stats / contrast / RNG vocabulary\n\n| Function | primary(pos) | metric-selector | base / eps | by | value | reference | paired | hdi_prob | alternative | seed / random_state | resample count | pos_label |\n|---|---|---|---|---|---|---|---|---|---|---|---|---|\n| shared.phenotype_distance (priv) | p, q POS | **DIVERGENT: metric=\"l1\"** | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 |\n| shared.kl_divergence (priv) | p, q POS | \u2014 | base=2.0, eps=1e-12 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 |\n| shared.js_divergence (priv) | p, q POS | \u2014 | base=2.0, eps=1e-12 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 |\n| shared.hdi (priv) | samples POS | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | hdi_prob=0.94 | \u2014 | \u2014 | \u2014 | \u2014 |\n| shared.summarize (priv) | samples POS | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | hdi_prob=0.94 | \u2014 | \u2014 | \u2014 | \u2014 |\n| shared.mann_whitney (priv) | a, b POS | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | alternative=\"two-sided\" | \u2014 | \u2014 | \u2014 |\n| tl.compare_groups | df POS | \u2014 | \u2014 | **by(req)** | value(req) | reference=None | paired=False | hdi_prob=0.94 | alternative=\"two-sided\" | \u2014 | \u2014 | \u2014 |\n| shared.auc_and_label_permutation (priv) | scores, labels POS | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | **seed=42** | n_perm=200_000 (max_exact=200_000) | pos_label=None |\n| shared.bootstrap_auc (priv) | scores, labels POS | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | **seed=42** | n_boot=5000 | pos_label=None |\n| diag.reconstruction_ppc | model POS | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | **seed=0** | n_samples=100 | \u2014 |\n| diag.permutation_null | adata POS | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | **seed=0** | n_permutations=1000 | \u2014 |\n| _compute._joint_draws (priv) | adata POS | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | random_state(kw) | n_samples(kw) | (gate_prob kw; **order/set differs from joint_distribution**) |\n\n### Consistency legend\n- **RNG:** `random_state=None` (engine+tl+pl) vs `seed=42` (stats) vs `seed=0` (diag) \u2014 3 spellings/defaults for one concept.\n- **distance selector:** `distance_metric=\"l1\"` (tl/pl flux, diag ppc) vs `metric=\"l1\"` (`_distance.phenotype_distance`, private).\n- **clones\u2194temperature order:** engine + `pl.phenotypic_flux` emit `clones` before `temperature`; every other metric emits `temperature` before `clones`.\n- **normalize vs normalized:** `normalize=` (distribution bool) only on `pl.phenotypic_flux`; `normalized=` (metric [0,1] scaling bool) everywhere else \u2014 a near-collision.\n- **device:** present on engine + all tl metrics; absent on every `pl.*` metric plot.\n- **splitby:** present on the tl\u2192pl box plots; absent on `pl.phenotypic_flux` and `pl.probability_ternary`.\n- **figsize defaults:** (6,3) clonotypic_entropy/flux \u00b7 (8,4) phenotypic_entropy/MI/_metric_boxplot \u00b7 (5,5) ternary.\n- **overloaded `splitby` (legacy):** on box plots the legacy `splitby` was the cohort-hue; on ternary/flux the legacy `splitby` was the covariate **column** paired with `conditions=`/`order=` **values** \u2014 two different concepts under one name.\n\n### Clean \u2014 verified consistent (no findings)\n`temperature=1.0`; the `n_samples=0` point-estimate convention (single exception: `reconstruction_ppc`); `hdi_prob=0.94`; `alternative=\"two-sided\"`; `use_logits` correctly engine-only; **no `point_estimate=` / `posterior=` survivors**; **no British-spelling residue** anywhere in the target surface.\n\n---\n\n## 2. CONSISTENCY \u2014 confirmed inconsistencies + canonical decision\n\nEach item: the divergence \u2192 the single canonical name/order/default/type to adopt \u2192 the functions that change. Duplicate audit ids are merged. Private-symbol renames are marked *(private \u2192 optional/lower stakes)*.\n\n### High\n\n**H1 \u2014 RNG fragmentation \u2192 `random_state`** *(C1 \u2261 S10, confirmed).*\n`random_state=None` (engine/tl/pl) vs `seed=42` (`_stats`) vs `seed=0` (diag). One concept, three names/defaults/types.\n**CANONICAL:** `random_state=None`, keyword-only, typed `int | numpy.Generator | torch.Generator | None`, placed as the **penultimate** compute arg (immediately before `device`) everywhere it appears. Drop the literal `42`/`0` defaults for `None` (seed internally from the passed generator).\n**Load-bearing (public):** `diag.reconstruction_ppc`, `diag.permutation_null` (rename `seed`\u2192`random_state`, default `None`). *(Private \u2192 optional):* `_stats.auc_and_label_permutation`, `_stats.bootstrap_auc` may keep `seed` internally but should be documented as the split.\n\n**H2 \u2014 `clones`/`temperature` order \u2192 temperature-before-clones** *(C2, confirmed; resolve with H8/C12).*\nEngine emits `(n_samples, use_logits, clones, temperature)` and `pl.phenotypic_flux` emits `(clones, normalize, temperature)` \u2014 `clones` before `temperature` \u2014 while every other metric emits `(n_samples, temperature, clones)`.\n**CANONICAL** shared block, left\u2192right: `covariate`/`cov_*`, `groupby`, `splitby`, `n_samples`, `temperature`, `clones`, `use_logits` (engine-only, after `clones`), `normalized`/`normalize_mode`/`distance_metric`, `random_state`, `device`. Reorder `tl.joint_distribution` and `pl.phenotypic_flux` to `temperature`\u2192`clones`.\n\n**H3 \u2014 `device` absent from every `pl.*` metric plot \u2192 add `device=None`** *(C3, revised).*\nConfirmed absent on `pl.clonotypic_entropy`, `pl.phenotypic_entropy`, `pl.mutual_information`, `pl.phenotypic_flux`. Revision: on the three box plots the dangling knob is that they expose posterior draws + `random_state` but cannot steer the draw device; on `pl.phenotypic_flux` the finding's \"draws-without-device\" premise is wrong (flux exposes **no** `n_samples`), so its dangling knob is `random_state` \u2014 fix flux by adding **both** `device` and `n_samples` (see S-flux below).\n**CANONICAL:** add `device=None` (`str | torch.device | None`, kw-only, penultimate before nothing/after `random_state`) to all four `pl.*` metric plots and thread it into the underlying `tl` call.\n\n**H4 \u2014 overloaded `splitby` on ternary/flux \u2192 `covariate_key` + `conditions`/`order`, NOT a cohort `splitby`** *(C4 revised \u2261 M2 missed).*\n`splitby` is genuinely absent from `pl.probability_ternary` and `pl.phenotypic_flux`, but the legacy notebooks passed `splitby=` there to name the covariate **column** (paired with `conditions=`/`order=` selecting the **values**), which is a *different* concept from the box-plot cohort-hue `splitby`. Adding a cohort-hue `splitby` here would re-introduce the exact name-collision the RULES forbid.\n**CANONICAL:** reserve `splitby` strictly for the cohort/faceting hue. For ternary/flux add `covariate_key=None` (the covariate **column** override) plus `conditions=None` (ternary: 1\u20132 covariate levels \u2192 start/end simplices) / reuse `order` (Sankey: the ordered value series). Document that legacy `splitby=` maps to `covariate_key`, not `splitby`, in the Phase-10 notebook rewrite.\n**Feasibility caveat (must be resolved in the contract):** the engine's joint is built from `p_ct` indexed by the trained covariate (`ct_to_cov`), so a `covariate_key` override is only implementable on an **empirical per-cell-probability aggregation path**, not the `p_ct` engine path. The contract must state which path multi-column flux/ternary use before `covariate_key` is added.\n\n### Medium\n\n**M1 \u2014 `n_samples=100` overloads the `n_samples=0` convention \u2192 rename to `n_sims`** *(C5, confirmed).*\n`diag.reconstruction_ppc.n_samples=100` counts simulated PPC datasets, colliding with the surface-wide `n_samples=0`=point / `>0`=draws convention.\n**CANONICAL:** rename to `n_sims=100` (`int`). Reserve `n_samples` (default `0`) exclusively for the posterior-draw convention.\n\n**M2 \u2014 distance selector `metric=` vs `distance_metric=`** *(C6, confirmed).*\n`_distance.phenotype_distance(metric=\"l1\")` (private dispatcher) vs `distance_metric=\"l1\"` on `tl.phenotypic_flux`, `pl.phenotypic_flux`, `diag.joint_distribution_ppc`.\n**CANONICAL:** `distance_metric` is the public name everywhere. Since `phenotype_distance` is private, keep its `metric` **only** as an internal dispatcher detail (documented) \u2014 no public churn required.\n\n**M3 \u2014 `normalize` vs `normalized` + `phenotype_colors` vs `palette` on `pl.phenotypic_flux`** *(C7, confirmed).*\n`normalize=True` (distribution-normalization bool feeding `_phenotype_mass_per_clone`) look-alikes `normalized=True` (metric [0,1] scaling) used everywhere else; `phenotype_colors=None` duplicates `palette`.\n**CANONICAL:** rename `phenotype_colors`\u2192`palette`. Rename the flux bool to `normalize_distributions=True` (or drop it if the engine always feeds normalized distributions). Never ship both `normalize` and `normalized` as look-alike names.\n\n**M4 \u2014 `order`/`hue_order` inconsistent across the box-plot family \u2192 expose both on all three** *(C9 \u2261 S2 \u2261 S13, confirmed).*\n`pl.clonotypic_entropy` has `hue_order` but no `order`; `pl.phenotypic_entropy` and `pl.mutual_information` have neither; the shared `_metric_boxplot` supports `order`; notebooks constantly pass `order=`.\n**CANONICAL:** expose `order=None` (x-axis category order) **and** `hue_order=None` (`splitby` level order), both `list[str] | None`, on all three metric box plots, wired to `_metric_boxplot.order` and the `splitby` hue.\n\n**M5 \u2014 `compare_groups.by` names the cohort \u2192 rename `by`\u2192`splitby`** *(C10, confirmed; defensible-either-way).*\n`by` is the column whose levels are contrasted (with `reference`/`paired`) \u2014 semantically the comparison **cohort**, i.e. the plotting surface's `splitby`.\n**CANONICAL:** rename `by`\u2192`splitby` so the aggregation-unit(`groupby`)/comparison-cohort(`splitby`) vocabulary is uniform metric\u2192plot\u2192contrast. (`by` is pandas-idiomatic; if the team prefers `by`, that is a documented, deliberate exception rather than drift.)\n\n**M6 \u2014 `pp.group_singletons` defaults `clonotype_key=\"trb\"`, `groupby=\"patient\"` \u2192 keep, document as pre-registration defaults** *(C11, revised).*\nValues confirmed, but the original \"force `None`/match `unique_clone_id`\" fix is wrong: `group_singletons` runs **before** `setup_anndata`, on the raw pre-registration column (`trb` \u2192 writes `target_col=\"trb_unique\"`, which only later becomes the registered `clonotype_key`), so `unique_clone_id` doesn't exist yet; its `groupby` is the collapse **unit** (consistent with `groupby`=aggregation-unit), only the *default* differs.\n**CANONICAL:** keep `clonotype_key=\"trb\"` and `groupby=\"patient\"` as intentional pre-registration defaults and **document** them as such; do not force `None`.\n\n**M7 \u2014 `_joint_draws` keyword order/set diverges from `joint_distribution`** *(C12, confirmed; private \u2192 lower stakes).*\n`_joint_draws(covariate, clones, n_samples, use_logits, temperature, gate_prob, \u2026)` vs `joint_distribution(covariate, groupby, n_samples, use_logits, clones, temperature, \u2026)`.\n**CANONICAL:** make `_joint_draws`' keyword order a strict subset-in-order of the public engine: `(covariate, n_samples, temperature, clones, use_logits, gate_prob, random_state, device)` \u2014 i.e. temperature-before-clones (dovetails H2). Resolve `gate_prob` at the public layer or document it as an internal-only extra.\n\n**M8 \u2014 `joint_distribution_ppc` arg-order + missing knobs** *(C13 ~ S15, confirmed).*\nOrders `(covariate, distance_metric, temperature)` (temperature after the selector) and omits `groupby`/`clones` (and `n_samples`/`random_state`) that its metric siblings expose.\n**CANONICAL:** reorder to `(covariate, groupby, clones, temperature, distance_metric)`; add `groupby=None`, `clones=None`. It is model-free/deterministic, so add `n_samples=0`/`random_state=None` **only if** a draw-based comparison is intended (flag as a contract decision, not a default addition).\n\n**M9 \u2014 `covariate=None` dual meaning in the `adata_or_jd` fast path** *(C19 \u2261 S14, confirmed).*\nAll signatures default `covariate=None`, but \u00a77.2(c)/\u00a77.9 let the precomputed-`jd` path read `covariate` as \"required/the one baked into this jd,\" contradicting the RULE that `covariate=None`=all covariate values.\n**CANONICAL:** `covariate=None` **always** means all covariate values on the adata path. On the precomputed-`jd` fast path, `covariate`/`n_samples`/`temperature`/`clones`/`random_state`/`device` are inert \u2014 **raise** a clear `ValueError` if any is set to a non-default \u2014 so `covariate=None`=all is preserved only where it is computed. Fix the entropy/MI docstrings and dispatch so `covariate` is never treated as required.\n\n### Low\n\n**L1 \u2014 figsize drift across box-plot siblings** *(C8, confirmed).* `pl.clonotypic_entropy=(6,3)` vs `(8,4)` for `pl.phenotypic_entropy`/`pl.mutual_information`/`_metric_boxplot`.\n**CANONICAL:** `figsize=(8,4)` for the metric box-plot family; set `pl.clonotypic_entropy` to `(8,4)`. Aspect-driven plots keep purpose-specific defaults, documented as intentional: ternary `(5,5)`, Sankey flux `(6,3)`.\n\n**L2 \u2014 2nd-positional violations \u2192 keyword-only** *(C14, confirmed).* `pl.resolve_palette(adata, columns)`, `pl._metric_boxplot(adata, function)`, `diag.reconstruction_ppc(model, adata=None)` each carry a second positional.\n**CANONICAL:** insert `*` after the single primary positional so `columns`, `function`, and `adata` are keyword-only. *Census note:* `_phenotype_mass_per_clone(adata, covariate, clones, normalize)` and `TCRIModule.prepare_two_level_params` also carry extra positionals (private, minor) \u2014 align in the same pass.\n\n**L3 \u2014 resample-count `n_permutations` vs `n_perm`** *(C15, confirmed).* `diag.permutation_null.n_permutations=1000` vs `_stats.auc_and_label_permutation.n_perm=200_000`.\n**CANONICAL:** `n_perm` for permutation counts everywhere; rename `permutation_null.n_permutations`\u2192`n_perm` (keep its `1000` default). `n_boot=5000` (bootstrap) is a distinct concept and keeps its name.\n\n**L4 \u2014 `show=` only on private `_finish`** *(C16, confirmed).* `_finish(show=None)` is unreachable from any public plot, though all expose `save=`/`return_*`.\n**CANONICAL:** adopt the scanpy triad uniformly \u2014 thread `show=None` (`bool | None`, kw-only) through **every** public `pl.*` entry point (with `save`/`return_*`), keeping it on `_finish`.\n\n**L5 \u2014 return-control name split** *(C17, confirmed; near-self-resolving).* `return_df` on DataFrame plots vs `return_axes` on figure plots.\n**CANONICAL:** keep the split (payloads genuinely differ) but **standardize which every plot exposes** \u2014 DataFrame/metric plots \u2192 `return_df`; figure-only plots (ternary, Sankey) \u2192 `return_axes` \u2014 and document it. No plot may omit both (all currently comply).\n\n**L6 \u2014 `eps` clip-floor drift + internal contradiction** *(C18 confirmed \u2261 M3-missed strengthened).* Four floors: `1e-8` (`TCRIModel.predict`, `_compute_logits_and_prior`), `1e-6` (`TCRIModule.prepare_two_level_params`), `1e-12` (`kl_divergence`, `js_divergence`), `1e-15` (hardcoded in `_mi_from_joint`). Sharper than drift: \u00a73.4 and appendix row 12 both assert \"one \u03b5=10\u207b\u00b9\u00b2 **library-wide**, matching entropy/MI,\" yet the MI kernel \u2014 the exact metric that claim name-checks \u2014 uses `1e-15`. The document contradicts its own stated invariant.\n**CANONICAL:** set `_mi_from_joint` `eps=1e-12` to satisfy the stated invariant (single probability-clip floor `1e-12` for the distance/MI/normalization paths). Keep context-specific floors (`predict`/`prepare_two_level_params`) only where numerically justified, and document why. Do not leave prose and kernel contradicting.\n\n### Documentation / cross-document consistency (verifier's missed items)\n\n**D1 \u2014 `\u00a711` and the implementation plan disagree on live callers** *(M1, missed).* API \u00a711 line 770 buckets `clonality` (tl+pl), `probability_distribution`, and `clone_fraction` under \"Deleted \u2026 0 live callers after census,\" but the plan (and the notebooks) show live callers: `tcri.pl.clonality` is called with full args (`groupby`/`splitby`/`order`/`palette`/`figsize`) across smith/renal/comparison/yost/zhang; `tcri.metrics.probability_distribution` is imported and called 4\u00d7 via the alias `pdistribution(psubset, method=\u2026)`. `compare_phenotypes` (1 live call in synthetic1) is dispositioned in the plan (DROP) but is **absent from \u00a711 entirely**.\n**CANONICAL:** reconcile \u00a711 with plan \u00a73 \u2014 move `clonality`, `probability_distribution`, `clone_fraction` out of the \"0 live callers\" bucket into a \"removed **with** replacement + in-PR notebook rewrite (Phase 6/10)\" category; add `compare_phenotypes` to \u00a711; and correct the false \"0 live callers after census\" label (the census missed alias-imported and package-qualified call-sites).\n\n**D2 \u2014 undispositioned legacy renames `method=` and `phenotype_names=`** *(M4, missed).* `method=` (e.g. `method=\"probabilistic\"`) is passed to `pl.clonotypic_entropy`, `pl.mutual_information`, `pl.flux`, `polar_plot`, and `pdistribution`, but only `posterior=`/`point_estimate=` are mapped in \u00a711 \u2014 `method=` is never named. `phenotype_names=` is the legacy positional/keyword on `probability_ternary`, and the kept \u00a78.7 signature renames it to `phenotypes` with no rename entry.\n**CANONICAL:** add disposition-map rows to \u00a711: `method=` \u2192 removed, expressed via `n_samples`/`use_logits`; `phenotype_names` \u2192 `phenotypes` on `pl.probability_ternary`.\n\n**D3 \u2014 caller-census counts are unreliable** *(M5, missed).* `polar_plot` is recorded as \"1 notebook use,\" but the notebooks contain 4 local `def polar_plot` redefinitions plus multiple call-sites; the same def-plus-call pattern holds for `clonality`. Neither audit disambiguated def-vs-call.\n**CANONICAL:** re-run the census distinguishing (a) package-qualified calls, (b) alias-imported calls, (c) notebook-local redefinitions; record accurate per-symbol counts; the Phase-10 rewrite must also strip shadow `def`s, not just swap call-sites.\n\n**D4 \u2014 `pl.flux` box-plot disposition unstated** *(S9, revised).* The old per-clone flux-**distance** box-plot `pl.flux(..., paint=, distance_metric='dkl', ...)` is distinct from the Sankey and is used in renal/yost/zhang, but \u00a711 only renames `flux`\u2192`tl.phenotypic_flux` and never states this box-plot's fate. The plan settles it as DROP (the Sankey is the flux plot); re-adding a box variant would contradict that decision.\n**CANONICAL:** state explicitly in \u00a711 that the `pl.flux` distance box-plot is **DROP** (Sankey is the flux plot) \u2014 a doc-completeness fix, not a new function.\n\n---\n\n## 3. SUFFICIENCY \u2014 confirmed missing knobs (exact argument to add)\n\nEvery confirmed gap, with the exact `name=default` (type) and the functions that receive it. Refuted gaps (S3/S8/S12) are excluded per \u00a7Dropped.\n\n| # | Argument to add | Default | Type | Functions receiving it | Audit id / status |\n|---|---|---|---|---|---|\n| U1 | `phenotypes` (rename from `phenotype_names`) | (req) | `list[str]` (3 axes) | `pl.probability_ternary` | S1/M4 confirmed |\n| U2 | `conditions` | `None` | `list[str] \\| None` (1\u20132 covariate levels \u2192 start/end) | `pl.probability_ternary` | S1/M2 confirmed |\n| U3 | `scale_function` | `None` | `Callable[[float], float] \\| None` (freq \u2192 marker size) | `pl.probability_ternary` | S1 confirmed (exercised in every zhang/yost ternary call) |\n| U4 | `color` | `None` | `str \\| None` | `pl.probability_ternary` | S1 confirmed |\n| U5 | `covariate` | `None` | `str \\| None` (value; `None`=all) | `pl.probability_ternary` | S1 confirmed (lower priority \u2014 historical default, not passed in notebooks) |\n| U6 | `top_n` | `None` | `int \\| None` | `pl.probability_ternary` | S1 confirmed (lower priority) |\n| U7 | `n_samples` | `0` | `int` | `pl.probability_ternary`, **`pl.phenotypic_flux`** | S1, S5 confirmed |\n| U8 | `temperature` | `1.0` | `float` | `pl.probability_ternary` | S1 confirmed (lower priority) |\n| U9 | `order` | `None` | `list[str] \\| None` (x-axis category order) | `pl.clonotypic_entropy`, `pl.phenotypic_entropy`, `pl.mutual_information` | S2 \u2261 M4/C9 confirmed |\n| U10 | `hue_order` | `None` | `list[str] \\| None` (`splitby` level order) | `pl.phenotypic_entropy`, `pl.mutual_information` (already on `pl.clonotypic_entropy`) | S13 \u2261 C9 confirmed |\n| U11 | `device` | `None` | `str \\| torch.device \\| None` | `pl.clonotypic_entropy`, `pl.phenotypic_entropy`, `pl.mutual_information`, `pl.phenotypic_flux` | C3 confirmed/revised |\n| U12 | `minimum_clone_size` | `None` | `int \\| None` (filter `clone_size` before aggregation) | `pl.clonotypic_entropy`, `pl.phenotypic_entropy`, `pl.mutual_information` (optionally the `tl` metrics) | S6 confirmed |\n| U13 | `phenotype_subset` | `None` | `list[str] \\| None` (restrict rendered phenotype nodes) | `pl.phenotypic_flux` | S5 confirmed |\n| U14 | `pair_on` | `None` | `str \\| None` (unit column aligning the two levels under `paired=True`) | `tl.compare_groups` | S7 confirmed (closes the incomplete `delta_entropy_table` subsumption) |\n| U15 | `covariate_key` | `None` | `str \\| None` (covariate **column** override) | `tl.joint_distribution`, `tl.phenotypic_flux`, `pl.phenotypic_flux`, `pl.probability_ternary` | S4/M2 revised \u2014 **conditional** on the empirical-vs-`p_ct` path decision (H4 caveat) |\n| U16 | metric pass-throughs: `temperature=1.0` (`float`), `normalized=True` (`bool`), `normalize_mode=\"min\"` (`str`), `clones=None` (`list[str] \\| None`), `splitby=None` (`str \\| None`) | as noted | as noted | `diag.permutation_null` (so the null matches the reported statistic) | S11 confirmed |\n| U17 | `clones` | `None` | `list[str] \\| None` | `diag.joint_distribution_ppc`, `diag.permutation_null` | S15 \u2261 C13 confirmed |\n| U18 | `groupby` | `None` | `str \\| None` | `diag.joint_distribution_ppc` | S15 \u2261 C13 confirmed |\n| U19 | `show` | `None` | `bool \\| None` | all public `pl.*` plots (scanpy triad) | C16 confirmed |\n\n**Not a new knob, but tied to sufficiency:** `tl.compare_groups` currently reproduces one-axis contrasts but not the paired-**delta-then-cohort-contrast** of `delta_entropy_table` in one call \u2014 `paired=True` aligns by `sample_id` (posterior-draw pairing), not pre/post-within-unit pairing, and there is no second cohort axis. `pair_on` (U14) supplies the within-unit pairing key; the \u00a711 claim that `delta_entropy_table` is \"expressed via `groupby` + `compare_groups`\" must be softened to \"with `pair_on` + an explicit cohort split\" or documented as a two-pass recipe.\n\n---\n\n## 4. CORRECTED SIGNATURES\n\nFinal signatures for every function that changes. Deltas from the current contract are noted after each block. `# \u2020` marks the conditional `covariate_key` (gated on the H4 empirical-path decision).\n\n**`tl.joint_distribution`** \u2014 reorder to temperature-before-clones (`use_logits` engine-only, after `clones`).\n```python\njoint_distribution(\n adata, *,\n covariate=None, # None \u2192 ALL covariate values (one shared draw)\n groupby=None,\n n_samples=0,\n temperature=1.0,\n clones=None,\n use_logits=True, # engine-only; was posterior=; alias cell_informed=\n random_state=None, # int | numpy.Generator | torch.Generator | None\n device=None,\n) -> pandas.DataFrame\n```\n\n**`tl.compare_groups`** \u2014 `by`\u2192`splitby`; add `pair_on`.\n```python\ncompare_groups(\n df, *,\n value,\n splitby, # was by=\n reference=None,\n paired=False,\n pair_on=None, # NEW: within-unit pairing key when paired=True\n hdi_prob=0.94,\n alternative=\"two-sided\",\n) -> pandas.DataFrame\n```\n\n**`tl.phenotypic_flux`** \u2014 add optional `covariate_key`.\n```python\nphenotypic_flux(\n adata, *,\n covariate_key=None, # \u2020 covariate COLUMN override (H4 caveat)\n cov_from, cov_to,\n groupby=None,\n n_samples=0,\n temperature=1.0,\n clones=None,\n distance_metric=\"l1\",\n random_state=None,\n device=None,\n) -> pandas.DataFrame\n```\n\n**`pl.clonotypic_entropy`** \u2014 figsize (6,3)\u2192(8,4); add `order`, `minimum_clone_size`, `device`, `show` (`hue_order` already present).\n```python\nclonotypic_entropy(\n adata, *,\n covariate=None, groupby=None, splitby=None,\n n_samples=0, temperature=1.0, clones=None,\n normalized=True, n_clones_ref=None, minimum_clone_size=None,\n order=None, hue_order=None,\n palette=None, ax=None, figsize=(8, 4),\n rotation=90, legend_fontsize=6, bbox_to_anchor=(1.15, 1.0),\n random_state=None, device=None,\n show=None, save=None, return_df=False,\n)\n```\n\n**`pl.phenotypic_entropy`** \u2014 add `order`, `hue_order`, `minimum_clone_size`, `device`, `show`.\n```python\nphenotypic_entropy(\n adata, *,\n covariate=None, groupby=None, splitby=None,\n n_samples=0, temperature=1.0, clones=None,\n normalized=True, minimum_clone_size=None,\n order=None, hue_order=None,\n palette=None, ax=None, figsize=(8, 4),\n rotation=90, legend_fontsize=6, bbox_to_anchor=(1.15, 1.0),\n random_state=None, device=None,\n show=None, save=None, return_df=False,\n)\n```\n\n**`pl.mutual_information`** \u2014 add `order`, `hue_order`, `minimum_clone_size`, `device`, `show`.\n```python\nmutual_information(\n adata, *,\n covariate=None, groupby=None, splitby=None,\n n_samples=0, temperature=1.0, clones=None,\n normalized=True, normalize_mode=\"min\", minimum_clone_size=None,\n order=None, hue_order=None,\n palette=None, ax=None, figsize=(8, 4), rotation=90,\n legend_fontsize=6, bbox_to_anchor=(1.15, 1.0),\n random_state=None, device=None,\n show=None, save=None, return_df=False,\n)\n```\n\n**`pl.phenotypic_flux`** \u2014 reorder temperature-before-clones; `normalize`\u2192`normalize_distributions`; `phenotype_colors`\u2192`palette`; add `covariate_key`, `n_samples`, `phenotype_subset`, `device`, `show`.\n```python\nphenotypic_flux(\n adata, *,\n covariate_key=None, # \u2020 covariate COLUMN override (H4 caveat)\n order, # ordered covariate value series\n groupby=None,\n n_samples=0, # NEW: drives ribbon uncertainty (feeds random_state)\n temperature=1.0,\n clones=None,\n normalize_distributions=True, # renamed from normalize=\n distance_metric=\"l1\",\n phenotype_subset=None, # NEW: restrict rendered phenotype nodes\n palette=None, # renamed from phenotype_colors=\n ax=None, figsize=(6, 3),\n show_legend=True, title=None,\n random_state=None, device=None,\n show=None, save=None, return_axes=False,\n)\n```\n\n**`pl.probability_ternary`** \u2014 `phenotype_names`\u2192`phenotypes`; add `covariate_key`, `covariate`, `conditions`, `n_samples`, `temperature`, `scale_function`, `top_n`, `color`, `random_state`, `show`. No cohort-hue `splitby` (H4).\n```python\nprobability_ternary(\n adata, *,\n phenotypes, # renamed from phenotype_names; the 3 simplex axes\n covariate_key=None, # \u2020 covariate COLUMN override (H4 caveat)\n covariate=None, # covariate VALUE (None \u2192 all)\n conditions=None, # 1\u20132 covariate levels \u2192 start/end simplices\n groupby=None, clones=None,\n n_samples=0, temperature=1.0,\n scale_function=None, # freq \u2192 marker size\n top_n=None, color=None,\n palette=None, ax=None, figsize=(5, 5),\n random_state=None,\n show=None, save=None, return_axes=False,\n)\n```\n\n**`diag.joint_distribution_ppc`** \u2014 reorder temperature-before-`distance_metric`; add `groupby`, `clones`.\n```python\njoint_distribution_ppc(\n adata, *,\n covariate=None,\n groupby=None,\n clones=None,\n temperature=1.0,\n distance_metric=\"l1\",\n # add n_samples=0 / random_state=None ONLY if draw-based comparison is adopted\n) -> pandas.DataFrame\n```\n\n**`diag.reconstruction_ppc`** \u2014 `adata` keyword-only; `n_samples`\u2192`n_sims`; `seed`\u2192`random_state` (default `None`).\n```python\nreconstruction_ppc(\n model, *,\n adata=None,\n n_sims=100, # was n_samples=100\n random_state=None, # was seed=0\n) -> pandas.DataFrame\n```\n\n**`diag.permutation_null`** \u2014 `n_permutations`\u2192`n_perm`; `seed`\u2192`random_state`; add metric pass-throughs + `clones`.\n```python\npermutation_null(\n adata, *,\n metric=\"mutual_information\",\n covariate=None, groupby=None, splitby=None, clones=None,\n temperature=1.0, normalized=True, normalize_mode=\"min\",\n n_perm=1000, # was n_permutations=1000\n random_state=None, # was seed=0\n) -> pandas.DataFrame\n```\n\n**`pl.resolve_palette`** \u2014 `columns` keyword-only.\n```python\nresolve_palette(adata, *, columns, palette=None) -> dict\n```\n\n**Private helpers (align in the same pass, lower stakes):**\n```python\n_metric_boxplot(adata, *, function, groupby=None, splitby=None,\n ylabel=\"\", order=None, hue_order=None,\n palette=None, s=20, ax=None, figsize=(8, 4)) -> (fig, ax) # function kw-only; add hue_order\n_joint_draws(adata, *, covariate=None, n_samples=0, temperature=1.0,\n clones=None, use_logits=True, gate_prob=None,\n random_state=None, device=None) # subset-in-order of joint_distribution\n# _stats.auc_and_label_permutation / bootstrap_auc: optionally seed\u2192random_state (private; document the split)\n# _distance.phenotype_distance: keep metric= as an internal dispatcher detail (no public change)\n```\nAlso: `_mi_from_joint` \u2014 change hardcoded `eps=1e-15`\u2192`1e-12` (L6). `pp.group_singletons` / `pp.clone_size` \u2014 **no signature change** (M6: defaults kept and documented as intentional pre-registration values).\n\n---\n\n## 5. READY-TO-APPLY PATCH LIST \u2014 edits to fold into `docs/contract/tcri_api_and_responsibilities.md`\n\n1. **\u00a77.1 engine** \u2014 reorder the `joint_distribution` block to `covariate, groupby, n_samples, temperature, clones, use_logits, random_state, device` (temperature before clones; `use_logits` after `clones`). Update the \u00a77.1(c) arguments table row order to match.\n2. **\u00a77.6 `compare_groups`** \u2014 rename `by`\u2192`splitby` (signature + the \"grouping column\" comment + Math/Return prose); add `pair_on=None` between `paired` and `hdi_prob`; soften the \"subsumes `delta_entropy_table`\" claim to require `pair_on` (+ a cohort split) or document the two-pass recipe.\n3. **\u00a77.5 `tl.phenotypic_flux`** \u2014 add optional `covariate_key=None` (gated on the H4 empirical-path decision; add a one-line note stating the chosen path).\n4. **\u00a77.2 / \u00a77.9** \u2014 pin `covariate=None`=all covariate values on the adata path; state that the precomputed-`jd` fast path **ignores** `covariate`/`n_samples`/`temperature`/`clones`/`random_state`/`device` and **raises** if any is non-default (remove any \"covariate required\" reading).\n5. **\u00a78.1 `pl.clonotypic_entropy`** \u2014 `figsize=(6,3)`\u2192`(8,4)`; add `order=None`, `minimum_clone_size=None`, `device=None`, `show=None`.\n6. **\u00a78.1 `pl.phenotypic_entropy`** \u2014 add `order=None`, `hue_order=None`, `minimum_clone_size=None`, `device=None`, `show=None`.\n7. **\u00a78.2 `pl.mutual_information`** \u2014 add `order=None`, `hue_order=None`, `minimum_clone_size=None`, `device=None`, `show=None`.\n8. **\u00a78.3 `pl.phenotypic_flux`** \u2014 reorder temperature-before-clones; rename `normalize`\u2192`normalize_distributions` and `phenotype_colors`\u2192`palette`; add `covariate_key=None` (\u2020), `n_samples=0`, `phenotype_subset=None`, `device=None`, `show=None`; update the Sankey prose to reference `n_samples`-driven ribbons.\n9. **\u00a78.7 `pl.probability_ternary`** \u2014 rename `phenotype_names`\u2192`phenotypes`; add `covariate_key=None` (\u2020), `covariate=None`, `conditions=None`, `n_samples=0`, `temperature=1.0`, `scale_function=None`, `top_n=None`, `color=None`, `random_state=None`, `show=None`; do **not** add a cohort `splitby`.\n10. **\u00a78.5** \u2014 make `_metric_boxplot`'s `function` keyword-only; add `hue_order=None`. **\u00a78.6** \u2014 make `resolve_palette`'s `columns` keyword-only.\n11. **\u00a78.5 `_finish` / all \u00a78 public plots** \u2014 thread `show=None` through every public `pl.*` entry point (scanpy `show`/`save`/`return_*` triad); document the `return_df` (DataFrame) vs `return_axes` (figure) split so no plot omits both.\n12. **\u00a79.1 `joint_distribution_ppc`** \u2014 reorder to `(covariate, groupby, clones, temperature, distance_metric)`; add `groupby=None`, `clones=None`; note `n_samples`/`random_state` deferred unless draw-based comparison is adopted.\n13. **\u00a79.1 `reconstruction_ppc`** \u2014 `adata` keyword-only; `n_samples`\u2192`n_sims`; `seed`\u2192`random_state` (default `None`).\n14. **\u00a79.1 `permutation_null`** \u2014 `n_permutations`\u2192`n_perm`; `seed`\u2192`random_state`; add `splitby=None`, `clones=None`, `temperature=1.0`, `normalized=True`, `normalize_mode=\"min\"` pass-throughs.\n15. **\u00a73.3** \u2014 optionally rename `_stats.auc_and_label_permutation`/`bootstrap_auc` `seed`\u2192`random_state`, or document that the private primitives keep `seed` while the public diagnostics standardize on `random_state`.\n16. **\u00a73.4 / \u00a77.4 / appendix row 12** \u2014 set `_mi_from_joint` `eps` to `1e-12`, OR amend the \"one \u03b5=10\u207b\u00b9\u00b2 library-wide, matching entropy/MI\" claim to carve out the MI-specific `1e-15` with a numerical justification. Resolve the prose\u2194kernel contradiction either way.\n17. **\u00a70.10** \u2014 state the canonical RNG type once: `random_state: int | numpy.Generator | torch.Generator | None`, keyword-only, penultimate before `device`; note the `_stats` `seed` exception if kept.\n18. **\u00a76.1 `pp.group_singletons`** \u2014 add a note that `clonotype_key=\"trb\"` and `groupby=\"patient\"` are intentional pre-registration defaults (operate on the raw column before `setup_anndata`); do not change them.\n19. **\u00a711 census (correctness fix)** \u2014 move `clonality` (tl+pl), `probability_distribution`, `clone_fraction` **out** of \"Deleted \u2026 0 live callers\" into a \"removed **with** replacement + in-PR notebook rewrite (Phase 6/10)\" category; **add `compare_phenotypes`** (currently absent) as DROP-with-rewrite; correct the false \"0 live callers\" label. Add an explicit disposition line for the **`pl.flux` distance box-plot: DROP** (Sankey is the flux plot).\n20. **\u00a711 renamed/removed map** \u2014 add rows: legacy `method=` \u2192 removed (expressed via `n_samples`/`use_logits`); `phenotype_names` \u2192 `phenotypes` on `pl.probability_ternary`. Note that the caller census must distinguish package-qualified calls, alias-imported calls, and notebook-local `def` redefinitions (`polar_plot`, `clonality`), and that Phase-10 must strip shadow definitions.\n21. **\u00a70.11 (new, cross-cutting note)** \u2014 record the H4 open decision: because `p_ct` is covariate-locked (`ct_to_cov`), any `covariate_key` column override on flux/ternary must be computed on the **empirical per-cell-probability path**, not the `p_ct` engine path; the contract must state which path these plots use before `covariate_key` (\u2020) ships.", "consistency": {"matrix": "## PARAMETER MATRIX\n\nCells show `name=default` (kw-only unless noted), `name(req)` = required keyword-only (no default), `POS` = positional, `\u2014` = argument absent, and a **DIVERGENT** call-out where the concept exists under a different name. Only the analysis/plot/diag/pp/stats/engine surface carries the shared vocabulary; pure internal helpers (`_console`, priors/module/classifier/training internals, `_register._write_*`, session utils, `_compute` reducers, `resolve_device`/`get_xp`/`asnumpy`) have **no** shared-vocabulary args and are omitted (all-`\u2014` rows).\n\n### Table A \u2014 engine / tl / pl / diag / pp core vocabulary\n\n| Function (ns/module) | primary(pos) | covariate | cov_from/cov_to | order | groupby | splitby | n_samples | temperature | clones | use_logits | normalized | normalize_mode | distance_metric | n_clones_ref | random_state | device | palette | ax | figsize | save | return |\n|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|\n| tl.joint_distribution | adata | covariate=None | \u2014 | \u2014 | groupby=None | \u2014 | n_samples=0 | temperature=1.0 | clones=None | use_logits=True | \u2014 | \u2014 | \u2014 | \u2014 | random_state=None | device=None | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 |\n| tl.clonotypic_entropy | adata_or_jd | covariate=None | \u2014 | \u2014 | groupby=None | \u2014 | n_samples=0 | temperature=1.0 | clones=None | \u2014 | normalized=True | \u2014 | \u2014 | n_clones_ref=None | random_state=None | device=None | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 |\n| tl.phenotypic_entropy | adata_or_jd | covariate=None | \u2014 | \u2014 | groupby=None | \u2014 | n_samples=0 | temperature=1.0 | clones=None | \u2014 | normalized=True | \u2014 | \u2014 | \u2014 | random_state=None | device=None | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 |\n| tl.mutual_information | adata_or_jd | covariate=None | \u2014 | \u2014 | groupby=None | \u2014 | n_samples=0 | temperature=1.0 | clones=None | \u2014 | normalized=True | normalize_mode=\"min\" | \u2014 | \u2014 | random_state=None | device=None | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 |\n| tl.phenotypic_flux | adata | \u2014 | cov_from(req), cov_to(req) | \u2014 | groupby=None | \u2014 | n_samples=0 | temperature=1.0 | clones=None | \u2014 | \u2014 | \u2014 | distance_metric=\"l1\" | \u2014 | random_state=None | device=None | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 |\n| tl.compare_groups | df | \u2014 | \u2014 | \u2014 | **DIVERGENT: by(req)** | **DIVERGENT: by(req)** | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 |\n| pl.clonotypic_entropy | adata | covariate=None | \u2014 | **DIVERGENT: hue_order=None** | groupby=None | splitby=None | n_samples=0 | temperature=1.0 | clones=None | \u2014 | normalized=True | \u2014 | \u2014 | n_clones_ref=None | random_state=None | **\u2014 (missing)** | palette=None | ax=None | figsize=(6,3) | save=None | return_df=False |\n| pl.phenotypic_entropy | adata | covariate=None | \u2014 | \u2014 | groupby=None | splitby=None | n_samples=0 | temperature=1.0 | clones=None | \u2014 | normalized=True | \u2014 | \u2014 | \u2014 | random_state=None | **\u2014 (missing)** | palette=None | ax=None | figsize=(8,4) | save=None | return_df=False |\n| pl.mutual_information | adata | covariate=None | \u2014 | \u2014 | groupby=None | splitby=None | n_samples=0 | temperature=1.0 | clones=None | \u2014 | normalized=True | normalize_mode=\"min\" | \u2014 | \u2014 | random_state=None | **\u2014 (missing)** | palette=None | ax=None | figsize=(8,4) | save=None | return_df=False |\n| pl.phenotypic_flux | adata | \u2014 | \u2014 | order(req) | groupby=None | **\u2014 (missing)** | **\u2014 (missing)** | temperature=1.0 | clones=None | \u2014 | **DIVERGENT: normalize=True** | \u2014 | distance_metric=\"l1\" | \u2014 | random_state=None | **\u2014 (missing)** | **DIVERGENT: phenotype_colors=None** | ax=None | figsize=(6,3) | save=None | return_axes=False |\n| pl.probability_ternary | adata | \u2014 | \u2014 | \u2014 | groupby=None | **\u2014 (missing)** | \u2014 | \u2014 | clones=None | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | palette=None | ax=None | figsize=(5,5) | save=None | return_axes=False |\n| pl._metric_boxplot (priv) | adata, function POS | \u2014 | \u2014 | order=None | groupby=None | splitby=None | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | palette=None | ax=None | figsize=(8,4) | \u2014 | \u2014 |\n| pl._finish (priv) | fig, ax POS | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | ax POS | \u2014 | save=None | return_axes=False; **show=None (only here)** |\n| pl.resolve_palette | adata, columns POS | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | palette=None | \u2014 | \u2014 | \u2014 | \u2014 |\n| diag.joint_distribution_ppc | adata | covariate=None | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | temperature=1.0 | \u2014 | \u2014 | \u2014 | \u2014 | distance_metric=\"l1\" | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 |\n| diag.phenotype_calibration | adata | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 (n_bins=10) |\n| diag.reconstruction_ppc | model, adata=None POS | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | **DIVERGENT: n_samples=100** | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | **DIVERGENT: seed=0** | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 |\n| diag.permutation_null | adata | covariate=None | \u2014 | \u2014 | groupby=None | \u2014 | **\u2014 (no n_samples; metric=\"mutual_information\")** | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | **DIVERGENT: seed=0** | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 (n_permutations=1000) |\n| diag.loss | model | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | ax=None | \u2014 | save=None | \u2014 (log_scale=False) |\n| diag.archetypes | model | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | ax=None | \u2014 | save=None | \u2014 |\n| pp.group_singletons | adata | \u2014 | \u2014 | \u2014 | **DIVERGENT default: groupby=\"patient\"** | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 (clonotype_key=\"trb\", target_col=\"trb_unique\", min_clone_size=10) |\n| pp.clone_size | adata | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 (key_added=\"clone_size\", return_counts=False) |\n\n### Table B \u2014 distance / stats / contrast / RNG vocabulary\n\n| Function | primary(pos) | metric-selector | base / eps | by | value | reference | paired | hdi_prob | alternative | seed / random_state | resample count | pos_label |\n|---|---|---|---|---|---|---|---|---|---|---|---|---|\n| shared.phenotype_distance | p, q POS | **DIVERGENT: metric=\"l1\"** | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 |\n| shared.kl_divergence | p, q POS | \u2014 | base=2.0, eps=1e-12 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 |\n| shared.js_divergence | p, q POS | \u2014 | base=2.0, eps=1e-12 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 |\n| shared.hdi | samples POS | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | hdi_prob=0.94 | \u2014 | \u2014 | \u2014 | \u2014 |\n| shared.summarize | samples POS | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | hdi_prob=0.94 | \u2014 | \u2014 | \u2014 | \u2014 |\n| shared.mann_whitney | a, b POS | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | alternative=\"two-sided\" | \u2014 | \u2014 | \u2014 |\n| tl.compare_groups | df POS | \u2014 | \u2014 | by(req) | value(req) | reference=None | paired=False | hdi_prob=0.94 | alternative=\"two-sided\" | \u2014 | \u2014 | \u2014 |\n| shared.auc_and_label_permutation | scores, labels POS | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | **seed=42** | n_perm=200_000 (max_exact=200_000) | pos_label=None |\n| shared.bootstrap_auc | scores, labels POS | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | **seed=42** | n_boot=5000 | pos_label=None |\n| diag.reconstruction_ppc | model POS | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | **seed=0** | n_samples=100 | \u2014 |\n| diag.permutation_null | adata POS | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | **seed=0** | n_permutations=1000 | \u2014 |\n| _compute._joint_draws | adata POS | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | \u2014 | random_state(kw) | n_samples(kw) | \u2014 (gate_prob kw; **order/set differs from tl.joint_distribution**) |\n\n### Consistency legend\n- **RNG:** `random_state=None` (engine+tl+pl) vs `seed=42` (stats) vs `seed=0` (diag) \u2014 3 spellings/defaults for one concept.\n- **distance selector:** `distance_metric=\"l1\"` (tl/pl flux, diag ppc) vs `metric=\"l1\"` (phenotype_distance).\n- **clones\u2194temperature order:** engine + pl.phenotypic_flux emit `clones` before `temperature`; every other metric emits `temperature` before `clones`.\n- **normalize vs normalized:** `normalize=` (bool, distribution) only on pl.phenotypic_flux; `normalized=` (bool, [0,1] metric scaling) everywhere else.\n- **device:** present on all tl metrics + engine, **absent on every pl.* metric plot** (which still expose n_samples/random_state).\n- **splitby:** present on tl\u2192pl box plots; **absent on pl.phenotypic_flux and pl.probability_ternary** (both faceted by cohort in notebooks).\n- **figsize defaults:** (6,3) clonotypic_entropy/flux \u00b7 (8,4) phenotypic_entropy/MI/_metric_boxplot \u00b7 (5,5) ternary.\n- **Clean (no findings):** `temperature=1.0`, `n_samples=0` convention (except reconstruction_ppc), `hdi_prob=0.94`, `alternative=\"two-sided\"` are consistent; `use_logits` correctly engine-only; **no `point_estimate=`/`posterior=` survivors** and **no British-spelling residue** anywhere in the target surface.", "findings": [{"category": "rng-name-default", "issue": "The random-number concept has three spellings and three defaults across the surface: random_state=None (tl.joint_distribution, all tl metrics, all pl metrics, _joint_draws), seed=42 (shared.auc_and_label_permutation, shared.bootstrap_auc), and seed=0 (diag.reconstruction_ppc, diag.permutation_null). Same concept, divergent name/default/type.", "severity": "high", "fix": "CANONICAL: random_state=None, keyword-only, typed int | numpy.Generator | torch.Generator | None, placed as the penultimate arg (before device) everywhere. Rename seed->random_state in auc_and_label_permutation, bootstrap_auc, reconstruction_ppc, permutation_null; drop the 42/0 literal defaults for None (seed internally from the passed generator).", "functions_affected": "shared.auc_and_label_permutation, shared.bootstrap_auc, diag.reconstruction_ppc, diag.permutation_null (align to tl/pl/engine random_state=None)"}, {"category": "arg-order", "issue": "The relative order of `clones` and `temperature` is inconsistent. tl.joint_distribution emits (n_samples, use_logits, clones, temperature) and pl.phenotypic_flux emits (clones, normalize, temperature) \u2014 clones BEFORE temperature \u2014 while every tl metric and pl metric box plot emits (n_samples, temperature, clones) \u2014 temperature BEFORE clones. Violates the same-relative-order rule.", "severity": "high", "fix": "CANONICAL order for the shared block: covariate/cov_*, groupby, splitby, n_samples, temperature, clones, use_logits (engine-only), normalized/normalize_mode/distance_metric, random_state, device. Reorder tl.joint_distribution and pl.phenotypic_flux so clones follows temperature.", "functions_affected": "tl.joint_distribution, pl.phenotypic_flux (align to the temperature-before-clones majority)"}, {"category": "missing-arg", "issue": "device=None is exposed on the engine and all tl metrics but is ABSENT from every pl.* metric plot (pl.clonotypic_entropy, pl.phenotypic_entropy, pl.mutual_information, pl.phenotypic_flux), even though those plots do expose n_samples>0 and random_state. A user can request posterior draws and set the seed but cannot steer the draw onto GPU/CPU.", "severity": "high", "fix": "CANONICAL: add device=None (keyword-only, last positional-of-vocabulary) to all pl.* metric plots and thread it into the underlying tl call, mirroring the tl signatures.", "functions_affected": "pl.clonotypic_entropy, pl.phenotypic_entropy, pl.mutual_information, pl.phenotypic_flux"}, {"category": "missing-splitby", "issue": "splitby is present on the pl metric box plots (clonotypic_entropy/phenotypic_entropy/mutual_information) and _metric_boxplot, but MISSING on pl.probability_ternary and pl.phenotypic_flux \u2014 both of which are faceted by cohort in the notebooks (probability_ternary is the heaviest plot, 19 calls, invoked with splitby+conditions; phenotypic_flux invoked with splitby).", "severity": "high", "fix": "CANONICAL: add splitby=None (immediately after groupby) to pl.probability_ternary and pl.phenotypic_flux, matching the distinct groupby(unit)/splitby(cohort) semantics used by the box-plot family.", "functions_affected": "pl.probability_ternary, pl.phenotypic_flux"}, {"category": "n_samples-convention", "issue": "diag.reconstruction_ppc uses n_samples=100, colliding with the surface-wide convention that n_samples=0 means deterministic posterior-mean and n_samples>0 means draws. Here n_samples is the number of simulated PPC datasets (a different concept) with a non-zero default, so a user reading it as the draw-count knob gets the wrong mental model.", "severity": "medium", "fix": "CANONICAL: rename to n_sims (or n_ppc_samples), keep default 100. Reserve `n_samples` (default 0) exclusively for the posterior-draw / point-estimate convention.", "functions_affected": "diag.reconstruction_ppc"}, {"category": "distance-metric-name", "issue": "The distance-selector concept is `distance_metric=\"l1\"` on tl.phenotypic_flux, pl.phenotypic_flux and diag.joint_distribution_ppc, but `metric=\"l1\"` on shared.phenotype_distance (the dispatcher). Same concept, two names.", "severity": "medium", "fix": "CANONICAL: distance_metric everywhere on the public/tl/pl/diag surface. Rename phenotype_distance's parameter to distance_metric (or keep `metric` only as an internal dispatcher detail not exposed to users).", "functions_affected": "shared.phenotype_distance (vs tl.phenotypic_flux, pl.phenotypic_flux, diag.joint_distribution_ppc)"}, {"category": "normalize-vs-normalized", "issue": "pl.phenotypic_flux uses normalize=True (a distribution-normalization bool) and phenotype_colors=None, while the rest of the surface uses normalized=True (a [0,1] metric-scaling bool) and palette=None. `normalize` vs `normalized` is a near-collision that reads as the same knob but is not; phenotype_colors duplicates palette.", "severity": "medium", "fix": "CANONICAL: rename phenotype_colors->palette. Disambiguate the bool: either keep normalized reserved for metric scaling and rename the flux bool to normalize_distributions=True, or drop it if the engine always feeds normalized distributions. Do not ship both `normalize` and `normalized` as look-alike names.", "functions_affected": "pl.phenotypic_flux"}, {"category": "figsize-default", "issue": "figsize defaults diverge across sibling box plots: pl.clonotypic_entropy=(6,3), pl.phenotypic_entropy=(8,4), pl.mutual_information=(8,4), _metric_boxplot=(8,4); pl.phenotypic_flux=(6,3); pl.probability_ternary=(5,5). The three metric box plots that share _metric_boxplot should share its default.", "severity": "low", "fix": "CANONICAL: figsize=(8,4) for the metric box-plot family (matching _metric_boxplot); set pl.clonotypic_entropy to (8,4). Aspect-driven plots (ternary (5,5), Sankey flux (6,3)) may keep purpose-specific defaults but should be documented as intentional.", "functions_affected": "pl.clonotypic_entropy (vs pl.phenotypic_entropy, pl.mutual_information, _metric_boxplot)"}, {"category": "order-hue_order", "issue": "Category-ordering controls are inconsistent across the box-plot family: pl.clonotypic_entropy exposes hue_order (but no order), pl.phenotypic_entropy and pl.mutual_information expose neither, while the shared _metric_boxplot exposes order. Notebooks call these plots with order=. So the wrapper drops an ordering knob its own engine supports, and only one of three siblings has hue_order.", "severity": "medium", "fix": "CANONICAL: expose both order=None (x-axis/covariate category order) and hue_order=None (splitby level order) on all three metric box plots, wired to _metric_boxplot.order and the splitby hue. Keep the names order/hue_order consistent everywhere.", "functions_affected": "pl.clonotypic_entropy, pl.phenotypic_entropy, pl.mutual_information"}, {"category": "contrast-cohort-name", "issue": "tl.compare_groups uses `by` for the column whose levels are contrasted (with reference/paired) \u2014 semantically the comparison COHORT, i.e. the same concept the plotting surface calls `splitby`. Divergent name for the cohort concept on the core stats surface.", "severity": "medium", "fix": "CANONICAL: rename by->splitby in compare_groups (keeping value/reference/paired/hdi_prob/alternative) so the aggregation-unit(groupby)/comparison-cohort(splitby) vocabulary is uniform from metric to plot to contrast.", "functions_affected": "tl.compare_groups"}, {"category": "pp-default-drift", "issue": "pp.group_singletons.clonotype_key defaults to \"trb\" while TCRIModel.setup_anndata.clonotype_key defaults to \"unique_clone_id\" (same concept, different default), and pp.group_singletons.groupby defaults to \"patient\" whereas groupby defaults to None everywhere else on the analysis surface.", "severity": "medium", "fix": "CANONICAL: clonotype_key default should be None (require explicit) or match the registered default \"unique_clone_id\"; groupby default should be None for cross-surface consistency (document if a pre-registration pp step legitimately needs a concrete default).", "functions_affected": "pp.group_singletons (vs TCRIModel.setup_anndata; vs all groupby=None callers)"}, {"category": "engine-helper-mismatch", "issue": "tl.joint_distribution emits (covariate, groupby, n_samples, use_logits, clones, temperature, random_state, device) while its own helper _compute._joint_draws emits (covariate, clones, n_samples, use_logits, temperature, gate_prob, random_state, device): clones jumps position, temperature/use_logits reorder, _joint_draws adds gate_prob and drops groupby. The extraction itself flags this order/set divergence.", "severity": "medium", "fix": "CANONICAL: make _joint_draws' keyword order a strict subset-in-order of joint_distribution (covariate, n_samples, temperature, clones, use_logits, gate_prob, random_state, device); resolve gate_prob at the public layer or document it as an internal-only extra.", "functions_affected": "tl.joint_distribution, _compute._joint_draws"}, {"category": "ppc-arg-order-and-gaps", "issue": "diag.joint_distribution_ppc orders args (covariate, distance_metric, temperature) \u2014 temperature AFTER the metric selector, whereas the canonical block places temperature before distance_metric \u2014 and it omits groupby/n_samples/random_state that its metric siblings expose.", "severity": "low", "fix": "CANONICAL: reorder to (covariate, temperature, distance_metric); add n_samples=0 and random_state=None if draw-based comparison is intended, and groupby=None for symmetry with permutation_null.", "functions_affected": "diag.joint_distribution_ppc"}, {"category": "positional-should-be-kwonly", "issue": "Rule requires only the primary data object to be positional. Violations: pl.resolve_palette(adata, columns) \u2014 columns is a 2nd positional; pl._metric_boxplot(adata, function) \u2014 function is a 2nd positional; diag.reconstruction_ppc(model, adata=None) \u2014 adata is a 2nd positional (model is primary).", "severity": "low", "fix": "CANONICAL: make columns, function, and adata keyword-only (place a `*` after the single primary positional). For reconstruction_ppc, adata=None stays keyword-only.", "functions_affected": "pl.resolve_palette, pl._metric_boxplot, diag.reconstruction_ppc"}, {"category": "resample-count-name", "issue": "The resample-count concept is n_permutations=1000 in diag.permutation_null but n_perm=200_000 in shared.auc_and_label_permutation \u2014 same concept (permutation count), two names. (n_boot=5000 in bootstrap_auc is a distinct bootstrap concept and can keep its name.)", "severity": "low", "fix": "CANONICAL: n_perm for permutation counts everywhere; rename permutation_null.n_permutations->n_perm.", "functions_affected": "diag.permutation_null (vs shared.auc_and_label_permutation)"}, {"category": "show-arg", "issue": "A scanpy-style show= toggle exists only on the private pl._finish (show=None); no public plotting function exposes show=, though they all expose save= and return_df/return_axes. Inconsistent display-control vocabulary.", "severity": "low", "fix": "CANONICAL: thread show=None through every public plotting function (scanpy convention: show/save/return_*), or drop show from _finish. Pick one policy and apply it to all pl.* entry points.", "functions_affected": "all pl.* public plots (vs pl._finish)"}, {"category": "return-control-name", "issue": "Return-control name splits by return type: return_df=False on metric box plots (clonotypic_entropy, phenotypic_entropy, mutual_information) vs return_axes=False on pl.phenotypic_flux, pl.probability_ternary, pl._finish. Same 'give me the underlying object' intent, two names.", "severity": "low", "fix": "CANONICAL: acceptable to keep return_df (DataFrame) vs return_axes (Axes) since the payloads differ, BUT standardize which every plot exposes: metric plots -> return_df; figure-only plots -> return_axes; document the split. Do not let a single plot omit both.", "functions_affected": "pl.phenotypic_flux, pl.probability_ternary (return_axes) vs pl.clonotypic_entropy/phenotypic_entropy/mutual_information (return_df)"}, {"category": "eps-default", "issue": "eps has four defaults for probability-space clipping: 1e-8 (TCRIModel.predict, _compute_logits_and_prior), 1e-6 (TCRIModule.prepare_two_level_params), 1e-12 (kl_divergence, js_divergence), 1e-15 (hardcoded in _mi_from_joint). Where interchangeable these should not drift.", "severity": "low", "fix": "CANONICAL: pick one probability-clip floor (e.g. eps=1e-12) for the distance/MI/normalization paths and reuse it; keep context-specific floors only where numerically justified and document why.", "functions_affected": "shared.kl_divergence, shared.js_divergence, tl._mi_from_joint, TCRIModel.predict, _compute_logits_and_prior, TCRIModule.prepare_two_level_params"}, {"category": "covariate-none-semantics", "issue": "Rule: covariate=None must mean 'ALL covariate values (one shared-draw pass)' consistently. All signatures default covariate=None, but the extraction flags a known tension in the adata_or_jd fast-path text where, when a precomputed joint (jd) is passed to tl.clonotypic_entropy/phenotypic_entropy/mutual_information, covariate may read as 'required' rather than 'all'.", "severity": "medium", "fix": "CANONICAL: covariate=None ALWAYS means all covariate values, in both the adata path and the adata_or_jd fast path; when a precomputed jd is supplied, covariate merely filters and None still selects all. Ensure the entropy/MI docstrings and dispatch never treat covariate as required.", "functions_affected": "tl.clonotypic_entropy, tl.phenotypic_entropy, tl.mutual_information (adata_or_jd fast path)"}], "summary": "Built a full parameter matrix (Table A: engine/tl/pl/diag/pp core vocabulary across 21 columns; Table B: distance/stats/contrast/RNG vocabulary) over the target surface; pure internal helpers with no shared args are omitted as all-'\u2014' rows. The surface is mostly coherent \u2014 temperature=1.0, the n_samples=0 point-estimate convention, hdi_prob=0.94, alternative=\\\"two-sided\\\" are uniform; use_logits is correctly engine-only; and there are NO point_estimate=/posterior= survivors and NO British-spelling residue. The real defects cluster in five areas. (1) RNG fragmentation: random_state=None vs seed=42 vs seed=0 \u2014 canonicalize to random_state=None everywhere. (2) Order drift: clones/temperature swap between the engine (+pl.phenotypic_flux) and every other metric \u2014 canonicalize to temperature-before-clones. (3) Missing shared args where applicable: device absent from all pl.* metric plots (though they expose draws+random_state); splitby absent from pl.probability_ternary and pl.phenotypic_flux (both cohort-faceted in notebooks); order/hue_order inconsistent across the box-plot family. (4) Name divergences for one concept: distance_metric vs metric (phenotype_distance), palette vs phenotype_colors and normalized vs normalize (pl.phenotypic_flux), by vs splitby (compare_groups), n_permutations vs n_perm. (5) Convention/default drift: reconstruction_ppc.n_samples=100 overloads the n_samples=0 convention (rename to n_sims); group_singletons defaults (clonotype_key=\\\"trb\\\", groupby=\\\"patient\\\") diverge from the registered/None defaults; figsize defaults split (6,3)/(8,4)/(5,5); a covariate=None-means-all vs adata_or_jd-fast-path 'required' tension must be resolved to 'all'. 20 findings total (4 high, 8 medium, 8 low), each with a single canonical name/default/order decision."}, "sufficiency": {"findings": [{"category": "sufficiency", "severity": "high", "functions_affected": "pl.probability_ternary (tcri/plotting/_ternary.py)", "issue": "The heaviest plotting fn (19 notebook calls) is stripped of nearly all its real knobs. Historical signature (176b62e:_plotting.py:392) is probability_ternary(adata, phenotype_names, covariate, splitby, conditions, n_samples=1, temperature=1, top_n=None, scale_function=None, color='k', save=None). Notebooks depend on: splitby+conditions to pick the cohort/timepoint pair drawn as start/end simplices (zhang/yost: splitby='timepoint_relapse', conditions=['CR_Pre-treatment','CR_Post-treatment']); scale_function=lambda freq: size (marker sizing by clone frequency, used in EVERY zhang/yost call); covariate to condition the posterior joint; n_samples+temperature for the probabilistic draw; top_n to cap displayed clones; color for marker color. The extracted surface exposes only phenotypes, groupby, clones, palette, ax, figsize, save, return_axes. groupby is a mis-mapping of splitby; conditions, scale_function, covariate, n_samples, temperature, top_n, color, random_state are all absent, making the notebook figures non-reproducible.", "fix": "Rename groupby->splitby; add conditions=None (list[str], 1 or 2 covariate/cohort levels -> start/end), covariate=None (str, value to condition on), scale_function=None (Callable[[float],float] freq->marker size), top_n=None (int), n_samples=0 (int), temperature=1.0 (float), color=None (str), random_state=None (int|Generator|None)."}, {"category": "sufficiency", "severity": "high", "functions_affected": "pl.clonotypic_entropy, pl.phenotypic_entropy, pl.mutual_information (tcri/plotting/_entropy.py,_mutual_information.py)", "issue": "The public per-phenotype/per-cohort box-plots drop `order=`, the x-axis category ordering. The private engine _metric_boxplot DOES take order= (extracted sig) and the old tcri_boxplot threaded it (_plotting.py:598,623-651), and notebooks pass it constantly: pl.clonotypic_entropy(...,order=['Naive/CM','Memory','Dsyfunctional','Activated']) (preprocess_conga, zhang), pl.phenotypic_entropy(...,order=['Yes_pre','Yes_post','No_pre','No_post']) and order=['Yes_bcc','No_bcc','Yes_scc','No_scc'] (yost). Without order= the categorical axis falls back to median-sort, so the published panels cannot be reproduced.", "fix": "Add order=None (list[str] | None, kwonly) to all three public pl metric plots and pass it through to _metric_boxplot; orders the phenotype (or splitby) categories on the x-axis."}, {"category": "sufficiency", "severity": "medium", "functions_affected": "tl.joint_distribution, tl.mutual_information, pl.probability_ternary", "issue": "`weighted=` (clone-size / abundance weighting of the joint) is removed everywhere. Old joint_distribution (_preprocessing.py:364,470-484) and joint_distribution_posterior (:266,323) took weighted=, threaded through mutual_information (_metrics.py:633) and ternary; yost_v2 explicitly calls mutual_information(...,weighted=False) and old pl.mutual_information defaulted weighted=True. The unweighted joint (per-clone uniform) and the abundance-weighted joint are genuinely different distributions and change MI/entropy/ternary mass. With no arg, only one weighting is reachable and old weighted results are non-reproducible.", "fix": "Confirm the intentional removal; if kept, add weighted=False (bool, kwonly) to joint_distribution and have MI/entropy/ternary thread it \u2014 OR document the canonical decision that the engine is always unweighted and note the behavioral change from the old weighted=True default."}, {"category": "sufficiency", "severity": "medium", "functions_affected": "tl.joint_distribution, tl.phenotypic_flux, pl.phenotypic_flux", "issue": "No way to select the covariate COLUMN at call time. `covariate=` is a VALUE (or None='all values') and the column is frozen at setup_anndata(covariate_key=). But notebooks run flux/Sankey across MULTIPLE different obs columns on the same adata: old pl.flux/phenotypic_flux/tl.flux used key=/splitby= = 'treatment','timepoint','timepoint_relapse','timepoint_instant' (yost, zhang) \u2014 only one of which can be the registered covariate_key. Reproducing the others requires a full re-run of setup_anndata, which is impractical.", "fix": "Add covariate_key=None (str | None, kwonly) to joint_distribution / tl.phenotypic_flux / pl.phenotypic_flux to override the registered covariate column that `covariate`/`cov_from`/`cov_to`/`order` index into (defaults to the registered covariate_key)."}, {"category": "sufficiency", "severity": "medium", "functions_affected": "pl.phenotypic_flux (tcri/plotting/_flux.py)", "issue": "The Sankey loses two used knobs. Old phenotypic_flux (176b62e:_plotting.py:287) had n_samples=50 (posterior draws that give the ribbon uncertainty) and yost_v2 passed phenotype_subset=phenos to restrict which phenotype nodes are drawn. The extracted surface has random_state but NO n_samples to drive it, and no phenotype_subset \u2014 so the Sankey is forced to a single point estimate over all phenotypes.", "fix": "Add n_samples=0 (int, kwonly) and phenotype_subset=None (list[str] | None, kwonly) to pl.phenotypic_flux (n_samples matches the shared engine name; phenotype_subset restricts the rendered phenotype columns)."}, {"category": "sufficiency", "severity": "medium", "functions_affected": "pl.clonotypic_entropy, pl.phenotypic_entropy, pl.mutual_information", "issue": "`minimum_clone_size=` is dropped. Notebooks filter by it (comparison.ipynb: pl.mutual_information(...,minimum_clone_size=1), phenotypic_entropy(...,minimum_clone_size=0)); implemented as adata=adata[obs.clone_size>=k]. Reproducible by manual slicing or by pre-building clones=, but the convenience knob listed in capabilities-to-preserve has no direct home and the recorded calls won't run as written.", "fix": "Add minimum_clone_size=None (int | None, kwonly) to the three pl metric plots (and optionally the tl metrics), filtering clones with clone_size below the threshold before aggregation."}, {"category": "sufficiency", "severity": "medium", "functions_affected": "tl.compare_groups (tcri/tools/_compare.py)", "issue": "compare_groups can reproduce mi_compare (one contrast axis via by=/reference=), but NOT the paired delta_entropy_table/flux_table pattern in one call. delta_entropy_table (_metrics.py:392) pairs post-pre WITHIN a unit (clone/patient) across the covariate AND then contrasts across a cohort (splitby='response'). paired=True has no pairing-key argument, so it cannot know which pre row matches which post row, and there is no second axis to also split by cohort. The signed-delta-then-cohort-contrast requires two manual passes.", "fix": "Add pair_on=None (str | None, kwonly) naming the unit column that aligns the two `by` levels under paired=True; document (or add a `split=`/second-axis) to reproduce paired-delta + cohort-contrast in one call."}, {"category": "sufficiency", "severity": "medium", "functions_affected": "pp.gene_entropy (unmapped; absent from extracted surface)", "issue": "pp.gene_entropy(adata, key_added='entropy', batch_key=None, agg_function=None) (_preprocessing.py:519) is used in 5 notebooks, twice with the batch_key knob (smith: gene_entropy(adata, batch_key='sample')). It has no home in the extracted surface, so its per-batch gene-entropy capability (and batch_key/agg_function knobs) is lost.", "fix": "Restore gene_entropy under pp with key_added='entropy' (str), batch_key=None (str | None), agg_function=None (Callable | None) \u2014 or record the decision to drop it."}, {"category": "sufficiency", "severity": "medium", "functions_affected": "pl.flux box-plot (unmapped; only the Sankey pl.phenotypic_flux survives)", "issue": "The old per-clone flux-DISTANCE box-plot pl.flux(adata, key, order, groupby, paint=, paint_dict=, method=, distance_metric=, paint_order=, palette=) (_plotting.py:1046) is distinct from the Sankey and is used in renal/yost/zhang (e.g. tcri.pl.flux(adata, key='treatment', order=['pre','post'], groupby='response_disease', distance_metric='dkl', paint='response')). The new surface has only the Sankey; the box-plot of flux distances across groups colored by a cohort (paint=) has no public home (private _metric_boxplot exists but isn't wired to flux, and lacks paint/distance_metric).", "fix": "Add a public pl.phenotypic_flux box variant (or return_df on the flux tool) exposing key/order (covariate), groupby (unit), splitby (=old paint, cohort hue), distance_metric='l1' (str), palette=None, figsize=."}, {"category": "sufficiency", "severity": "medium", "functions_affected": "_stats.auc_and_label_permutation, _stats.bootstrap_auc, diag.reconstruction_ppc, diag.permutation_null", "issue": "Reproducibility knob is named inconsistently: the tl/pl surface uses random_state everywhere, but these four use seed= (auc_and_label_permutation seed=42, bootstrap_auc seed=42, reconstruction_ppc seed=0, permutation_null seed=0). Same concept, two names \u2014 violates the 'same CONCEPT -> same NAME' rule and the shared random_state vocabulary.", "fix": "Rename seed -> random_state (int | numpy.Generator | None) on all user-facing diagnostics (reconstruction_ppc, permutation_null) at minimum; optionally keep seed on the low-level _stats primitives but document the split."}, {"category": "sufficiency", "severity": "medium", "functions_affected": "diag.permutation_null (tcri/diagnostics/_ppc.py)", "issue": "permutation_null(metric='mutual_information', covariate, groupby, n_permutations, seed) cannot thread metric-specific parameters to the metric it permutes. mutual_information needs temperature/normalized/normalize_mode; entropy metrics need temperature/normalized/n_clones_ref; all support clones=. None are exposable, so the null is computed with fixed defaults that may not match the observed statistic the user actually reported.", "fix": "Add temperature=1.0 (float), normalized=True (bool), normalize_mode='min' (str), clones=None (list[str] | None), splitby=None (str | None) pass-throughs (or a metric_kwargs dict) and rename seed->random_state."}, {"category": "sufficiency", "severity": "low", "functions_affected": "tl.clonotypic_entropy, tl.phenotypic_entropy (and pl twins)", "issue": "`base=` (log base) and `decimals=` (rounding) are dropped. Old clonotypic/phenotypic entropy took base=2 (_metrics.py:183,235) and the plural tl.phenotypic_entropies was called with decimals= in notebooks. base is base-invariant only when normalized=True; with normalized=False the reported entropy scale changes. Cosmetic-to-minor, but a listed capability with no home.", "fix": "Add base=2 (int, kwonly) to the entropy metrics (matters when normalized=False); decimals is cosmetic \u2014 drop or add decimals=None (int | None) for parity."}, {"category": "sufficiency", "severity": "low", "functions_affected": "pl.phenotypic_entropy, pl.mutual_information", "issue": "`hue_order=` is present on pl.clonotypic_entropy but absent on pl.phenotypic_entropy and pl.mutual_information, though all three draw a splitby hue. Same concept should appear consistently, so the splitby legend/dodge order can't be fixed on two of the three twins.", "fix": "Add hue_order=None (list[str] | None, kwonly) to pl.phenotypic_entropy and pl.mutual_information to match pl.clonotypic_entropy."}, {"category": "sufficiency", "severity": "low", "functions_affected": "tl.clonotypic_entropy, tl.mutual_information (adata_or_jd fast-path)", "issue": "covariate=None is documented as 'ALL covariate values (one shared-draw pass)' but when adata_or_jd is a precomputed jd DataFrame the covariate/groupby/n_samples/temperature/clones/random_state/device args are inert and covariate=None effectively means 'the single covariate already baked into this jd'. Same arg, two meanings depending on input type \u2014 the known tension, and it should be pinned down so callers know the jd path ignores covariate.", "fix": "Decide and document: when a jd is passed, covariate/n_samples/etc. must be None and are ignored (raise if set), preserving covariate=None == 'all values' only on the adata path."}, {"category": "sufficiency", "severity": "low", "functions_affected": "diag.joint_distribution_ppc, diag.permutation_null", "issue": "Both compute per-clone quantities but omit clones= subsetting; joint_distribution_ppc also omits groupby/n_samples/random_state (it is deterministic and whole-repertoire only). This blocks per-cohort or per-clone-subset PPCs that the metric functions elsewhere support via clones=/groupby=.", "fix": "Add clones=None (list[str] | None) and groupby=None (str | None) to joint_distribution_ppc; add clones=None to permutation_null (kwonly)."}, {"category": "sufficiency", "severity": "low", "functions_affected": "Auxiliary plots/tools absent from extracted surface: pl.top_clone_umap, pl.polar_plot, tl.probability_distribution, pl.compare_phenotypes, tl/pl.clonality, tl.clone_fraction", "issue": "Several notebook-exercised capabilities have no new home: pl.top_clone_umap(top_n,bg_size,bg_alpha,fg_size,figsize) (5+ calls across renal/yost/zhang/smith), pl.polar_plot(statistic,splitby,color_dict,title) (colitis/smith/zhang), tl.probability_distribution(method) (4 calls), pl.compare_phenotypes (synthetic), and clonality/clone_fraction (imported). These are whole-function omissions rather than arg gaps, but they are listed capabilities-to-preserve.", "fix": "Decide per function whether to port to the new surface (with groupby/splitby/palette/save parity) or explicitly document as deprecated; at minimum port top_clone_umap and polar_plot given their notebook frequency."}], "summary": "Argument-sufficiency audit of the extracted TCRi surface against the old-notebook capabilities. The surface is mostly sufficient for the core engine (joint_distribution/entropy/MI/flux carry covariate/n_samples/temperature/clones/random_state/device consistently, covariate=None='all', use_logits replaces posterior), but 16 knob gaps remain. HIGH: (1) pl.probability_ternary \u2014 the heaviest plotting fn (19 calls) is stripped of splitby, conditions, covariate, n_samples, temperature, scale_function, top_n, color, random_state, making its figures non-reproducible; (2) pl.clonotypic_entropy/phenotypic_entropy/mutual_information drop the x-axis order= that the private _metric_boxplot supports and notebooks constantly pass. MEDIUM: weighted= (abundance weighting) removed from joint_distribution/MI/ternary; no call-time covariate-COLUMN selector though notebooks Sankey/flux across several different obs columns; pl.phenotypic_flux lacks n_samples= and phenotype_subset=; minimum_clone_size= dropped from pl metric plots; compare_groups paired mode lacks a pairing-key/second axis to reproduce paired-delta+cohort-contrast (delta_entropy_table); pp.gene_entropy(batch_key) and the pl.flux distance box-plot (paint=) have no home; seed vs random_state naming split in diagnostics; permutation_null can't thread metric params. LOW: base=/decimals= on entropy, hue_order= inconsistency, adata_or_jd covariate=None dual meaning, missing clones=/groupby= on diag PPCs, and absent auxiliaries (top_clone_umap, polar_plot, probability_distribution, compare_phenotypes, clonality, clone_fraction)."}, "verify": {"verdicts": [{"finding": "C1 rng-name-default: seed=42 (auc_and_label_permutation, bootstrap_auc), seed=0 (reconstruction_ppc, permutation_null) vs random_state=None everywhere else", "status": "confirmed", "note": "Exact in doc: lines 266,267 (seed=42), 731,732 (seed=0), \u00a70.10 random_state on tl/pl/engine. Fix is correct. Nuance: auc_and_label_permutation/bootstrap_auc are PRIVATE (_stats.py, not re-exported), so renaming them is optional; the load-bearing renames are the PUBLIC diag.reconstruction_ppc/permutation_null. Duplicate of sufficiency S10."}, {"finding": "C2 arg-order: clones-before-temperature in joint_distribution & phenotypic_flux vs temperature-before-clones in all metrics", "status": "confirmed", "note": "joint_distribution (lines 397-399): use_logits,clones,temperature. pl.phenotypic_flux (674-675): clones,normalize,temperature. Metrics (461,497,519): temperature,clones. Real divergence. Fix (temperature-before-clones majority) is correct and dovetails with C12's canonical _joint_draws order."}, {"finding": "C3 missing device on all pl.* metric plots", "status": "revised", "note": "device absent from pl.clonotypic_entropy/phenotypic_entropy/mutual_information/phenotypic_flux \u2014 CONFIRMED (no device in \u00a78.1-8.3). But the finding's premise is wrong for flux: pl.phenotypic_flux does NOT expose n_samples (\u00a78.3 has random_state but no n_samples), so on flux the dangling knob is random_state, not 'draws-without-device'. Fix (add device) correct for the three box plots; for flux, pair with S5 (add n_samples too)."}, {"finding": "C4 missing-splitby on probability_ternary & phenotypic_flux", "status": "revised", "note": "Real gap (both lack the arg notebooks pass) BUT the proposed fix mis-names it. Legacy notebooks pass splitby= to BOTH \u2014 however there splitby is the covariate COLUMN (splitby=\"treatment\"/\"timepoint\" with order/conditions selecting VALUES), NOT the box-plot cohort-hue splitby. Adding a hue-semantics 'splitby' would re-introduce the very name-collision the RULES forbid. Correct fix: covariate_key (column) + conditions (ternary values)/order (Sankey values); reserve splitby for cohort hue. See missed item on overloaded splitby. Overlaps/conflicts with S1's 'rename groupby->splitby'."}, {"finding": "C5 n_samples-convention: reconstruction_ppc n_samples=100 collides with n_samples=0/point-vs-draws convention", "status": "confirmed", "note": "Line 731: n_samples=100 = number of simulated PPC datasets, a different concept than \u00a70.5's draw-count. Real collision. Rename to n_sims/n_ppc_samples is correct."}, {"finding": "C6 distance-metric-name: metric= on phenotype_distance vs distance_metric= elsewhere", "status": "confirmed", "note": "Line 278 phenotype_distance(...metric=\"l1\") vs distance_metric on tl/pl flux (552,676) and diag ppc (729). Real. phenotype_distance is PRIVATE (_distance.py), so 'keep metric as internal dispatcher detail' is the lower-churn valid option the fix already offers."}, {"finding": "C7 normalize vs normalized + phenotype_colors vs palette on pl.phenotypic_flux", "status": "confirmed", "note": "Line 675 normalize=True, 677 phenotype_colors=None. normalize (distribution-normalization, feeds _phenotype_mass_per_clone line 688) genuinely differs from normalized (metric scaling) \u2014 a real look-alike collision. Notebooks use phenotype_colors=main_color_map, confirming the duplicate-of-palette. Fix correct."}, {"finding": "C8 figsize-default drift across box-plot siblings", "status": "confirmed", "note": "clonotypic_entropy (6,3) line 631 vs phenotypic_entropy/mutual_information/_metric_boxplot (8,4). Set clonotypic_entropy to (8,4). Correct, low severity."}, {"finding": "C9 order/hue_order inconsistency across metric box plots", "status": "confirmed", "note": "clonotypic_entropy has hue_order but no order (631); phenotypic_entropy/mutual_information have neither; _metric_boxplot has order (694). Notebooks pass order= (verified: preprocess_conga, zhang, yost) AND minimum_clone_size. Expose order+hue_order on all three. Merges S2 (order) and S13 (hue_order)."}, {"finding": "C10 compare_groups by= vs splitby=", "status": "confirmed", "note": "Line 586 by= is the contrast cohort (e.g. response) = splitby concept; compare_groups IS the stats test the RULES tie to splitby. Rename by->splitby is consistent with the rulebook. Judgment call: 'by' is pandas-idiomatic, so this is a defensible-either-way naming choice rather than a hard bug."}, {"finding": "C11 pp.group_singletons defaults clonotype_key=\"trb\", groupby=\"patient\"", "status": "revised", "note": "Values confirmed (line 364). But the fix (force None / match unique_clone_id) is likely wrong: group_singletons runs BEFORE setup_anndata (\u00a76.1), operating on the RAW pre-registration column (reads 'trb', writes target_col='trb_unique' which later becomes clonotype_key), so 'unique_clone_id' doesn't exist yet. And its groupby is the collapse-UNIT (consistent with groupby=aggregation-unit), only the default differs. Better fix: document these as intentional pre-registration defaults; don't force None."}, {"finding": "C12 engine-helper-mismatch: _joint_draws vs joint_distribution keyword order/set", "status": "confirmed", "note": "Line 296 _joint_draws(covariate,clones,n_samples,use_logits,temperature,gate_prob,...) vs joint_distribution(covariate,groupby,n_samples,use_logits,clones,temperature,...). Real order/set divergence; _joint_draws is PRIVATE so lower stakes. Fix's canonical order (temperature before clones) is consistent with C2 \u2014 resolve C2 and C12 together."}, {"finding": "C13 joint_distribution_ppc arg-order (temperature after distance_metric) + missing groupby/n_samples/random_state", "status": "confirmed", "note": "Line 729: (covariate, distance_metric, temperature). Reorder to temperature-before-distance_metric per canonical block. It is model-free/deterministic, so n_samples/random_state only if draw-based comparison is added \u2014 the fix hedges correctly. Overlaps S15."}, {"finding": "C14 positional-should-be-kwonly: resolve_palette(adata,columns), _metric_boxplot(adata,function), reconstruction_ppc(model,adata)", "status": "confirmed", "note": "All three have a 2nd positional (lines 702,694,731). resolve_palette is public (clearest violation); the other two are private. Fix correct. Census incomplete: _phenotype_mass_per_clone(adata,covariate,clones,normalize) line 688 and prepare_two_level_params also carry extra positionals (private, minor) \u2014 see missed."}, {"finding": "C15 resample-count-name: n_permutations=1000 vs n_perm=200_000", "status": "confirmed", "note": "Line 732 permutation_null n_permutations vs line 266 n_perm. Same concept; n_boot correctly excluded as distinct. Rename to n_perm. Correct. Duplicate-adjacent to C1/S10 (both about diagnostics arg naming)."}, {"finding": "C16 show-arg only on private _finish", "status": "confirmed", "note": "Line 695 _finish(show=None); no public plot exposes show (verified \u00a78.1-8.7 have save/return_* but no show). show on _finish is thus unreachable from the public API. Fix (thread show through all pl, or drop it) correct."}, {"finding": "C17 return-control-name: return_df vs return_axes", "status": "confirmed", "note": "return_df on metric plots, return_axes on figure plots (lines 632,679,715). Real split. The finding itself concludes the split is acceptable (payloads differ) and only asks to standardize/document which each exposes \u2014 so this is near-self-resolving, effectively lowest priority."}, {"finding": "C18 eps-default: 1e-8 / 1e-6 / 1e-12 / 1e-15 across the surface", "status": "confirmed", "note": "Confirmed all four (predict/_compute_logits 1e-8 lines 317,382; prepare_two_level_params 1e-6 line 330; kl/js 1e-12 lines 276-277; _mi_from_joint 1e-15 line 528). Fix correct. STRENGTHEN: the doc self-contradicts \u2014 \u00a73.4 line 274 and appendix row 12 line 793 both assert 'one \u03b5=1e-12 library-wide, matching entropy/MI', yet the MI kernel uses 1e-15. That internal contradiction (not just drift) is the sharper issue \u2014 see missed."}, {"finding": "C19 covariate=None semantics tension in adata_or_jd fast path", "status": "confirmed", "note": "Real internal tension: signatures default covariate=None but \u00a77.2(c) line 478 says the adata path treats 'covariate required'. Fix (covariate=None ALWAYS = all covariate values; jd path merely filters) correct. Duplicate of S14."}, {"finding": "S1 probability_ternary stripped of splitby/conditions/covariate/n_samples/temperature/scale_function/top_n/color/random_state", "status": "confirmed", "note": "Strongly verified against notebooks (24 calls, doc header \u00a78.7 says 24 not the finding's 19): every yost/zhang call passes splitby=, conditions=, scale_function=, color=, and phenotype_names positionally. \u00a78.7 exposes none of these. REVISE priority: the knobs ACTUALLY exercised are splitby/conditions/scale_function/color/phenotype_names(positional); covariate/n_samples/temperature/top_n exist in the historical DEFAULT sig but are not passed in any notebook, so their absence is lower-priority. Also: 'rename groupby->splitby' is imprecise \u2014 legacy splitby here is the covariate COLUMN + conditions=values (see C4/missed)."}, {"finding": "S2 order= dropped from the three pl metric box plots", "status": "confirmed", "note": "Verified: order= passed on clonotypic_entropy (preprocess_conga, zhang) and phenotypic_entropy (yost, zhang); _metric_boxplot supports order (line 694). Duplicate of C9 \u2014 merge."}, {"finding": "S3 weighted= removed from joint_distribution/MI/ternary", "status": "refuted", "note": "Not a sufficiency gap: \u00a70.8 (lines 107-111) DELIBERATELY drops weighted (uniform-clonotype prior) with an explicit changelog/behavior-change note, echoed in the plan (line 300). The finding's own ask ('confirm intentional removal OR document') is already satisfied by the contract. It is a documented behavior change, not an oversight."}, {"finding": "S4 no call-time covariate COLUMN selector (covariate_key)", "status": "revised", "note": "Real capability gap CONFIRMED from notebooks: pl.flux uses key='treatment'/'timepoint_relapse', phenotypic_flux/ternary use splitby= across treatment/timepoint/timepoint_relapse \u2014 multiple columns per adata. BUT feasibility caveat the finding misses: the new engine's joint is built from p_ct indexed by the TRAINED covariate (ct_to_cov), so a covariate_key override is only implementable on an EMPIRICAL per-cell-probability aggregation path, not the p_ct engine path. The contract must decide empirical-vs-engine for multi-column flux/ternary before adding covariate_key. See missed item."}, {"finding": "S5 pl.phenotypic_flux lacks n_samples= and phenotype_subset=", "status": "confirmed", "note": "Verified: \u00a78.3 has random_state but no n_samples and no phenotype_subset; yost_v2 calls phenotypic_flux with n_samples and phenotype_subset=phenos. random_state without n_samples is a dead knob (ties to C3). Fix correct."}, {"finding": "S6 minimum_clone_size= dropped from pl metric plots", "status": "confirmed", "note": "Verified: comparison.ipynb calls pl.phenotypic_entropy(...minimum_clone_size=3), pl.mutual_information(...minimum_clone_size=1). Absent from \u00a78.1-8.2. Reproducible via clones=/slicing but the recorded calls won't run. Add as kwonly. Correct."}, {"finding": "S7 compare_groups paired mode lacks a pairing key / second axis (delta_entropy_table)", "status": "confirmed", "note": "\u00a77.6 paired=True aligns by sample_id (posterior-draw pairing), which is NOT the pre/post-within-unit pairing delta_entropy_table did, and there is no second axis for the cohort split. The plan (line 118) and API \u00a711 CLAIM delta_entropy_table is 'expressed via groupby + compare_groups' \u2014 this finding correctly shows that subsumption is incomplete for the paired-delta-then-cohort-contrast. Genuine over-claim. Fix (pair_on= / second axis) correct."}, {"finding": "S8 pp.gene_entropy unmapped (batch_key lost)", "status": "refuted", "note": "Not unmapped: relocated to examples/ (API line 173/773; plan line 127 'relocate to examples/, NOT a 0-caller'). Moved wholesale, so batch_key/agg_function travel with it (smith's gene_entropy(batch_key='sample') verified). Framing 'no home in extracted surface' is inaccurate. Only debatable point is examples/ vs pp \u2014 a placement judgment, not a lost capability."}, {"finding": "S9 pl.flux distance box-plot (paint=) has no public home", "status": "revised", "note": "Verified the plot exists and is used (renal/yost/zhang: pl.flux with key/order/groupby/paint=/distance_metric='dkl'). But it is DISPOSITIONED, not overlooked: plan line 129 'pl.flux boxplot (5 uses) DROP \u2014 sankey is the flux plot' + changelog line 300. So the finding's fix (add a box variant) CONTRADICTS the settled decision to drop it. Real residual issue: API doc \u00a711 never states this box-plot's disposition (only renames flux->tl.phenotypic_flux), a doc-completeness gap."}, {"finding": "S10 seed vs random_state naming in diagnostics/_stats", "status": "confirmed", "note": "Same as C1. Load-bearing renames are the public diag.reconstruction_ppc/permutation_null; the _stats primitives are private. Merge with C1/C15."}, {"finding": "S11 permutation_null can't thread metric params (temperature/normalized/normalize_mode/clones/splitby)", "status": "confirmed", "note": "\u00a79.1 line 732 signature has no pass-through, so the null uses metric defaults that may not match the user's reported statistic. Real. Fix (add pass-throughs or metric_kwargs, rename seed) correct."}, {"finding": "S12 base=/decimals= dropped from entropy metrics", "status": "refuted", "note": "base= removal is INTENTIONAL: \u00a77.2/7.3 'Base fixed to 2' and \u00a73.4/appendix row 12 standardize on log2/bits library-wide as an explicit unit-consistency fix. Re-adding base= would undermine that fix. decimals= is cosmetic (drop is fine). The finding's proposed base=2 arg works against the contract's stated goal."}, {"finding": "S13 hue_order absent on pl.phenotypic_entropy/mutual_information", "status": "confirmed", "note": "Verified (\u00a78.1-8.2 lack hue_order; only clonotypic_entropy has it, line 631). Duplicate of C9 \u2014 merge."}, {"finding": "S14 adata_or_jd covariate=None dual meaning", "status": "confirmed", "note": "Same tension as C19 (line 478). Fix (jd path ignores covariate/n_samples/etc.; raise if set; covariate=None='all' only on adata path) correct. Merge with C19."}, {"finding": "S15 diag PPCs omit clones=/groupby=", "status": "confirmed", "note": "\u00a79.1: joint_distribution_ppc omits groupby/n_samples/random_state/clones (line 729); permutation_null omits clones (line 732). Real. Overlaps C13. Add clones=/groupby= for per-subset PPCs."}, {"finding": "S16 auxiliaries absent: top_clone_umap, polar_plot, probability_distribution, compare_phenotypes, clonality, clone_fraction", "status": "revised", "note": "Every one IS dispositioned in the implementation plan \u00a73: clonality->DROP (line 114), compare_phenotypes->DROP (117), gene handled, top_clone_umap->examples (301), polar_plot->DROP (128), probability_distribution->delete (272), clone_fraction->delete (272). So 'no new home' is inaccurate, and the fix 'port top_clone_umap and polar_plot' CONTRADICTS the plan (polar_plot is explicitly DROP-with-rewrite). Genuine residual: the API-doc \u00a711 census mislabels clonality/probability_distribution as '0 live callers' and omits compare_phenotypes entirely \u2014 but that is the MISSED cross-doc contradiction below, mischaracterized/underrated (LOW) here. clonality is CALLED (tcri.pl.clonality) with rich args, not merely 'imported'."}], "missed": [{"category": "cross-doc-contradiction / deletion-census", "issue": "The two contract docs disagree on which symbols have live callers. API doc \u00a711 line 770 buckets clonality (tl+pl), probability_distribution, and clone_fraction under 'Deleted (dead/broken/out-of-scope, 0 live callers after census)'. But (a) the implementation plan explicitly says the opposite \u2014 clonality and compare_phenotypes HAVE live notebook callers and must be removed WITH their replacement so no caller is orphaned (plan lines 114, 263, 273, 294, 369); and (b) the notebooks confirm live PACKAGE calls: tcri.pl.clonality is called with full args (groupby/splitby/order/palette/figsize) in smith/renal/comparison/yost/zhang, and tcri.metrics.probability_distribution is imported and called 4x via the alias pdistribution(psubset, method=method). Additionally compare_phenotypes (1 live call, tcri.pl.compare_phenotypes in synthetic1) is dispositioned in the plan (DROP, line 117) but is ABSENT from the API doc \u00a711 surface-delta lists entirely. Both audits under-caught this: the consistency audit only checked arg names/order; the sufficiency audit lumped these into one LOW 'auxiliaries' finding, mislabeled clonality as merely 'imported', and wrongly framed them as having 'no home'.", "fix": "Reconcile API doc \u00a711 with plan \u00a73. Move clonality, probability_distribution (the tl/metrics one that is actually called), and compare_phenotypes OUT of the '0 live callers' Deleted bucket into a 'removed WITH replacement + in-PR notebook rewrite (Phase 6/10)' category, matching the plan. Add compare_phenotypes to \u00a711 (currently missing). Correct the '0 live callers after census' label since it is factually false for these symbols; the census clearly missed alias-imported and package-qualified call-sites."}, {"category": "overloaded-concept / name-collision", "issue": "Legacy `splitby` is overloaded and neither audit distinguished the two meanings. On the box plots splitby = cohort hue (correctly mapped to the new splitby). But on probability_ternary and phenotypic_flux(Sankey) the legacy `splitby` names the covariate COLUMN, paired with conditions=/order= that select the VALUES (verified: probability_ternary(...splitby=\"timepoint_relapse\", conditions=[\"CR_Pre-treatment\",...]); pl.phenotypic_flux(...splitby=\"treatment\", order=[\"pre\",\"post\"])). Consistency C4 proposes adding a box-plot-semantics `splitby` to these two (wrong role -> reintroduces the collision the RULES forbid), while sufficiency S1 proposes renaming ternary's groupby->splitby (also conflates column-vs-hue). Both are off.", "fix": "For probability_ternary and phenotypic_flux, add covariate_key=None (the covariate COLUMN) + conditions=None (ternary: 1-2 values -> start/end simplices) / rely on order= (Sankey values); RESERVE `splitby` strictly for the cohort hue. Document that the legacy `splitby=` notebook usage maps to covariate_key, not splitby, in the Phase-10 rewrite. This also subsumes S4 (the covariate-column selector) \u2014 but note the feasibility caveat: the p_ct engine is covariate-locked (ct_to_cov), so multi-column flux/ternary must be computed on the empirical per-cell-probability path, not the p_ct joint engine; the contract does not yet state which path these plots use."}, {"category": "internal-contradiction", "issue": "The eps-clip invariant is self-contradictory in the doc (sharper than consistency C18's 'drift'). \u00a73.4 line 274 asserts 'one base (bits, log2) and one \u03b5=10^-12 library-wide, matching entropy/MI', and appendix row 12 line 793 repeats 'one \u03b5=10^-12'. But the MI kernel _mi_from_joint in \u00a77.4 line 528 hardcodes \u03b5=10^-15. So the document violates its own explicitly stated library-wide invariant for MI, the exact metric \u00a73.4 name-checks.", "fix": "Set the MI kernel _mi_from_joint eps to 1e-12 to satisfy the stated invariant, OR amend the \u00a73.4/appendix 'library-wide 1e-12 matching entropy/MI' claim to carve out the MI-specific 1e-15 with a numerical justification. Do not leave the prose and the kernel contradicting each other."}, {"category": "undispositioned-rename", "issue": "Two legacy parameters used in the notebooks are absent from both docs' rename/disposition maps, leaving the Phase-10 notebook rewrite underspecified. (1) `method=` (e.g. method=\"probabilistic\", also method=method) is passed to pl.clonotypic_entropy, pl.mutual_information, pl.flux, polar_plot, and pdistribution, but only posterior=/point_estimate= are mapped (plan line 112) \u2014 method= itself is never named. (2) `phenotype_names=` is the legacy positional/keyword on probability_ternary (verified in every zhang/smith call) but the kept \u00a78.7 signature renames it to `phenotypes` with no rename entry in \u00a711 or plan \u00a73.", "fix": "Add explicit disposition-map rows: `method=` (old posterior/deterministic toggle) -> removed, expressed via n_samples/use_logits; and `phenotype_names` -> `phenotypes` on pl.probability_ternary. Without these the notebook-rewrite mapping (plan Phase 10) is ambiguous for calls that pass method= or phenotype_names=."}, {"category": "caller-count-reliability", "issue": "The \u00a711/plan caller-census counts are unreliable in at least one more place, undermining confidence in the deletion-safety gate. polar_plot is recorded as '1 notebook use' (API line 773, plan line 128), but the notebooks contain 4 local `def polar_plot` redefinitions plus multiple call-sites across smith/comparison/yost/zhang; the '1 caller' figure does not reflect the actual notebook footprint (whether package-call or local-def, the rewrite scope is larger than 1). Same pattern (local def + call) holds for clonality. Neither audit checked def-vs-call disambiguation.", "fix": "Re-run the caller census distinguishing (a) package-qualified calls, (b) alias-imported calls, and (c) notebook-local redefinitions, and record accurate per-symbol counts in the disposition map. The presence of local `def clonality`/`def polar_plot` in notebooks means the Phase-10 rewrite must also strip the shadow definitions, not just swap call-sites."}], "summary": "Verified all 19 consistency + 16 sufficiency findings against the two contract docs and the six example notebooks (read directly). Most consistency findings are REAL and correctly fixed: arg-order clones/temperature (C2), reconstruction_ppc n_samples=100 collision (C5), metric vs distance_metric (C6), normalize/normalized + phenotype_colors (C7), figsize drift (C8), order/hue_order (C9), _joint_draws mismatch (C12), joint_distribution_ppc order/gaps (C13), 2nd-positional violations (C14), n_permutations vs n_perm (C15), show= plumbing (C16), eps drift (C18), covariate=None tension (C19), and seed-vs-random_state (C1). Revisions: C3 (flux does NOT expose n_samples, so its dangling knob is random_state), C4 (the 'add splitby' fix mis-names a covariate-COLUMN role as cohort-hue), C11 (group_singletons 'trb'/'patient' defaults are defensible pre-registration values, not a bug). Strongest sufficiency findings confirmed: probability_ternary stripped of splitby/conditions/scale_function/color (S1, verified across 24 notebook calls), order= (S2), pl.phenotypic_flux missing n_samples/phenotype_subset (S5), minimum_clone_size (S6), compare_groups incomplete subsumption of delta_entropy_table (S7), permutation_null can't thread metric params (S11). Refuted as gaps: weighted= (S3, deliberately dropped and documented), gene_entropy (S8, relocated to examples/), base= (S12, intentional single-base fix), and the S16 'auxiliaries' bundle (all dispositioned in the plan; framing 'no home / port them' contradicts the plan's DROP decisions). Numerous duplicates to merge: C1==S10 (seed naming), C9==S2==S13 (order/hue_order), C19==S14 (covariate=None dual meaning), C13~S15 (diag PPC gaps). The single most important thing BOTH audits missed: the two contract docs CONTRADICT each other \u2014 API doc \u00a711 lists clonality/probability_distribution/clone_fraction as '0 live callers' deletes, but the implementation plan (and the notebooks: 4+ tcri.pl.clonality calls, 4 aliased probability_distribution calls, 1 compare_phenotypes call) shows they have live callers and must be dropped WITH a replacement; compare_phenotypes is absent from \u00a711 entirely. Other missed items: legacy `splitby` is overloaded (cohort-hue vs covariate-column) which mis-guides both C4 and S1; the doc self-contradicts on eps (\u00a73.4/appendix say 1e-12 'library-wide' but the MI kernel uses 1e-15); and method=/phenotype_names renames are undispositioned."}, "sigs": {"functions": [{"name": "info", "module": "tcri/_console.py", "namespace": "shared", "responsibility": "scanpy.logging.info wrapper; silenced by scanpy verbosity.", "args": [{"name": "msg"}, {"name": "deep", "default": "None", "kwonly": true}]}, {"name": "warning", "module": "tcri/_console.py", "namespace": "shared", "responsibility": "scanpy.logging.warning wrapper.", "args": [{"name": "msg"}]}, {"name": "success", "module": "tcri/_console.py", "namespace": "shared", "responsibility": "scanpy.logging.hint wrapper.", "args": [{"name": "msg"}]}, {"name": "done", "module": "tcri/_console.py", "namespace": "shared", "responsibility": "Terminal completion line.", "args": [{"name": "msg", "default": "\"done\""}]}, {"name": "stars", "module": "tcri/_stats.py", "namespace": "shared", "responsibility": "Map p-value to star string (****,***,**,*,ns) by fixed thresholds.", "args": [{"name": "p"}]}, {"name": "mann_whitney", "module": "tcri/_stats.py", "namespace": "shared", "responsibility": "Mann-Whitney U with two-sided p from rank-sum null; returns (U, p).", "args": [{"name": "a"}, {"name": "b"}, {"name": "alternative", "default": "\"two-sided\"", "kwonly": true}]}, {"name": "prob_direction", "module": "tcri/_stats.py", "namespace": "shared", "responsibility": "Signed-contrast only: from paired diff-draw vector return (p_gt, p_lt).", "args": [{"name": "delta"}]}, {"name": "hdi", "module": "tcri/_stats.py", "namespace": "shared", "responsibility": "True highest-density interval (narrowest window of ceil(hdi_prob*N) points); returns (low, high).", "args": [{"name": "samples"}, {"name": "hdi_prob", "default": "0.94", "kwonly": true}]}, {"name": "summarize", "module": "tcri/_stats.py", "namespace": "shared", "responsibility": "Reduce a raw metric draw vector to {mean, sd, hdi_low, hdi_high}; no p_gt.", "args": [{"name": "samples"}, {"name": "hdi_prob", "default": "0.94", "kwonly": true}]}, {"name": "auc_and_label_permutation", "module": "tcri/_stats.py", "namespace": "shared", "responsibility": "Observed ROC-AUC + two-sided permutation p (exact when C(n,k)<=max_exact else MC); returns (auc, p, perm_stats, mode).", "args": [{"name": "scores"}, {"name": "labels"}, {"name": "pos_label", "default": "None", "kwonly": true}, {"name": "n_perm", "default": "200_000", "kwonly": true}, {"name": "seed", "default": "42", "kwonly": true}, {"name": "max_exact", "default": "200_000", "kwonly": true}]}, {"name": "bootstrap_auc", "module": "tcri/_stats.py", "namespace": "shared", "responsibility": "Bootstrap AUROC (reject draws missing a class); returns np.array([lo, hi]) 2.5/97.5 quantiles.", "args": [{"name": "scores"}, {"name": "labels"}, {"name": "pos_label", "default": "None", "kwonly": true}, {"name": "n_boot", "default": "5000", "kwonly": true}, {"name": "seed", "default": "42", "kwonly": true}]}, {"name": "l1_distance", "module": "tcri/_distance.py", "namespace": "shared", "responsibility": "Sum|p_i-q_i|; symmetric, range [0,2]; renormalizes inputs; returns float.", "args": [{"name": "p"}, {"name": "q"}]}, {"name": "kl_divergence", "module": "tcri/_distance.py", "namespace": "shared", "responsibility": "KL(p||q) in bits (log2); clip to [eps,1] then renormalize; asymmetric [0,inf); returns float.", "args": [{"name": "p"}, {"name": "q"}, {"name": "base", "default": "2.0", "kwonly": true}, {"name": "eps", "default": "1e-12", "kwonly": true}]}, {"name": "js_divergence", "module": "tcri/_distance.py", "namespace": "shared", "responsibility": "NEW Jensen-Shannon = 1/2 KL(p||r)+1/2 KL(q||r), r=1/2(p+q); symmetric, bounded [0,1] bit; returns float.", "args": [{"name": "p"}, {"name": "q"}, {"name": "base", "default": "2.0", "kwonly": true}, {"name": "eps", "default": "1e-12", "kwonly": true}]}, {"name": "phenotype_distance", "module": "tcri/_distance.py", "namespace": "shared", "responsibility": "Dispatcher: 'l1'->l1_distance; 'kl'/'dkl'->kl_divergence; 'jsd'->js_divergence; callable; else ValueError.", "args": [{"name": "p"}, {"name": "q"}, {"name": "metric", "default": "\"l1\"", "kwonly": true}]}, {"name": "resolve_device", "module": "tcri/_compute/_xp.py", "namespace": "shared", "responsibility": "Device ladder: None/'cpu'->cpu; 'mps'->cpu; 'cuda'/'gpu'/'auto'->GPU iff backend imports and getDeviceCount()>0 else CPU (explicit 'cuda' warns).", "args": [{"name": "device"}]}, {"name": "get_xp", "module": "tcri/_compute/_xp.py", "namespace": "shared", "responsibility": "Return the array module: torch(-cuda) preferred, cupy optional, numpy default; GPU libs imported lazily inside.", "args": [{"name": "device"}]}, {"name": "asnumpy", "module": "tcri/_compute/_xp.py", "namespace": "shared", "responsibility": "Host-boundary shim: cupy.asnumpy(x)/x.cpu().numpy()/np.asarray(x).", "args": [{"name": "x"}]}, {"name": "_joint_draws", "module": "tcri/_compute/_joint.py", "namespace": "shared", "responsibility": "Return [max(n_samples,1), n_clones, P] joint stack via batched Dirichlet draw + gated softmax + scatter-add reduction; validates on-device; returns np.ndarray. (Impl-plan \u00a72 lists args as (adata, covariate, n_samples, *, use_logits, temperature, device, random_state) \u2014 differs from the API doc arg order/set below.)", "args": [{"name": "adata"}, {"name": "covariate", "kwonly": true}, {"name": "clones", "kwonly": true}, {"name": "n_samples", "kwonly": true}, {"name": "use_logits", "kwonly": true}, {"name": "temperature", "kwonly": true}, {"name": "gate_prob", "kwonly": true}, {"name": "random_state", "kwonly": true}, {"name": "device", "kwonly": true}]}, {"name": "entropy", "module": "tcri/_compute/_reduce.py", "namespace": "shared", "responsibility": "Batched xlogx entropy reduction over the [S,n_clones,P] stack (no explicit signature given in doc).", "args": []}, {"name": "mutual_information", "module": "tcri/_compute/_reduce.py", "namespace": "shared", "responsibility": "Batched MI (joint vs outer-product) reduction over the stack (no explicit signature given in doc).", "args": []}, {"name": "distance", "module": "tcri/_compute/_reduce.py", "namespace": "shared", "responsibility": "Batched distance reduction over the stack plus summarize/hdi over the sample axis (no explicit signature given in doc).", "args": []}, {"name": "umap", "module": "tcri/_compute/_embedding.py", "namespace": "shared", "responsibility": "Latent UMAP behind _use_gpu gate (cuML on GPU, umap-learn CPU), lazy import (impl-plan \u00a72 only; _embedding.py absent from API-doc \u00a71 package tree \u2014 inconsistency).", "args": []}, {"name": "_use_gpu", "module": "tcri/_compute/_embedding.py", "namespace": "shared", "responsibility": "Gate deciding GPU vs CPU embedding path (API-doc \u00a74.3; no explicit signature \u2014 argument inferred as device).", "args": [{"name": "device"}]}, {"name": "TCRIModel.setup_anndata", "module": "tcri/model/_model.py", "namespace": "ml", "responsibility": "Registration only: register clonotype/phenotype/covariate/batch/count fields with scvi, store layer, write+register obs['indices']; removes uns['tcri_manager'] stash; no analysis/label obs mutation.", "args": [{"name": "cls"}, {"name": "adata"}, {"name": "layer", "default": "None", "kwonly": true}, {"name": "clonotype_key", "default": "\"unique_clone_id\"", "kwonly": true}, {"name": "phenotype_key", "default": "\"phenotype_col\"", "kwonly": true}, {"name": "covariate_key", "default": "\"timepoint\"", "kwonly": true}, {"name": "batch_key", "default": "\"patient\"", "kwonly": true}, {"name": "**kwargs"}]}, {"name": "TCRIModel.__init__", "module": "tcri/model/_model.py", "namespace": "ml", "responsibility": "Build empirical clone->phenotype prior + KMeans archetypes + index maps + class weights, then construct/prime TCRIModule; gate_prob=None => ungated.", "args": [{"name": "self"}, {"name": "adata"}, {"name": "n_latent", "default": "128", "kwonly": true}, {"name": "n_hidden", "default": "128", "kwonly": true}, {"name": "n_layers", "default": "3", "kwonly": true}, {"name": "classifier_n_layers", "default": "3", "kwonly": true}, {"name": "global_scale", "default": "5.0", "kwonly": true}, {"name": "local_scale", "default": "3.0", "kwonly": true}, {"name": "prior_temperature", "default": "1.0", "kwonly": true}, {"name": "guide_temperature", "default": "1.0", "kwonly": true}, {"name": "use_enumeration", "default": "False", "kwonly": true}, {"name": "patience", "default": "300", "kwonly": true}, {"name": "classifier_hidden", "default": "128", "kwonly": true}, {"name": "classifier_dropout", "default": "0.1", "kwonly": true}, {"name": "n_pseudo_obs", "default": "10", "kwonly": true}, {"name": "K", "default": "10", "kwonly": true}, {"name": "phenotype_weights", "default": "None", "kwonly": true}, {"name": "gate_prob", "default": "None", "kwonly": true}, {"name": "kl_weight_max", "default": "1.0", "kwonly": true}, {"name": "guide_init_scale", "default": "10.0", "kwonly": true}, {"name": "classifier_temperature", "default": "1.0", "kwonly": true}, {"name": "**kwargs"}]}, {"name": "TCRIModel.train", "module": "tcri/model/_model.py", "namespace": "ml", "responsibility": "0.9/0.1 split, UnifiedTrainingPlan, TrainRunner with elbo_validation early stopping.", "args": [{"name": "self"}, {"name": "max_epochs", "default": "1000", "kwonly": true}, {"name": "batch_size", "default": "1000", "kwonly": true}, {"name": "lr", "default": "1e-3", "kwonly": true}, {"name": "reconstruction_loss_scale", "default": "1e-3", "kwonly": true}, {"name": "n_steps_kl_warmup", "default": "2000", "kwonly": true}, {"name": "**kwargs"}]}, {"name": "TCRIModel.get_latent_representation", "module": "tcri/model/_model.py", "namespace": "ml", "responsibility": "Batched encode to (n_cells, n_latent) posterior-mean latent; returns np.ndarray.", "args": [{"name": "self"}, {"name": "adata", "default": "None"}, {"name": "indices", "default": "None", "kwonly": true}, {"name": "batch_size", "default": "None", "kwonly": true}]}, {"name": "TCRIModel.predict", "module": "tcri/model/_model.py", "namespace": "ml", "responsibility": "(renamed from get_cell_phenotype_probs) Per-cell phenotype-probability DataFrame combining classifier logits with log p_ct (gate/additive); order-preserving loader; returns pd.DataFrame.", "args": [{"name": "self"}, {"name": "adata", "default": "None"}, {"name": "batch_size", "default": "256", "kwonly": true}, {"name": "eps", "default": "1e-8", "kwonly": true}]}, {"name": "TCRIModel.get_p_ct", "module": "tcri/model/_model.py", "namespace": "ml", "responsibility": "Return learned (ct_count, P) posterior mean m=normalize(q_p_ct_raw); at guide_temperature=1.0 equals uns[K.P_CT]; returns np.ndarray.", "args": [{"name": "self"}, {"name": "guide_temperature", "default": "1.0", "kwonly": true}]}, {"name": "TCRIModel.to_anndata", "module": "tcri/model/_model.py", "namespace": "ml", "responsibility": "(replaces register_model) Thin writer of the canonical minimum: metadata+categories, X_tcri latent, obsm logits, predict() probs+argmax labels, p_ct (+ct_to_cov/ct_to_c/per-cell arrays), local_scale, gate_prob, classifier_temperature; returns AnnData.", "args": [{"name": "self"}, {"name": "adata", "default": "None"}, {"name": "latent_key", "default": "\"X_tcri\"", "kwonly": true}, {"name": "logits_key", "default": "\"X_tcri_logits\"", "kwonly": true}, {"name": "predictions_key", "default": "\"X_tcri_probabilities\"", "kwonly": true}, {"name": "label_key", "default": "\"tcri_phenotype\"", "kwonly": true}]}, {"name": "TCRIModule.__init__", "module": "tcri/model/_module.py", "namespace": "ml", "responsibility": "Construct encoder/decoder/classifier/VampPrior, px_r; register empty two-level buffers + class weights (mixture_concentration required). Internal.", "args": [{"name": "self"}, {"name": "n_input"}, {"name": "n_latent"}, {"name": "P"}, {"name": "n_batch"}, {"name": "global_scale", "default": "10.0", "kwonly": true}, {"name": "local_scale", "default": "5.0", "kwonly": true}, {"name": "prior_temperature", "default": "1.0", "kwonly": true}, {"name": "guide_temperature", "default": "1.0", "kwonly": true}, {"name": "gate_prob", "default": "0.5", "kwonly": true}, {"name": "mixture_concentration", "default": "None", "kwonly": true}, {"name": "n_pseudo_obs", "default": "10", "kwonly": true}, {"name": "use_enumeration", "default": "False", "kwonly": true}, {"name": "classifier_hidden", "default": "128", "kwonly": true}, {"name": "classifier_dropout", "default": "0.1", "kwonly": true}, {"name": "classifier_n_layers", "default": "3", "kwonly": true}, {"name": "n_hidden", "default": "128", "kwonly": true}, {"name": "n_layers", "default": "3", "kwonly": true}, {"name": "class_weights", "default": "None", "kwonly": true}, {"name": "kl_weight_max", "default": "1.0", "kwonly": true}, {"name": "guide_init_scale", "default": "10.0", "kwonly": true}, {"name": "classifier_temperature", "default": "1.0", "kwonly": true}]}, {"name": "TCRIModule.prepare_two_level_params", "module": "tcri/model/_module.py", "namespace": "ml", "responsibility": "Normalize/temperature the clone-phenotype prior; register two-level index buffers. Internal.", "args": [{"name": "self"}, {"name": "clone_phen_prior_mat"}, {"name": "ct_to_c"}, {"name": "ct_to_cov"}, {"name": "ct_array_for_cells"}, {"name": "cov_array_for_cells"}, {"name": "eps", "default": "1e-6", "kwonly": true}]}, {"name": "TCRIModule.model", "module": "tcri/model/_module.py", "namespace": "ml", "responsibility": "Generative: sample p_c (MixtureDirichlet), p_ct (Dirichlet centered at p_c), latent z (VampPrior), ZINB gene obs. Internal. Signature shows trailing unspecified args.", "args": [{"name": "self"}, {"name": "x"}, {"name": "batch_idx"}, {"name": "log_library"}, {"name": "..."}]}, {"name": "TCRIModule.guide", "module": "tcri/model/_module.py", "namespace": "ml", "responsibility": "Guide: learnable Dirichlet q(p_c), q(p_ct) with clamp(min=1e-3) on concentration; Normal q(z) from encoder. Internal. Signature shows trailing unspecified args.", "args": [{"name": "self"}, {"name": "x"}, {"name": "batch_idx"}, {"name": "log_library"}, {"name": "..."}]}, {"name": "TCRIModule.get_latent", "module": "tcri/model/_module.py", "namespace": "ml", "responsibility": "Encode a batch to posterior-mean latent z_loc; returns torch.Tensor. Internal.", "args": [{"name": "self"}, {"name": "tensor_dict"}]}, {"name": "TCRIModule.get_p_ct", "module": "tcri/model/_module.py", "namespace": "ml", "responsibility": "Read q_p_ct_raw from process-global param store; return row-normalized (ct_count, P) posterior mean; returns torch.Tensor. Internal.", "args": [{"name": "self"}]}, {"name": "TCRIModule.use_gate", "module": "tcri/model/_module.py", "namespace": "ml", "responsibility": "Property: True when gate_prob is not None. Internal.", "args": [{"name": "self"}]}, {"name": "TCRIModule._get_fn_args_from_batch", "module": "tcri/model/_module.py", "namespace": "ml", "responsibility": "Staticmethod: extract (x, batch_idx, log_library) from a scvi batch dict; returns tuple. Internal.", "args": [{"name": "tensor_dict"}]}, {"name": "MixtureDirichlet.__init__", "module": "tcri/model/_priors.py", "namespace": "ml", "responsibility": "Clonotype prior p_c distribution. Internal.", "args": [{"name": "self"}, {"name": "mixture_weights"}, {"name": "concentration"}, {"name": "validate_args", "default": "None"}]}, {"name": "MixtureDirichlet.sample", "module": "tcri/model/_priors.py", "namespace": "ml", "responsibility": "Draw a sample from the mixture. Internal (no explicit arg list in doc).", "args": [{"name": "self"}]}, {"name": "MixtureDirichlet.log_prob", "module": "tcri/model/_priors.py", "namespace": "ml", "responsibility": "Log-sum-exp of component Dirichlet log-probs. Internal (no explicit arg list in doc).", "args": [{"name": "self"}]}, {"name": "MixtureDirichlet.score_parts", "module": "tcri/model/_priors.py", "namespace": "ml", "responsibility": "Returns (log_prob, 0, 0) -> reparam-free. Internal (no explicit arg list in doc).", "args": [{"name": "self"}]}, {"name": "MixtureDirichlet.__call__", "module": "tcri/model/_priors.py", "namespace": "ml", "responsibility": "Alias for sample. Internal (no explicit arg list in doc).", "args": [{"name": "self"}]}, {"name": "VampPrior.__init__", "module": "tcri/model/_priors.py", "namespace": "ml", "responsibility": "VampPrior over z. Internal.", "args": [{"name": "self"}, {"name": "pseudo_inputs"}, {"name": "encoder"}]}, {"name": "VampPrior.get_mixture", "module": "tcri/model/_priors.py", "namespace": "ml", "responsibility": "Uniform MixtureSameFamily of Independent Normals. Internal (no explicit arg list in doc).", "args": [{"name": "self"}]}, {"name": "VampPrior.log_prob", "module": "tcri/model/_priors.py", "namespace": "ml", "responsibility": "Log density of z under the VampPrior. Internal.", "args": [{"name": "self"}, {"name": "z"}]}, {"name": "VampPrior.sample", "module": "tcri/model/_priors.py", "namespace": "ml", "responsibility": "Draw z from the VampPrior. Internal (no explicit arg list in doc).", "args": [{"name": "self"}]}, {"name": "PhenotypeClassifier.__init__", "module": "tcri/model/_classifier.py", "namespace": "ml", "responsibility": "MLP phenotype classifier. Internal.", "args": [{"name": "self"}, {"name": "n_latent"}, {"name": "classifier_hidden"}, {"name": "P"}, {"name": "num_layers", "default": "3", "kwonly": true}, {"name": "dropout_rate", "default": "0.1", "kwonly": true}, {"name": "temperature", "default": "1.0", "kwonly": true}]}, {"name": "PhenotypeClassifier.forward", "module": "tcri/model/_classifier.py", "namespace": "ml", "responsibility": "Return MLP logits divided by temperature (tau_cls). Internal.", "args": [{"name": "self"}, {"name": "x"}]}, {"name": "UnifiedTrainingPlan.__init__", "module": "tcri/model/_training.py", "namespace": "ml", "responsibility": "Pyro training plan. Internal.", "args": [{"name": "self"}, {"name": "module"}, {"name": "n_steps_kl_warmup", "default": "1000", "kwonly": true}, {"name": "reconstruction_loss_scale", "default": "1e-2", "kwonly": true}, {"name": "num_particles", "default": "5", "kwonly": true}, {"name": "optimizer_config", "default": "None", "kwonly": true}, {"name": "class_weights", "default": "None", "kwonly": true}, {"name": "**kwargs"}]}, {"name": "UnifiedTrainingPlan.loss", "module": "tcri/model/_training.py", "namespace": "ml", "responsibility": "Loss property. Internal (no explicit arg list in doc).", "args": [{"name": "self"}]}, {"name": "UnifiedTrainingPlan.configure_optimizers", "module": "tcri/model/_training.py", "namespace": "ml", "responsibility": "Configure optimizers. Internal (no explicit arg list in doc).", "args": [{"name": "self"}]}, {"name": "UnifiedTrainingPlan.training_step", "module": "tcri/model/_training.py", "namespace": "ml", "responsibility": "One training step. Internal (no explicit arg list in doc).", "args": [{"name": "self"}]}, {"name": "UnifiedTrainingPlan.validation_step", "module": "tcri/model/_training.py", "namespace": "ml", "responsibility": "Validation step; logs elbo_validation. Internal (no explicit arg list in doc).", "args": [{"name": "self"}]}, {"name": "build_archetypes", "module": "tcri/model/_training.py", "namespace": "ml", "responsibility": "KMeans-cluster clone->phenotype rows into K normalized archetype centroids; returns (centers, labels). K default 10 (aligned to TCRIModel).", "args": [{"name": "c2p_mat"}, {"name": "K", "default": "10", "kwonly": true}]}, {"name": "group_singletons", "module": "tcri/preprocessing/_clones.py", "namespace": "pp", "responsibility": "Collapse clones smaller than min_clone_size (per groupby) into 'Singleton_{group}' labels in target_col; must run before setup_anndata; returns AnnData.", "args": [{"name": "adata"}, {"name": "clonotype_key", "default": "\"trb\"", "kwonly": true}, {"name": "groupby", "default": "\"patient\"", "kwonly": true}, {"name": "target_col", "default": "\"trb_unique\"", "kwonly": true}, {"name": "min_clone_size", "default": "10", "kwonly": true}]}, {"name": "clone_size", "module": "tcri/preprocessing/_clones.py", "namespace": "pp", "responsibility": "Per-clone cell counts written per cell into obs[key_added]; reads uns[K.METADATA]['clone_col'].", "args": [{"name": "adata"}, {"name": "key_added", "default": "\"clone_size\"", "kwonly": true}, {"name": "return_counts", "default": "False", "kwonly": true}]}, {"name": "collapse_singleton", "module": "tcri/preprocessing/_clones.py", "namespace": "pp", "responsibility": "Private inner helper of group_singletons (per-row relabel).", "args": [{"name": "row"}]}, {"name": "_write_metadata", "module": "tcri/preprocessing/_register.py", "namespace": "pp", "responsibility": "Write uns[K.METADATA] (single scheme) + the three category lists from the registry. Private to_anndata writer.", "args": [{"name": "adata"}, {"name": "model"}]}, {"name": "_register_clonotype_key", "module": "tcri/preprocessing/_register.py", "namespace": "pp", "responsibility": "Register the clonotype obs column + ordered categories. Private.", "args": [{"name": "adata"}, {"name": "clonotype_key"}, {"name": "order", "default": "None", "kwonly": true}]}, {"name": "_register_phenotype_key", "module": "tcri/preprocessing/_register.py", "namespace": "pp", "responsibility": "Register the phenotype obs column + ordered categories. Private.", "args": [{"name": "adata"}, {"name": "phenotype_key"}, {"name": "order", "default": "None", "kwonly": true}]}, {"name": "_write_latent", "module": "tcri/preprocessing/_register.py", "namespace": "pp", "responsibility": "Encoder posterior-mean latent -> obsm. Private.", "args": [{"name": "adata"}, {"name": "model"}, {"name": "latent_key", "default": "\"X_tcri\"", "kwonly": true}, {"name": "batch_size", "default": "256", "kwonly": true}]}, {"name": "_write_logits", "module": "tcri/preprocessing/_register.py", "namespace": "pp", "responsibility": "Per-cell classifier logits -> obsm[K.X_LOGITS] (canonical; required by default engine path). Private.", "args": [{"name": "adata"}, {"name": "model"}, {"name": "logits_key", "default": "\"X_tcri_logits\"", "kwonly": true}, {"name": "batch_size", "default": "256", "kwonly": true}]}, {"name": "_write_predictions", "module": "tcri/preprocessing/_register.py", "namespace": "pp", "responsibility": "predict() probs -> obsm; argmax hard labels -> obs. Private.", "args": [{"name": "adata"}, {"name": "model"}, {"name": "predictions_key", "default": "\"X_tcri_probabilities\"", "kwonly": true}, {"name": "label_key", "default": "\"tcri_phenotype\"", "kwonly": true}, {"name": "batch_size", "default": "256", "kwonly": true}]}, {"name": "_write_p_ct", "module": "tcri/preprocessing/_register.py", "namespace": "pp", "responsibility": "Write p_ct, ct_to_cov, ct_to_c, per-cell ct/cov arrays, local_scale, gate_prob, classifier_temperature -> uns. Private.", "args": [{"name": "adata"}, {"name": "model"}]}, {"name": "_compute_logits_and_prior", "module": "tcri/preprocessing/_register.py", "namespace": "pp", "responsibility": "Run encoder+classifier to extract per-cell logits and log p_ct from get_p_ct(); returns tuple[np.ndarray, np.ndarray]. Private.", "args": [{"name": "model"}, {"name": "adata"}, {"name": "batch_size", "default": "256", "kwonly": true}, {"name": "eps", "default": "1e-8", "kwonly": true}]}, {"name": "joint_distribution", "module": "tcri/tools/_joint.py", "namespace": "tl", "responsibility": "THE ENGINE: clone x phenotype distribution at a covariate from the learned posterior of p_ct (mean or draws); use_logits folds per-cell logits (gate-aware); re-exported as tcri.joint_distribution; returns pandas.DataFrame.", "args": [{"name": "adata"}, {"name": "covariate", "default": "None", "kwonly": true}, {"name": "groupby", "default": "None", "kwonly": true}, {"name": "n_samples", "default": "0", "kwonly": true}, {"name": "use_logits", "default": "True", "kwonly": true}, {"name": "clones", "default": "None", "kwonly": true}, {"name": "temperature", "default": "1.0", "kwonly": true}, {"name": "random_state", "default": "None", "kwonly": true, "type": "int | numpy.Generator | torch.Generator | None"}, {"name": "device", "default": "None", "kwonly": true}]}, {"name": "clonotypic_entropy", "module": "tcri/tools/_entropy.py", "namespace": "tl", "responsibility": "Per phenotype, normalized Shannon entropy of the distribution over clonotypes carrying it, H[P(c|phi,m)] (spread across clones); repertoire-level; returns float | pandas.Series | pandas.DataFrame.", "args": [{"name": "adata_or_jd"}, {"name": "covariate", "default": "None", "kwonly": true}, {"name": "groupby", "default": "None", "kwonly": true}, {"name": "n_samples", "default": "0", "kwonly": true}, {"name": "temperature", "default": "1.0", "kwonly": true}, {"name": "clones", "default": "None", "kwonly": true}, {"name": "normalized", "default": "True", "kwonly": true}, {"name": "n_clones_ref", "default": "None", "kwonly": true}, {"name": "random_state", "default": "None", "kwonly": true, "type": "int | numpy.Generator | torch.Generator | None"}, {"name": "device", "default": "None", "kwonly": true}]}, {"name": "phenotypic_entropy", "module": "tcri/tools/_entropy.py", "namespace": "tl", "responsibility": "Per clonotype, normalized Shannon entropy of its phenotype distribution H[P(phi|c,m)] (plasticity vs commitment); zero-support clone -> NaN; returns float | pandas.Series | pandas.DataFrame.", "args": [{"name": "adata_or_jd"}, {"name": "covariate", "default": "None", "kwonly": true}, {"name": "groupby", "default": "None", "kwonly": true}, {"name": "n_samples", "default": "0", "kwonly": true}, {"name": "temperature", "default": "1.0", "kwonly": true}, {"name": "clones", "default": "None", "kwonly": true}, {"name": "normalized", "default": "True", "kwonly": true}, {"name": "random_state", "default": "None", "kwonly": true, "type": "int | numpy.Generator | torch.Generator | None"}, {"name": "device", "default": "None", "kwonly": true}]}, {"name": "mutual_information", "module": "tcri/tools/_mutual_information.py", "namespace": "tl", "responsibility": "I(c;phi|m) in bits (clone-phenotype coupling), optionally normalized to [0,1] by normalize_mode; returns float | numpy.ndarray | pandas.DataFrame.", "args": [{"name": "adata_or_jd"}, {"name": "covariate", "default": "None", "kwonly": true}, {"name": "groupby", "default": "None", "kwonly": true}, {"name": "n_samples", "default": "0", "kwonly": true}, {"name": "temperature", "default": "1.0", "kwonly": true}, {"name": "clones", "default": "None", "kwonly": true}, {"name": "normalized", "default": "True", "kwonly": true}, {"name": "normalize_mode", "default": "\"min\"", "kwonly": true}, {"name": "random_state", "default": "None", "kwonly": true, "type": "int | numpy.Generator | torch.Generator | None"}, {"name": "device", "default": "None", "kwonly": true}]}, {"name": "_mi_from_joint", "module": "tcri/tools/_mutual_information.py", "namespace": "tl", "responsibility": "Private kernel: normalize joint, compute I in bits (eps=1e-15); mode='min'->I/min(H_c,H_p), 'average'->I/(1/2(H_c+H_p)); returns float.", "args": [{"name": "pxy"}, {"name": "normalized", "kwonly": true}, {"name": "mode", "default": "\"min\"", "kwonly": true}]}, {"name": "phenotypic_flux", "module": "tcri/tools/_flux.py", "namespace": "tl", "responsibility": "(renamed from flux) Per-clonotype distance between a clone's phenotype distribution at cov_from vs cov_to over the clone intersection; returns pandas.Series | numpy.ndarray | pandas.DataFrame.", "args": [{"name": "adata"}, {"name": "cov_from", "kwonly": true}, {"name": "cov_to", "kwonly": true}, {"name": "groupby", "default": "None", "kwonly": true}, {"name": "n_samples", "default": "0", "kwonly": true}, {"name": "temperature", "default": "1.0", "kwonly": true}, {"name": "clones", "default": "None", "kwonly": true}, {"name": "distance_metric", "default": "\"l1\"", "kwonly": true}, {"name": "random_state", "default": "None", "kwonly": true, "type": "int | numpy.Generator | torch.Generator | None"}, {"name": "device", "default": "None", "kwonly": true}]}, {"name": "compare_groups", "module": "tcri/tools/_compare.py", "namespace": "tl", "responsibility": "Public replacement for mi_compare/delta_entropy_table/flux_table: turn a tidy groupby result or paired posterior draws into group contrasts (MWU, delta, and p_gt/p_lt+HDI on signed delta); returns pandas.DataFrame.", "args": [{"name": "df"}, {"name": "value", "kwonly": true}, {"name": "by", "kwonly": true}, {"name": "reference", "default": "None", "kwonly": true}, {"name": "paired", "default": "False", "kwonly": true}, {"name": "hdi_prob", "default": "0.94", "kwonly": true}, {"name": "alternative", "default": "\"two-sided\"", "kwonly": true}]}, {"name": "clonotypic_entropy", "module": "tcri/plotting/_entropy.py", "namespace": "pl", "responsibility": "(renamed from clonotypic_entropy_by_phenotype) Box-and-dot plot of clonotypic entropy per phenotype across covariate values, per-group dots, significance brackets.", "args": [{"name": "adata"}, {"name": "covariate", "default": "None", "kwonly": true}, {"name": "groupby", "default": "None", "kwonly": true}, {"name": "splitby", "default": "None", "kwonly": true}, {"name": "n_samples", "default": "0", "kwonly": true}, {"name": "temperature", "default": "1.0", "kwonly": true}, {"name": "clones", "default": "None", "kwonly": true}, {"name": "normalized", "default": "True", "kwonly": true}, {"name": "n_clones_ref", "default": "None", "kwonly": true}, {"name": "palette", "default": "None", "kwonly": true}, {"name": "hue_order", "default": "None", "kwonly": true}, {"name": "ax", "default": "None", "kwonly": true}, {"name": "figsize", "default": "(6, 3)", "kwonly": true}, {"name": "rotation", "default": "90", "kwonly": true}, {"name": "legend_fontsize", "default": "6", "kwonly": true}, {"name": "bbox_to_anchor", "default": "(1.15, 1.0)", "kwonly": true}, {"name": "random_state", "default": "None", "kwonly": true}, {"name": "save", "default": "None", "kwonly": true}, {"name": "return_df", "default": "False", "kwonly": true}]}, {"name": "phenotypic_entropy", "module": "tcri/plotting/_entropy.py", "namespace": "pl", "responsibility": "[FIXED] Box/strip plot of phenotypic entropy per covariate/group.", "args": [{"name": "adata"}, {"name": "covariate", "default": "None", "kwonly": true}, {"name": "groupby", "default": "None", "kwonly": true}, {"name": "splitby", "default": "None", "kwonly": true}, {"name": "n_samples", "default": "0", "kwonly": true}, {"name": "temperature", "default": "1.0", "kwonly": true}, {"name": "clones", "default": "None", "kwonly": true}, {"name": "normalized", "default": "True", "kwonly": true}, {"name": "palette", "default": "None", "kwonly": true}, {"name": "ax", "default": "None", "kwonly": true}, {"name": "figsize", "default": "(8, 4)", "kwonly": true}, {"name": "rotation", "default": "90", "kwonly": true}, {"name": "legend_fontsize", "default": "6", "kwonly": true}, {"name": "bbox_to_anchor", "default": "(1.15, 1.0)", "kwonly": true}, {"name": "random_state", "default": "None", "kwonly": true}, {"name": "save", "default": "None", "kwonly": true}, {"name": "return_df", "default": "False", "kwonly": true}]}, {"name": "mutual_information", "module": "tcri/plotting/_mutual_information.py", "namespace": "pl", "responsibility": "[FIXED; subsumes mi_compare] Box/strip plot of clone x phenotype MI per covariate; groupby supplies per-group points and drives AUROC/MWU/label-permutation stats; weighted removed; default normalize_mode='min'.", "args": [{"name": "adata"}, {"name": "covariate", "default": "None", "kwonly": true}, {"name": "groupby", "default": "None", "kwonly": true}, {"name": "splitby", "default": "None", "kwonly": true}, {"name": "n_samples", "default": "0", "kwonly": true}, {"name": "temperature", "default": "1.0", "kwonly": true}, {"name": "clones", "default": "None", "kwonly": true}, {"name": "normalized", "default": "True", "kwonly": true}, {"name": "normalize_mode", "default": "\"min\"", "kwonly": true}, {"name": "palette", "default": "None", "kwonly": true}, {"name": "ax", "default": "None", "kwonly": true}, {"name": "figsize", "default": "(8, 4)", "kwonly": true}, {"name": "rotation", "default": "90", "kwonly": true}, {"name": "legend_fontsize", "default": "6", "kwonly": true}, {"name": "bbox_to_anchor", "default": "(1.15, 1.0)", "kwonly": true}, {"name": "random_state", "default": "None", "kwonly": true}, {"name": "save", "default": "None", "kwonly": true}, {"name": "return_df", "default": "False", "kwonly": true}]}, {"name": "phenotypic_flux", "module": "tcri/plotting/_flux.py", "namespace": "pl", "responsibility": "The Sankey: phenotype-distribution flow across the ordered 'order' sequence, calling tl.phenotypic_flux pairwise between consecutive values under one shared draw stack.", "args": [{"name": "adata"}, {"name": "order", "kwonly": true}, {"name": "groupby", "default": "None", "kwonly": true}, {"name": "clones", "default": "None", "kwonly": true}, {"name": "normalize", "default": "True", "kwonly": true}, {"name": "temperature", "default": "1.0", "kwonly": true}, {"name": "distance_metric", "default": "\"l1\"", "kwonly": true}, {"name": "phenotype_colors", "default": "None", "kwonly": true}, {"name": "ax", "default": "None", "kwonly": true}, {"name": "figsize", "default": "(6, 3)", "kwonly": true}, {"name": "show_legend", "default": "True", "kwonly": true}, {"name": "title", "default": "None", "kwonly": true}, {"name": "random_state", "default": "None", "kwonly": true}, {"name": "save", "default": "None", "kwonly": true}, {"name": "return_axes", "default": "False", "kwonly": true}]}, {"name": "SankeyNode.__init__", "module": "tcri/plotting/_sankey.py", "namespace": "pl", "responsibility": "Private Sankey drawing node.", "args": [{"name": "self"}, {"name": "x"}, {"name": "y"}, {"name": "val"}, {"name": "dx", "default": "0.2", "kwonly": true}, {"name": "color", "default": "None", "kwonly": true}, {"name": "**kwargs"}]}, {"name": "SankeyNode.plot", "module": "tcri/plotting/_sankey.py", "namespace": "pl", "responsibility": "Draw the node on ax. Private.", "args": [{"name": "self"}, {"name": "ax"}]}, {"name": "SankeyNode.plot_node_connection", "module": "tcri/plotting/_sankey.py", "namespace": "pl", "responsibility": "Curved, color-interpolated ribbon to destination_node. Private.", "args": [{"name": "self"}, {"name": "destination_node"}, {"name": "ax"}, {"name": "**kwargs"}]}, {"name": "_phenotype_mass_per_clone", "module": "tcri/plotting/_sankey.py", "namespace": "pl", "responsibility": "Return {clone -> phenotype-mass vector} at one covariate; returns dict[str, np.ndarray]. Private.", "args": [{"name": "adata"}, {"name": "covariate"}, {"name": "clones"}, {"name": "normalize"}]}, {"name": "_metric_boxplot", "module": "tcri/plotting/_base.py", "namespace": "pl", "responsibility": "(renamed from tcri_boxplot) Generic per-phenotype metric box/strip engine across groupby/splitby strata; computes each stratum by full-space restriction; returns (fig, ax). Private.", "args": [{"name": "adata"}, {"name": "function"}, {"name": "groupby", "default": "None", "kwonly": true}, {"name": "splitby", "default": "None", "kwonly": true}, {"name": "ylabel", "default": "\"\"", "kwonly": true}, {"name": "order", "default": "None", "kwonly": true}, {"name": "palette", "default": "None", "kwonly": true}, {"name": "s", "default": "20", "kwonly": true}, {"name": "ax", "default": "None", "kwonly": true}, {"name": "figsize", "default": "(8, 4)", "kwonly": true}]}, {"name": "_finish", "module": "tcri/plotting/_base.py", "namespace": "pl", "responsibility": "scanpy-style show/save/return finalizer. Private.", "args": [{"name": "fig"}, {"name": "ax"}, {"name": "save", "default": "None", "kwonly": true}, {"name": "show", "default": "None", "kwonly": true}, {"name": "return_axes", "default": "False", "kwonly": true}]}, {"name": "resolve_palette", "module": "tcri/plotting/_colors.py", "namespace": "pl", "responsibility": "(renamed from set_color_palette) Assign tcri_colors to each obs column's categories, store in uns['_colors'], return the map (mutates in place); returns dict.", "args": [{"name": "adata"}, {"name": "columns"}, {"name": "palette", "default": "None", "kwonly": true}]}, {"name": "probability_ternary", "module": "tcri/plotting/_ternary.py", "namespace": "pl", "responsibility": "Ternary phenotype-simplex scatter of per-cell/per-clone phenotype probabilities; kept public; reads K.X_PROBABILITIES; weighted removed.", "args": [{"name": "adata"}, {"name": "phenotypes", "kwonly": true}, {"name": "groupby", "default": "None", "kwonly": true}, {"name": "clones", "default": "None", "kwonly": true}, {"name": "palette", "default": "None", "kwonly": true}, {"name": "ax", "default": "None", "kwonly": true}, {"name": "figsize", "default": "(5, 5)", "kwonly": true}, {"name": "save", "default": "None", "kwonly": true}, {"name": "return_axes", "default": "False", "kwonly": true}]}, {"name": "joint_distribution_ppc", "module": "tcri/diagnostics/_ppc.py", "namespace": "diag", "responsibility": "(fixed compare_joint_distribution) Model vs empirical per-clone phenotype frequencies; per-clone L1 or KL(emp||model) + per-covariate aggregate; model-free (adata only); returns pandas.DataFrame.", "args": [{"name": "adata"}, {"name": "covariate", "default": "None", "kwonly": true}, {"name": "distance_metric", "default": "\"l1\"", "kwonly": true}, {"name": "temperature", "default": "1.0", "kwonly": true}]}, {"name": "phenotype_calibration", "module": "tcri/diagnostics/_ppc.py", "namespace": "diag", "responsibility": "Reliability of predict() probabilities: bin by predicted max-prob, compare mean predicted prob to empirical accuracy, ECE; adata only; returns (bin, mean_pred, emp_freq, count) + scalar ECE.", "args": [{"name": "adata"}, {"name": "n_bins", "default": "10", "kwonly": true}]}, {"name": "reconstruction_ppc", "module": "tcri/diagnostics/_ppc.py", "namespace": "diag", "responsibility": "ZINB reconstruction PPC: simulate from fitted decoder (mu,theta,pi_dropout), compare library size / per-gene dropout / mean-variance vs observed; model REQUIRED; returns pandas.DataFrame.", "args": [{"name": "model"}, {"name": "adata", "default": "None"}, {"name": "n_samples", "default": "100", "kwonly": true}, {"name": "seed", "default": "0", "kwonly": true}]}, {"name": "permutation_null", "module": "tcri/diagnostics/_ppc.py", "namespace": "diag", "responsibility": "Permute phenotype labels within each covariate R times, recompute the metric for a null; p=mean(null>=obs), z-score; adata only; one shared draw stack; returns per-stratum (observed, null_mean, null_sd, z, p).", "args": [{"name": "adata"}, {"name": "metric", "default": "\"mutual_information\"", "kwonly": true}, {"name": "covariate", "default": "None", "kwonly": true}, {"name": "groupby", "default": "None", "kwonly": true}, {"name": "n_permutations", "default": "1000", "kwonly": true}, {"name": "seed", "default": "0", "kwonly": true}]}, {"name": "loss", "module": "tcri/diagnostics/_training.py", "namespace": "diag", "responsibility": "(relocated plot_loss) Plot training/validation ELBO and prior-KL from model.history_.", "args": [{"name": "model"}, {"name": "log_scale", "default": "False", "kwonly": true}, {"name": "ax", "default": "None", "kwonly": true}, {"name": "save", "default": "None", "kwonly": true}]}, {"name": "archetypes", "module": "tcri/diagnostics/_training.py", "namespace": "diag", "responsibility": "(relocated plot_archetypes) Cluster-ordered clone-phenotype heatmap + archetype centroids, ordered by build_archetypes labels.", "args": [{"name": "model"}, {"name": "ax", "default": "None", "kwonly": true}, {"name": "save", "default": "None", "kwonly": true}]}, {"name": "save_tcri_session", "module": "tcri/utils/_session.py", "namespace": "ut", "responsibility": "Persist a trained session: scvi model (weights+registry, no embedded adata), Pyro param store, setup.json, the h5ad; returns dict.", "args": [{"name": "model"}, {"name": "adata"}, {"name": "out_dir"}, {"name": "save_adata", "default": "True", "kwonly": true}, {"name": "compression", "default": "\"gzip\"", "kwonly": true}]}, {"name": "load_tcri_session", "module": "tcri/utils/_session.py", "namespace": "ut", "responsibility": "Reconstruct TCRIModel + AnnData: read h5ad, restore setup/category order, re-run setup_anndata, load model + Pyro params (sets global store before any get_p_ct/to_anndata); returns (TCRIModel, AnnData).", "args": [{"name": "run_dir"}, {"name": "adata_path", "default": "None", "kwonly": true}, {"name": "map_location", "default": "None", "kwonly": true}, {"name": "layer", "default": "None", "kwonly": true}]}, {"name": "_to_jsonable", "module": "tcri/utils/_session.py", "namespace": "ut", "responsibility": "Private: coerce a value to a JSON-serializable form.", "args": [{"name": "x"}]}, {"name": "_collect_setup_from_adata_or_model", "module": "tcri/utils/_session.py", "namespace": "ut", "responsibility": "Private: gather setup info from adata or model.", "args": [{"name": "adata"}, {"name": "model"}]}, {"name": "_restore_category_order", "module": "tcri/utils/_session.py", "namespace": "ut", "responsibility": "Private: restore category ordering on adata from saved setup.", "args": [{"name": "adata"}, {"name": "setup"}]}, {"name": "_resolve_TCRIModel", "module": "tcri/utils/_session.py", "namespace": "ut", "responsibility": "Private: resolve the TCRIModel class.", "args": []}, {"name": "_disable_scvi_onload_train", "module": "tcri/utils/_session.py", "namespace": "ut", "responsibility": "Private: disable scvi on-load training behavior.", "args": []}, {"name": "_ensure_pyro_posterior_params", "module": "tcri/utils/_session.py", "namespace": "ut", "responsibility": "Private: guarantee q_p_ct_raw exists; if missing, warn + re-init to uniform 1/P.", "args": [{"name": "model"}, {"name": "adata"}]}, {"name": "_pyro_load", "module": "tcri/utils/_session.py", "namespace": "ut", "responsibility": "Private: load Pyro param store from path.", "args": [{"name": "path"}, {"name": "map_location", "default": "None", "kwonly": true}]}, {"name": "_ensure_dir", "module": "tcri/utils/_session.py", "namespace": "ut", "responsibility": "Private: ensure a directory exists.", "args": [{"name": "path"}]}]}, "nbUsage": {"usage": [{"function": "tcri.pp.joint_distribution", "args_seen": ["", "", "method", "temperature", "n_samples", "clones", "weighted", "covariate_label", "adata"], "n_calls": 76}, {"function": "tcri.pp.register_clonotype_key", "args_seen": ["", ""], "n_calls": 23}, {"function": "tcri.pp.clone_size", "args_seen": [""], "n_calls": 23}, {"function": "tcri.pp.register_phenotype_key", "args_seen": ["", ""], "n_calls": 11}, {"function": "tcri.pp.gene_entropy", "args_seen": ["", "batch_key"], "n_calls": 5}, {"function": "tcri.pp.register_model", "args_seen": ["", "", "gate_prob"], "n_calls": 2}, {"function": "tcri.pp.group_singletons", "args_seen": ["", "min_clone_size"], "n_calls": 1}, {"function": "tcri.tl.phenotypic_entropies", "args_seen": ["", "normalized", "method", "decimals"], "n_calls": 17}, {"function": "tcri.tl.mutual_information", "args_seen": ["", "covariate", "method", "temperature", "clones", "n_samples", "weighted", "normalized", "normalize_mode", "adata"], "n_calls": 15}, {"function": "tcri.tl.clonotypic_entropies", "args_seen": ["", "normalized", "method"], "n_calls": 8}, {"function": "tcri.tl.flux", "args_seen": ["", "from_this", "to_that", "distance_metric", "key", "clones", "temperature"], "n_calls": 4}, {"function": "tcri.tl.probability_distribution", "args_seen": ["", "method"], "n_calls": 4}, {"function": "tcri.tl.phenotypic_entropy (singular)", "args_seen": ["", "base", "normalized", "temperature", "n_samples"], "n_calls": 2}, {"function": "tcri.tl.clonotypic_entropy (singular)", "args_seen": ["", "temperature", "clones", "n_samples", "normalized"], "n_calls": 1}, {"function": "tcri.tl.phenotypic_entropy_delta", "args_seen": ["", "", "", "", ""], "n_calls": 1}, {"function": "tcri.tl.clonality (imported as clonality_tl, not called in uncommented cells)", "args_seen": [], "n_calls": 0}, {"function": "tcri.tl.clone_fraction (imported as clone_fraction_tl, not called)", "args_seen": [], "n_calls": 0}, {"function": "tcri.pl.probability_ternary", "args_seen": ["", "", "splitby", "conditions", "color", "scale_function", "phenotype_names", "save", "nt", "method", "top_n"], "n_calls": 19}, {"function": "tcri.pl.phenotypic_entropy", "args_seen": ["", "groupby", "splitby", "figsize", "order", "normalized", "save", "rotation", "minimum_clone_size", "return_df"], "n_calls": 9}, {"function": "tcri.pl.phenotypic_flux", "args_seen": ["", "splitby", "order", "nt", "normalize", "figsize", "phenotype_colors", "save"], "n_calls": 8}, {"function": "tcri.pl.clonotypic_entropy", "args_seen": ["", "groupby", "method", "normalized", "figsize", "order", "splitby"], "n_calls": 4}, {"function": "tcri.pl.clonality", "args_seen": ["", "groupby", "splitby", "figsize", "order"], "n_calls": 4}, {"function": "tcri.pl.mutual_information", "args_seen": ["", "groupby", "splitby", "figsize", "method", "minimum_clone_size", "return_df", "bbox_to_anchor"], "n_calls": 3}, {"function": "tcri.pl.flux", "args_seen": ["", "key", "order", "groupby", "method", "distance_metric", "paint", "figsize", "palette"], "n_calls": 3}, {"function": "tcri.pl.phenotypic_entropy_delta", "args_seen": ["", "groupby", "key", "from_this", "to_that", "figsize", "palette", "save"], "n_calls": 3}, {"function": "tcri.pl.top_clone_umap", "args_seen": ["", "bg_size", "bg_alpha", "figsize", "top_n", "fg_size"], "n_calls": 2}, {"function": "tcri.pl.polar_plot", "args_seen": ["", "splitby", "statistic", "title", "save", "color_dict"], "n_calls": 1}, {"function": "tcri.pl.compare_phenotypes", "args_seen": ["", "", ""], "n_calls": 1}, {"function": "tcri.pl.tcri_colors (palette: indexed tcri_colors[i] and .copy())", "args_seen": [], "n_calls": 12}, {"function": "tcri.ml.TCRIModel (constructor)", "args_seen": ["", "n_latent", "n_hidden", "n_layers", "global_scale", "local_scale", "use_enumeration", "patience", "classifier_hidden", "classifier_dropout", "classifer_dropout(typo)", "classifier_n_layers", "classifier_n_heads", "sharp_temperature", "K", "gate_prob", "kl_weight"], "n_calls": 2}, {"function": "tcri.ml.TCRIModel.setup_anndata", "args_seen": ["", "layer", "clonotype_key", "phenotype_key", "covariate_key", "batch_key"], "n_calls": 2}, {"function": "tcri.ml.JointProbabilityDistribution", "args_seen": ["", "", "batch_size"], "n_calls": 1}, {"function": "tcri.ut.Tcell", "args_seen": ["phenotypes", "phenotypes_and_counts", "TRB", "use_genes"], "n_calls": 2}, {"function": "tcri.ut.CellRepertoire", "args_seen": ["clones_and_phenos", "phenotypes", "use_genes", "use_chain", "seq_type", "chains_to_use", "name"], "n_calls": 1}, {"function": "tcri.ut.Phenotypes", "args_seen": [], "n_calls": 1}, {"function": "tcri.ut.plot_pheno_sankey", "args_seen": ["phenotypes", "cell_repertoires", "clones", "times", "xlim", "time_rescale", "normalize", "xlabel", "return_axes", "show_legend", "figsize", "phenotype_colors"], "n_calls": 1}], "capabilities_to_preserve": ["joint_distribution as the central primitive: callable as tcri.pp.joint_distribution(adata, covariate) positionally, dominating usage (76 call sites). Must return a tidy per-clone x per-phenotype DataFrame that downstream code binds directly (xdata/jd) and feeds into ternary/entropy/MI functions.", "method= toggle between 'probabilistic' (model-posterior) and empirical/raw distributions, threaded consistently through joint_distribution, tl.mutual_information, tl.clonotypic_entropies, tl.phenotypic_entropies, tl.probability_distribution, tl.flux, pl.clonotypic_entropy, pl.mutual_information, pl.flux.", "temperature= softmax-sharpening knob on joint_distribution, tl.mutual_information, tl.clonotypic_entropy, tl.phenotypic_entropy, tl.flux (default effectively 1.0).", "n_samples= posterior-sampling count on joint_distribution, tl.mutual_information, tl.clonotypic_entropy, tl.phenotypic_entropy.", "clones= restrict computation to an explicit clonotype subset on joint_distribution, tl.mutual_information, tl.clonotypic_entropy, tl.flux.", "weighted= (clone-size weighting) on tl.mutual_information and joint_distribution.", "normalized= on tl.clonotypic_entropies, tl.phenotypic_entropies, tl.mutual_information, pl.clonotypic_entropy, pl.phenotypic_entropy; plus normalize_mode= on mutual_information and base=/decimals= rounding controls on the entropy metrics.", "Both PLURAL aggregate forms (tl.clonotypic_entropies, tl.phenotypic_entropies -> table over groups) and SINGULAR per-item forms (tl.clonotypic_entropy, tl.phenotypic_entropy) of the entropy metrics; notebooks use both.", "covariate_label= / covariate= selection of a single covariate value or covariate key for joint_distribution and mutual_information.", "groupby= + splitby= grouping/faceting semantics shared across all pl.* metrics (clonality, clonotypic_entropy, phenotypic_entropy, mutual_information, flux) and via splitby+conditions in probability_ternary/phenotypic_flux.", "Temporal delta / change metrics: pl.phenotypic_entropy_delta and tl.phenotypic_entropy_delta(groupby, key, from_this, to_that); tl.flux(from_this, to_that, key, distance_metric). distance_metric= (e.g. 'l1') and paint=/paint_dict= for pl.flux.", "probability_ternary (19 calls, heaviest plotting fn): 3 phenotype_names -> ternary corners, splitby+conditions to pick timepoint/response group, nt= boolean toggle, method='probabilistic', scale_function= callable mapping clone frequency -> marker size, color=, top_n=, save=. Must be reproducible on the new surface.", "phenotypic_flux Sankey/flow plot: splitby, order, nt, normalize, phenotype_colors, save.", "return_df= on plotting functions (pl.phenotypic_entropy, pl.mutual_information) to retrieve the underlying table instead of only a figure; minimum_clone_size= filtering; bbox_to_anchor/figsize/order/rotation/save styling knobs.", "Registration/setup preprocessing workflow that gates all metrics: pp.register_clonotype_key(adata, key), pp.register_phenotype_key(adata, key), pp.register_model(adata, model, gate_prob=), pp.clone_size(adata), pp.group_singletons(min_clone_size=), pp.gene_entropy(batch_key=).", "Probabilistic model surface ml.TCRIModel with architecture/inference knobs (n_latent, n_hidden, n_layers, global_scale, local_scale, use_enumeration, patience, classifier_hidden/dropout/n_layers/n_heads, sharp_temperature, K, gate_prob, kl_weight), TCRIModel.setup_anndata(layer, clonotype_key, phenotype_key, covariate_key, batch_key), model training, and its role backing method='probabilistic' joint_distribution.", "Color palette accessor tcri.pl.tcri_colors: an indexable sequence (tcri_colors[i]) supporting .copy(); used to hand explicit colors to ternary/flux/entropy plots.", "Auxiliary plots: pl.top_clone_umap(top_n, bg_size, bg_alpha, fg_size, figsize) UMAP overlay of top clones; pl.polar_plot(statistic, splitby, color_dict, title); pl.compare_phenotypes; pl.clonality(groupby, splitby, order).", "Old metrics-module functions defined in current source but NOT exercised by notebooks yet expected to remain reproducible per task: tl.mi_compare(groupby, groups, treatment, n_samples,...), tl.delta_entropy_table, tl.flux_table, tl.delta_clonotypic_entropy, plus pl.clonotypic_entropy_by_phenotype.", "Legacy object-oriented repertoire model (tcri.utils/ut) used in yost_v2: Phenotypes, CellRepertoire(clones_and_phenos, phenotypes, use_genes, use_chain, seq_type, chains_to_use, name), Tcell(phenotypes, phenotypes_and_counts, TRB, use_genes), plot_pheno_sankey(...), and imported helpers plot_pheno_ternary_change_plots, draw_clone_bars, probabilities, set_ternary_corner_label, ternary_plot_projection."], "notes": "Scanned all 8 example notebooks (colitis/comparison, colitis/preprocess_conga, renal/renal, smith/smith, synthetic/synthetic1, yost/yost, yost/yost_v2, zhang/zhang) via AST parse (0 parse failures) plus an alias map that resolves `from tcri. import X as Y` bare calls. Only import alias is plain `import tcri` (no `tc`/`import tcri as tc` anywhere). Old module names map to the target surface: tcri.metrics->tl, tcri.preprocessing->pp, tcri.utils->ut, tcri.pl->pl, tcri.model->ml; notebooks call BOTH old paths (tcri.metrics.mutual_information, tcri.tl.mutual_information) so treat metrics==tl and preprocessing==pp. args_seen lists distinct keyword names observed across all call sites (positional args marked ); n_calls is total across notebooks. Ranking: joint_distribution (76) >> register_clonotype_key & clone_size (23) > probability_ternary (19) > phenotypic_entropies (17) > mutual_information tl (15) > register_phenotype_key (11). IMPORTANT status check against current /Users/ceglian/Codebase/GitHub/tcri/tcri source: probability_ternary, phenotypic_entropy_delta, phenotypic_entropies (plural), and clonotypic_entropies (plural) are CALLED by the notebooks but have NO `def` in the current tcri/ tree (removed/renamed) -- these are exactly the surfaces the new API must reproduce. Conversely mi_compare (tcri/metrics/_metrics.py:93 and duplicated in plotting/_plotting.py:95), delta_entropy_table (_metrics.py:392), flux_table (_metrics.py:765), delta_clonotypic_entropy (_metrics.py:324), and clonotypic_entropy_by_phenotype (plotting/_plotting.py:664) DO exist in current source but are NOT invoked in any notebook; clonotypic_entropy_by_phenotype is additionally copy-pasted as a local def inside yost_v2. Shadowing caveat: yost_v2, smith, and colitis/comparison define local wrappers named mutual_information/phenotypic_entropy/phenotypic_entropy_delta that shadow the tcri imports, but those wrappers call the tcri function underneath, so the recorded kwargs reflect the real tcri signatures. A typo kwarg `classifer_dropout` (missing 'i') appears alongside the correct `classifier_dropout` in TCRIModel construction -- treat as the intended classifier_dropout. clonality_tl and clone_fraction_tl are imported in several notebooks but never called in uncommented cells (n_calls=0). yost_v2 is the only notebook exercising the ml.TCRIModel training path and the legacy ut.CellRepertoire/Tcell/plot_pheno_sankey object model."}} \ No newline at end of file diff --git a/docs/contract/_quarantine/tcri_arg_consistency_sufficiency.md b/docs/contract/_quarantine/tcri_arg_consistency_sufficiency.md new file mode 100644 index 0000000..2fffff2 --- /dev/null +++ b/docs/contract/_quarantine/tcri_arg_consistency_sufficiency.md @@ -0,0 +1,417 @@ +# TCRI — Argument Consistency & Sufficiency (Final) + +**Scope.** The analysis/plot/diagnostic/preprocessing/stats/engine surface defined in `docs/contract/tcri_api_and_responsibilities.md` (§0–§12). Pure internal helpers with no shared-vocabulary arguments (`_console`, priors/module/classifier/training internals, `_register._write_*`, session utils, `_compute` reducers, `resolve_device`/`get_xp`/`asnumpy`) are omitted as all-`—` rows. + +**Method.** A full parameter matrix was extracted, every shared concept was checked against the four RULES (same concept → same name/order/default/type; kw-only except the primary positional; American spelling; `groupby`≠`splitby`; `covariate=None`=all; `n_samples=0`=point estimate), then an argument-sufficiency pass compared the surface to the six example notebooks. A verifier then adjudicated every finding against both contract docs and the notebooks. + +**This document folds in the verifier's verdicts:** refuted findings are dropped (listed once, for the record); revised findings are restated in their corrected form; the verifier's missed items are added. Each item carries its audit id and status: **[confirmed]**, **[revised]**, or **[missed]**. + +**Dropped as refuted (not actionable — recorded so they are not re-raised):** +- **S3 `weighted=`** — deliberately removed and documented (§0.8, appendix row-and-changelog). Behavior change, not an oversight. +- **S8 `pp.gene_entropy(batch_key=)`** — relocated to `examples/` with its notebook (§11), so `batch_key`/`agg_function` travel with it. Not lost. +- **S12 `base=`** — single-base (`log2`/bits) is an intentional unit-consistency fix (§3.4, §7.2/7.3, appendix row 12). Re-adding `base=` would undo it. `decimals=` is cosmetic; drop. + +--- + +## 1. PARAMETER MATRIX + +Cells show `name=default` (keyword-only unless noted), `name(req)` = required keyword-only (no default), `POS` = positional, `—` = absent, **DIVERGENT** = the concept exists under a different name/default, **— (missing)** = the concept is applicable but absent. + +### Table A — engine / tl / pl / diag / pp core vocabulary + +| Function (ns/module) | primary(pos) | covariate | cov_from/cov_to | order | groupby | splitby | n_samples | temperature | clones | use_logits | normalized | normalize_mode | distance_metric | n_clones_ref | random_state | device | palette | ax | figsize | save | return | +|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---| +| tl.joint_distribution | adata | covariate=None | — | — | groupby=None | — | n_samples=0 | temperature=1.0 | clones=None | use_logits=True | — | — | — | — | random_state=None | device=None | — | — | — | — | — | +| tl.clonotypic_entropy | adata_or_jd | covariate=None | — | — | groupby=None | — | n_samples=0 | temperature=1.0 | clones=None | — | normalized=True | — | — | n_clones_ref=None | random_state=None | device=None | — | — | — | — | — | +| tl.phenotypic_entropy | adata_or_jd | covariate=None | — | — | groupby=None | — | n_samples=0 | temperature=1.0 | clones=None | — | normalized=True | — | — | — | random_state=None | device=None | — | — | — | — | — | +| tl.mutual_information | adata_or_jd | covariate=None | — | — | groupby=None | — | n_samples=0 | temperature=1.0 | clones=None | — | normalized=True | normalize_mode="min" | — | — | random_state=None | device=None | — | — | — | — | — | +| tl.phenotypic_flux | adata | — | cov_from(req), cov_to(req) | — | groupby=None | — | n_samples=0 | temperature=1.0 | clones=None | — | — | — | distance_metric="l1" | — | random_state=None | device=None | — | — | — | — | — | +| tl.compare_groups | df | — | — | — | **DIVERGENT: by(req)** | **DIVERGENT: by(req)** | — | — | — | — | — | — | — | — | — | — | — | — | — | — | — | +| pl.clonotypic_entropy | adata | covariate=None | — | **— (missing)** | groupby=None | splitby=None | n_samples=0 | temperature=1.0 | clones=None | — | normalized=True | — | — | n_clones_ref=None | random_state=None | **— (missing)** | palette=None | ax=None | figsize=(6,3) | save=None | return_df=False | +| pl.phenotypic_entropy | adata | covariate=None | — | **— (missing)** | groupby=None | splitby=None | n_samples=0 | temperature=1.0 | clones=None | — | normalized=True | — | — | — | random_state=None | **— (missing)** | palette=None | ax=None | figsize=(8,4) | save=None | return_df=False | +| pl.mutual_information | adata | covariate=None | — | **— (missing)** | groupby=None | splitby=None | n_samples=0 | temperature=1.0 | clones=None | — | normalized=True | normalize_mode="min" | — | — | random_state=None | **— (missing)** | palette=None | ax=None | figsize=(8,4) | save=None | return_df=False | +| pl.phenotypic_flux | adata | — | — | order(req) | groupby=None | **— (missing)** | **— (missing n_samples)** | temperature=1.0 | clones=None | — | **DIVERGENT: normalize=True** | — | distance_metric="l1" | — | random_state=None | **— (missing)** | **DIVERGENT: phenotype_colors=None** | ax=None | figsize=(6,3) | save=None | return_axes=False | +| pl.probability_ternary | adata | **— (missing)** | — | — | groupby=None | **— (missing)** | **— (missing)** | **— (missing)** | clones=None | — | — | — | — | — | **— (missing)** | — | palette=None | ax=None | figsize=(5,5) | save=None | return_axes=False | +| pl._metric_boxplot (priv) | adata, function POS | — | — | order=None | groupby=None | splitby=None | — | — | — | — | — | — | — | — | — | — | palette=None | ax=None | figsize=(8,4) | — | — | +| pl._finish (priv) | fig, ax POS | — | — | — | — | — | — | — | — | — | — | — | — | — | — | — | — | ax POS | — | save=None | return_axes=False; **show=None (only here)** | +| pl.resolve_palette | adata, columns POS | — | — | — | — | — | — | — | — | — | — | — | — | — | — | — | palette=None | — | — | — | — | +| diag.joint_distribution_ppc | adata | covariate=None | — | — | **— (missing)** | — | **— (missing)** | temperature=1.0 | — | — | — | — | distance_metric="l1" | — | **— (missing)** | — | — | — | — | — | — | +| diag.phenotype_calibration | adata | — | — | — | — | — | — | — | — | — | — | — | — | — | — | — | — | — | — | — | (n_bins=10) | +| diag.reconstruction_ppc | model, adata=None POS | — | — | — | — | — | **DIVERGENT: n_samples=100** | — | — | — | — | — | — | — | **DIVERGENT: seed=0** | — | — | — | — | — | — | +| diag.permutation_null | adata | covariate=None | — | — | groupby=None | **— (missing)** | **— (metric="mutual_information"; no metric passthrough)** | — | — | — | — | — | — | — | **DIVERGENT: seed=0** | — | — | — | — | — | (n_permutations=1000) | +| diag.loss | model | — | — | — | — | — | — | — | — | — | — | — | — | — | — | — | — | ax=None | — | save=None | (log_scale=False) | +| diag.archetypes | model | — | — | — | — | — | — | — | — | — | — | — | — | — | — | — | — | ax=None | — | save=None | — | +| pp.group_singletons | adata | — | — | — | **DIVERGENT default: groupby="patient"** | — | — | — | — | — | — | — | — | — | — | — | — | — | — | — | (clonotype_key="trb", target_col="trb_unique", min_clone_size=10) | +| pp.clone_size | adata | — | — | — | — | — | — | — | — | — | — | — | — | — | — | — | — | — | — | — | (key_added="clone_size", return_counts=False) | + +### Table B — distance / stats / contrast / RNG vocabulary + +| Function | primary(pos) | metric-selector | base / eps | by | value | reference | paired | hdi_prob | alternative | seed / random_state | resample count | pos_label | +|---|---|---|---|---|---|---|---|---|---|---|---|---| +| shared.phenotype_distance (priv) | p, q POS | **DIVERGENT: metric="l1"** | — | — | — | — | — | — | — | — | — | — | +| shared.kl_divergence (priv) | p, q POS | — | base=2.0, eps=1e-12 | — | — | — | — | — | — | — | — | — | +| shared.js_divergence (priv) | p, q POS | — | base=2.0, eps=1e-12 | — | — | — | — | — | — | — | — | — | +| shared.hdi (priv) | samples POS | — | — | — | — | — | — | hdi_prob=0.94 | — | — | — | — | +| shared.summarize (priv) | samples POS | — | — | — | — | — | — | hdi_prob=0.94 | — | — | — | — | +| shared.mann_whitney (priv) | a, b POS | — | — | — | — | — | — | — | alternative="two-sided" | — | — | — | +| tl.compare_groups | df POS | — | — | **by(req)** | value(req) | reference=None | paired=False | hdi_prob=0.94 | alternative="two-sided" | — | — | — | +| shared.auc_and_label_permutation (priv) | scores, labels POS | — | — | — | — | — | — | — | — | **seed=42** | n_perm=200_000 (max_exact=200_000) | pos_label=None | +| shared.bootstrap_auc (priv) | scores, labels POS | — | — | — | — | — | — | — | — | **seed=42** | n_boot=5000 | pos_label=None | +| diag.reconstruction_ppc | model POS | — | — | — | — | — | — | — | — | **seed=0** | n_samples=100 | — | +| diag.permutation_null | adata POS | — | — | — | — | — | — | — | — | **seed=0** | n_permutations=1000 | — | +| _compute._joint_draws (priv) | adata POS | — | — | — | — | — | — | — | — | random_state(kw) | n_samples(kw) | (gate_prob kw; **order/set differs from joint_distribution**) | + +### Consistency legend +- **RNG:** `random_state=None` (engine+tl+pl) vs `seed=42` (stats) vs `seed=0` (diag) — 3 spellings/defaults for one concept. +- **distance selector:** `distance_metric="l1"` (tl/pl flux, diag ppc) vs `metric="l1"` (`_distance.phenotype_distance`, private). +- **clones↔temperature order:** engine + `pl.phenotypic_flux` emit `clones` before `temperature`; every other metric emits `temperature` before `clones`. +- **normalize vs normalized:** `normalize=` (distribution bool) only on `pl.phenotypic_flux`; `normalized=` (metric [0,1] scaling bool) everywhere else — a near-collision. +- **device:** present on engine + all tl metrics; absent on every `pl.*` metric plot. +- **splitby:** present on the tl→pl box plots; absent on `pl.phenotypic_flux` and `pl.probability_ternary`. +- **figsize defaults:** (6,3) clonotypic_entropy/flux · (8,4) phenotypic_entropy/MI/_metric_boxplot · (5,5) ternary. +- **overloaded `splitby` (legacy):** on box plots the legacy `splitby` was the cohort-hue; on ternary/flux the legacy `splitby` was the covariate **column** paired with `conditions=`/`order=` **values** — two different concepts under one name. + +### Clean — verified consistent (no findings) +`temperature=1.0`; the `n_samples=0` point-estimate convention (single exception: `reconstruction_ppc`); `hdi_prob=0.94`; `alternative="two-sided"`; `use_logits` correctly engine-only; **no `point_estimate=` / `posterior=` survivors**; **no British-spelling residue** anywhere in the target surface. + +--- + +## 2. CONSISTENCY — confirmed inconsistencies + canonical decision + +Each item: the divergence → the single canonical name/order/default/type to adopt → the functions that change. Duplicate audit ids are merged. Private-symbol renames are marked *(private → optional/lower stakes)*. + +### High + +**H1 — RNG fragmentation → `random_state`** *(C1 ≡ S10, confirmed).* +`random_state=None` (engine/tl/pl) vs `seed=42` (`_stats`) vs `seed=0` (diag). One concept, three names/defaults/types. +**CANONICAL:** `random_state=None`, keyword-only, typed `int | numpy.Generator | torch.Generator | None`, placed as the **penultimate** compute arg (immediately before `device`) everywhere it appears. Drop the literal `42`/`0` defaults for `None` (seed internally from the passed generator). +**Load-bearing (public):** `diag.reconstruction_ppc`, `diag.permutation_null` (rename `seed`→`random_state`, default `None`). *(Private → optional):* `_stats.auc_and_label_permutation`, `_stats.bootstrap_auc` may keep `seed` internally but should be documented as the split. + +**H2 — `clones`/`temperature` order → temperature-before-clones** *(C2, confirmed; resolve with H8/C12).* +Engine emits `(n_samples, use_logits, clones, temperature)` and `pl.phenotypic_flux` emits `(clones, normalize, temperature)` — `clones` before `temperature` — while every other metric emits `(n_samples, temperature, clones)`. +**CANONICAL** shared block, left→right: `covariate`/`cov_*`, `groupby`, `splitby`, `n_samples`, `temperature`, `clones`, `use_logits` (engine-only, after `clones`), `normalized`/`normalize_mode`/`distance_metric`, `random_state`, `device`. Reorder `tl.joint_distribution` and `pl.phenotypic_flux` to `temperature`→`clones`. + +**H3 — `device` absent from every `pl.*` metric plot → add `device=None`** *(C3, revised).* +Confirmed absent on `pl.clonotypic_entropy`, `pl.phenotypic_entropy`, `pl.mutual_information`, `pl.phenotypic_flux`. Revision: on the three box plots the dangling knob is that they expose posterior draws + `random_state` but cannot steer the draw device; on `pl.phenotypic_flux` the finding's "draws-without-device" premise is wrong (flux exposes **no** `n_samples`), so its dangling knob is `random_state` — fix flux by adding **both** `device` and `n_samples` (see S-flux below). +**CANONICAL:** add `device=None` (`str | torch.device | None`, kw-only, penultimate before nothing/after `random_state`) to all four `pl.*` metric plots and thread it into the underlying `tl` call. + +**H4 — overloaded `splitby` on ternary/flux → `covariate_key` + `conditions`/`order`, NOT a cohort `splitby`** *(C4 revised ≡ M2 missed).* +`splitby` is genuinely absent from `pl.probability_ternary` and `pl.phenotypic_flux`, but the legacy notebooks passed `splitby=` there to name the covariate **column** (paired with `conditions=`/`order=` selecting the **values**), which is a *different* concept from the box-plot cohort-hue `splitby`. Adding a cohort-hue `splitby` here would re-introduce the exact name-collision the RULES forbid. +**CANONICAL:** reserve `splitby` strictly for the cohort/faceting hue. For ternary/flux add `covariate_key=None` (the covariate **column** override) plus `conditions=None` (ternary: 1–2 covariate levels → start/end simplices) / reuse `order` (Sankey: the ordered value series). Document that legacy `splitby=` maps to `covariate_key`, not `splitby`, in the Phase-10 notebook rewrite. +**Feasibility caveat (must be resolved in the contract):** the engine's joint is built from `p_ct` indexed by the trained covariate (`ct_to_cov`), so a `covariate_key` override is only implementable on an **empirical per-cell-probability aggregation path**, not the `p_ct` engine path. The contract must state which path multi-column flux/ternary use before `covariate_key` is added. + +### Medium + +**M1 — `n_samples=100` overloads the `n_samples=0` convention → rename to `n_sims`** *(C5, confirmed).* +`diag.reconstruction_ppc.n_samples=100` counts simulated PPC datasets, colliding with the surface-wide `n_samples=0`=point / `>0`=draws convention. +**CANONICAL:** rename to `n_sims=100` (`int`). Reserve `n_samples` (default `0`) exclusively for the posterior-draw convention. + +**M2 — distance selector `metric=` vs `distance_metric=`** *(C6, confirmed).* +`_distance.phenotype_distance(metric="l1")` (private dispatcher) vs `distance_metric="l1"` on `tl.phenotypic_flux`, `pl.phenotypic_flux`, `diag.joint_distribution_ppc`. +**CANONICAL:** `distance_metric` is the public name everywhere. Since `phenotype_distance` is private, keep its `metric` **only** as an internal dispatcher detail (documented) — no public churn required. + +**M3 — `normalize` vs `normalized` + `phenotype_colors` vs `palette` on `pl.phenotypic_flux`** *(C7, confirmed).* +`normalize=True` (distribution-normalization bool feeding `_phenotype_mass_per_clone`) look-alikes `normalized=True` (metric [0,1] scaling) used everywhere else; `phenotype_colors=None` duplicates `palette`. +**CANONICAL:** rename `phenotype_colors`→`palette`. Rename the flux bool to `normalize_distributions=True` (or drop it if the engine always feeds normalized distributions). Never ship both `normalize` and `normalized` as look-alike names. + +**M4 — `order`/`hue_order` inconsistent across the box-plot family → expose both on all three** *(C9 ≡ S2 ≡ S13, confirmed).* +`pl.clonotypic_entropy` has `hue_order` but no `order`; `pl.phenotypic_entropy` and `pl.mutual_information` have neither; the shared `_metric_boxplot` supports `order`; notebooks constantly pass `order=`. +**CANONICAL:** expose `order=None` (x-axis category order) **and** `hue_order=None` (`splitby` level order), both `list[str] | None`, on all three metric box plots, wired to `_metric_boxplot.order` and the `splitby` hue. + +**M5 — `compare_groups.by` names the cohort → rename `by`→`splitby`** *(C10, confirmed; defensible-either-way).* +`by` is the column whose levels are contrasted (with `reference`/`paired`) — semantically the comparison **cohort**, i.e. the plotting surface's `splitby`. +**CANONICAL:** rename `by`→`splitby` so the aggregation-unit(`groupby`)/comparison-cohort(`splitby`) vocabulary is uniform metric→plot→contrast. (`by` is pandas-idiomatic; if the team prefers `by`, that is a documented, deliberate exception rather than drift.) + +**M6 — `pp.group_singletons` defaults `clonotype_key="trb"`, `groupby="patient"` → keep, document as pre-registration defaults** *(C11, revised).* +Values confirmed, but the original "force `None`/match `unique_clone_id`" fix is wrong: `group_singletons` runs **before** `setup_anndata`, on the raw pre-registration column (`trb` → writes `target_col="trb_unique"`, which only later becomes the registered `clonotype_key`), so `unique_clone_id` doesn't exist yet; its `groupby` is the collapse **unit** (consistent with `groupby`=aggregation-unit), only the *default* differs. +**CANONICAL:** keep `clonotype_key="trb"` and `groupby="patient"` as intentional pre-registration defaults and **document** them as such; do not force `None`. + +**M7 — `_joint_draws` keyword order/set diverges from `joint_distribution`** *(C12, confirmed; private → lower stakes).* +`_joint_draws(covariate, clones, n_samples, use_logits, temperature, gate_prob, …)` vs `joint_distribution(covariate, groupby, n_samples, use_logits, clones, temperature, …)`. +**CANONICAL:** make `_joint_draws`' keyword order a strict subset-in-order of the public engine: `(covariate, n_samples, temperature, clones, use_logits, gate_prob, random_state, device)` — i.e. temperature-before-clones (dovetails H2). Resolve `gate_prob` at the public layer or document it as an internal-only extra. + +**M8 — `joint_distribution_ppc` arg-order + missing knobs** *(C13 ~ S15, confirmed).* +Orders `(covariate, distance_metric, temperature)` (temperature after the selector) and omits `groupby`/`clones` (and `n_samples`/`random_state`) that its metric siblings expose. +**CANONICAL:** reorder to `(covariate, groupby, clones, temperature, distance_metric)`; add `groupby=None`, `clones=None`. It is model-free/deterministic, so add `n_samples=0`/`random_state=None` **only if** a draw-based comparison is intended (flag as a contract decision, not a default addition). + +**M9 — `covariate=None` dual meaning in the `adata_or_jd` fast path** *(C19 ≡ S14, confirmed).* +All signatures default `covariate=None`, but §7.2(c)/§7.9 let the precomputed-`jd` path read `covariate` as "required/the one baked into this jd," contradicting the RULE that `covariate=None`=all covariate values. +**CANONICAL:** `covariate=None` **always** means all covariate values on the adata path. On the precomputed-`jd` fast path, `covariate`/`n_samples`/`temperature`/`clones`/`random_state`/`device` are inert — **raise** a clear `ValueError` if any is set to a non-default — so `covariate=None`=all is preserved only where it is computed. Fix the entropy/MI docstrings and dispatch so `covariate` is never treated as required. + +### Low + +**L1 — figsize drift across box-plot siblings** *(C8, confirmed).* `pl.clonotypic_entropy=(6,3)` vs `(8,4)` for `pl.phenotypic_entropy`/`pl.mutual_information`/`_metric_boxplot`. +**CANONICAL:** `figsize=(8,4)` for the metric box-plot family; set `pl.clonotypic_entropy` to `(8,4)`. Aspect-driven plots keep purpose-specific defaults, documented as intentional: ternary `(5,5)`, Sankey flux `(6,3)`. + +**L2 — 2nd-positional violations → keyword-only** *(C14, confirmed).* `pl.resolve_palette(adata, columns)`, `pl._metric_boxplot(adata, function)`, `diag.reconstruction_ppc(model, adata=None)` each carry a second positional. +**CANONICAL:** insert `*` after the single primary positional so `columns`, `function`, and `adata` are keyword-only. *Census note:* `_phenotype_mass_per_clone(adata, covariate, clones, normalize)` and `TCRIModule.prepare_two_level_params` also carry extra positionals (private, minor) — align in the same pass. + +**L3 — resample-count `n_permutations` vs `n_perm`** *(C15, confirmed).* `diag.permutation_null.n_permutations=1000` vs `_stats.auc_and_label_permutation.n_perm=200_000`. +**CANONICAL:** `n_perm` for permutation counts everywhere; rename `permutation_null.n_permutations`→`n_perm` (keep its `1000` default). `n_boot=5000` (bootstrap) is a distinct concept and keeps its name. + +**L4 — `show=` only on private `_finish`** *(C16, confirmed).* `_finish(show=None)` is unreachable from any public plot, though all expose `save=`/`return_*`. +**CANONICAL:** adopt the scanpy triad uniformly — thread `show=None` (`bool | None`, kw-only) through **every** public `pl.*` entry point (with `save`/`return_*`), keeping it on `_finish`. + +**L5 — return-control name split** *(C17, confirmed; near-self-resolving).* `return_df` on DataFrame plots vs `return_axes` on figure plots. +**CANONICAL:** keep the split (payloads genuinely differ) but **standardize which every plot exposes** — DataFrame/metric plots → `return_df`; figure-only plots (ternary, Sankey) → `return_axes` — and document it. No plot may omit both (all currently comply). + +**L6 — `eps` clip-floor drift + internal contradiction** *(C18 confirmed ≡ M3-missed strengthened).* Four floors: `1e-8` (`TCRIModel.predict`, `_compute_logits_and_prior`), `1e-6` (`TCRIModule.prepare_two_level_params`), `1e-12` (`kl_divergence`, `js_divergence`), `1e-15` (hardcoded in `_mi_from_joint`). Sharper than drift: §3.4 and appendix row 12 both assert "one ε=10⁻¹² **library-wide**, matching entropy/MI," yet the MI kernel — the exact metric that claim name-checks — uses `1e-15`. The document contradicts its own stated invariant. +**CANONICAL:** set `_mi_from_joint` `eps=1e-12` to satisfy the stated invariant (single probability-clip floor `1e-12` for the distance/MI/normalization paths). Keep context-specific floors (`predict`/`prepare_two_level_params`) only where numerically justified, and document why. Do not leave prose and kernel contradicting. + +### Documentation / cross-document consistency (verifier's missed items) + +**D1 — `§11` and the implementation plan disagree on live callers** *(M1, missed).* API §11 line 770 buckets `clonality` (tl+pl), `probability_distribution`, and `clone_fraction` under "Deleted … 0 live callers after census," but the plan (and the notebooks) show live callers: `tcri.pl.clonality` is called with full args (`groupby`/`splitby`/`order`/`palette`/`figsize`) across smith/renal/comparison/yost/zhang; `tcri.metrics.probability_distribution` is imported and called 4× via the alias `pdistribution(psubset, method=…)`. `compare_phenotypes` (1 live call in synthetic1) is dispositioned in the plan (DROP) but is **absent from §11 entirely**. +**CANONICAL:** reconcile §11 with plan §3 — move `clonality`, `probability_distribution`, `clone_fraction` out of the "0 live callers" bucket into a "removed **with** replacement + in-PR notebook rewrite (Phase 6/10)" category; add `compare_phenotypes` to §11; and correct the false "0 live callers after census" label (the census missed alias-imported and package-qualified call-sites). + +**D2 — undispositioned legacy renames `method=` and `phenotype_names=`** *(M4, missed).* `method=` (e.g. `method="probabilistic"`) is passed to `pl.clonotypic_entropy`, `pl.mutual_information`, `pl.flux`, `polar_plot`, and `pdistribution`, but only `posterior=`/`point_estimate=` are mapped in §11 — `method=` is never named. `phenotype_names=` is the legacy positional/keyword on `probability_ternary`, and the kept §8.7 signature renames it to `phenotypes` with no rename entry. +**CANONICAL:** add disposition-map rows to §11: `method=` → removed, expressed via `n_samples`/`use_logits`; `phenotype_names` → `phenotypes` on `pl.probability_ternary`. + +**D3 — caller-census counts are unreliable** *(M5, missed).* `polar_plot` is recorded as "1 notebook use," but the notebooks contain 4 local `def polar_plot` redefinitions plus multiple call-sites; the same def-plus-call pattern holds for `clonality`. Neither audit disambiguated def-vs-call. +**CANONICAL:** re-run the census distinguishing (a) package-qualified calls, (b) alias-imported calls, (c) notebook-local redefinitions; record accurate per-symbol counts; the Phase-10 rewrite must also strip shadow `def`s, not just swap call-sites. + +**D4 — `pl.flux` box-plot disposition unstated** *(S9, revised).* The old per-clone flux-**distance** box-plot `pl.flux(..., paint=, distance_metric='dkl', ...)` is distinct from the Sankey and is used in renal/yost/zhang, but §11 only renames `flux`→`tl.phenotypic_flux` and never states this box-plot's fate. The plan settles it as DROP (the Sankey is the flux plot); re-adding a box variant would contradict that decision. +**CANONICAL:** state explicitly in §11 that the `pl.flux` distance box-plot is **DROP** (Sankey is the flux plot) — a doc-completeness fix, not a new function. + +--- + +## 3. SUFFICIENCY — confirmed missing knobs (exact argument to add) + +Every confirmed gap, with the exact `name=default` (type) and the functions that receive it. Refuted gaps (S3/S8/S12) are excluded per §Dropped. + +| # | Argument to add | Default | Type | Functions receiving it | Audit id / status | +|---|---|---|---|---|---| +| U1 | `phenotypes` (rename from `phenotype_names`) | (req) | `list[str]` (3 axes) | `pl.probability_ternary` | S1/M4 confirmed | +| U2 | `conditions` | `None` | `list[str] \| None` (1–2 covariate levels → start/end) | `pl.probability_ternary` | S1/M2 confirmed | +| U3 | `scale_function` | `None` | `Callable[[float], float] \| None` (freq → marker size) | `pl.probability_ternary` | S1 confirmed (exercised in every zhang/yost ternary call) | +| U4 | `color` | `None` | `str \| None` | `pl.probability_ternary` | S1 confirmed | +| U5 | `covariate` | `None` | `str \| None` (value; `None`=all) | `pl.probability_ternary` | S1 confirmed (lower priority — historical default, not passed in notebooks) | +| U6 | `top_n` | `None` | `int \| None` | `pl.probability_ternary` | S1 confirmed (lower priority) | +| U7 | `n_samples` | `0` | `int` | `pl.probability_ternary`, **`pl.phenotypic_flux`** | S1, S5 confirmed | +| U8 | `temperature` | `1.0` | `float` | `pl.probability_ternary` | S1 confirmed (lower priority) | +| U9 | `order` | `None` | `list[str] \| None` (x-axis category order) | `pl.clonotypic_entropy`, `pl.phenotypic_entropy`, `pl.mutual_information` | S2 ≡ M4/C9 confirmed | +| U10 | `hue_order` | `None` | `list[str] \| None` (`splitby` level order) | `pl.phenotypic_entropy`, `pl.mutual_information` (already on `pl.clonotypic_entropy`) | S13 ≡ C9 confirmed | +| U11 | `device` | `None` | `str \| torch.device \| None` | `pl.clonotypic_entropy`, `pl.phenotypic_entropy`, `pl.mutual_information`, `pl.phenotypic_flux` | C3 confirmed/revised | +| U12 | `minimum_clone_size` | `None` | `int \| None` (filter `clone_size` before aggregation) | `pl.clonotypic_entropy`, `pl.phenotypic_entropy`, `pl.mutual_information` (optionally the `tl` metrics) | S6 confirmed | +| U13 | `phenotype_subset` | `None` | `list[str] \| None` (restrict rendered phenotype nodes) | `pl.phenotypic_flux` | S5 confirmed | +| U14 | `pair_on` | `None` | `str \| None` (unit column aligning the two levels under `paired=True`) | `tl.compare_groups` | S7 confirmed (closes the incomplete `delta_entropy_table` subsumption) | +| U15 | `covariate_key` | `None` | `str \| None` (covariate **column** override) | `tl.joint_distribution`, `tl.phenotypic_flux`, `pl.phenotypic_flux`, `pl.probability_ternary` | S4/M2 revised — **conditional** on the empirical-vs-`p_ct` path decision (H4 caveat) | +| U16 | metric pass-throughs: `temperature=1.0` (`float`), `normalized=True` (`bool`), `normalize_mode="min"` (`str`), `clones=None` (`list[str] \| None`), `splitby=None` (`str \| None`) | as noted | as noted | `diag.permutation_null` (so the null matches the reported statistic) | S11 confirmed | +| U17 | `clones` | `None` | `list[str] \| None` | `diag.joint_distribution_ppc`, `diag.permutation_null` | S15 ≡ C13 confirmed | +| U18 | `groupby` | `None` | `str \| None` | `diag.joint_distribution_ppc` | S15 ≡ C13 confirmed | +| U19 | `show` | `None` | `bool \| None` | all public `pl.*` plots (scanpy triad) | C16 confirmed | + +**Not a new knob, but tied to sufficiency:** `tl.compare_groups` currently reproduces one-axis contrasts but not the paired-**delta-then-cohort-contrast** of `delta_entropy_table` in one call — `paired=True` aligns by `sample_id` (posterior-draw pairing), not pre/post-within-unit pairing, and there is no second cohort axis. `pair_on` (U14) supplies the within-unit pairing key; the §11 claim that `delta_entropy_table` is "expressed via `groupby` + `compare_groups`" must be softened to "with `pair_on` + an explicit cohort split" or documented as a two-pass recipe. + +--- + +## 4. CORRECTED SIGNATURES + +Final signatures for every function that changes. Deltas from the current contract are noted after each block. `# †` marks the conditional `covariate_key` (gated on the H4 empirical-path decision). + +**`tl.joint_distribution`** — reorder to temperature-before-clones (`use_logits` engine-only, after `clones`). +```python +joint_distribution( + adata, *, + covariate=None, # None → ALL covariate values (one shared draw) + groupby=None, + n_samples=0, + temperature=1.0, + clones=None, + use_logits=True, # engine-only; was posterior=; alias cell_informed= + random_state=None, # int | numpy.Generator | torch.Generator | None + device=None, +) -> pandas.DataFrame +``` + +**`tl.compare_groups`** — `by`→`splitby`; add `pair_on`. +```python +compare_groups( + df, *, + value, + splitby, # was by= + reference=None, + paired=False, + pair_on=None, # NEW: within-unit pairing key when paired=True + hdi_prob=0.94, + alternative="two-sided", +) -> pandas.DataFrame +``` + +**`tl.phenotypic_flux`** — add optional `covariate_key`. +```python +phenotypic_flux( + adata, *, + covariate_key=None, # † covariate COLUMN override (H4 caveat) + cov_from, cov_to, + groupby=None, + n_samples=0, + temperature=1.0, + clones=None, + distance_metric="l1", + random_state=None, + device=None, +) -> pandas.DataFrame +``` + +**`pl.clonotypic_entropy`** — figsize (6,3)→(8,4); add `order`, `minimum_clone_size`, `device`, `show` (`hue_order` already present). +```python +clonotypic_entropy( + adata, *, + covariate=None, groupby=None, splitby=None, + n_samples=0, temperature=1.0, clones=None, + normalized=True, n_clones_ref=None, minimum_clone_size=None, + order=None, hue_order=None, + palette=None, ax=None, figsize=(8, 4), + rotation=90, legend_fontsize=6, bbox_to_anchor=(1.15, 1.0), + random_state=None, device=None, + show=None, save=None, return_df=False, +) +``` + +**`pl.phenotypic_entropy`** — add `order`, `hue_order`, `minimum_clone_size`, `device`, `show`. +```python +phenotypic_entropy( + adata, *, + covariate=None, groupby=None, splitby=None, + n_samples=0, temperature=1.0, clones=None, + normalized=True, minimum_clone_size=None, + order=None, hue_order=None, + palette=None, ax=None, figsize=(8, 4), + rotation=90, legend_fontsize=6, bbox_to_anchor=(1.15, 1.0), + random_state=None, device=None, + show=None, save=None, return_df=False, +) +``` + +**`pl.mutual_information`** — add `order`, `hue_order`, `minimum_clone_size`, `device`, `show`. +```python +mutual_information( + adata, *, + covariate=None, groupby=None, splitby=None, + n_samples=0, temperature=1.0, clones=None, + normalized=True, normalize_mode="min", minimum_clone_size=None, + order=None, hue_order=None, + palette=None, ax=None, figsize=(8, 4), rotation=90, + legend_fontsize=6, bbox_to_anchor=(1.15, 1.0), + random_state=None, device=None, + show=None, save=None, return_df=False, +) +``` + +**`pl.phenotypic_flux`** — reorder temperature-before-clones; `normalize`→`normalize_distributions`; `phenotype_colors`→`palette`; add `covariate_key`, `n_samples`, `phenotype_subset`, `device`, `show`. +```python +phenotypic_flux( + adata, *, + covariate_key=None, # † covariate COLUMN override (H4 caveat) + order, # ordered covariate value series + groupby=None, + n_samples=0, # NEW: drives ribbon uncertainty (feeds random_state) + temperature=1.0, + clones=None, + normalize_distributions=True, # renamed from normalize= + distance_metric="l1", + phenotype_subset=None, # NEW: restrict rendered phenotype nodes + palette=None, # renamed from phenotype_colors= + ax=None, figsize=(6, 3), + show_legend=True, title=None, + random_state=None, device=None, + show=None, save=None, return_axes=False, +) +``` + +**`pl.probability_ternary`** — `phenotype_names`→`phenotypes`; add `covariate_key`, `covariate`, `conditions`, `n_samples`, `temperature`, `scale_function`, `top_n`, `color`, `random_state`, `show`. No cohort-hue `splitby` (H4). +```python +probability_ternary( + adata, *, + phenotypes, # renamed from phenotype_names; the 3 simplex axes + covariate_key=None, # † covariate COLUMN override (H4 caveat) + covariate=None, # covariate VALUE (None → all) + conditions=None, # 1–2 covariate levels → start/end simplices + groupby=None, clones=None, + n_samples=0, temperature=1.0, + scale_function=None, # freq → marker size + top_n=None, color=None, + palette=None, ax=None, figsize=(5, 5), + random_state=None, + show=None, save=None, return_axes=False, +) +``` + +**`diag.joint_distribution_ppc`** — reorder temperature-before-`distance_metric`; add `groupby`, `clones`. +```python +joint_distribution_ppc( + adata, *, + covariate=None, + groupby=None, + clones=None, + temperature=1.0, + distance_metric="l1", + # add n_samples=0 / random_state=None ONLY if draw-based comparison is adopted +) -> pandas.DataFrame +``` + +**`diag.reconstruction_ppc`** — `adata` keyword-only; `n_samples`→`n_sims`; `seed`→`random_state` (default `None`). +```python +reconstruction_ppc( + model, *, + adata=None, + n_sims=100, # was n_samples=100 + random_state=None, # was seed=0 +) -> pandas.DataFrame +``` + +**`diag.permutation_null`** — `n_permutations`→`n_perm`; `seed`→`random_state`; add metric pass-throughs + `clones`. +```python +permutation_null( + adata, *, + metric="mutual_information", + covariate=None, groupby=None, splitby=None, clones=None, + temperature=1.0, normalized=True, normalize_mode="min", + n_perm=1000, # was n_permutations=1000 + random_state=None, # was seed=0 +) -> pandas.DataFrame +``` + +**`pl.resolve_palette`** — `columns` keyword-only. +```python +resolve_palette(adata, *, columns, palette=None) -> dict +``` + +**Private helpers (align in the same pass, lower stakes):** +```python +_metric_boxplot(adata, *, function, groupby=None, splitby=None, + ylabel="", order=None, hue_order=None, + palette=None, s=20, ax=None, figsize=(8, 4)) -> (fig, ax) # function kw-only; add hue_order +_joint_draws(adata, *, covariate=None, n_samples=0, temperature=1.0, + clones=None, use_logits=True, gate_prob=None, + random_state=None, device=None) # subset-in-order of joint_distribution +# _stats.auc_and_label_permutation / bootstrap_auc: optionally seed→random_state (private; document the split) +# _distance.phenotype_distance: keep metric= as an internal dispatcher detail (no public change) +``` +Also: `_mi_from_joint` — change hardcoded `eps=1e-15`→`1e-12` (L6). `pp.group_singletons` / `pp.clone_size` — **no signature change** (M6: defaults kept and documented as intentional pre-registration values). + +--- + +## 5. READY-TO-APPLY PATCH LIST — edits to fold into `docs/contract/tcri_api_and_responsibilities.md` + +1. **§7.1 engine** — reorder the `joint_distribution` block to `covariate, groupby, n_samples, temperature, clones, use_logits, random_state, device` (temperature before clones; `use_logits` after `clones`). Update the §7.1(c) arguments table row order to match. +2. **§7.6 `compare_groups`** — rename `by`→`splitby` (signature + the "grouping column" comment + Math/Return prose); add `pair_on=None` between `paired` and `hdi_prob`; soften the "subsumes `delta_entropy_table`" claim to require `pair_on` (+ a cohort split) or document the two-pass recipe. +3. **§7.5 `tl.phenotypic_flux`** — add optional `covariate_key=None` (gated on the H4 empirical-path decision; add a one-line note stating the chosen path). +4. **§7.2 / §7.9** — pin `covariate=None`=all covariate values on the adata path; state that the precomputed-`jd` fast path **ignores** `covariate`/`n_samples`/`temperature`/`clones`/`random_state`/`device` and **raises** if any is non-default (remove any "covariate required" reading). +5. **§8.1 `pl.clonotypic_entropy`** — `figsize=(6,3)`→`(8,4)`; add `order=None`, `minimum_clone_size=None`, `device=None`, `show=None`. +6. **§8.1 `pl.phenotypic_entropy`** — add `order=None`, `hue_order=None`, `minimum_clone_size=None`, `device=None`, `show=None`. +7. **§8.2 `pl.mutual_information`** — add `order=None`, `hue_order=None`, `minimum_clone_size=None`, `device=None`, `show=None`. +8. **§8.3 `pl.phenotypic_flux`** — reorder temperature-before-clones; rename `normalize`→`normalize_distributions` and `phenotype_colors`→`palette`; add `covariate_key=None` (†), `n_samples=0`, `phenotype_subset=None`, `device=None`, `show=None`; update the Sankey prose to reference `n_samples`-driven ribbons. +9. **§8.7 `pl.probability_ternary`** — rename `phenotype_names`→`phenotypes`; add `covariate_key=None` (†), `covariate=None`, `conditions=None`, `n_samples=0`, `temperature=1.0`, `scale_function=None`, `top_n=None`, `color=None`, `random_state=None`, `show=None`; do **not** add a cohort `splitby`. +10. **§8.5** — make `_metric_boxplot`'s `function` keyword-only; add `hue_order=None`. **§8.6** — make `resolve_palette`'s `columns` keyword-only. +11. **§8.5 `_finish` / all §8 public plots** — thread `show=None` through every public `pl.*` entry point (scanpy `show`/`save`/`return_*` triad); document the `return_df` (DataFrame) vs `return_axes` (figure) split so no plot omits both. +12. **§9.1 `joint_distribution_ppc`** — reorder to `(covariate, groupby, clones, temperature, distance_metric)`; add `groupby=None`, `clones=None`; note `n_samples`/`random_state` deferred unless draw-based comparison is adopted. +13. **§9.1 `reconstruction_ppc`** — `adata` keyword-only; `n_samples`→`n_sims`; `seed`→`random_state` (default `None`). +14. **§9.1 `permutation_null`** — `n_permutations`→`n_perm`; `seed`→`random_state`; add `splitby=None`, `clones=None`, `temperature=1.0`, `normalized=True`, `normalize_mode="min"` pass-throughs. +15. **§3.3** — optionally rename `_stats.auc_and_label_permutation`/`bootstrap_auc` `seed`→`random_state`, or document that the private primitives keep `seed` while the public diagnostics standardize on `random_state`. +16. **§3.4 / §7.4 / appendix row 12** — set `_mi_from_joint` `eps` to `1e-12`, OR amend the "one ε=10⁻¹² library-wide, matching entropy/MI" claim to carve out the MI-specific `1e-15` with a numerical justification. Resolve the prose↔kernel contradiction either way. +17. **§0.10** — state the canonical RNG type once: `random_state: int | numpy.Generator | torch.Generator | None`, keyword-only, penultimate before `device`; note the `_stats` `seed` exception if kept. +18. **§6.1 `pp.group_singletons`** — add a note that `clonotype_key="trb"` and `groupby="patient"` are intentional pre-registration defaults (operate on the raw column before `setup_anndata`); do not change them. +19. **§11 census (correctness fix)** — move `clonality` (tl+pl), `probability_distribution`, `clone_fraction` **out** of "Deleted … 0 live callers" into a "removed **with** replacement + in-PR notebook rewrite (Phase 6/10)" category; **add `compare_phenotypes`** (currently absent) as DROP-with-rewrite; correct the false "0 live callers" label. Add an explicit disposition line for the **`pl.flux` distance box-plot: DROP** (Sankey is the flux plot). +20. **§11 renamed/removed map** — add rows: legacy `method=` → removed (expressed via `n_samples`/`use_logits`); `phenotype_names` → `phenotypes` on `pl.probability_ternary`. Note that the caller census must distinguish package-qualified calls, alias-imported calls, and notebook-local `def` redefinitions (`polar_plot`, `clonality`), and that Phase-10 must strip shadow definitions. +21. **§0.11 (new, cross-cutting note)** — record the H4 open decision: because `p_ct` is covariate-locked (`ct_to_cov`), any `covariate_key` column override on flux/ternary must be computed on the **empirical per-cell-probability path**, not the `p_ct` engine path; the contract must state which path these plots use before `covariate_key` (†) ships. \ No newline at end of file diff --git a/docs/contract/_quarantine/tcri_implementation_plan.md b/docs/contract/_quarantine/tcri_implementation_plan.md new file mode 100644 index 0000000..39a8b33 --- /dev/null +++ b/docs/contract/_quarantine/tcri_implementation_plan.md @@ -0,0 +1,466 @@ +# TCRI Refactor — Implementation Plan (Final) + +**From** the current five-monolith package **to** the settled Door‑A, scverse‑ecosystem target (grafiti‑mirrored, one‑file‑per‑topic), with every audit fix folded in. + +**Source of truth.** The *SETTLED DESIGN* block is authoritative; the *prior‑vs‑mean RESOLUTION* settles the one parked statistical question and is adopted verbatim. Where the settled design diverges from `docs/contract/tcri_function_inventory.md` §3/§9 (engine lives in `tl` not `pp`; `groupby` subsumes all `*_table`/`*_compare`/`*_delta`; `flux → phenotypic_flux`; `register_model → model.to_anndata`; `tl` mirrors grafiti's `tools/`), this plan follows the settled design and flags the divergence inline. This document supersedes the earlier draft wherever the audit corrected it; the corrections are integrated at the point they bite, and every finding is cross‑referenced in **Appendix A** so nothing is lost. + +**How to read.** §1 fixes the invariants and records the resolved decisions that unblock the engine/metric phases. §2 is the target tree. §3–4 are the rename/disposition map and the shared‑helper extraction. §5 is the model→AnnData streamline. §6 is the engine + metric **numeric contract** (all math/stats fixes live here). §7 is the GPU/optimization architecture, grounded in grafiti. §8 is the ordered, independently‑landable PR sequence. §9 is testing + scverse‑CI. §10 is the ordering‑hazard graph. §11 is risks & open items. §12 is the per‑PR checklist. + +--- + +## 1. Invariants and resolved decisions + +### 1.1 Invariants held at every PR boundary +- **`import tcri` stays green** and the public handles `tl / pp / pl / ml / diag / ut` remain importable. No PR leaves `main` with a broken import or red CI. **Corollary (audit):** "0‑caller" means *no call‑sites **and** no import‑sites, in package **and** in `example/`+`docs/` notebooks* — a symbol imported at module top (e.g. `utils.probabilities` at `_plotting.py:18`) is a live dependency even with zero calls. +- **One behavior change per PR.** Mechanical moves (splits, helper extraction) never change numerics; numeric changes (engine, metrics) never also move files. Every diff stays reviewable and every regression bisectable. +- **`_keys.py` is the only place a key string is written**, from Phase 1 on. No new `uns/obsm/obs` key literal may appear outside it. +- **No `import *`.** Each split lands with an explicit `__all__` and named re‑exports; the top‑level `__init__` flip is the last PR. +- **`setup_anndata` performs no *analysis/label* obs mutation.** *(Corrected from the draft's "no obs mutation.")* It must still write and register the `obs['indices']` field the training step consumes (`batch['indices']`, `_model.py:603/641/678/684`). The invariant forbids writing *results* (probabilities, hard labels, latent) into `obs`, not the registration glue column. +- **GPU is never on the import path and never in `install_requires`.** Every accelerated path is pure opt‑in behind a device seam with a fully‑functional CPU fallback (§7). + +### 1.2 Resolved decisions (adopted from the prior‑vs‑mean RESOLUTION + math/stats audit) + +**(R1) Point estimate = closed‑form variational posterior mean.** `n_samples=0` returns `E_q[p_ct] = normalize(q_p_ct_raw) = TCRIModel.get_p_ct() = adata.uns[K.P_CT]`, read directly, **never sampled**. Because the guide is `Dirichlet(local_scale · m)` with `m` on the simplex, `Σα = local_scale` and the mean is exactly `m` — `local_scale` cancels. The **prior/archetype path is dropped** (it is guide‑init/generative anchor built from leaked hard labels, not what training learned); the **MAP/mode is rejected** (`α_k = local_scale·m_k` is routinely `< 1`, so the mode sits on the simplex boundary / is undefined); **mean‑of‑draws is rejected** (Rao‑Blackwell‑dominated, adds only MC noise to a closed‑form quantity). + +**(R2) Rename the mis‑named `posterior=` axis to `use_logits` — a classifier‑mixing switch, not a prior/posterior switch.** There was never a live prior branch; what actually differs is *with logits vs without*. The **engine** `joint_distribution` keeps exactly one such flag, `use_logits` (alias `cell_informed`), replacing both `posterior=` and `combine_with_logits=`: + - `use_logits=True` (default) folds per‑cell classifier logits into `log(base)` exactly like `model.predict` (gate‑aware, §5). + - `use_logits=False` returns the ct‑level table directly. + Both branches use the posterior **mean** (`n_samples=0`) or **draws** (`n_samples>0`) of `p_ct`; neither ever touches the generative prior. **The four metrics do NOT expose `use_logits`, `posterior=`, or `point_estimate=`** — given an `adata` they always compute the cell‑informed joint. `point_estimate=` is deleted. + +**(R3) `n_samples=0` is a *plug‑in* estimator, not the posterior mean of the metric.** Entropy, MI, and l1/KL flux are **nonlinear** functionals, so `metric(E_q[p]) ≠ E_q[metric(p)]` by a **Jensen gap that is not Monte‑Carlo noise**: entropy (concave) plug‑in **over‑estimates**, flux (convex) plug‑in **under‑estimates**, MI is sign‑indeterminate. Therefore: + - Document `n_samples=0` as **`metric‑at‑posterior‑mean`** and the `n_samples>0` `mean` column as **`E_q[metric]`** — two different, clearly‑labeled estimators. + - **No test may assert `n_samples=0 == mean(n_samples>0)`.** The determinism/reproducibility guarantee (bit‑identical repeated `n_samples=0` calls) still holds and is tested; the *equality across estimators* claim is dropped. + +**(R4) `n_samples>0` draws must reproduce the guide's clamped concentration.** Draw from `Dirichlet(clamp(local_scale · m̃, min=1e-3))`, reusing the guide's floor (`_model.py:490`). The three inconsistent current variants (`local_scale·m` unclamped; `local_scale·m + 1e-8`; the clamped guide form) are unified onto the guide form so reported HDIs summarize the distribution the model actually learned. Note that for committed clones (`m_k < 1e-3/local_scale`) the clamp makes the draw‑mean differ slightly from `m`; this is documented, and `n_samples=0` remains the closed‑form `m` by definition. + +**(R5) `p_gt` (P(>0)) is only meaningful on a signed contrast.** Entropy, MI, and flux are all `≥ 0`, so `P(draw > 0) ≈ 1` and is vacuous. **Remove `p_gt` from the single‑metric `n_samples>0` summary** (emit `mean`, `sd`, `hdi_low`, `hdi_high` only). `p_gt` is computed **only** on a paired between‑group difference vector by the comparison helper (§4, §6). + +**(R6) Temperature is a single analysis knob applied one consistent way.** `m̃ = softmax(log(m + 1e-8) / T)` (identity at `T=1`), applied identically in the mean and draw branches. To avoid double‑tempering, **`to_anndata` persists the *raw* posterior mean** `m = normalize(q_p_ct_raw)` (not the guide‑temperature‑adjusted vector), so analysis `temperature` is the sole tempering knob. At `T=1`, `use_logits=True` reproduces `model.predict` exactly. `temperature≠1` re‑centers the sampled distribution away from the learned posterior — documented in the metric docstrings. + +**(R7) Reproducibility is via a seeded torch generator.** All draws move to a seeded `torch.Generator` (device‑aware; also seed cupy/torch.cuda RNG on GPU). Add `random_state: int | torch.Generator | None` to `joint_distribution` and every sampling metric/diag function; thread it into the Dirichlet draw. `np.random.seed` alone was a **no‑op** for the torch draws and is retired. `n_samples=0` is deterministic regardless. + +**(R8) `use_logits=True` parity with `predict` requires three model facts persisted to `adata`.** The gate‑aware combination `gate_prob·logits + (1-gate_prob)·log(base)` and the classifier temperature are **model attributes**. `to_anndata` must persist `obsm[K.X_LOGITS]`, `uns[K.GATE_PROB]`, and `uns[K.CLASSIFIER_TEMPERATURE]` (§5). Without them the engine can only reproduce the additive rule — the exact disagreement the refactor set out to remove. + +--- + +## 2. Target layout (grafiti‑mirrored, authoritative) + +``` +tcri/ + __init__.py # explicit re-export; sys.modules aliases tl/pp/pl/ml/diag/ut; + # top-level tcri.joint_distribution; NO import * + _keys.py # single source of every uns/obsm/obs key string (adopt in Phase 1) + # NOW INCLUDES: X_LOGITS, GATE_PROB, CLASSIFIER_TEMPERATURE, + # LOCAL_SCALE, P_CT, X_TCRI, X_PROBABILITIES, CLONE_COL/PHENO_COL + _console.py # leveled, silenceable logging over scanpy logging; NO ANSI, NO _ascii_hist + _stats.py # stars, auc_and_label_permutation, bootstrap_auc, + # + posterior-comparison primitives: mann_whitney, prob_gt_zero, hdi (TRUE HDI) + _distance.py # kl_divergence (log2/bits), l1_distance, jensen_shannon, phenotype_distance dispatcher + _compute/ # NEW — device seam + batched numeric core (grafiti/_compute parity) + _xp.py # resolve_device / get_xp / asnumpy (torch-first, cupy optional, cpu default) + _joint.py # _joint_draws(adata, covariate, n_samples, *, use_logits, temperature, device, + # random_state) -> ndarray[n_samples, n_clones, P] (scatter-add reduction) + _reduce.py # batched entropy / mutual_information / flux over the [S, n_clones, P] stack + _embedding.py # umap() behind _use_gpu gate (cuML on GPU, umap-learn CPU); lazy import + model/ # ml + _model.py # TCRIModel: setup_anndata, train, get_latent_representation, + # predict (was get_cell_phenotype_probs), get_p_ct, to_anndata + _module.py # TCRIModule (pyro model/guide) + _priors.py # MixtureDirichlet, VampPrior + _classifier.py # PhenotypeClassifier + _training.py # UnifiedTrainingPlan, build_archetypes (returns centers AND labels) + preprocessing/ # pp (shrinks to clone utilities; engine moved to tools/) + _clones.py # group_singletons (must precede setup_anndata), clone_size + tools/ # tl (metrics + engine; mirrors grafiti tools/) + _joint.py # joint_distribution (thin DataFrame wrapper over _compute._joint; + # re-exported as tcri.joint_distribution) + _entropy.py # clonotypic_entropy, phenotypic_entropy + _mutual_information.py # mutual_information (+ private _mi_from_joint) + _flux.py # phenotypic_flux (was flux; cov_from / cov_to) + _compare.py # compare_groups (PUBLIC mid-level stats helper; replaces *_compare/*_delta) + plotting/ # pl (twins mirror tl by filename) + _base.py # _metric_boxplot (was tcri_boxplot; groupby + splitby), _finish + _colors.py # tcri_colors, resolve_palette (was set_color_palette) + _entropy.py # clonotypic_entropy [FIX], phenotypic_entropy [FIX] + _mutual_information.py # mutual_information [FIX] + _flux.py # phenotypic_flux (sankey) + _sankey.py # SankeyNode, _phenotype_mass_per_clone + diagnostics/ # diag (NEW — PPCs + model validation; returns DataFrames) + _ppc.py # joint-distribution PPC (fixed compare_joint_distribution), + # calibration, reconstruction PPC (model-required), permutation-null + _training.py # loss curves (was plot_loss), archetypes (was plot_archetypes) + utils/ # ut + _session.py # save/load_tcri_session, _to_jsonable (plain h5ad; no manager hack) +examples/ # top_clone_umap, clone_size_umap, phenotype_probabilities, + # probability_ternary, gene_entropy (relocated), rewritten notebooks +docs/ # model PGM (build_nested_tcri_pgm) lives here, out of the package +``` + +Divergences from inventory §3, called out: **(a)** `tl` package is `tools/` (grafiti parity), not `metrics/`; **(b)** the engine's *numeric core* is `_compute/_joint.py` with a thin `tools/_joint.py` DataFrame wrapper — **not** `preprocessing/_engine.py`; **(c)** `preprocessing/` loses `_register.py`/`_engine.py` (registration collapses into `model.to_anndata`), leaving `pp` = `_clones.py`; **(d)** no `_tables.py` — `groupby` + `tl.compare_groups` subsume it; **(e)** a **new `_compute/` package** (not in the original inventory) is the device seam that makes the GPU wins additive. + +--- + +## 3. Rename & disposition map + +Freeze the map **before** Phase 5 (the first breaking PR). Renames are breaking; pre‑1.0 we pay once. The deltas this plan enforces — **including the four settled‑design corrections the draft had wrong** (marked ⚠): + +| current | → target | lands in | +|---|---|---| +| `metrics/` package | `tools/` package (aliased `tl`) | Phase 5 | +| `joint_distribution` + `joint_distribution_posterior` | `tl.joint_distribution` (engine) → top‑level `tcri.joint_distribution` | Phase 5 | +| `posterior=` / `combine_with_logits=` (engine) | **`use_logits=`** (alias `cell_informed`), default `True` (R2) | Phase 5 | +| `point_estimate=`; public `posterior=` on metrics | **removed** — `n_samples` is the only point/draws knob (R1–R2) | Phase 5–6 | +| `flux` / `from_this` / `to_that` | `phenotypic_flux` / `cov_from` / `cov_to` | Phase 6 | +| ⚠ `clonality` | **DROP entirely — do NOT merge into `clonotypic_entropy`** (generic repertoire stat) | Phase 6 | +| ⚠ `clonotypic_entropy_base` | **DROP** (not merged); log base standardized via `_distance` (bits) | Phase 6 | +| ⚠ `ridge_delta_entropy` | **DROP** (not keep‑and‑fix) | Phase 7 | +| ⚠ `compare_phenotypes` | **DROP** (not move‑to‑examples) | Phase 6 | +| `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 (notebook rewrite, Phase 10) | Phase 6/10 | +| `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_clonotype_key` / `register_phenotype_key` (34 notebook uses) | folded into `setup_anndata`/`to_anndata`; **notebook‑rewrite mapping batched to Phase 10** (writer retired only when every reader is migrated) | Phase 4/10 | +| `pl.clonotypic_entropy_by_phenotype` | `pl.clonotypic_entropy` | Phase 7 | +| `plot_pheno_sankey` | `pl.phenotypic_flux` (sankey) | Phase 7 | +| `plot_phenotype_probabilities` | `examples/phenotype_probabilities` — **reads `K.X_PROBABILITIES` ('X_tcri_probabilities')**, not the retired `X_tcri_phenotypes` slot | Phase 10 | +| `probability_ternary` (24 notebook uses) | **keep** as `pl.probability_ternary` (phenotype‑simplex viz) — *explicitly dispositioned; was undispositioned in the draft* | Phase 7 | +| `gene_entropy` (5 notebook uses) | **relocate to `examples/`** — *NOT a 0‑caller; removed from the Phase‑2 delete list* | Phase 10 | +| `polar_plot` (1 notebook use) | **DROP** with an in‑PR notebook‑cell rewrite (not a silent Phase‑2 delete) | Phase 2/10 | +| `pl.flux` boxplot (5 notebook uses) | **DROP** (sankey is the flux plot) — notebook rewrite | Phase 7/10 | +| `tcri_boxplot` | `_metric_boxplot` (private; keeps a **`splitby=` axis**, §Phase 7) | Phase 7 | +| `set_color_palette` | `resolve_palette` | Phase 7 | +| `dkl`, `flux.dkl_func` | `_distance.kl_divergence` (log2/bits, single eps) + `_distance.jensen_shannon` | Phase 1 | +| `Δ` (unicode) | `delta` (ASCII, greppable) | Phase 6 | +| `c2p_mat` | `clone_phenotype_prior` | Phase 3 | +| `centropy` / `pentropy` / `*_tl` leaked aliases | removed via `__all__` | Phase 11 | +| `uns["tcri_clone_key"]`/`["tcri_phenotype_key"]` **and** `uns["tcri_metadata"][...]` (two schemes) | one scheme via `_keys.py` (single `tcri_metadata`) | Phase 1 | + +--- + +## 4. Shared‑helper extraction (`_keys` / `_console` / `_stats` / `_distance`) + +Lands as **Phase 1**, before any file move, because every later phase imports these. Pure internal dedup; public API unchanged except for the new public `tl.compare_groups` (Phase 6, built on `_stats`). + +| new module | absorbs | notes / fixes folded in | +|---|---|---| +| `tcri/_keys.py` | every `uns/obsm/obs` key literal, both current schemes | constants only; unifies `tcri_clone_key`/`tcri_metadata['clone_col']` to one `tcri_metadata` scheme; **adds `X_LOGITS`, `GATE_PROB`, `CLASSIFIER_TEMPERATURE`, `LOCAL_SCALE`** (R4/R6/R8); `clone_size` and any other reader flips to `K.CLONE_COL` in the **same** change that retires the writer (no orphaned reader) | +| `tcri/_console.py` | triplicated `_ok/_info/_warn/_fin` (metrics/pp/pl copies) | reimplement over **scanpy's `logging`/verbosity**; drop raw ANSI; **drop `_ascii_hist` and every `graph=`/ASCII‑histogram code path** | +| `tcri/_stats.py` | `stars`, `auc_and_label_permutation`, `bootstrap_auc` (from `utils/_utils.py`) | **add** `mann_whitney`, `prob_gt_zero` (Bayesian P(>0), for signed contrasts only — R5), **`hdi` implemented as a TRUE highest‑density interval** (not the mislabeled equal‑tailed `percentile[2.5,97.5]`); document boundary instability for bounded skewed posteriors | +| `tcri/_distance.py` | module‑level dead `dkl` + `flux.dkl_func` | one `kl_divergence` + `l1_distance` + **`jensen_shannon`** + `phenotype_distance(metric=...)`; **fix the mixed‑units bug — standardize on `log2` (bits) across entropy/MI/KL with one eps**; document `dkl` as directional/unbounded and recommend JSD (bounded ≤1 bit) for symmetric shift; `l1` stays the safe bounded default | + +`_mi_from_joint` stays module‑private in `tools/_mutual_information.py`. The **public comparison surface** is `tl.compare_groups` (§6) — not the private `_stats` primitives — so "comparisons via `groupby` + stats" is programmatically reproducible. + +--- + +## 5. Model → AnnData streamline + +The single highest‑risk behavior change, and the one that makes the metric↔model agreement guarantee (R8) achievable. It kills the `AnnDataManager`‑in‑`uns` hack and fixes the write‑set the draft under‑specified. + +**5.1 `setup_anndata` — registration only (no *analysis* obs mutation).** +- Registers fields via scvi `REGISTRY_KEYS`; **keeps writing/registering `obs['indices']`** (training glue, `_model.py:678/684`) — this is not an analysis mutation and must not be removed. +- **Removes the manager stash `adata.uns['tcri_manager'] = adata_manager` (`_model.py:697`)** — *this is where the stash actually lives, not in `register_model`.* Removing it here is what lets `write_adata_safely`/`_pop_nonserializables` be deleted. +- `group_singletons` stays a **separate `pp` step that must run BEFORE `setup_anndata`** (it relabels clones; running it after desyncs `ct_to_c`/`p_ct` from `obs`). Enforced: `setup_anndata` errors if a later relabel is detected. + +**5.2 `model.to_anndata(adata)` — thin, canonical write‑set.** Writes **only** the canonical minimum via `_keys`, and the canonical minimum now **explicitly includes the three items the engine needs for `predict` parity** (correcting the draft's "nothing else"): + +| slot | key | why it is canonical | +|---|---|---| +| metadata/categories | `K.META`, covariate/phenotype/ct category maps | registry provenance | +| latent | `obsm[K.X_TCRI]` | embedding | +| phenotype probs + hard labels | `obsm[K.X_PROBABILITIES]`, `obs[...]` | `predict()` output; standard slot (retire `X_tcri_phenotypes`) | +| ct‑level prior mean | `uns[K.P_CT]` = **raw** `normalize(q_p_ct_raw)` (R6) | `n_samples=0` closed‑form mean | +| **per‑cell logits** | **`obsm[K.X_LOGITS]`** | `use_logits=True` engine path (R8) — hard‑required | +| **gate probability** | **`uns[K.GATE_PROB]`** (scalar or `None`) | gate‑aware combine parity with `predict` (R8) | +| **classifier temperature** | **`uns[K.CLASSIFIER_TEMPERATURE]`** | matches `predict`'s logit scaling (R8) | +| **local scale** | **`uns[K.LOCAL_SCALE]`** | draw variance for `n_samples>0`; engine **raises** (never defaults to 1.0) if missing when `n_samples>0` (R4/R8) | + +- **Stops writing `uns['tcri_manager']`** (already removed in 5.1) → deletes the `write_adata_safely`/`_pop_nonserializables` hack. +- `predict` (renamed from `get_cell_phenotype_probs`): returns a per‑cell phenotype‑prob `DataFrame`; **asserts the inference `DataLoader` is order‑preserving (`shuffle=False`) and indexes by `adata.obs_names`** (or carries the registered `indices` field and reindexes) so ct assignment and barcode labels cannot drift. + +**5.3 Session IO.** `write_adata_safely → save_tcri_session` writes a **plain h5ad** (nothing non‑picklable in `uns` anymore); `load_tcri_session` rebuilds the registry by re‑running `setup_anndata`. `get_p_ct` reads the **process‑global** pyro param store (`q_p_ct_raw`); load must set the store immediately before any `get_p_ct`/`to_anndata` call, and multi‑model/round‑trip diag workflows are documented as single‑model‑per‑process unless params are namespaced (§11). + +--- + +## 6. Engine + metric numeric contract + +This section is the substrate all metrics read; it folds in every math/stats and missing‑link correction. + +**6.1 Engine `tools/_joint.py::joint_distribution`** (thin DataFrame wrapper over `_compute/_joint.py::_joint_draws`, §7): + +``` +joint_distribution(adata, *, covariate=None, groupby=None, n_samples=0, + use_logits=True, clones=None, temperature=1.0, + random_state=None) -> pandas.DataFrame +``` +- Unifies `joint_distribution` + `joint_distribution_posterior`; re‑exported as `tcri.joint_distribution`. Provenance in a **serializable form** (a `params` column or a companion `uns` sidecar), **not** only `df.attrs` (R‑forward: h5ad round‑trips must not silently drop it, §11). +- `n_samples=0`: closed‑form path. `use_logits=False` → tempered `m̃` rows (`== uns[K.P_CT]` at `T=1`); `use_logits=True` (default) → per‑cell `softmax((logits + gate‑combine(log m̃))/T)` aggregated per clone, **identical to `model.predict`** (R6/R8). +- `n_samples>0`: draw `p_ct ~ Dirichlet(clamp(local_scale·m̃, 1e-3))` **once per sample via a seeded torch generator** (R4/R7), feed each draw through the same temperature + combine + scatter‑sum; stack a `sample_id` axis. **All clones within one draw share the SAME `p_ct` draw** (one coherent joint per `sample_id`) — never independent per‑clone draws. +- **`covariate=None` computes the joint across ALL covariate values in one pass from a single `p_ct` draw** (the all‑timepoints path the sankey and multi‑covariate metrics need). +- **Draw‑once invariant:** for `n_samples>0`, the number of Dirichlet draws is exactly `n_samples`, **independent of `#groups` and `#covariates`** — draws are reused across groups/covariates by cell‑masking, not re‑drawn. Enforced by a draw‑counter test (§9). +- **`groupby` is implemented by cell/clone RESTRICTION over the FULL `adata`** (positional masks into full‑space `uns` arrays + `clones=`), **never by passing a sliced `AnnData` to the engine** — this avoids the hard full‑space‑vs‑subset alignment guard that today's `tcri_boxplot` slicing would trip. `_metric_boxplot` is rewritten off the slice‑and‑call pattern. +- **`groupby` requires the cell‑informed path** (or group keys that are clone‑nested / constant within `clone × covariate`); it is **ill‑defined on the ct‑level table** for non‑clone‑determined columns. The design assumes **clones are disjoint across `groupby` groups** (a TCR clone does not span two patients) — now stated explicitly, with a validation that errors when a group split would bisect a clone's cells. +- **Engine bug fixes folded in:** weighting keyed on the **`ct` index** (not the clone index); consistent normalization (row‑ vs whole‑table) across the two old functions; all‑zero‑clone reindex yields **NaN, not inflated uniform entropy**; **torch‑seeded** determinism. + +**6.2 Four metrics** (`tools/_entropy.py`, `_mutual_information.py`, `_flux.py`), uniform signatures; **none expose `use_logits`/`posterior=`/`point_estimate=`**: +``` +tl.clonotypic_entropy(adata_or_jd, *, covariate=None, groupby=None, n_samples=0, temperature=1.0, + clones=None, normalized=True, random_state=None) +tl.phenotypic_entropy (adata_or_jd, *, covariate=None, groupby=None, n_samples=0, temperature=1.0, + clones=None, normalized=True, random_state=None) +tl.mutual_information (adata_or_jd, *, covariate=None, groupby=None, n_samples=0, temperature=1.0, + clones=None, normalized=True, normalize_mode='min', random_state=None) +tl.phenotypic_flux (adata, *, cov_from, cov_to, groupby=None, n_samples=0, temperature=1.0, + clones=None, distance_metric='l1', random_state=None) +``` +- **`mutual_information` default `normalize_mode='min'` (was `'average'`).** Under the settled uniform‑clone prior, `P(c)=1/C` pins the clone marginal entropy `H_c` to `log2(C)` (structural, uninformative), which throttles `'average'` normalization by `~1/log2(C)` and makes it non‑comparable across groups with different clone counts — breaking the whole `groupby` workflow. `'min'` (`I/H_p`, coefficient of constraint) is in `[0,1]`, reaches 1 when clone determines phenotype, and is `C`‑independent. `'average'` is documented as not‑recommended (or dropped). +- **`clonotypic_entropy` normalizer comparability:** divide by `log2` of the number of clones with **genuine support**, dropping requested‑but‑absent (all‑zero, reindexed) clones from `C`; expose a **common‑denominator option** (`log2` of a fixed reference clone count) for cross‑group plots. Document that normalized clonotypic entropy uses group‑specific denominators by default. +- **Dual input.** Each accepts `(adata + covariate)` [compute joint internally] **or** a precomputed joint `DataFrame` [fast path]. **A bare precomputed joint forces `n_samples=0` and `groupby=None`** (it carries no `p_ct`/`local_scale`/logits/cells to resample or re‑partition); `clones=` just re‑filters rows; **raise a clear error** for `n_samples>0`/`groupby` on a bare jd. `phenotypic_flux` correctly takes `adata` only (it needs two joints). Metrics **propagate the input joint's provenance** into their output. +- **Return‑shape rule (uniform):** no `groupby` & `n_samples=0` → scalar/Series; `groupby` set → tidy `DataFrame` (row per group [× phenotype/clone]); `n_samples>0` → adds a `sample_id` axis and, on reduction, summary columns **`mean`, `sd`, `hdi_low`, `hdi_high`** (**no `p_gt`** — R5). Metrics reduce the stack by **iterating the `sample_id` level** (per‑draw full‑joint metric, then summarize). +- **`n_samples>0` intervals are partial posterior** (only `p_ct` uncertainty; classifier logits held at their posterior‑mean encoding) — documented so users don't read them as full posterior‑predictive uncertainty. +- **h5ad‑serializable returns (build‑toward `@tl_result`):** flat columns, **no object‑array columns and no `df.attrs`‑only provenance**; per‑draw vectors go in a separate long frame or a `uns` sidecar. Define the cache key as a hash of `(covariate, groupby, n_samples, temperature, clones, normalized, normalize_mode, distance_metric, random_state)`. + +**6.3 Comparisons via `groupby` + `tl.compare_groups` (public).** The deleted `*_compare`/`*_delta` functions are replaced by a **public mid‑level helper**, not private primitives: +``` +tl.compare_groups(df, *, value, by, reference=None, paired=False, hdi_prob=0.94) + -> tidy DataFrame with per-pair: mean_a, mean_b, delta, U, p (Mann-Whitney), + p_gt (Bayesian P(delta>0)), hdi_low, hdi_high +``` +This is where `p_gt`/HDI live (on the **signed** `delta`, R5). The docstring shows the recipe that recreates `mi_compare`'s per‑pair output, so "groupby + stats subsumes the tables/deltas" is actually reproducible. + +--- + +## 7. GPU / optimization architecture + +Grounded 1:1 in grafiti's `_compute/` wins. The **one architecturally load‑bearing decision**: write the engine's numeric core as a **batched, device‑routable function returning a `[n_samples, n_clones, P]` array**, with pandas only at the boundary — so every win below is designed *in*, not retrofitted later. + +**7.1 Device seam — `tcri/_compute/_xp.py` (copy grafiti's 58‑line reference).** `resolve_device` / `get_xp` / `asnumpy`. Because **torch≥2.4.1 is already a hard dep**, a **torch‑tensor core is the first backend** (torch.cuda when present — zero new deps), with **cupy optional** as a numpy‑style second backend. Every accelerated function returns a plain numpy array via `asnumpy`. Device ladder: `None/'cpu'→cpu`; `'mps'→cpu`; `'auto'/'gpu'/'cuda'→GPU only if the lib imports **and** `getDeviceCount()>0`, else CPU (explicit `'cuda'` warns on fallback, `'auto'/'gpu'` silent). + +**7.2 Batched engine core — `_compute/_joint.py::_joint_draws(...) -> ndarray[n_samples, n_clones, P]`.** Precompute clone integer codes **once**; batch the Dirichlet draw and the softmax on the leading sample axis; scatter‑sum by clone. `tools/joint_distribution` is a thin single‑draw/summary DataFrame wrapper over this stack. Metrics consume the stack and reduce vectorized (`_compute/_reduce.py`). + +**7.3 The wins, priority‑ordered, with expected gains:** + +| # | operation | current hot path | fix | expected gain | prio | +|---|---|---|---|---|---| +| P0 | joint‑by‑clone reduction | `pd.DataFrame(...).groupby(level=0).sum()` rebuilt every draw (`_preprocessing.py:320-322`) over 1e5–1e6 cells | precompute clone codes once; **scatter‑add** (`np.add.at`/`np.bincount` CPU, `torch.index_add_`/`cupy.bincount` GPU) batched across all `n_samples` — grafiti `contingency.py`/`edge_tensor` verbatim | **10–50× on CPU** (pandas groupby → integer‑keyed bincount), multiplied again on GPU; the single biggest win | P0 | +| P0/P1 | Dirichlet sampling loop + softmax | `for i in range(n_samples): joint_distribution_posterior(...)` (`_metrics.py:296/363/548/741/987`), redrawing the **full** `p_ct` each call | draw all `n_samples` at once, **restricted to the covariate's ct rows first**; batch `softmax((logits+log b)/T)` over the leading axis; route via torch.cuda | removes the `n_samples`‑fold Python loop + per‑iteration `.uns`/DataFrame setup; collapses 200–1000 iterations into a few kernels | P0/P1 | +| P1 | entropy / MI reductions | `scipy.stats.entropy` per draw + `jd.loc[cl]` per clone (`_metrics.py:315/562-569/744`) | batched `xlogx` reduction over `[S, n_clones, P]`; MI as joint‑vs‑outer‑product; `nanmean`/HDI over the sample axis — grafiti `joint.py::_entropy/_mi`, float64 accumulators | removes two nested Python loops + slow pandas `.loc`; medium‑high | P1 | +| P1 | share the joint across metrics | each of the 4 metrics rebuilds the joint for the same `(covariate, n_samples)` | compute the `[S, n_clones, P]` stack **once per covariate** and reduce it in all four metrics; **build `groupby` groups in one batched pass** | divides the dominant cost by `#metrics`; realizes the draw‑once invariant (§6.1) | P1 | +| — | deterministic point estimate | `n_samples=0` still enters the sampler / returns one draw today | zero‑draw read of `uns[K.P_CT]` (R1) | correctness **and** cheapest path (no Monte‑Carlo at all) | — | +| P2 | latent UMAP | `umap.UMAP` with `import umap` at module top (`_preprocessing.py:20`) | `cuml.manifold.UMAP` behind `_use_gpu`, umap‑learn CPU fallback, **lazy import moved inside the function**; return float64 ndarray | cuML UMAP commonly **10–50×** umap‑learn, but runs once per analysis → lower total impact; also fixes the module‑top heavy import | P2 | + +**7.4 Grafiti guardrails (replicate ALL):** (1) every GPU lib imported **lazily inside** the function — `import tcri` never touches cupy/cuml/torch.cuda (the current `import umap` at module top already violates this and is fixed); (2) GPU deps never in `install_requires`; (3) permissive device ladder with `getDeviceCount()>0` verification; (4) `asnumpy` at every return boundary; (5) GPU body wrapped in `try/except` that degrades to CPU and reports which backend ran; (6) **float64 accumulators** so GPU matches CPU; (7) **validate the joint** (finiteness, nonnegativity, per‑row sum≈1 — per‑row, to catch cancelling defects) **on‑device before compute**; (8) **chunk the batched reduction** over cells/draws (grafiti Moran's‑I `chunk_size=256`, KDE `blk=8192`) to bound device memory on the large `[n_samples, n_cells, P]` tensor and avoid OOM. + +**7.5 Seeding (R7).** Draws move to a seeded `torch.Generator` (and cupy/torch.cuda RNG on GPU); `random_state` is threaded from every public sampling function. The old `np.random.seed` was a no‑op for the torch draws. + +--- + +## 8. Ordered PR sequence + +Each PR is independently landable with green CI, lowest‑risk first. + +### Phase 0 — Contract freeze + CI scaffolding + caller census *(docs/tests only; zero code change)* +- Frozen contract: `tcri/_contract.pyi` + `tests/test_contract_conformance.py`, ported from grafiti's `_pyi_gen.py`/`test_contract_conformance.py`. Seed `IMPLEMENTED = {}`; each target function flips to implemented as its phase lands. Markdown→`.pyi`→live‑signature drift fails CI. +- **Caller census over `example/` + `docs/` notebooks** (not package source only). Reclassify the draft's "12 confirmed dead": `gene_entropy` (5 uses), `polar_plot` (1), `register_*_key` (34), `pl.flux` boxplot (5), `probability_ternary` (24) are **not** safe deletes. Land the corrected disposition map (§3) into `docs/contract/`. +- **Risk:** none. **Depends on:** nothing. + +### Phase 1 — Shared helpers + `_keys` adoption *(internal dedup; API unchanged)* +- Create `_keys.py` (incl. `X_LOGITS/GATE_PROB/CLASSIFIER_TEMPERATURE/LOCAL_SCALE`), `_console.py`, `_stats.py` (true HDI, `prob_gt_zero`), `_distance.py` (bits/log2, JSD) (§4). +- **Adopt `_keys` at every read/write site.** Migrate `clone_size` to `K.CLONE_COL` in the **same** change that retires `tcri_clone_key` (no orphaned reader). +- **Risk:** low (mechanical, no numerics). **Depends on:** Phase 0. **Hazard:** must precede Phases 4/5. + +### Phase 2 — Safe deletions *(only symbols unreferenced in package AND notebooks)* +- Delete the genuinely dead: `get_latent_embedding`, `group_small_clones`, `register_probability_columns`, `remove_meaningless_genes`, `metrics._ent`, `clone_fraction`, module‑level `dkl`, **`utils.probabilities` (and the `_plotting.py:18` import in the SAME PR)**, `pl.probability_distribution`, `pl.bayesian_mutual_information`, `SankeyNode.hex_to_rgb`. +- **Removed from the draft's delete list:** `gene_entropy` (→ examples, Phase 10), `polar_plot` (→ drop with notebook‑cell rewrite). **Not deleted here:** `compare_joint_distribution` (→ diag, Phase 8), `pl.mutual_information`/`pl.phenotypic_entropy` (keep+fix, Phase 7), and the consolidated‑away `*_table`/`*_delta`/`clonality`/`clonotypic_entropy_base`/`ridge_delta_entropy` (removed **with** their replacement in Phase 6, so no caller is orphaned). +- **Risk:** very low. **Depends on:** Phase 1. + +### Phase 3 — Model module split *(mechanical; no behavior change)* +- Split `model/_model.py` (1074 ln) → `_model.py` + `_module.py` + `_priors.py` + `_classifier.py` + `_training.py`. Rename `c2p_mat → clone_phenotype_prior`. +- **`build_archetypes` keeps returning `(centers, labels)`** (labels drive `diag.archetypes`' cluster ordering); persist labels on the model/`uns`. Reconcile the default‑`K` mismatch (`build_archetypes` default `K=4` vs model `K=10`). +- Explicit `__all__` per module. **Risk:** low. **Depends on:** Phase 1. **Verify:** `test_model_setup`, `test_pyro_params` unchanged‑green. + +### Phase 4 — Model→AnnData streamline *(behavior change; kills the manager hack)* +- Implement §5: `setup_anndata` registration‑only (keeps `obs['indices']`; **removes the `_model.py:697` manager stash**); `group_singletons` enforced to precede `setup_anndata`. +- `register_model → model.to_anndata` writing the canonical set **including `X_LOGITS`, `GATE_PROB`, `CLASSIFIER_TEMPERATURE`, `LOCAL_SCALE`, and the raw `P_CT`** (R6/R8). `get_cell_phenotype_probs → predict` (order‑preserving loader, `obs_names` index). +- `write_adata_safely → save_tcri_session` (plain h5ad); load rebuilds the registry via `setup_anndata` and sets the pyro param store before any `get_p_ct`. +- **Risk:** HIGH. **Depends on:** Phase 1, Phase 3. **Gate:** `test_session_round_trip` rewritten — proves save/load reproduces `p_ct` + latent + `predict` probs with **no `tcri_manager` in `uns`**, `setup_anndata` leaves analysis `obs` untouched, and `to_anndata` writes **exactly** the canonical key set (asserts logits/gate/cls‑temp/local‑scale present). + +### Phase 5 — Engine consolidation *(the substrate; §6.1, §7.2)* +- Create `tools/` (aliased `tl`) and `_compute/` (`_xp.py`, `_joint.py`, `_reduce.py`). Implement `_joint_draws` (batched, device‑routable, scatter‑add) and the thin `tools/joint_distribution` wrapper; re‑export as `tcri.joint_distribution`. +- Signature per §6.1 with **`use_logits`** (renamed from `posterior=`), `random_state`, `covariate=None` → all‑covariates one‑pass, draw‑once invariant, clamped‑Dirichlet draws, single‑knob temperature, groupby by full‑space restriction. Fold in the weighting/normalization/zero‑clone/seed bug fixes. +- **Risk:** HIGH — every metric reads this. **Depends on:** Phase 4. **Verify:** `test_tools/test_joint`: `use_logits=False & n_samples=0 == tempered uns[K.P_CT]` **exactly**; `use_logits=True & n_samples=0 & T=1 == model.predict` aggregation; repeated `n_samples=0` **bit‑identical**; `n_samples>0` **torch‑seeded reproducible** and drawn from `Dirichlet(clamp(local_scale·m̃,1e-3))`; **draw‑counter == n_samples independent of #groups/#covariates**; weighting keyed on `ct`; serializable provenance. + +### Phase 6 — Metric‑API consolidation *(four metrics + `compare_groups`)* +- Populate `tools/_entropy.py`, `_mutual_information.py`, `_flux.py` per §6.2 (dual input, `normalize_mode='min'` default, support‑only clonotypic denominator + common‑denominator option, coherent per‑`sample_id` draws, `mean/sd/hdi` summary with **no `p_gt`**, serializable returns). Add public **`tl.compare_groups`** (§6.3). +- **Delete** `mi_compare`, `delta_entropy_table`, `flux_table`, `delta_clonotypic_entropy`, `phenotypic_entropy_delta`, `clonotypic_entropy_base`, `clonality`, `ridge_delta_entropy`, `compare_phenotypes`, the plural `*_entropies`, and `metrics.dkl`. Delete `metrics/` after migration. +- **Risk:** HIGH. **Depends on:** Phase 5. **Verify:** ranges (`[0,1]` normalized), `n_samples=0` determinism, `groupby` tidy shape, **dual‑input equivalence at `n_samples=0` only**, all‑zero‑clone → NaN, `phenotypic_flux` `cov_from/cov_to` + seeded draws, `compare_groups` recreates `mi_compare`'s per‑pair output, **no test asserts `n_samples=0 == mean(n_samples>0)`** (R3). + +### Phase 7 — Plotting split + pl twins *(fix the broken core plots)* +- Split `plotting/_plotting.py` (1437 ln) → `_base.py` (`_metric_boxplot` **with a `splitby=` axis**, `_finish`), `_colors.py` (`resolve_palette`), `_entropy.py`, `_mutual_information.py`, `_flux.py`, `_sankey.py`. +- Ship the four tl↔pl twins: `pl.clonotypic_entropy` (was `_by_phenotype`), `pl.phenotypic_entropy` **[FIX]**, `pl.mutual_information` **[FIX]**, `pl.phenotypic_flux` (sankey). Keep `pl.probability_ternary`. pl functions are **cache renderers** (no metric math). **Retain `splitby`** (116 notebook uses) so two‑axis figures (dots=patient, boxes=response, x=phenotype) survive; document per‑figure the `groupby`+`splitby` recipe. +- **Changelog behavior notes:** removing `pl.mutual_information`'s `weighted=True` default enshrines the uniform‑clonotype prior → user‑visible MI numbers change (cell‑weighted → per‑clonotype); `pl.flux` boxplot and `clonality` plot dropped. +- **Move to `examples/`:** `top_clone_umap`, `clone_size_umap`. **Risk:** medium. **Depends on:** Phase 6, Phase 1. **Verify:** each twin returns a `Figure`/`Axes` from a tidy tl result; sankey renders. + +### Phase 8 — `diag/` seeding *(new; additive)* +- `diag/_ppc.py`: the **fixed** `compare_joint_distribution` (no undefined‑global `NameError`; model `p(clone,phenotype)` vs empirical counts), phenotype‑probability calibration, reconstruction PPC, entropy/MI vs permutation null. **All return DataFrames.** Make the **model requirement explicit per function**: `joint_distribution_ppc`/calibration/permutation‑null run **adata‑only**; `reconstruction_ppc` **requires the live model** (ZINB decoder lives on the module). `diag/_training.py`: `plot_loss → loss curves`, `plot_archetypes → archetypes` (consumes `build_archetypes` labels). +- **Risk:** low‑medium. **Depends on:** Phase 4 (finalized model), Phase 5 (engine). **Verify:** each PPC returns the expected columns on the `trained_model` fixture; permutation‑null seeded; single‑model‑per‑process param‑store scoping honored. + +### Phase 9 — PGM → docs; utils finalize +- Move `build_nested_tcri_pgm`/`draw_tcri_pgm_nested` **out of the package into `docs/`**; drop `daft` from runtime deps → docs extras only. `utils/_utils.py → utils/_session.py` (session‑io + `_to_jsonable` only). +- **Risk:** low. **Depends on:** Phase 1, Phase 8. + +### Phase 10 — Notebook + examples rewrite *(the real end‑to‑end acceptance test)* +- Rewrite `example/` notebooks (`smith`, `renal`, `yost`, `zhang`, `synthetic`, `colitis`) to the new API: `setup_anndata → TCRIModel → train → model.to_anndata`; `tcri.joint_distribution`; four metrics with `groupby`/`n_samples`; `tl.compare_groups`; four pl twins; `diag` checks. Apply the **notebook‑rewrite mappings** for the retired `register_*_key` (34), `gene_entropy` (→ examples), `polar_plot`, `pl.flux` boxplot, and the `X_tcri_phenotypes → K.X_PROBABILITIES` key fix in the moved `phenotype_probabilities` example. One canonical end‑to‑end tutorial notebook under `docs/`. +- **Risk:** low (docs), high value. **Depends on:** Phases 4–8. + +### Phase 11 — Public API finalize + scverse ecosystem CI +- `tcri/__init__.py`: **explicit** named re‑exports, `sys.modules` aliases for `tl/pp/pl/ml/diag/ut`, top‑level `tcri.joint_distribution`, **remove `import *`**, kill leaked aliases (`centropy`/`pentropy`/`*_tl`) via `__all__`. Flip **all** target functions to `IMPLEMENTED`; drift now hard‑fails CI. Turn on the full scverse‑ecosystem gate (§9.2). +- **Risk:** low‑medium. **Depends on:** all prior phases. + +--- + +## 9. Testing + scverse‑CI strategy + +### 9.1 Tests per phase + +| phase | required tests | +|---|---| +| 0 | `test_contract_conformance` (markdown→`.pyi`→signature drift); import‑smoke py3.10/3.11; **notebook caller‑census fixture** feeding the disposition map | +| 1 | `_stats` (`stars`, AUC/permutation, **true HDI vs equal‑tailed**, `prob_gt_zero` on a signed vector), `_distance` (`kl_divergence` bits/symmetry, JSD bound), `_console` silence flag, `_keys` "no stray literal" grep | +| 2 | every deleted symbol absent from `__all__` **and** unreferenced (import‑graph test, package **+** notebooks); `import tcri` green after `utils.probabilities` + its `_plotting.py:18` import go together | +| 3 | `test_model_setup`, `test_pyro_params` green through the split; `build_archetypes` returns `(centers, labels)`; submodule import smoke | +| 4 | **`test_session_round_trip` rewritten** — save/load reproduces `p_ct`+latent+`predict` probs with **no `tcri_manager`**; `setup_anndata` leaves analysis `obs` untouched but keeps `obs['indices']`; `to_anndata` writes **exactly** the canonical set incl. `X_LOGITS/GATE_PROB/CLASSIFIER_TEMPERATURE/LOCAL_SCALE`; `predict` order‑preserving | +| 5 | `test_tools/test_joint`: `use_logits=False,n=0 == tempered uns[P_CT]`; `use_logits=True,n=0,T=1 == predict` aggregation (**engine==predict** agreement test); repeat `n=0` bit‑identical; `n>0` torch‑seeded + clamped‑Dirichlet; **draw‑count == n_samples ⟂ #groups/#covariates**; weighting on `ct`; `groupby` via full‑space restriction does not trip the alignment guard | +| 6 | ranges, `n=0` determinism, `groupby` tidy shape, **dual‑input equivalence (n=0 only)**, all‑zero‑clone→NaN, `normalize_mode='min'` C‑independence, `compare_groups` recreates `mi_compare`; **golden regression** computing MI/entropy from a fixed `uns[P_CT]` by hand; **no `n=0==mean(n>0)` assertion** | +| 7 | each pl twin returns `Figure`/`Axes` from a tidy tl result without metric math; `splitby` two‑axis render; sankey renders; `probability_ternary` renders | +| 8 | each PPC returns expected columns on `trained_model`; `reconstruction_ppc` requires model, others adata‑only; permutation‑null seeded | +| 9 | no `daft`/PGM import in the installed package (import‑graph test) | +| 10 | **notebook execution** (`pytest --nbmake`) on the synthetic tutorial end‑to‑end; retired‑symbol mappings applied | +| 11 | contract test with **all** functions `IMPLEMENTED`; `__all__` completeness (public names ↔ contract); "no `import *`" AST test | + +Reuse `conftest.py` fixtures (`synthetic_adata`, `trained_model`, `mock_adata`); extend `mock_adata` to the unified `_keys` scheme in Phase 1 and to `to_anndata`'s canonical key set (incl. logits/gate/cls‑temp/local‑scale) in Phase 4. + +### 9.2 scverse‑ecosystem CI + +Bring `.github/workflows/tests.yml` to the cookiecutter‑scverse bar, layered so each phase stays green: +- **Matrix** py3.10 + py3.11 (extend to 3.12 before 1.0), `pip install -e ".[test]"`, `pytest tests/ -v --cov=tcri`. +- **Import‑smoke job** (grafiti pattern): `python -c "import tcri; from tcri.model._model import TCRIModel; from tcri.tools._joint import joint_distribution"` — **and assert no GPU lib (cupy/cuml/torch.cuda) was imported** (guardrail #1). +- **Lint/format gate:** `pre-commit` (ruff + ruff‑format) — add in Phase 1, enforce from Phase 3. +- **Contract‑conformance job** (Phase 0 on) — the markdown→`.pyi`→signature guardrail that lets views land independently. +- **CPU‑only correctness job** — the full suite must pass with no GPU present; a separate optional GPU job (if a runner is available) asserts float64 GPU≈CPU parity. +- **Notebook‑execution job** (`nbmake`) from Phase 10. +- **Docs build** (sphinx + numpydoc; `.readthedocs.yaml` present) must pass; public API fully docstringed/typed. +- **Ecosystem checklist** (final gate, Phase 11): AnnData‑native `setup_anndata`, no `import *`, typed public surface, tutorial notebook, `diag` returns data not plots, GPU strictly opt‑in. + +--- + +## 10. Dependencies & ordering hazards + +``` +0 ─▶ 1 ─┬─▶ 2 + ├─▶ 3 ─▶ 4 ─▶ 5 ─▶ 6 ─▶ 7 ─▶ 8 ─▶ 9 ─▶ 10 ─▶ 11 + └─────────────────────(1 gates 4,5,6) +``` +Hard, load‑bearing constraints: +1. **`_keys` (1) before `to_anndata` (4) and the engine (5).** Both rewrite key IO; centralizing strings first keeps diffs local and prevents two‑scheme drift. `clone_size`'s reader flip and `tcri_clone_key`'s writer retirement must land **together**. +2. **Model split (3) before `to_anndata` (4).** `to_anndata` is a `TCRIModel` method; split first so the behavior‑change diff is readable. +3. **`to_anndata` (4) before the engine (5).** The engine reads exactly the canonical keys/**logits/gate/cls‑temp/local‑scale** `to_anndata` writes; if the write‑set changes after the engine exists, `use_logits=True` breaks silently. The **manager‑in‑`uns` removal (at `setup_anndata:697`) is the single highest‑risk hazard** — it changes the session round‑trip; `test_session_round_trip` is rewritten in the same PR. +4. **Engine (5) before metrics (6).** Metrics are thin readers; building them against a pre‑fix engine would bake in the weighting/normalization/seed/Jensen bugs. +5. **Metrics (6) before pl twins (7).** pl are cache renderers. +6. **Delete `*_table`/`*_delta`/`clonality`/`_base`/`ridge_delta_entropy`/`compare_phenotypes` WITH their replacement (6), never before** — they have live notebook callers; removing ahead of `groupby`/`compare_groups` orphans callers and reds the notebook‑execution job. +7. **Rename freeze before Phase 5.** `from_this/to_that → cov_from/cov_to`, `flux → phenotypic_flux`, `metrics/ → tools/`, `posterior= → use_logits`, dropping `point_estimate=` are all breaking; batch them so notebooks are rewritten once (Phase 10). +8. **`group_singletons` before `setup_anndata`** (clone‑relabel ordering) — enforced by `setup_anndata`. +9. **Deferred, designed‑for:** `@tl_result` uns‑cache + `get.py`. Every tl function returns an **h5ad‑serializable** tidy result (flat columns, no object arrays, serializable provenance) so the decorator is a one‑line wrap; until then pl recomputes via the engine. Do **not** block the refactor on the cache. + +**Independent / parallelizable:** Phase 2 (safe deletions) and Phase 8 (`diag`, once 4–5 land) proceed alongside neighbors without contending for the same files. + +--- + +## 11. Risks & open items + +- **Highest‑risk hazard: the manager‑stash removal + write‑set change (Phase 4).** It reshapes the train→`to_anndata`→session round‑trip and adds four persisted keys. Mitigation: rewrite `test_session_round_trip` in‑PR; assert the exact canonical set; keep the CPU‑only correctness job as the tripwire. +- **`use_logits=True` parity depends on persisted model facts (R8).** If any of `X_LOGITS`/`GATE_PROB`/`CLASSIFIER_TEMPERATURE`/`LOCAL_SCALE` is dropped, parity degrades silently to the additive rule. Mitigation: the engine==predict agreement test (Phase 5) and the canonical‑set assertion (Phase 4). Note `gate_prob` defaults to `None` (models ungated by default), so the bug is latent until someone trains with a gate — the guarantee still must hold. +- **Estimator semantics (R3) are a documentation risk, not a code risk.** `n_samples=0` (plug‑in) and the `n_samples>0` `mean` are *different estimators*; users will expect them to match. Mitigation: docstrings state the Jensen gap explicitly; no conformance test asserts equality. +- **Global pyro param store.** `get_p_ct`/`to_anndata`/diag PPCs read the process‑global `q_p_ct_raw`; loading two models in one process clobbers it. **Open item:** namespace params per model or ship a documented single‑model‑per‑process contract with load setting the store immediately before use. Pre‑existing, not introduced by the refactor, but the multi‑model diag/round‑trip workflows expose it. +- **`temperature≠1` re‑centers the sampled distribution** away from the learned posterior (R6); intervals under non‑default `T` are not the model posterior. Documented; `T=1` is the parity point. +- **Partial posterior.** `n_samples>0` intervals capture only `p_ct` uncertainty (classifier logits fixed). Documented so they are not read as full posterior‑predictive uncertainty. +- **GPU memory on large `[n_samples, n_cells, P]`.** Without chunked reductions (guardrail #8) the batching win becomes an OOM risk on realistic `n_samples × n_cells`. Mitigation: chunk over cells/draws; try/except degrade to CPU. +- **`probability_ternary` / `splitby` disposition** are decided here (keep) but should be confirmed against the final notebook set in Phase 0's census before Phase 7 fixes the `pl` signatures. +- **`@tl_result` cache is deferred.** Return shapes are constrained now to be h5ad‑serializable so the later wrap is a one‑liner; the cache‑key scheme is specified but not implemented. +- **cuML/rapids GPU parity** for UMAP is layout‑different (both valid, as grafiti documents); not bit‑identical to CPU. No test asserts embedding bit‑identity. +- **Open (deliberately deferred):** py3.12 matrix, the `get.py` accessor surface, and any KL/JSD default change for `phenotypic_flux` beyond offering JSD (l1 remains the bounded default). + +--- + +## 12. Rollout checklist (per PR) + +- [ ] One behavior change (or zero, if a mechanical move). +- [ ] `import tcri` green; `tl/pp/pl/ml/diag/ut` handles intact; **no GPU lib imported at import time**. +- [ ] No new key‑string literal outside `_keys.py`; no new `import *`. +- [ ] Deleted symbols removed from `__all__` and unreferenced **in package AND notebooks** (import‑site + call‑site grep). +- [ ] Contract test updated (`IMPLEMENTED` flipped for functions this PR lands). +- [ ] Phase‑specific tests from §9.1 present and green on py3.10/3.11 (CPU‑only job passes). +- [ ] Rename/disposition rows for this phase applied everywhere (code + fixtures; notebooks batched to Phase 10). +- [ ] Any persisted‑key change reflected in `_keys`, `to_anndata`'s canonical set, and the round‑trip test. + +--- + +## Appendix A — Audit findings by severity (traceability) + +Every finding folded into the body above, cross‑referenced so nothing is lost. **Src:** PC = plan‑correctness, MS = math/stats, ML = missing‑links, GPU = GPU/optimization, PVM = prior‑vs‑mean resolution. + +### A.1 Blocking / High + +| # | src | finding | resolved in | +|---|---|---|---| +| B1 | PC/ML | `to_anndata` "canonical minimum" drops `obsm['X_tcri_logits']` that the default engine path hard‑requires | §5.2, R8, Phase 4/5 | +| B2 | PC/ML | `gate_prob` (and `classifier_temperature`) are model attributes never persisted → gate‑aware parity infeasible from `adata` | §5.2, R8, Phase 4 | +| B3 | PC | `setup_anndata` cannot be "no obs mutation" — must keep registered `obs['indices']` | §1.1, §5.1, Phase 4 | +| B4 | PC | Deleting `utils.probabilities` (Phase 2) breaks `import tcri` via `_plotting.py:18` import | §1.1, Phase 2 | +| B5 | PC | `groupby` by slicing `adata` trips the full‑space‑vs‑subset alignment guard | §6.1, Phase 5 | +| B6 | MS | Jensen gap: `metric(E[p]) ≠ E[metric(p)]`; `n=0` (plug‑in) vs `mean(n>0)` disagree — drop the equality test | R3, §1.2, Phase 6 | +| B7 | MS | `p_gt`/P(>0) vacuous per‑metric (all metrics ≥0) — only on signed contrasts | R5, §6.2/6.3 | +| B8 | MS | `n>0` draws must use the guide's **clamped** `Dirichlet(clamp(local_scale·m,1e-3))` | R4, §6.1, Phase 5 | +| B9 | MS | MI `normalize_mode='average'` breaks under uniform‑clone prior (`H_c=log2 C`) — default `'min'` | §6.2, Phase 6 | +| B10 | ML | `_stats` is private — no public comparison entry point; add `tl.compare_groups` | §4, §6.3, Phase 6 | +| B11 | ML | "safe deletion" census run on package only — `gene_entropy`/`polar_plot`/`register_*_key` etc. have notebook callers | §3, Phase 0/2/10 | +| B12 | ML | `covariate × groupby` semantics unspecified; requires cell‑informed path / clone‑disjoint groups | §6.1, Phase 5/6 | +| B13 | ML | No mechanism shares one `p_ct` draw across groups/covariates → O(groups×cov×n_samples) | §6.1 draw‑once, §7, Phase 5 | +| B14 | GPU | Engine must be a batched device‑routable core returning `[n_samples,n_clones,P]`, DataFrame at boundary | §7.2, Phase 5 | + +### A.2 Medium + +| # | src | finding | resolved in | +|---|---|---|---| +| M1 | PC | `n=0 == uns[P_CT]` holds only for `use_logits=False`; default folds logits — scope the test | R2, §6.1, Phase 5 | +| M2 | PC | `groupby` unrepresentable in ct‑level table for non‑clone‑determined columns | §6.1 | +| M3 | PC/MS | `adata_or_jd` fast path incompatible with `n>0`/`groupby` — restrict to `n=0`,`groupby=None` | §6.2 | +| M4 | PC | `clone_size` reads retired `tcri_clone_key` — migrate reader with writer | §4, Phase 1 | +| M5 | PC | `build_archetypes` must keep returning `labels`; default‑K mismatch (4 vs 10) | §Phase 3, §Phase 8 | +| M6 | MS | `clonotypic_entropy` `log2(C)` denominator group‑dependent, inflated by absent clones | §6.2, Phase 6 | +| M7 | MS | "HDI" is actually equal‑tailed — implement true HDI or rename | §4 (`_stats`), Phase 1/6 | +| M8 | MS/ML | Dual‑input + `n>0` ill‑defined — force `n=0` on bare jd | §6.2 | +| M9 | MS/ML | Temperature re‑centers draws / double‑tempering with guide_temperature — single knob, store raw mean | R6, §5.2, §6.1 | +| M10 | ML | `splitby` (116 uses) dropped — retain a `splitby` axis on pl twins | §Phase 7 | +| M11 | ML | Undispositioned public symbols (`probability_ternary`×24, plural `*_entropies`, `phenotypic_entropy_delta`, `pl.flux` boxplot) | §3 | +| M12 | ML | No `seed`/`random_state` on engine/metrics — add and thread to torch | R7, §6, Phase 5/6 | +| M13 | ML | Normalized entropy/MI non‑comparable across groups (group‑specific denominators) | §6.2 common‑denominator option | +| M14 | ML | `@tl_result` h5ad‑hostile returns (MultiIndex, object arrays, `.attrs`) | §6.2, §10(9) | +| M15 | ML | `n>0` reduction must keep draws coherent (same `p_ct` per `sample_id`) | §6.1/6.2, Phase 5/6 | +| M16 | GPU | Seed torch (not `np.random`) for reproducible draws | R7, §7.5 | + +### A.3 Low + +| # | src | finding | resolved in | +|---|---|---|---| +| L1 | PC | `predict()` DataFrame indexing assumes order‑preserving loader | §5.2, Phase 4 | +| L2 | PC | `pl.mutual_information` `weighted=True` default removal changes displayed numbers | §Phase 7 changelog | +| L3 | PC | `tcri_manager` stash is in `setup_anndata:697`, not `register_model` | §5.1, Phase 4 | +| L4 | PC/ML | moved `plot_phenotype_probabilities` reads `X_tcri_phenotypes`, not `X_tcri_probabilities` | §3, Phase 10 | +| L5 | PC | `get_p_ct` reads the global pyro param store — multi‑model clobber | §5.3, §11, Phase 8 | +| L6 | MS | KL flux asymmetric/unbounded, base mismatch — bits/log2, offer JSD, l1 default | §4 (`_distance`) | +| L7 | MS | `n>0` intervals are partial posterior (p_ct only) — document | §6.2, §11 | +| L8 | MS | `local_scale` uns fallback `1.0` corrupts draw variance — `to_anndata` always writes it; engine raises if missing at `n>0` | §5.2, §6.1, Phase 4/5 | +| L9 | ML | `group_singletons` ordering vs `setup_anndata` unspecified | §5.1, Phase 4 | +| L10 | ML | `diag` model‑required vs adata‑only inconsistent; missing golden + engine==predict agreement tests | §Phase 8, §9.1 | +| L11 | PVM | current `softmax(log p_ct/T)` path is already the posterior mean at `T=1`; "prior/non‑posterior" label is a misnomer — collapse the axis to `use_logits` | R1/R2, §6.1 | +| L12 | GPU | module‑top `import umap` violates lazy‑import guardrail; cuML UMAP behind `_use_gpu` | §7.3/7.4, Phase 5+ | \ No newline at end of file diff --git a/docs/contract/_quarantine/tcri_refactor_audit_data.json b/docs/contract/_quarantine/tcri_refactor_audit_data.json new file mode 100644 index 0000000..27b1e25 --- /dev/null +++ b/docs/contract/_quarantine/tcri_refactor_audit_data.json @@ -0,0 +1 @@ +{"finalApiDoc": "# TCRI \u2014 API & Function Responsibilities (Final)\n\n*The definitive, publishable specification for the refactored, grafiti-mirrored `tcri` package (Door A: standalone scverse package). It folds the API-surface draft and the math/stats draft into a single artifact and applies every fix from the plan-correctness, math/stats, prior-vs-mean, missing-links, and GPU/optimization audits **inline** \u2014 no known-wrong math survives below. For each function you get: exact final signature and module home; responsibility; the precise math/stats it performs; how every argument dictates that math; and the exact return shape per argument combination.*\n\n---\n\n## 0. Conventions, notation, and resolved decisions\n\n### 0.1 Layout principle\n\nMirrors `grafiti`: one file per topic (never a monolith), private cross-cutting helper modules, explicit `__all__` re-export, **never `import *`**. Five view directories map to short handles exactly as grafiti does \u2014 `model\u2192ml`, `tools\u2192tl`, `preprocessing\u2192pp`, `plotting\u2192pl`, `diagnostics\u2192diag` \u2014 plus `utils\u2192ut` (tcri-specific session I/O) and a new private `_compute/` numeric+device seam. The `tl` view lives in `tcri/tools/` (grafiti `tools\u2194tl`), superseding the earlier working name `metrics/`.\n\n### 0.2 Notation\n\n| Symbol | Meaning |\n|---|---|\n| $P$ | number of phenotypes (columns of every joint) |\n| $C$ | number of clonotypes (rows of a joint at one covariate) |\n| $M$ | number of covariate values (e.g. timepoints) |\n| $\\mathrm{ct}$ | a $(\\text{clonotype},\\text{covariate})$ group; `ct_count` = number of them |\n| $p_{ct}\\in\\Delta^{P}$ | learned per-`ct` phenotype distribution; `uns[K.P_CT]`, shape `(ct_count, P)` |\n| $m$ | $=\\text{normalize}(q\\_p\\_ct\\_raw)=\\mathbb{E}_q[p_{ct}]=$ `uns[K.P_CT]` |\n| $\\ell_i\\in\\mathbb{R}^P$ | per-cell classifier logits, `obsm[K.X_LOGITS]` (already scaled by classifier temperature) |\n| $s$ | local scale, `uns[K.LOCAL_SCALE]` (Dirichlet total concentration; controls draw spread only) |\n| $g$ | `gate_prob`, `uns[K.GATE_PROB]` (scalar $\\in[0,1]$, or `None` \u2192 additive) |\n| $\\tau_{\\text{cls}}$ | classifier temperature, `uns[K.CLASSIFIER_TEMPERATURE]` |\n| $T$ | analysis-time `temperature` argument |\n| $\\varepsilon$ | numerical floor; values stated in situ |\n\nIndex maps: `uns[K.CT_TO_COV]`, `uns[K.CT_TO_C]`, `uns[K.CT_ARRAY]` (per-cell ct index), `uns[K.COV_ARRAY]` (per-cell covariate index).\n\n### 0.3 The one substrate every metric reads\n\nTraining learns a variational Dirichlet posterior per `ct` row:\n\n$$q(p_{ct})=\\mathrm{Dirichlet}(\\alpha),\\qquad \\alpha=\\mathrm{clamp}(s\\cdot m,\\ \\min=10^{-3}),\\quad m=\\text{normalize}(q\\_p\\_ct\\_raw)\\in\\Delta^{P}.$$\n\nWhere the clamp is **inactive** (the common case), $\\alpha=s\\cdot m$, so $\\sum_k\\alpha_k=s$ and the mean is closed-form:\n\n$$\\mathbb{E}_q[p_{ct}]=\\frac{\\alpha}{\\sum_k\\alpha_k}=m=\\texttt{get\\_p\\_ct()}=\\texttt{uns[K.P\\_CT]}.$$\n\nThe local scale $s$ **cancels in the mean** and matters only for the spread of draws.\n\n> **Persisted-substrate decision (fixes the double-tempering bug).** `uns[K.P_CT]` stores the **raw** normalized posterior mean $m=\\text{normalize}(q\\_p\\_ct\\_raw)$ with **no** guide-temperature and **no** analysis-temperature baked in. `TCRIModel.get_p_ct()` returns exactly this at its default `guide_temperature=1.0`. The analysis-time `temperature` argument (\u00a70.9) is therefore the **single** tempering knob; it is never composed on top of a pre-tempered vector.\n\n### 0.4 RESOLVED \u2014 the point estimate is the closed-form posterior MEAN (prior vs mean vs MAP vs draw)\n\nThis is the audit's required decision. `n_samples=0` returns the **closed-form variational posterior mean** $\\mathbb{E}_q[p_{ct}]=m=\\texttt{uns[K.P\\_CT]}=\\texttt{get\\_p\\_ct()}$ (temperature-adjusted per \u00a70.9). It is read directly and **never sampled**.\n\nOptions analyzed and their disposition:\n\n| Option | Verdict | Reason |\n|---|---|---|\n| **(a) Closed-form posterior mean** $\\mathbb{E}_q[p_{ct}]=m$ | **ADOPTED** | Exact, closed-form, deterministic, reproducible; already stored in `uns[K.P_CT]` (zero sampling cost). It is the Bayes point estimator under squared-error loss and lies in the simplex interior. Invariant to $s$ (which only sets spread). |\n| (b) Generative prior `clone_phen_prior` / archetype `mixture_concentration` | **REJECTED** | It is guide **initialization** / generative anchor, not what training learned; built from argmax hard labels (leakage/circularity); indexed at clone/archetype level, not the `ct` level metrics need. No metric may read it. |\n| (c) MAP / posterior mode $(\\alpha-1)/(\\sum\\alpha-P)$ | **REJECTED** | $\\alpha_k=s\\,m_k$ is routinely $<1$ (small $s$, $m_k<1$), so the mode sits on the simplex boundary / is undefined \u2014 unstable, discontinuous in $s$. |\n| (d) Mean of many Dirichlet draws | **REJECTED** | Converges to (a) only up to Monte-Carlo noise; a strictly Rao-Blackwell-dominated, non-reproducible estimator of a quantity available in closed form. |\n\n**Fate of the `posterior=` argument.** The *only* real difference between today's two engines was never prior-vs-posterior \u2014 both already used the posterior mean of $p_{ct}$; neither ever touched the generative prior. The real axis is **whether per-cell classifier logits $\\ell_i$ are folded in**. Therefore:\n\n- **`posterior=` is DELETED from all four metrics and from `phenotypic_flux`.** They always use the learned posterior (mean at `n_samples=0`, draws at `n_samples>0`) and, given an `adata`, compute the joint with logits folded in.\n- **On the engine `joint_distribution`, the flag survives but is REDEFINED and RENAMED to `use_logits`** (alias `cell_informed=`), replacing both `posterior=` and the old `combine_with_logits=`. It is a **classifier-mixing switch**, not a prior/posterior switch: `use_logits=True` folds per-cell logits into $\\log(\\text{base})$ exactly like `predict()`; `use_logits=False` returns the `ct`-level table directly. **Both branches use the posterior mean/draws of $p_{ct}$; neither ever touches the generative prior.** The dead `mutual_information(posterior=False) \u2192 NotImplementedError` branch and the \"prior path\" are removed, not implemented.\n\n### 0.5 Uniform sampling convention (`n_samples`)\n\n| `n_samples` | Operation |\n|---|---|\n| `0` | **Deterministic point estimate.** Use the posterior mean $m$ (temperature-adjusted); **no Dirichlet draw**; bit-reproducible on repeat calls. Fixes today's latent bug where `mutual_information`/`flux` at `n_samples=0` returned *one random draw*. |\n| `N>0` | **$N$ i.i.d. posterior draws** from the **exact guide posterior** $p_{ct}^{(s)}\\sim\\mathrm{Dirichlet}\\big(\\mathrm{clamp}(s\\cdot\\tilde p_{ct},\\ \\min=10^{-3})\\big)$, seeded (\u00a70.11). Adds a sample axis; posterior mean/HDI of the functional fall out. |\n\nThe old `point_estimate=` argument is **deleted**; `n_samples` is the only point-vs-draws knob.\n\n> **Clamp fix (blocking).** Draws use the guide's clamped concentration `clamp(local_scale\u00b7m\u0303, 1e-3)` \u2014 **not** the bare `local_scale\u00b7m\u0303` or `local_scale\u00b7m\u0303+1e-8` variants that appear in today's three inconsistent engines and summarize a distribution the model never learned. See \u00a70.10 for the induced (documented, intentional) mean discrepancy on committed clones.\n\n### 0.6 Estimator honesty \u2014 plug-in vs posterior-mean (Jensen gap)\n\nEntropy, MI, and KL/L1 flux are **nonlinear** functionals of $p$, so $\\text{metric}(\\mathbb{E}_q[p])\\neq\\mathbb{E}_q[\\text{metric}(p)]$; the difference is a **Jensen gap**, *not* Monte-Carlo noise. Consequently the two numbers below are **different estimators by design** and must be documented and tested as such:\n\n- **`n_samples=0`** computes the **plug-in-at-posterior-mean** estimator $\\text{metric}(m\u0303)$.\n- The **`mean` summary column of `n_samples>0`** estimates the **posterior mean of the functional** $\\mathbb{E}_q[\\text{metric}(p)]$.\n\nDirections of the gap (Shannon entropy concave; L1/KL flux convex):\n\n| Metric | Relationship |\n|---|---|\n| clonotypic / phenotypic entropy | plug-in $\\ge$ posterior-mean (over-estimates) |\n| phenotypic flux (L1, KL) | plug-in $\\le$ posterior-mean (under-estimates); a clone with no real shift reads exactly $0$ at `n_samples=0` but strictly $>0$ in the `n_samples>0` mean |\n| mutual information $=H(\\phi)-H(\\phi\\mid c)$ | gaps partly cancel; **sign indeterminate**, magnitudes differ |\n\nAdditionally, because draws use the **clamped** concentration while the `n_samples=0` base is the **unclamped** mean $m$, the two also differ on **committed clones** (where $s\\,m_k<10^{-3}$) by a second, clamp-induced term. **No conformance test may assert `n_samples=0 == mean(n_samples>0)`.** Docstrings label the two estimators distinctly.\n\n### 0.7 Uniform return-shape rule\n\n| `groupby` | `n_samples` | Return |\n|---|---|---|\n| unset | `0` | scalar (MI, single-clone flux) or `Series` (entropies over phenotypes/clones) |\n| unset | `N>0` | draw array with a sample axis **+** summary columns `mean, sd, hdi_low, hdi_high` |\n| set | `0` | tidy `DataFrame`, one row per group [\u00d7 phenotype / \u00d7 clone] |\n| set | `N>0` | tidy `DataFrame`, one row per group [\u00d7 phenotype / \u00d7 clone] + `mean, sd, hdi_low, hdi_high` |\n\n> **`p_gt` fix (blocking).** `p_gt`/`P(>0)` is **removed from every single-metric summary**. Entropy, MI, L1 and KL flux are all **$\\ge 0$**, so `P(draw>0) \u2248 1` always and is vacuous. A signed-direction probability is emitted **only** by `tl.compare_groups` on a between-group **difference** $\\Delta$ (\u00a77.6), where it is meaningful.\n\n> **HDI fix.** The interval columns are a **true highest-density interval** (`hdi_low/hdi_high`), i.e. the narrowest interval containing `hdi_prob` mass \u2014 **not** the equal-tailed `np.percentile(x,[2.5,97.5])` mislabeled \"HDI\" in today's code. For the bounded, right-skewed entropy/flux posteriors (mass piled against the boundary for committed clones) the equal-tailed interval is materially wrong. HDIs from few hundred draws near a boundary are documented as unstable.\n\n> **Draw-coherence rule (correctness).** For `n_samples>0`, **all clones within one sample share the same $p_{ct}$ draw** (one coherent joint per sample). Metrics iterate the `sample_id` level and compute the full-joint metric per draw, then summarize \u2014 never independent per-clone draws.\n\n### 0.8 The `weighted` axis \u2014 removed (uniform-clonotype prior)\n\n`weighted` is **dropped from every public signature** (engine and all four metrics). Each clonotype is one unit on the simplex regardless of cell count \u2014 every entropy/MI/flux is a **repertoire-level (per-clonotype) statistic**, stated in each docstring. This removes (i) the current weight-lookup bug (a `ct`-indexed `Counter` keyed with clone indices) and (ii) the inconsistency where `joint_distribution` normalized the whole table to sum 1 while `joint_distribution_posterior` returned un-normalized counts.\n\n> **Behavior-change note (changelog).** `pl.mutual_information` currently defaults `weighted=True`; removing it flips displayed MI from cell-weighted to per-clonotype. `tl.mutual_information` never accepted `weighted`, so only the plotting default changes user-visible numbers.\n\n### 0.9 Temperature \u2014 single knob, one consistent placement\n\n$T$ power-tempers the base **once**, identically in the mean and draw paths:\n\n$$\\tilde p_{ct}=\\mathrm{softmax}\\!\\Big(\\tfrac1T\\log(m+\\varepsilon)\\Big)=\\frac{m^{1/T}}{\\sum_\\phi m_\\phi^{1/T}},\\qquad \\varepsilon=10^{-8}.$$\n\n$T=1$ is the identity (renormalization only); $T<1$ sharpens, $T>1$ flattens. This fixes today's split where the two engines tempered at different stages.\n\n- For `use_logits=True`, the combined per-cell logit is divided by $T$ **once**: $P(\\phi\\mid i)=\\mathrm{softmax}\\big(\\text{combine}(\\ell_i,\\log\\tilde b)/T\\big)$. At **$T=1$** this reproduces `predict()` **bit-for-bit** (classifier temperature $\\tau_{\\text{cls}}$ is already baked into $\\ell_i$; no second division). $T\\neq1$ is an analysis-time temper that intentionally diverges from `predict()`; documented.\n- For `n_samples>0`, draws are centered on the **re-tempered** $\\tilde p_{ct}$; docstrings state that $T\\neq1$ makes the sampled distribution a re-tempered object, not the raw learned posterior.\n\n### 0.10 Reproducibility / seeding\n\n`random_state` (`int | numpy.Generator | torch.Generator | None`) is added to `joint_distribution`, all four metrics, `compare_groups`'s bootstrap, and the sampling `diag` functions. It seeds a **`torch.Generator`** (and, on the GPU path, the CUDA RNG) because draws are `torch` Dirichlet \u2014 fixing the standing no-op where `seed` only touched `np.random`. `n_samples=0` is deterministic regardless of `random_state`.\n\n---\n\n## 1. Package tree\n\n```\ntcri/\n __init__.py # explicit re-export + sys.modules aliases (tl/pp/pl/ml/ut/diag); top-level joint_distribution; NO import *\n _keys.py # single source of every uns/obsm/obs key string (constants only)\n _console.py # leveled, silenceable logging over scanpy.logging (no raw ANSI, no _ascii_hist)\n _stats.py # stars, AUROC+permutation, bootstrap, MWU, prob_direction, hdi, summarize\n _distance.py # kl_divergence, l1_distance, js_divergence, phenotype_distance dispatcher\n _compute/ # NEW private numeric+device seam (grafiti-mirrored)\n _xp.py # resolve_device, get_xp, asnumpy (torch-first, cupy optional, CPU default)\n _joint.py # _joint_draws(...) -> ndarray[n_samples, n_clones, P] (scatter-add core)\n _reduce.py # batched entropy / mutual-information / distance reductions over the stack\n model/ # ml\n _model.py # TCRIModel\n _module.py # TCRIModule (pyro model/guide, get_latent, get_p_ct)\n _priors.py # MixtureDirichlet, VampPrior\n _classifier.py # PhenotypeClassifier\n _training.py # UnifiedTrainingPlan, build_archetypes\n preprocessing/ # pp\n _register.py # registration writers behind TCRIModel.to_anndata (all private)\n _clones.py # group_singletons, clone_size\n tools/ # tl (mirrors pl by filename)\n _joint.py # joint_distribution (THE ENGINE)\n _entropy.py # clonotypic_entropy, phenotypic_entropy\n _mutual_information.py # mutual_information (+ private _mi_from_joint)\n _flux.py # phenotypic_flux\n _compare.py # compare_groups (public group-comparison orchestrator)\n plotting/ # pl\n _base.py # _metric_boxplot, _finish\n _colors.py # tcri_colors, resolve_palette\n _entropy.py # clonotypic_entropy, phenotypic_entropy\n _mutual_information.py # mutual_information\n _flux.py # phenotypic_flux (sankey)\n _ternary.py # probability_ternary\n _sankey.py # SankeyNode, _phenotype_mass_per_clone (private)\n diagnostics/ # diag (NEW)\n _ppc.py # joint_distribution_ppc, phenotype_calibration, reconstruction_ppc, permutation_null\n _training.py # loss, archetypes\n utils/ # ut\n _session.py # save_tcri_session, load_tcri_session (+ private helpers)\n```\n\n`examples/` (outside the package): `top_clone_umap`, `clone_size_umap`, `plot_phenotype_probabilities`, `gene_entropy`, `polar_plot`, rewritten notebooks. `docs/`: the model PGM (`build_nested_tcri_pgm` / `draw_tcri_pgm_nested`).\n\n---\n\n## 2. Top-level `__init__.py` and the `__all__` story\n\n`__all__` is declared at **both** levels (grafiti pattern). Every impl module declares its own `__all__`; every view `__init__` imports symbols by name and re-declares an aggregate `__all__`. The root imports the six view packages, aliases them into `sys.modules`, and re-exports `joint_distribution` for prominence. **No `import *` anywhere** \u2014 numpy/pandas/torch and every `_helper` stay unexported; GPU libs are never imported at module top (\u00a74.3).\n\n```python\n# tcri/__init__.py\nfrom importlib.metadata import PackageNotFoundError, version as _version\ntry:\n __version__ = _version(\"tcri\")\nexcept PackageNotFoundError:\n __version__ = \"0.0.0+unknown\"\n\nfrom . import tools as tl\nfrom . import preprocessing as pp\nfrom . import plotting as pl\nfrom . import model as ml\nfrom . import diagnostics as diag\nfrom . import utils as ut\nfrom .tools import joint_distribution # tcri.joint_distribution\n\nimport sys\nsys.modules.update({f\"{__name__}.{m}\": globals()[m]\n for m in (\"tl\", \"pp\", \"pl\", \"ml\", \"diag\", \"ut\")})\n\n__all__ = [\"tl\", \"pp\", \"pl\", \"ml\", \"diag\", \"ut\", \"joint_distribution\", \"__version__\"]\n```\n\n| View `__init__` | `__all__` |\n|---|---|\n| `tools/__init__.py` (`tl`) | `joint_distribution`, `clonotypic_entropy`, `phenotypic_entropy`, `mutual_information`, `phenotypic_flux`, `compare_groups` |\n| `preprocessing/__init__.py` (`pp`) | `group_singletons`, `clone_size` |\n| `plotting/__init__.py` (`pl`) | `clonotypic_entropy`, `phenotypic_entropy`, `mutual_information`, `phenotypic_flux`, `probability_ternary`, `tcri_colors`, `resolve_palette` |\n| `model/__init__.py` (`ml`) | `TCRIModel` |\n| `diagnostics/__init__.py` (`diag`) | `joint_distribution_ppc`, `phenotype_calibration`, `reconstruction_ppc`, `permutation_null`, `loss`, `archetypes` |\n| `utils/__init__.py` (`ut`) | `save_tcri_session`, `load_tcri_session` |\n\nPrivate modules (`_keys`, `_console`, `_stats`, `_distance`, `_compute/*`) and every private symbol (`_mi_from_joint`, `_metric_boxplot`, `_finish`, `SankeyNode`, `_phenotype_mass_per_clone`, all `_register.py` writers) are **absent** from every `__all__`.\n\n---\n\n## 3. Shared private helper modules\n\n### 3.1 `tcri/_keys.py` \u2014 canonical AnnData key registry (constants only)\n\nImport as `from tcri import _keys as K`; no key literal lives anywhere else.\n\n| Constant | Value | Slot | Meaning |\n|---|---|---|---|\n| `K.METADATA` | `\"tcri_metadata\"` | `uns` | dict: `covariate_col`, `clone_col`, `phenotype_col`, `batch_col` (single scheme; replaces dual `tcri_clone_key`/`tcri_phenotype_key`) |\n| `K.PHENOTYPE_CATEGORIES` | `\"tcri_phenotype_categories\"` | `uns` | ordered phenotype categories |\n| `K.CLONOTYPE_CATEGORIES` | `\"tcri_clonotype_categories\"` | `uns` | ordered clonotype categories |\n| `K.COVARIATE_CATEGORIES` | `\"tcri_covariate_categories\"` | `uns` | ordered covariate categories |\n| `K.P_CT` | `\"tcri_p_ct\"` | `uns` | `(ct_count, P)` **raw** posterior mean $m=\\mathbb{E}_q[p_{ct}]$ (no temper baked in) |\n| `K.CT_TO_COV` | `\"tcri_ct_to_cov\"` | `uns` | `(ct_count,)` ct\u2192covariate index |\n| `K.CT_TO_C` | `\"tcri_ct_to_c\"` | `uns` | `(ct_count,)` ct\u2192clonotype index |\n| `K.CT_ARRAY` | `\"tcri_ct_array_for_cells\"` | `uns` | `(n_obs,)` per-cell ct index |\n| `K.COV_ARRAY` | `\"tcri_cov_array_for_cells\"` | `uns` | `(n_obs,)` per-cell covariate index |\n| `K.LOCAL_SCALE` | `\"tcri_local_scale\"` | `uns` | scalar Dirichlet total concentration $s$ (draws only) |\n| **`K.GATE_PROB`** | `\"tcri_gate_prob\"` | `uns` | **NEW** \u2014 scalar gate $g\\in[0,1]$ or `None`; required for gate-aware `use_logits=True` parity with `predict()` |\n| **`K.CLASSIFIER_TEMPERATURE`** | `\"tcri_classifier_temperature\"` | `uns` | **NEW** \u2014 $\\tau_{\\text{cls}}$; persisted for reproducibility/auditing (already baked into stored logits) |\n| `K.X_LATENT` | `\"X_tcri\"` | `obsm` | `(n_obs, n_latent)` encoder posterior-mean latent |\n| **`K.X_LOGITS`** | `\"X_tcri_logits\"` | `obsm` | `(n_obs, P)` classifier logits \u2014 **part of the canonical write-set** (\u00a75.1); the `use_logits=True` engine path requires it |\n| `K.X_PROBABILITIES` | `\"X_tcri_probabilities\"` | `obsm` | `(n_obs, P)` per-cell phenotype probabilities (`predict()`) |\n| `K.PHENOTYPE_LABEL` | `\"tcri_phenotype\"` | `obs` | per-cell argmax hard label |\n| `K.CLONE_SIZE` | `\"clone_size\"` | `obs` | per-cell clone cell-count |\n| `K.OBS_INDICES` | `\"indices\"` | `obs` | per-cell integer index \u2014 **registration glue** written by `setup_anndata` (\u00a75.1); not analysis output |\n\n> The AnnDataManager is **no longer** stashed in `uns` (`tcri_manager` retired) \u2014 deleting the `write_adata_safely` / `_pop_nonserializables` hack. The stash lived in `setup_anndata`, so it is removed there (\u00a75.1).\n\n### 3.2 `tcri/_console.py` \u2014 leveled, silenceable logging (private)\n\nThin wrappers over `scanpy.logging`; respects scanpy verbosity. Raw ANSI prints and `_ascii_hist` (and every `graph=`/ASCII-histogram code path) are **deleted**.\n\n| Signature | Responsibility |\n|---|---|\n| `info(msg, *, deep=None)` | `scanpy.logging.info`; silenced by scanpy verbosity. |\n| `warning(msg)` | `scanpy.logging.warning`. |\n| `success(msg)` | `scanpy.logging.hint`. |\n| `done(msg=\"done\")` | terminal completion line. |\n\n### 3.3 `tcri/_stats.py` \u2014 significance + posterior-comparison statistics (private)\n\n| Signature | Responsibility / math |\n|---|---|\n| `stars(p)` | $p<10^{-4}\\to$`****`; $<10^{-3}\\to$`***`; $<10^{-2}\\to$`**`; $<0.05\\to$`*`; else `ns`. |\n| `mann_whitney(a, b, *, alternative=\"two-sided\")` | Mann\u2013Whitney $U=\\sum_{i,j}\\mathbb1[a_i>b_j]+\\tfrac12\\mathbb1[a_i=b_j]$; two-sided $p$ from the rank-sum null (`scipy.stats.mannwhitneyu`). Returns `(U, p)`. |\n| `prob_direction(delta)` | **Signed-contrast only.** Given a paired difference-draw vector $\\Delta$: $p_{gt}=\\frac1N\\sum_s\\mathbb1[\\Delta^{(s)}>0]$, $p_{lt}=\\frac1N\\sum_s\\mathbb1[\\Delta^{(s)}<0]$. Returns `(p_gt, p_lt)`. |\n| `hdi(samples, *, hdi_prob=0.94)` | **True** highest-density interval: over sorted samples, the **narrowest** window containing $\\lceil hdi\\_prob\\cdot N\\rceil$ points. Returns `(low, high)`. Documented unstable near a bounded posterior's boundary for small $N$. |\n| `summarize(samples, *, hdi_prob=0.94)` | Reducer for a **raw metric** draw vector \u2192 `{mean, sd, hdi_low, hdi_high}`. **No `p_gt`** (vacuous for $\\ge0$ metrics, \u00a70.7). |\n| `auc_and_label_permutation(scores, labels, *, pos_label=None, n_perm=200_000, seed=42, max_exact=200_000)` | Observed ROC-AUC + two-sided permutation $p$: exact enumeration when $\\binom{n}{k}\\le$`max_exact`, else Monte-Carlo; $p_{\\text{perm}}=\\text{mean}(|\\mathrm{AUC}_{\\text{perm}}-0.5|\\ge|\\mathrm{AUC}_{\\text{obs}}-0.5|)$. Returns `(auc, p, perm_stats, mode)`. |\n| `bootstrap_auc(scores, labels, *, pos_label=None, n_boot=5000, seed=42)` | Resample cells with replacement (reject draws missing a class), recompute AUROC, return the 2.5/97.5 quantiles. Returns `np.array([lo, hi])`. |\n\n### 3.4 `tcri/_distance.py` \u2014 phenotype-distribution distances (private)\n\nDedupes the old module-level `dkl` and `flux.dkl_func`; **one base (bits, $\\log_2$) and one $\\varepsilon=10^{-12}$ library-wide**, matching entropy/MI.\n\n| Signature | Responsibility / math |\n|---|---|\n| `l1_distance(p, q)` | $\\sum_i|p_i-q_i|$; symmetric; range $[0,2]$ on the simplex. Defensively renormalizes inputs. Returns float. |\n| `kl_divergence(p, q, *, base=2.0, eps=1e-12)` | $\\mathrm{KL}(p\\Vert q)=\\sum_i p_i\\log_2\\frac{p_i}{q_i}$; clip to $[\\varepsilon,1]$ then renormalize. **Asymmetric**, range $[0,\\infty)$. **Single base fixed to $\\log_2$ / single $\\varepsilon$** (repairs the natural-log-vs-$\\log_2$ and $10^{-10}$-vs-$10^{-15}$ divergence between the two dead copies). Returns float. |\n| `js_divergence(p, q, *, base=2.0, eps=1e-12)` | **NEW** Jensen\u2013Shannon $=\\tfrac12\\mathrm{KL}(p\\Vert r)+\\tfrac12\\mathrm{KL}(q\\Vert r)$, $r=\\tfrac12(p+q)$; **symmetric, bounded $[0,1]$ bit** \u2014 the recommended symmetric shift measure. Returns float. |\n| `phenotype_distance(p, q, *, metric=\"l1\")` | Dispatcher: `\"l1\"`\u2192`l1_distance`; `\"kl\"`/`\"dkl\"`\u2192`kl_divergence` (directional, unbounded, bits); `\"jsd\"`\u2192`js_divergence`; a callable `f(p,q)`; else `ValueError`. |\n\n---\n\n## 4. `tcri/_compute/` \u2014 numeric + device seam (NEW, private)\n\nThe engine's numeric core is written **once** as a batched, device-routable function so the acceleration is additive and reversible, and so the four metrics share one joint-draw stack.\n\n### 4.1 `_xp.py` \u2014 the device seam (grafiti reference, copied 1:1)\n\n| Signature | Responsibility |\n|---|---|\n| `resolve_device(device)` | `None`/`\"cpu\"`\u2192`\"cpu\"`; `\"mps\"`\u2192`\"cpu\"`; `\"cuda\"`/`\"gpu\"`/`\"auto\"`\u2192GPU **iff** the backend imports AND a device is present (`getDeviceCount()>0`), else CPU. Explicit `\"cuda\"` warns on fallback; `\"auto\"`/`\"gpu\"` silent; unknown warns. |\n| `get_xp(device)` | Return the array module \u2014 torch(-cuda) preferred (already a hard dep \u2192 zero new deps), cupy optional, numpy default. GPU libs imported **lazily inside** the function. |\n| `asnumpy(x)` | Host-boundary shim: `cupy.asnumpy(x)` / `x.cpu().numpy()` / `np.asarray(x)`. Every accelerated function returns a plain numpy array. |\n\n### 4.2 `_joint.py` / `_reduce.py` \u2014 the batched core\n\n- **`_joint_draws(adata, *, covariate, clones, n_samples, use_logits, temperature, gate_prob, random_state, device) -> np.ndarray`** \u2014 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\u00b7m\u0303, 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` \u2014 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.\n- **`_reduce.py`** \u2014 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.\n\n### 4.3 GPU guardrails (replicated uniformly from grafiti)\n\nLazy GPU imports (never at module top \u2014 `import tcri` never touches a GPU lib; the old module-top `import umap` is moved inside its function); GPU deps never in `install_requires` (CPU path always fully functional); permissive device ladder with device-count verification; `asnumpy`/`output_type=\"numpy\"` at every boundary; `try/except` degrade-to-CPU reporting which backend ran; `float64` where parity must hold; validate-before-compute on **per-row** invariants; chunked reductions to bound host+device memory on the large `[n_samples, n_cells, P]` tensor. `n_samples=0` performs **zero** draws (closed-form read of `uns[K.P_CT]`). cuML UMAP for the one-off latent embedding sits behind a `_use_gpu(device)` gate with the `umap-learn` CPU fallback; layouts differ (both valid), documented, not claimed bit-identical.\n\n---\n\n## 5. `tcri.ml` \u2014 model (`model/`)\n\n### 5.1 `model/_model.py`\n\n**`class TCRIModel(BaseModelClass)`** \u2014 register \u2192 build \u2192 train \u2192 extract \u2192 write.\n\n| Method (signature) | Responsibility |\n|---|---|\n| `@classmethod setup_anndata(cls, adata, *, layer=None, clonotype_key=\"unique_clone_id\", phenotype_key=\"phenotype_col\", covariate_key=\"timepoint\", batch_key=\"patient\", **kwargs)` | **Registration only** \u2014 register clonotype/phenotype/covariate/batch/count fields with scvi and store the layer. **Writes `obs[\"indices\"]=range(n)` and registers it** (`CategoricalObsField`) \u2014 this is registration glue that `training_step`/`validation_step` consume via `batch[\"indices\"]`; it is **not** analysis output. Invariant is **\"no analysis/label `obs` mutation\"** (labels/probabilities are written only by `to_anndata`). **Removes the `uns[\"tcri_manager\"]` stash** (was here), deleting the need for `write_adata_safely`. |\n| `__init__(self, adata, *, n_latent=128, n_hidden=128, n_layers=3, classifier_n_layers=3, global_scale=5.0, local_scale=3.0, prior_temperature=1.0, guide_temperature=1.0, use_enumeration=False, patience=300, classifier_hidden=128, classifier_dropout=0.1, n_pseudo_obs=10, K=10, phenotype_weights=None, gate_prob=None, kl_weight_max=1.0, guide_init_scale=10.0, classifier_temperature=1.0, **kwargs)` | Build the empirical clone\u2192phenotype prior + KMeans archetypes + clonotype/covariate index maps + class weights, then construct/prime `TCRIModule`. Note `gate_prob=None` default \u21d2 ungated model; the gate-parity guarantee is only exercised when a gate is trained. |\n| `train(self, *, max_epochs=1000, batch_size=1000, lr=1e-3, reconstruction_loss_scale=1e-3, n_steps_kl_warmup=2000, **kwargs)` | 0.9/0.1 split, `UnifiedTrainingPlan`, `TrainRunner` with `elbo_validation` early stopping. |\n| `get_latent_representation(self, adata=None, *, indices=None, batch_size=None) -> np.ndarray` | Batched encode to the `(n_cells, n_latent)` posterior-mean latent. |\n| `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$** (\u00a70.9, \u00a77.1). Uses an **order-preserving loader** (shuffle=False / sequential sampler) and the registered `indices` field so ct-lookup and barcode labels cannot drift. |\n| `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. |\n| `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 \u2014 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. |\n\n> Relocated off the model: `plot_archetypes`\u2192`diag.archetypes`; `plot_loss`\u2192`diag.loss`. `boost_phenotype_prior`, `use_gate` remain internal.\n\n### 5.2 `model/_module.py`\n\n**`class TCRIModule(PyroBaseModuleClass)`** \u2014 Pyro CVAE with hierarchical clonotype\u2192(clonotype\u00d7covariate) Dirichlet priors and a phenotype classifier. *(internal)*\n\n| Member (signature) | Responsibility |\n|---|---|\n| `__init__(self, n_input, n_latent, P, n_batch, *, global_scale=10.0, local_scale=5.0, prior_temperature=1.0, guide_temperature=1.0, gate_prob=0.5, mixture_concentration=None, n_pseudo_obs=10, use_enumeration=False, classifier_hidden=128, classifier_dropout=0.1, classifier_n_layers=3, n_hidden=128, n_layers=3, class_weights=None, kl_weight_max=1.0, guide_init_scale=10.0, classifier_temperature=1.0)` | Construct encoder/decoder/classifier/VampPrior, `px_r`; register empty two-level buffers + class weights (`mixture_concentration` required). |\n| `prepare_two_level_params(self, clone_phen_prior_mat, ct_to_c, ct_to_cov, ct_array_for_cells, cov_array_for_cells, *, eps=1e-6)` | Normalize/temperature the clone-phenotype prior; register two-level index buffers. |\n| `model(self, x, batch_idx, log_library, ...)` | Generative: sample $p_c$ (MixtureDirichlet), $p_{ct}$ (Dirichlet centered at $p_c$), latent $z$ (VampPrior), ZINB gene obs. |\n| `guide(self, x, batch_idx, log_library, ...)` | Guide: learnable Dirichlet params $q(p_c)$, $q(p_{ct})$ with **`clamp(min=1e-3)`** on the concentration (the floor the draw path must reproduce, \u00a70.5); Normal $q(z)$ from the encoder. |\n| `get_latent(self, tensor_dict) -> torch.Tensor` | Encode a batch to posterior-mean latent $z_{\\text{loc}}$. |\n| `get_p_ct(self) -> torch.Tensor` | Read `q_p_ct_raw` from the (process-global) param store; return the row-normalized `(ct_count, P)` posterior mean. |\n| `use_gate(self) -> bool` (property) | `True` when `gate_prob is not None`. |\n| `@staticmethod _get_fn_args_from_batch(tensor_dict) -> tuple` | Extract `(x, batch_idx, log_library)` from a scvi batch dict. |\n\n> **Param-store caveat (documented).** `get_p_ct` reads the **process-global** Pyro param store (`q_p_ct_raw`); loading two sessions in one process clobbers it. `to_anndata` and every `diag` PPC must be called immediately after the intended model's params are set; `load_tcri_session` sets the store before any `get_p_ct`/`to_anndata` call. Single-model-per-process otherwise.\n\n### 5.3 `model/_priors.py`\n\n**`class MixtureDirichlet(dist.TorchDistribution)`** *(internal)* \u2014 clonotype prior $p_c$. Members: `__init__(self, mixture_weights, concentration, validate_args=None)`; `sample`; `log_prob` (log-sum-exp of component Dirichlet log-probs); `score_parts` (returns `(log_prob, 0, 0)` \u2192 reparam-free); `__call__` (alias for `sample`).\n\n**`class VampPrior(torch.nn.Module)`** *(internal)* \u2014 VampPrior over $z$. Members: `__init__(self, pseudo_inputs, encoder)`; `get_mixture` (uniform `MixtureSameFamily` of `Independent` Normals); `log_prob(self, z)`; `sample`.\n\n### 5.4 `model/_classifier.py`\n\n**`class PhenotypeClassifier(nn.Module)`** *(internal)* \u2014 `__init__(self, n_latent, classifier_hidden, P, *, num_layers=3, dropout_rate=0.1, temperature=1.0)`; `forward(self, x)` returns MLP logits divided by `temperature` ($\\tau_{\\text{cls}}$, baked into the stored logits).\n\n### 5.5 `model/_training.py`\n\n**`class UnifiedTrainingPlan(PyroTrainingPlan)`** *(internal)* \u2014 `__init__(self, module, *, n_steps_kl_warmup=1000, reconstruction_loss_scale=1e-2, num_particles=5, optimizer_config=None, class_weights=None, **kwargs)`; `loss` (property); `configure_optimizers`; `training_step`; `validation_step` (logs `elbo_validation`).\n\n**Module function:** `build_archetypes(c2p_mat, *, K=10) -> tuple[np.ndarray, np.ndarray]` \u2014 KMeans-cluster clone\u2192phenotype rows into `K` normalized archetype centroids. **Returns `(centers, labels)`** \u2014 labels are retained so `diag.archetypes` can reproduce the cluster-ordered heatmap. `K` default is **10** (aligned to `TCRIModel`, repairing the former `K=4` default mismatch).\n\n---\n\n## 6. `tcri.pp` \u2014 preprocessing (`preprocessing/`)\n\n### 6.1 `preprocessing/_clones.py` \u2014 public\n\n| Signature | Responsibility |\n|---|---|\n| `group_singletons(adata, *, clonotype_key=\"trb\", groupby=\"patient\", target_col=\"trb_unique\", min_clone_size=10) -> AnnData` | Collapse clones smaller than `min_clone_size` (per `groupby`) into `\"Singleton_{group}\"` labels in `target_col`. **Ordering invariant (documented + enforced):** any clone relabeling must run **before** `setup_anndata`/`train`, else the learned clonotype categories and `p_ct`'s `ct_to_c` map desync from `obs`; `setup_anndata` refuses registration if a later relabel is detected. |\n| `clone_size(adata, *, key_added=\"clone_size\", return_counts=False)` | Per-clone cell counts, written per cell into `obs[key_added]`. **Reads `uns[K.METADATA][\"clone_col\"]`** (migrated off the retired `tcri_clone_key` in the same change that stops writing it). |\n\n**Private inner:** `group_singletons.collapse_singleton(row)`.\n\n### 6.2 `preprocessing/_register.py` \u2014 private (the `to_anndata` writers)\n\nCalled only by `TCRIModel.to_anndata`; folds in the old `register_phenotype_key` / `register_clonotype_key` / `_compute_logits_and_prior`.\n\n| Signature | Responsibility |\n|---|---|\n| `_write_metadata(adata, model)` | `uns[K.METADATA]` (single scheme) + the three category lists from the registry. |\n| `_register_clonotype_key(adata, clonotype_key, *, order=None)` | Register the clonotype `obs` column + ordered categories. |\n| `_register_phenotype_key(adata, phenotype_key, *, order=None)` | Register the phenotype `obs` column + ordered categories. |\n| `_write_latent(adata, model, *, latent_key=\"X_tcri\", batch_size=256)` | Encoder posterior-mean latent \u2192 `obsm`. |\n| `_write_logits(adata, model, *, logits_key=\"X_tcri_logits\", batch_size=256)` | **Per-cell classifier logits \u2192 `obsm[K.X_LOGITS]`** (canonical; required by the default engine path). |\n| `_write_predictions(adata, model, *, predictions_key=\"X_tcri_probabilities\", label_key=\"tcri_phenotype\", batch_size=256)` | `predict()` probs \u2192 `obsm`; argmax hard labels \u2192 `obs`. |\n| `_write_p_ct(adata, model)` | `p_ct`, `ct_to_cov`, `ct_to_c`, per-cell ct/cov arrays, `local_scale`, **`gate_prob`**, **`classifier_temperature`** \u2192 `uns`. |\n| `_compute_logits_and_prior(model, adata, *, batch_size=256, eps=1e-8) -> tuple[np.ndarray, np.ndarray]` | Run encoder+classifier to extract per-cell logits and $\\log p_{ct}$ from `get_p_ct()`. |\n\n---\n\n## 7. `tcri.tl` \u2014 tools / metrics (`tools/`)\n\n### 7.1 Engine \u2014 `tools/_joint.py`\n\n```python\njoint_distribution(\n adata, *,\n covariate=None, # None \u2192 ALL covariate values in one pass (shared draw)\n groupby=None,\n n_samples=0,\n use_logits=True, # was posterior=; alias cell_informed=; classifier-mixing switch\n clones=None,\n temperature=1.0,\n random_state=None,\n device=None,\n) -> pandas.DataFrame\n```\nRe-exported top-level as `tcri.joint_distribution`. Unifies today's `joint_distribution` + `joint_distribution_posterior`.\n\n**(a) Responsibility.** Produce the clone\u00d7phenotype distribution (the substrate every metric consumes) at a covariate value from the learned variational posterior of $p_{ct}$ \u2014 a deterministic posterior-mean table or posterior draws. Provenance in `df.attrs[\"params\"]` (and, for cache-friendliness, in a flat `_provenance` sidecar column, \u00a77.7).\n\n**(b) Math.** Select `ct` rows for covariate $m$ via `ct_to_cov`; each clone maps to exactly one `ct`, so rows index by clonotype. Temperature-temper the base **once** (\u00a70.9): $\\tilde p_{ct}=\\mathrm{softmax}(\\tfrac1T\\log(m+10^{-8}))$.\n\n*Base per point/draw:*\n- `n_samples=0`: base $b=\\tilde p_{ct}$ (posterior mean; deterministic).\n- `n_samples=N`: bases $b^{(s)}\\sim\\mathrm{Dirichlet}\\big(\\mathrm{clamp}(s\\cdot\\tilde p_{ct},\\ \\min=10^{-3})\\big)$, $s=$`local_scale`, seeded (\u00a70.10). One coherent joint per sample (\u00a70.7).\n\n*`use_logits=False` (ct-level table):* row $c=b$. For `n_samples=0` this is exactly $\\tilde p_{ct}$; at $T=1$ it equals `uns[K.P_CT]` restricted to the covariate \u2014 the clean closed-form identity used in tests.\n\n*`use_logits=True` (fold per-cell logits; matches `predict()`):* per cell $i$ with clonotype $c(i)$, group $\\mathrm{ct}(i)$, base $b_{\\mathrm{ct}(i)}$:\n$$P(\\phi\\mid i)=\\mathrm{softmax}\\!\\Big(\\tfrac1T\\,\\mathrm{combine}\\big(\\ell_i,\\ \\log(b_{\\mathrm{ct}(i)}+\\varepsilon)\\big)\\Big),\\quad\n\\mathrm{combine}=\\begin{cases}g\\,\\ell_i+(1-g)\\log b & g=\\texttt{gate\\_prob}\\neq\\text{None}\\\\ \\ell_i+\\log b & \\text{additive}\\end{cases}$$\nthen $J[c,\\phi]=\\sum_{i\\in c}P(\\phi\\mid i)$, row-normalize. At $T=1$, `n_samples=0`, this equals `predict()` aggregated per clone \u2014 **bit-for-bit**, gate-aware (fixes the standing disagreement where metrics used one Dirichlet *draw* and *never* applied the gate).\n\n*`covariate=None`:* compute the joint for **all** covariate values from a **single shared draw** per sample (the draw-once invariant, \u00a77.8), stacking a covariate axis.\n\n*`groupby=g`:* restrict by cell/clone masks into the **full-space** `uns` arrays (never by slicing the AnnData \u2014 see the guard note below), computing per group value and stacking a group axis.\n\n**(c) Arguments \u2192 math.**\n\n| Argument | Effect |\n|---|---|\n| `covariate` | Selects `ct` rows via `ct_to_cov`. `None` \u2192 all covariates in one shared-draw pass. |\n| `groupby` | Separate joint per group value (adds a group axis), implemented by **restriction over full adata**, not slicing. Requires the cell-informed path or a clone-constant key (see semantics note). |\n| `n_samples` | `0` \u2192 posterior-mean table (deterministic); `N` \u2192 $N$ clamped-Dirichlet draws. Only place `local_scale` enters. |\n| `use_logits` | `True` \u2192 fold logits with $\\log b$ (gate-aware), aggregate per clone, row-normalize; `False` \u2192 `ct`-level $\\tilde p_{ct}$ rows. Neither is the generative prior. |\n| `clones` | Filters rows to the listed clonotypes; with `use_logits=True` also restricts aggregated cells; final reindex to the exact list (absent clones \u2192 dropped, **not** all-zero rows \u2014 see \u00a77.2 fix). |\n| `temperature` | $T$ tempers the base once (\u00a70.9). $T=1$ identity; at $T=1$, `use_logits=True` reproduces `predict()`. |\n| `random_state` | Seeds the torch (and CUDA) Dirichlet generator for `n_samples>0`. Ignored at `n_samples=0`. |\n| `device` | Routes the numeric core through `_compute/_xp` (CPU / torch-CUDA / cupy); result is always host numpy. |\n\n**(d) Return shape** \u2014 `pandas.DataFrame`, columns = phenotype categories.\n\n| `covariate` | `groupby` | `n_samples` | Index / axes |\n|---|---|---|---|\n| set | unset | `0` | rows = clonotype id; `(C_m, P)` |\n| set | unset | `N>0` | MultiIndex (clonotype, `sample_id`); `(C_m\u00b7N, P)` |\n| set | set | `0` | MultiIndex (group, clonotype) |\n| set | set | `N>0` | MultiIndex (group, clonotype, `sample_id`) |\n| `None` | \u2014 | \u2014 | adds a leading covariate level to any of the above |\n\n> **groupby \u2194 alignment guard (blocking fix).** `joint_distribution_posterior` hard-raises if per-cell `uns[...array_for_cells]` lengths $\\neq$ `n_obs`. Passing a **sliced** AnnData (today's `tcri_boxplot` pattern) trips this. groupby is therefore implemented by **positional cell/clone masks into the full-space `uns` arrays** + `clones=`, never by handing a slice to the engine. `_metric_boxplot` (\u00a78.5) is rewritten off the slice-and-call pattern.\n\n> **groupby \u2194 covariate semantics.** `p_ct` is indexed by `ct=(clonotype, covariate)` only. A `groupby` key that is **not** functionally determined by clonotype-at-fixed-covariate (e.g. a tissue cross-cutting one clone) is unrepresentable in the `use_logits=False` table and is only well-defined on the cell-informed `use_logits=True` path. The engine **requires the cell-informed path for such keys**, or requires the key be clone-nested / constant within a clone\u00d7covariate; it errors/warns on `use_logits=False` + a non-clone-determined groupby. The whole per-group scheme assumes **clones are disjoint across groups** (a TCR clone never spans two patients) \u2014 stated explicitly.\n\n`__all__ = [\"joint_distribution\"]`\n\n### 7.2 `tools/_entropy.py` \u2014 `clonotypic_entropy`\n\n```python\nclonotypic_entropy(\n adata_or_jd, *,\n covariate=None, groupby=None,\n n_samples=0, temperature=1.0,\n clones=None, normalised=True, n_clones_ref=None,\n random_state=None, device=None,\n) -> float | pandas.Series | pandas.DataFrame\n```\n\n**(a) Responsibility.** For each phenotype $\\phi$ (at covariate $m$), the normalized Shannon entropy of the distribution over clonotypes carrying that phenotype, $H[P(c\\mid\\phi,m)]$ \u2014 spread of a phenotype across clones. **Repertoire-level (uniform-clonotype prior; \u00a70.8).**\n\n**(b) Math.** From joint $J$, take column $\\phi$ over the **supported** clones only (absent/all-zero clones excluded \u2014 see fix), normalize, take entropy in bits:\n$$v_c=\\frac{J[c,\\phi]}{\\sum_{c'\\in\\text{supp}}J[c',\\phi]},\\qquad H_\\phi=-\\sum_{c\\in\\text{supp}} v_c\\log_2 v_c.$$\nIf `normalised`: divide by $\\log_2 C_{\\text{den}}$ where $C_{\\text{den}}$ = number of **supported** clones (default), or `n_clones_ref` if given (a fixed reference for cross-group comparability). No division when $C_{\\text{den}}\\le1$. Base fixed to 2.\n\n- **Estimator (\u00a70.6):** `n_samples=0` = plug-in $H_\\phi(m\u0303)$; `n_samples>0` `mean` = $\\mathbb{E}_q[H_\\phi]$ (plug-in $\\ge$ posterior-mean for entropy). Documented as distinct.\n- **Fixes:** deterministic `n_samples=0` (no forced draw); `normalised` exposed (was hard-normalized); base fixed to 2; **absent/zero-support clones excluded before normalizing** (no $\\varepsilon$-clip fabricating uniform mass or inflating $C$).\n\n**(c) Arguments \u2192 math.**\n\n| Argument | Effect |\n|---|---|\n| `adata_or_jd` | AnnData \u2192 compute $J$ internally via \u00a77.1 (`use_logits=True`; `covariate` required); precomputed joint DataFrame \u2192 skip to entropy (fast path \u2014 valid only at `n_samples=0`, `groupby=None`; `clones` just re-filters; else `ValueError`, \u00a77.9). |\n| `covariate` / `groupby` | Condition $m$; per-group entropy \u2192 tidy rows (group \u00d7 phenotype). |\n| `n_samples` | `0` \u2192 plug-in per phenotype; `N` \u2192 per-draw + summary. |\n| `temperature` | Tempers $J$ before the column is read. |\n| `clones` | Restricts the clone set \u2192 changes support and the default $\\log_2 C_{\\text{den}}$. |\n| `normalised` / `n_clones_ref` | `True` \u2192 divide by $\\log_2 C_{\\text{den}}$ (range $[0,1]$); `n_clones_ref` fixes the denominator for comparability; `False` \u2192 raw bits. |\n| `random_state` / `device` | Seeding / backend routing for `n_samples>0`. |\n\n**(d) Return shape.** `n_samples=0`, no `groupby` \u2192 `Series` over phenotypes; `n_samples>0`, no `groupby` \u2192 per-phenotype `mean, sd, hdi_low, hdi_high`; `groupby` \u2192 tidy DataFrame row per (group, phenotype) [+ summary]. **Absent phenotype \u2192 `NaN`, not 0.**\n\n> **Comparability note.** Because the default denominator is group-specific, normalized clonotypic entropy is **within-group** unless `n_clones_ref` (a common denominator) is supplied. The `pl` twin defaults cross-group plots to a common `n_clones_ref`.\n\n### 7.3 `tools/_entropy.py` \u2014 `phenotypic_entropy`\n\n```python\nphenotypic_entropy(\n adata_or_jd, *,\n covariate=None, groupby=None,\n n_samples=0, temperature=1.0,\n clones=None, normalised=True,\n random_state=None, device=None,\n) -> float | pandas.Series | pandas.DataFrame\n```\n\n**(a) Responsibility.** For each clonotype $c$, the normalized Shannon entropy of its phenotype distribution $H[P(\\phi\\mid c,m)]$ \u2014 plasticity vs commitment.\n\n**(b) Math.** Row $c$ = $P(\\phi\\mid c)$; normalize over phenotypes, entropy in bits, divide by $\\log_2 P$ if `normalised` and $P>1$:\n$$p_\\phi=\\frac{J[c,\\phi]}{\\sum_{\\phi'}J[c,\\phi']},\\quad H_c=-\\sum_\\phi p_\\phi\\log_2 p_\\phi.$$\nEstimator convention as \u00a70.6 (plug-in at `n_samples=0`). **Critical bug fix:** a clone present in `obs` but with **zero posterior mass** returns **`NaN`** and is excluded \u2014 it is **not** reindexed to zeros, $\\varepsilon$-clipped, and normalized to uniform \u2192 spurious $H=1.0$. Base fixed to 2; `normalised` exposed.\n\n**(c) Arguments \u2192 math.** As \u00a77.2, except the normalizer is $\\log_2 P$ (depends on $P$, not clone count), so `clones` does not change the divisor; $P\\le1\\Rightarrow$ divisor 1.\n\n**(d) Return shape.** `n_samples=0`, no `groupby` \u2192 `Series` over clonotypes; `n_samples>0` \u2192 per-clone `mean, sd, hdi_low, hdi_high`; `groupby` \u2192 tidy DataFrame row per (group, clone) [+ summary].\n\n### 7.4 `tools/_mutual_information.py` \u2014 `mutual_information` (+ kernel)\n\n```python\nmutual_information(\n adata_or_jd, *,\n covariate=None, groupby=None,\n n_samples=0, temperature=1.0,\n clones=None, normalised=True,\n normalise_mode=\"min\", # CHANGED default (was \"average\")\n random_state=None, device=None,\n) -> float | numpy.ndarray | pandas.DataFrame\n```\n\n**(a) Responsibility.** $I(c;\\phi\\mid m)$ in bits \u2014 strength of clone\u2013phenotype coupling. Optionally normalized to $[0,1]$.\n\n**(b) Math \u2014 kernel `_mi_from_joint(pxy, *, normalised, mode=\"min\") -> float`.** Normalize the whole joint: $\\text{pxy}=J/\\sum J$; marginals $p_x=\\sum_\\phi\\text{pxy}$, $p_y=\\sum_c\\text{pxy}$.\n$$I=\\sum_{c,\\phi}\\text{pxy}\\,\\log_2\\frac{\\text{pxy}+\\varepsilon}{p_x p_y^\\top+\\varepsilon}\\ \\text{(bits)},\\quad \\varepsilon=10^{-15}.$$\nWith $H_c=-\\sum p_x\\log_2 p_x$, $H_p=-\\sum p_y\\log_2 p_y$:\n$$I_{\\text{norm}}=\\frac{I}{D},\\quad D=\\begin{cases}\\min(H_c,H_p) & \\texttt{mode=\"min\"}\\ \\text{(coefficient of constraint; default)}\\\\[2pt]\\tfrac12(H_c+H_p) & \\texttt{mode=\"average\"}\\end{cases}$$\nreturning 0 if $D\\le0$.\n\n> **Default `normalise_mode=\"min\"` (blocking fix).** Under the uniform-clonotype prior (\u00a70.8), each row sums to 1 and the table is divided by its sum, so $p_x=P(c)=1/C$ **exactly** and $H_c=\\log_2 C$ is **structural and maximal**. `\"average\"` denom $=\\tfrac12(\\log_2 C+H_p)$ throttles normalized MI by $\\sim1/\\log_2 C$ and shrinks its ceiling as $C$ grows \u2014 non-comparable across groups/covariates with different $C$, breaking the groupby-comparison workflow. `\"min\"` gives $I/H_p$ (reaches 1 when clone determines phenotype, $C$-independent) and is the default. Docstring states $H_c=\\log_2 C$ is not a meaningful normalizer here.\n\n**Fixes:** `n_samples=0` = deterministic plug-in $I(m\u0303)$ (was one random draw); `posterior=False \u2192 NotImplementedError` deleted (\u00a70.4). Estimator honesty per \u00a70.6 (MI Jensen-gap sign indeterminate).\n\n**(c) Arguments \u2192 math.** As the shared table; additionally `normalise_mode` selects $D$. `clones` restricts rows; `normalised` toggles $I$ vs $I/D$.\n\n**(d) Return shape.** `n_samples=0`, no `groupby` \u2192 scalar `float`; `n_samples>0`, no `groupby` \u2192 `(N,)` array + `mean, sd, hdi_low, hdi_high`; `groupby` \u2192 tidy DataFrame, one row per group, column `MI` [+ summary]. Fast path (precomputed jd) valid only at `n_samples=0`, `groupby=None`.\n\n`__all__ = [\"mutual_information\"]`\n\n### 7.5 `tools/_flux.py` \u2014 `phenotypic_flux` (renamed from `flux`)\n\n```python\nphenotypic_flux(\n adata, *,\n cov_from, cov_to,\n groupby=None,\n n_samples=0, temperature=1.0,\n clones=None,\n distance_metric=\"l1\",\n random_state=None, device=None,\n) -> pandas.Series | numpy.ndarray | pandas.DataFrame\n```\n\n**(a) Responsibility.** Per-clonotype distance between a clone's phenotype distribution at `cov_from` vs `cov_to`, over the clone intersection.\n\n**(b) Math.** $J_{\\text{from}}$, $J_{\\text{to}}$ (rows $P(\\phi\\mid c)$); over common clones $c$, $p=J_{\\text{from}}[c]$, $q=J_{\\text{to}}[c]$:\n$$d_c=\\text{phenotype\\_distance}(p,q,\\ \\text{metric}=\\texttt{distance\\_metric}),$$\ndispatched through `_distance` (\u00a73.4): `\"l1\"` (default, bounded $[0,2]$), `\"kl\"`/`\"dkl\"` (directional, unbounded, **bits**), `\"jsd\"` (symmetric, bounded $[0,1]$ bit), or callable.\n\n**Fixes carried by the rewrite:** (1) the `posterior=False` dead branch is gone (no more `TypeError` from unsupported `silent=`/`combine_with_logits=` kwargs); (2) **reproducibility** \u2014 draws use a seeded **torch** generator (old `seed=` only touched NumPy \u2192 no-op); (3) **unit consistency** \u2014 KL is $\\log_2$/bits, single $\\varepsilon$; (4) the `flux_table` column-misalignment (`clones_g.index(cl)` vs `common`-ordered columns) is removed by returning a tidy per-(group,clone) frame keyed by clone id. **Estimator honesty (\u00a70.6):** `n_samples=0` = plug-in $d_c(m\u0303)$ (convex \u2192 under-estimates $\\mathbb{E}_q[d_c]$; a clone with no real shift reads exactly 0 at `n_samples=0` but $>0$ in the `n_samples>0` mean).\n\n**(c) Arguments \u2192 math.**\n\n| Argument | Effect |\n|---|---|\n| `cov_from`, `cov_to` | The two conditions compared (was `from_this`/`to_that`). |\n| `groupby` | Per-group flux \u2192 tidy rows (group \u00d7 clone) with a `clone_size` column (replaces `flux_table`), via full-space restriction. |\n| `n_samples` | `0` \u2192 deterministic per-clone plug-in `Series`; `N` \u2192 $N$ redrawn distance vectors + summary. |\n| `temperature` | Tempers both joints identically before differencing. |\n| `clones` | Restricts both sides; distances over the intersection. |\n| `distance_metric` | `\"l1\"` / `\"kl\"` / `\"jsd\"` / callable. |\n| `random_state` / `device` | Seeding / backend. |\n\n**(d) Return shape.** `n_samples=0`, no `groupby` \u2192 `Series` over common clones; `n_samples>0` \u2192 per-clone `mean, sd, hdi_low, hdi_high`; `groupby` \u2192 tidy DataFrame row per (group, clone) + `clone_size` [+ summary].\n\n`__all__ = [\"phenotypic_flux\"]`\n\n### 7.6 `tools/_compare.py` \u2014 `compare_groups` (public group-comparison orchestrator)\n\n```python\ncompare_groups(\n df, *,\n value, # column holding the per-unit metric value or draw vector\n by, # grouping column (e.g. \"response\")\n reference=None, # baseline level; None \u2192 all pairwise\n paired=False, # True \u2192 paired posterior-draw contrast (uses prob_direction)\n hdi_prob=0.94,\n alternative=\"two-sided\",\n) -> pandas.DataFrame\n```\n\n**Responsibility.** The **public** replacement for the deleted `mi_compare` / `delta_entropy_table` / `flux_table`: turn a tidy `groupby` result (per-unit point estimates, e.g. per patient) or paired posterior-draw vectors into group contrasts. This closes the audit gap where \"`groupby` + `_stats` subsumes `*_compare`/`*_delta`\" was non-functional because `_stats` is private.\n\n**Math.** For each contrast (`reference` vs other, or all pairs):\n- **Unpaired point estimates:** Mann\u2013Whitney $U$ + two-sided $p$ (`_stats.mann_whitney`), group means, and $\\Delta=\\text{mean}_B-\\text{mean}_A$.\n- **Paired posterior draws** (`paired=True`, one draw vector per group per unit, aligned by `sample_id`): the signed difference $\\Delta^{(s)}=\\text{metric}_B^{(s)}-\\text{metric}_A^{(s)}$, then `mean(\u0394)`, `hdi(\u0394)`, and **`p_gt`/`p_lt` via `prob_direction`** \u2014 the **only** place a direction probability is emitted (\u00a70.7).\n\n**Return.** Tidy DataFrame, one row per contrast: `group_a, group_b, mean_a, mean_b, delta, U, p, p_gt, hdi_low, hdi_high, stars`. Recreates `mi_compare`'s per-pair output exactly.\n\n`tools/__init__.py __all__ = [\"joint_distribution\", \"clonotypic_entropy\", \"phenotypic_entropy\", \"mutual_information\", \"phenotypic_flux\", \"compare_groups\"]`\n\n### 7.7 h5ad-serializable return shapes (forward-compat with the deferred `@tl_result` uns-cache)\n\nEvery `tl` return frame is constrained **now** to survive an h5ad round-trip so the deferred `@tl_result` cache is a one-line wrapper later: **flat columns only, no object-dtype \"samples\" columns**, and provenance in a serializable `_provenance` column (JSON string) **plus** `df.attrs[\"params\"]` (attrs are convenience-only; the column is the durable copy). Per-draw values live in a **separate long frame** carrying an explicit `sample_id` level, never as numpy vectors embedded in object columns (which `AnnData.write` drops). Cache key = hash of `(covariate, groupby, n_samples, temperature, clones, normalised, normalise_mode, distance_metric, use_logits, random_state)`.\n\n### 7.8 Draw-once efficiency invariant\n\nFor `n_samples>0`, the engine draws the `p_ct` table **once per sample** and **reuses that draw across all covariates, groups, and clones**; groups are formed by cell/clone masking, not re-drawing. `covariate=None`, the flux sankey's pairwise series, per-patient analyses, and `diag.permutation_null` all consume one shared draw stack. A test/counter asserts the number of Dirichlet draws equals `n_samples`, independent of `#groups` and `#covariates`.\n\n### 7.9 Precomputed-joint fast path constraints\n\nA bare precomputed joint carries no `p_ct`/`local_scale`/logits/cells, so it is valid **only for `n_samples=0` and `groupby=None`** (`clones=` merely re-filters rows). `n_samples>0` or `groupby` on a bare joint raises a clear `ValueError`; a jd that already contains a `sample_id` axis is accepted, with `n_samples` treated as validation, not resampling. The adata-path == precomputed-jd equivalence test is scoped to `n_samples=0`. (`phenotypic_flux` takes adata only \u2014 it needs two joints \u2014 so it is unaffected.) Metrics propagate the input jd's `_provenance` into their output.\n\n---\n\n## 8. `tcri.pl` \u2014 plotting (`plotting/`)\n\nTwins mirror `tl` by filename and function name. Each renderer accepts its `tl` twin's metric arguments (computing the joint internally when needed) plus rendering args. Cross-group comparison is driven by **`groupby`** (dots) and **`splitby`** (box hue) \u2014 **both axes are retained** because `splitby` has 116 example call sites and most published figures carry two categorical axes simultaneously (e.g. dots = patient, boxes = response, x = phenotype). Statistics come from `_stats` / `compare_groups`.\n\n### 8.1 `plotting/_entropy.py`\n\n```python\nclonotypic_entropy(\n adata, *,\n covariate=None, groupby=None, splitby=None,\n n_samples=0, temperature=1.0, clones=None, normalised=True, n_clones_ref=None,\n palette=None, hue_order=None, ax=None, figsize=(6, 3),\n rotation=90, legend_fontsize=6, bbox_to_anchor=(1.15, 1.0),\n random_state=None, save=None, return_df=False,\n)\n```\n**(renamed from `clonotypic_entropy_by_phenotype`)** Box-and-dot plot of clonotypic entropy per phenotype across covariate values, per-group dots, significance brackets. Cross-group plots default to a common `n_clones_ref` for comparability (\u00a77.2).\n\n```python\nphenotypic_entropy(\n adata, *,\n covariate=None, groupby=None, splitby=None,\n n_samples=0, temperature=1.0, clones=None, normalised=True,\n palette=None, ax=None, figsize=(8, 4),\n rotation=90, legend_fontsize=6, bbox_to_anchor=(1.15, 1.0),\n random_state=None, save=None, return_df=False,\n)\n```\n**[FIXED]** Box/strip plot of phenotypic entropy per covariate/group.\n\n`__all__ = [\"clonotypic_entropy\", \"phenotypic_entropy\"]`\n\n### 8.2 `plotting/_mutual_information.py`\n\n```python\nmutual_information(\n adata, *,\n covariate=None, groupby=None, splitby=None,\n n_samples=0, temperature=1.0, clones=None,\n normalised=True, normalise_mode=\"min\",\n palette=None, ax=None, figsize=(8, 4), rotation=90,\n legend_fontsize=6, bbox_to_anchor=(1.15, 1.0),\n random_state=None, save=None, return_df=False,\n)\n```\n**[FIXED; subsumes `mi_compare`]** Box/strip plot of clone\u00d7phenotype MI per covariate; `groupby` (e.g. `\"patient\"`) supplies per-group points and drives AUROC/MWU/label-permutation stats. `weighted` removed (behavior-change note, \u00a70.8); default `normalise_mode=\"min\"` (\u00a77.4).\n\n`__all__ = [\"mutual_information\"]`\n\n### 8.3 `plotting/_flux.py`\n\n```python\nphenotypic_flux(\n adata, *,\n order,\n groupby=None, clones=None,\n normalize=True, temperature=1.0,\n distance_metric=\"l1\",\n phenotype_colors=None, ax=None, figsize=(6, 3),\n show_legend=True, title=None, random_state=None,\n save=None, return_axes=False,\n)\n```\nThe Sankey. Draws phenotype-distribution flow across the ordered `order` sequence of covariate values, calling `tl.phenotypic_flux` pairwise between consecutive values under one shared draw stack (\u00a77.8). `order` replaces pairwise `cov_from`/`cov_to` because a Sankey spans the full ordered series.\n\n`__all__ = [\"phenotypic_flux\"]`\n\n### 8.4 `plotting/_sankey.py` \u2014 private drawing primitives\n\n**`class SankeyNode`** *(internal)* \u2014 `__init__(self, x, y, val, *, dx=0.2, color=None, **kwargs)`; `plot(self, ax)`; `plot_node_connection(self, destination_node, ax, **kwargs)` (curved, color-interpolated ribbon). `_phenotype_mass_per_clone(adata, covariate, clones, normalize) -> dict[str, np.ndarray]` \u2014 `{clone \u2192 phenotype-mass vector}` at one covariate. `SankeyNode.hex_to_rgb` is **deleted** (0 callers; ribbons use `mcolors.to_rgb`).\n\n### 8.5 `plotting/_base.py` \u2014 private plotting engine\n\n| Signature | Responsibility |\n|---|---|\n| `_metric_boxplot(adata, function, *, groupby=None, splitby=None, ylabel=\"\", order=None, palette=None, s=20, ax=None, figsize=(8, 4)) -> (fig, ax)` | **(renamed from `tcri_boxplot`)** Generic per-phenotype metric box/strip engine across `groupby`/`splitby` strata. **Rewritten to compute each stratum by full-space restriction (`clones=`/masks), never `function(adata[mask])`** \u2014 so engine-backed metrics never trip the alignment guard (\u00a77.1). |\n| `_finish(fig, ax, *, save=None, show=None, return_axes=False)` | scanpy-style show/save/return finalizer. |\n\n### 8.6 `plotting/_colors.py`\n\n| Symbol | Responsibility |\n|---|---|\n| `tcri_colors` (`list[str]`) | Canonical categorical hex palette. |\n| `resolve_palette(adata, columns, *, palette=None) -> dict` | **(renamed from `set_color_palette`)** Assign `tcri_colors` to each `obs` column's categories, store in `uns[\"_colors\"]`, return the map. **Fixes the \"writes on `adata.copy()`\" bug (mutates in place).** |\n\n`__all__ = [\"tcri_colors\", \"resolve_palette\"]`\n\n### 8.7 `plotting/_ternary.py` \u2014 `probability_ternary` (dispositioned; 24 live callers)\n\n```python\nprobability_ternary(\n adata, *,\n phenotypes, # the 3 phenotype axes of the simplex\n groupby=None, clones=None,\n palette=None, ax=None, figsize=(5, 5),\n save=None, return_axes=False,\n)\n```\nTernary phenotype-simplex scatter of per-cell/per-clone phenotype probabilities. **Kept public** (heavily used in notebooks) and migrated onto `K.X_PROBABILITIES` and the single metadata scheme; `weighted` removed. `__all__ = [\"probability_ternary\"]`.\n\n---\n\n## 9. `tcri.diag` \u2014 diagnostics (`diagnostics/`) \u2014 NEW\n\nRead-only checks on the finalized model. PPCs return `DataFrame`s; the two relocated model plots render figures. **`model` is required exactly where the live decoder/param store is needed, optional where `adata` suffices** (stated per function).\n\n### 9.1 `diagnostics/_ppc.py`\n\n| Signature | Responsibility / math |\n|---|---|\n| `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`). |\n| `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`. |\n| `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\u2013variance vs observed. **`model` REQUIRED** (live decoder lives on the module, not in `adata`). Returns statistic \u00d7 {observed, simulated, discrepancy}. |\n| `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 (\u00a77.8). Returns per stratum: `observed, null_mean, null_sd, z, p`. |\n\n`__all__ = [\"joint_distribution_ppc\", \"phenotype_calibration\", \"reconstruction_ppc\", \"permutation_null\"]`\n\n### 9.2 `diagnostics/_training.py`\n\n| Signature | Responsibility |\n|---|---|\n| `loss(model, *, log_scale=False, ax=None, save=None)` | **(relocated `plot_loss`)** Plot training/validation ELBO and prior-KL from `model.history_`. |\n| `archetypes(model, *, ax=None, save=None)` | **(relocated `plot_archetypes`)** Cluster-ordered clone-phenotype heatmap + archetype centroids, ordered by the `labels` from `build_archetypes` (retained, \u00a75.5). |\n\n`diagnostics/__init__.py __all__ = [\"joint_distribution_ppc\", \"phenotype_calibration\", \"reconstruction_ppc\", \"permutation_null\", \"loss\", \"archetypes\"]`\n\n> The model PGM (`build_nested_tcri_pgm` / `draw_tcri_pgm_nested`) is moved **out of the package** to `docs/`.\n\n---\n\n## 10. `tcri.ut` \u2014 utilities (`utils/`)\n\n### 10.1 `utils/_session.py` \u2014 public\n\n| Signature | Responsibility |\n|---|---|\n| `save_tcri_session(model, adata, out_dir, *, save_adata=True, compression=\"gzip\") -> dict` | Persist a trained session: scvi model (weights + registry, no embedded adata), Pyro param store, `setup.json`, the h5ad. |\n| `load_tcri_session(run_dir, *, adata_path=None, map_location=None, layer=None) -> (TCRIModel, AnnData)` | Reconstruct `TCRIModel` + `AnnData`: read h5ad, restore setup/category order, re-run `setup_anndata`, load model + Pyro params. **Sets the global Pyro store before any `get_p_ct`/`to_anndata` call** (param-store caveat, \u00a75.2). |\n\n### 10.2 `utils/_session.py` \u2014 private helpers (not re-exported)\n\n`_to_jsonable(x)`; `_collect_setup_from_adata_or_model(adata, model)`; `_restore_category_order(adata, setup)`; `_resolve_TCRIModel()`; `_disable_scvi_onload_train()`; `_ensure_pyro_posterior_params(model, adata)` (guarantees `q_p_ct_raw`; if missing, warn + re-init to uniform $1/P$); `_pyro_load(path, *, map_location=None)`; `_ensure_dir(path)`.\n\n> **Removed from utils:** `write_adata_safely`, `_pop_nonserializables` (manager stash retired at `setup_anndata`); `probabilities` (dead: read a never-written `uns` key \u2014 **and its module-top import in `_plotting.py` is removed in the same PR**, \u00a711); `build_nested_tcri_pgm`/`draw_tcri_pgm_nested` (\u2192 `docs/`); `stars`/`auc_and_label_permutation`/`bootstrap_auc` (\u2192 `_stats.py`).\n\n---\n\n## 11. Surface deltas (removed / renamed / moved) and deletion-safety census\n\n**Deletion-safety rule (applied).** Every \"safe deletion\" is gated on a caller census over **`example/` + `docs/` notebooks**, not package source alone (notebook execution is itself an acceptance gate). Symbols with live notebook callers are **moved with their notebook, or the notebook is rewritten in the same PR** \u2014 never hard-deleted on package-only evidence. Every \"0-caller deletion\" PR also greps **import-sites** (not just call-sites) first.\n\n- **Deleted (dead/broken/out-of-scope, 0 live callers after census):** `clonality` (tl + pl), `clonotypic_entropy_base`, `delta_clonotypic_entropy`, `delta_entropy_table`, `mi_compare` (tl + pl), `flux_table`, `bayesian_mutual_information`, `probability_distribution`, `classify_phenotypes`, `get_latent_embedding`, `group_small_clones`, `register_probability_columns`, `remove_meaningless_genes`, `clone_fraction`, `_ent`, `ridge_delta_entropy`, `dkl` (\u2192 `_distance.kl_divergence`), `probabilities` (**and its `_plotting.py` import**), `SankeyNode.hex_to_rgb`, `_ascii_hist` (+ all `graph=`/ASCII paths), `write_adata_safely`, `_pop_nonserializables`, and the retired `uns` keys `tcri_manager`, `tcri_clone_key`, `tcri_phenotype_key`, the `X_tcri_phenotypes` obsm slot.\n- **Renamed:** `flux`\u2192`tl.phenotypic_flux`; `get_cell_phenotype_probs`\u2192`TCRIModel.predict`; `register_model`\u2192`TCRIModel.to_anndata`; `clonotypic_entropy_by_phenotype`\u2192`pl.clonotypic_entropy`; `tcri_boxplot`\u2192`_base._metric_boxplot`; `set_color_palette`\u2192`resolve_palette`; params `from_this`/`to_that`\u2192`cov_from`/`cov_to`; engine `posterior=`\u2192`use_logits=` (alias `cell_informed=`); `point_estimate=`\u2192removed (use `n_samples`); `weighted=`\u2192removed (uniform-clonotype prior).\n- **Made private (with a public migration path):** `register_clonotype_key` / `register_phenotype_key` \u2192 `preprocessing/_register._register_*_key` (34 notebook callers \u2192 notebooks rewritten to `setup_anndata`/`to_anndata`, or a documented public key-registration shim provided in the same PR).\n- **Moved to `examples/` (with notebooks):** `top_clone_umap`, `clone_size_umap`, `plot_phenotype_probabilities` (updated to read `K.X_PROBABILITIES`), `gene_entropy` (5 callers), `polar_plot` (1 caller). **Moved to `docs/`:** `build_nested_tcri_pgm`, `draw_tcri_pgm_nested`. **Moved to `diag/`:** `compare_joint_distribution`(\u2192`joint_distribution_ppc`), `plot_loss`(\u2192`loss`), `plot_archetypes`(\u2192`archetypes`).\n- **Subsumed by `groupby` + `compare_groups` (removed, with migration recipe):** the plural batch wrappers `clonotypic_entropies` / `phenotypic_entropies`, `pl.phenotypic_entropy_delta`, and every `*_compare` / `*_delta` / `*_table` variant.\n\n---\n\n## 12. Appendix \u2014 current \u2192 target math/stats deltas (what changed and why)\n\n| # | Site | Current | Target | Rationale |\n|---|---|---|---|---|\n| 1 | engine `n_samples=0` | `joint_distribution_posterior` always draws 1 Dirichlet sample | closed-form posterior **mean** $m$, no draw | reproducible, Rao-Blackwell (\u00a70.4) |\n| 2 | `mutual_information`/`flux` `n_samples=0` | returns one random draw | deterministic plug-in point estimate | latent bug (\u00a70.5) |\n| 3 | `posterior=` semantics | conflates draw-vs-mean **and** logit-folding; MI `posterior=False` raises | axis renamed `use_logits`, means *fold per-cell logits* only; both branches use the posterior, never the generative prior | \u00a70.4 |\n| 4 | metric \u2194 model agreement | metrics use a Dirichlet **draw** and **never** apply the gate | `use_logits=True` at $T=1$ uses the same gate-aware, mean-prior rule as `predict()` (needs persisted `X_tcri_logits`, `gate_prob`, `classifier_temperature`) | removes silent disagreement (\u00a70.9, \u00a75.1) |\n| 5 | **plug-in vs posterior-mean** | drafts equate `n_samples=0` with `mean(n_samples>0)` | documented as **different estimators** (Jensen gap: entropy plug-in $\\ge$ mean; flux plug-in $\\le$ mean; MI indeterminate); **no equality test** | \u00a70.6 |\n| 6 | **`p_gt` summary** | attached to every `n_samples>0` metric | **removed** from single-metric summaries; emitted only by `compare_groups` on a signed $\\Delta$ | metrics are $\\ge0$ \u21d2 $P(>0)\\approx1$ (\u00a70.7) |\n| 7 | **posterior draw concentration** | three inconsistent variants: `clamp(s\u00b7m,1e-3)` (guide) vs `s\u00b7p_ct` vs `s\u00b7p_ct+1e-8` | draw from the **exact guide** `Dirichlet(clamp(s\u00b7m\u0303, 1e-3))` | HDIs must summarize the learned posterior (\u00a70.5) |\n| 8 | **MI `normalise_mode` default** | `\"average\"` \u21d2 denom $\\tfrac12(\\log_2C+H_p)$, $C$-dependent | **`\"min\"`** \u21d2 $I/H_p$, $C$-independent; document $H_c=\\log_2C$ structural | cross-group comparability (\u00a77.4) |\n| 9 | clonotypic-entropy denominator | $\\log_2$ of raw reindexed row count (inflated by absent clones) | $\\log_2$ of **supported** clones; optional fixed `n_clones_ref` | comparability (\u00a77.2) |\n| 10 | phenotypic-entropy zero clone | zero-mass clone \u2192 uniform \u2192 $H=1.0$ | zero-support clone \u2192 **`NaN`/excluded** | \u00a77.3 |\n| 11 | flux `seed` | seeds NumPy only; torch draws unaffected | seed a **torch (+CUDA) Generator**; `random_state` on engine/metrics | reproducibility (\u00a70.10) |\n| 12 | KL base/$\\varepsilon$ | natural log in flux; $\\log_2$ elsewhere; mixed $\\varepsilon$ | one base ($\\log_2$/bits), one $\\varepsilon=10^{-12}$; add bounded symmetric `jsd` | unit consistency (\u00a73.4) |\n| 13 | \"HDI\" | equal-tailed percentiles labeled HDI | **true** highest-density interval `hdi_low/hdi_high` | correct for skewed bounded posteriors (\u00a70.7) |\n| 14 | temperature | applied at different stages in the two engines; double-tempered with `guide_temperature` | **single** power-temper of the base; `uns[K.P_CT]` stores the **raw** mean; $T=1$ reproduces `predict()` | \u00a70.9 |\n| 15 | joint-distribution PPC | references undefined global `model` \u2192 `NameError` | reads cols from `uns[K.METADATA]` | \u00a79.1 |\n| 16 | groupby via slicing | `function(adata[mask])` trips the full-space alignment guard | full-space cell/clone **restriction**; `_metric_boxplot` rewritten | \u00a77.1, \u00a78.5 |\n| 17 | `local_scale` fallback | `uns.get(\"tcri_local_scale\", 1.0)` silently corrupts draw variance if unwritten | `to_anndata` always writes `K.LOCAL_SCALE`; engine **raises** (no `1.0` default) when missing at `n_samples>0` | draw-variance integrity (\u00a75.1) |\n| 18 | partial posterior | intervals silently read as full predictive uncertainty | documented: `n_samples>0` captures **`p_ct` uncertainty only** (classifier logits fixed at their posterior-mean encoding) | \u00a70.6 |\n\n---\n\n*Source of truth cross-checked against `tcri/model/_model.py`, `tcri/model/_module.py`, `tcri/preprocessing/_preprocessing.py`, `tcri/metrics/_metrics.py`, `tcri/plotting/_plotting.py`, `tcri/plotting/_sankey.py`, `tcri/utils/_utils.py`, and the grafiti reference at `/Users/ceglian/Codebase/GitHub/grafiti/grafiti`. Intended document home: `/Users/ceglian/Codebase/GitHub/tcri/docs/contract/tcri_api_and_responsibilities.md`.*", "finalPlanDoc": "# TCRI Refactor \u2014 Implementation Plan (Final)\n\n**From** the current five-monolith package **to** the settled Door\u2011A, scverse\u2011ecosystem target (grafiti\u2011mirrored, one\u2011file\u2011per\u2011topic), with every audit fix folded in.\n\n**Source of truth.** The *SETTLED DESIGN* block is authoritative; the *prior\u2011vs\u2011mean RESOLUTION* settles the one parked statistical question and is adopted verbatim. Where the settled design diverges from `docs/contract/tcri_function_inventory.md` \u00a73/\u00a79 (engine lives in `tl` not `pp`; `groupby` subsumes all `*_table`/`*_compare`/`*_delta`; `flux \u2192 phenotypic_flux`; `register_model \u2192 model.to_anndata`; `tl` mirrors grafiti's `tools/`), this plan follows the settled design and flags the divergence inline. This document supersedes the earlier draft wherever the audit corrected it; the corrections are integrated at the point they bite, and every finding is cross\u2011referenced in **Appendix A** so nothing is lost.\n\n**How to read.** \u00a71 fixes the invariants and records the resolved decisions that unblock the engine/metric phases. \u00a72 is the target tree. \u00a73\u20134 are the rename/disposition map and the shared\u2011helper extraction. \u00a75 is the model\u2192AnnData streamline. \u00a76 is the engine + metric **numeric contract** (all math/stats fixes live here). \u00a77 is the GPU/optimization architecture, grounded in grafiti. \u00a78 is the ordered, independently\u2011landable PR sequence. \u00a79 is testing + scverse\u2011CI. \u00a710 is the ordering\u2011hazard graph. \u00a711 is risks & open items. \u00a712 is the per\u2011PR checklist.\n\n---\n\n## 1. Invariants and resolved decisions\n\n### 1.1 Invariants held at every PR boundary\n- **`import tcri` stays green** and the public handles `tl / pp / pl / ml / diag / ut` remain importable. No PR leaves `main` with a broken import or red CI. **Corollary (audit):** \"0\u2011caller\" means *no call\u2011sites **and** no import\u2011sites, in package **and** in `example/`+`docs/` notebooks* \u2014 a symbol imported at module top (e.g. `utils.probabilities` at `_plotting.py:18`) is a live dependency even with zero calls.\n- **One behavior change per PR.** Mechanical moves (splits, helper extraction) never change numerics; numeric changes (engine, metrics) never also move files. Every diff stays reviewable and every regression bisectable.\n- **`_keys.py` is the only place a key string is written**, from Phase 1 on. No new `uns/obsm/obs` key literal may appear outside it.\n- **No `import *`.** Each split lands with an explicit `__all__` and named re\u2011exports; the top\u2011level `__init__` flip is the last PR.\n- **`setup_anndata` performs no *analysis/label* obs mutation.** *(Corrected from the draft's \"no obs mutation.\")* It must still write and register the `obs['indices']` field the training step consumes (`batch['indices']`, `_model.py:603/641/678/684`). The invariant forbids writing *results* (probabilities, hard labels, latent) into `obs`, not the registration glue column.\n- **GPU is never on the import path and never in `install_requires`.** Every accelerated path is pure opt\u2011in behind a device seam with a fully\u2011functional CPU fallback (\u00a77).\n\n### 1.2 Resolved decisions (adopted from the prior\u2011vs\u2011mean RESOLUTION + math/stats audit)\n\n**(R1) Point estimate = closed\u2011form variational posterior mean.** `n_samples=0` returns `E_q[p_ct] = normalize(q_p_ct_raw) = TCRIModel.get_p_ct() = adata.uns[K.P_CT]`, read directly, **never sampled**. Because the guide is `Dirichlet(local_scale \u00b7 m)` with `m` on the simplex, `\u03a3\u03b1 = local_scale` and the mean is exactly `m` \u2014 `local_scale` cancels. The **prior/archetype path is dropped** (it is guide\u2011init/generative anchor built from leaked hard labels, not what training learned); the **MAP/mode is rejected** (`\u03b1_k = local_scale\u00b7m_k` is routinely `< 1`, so the mode sits on the simplex boundary / is undefined); **mean\u2011of\u2011draws is rejected** (Rao\u2011Blackwell\u2011dominated, adds only MC noise to a closed\u2011form quantity).\n\n**(R2) Rename the mis\u2011named `posterior=` axis to `use_logits` \u2014 a classifier\u2011mixing switch, not a prior/posterior switch.** There was never a live prior branch; what actually differs is *with logits vs without*. The **engine** `joint_distribution` keeps exactly one such flag, `use_logits` (alias `cell_informed`), replacing both `posterior=` and `combine_with_logits=`:\n - `use_logits=True` (default) folds per\u2011cell classifier logits into `log(base)` exactly like `model.predict` (gate\u2011aware, \u00a75).\n - `use_logits=False` returns the ct\u2011level table directly.\n Both branches use the posterior **mean** (`n_samples=0`) or **draws** (`n_samples>0`) of `p_ct`; neither ever touches the generative prior. **The four metrics do NOT expose `use_logits`, `posterior=`, or `point_estimate=`** \u2014 given an `adata` they always compute the cell\u2011informed joint. `point_estimate=` is deleted.\n\n**(R3) `n_samples=0` is a *plug\u2011in* estimator, not the posterior mean of the metric.** Entropy, MI, and l1/KL flux are **nonlinear** functionals, so `metric(E_q[p]) \u2260 E_q[metric(p)]` by a **Jensen gap that is not Monte\u2011Carlo noise**: entropy (concave) plug\u2011in **over\u2011estimates**, flux (convex) plug\u2011in **under\u2011estimates**, MI is sign\u2011indeterminate. Therefore:\n - Document `n_samples=0` as **`metric\u2011at\u2011posterior\u2011mean`** and the `n_samples>0` `mean` column as **`E_q[metric]`** \u2014 two different, clearly\u2011labeled estimators.\n - **No test may assert `n_samples=0 == mean(n_samples>0)`.** The determinism/reproducibility guarantee (bit\u2011identical repeated `n_samples=0` calls) still holds and is tested; the *equality across estimators* claim is dropped.\n\n**(R4) `n_samples>0` draws must reproduce the guide's clamped concentration.** Draw from `Dirichlet(clamp(local_scale \u00b7 m\u0303, min=1e-3))`, reusing the guide's floor (`_model.py:490`). The three inconsistent current variants (`local_scale\u00b7m` unclamped; `local_scale\u00b7m + 1e-8`; the clamped guide form) are unified onto the guide form so reported HDIs summarize the distribution the model actually learned. Note that for committed clones (`m_k < 1e-3/local_scale`) the clamp makes the draw\u2011mean differ slightly from `m`; this is documented, and `n_samples=0` remains the closed\u2011form `m` by definition.\n\n**(R5) `p_gt` (P(>0)) is only meaningful on a signed contrast.** Entropy, MI, and flux are all `\u2265 0`, so `P(draw > 0) \u2248 1` and is vacuous. **Remove `p_gt` from the single\u2011metric `n_samples>0` summary** (emit `mean`, `sd`, `hdi_low`, `hdi_high` only). `p_gt` is computed **only** on a paired between\u2011group difference vector by the comparison helper (\u00a74, \u00a76).\n\n**(R6) Temperature is a single analysis knob applied one consistent way.** `m\u0303 = softmax(log(m + 1e-8) / T)` (identity at `T=1`), applied identically in the mean and draw branches. To avoid double\u2011tempering, **`to_anndata` persists the *raw* posterior mean** `m = normalize(q_p_ct_raw)` (not the guide\u2011temperature\u2011adjusted vector), so analysis `temperature` is the sole tempering knob. At `T=1`, `use_logits=True` reproduces `model.predict` exactly. `temperature\u22601` re\u2011centers the sampled distribution away from the learned posterior \u2014 documented in the metric docstrings.\n\n**(R7) Reproducibility is via a seeded torch generator.** All draws move to a seeded `torch.Generator` (device\u2011aware; also seed cupy/torch.cuda RNG on GPU). Add `random_state: int | torch.Generator | None` to `joint_distribution` and every sampling metric/diag function; thread it into the Dirichlet draw. `np.random.seed` alone was a **no\u2011op** for the torch draws and is retired. `n_samples=0` is deterministic regardless.\n\n**(R8) `use_logits=True` parity with `predict` requires three model facts persisted to `adata`.** The gate\u2011aware combination `gate_prob\u00b7logits + (1-gate_prob)\u00b7log(base)` and the classifier temperature are **model attributes**. `to_anndata` must persist `obsm[K.X_LOGITS]`, `uns[K.GATE_PROB]`, and `uns[K.CLASSIFIER_TEMPERATURE]` (\u00a75). Without them the engine can only reproduce the additive rule \u2014 the exact disagreement the refactor set out to remove.\n\n---\n\n## 2. Target layout (grafiti\u2011mirrored, authoritative)\n\n```\ntcri/\n __init__.py # explicit re-export; sys.modules aliases tl/pp/pl/ml/diag/ut;\n # top-level tcri.joint_distribution; NO import *\n _keys.py # single source of every uns/obsm/obs key string (adopt in Phase 1)\n # NOW INCLUDES: X_LOGITS, GATE_PROB, CLASSIFIER_TEMPERATURE,\n # LOCAL_SCALE, P_CT, X_TCRI, X_PROBABILITIES, CLONE_COL/PHENO_COL\n _console.py # leveled, silenceable logging over scanpy logging; NO ANSI, NO _ascii_hist\n _stats.py # stars, auc_and_label_permutation, bootstrap_auc,\n # + posterior-comparison primitives: mann_whitney, prob_gt_zero, hdi (TRUE HDI)\n _distance.py # kl_divergence (log2/bits), l1_distance, jensen_shannon, phenotype_distance dispatcher\n _compute/ # NEW \u2014 device seam + batched numeric core (grafiti/_compute parity)\n _xp.py # resolve_device / get_xp / asnumpy (torch-first, cupy optional, cpu default)\n _joint.py # _joint_draws(adata, covariate, n_samples, *, use_logits, temperature, device,\n # random_state) -> ndarray[n_samples, n_clones, P] (scatter-add reduction)\n _reduce.py # batched entropy / mutual_information / flux over the [S, n_clones, P] stack\n _embedding.py # umap() behind _use_gpu gate (cuML on GPU, umap-learn CPU); lazy import\n model/ # ml\n _model.py # TCRIModel: setup_anndata, train, get_latent_representation,\n # predict (was get_cell_phenotype_probs), get_p_ct, to_anndata\n _module.py # TCRIModule (pyro model/guide)\n _priors.py # MixtureDirichlet, VampPrior\n _classifier.py # PhenotypeClassifier\n _training.py # UnifiedTrainingPlan, build_archetypes (returns centers AND labels)\n preprocessing/ # pp (shrinks to clone utilities; engine moved to tools/)\n _clones.py # group_singletons (must precede setup_anndata), clone_size\n tools/ # tl (metrics + engine; mirrors grafiti tools/)\n _joint.py # joint_distribution (thin DataFrame wrapper over _compute._joint;\n # re-exported as tcri.joint_distribution)\n _entropy.py # clonotypic_entropy, phenotypic_entropy\n _mutual_information.py # mutual_information (+ private _mi_from_joint)\n _flux.py # phenotypic_flux (was flux; cov_from / cov_to)\n _compare.py # compare_groups (PUBLIC mid-level stats helper; replaces *_compare/*_delta)\n plotting/ # pl (twins mirror tl by filename)\n _base.py # _metric_boxplot (was tcri_boxplot; groupby + splitby), _finish\n _colors.py # tcri_colors, resolve_palette (was set_color_palette)\n _entropy.py # clonotypic_entropy [FIX], phenotypic_entropy [FIX]\n _mutual_information.py # mutual_information [FIX]\n _flux.py # phenotypic_flux (sankey)\n _sankey.py # SankeyNode, _phenotype_mass_per_clone\n diagnostics/ # diag (NEW \u2014 PPCs + model validation; returns DataFrames)\n _ppc.py # joint-distribution PPC (fixed compare_joint_distribution),\n # calibration, reconstruction PPC (model-required), permutation-null\n _training.py # loss curves (was plot_loss), archetypes (was plot_archetypes)\n utils/ # ut\n _session.py # save/load_tcri_session, _to_jsonable (plain h5ad; no manager hack)\nexamples/ # top_clone_umap, clone_size_umap, phenotype_probabilities,\n # probability_ternary, gene_entropy (relocated), rewritten notebooks\ndocs/ # model PGM (build_nested_tcri_pgm) lives here, out of the package\n```\n\nDivergences from inventory \u00a73, called out: **(a)** `tl` package is `tools/` (grafiti parity), not `metrics/`; **(b)** the engine's *numeric core* is `_compute/_joint.py` with a thin `tools/_joint.py` DataFrame wrapper \u2014 **not** `preprocessing/_engine.py`; **(c)** `preprocessing/` loses `_register.py`/`_engine.py` (registration collapses into `model.to_anndata`), leaving `pp` = `_clones.py`; **(d)** no `_tables.py` \u2014 `groupby` + `tl.compare_groups` subsume it; **(e)** a **new `_compute/` package** (not in the original inventory) is the device seam that makes the GPU wins additive.\n\n---\n\n## 3. Rename & disposition map\n\nFreeze the map **before** Phase 5 (the first breaking PR). Renames are breaking; pre\u20111.0 we pay once. The deltas this plan enforces \u2014 **including the four settled\u2011design corrections the draft had wrong** (marked \u26a0):\n\n| current | \u2192 target | lands in |\n|---|---|---|\n| `metrics/` package | `tools/` package (aliased `tl`) | Phase 5 |\n| `joint_distribution` + `joint_distribution_posterior` | `tl.joint_distribution` (engine) \u2192 top\u2011level `tcri.joint_distribution` | Phase 5 |\n| `posterior=` / `combine_with_logits=` (engine) | **`use_logits=`** (alias `cell_informed`), default `True` (R2) | Phase 5 |\n| `point_estimate=`; public `posterior=` on metrics | **removed** \u2014 `n_samples` is the only point/draws knob (R1\u2013R2) | Phase 5\u20136 |\n| `flux` / `from_this` / `to_that` | `phenotypic_flux` / `cov_from` / `cov_to` | Phase 6 |\n| \u26a0 `clonality` | **DROP entirely \u2014 do NOT merge into `clonotypic_entropy`** (generic repertoire stat) | Phase 6 |\n| \u26a0 `clonotypic_entropy_base` | **DROP** (not merged); log base standardized via `_distance` (bits) | Phase 6 |\n| \u26a0 `ridge_delta_entropy` | **DROP** (not keep\u2011and\u2011fix) | Phase 7 |\n| \u26a0 `compare_phenotypes` | **DROP** (not move\u2011to\u2011examples) | Phase 6 |\n| `mi_compare`, `delta_entropy_table`, `flux_table`, `delta_clonotypic_entropy`, `phenotypic_entropy_delta` | **deleted** \u2014 expressed via `groupby=` + `tl.compare_groups` | Phase 6 |\n| `tl.phenotypic_entropies` / `tl.clonotypic_entropies` (plural batch forms) | **deleted** \u2014 subsumed by `groupby=` on the singular metric (notebook rewrite, Phase 10) | Phase 6/10 |\n| `get_cell_phenotype_probs` | `predict` (scvi/CellAssign idiom; order\u2011preserving loader, indexed by `obs_names`) | Phase 4 |\n| `register_model` (+ `classify_phenotypes`, `register_*_key`) | `model.to_anndata` (thin) | Phase 4 |\n| `register_clonotype_key` / `register_phenotype_key` (34 notebook uses) | folded into `setup_anndata`/`to_anndata`; **notebook\u2011rewrite mapping batched to Phase 10** (writer retired only when every reader is migrated) | Phase 4/10 |\n| `pl.clonotypic_entropy_by_phenotype` | `pl.clonotypic_entropy` | Phase 7 |\n| `plot_pheno_sankey` | `pl.phenotypic_flux` (sankey) | Phase 7 |\n| `plot_phenotype_probabilities` | `examples/phenotype_probabilities` \u2014 **reads `K.X_PROBABILITIES` ('X_tcri_probabilities')**, not the retired `X_tcri_phenotypes` slot | Phase 10 |\n| `probability_ternary` (24 notebook uses) | **keep** as `pl.probability_ternary` (phenotype\u2011simplex viz) \u2014 *explicitly dispositioned; was undispositioned in the draft* | Phase 7 |\n| `gene_entropy` (5 notebook uses) | **relocate to `examples/`** \u2014 *NOT a 0\u2011caller; removed from the Phase\u20112 delete list* | Phase 10 |\n| `polar_plot` (1 notebook use) | **DROP** with an in\u2011PR notebook\u2011cell rewrite (not a silent Phase\u20112 delete) | Phase 2/10 |\n| `pl.flux` boxplot (5 notebook uses) | **DROP** (sankey is the flux plot) \u2014 notebook rewrite | Phase 7/10 |\n| `tcri_boxplot` | `_metric_boxplot` (private; keeps a **`splitby=` axis**, \u00a7Phase 7) | Phase 7 |\n| `set_color_palette` | `resolve_palette` | Phase 7 |\n| `dkl`, `flux.dkl_func` | `_distance.kl_divergence` (log2/bits, single eps) + `_distance.jensen_shannon` | Phase 1 |\n| `\u0394` (unicode) | `delta` (ASCII, greppable) | Phase 6 |\n| `c2p_mat` | `clone_phenotype_prior` | Phase 3 |\n| `centropy` / `pentropy` / `*_tl` leaked aliases | removed via `__all__` | Phase 11 |\n| `uns[\"tcri_clone_key\"]`/`[\"tcri_phenotype_key\"]` **and** `uns[\"tcri_metadata\"][...]` (two schemes) | one scheme via `_keys.py` (single `tcri_metadata`) | Phase 1 |\n\n---\n\n## 4. Shared\u2011helper extraction (`_keys` / `_console` / `_stats` / `_distance`)\n\nLands as **Phase 1**, before any file move, because every later phase imports these. Pure internal dedup; public API unchanged except for the new public `tl.compare_groups` (Phase 6, built on `_stats`).\n\n| new module | absorbs | notes / fixes folded in |\n|---|---|---|\n| `tcri/_keys.py` | every `uns/obsm/obs` key literal, both current schemes | constants only; unifies `tcri_clone_key`/`tcri_metadata['clone_col']` to one `tcri_metadata` scheme; **adds `X_LOGITS`, `GATE_PROB`, `CLASSIFIER_TEMPERATURE`, `LOCAL_SCALE`** (R4/R6/R8); `clone_size` and any other reader flips to `K.CLONE_COL` in the **same** change that retires the writer (no orphaned reader) |\n| `tcri/_console.py` | triplicated `_ok/_info/_warn/_fin` (metrics/pp/pl copies) | reimplement over **scanpy's `logging`/verbosity**; drop raw ANSI; **drop `_ascii_hist` and every `graph=`/ASCII\u2011histogram code path** |\n| `tcri/_stats.py` | `stars`, `auc_and_label_permutation`, `bootstrap_auc` (from `utils/_utils.py`) | **add** `mann_whitney`, `prob_gt_zero` (Bayesian P(>0), for signed contrasts only \u2014 R5), **`hdi` implemented as a TRUE highest\u2011density interval** (not the mislabeled equal\u2011tailed `percentile[2.5,97.5]`); document boundary instability for bounded skewed posteriors |\n| `tcri/_distance.py` | module\u2011level dead `dkl` + `flux.dkl_func` | one `kl_divergence` + `l1_distance` + **`jensen_shannon`** + `phenotype_distance(metric=...)`; **fix the mixed\u2011units bug \u2014 standardize on `log2` (bits) across entropy/MI/KL with one eps**; document `dkl` as directional/unbounded and recommend JSD (bounded \u22641 bit) for symmetric shift; `l1` stays the safe bounded default |\n\n`_mi_from_joint` stays module\u2011private in `tools/_mutual_information.py`. The **public comparison surface** is `tl.compare_groups` (\u00a76) \u2014 not the private `_stats` primitives \u2014 so \"comparisons via `groupby` + stats\" is programmatically reproducible.\n\n---\n\n## 5. Model \u2192 AnnData streamline\n\nThe single highest\u2011risk behavior change, and the one that makes the metric\u2194model agreement guarantee (R8) achievable. It kills the `AnnDataManager`\u2011in\u2011`uns` hack and fixes the write\u2011set the draft under\u2011specified.\n\n**5.1 `setup_anndata` \u2014 registration only (no *analysis* obs mutation).**\n- Registers fields via scvi `REGISTRY_KEYS`; **keeps writing/registering `obs['indices']`** (training glue, `_model.py:678/684`) \u2014 this is not an analysis mutation and must not be removed.\n- **Removes the manager stash `adata.uns['tcri_manager'] = adata_manager` (`_model.py:697`)** \u2014 *this is where the stash actually lives, not in `register_model`.* Removing it here is what lets `write_adata_safely`/`_pop_nonserializables` be deleted.\n- `group_singletons` stays a **separate `pp` step that must run BEFORE `setup_anndata`** (it relabels clones; running it after desyncs `ct_to_c`/`p_ct` from `obs`). Enforced: `setup_anndata` errors if a later relabel is detected.\n\n**5.2 `model.to_anndata(adata)` \u2014 thin, canonical write\u2011set.** Writes **only** the canonical minimum via `_keys`, and the canonical minimum now **explicitly includes the three items the engine needs for `predict` parity** (correcting the draft's \"nothing else\"):\n\n| slot | key | why it is canonical |\n|---|---|---|\n| metadata/categories | `K.META`, covariate/phenotype/ct category maps | registry provenance |\n| latent | `obsm[K.X_TCRI]` | embedding |\n| phenotype probs + hard labels | `obsm[K.X_PROBABILITIES]`, `obs[...]` | `predict()` output; standard slot (retire `X_tcri_phenotypes`) |\n| ct\u2011level prior mean | `uns[K.P_CT]` = **raw** `normalize(q_p_ct_raw)` (R6) | `n_samples=0` closed\u2011form mean |\n| **per\u2011cell logits** | **`obsm[K.X_LOGITS]`** | `use_logits=True` engine path (R8) \u2014 hard\u2011required |\n| **gate probability** | **`uns[K.GATE_PROB]`** (scalar or `None`) | gate\u2011aware combine parity with `predict` (R8) |\n| **classifier temperature** | **`uns[K.CLASSIFIER_TEMPERATURE]`** | matches `predict`'s logit scaling (R8) |\n| **local scale** | **`uns[K.LOCAL_SCALE]`** | draw variance for `n_samples>0`; engine **raises** (never defaults to 1.0) if missing when `n_samples>0` (R4/R8) |\n\n- **Stops writing `uns['tcri_manager']`** (already removed in 5.1) \u2192 deletes the `write_adata_safely`/`_pop_nonserializables` hack.\n- `predict` (renamed from `get_cell_phenotype_probs`): returns a per\u2011cell phenotype\u2011prob `DataFrame`; **asserts the inference `DataLoader` is order\u2011preserving (`shuffle=False`) and indexes by `adata.obs_names`** (or carries the registered `indices` field and reindexes) so ct assignment and barcode labels cannot drift.\n\n**5.3 Session IO.** `write_adata_safely \u2192 save_tcri_session` writes a **plain h5ad** (nothing non\u2011picklable in `uns` anymore); `load_tcri_session` rebuilds the registry by re\u2011running `setup_anndata`. `get_p_ct` reads the **process\u2011global** pyro param store (`q_p_ct_raw`); load must set the store immediately before any `get_p_ct`/`to_anndata` call, and multi\u2011model/round\u2011trip diag workflows are documented as single\u2011model\u2011per\u2011process unless params are namespaced (\u00a711).\n\n---\n\n## 6. Engine + metric numeric contract\n\nThis section is the substrate all metrics read; it folds in every math/stats and missing\u2011link correction.\n\n**6.1 Engine `tools/_joint.py::joint_distribution`** (thin DataFrame wrapper over `_compute/_joint.py::_joint_draws`, \u00a77):\n\n```\njoint_distribution(adata, *, covariate=None, groupby=None, n_samples=0,\n use_logits=True, clones=None, temperature=1.0,\n random_state=None) -> pandas.DataFrame\n```\n- Unifies `joint_distribution` + `joint_distribution_posterior`; re\u2011exported as `tcri.joint_distribution`. Provenance in a **serializable form** (a `params` column or a companion `uns` sidecar), **not** only `df.attrs` (R\u2011forward: h5ad round\u2011trips must not silently drop it, \u00a711).\n- `n_samples=0`: closed\u2011form path. `use_logits=False` \u2192 tempered `m\u0303` rows (`== uns[K.P_CT]` at `T=1`); `use_logits=True` (default) \u2192 per\u2011cell `softmax((logits + gate\u2011combine(log m\u0303))/T)` aggregated per clone, **identical to `model.predict`** (R6/R8).\n- `n_samples>0`: draw `p_ct ~ Dirichlet(clamp(local_scale\u00b7m\u0303, 1e-3))` **once per sample via a seeded torch generator** (R4/R7), feed each draw through the same temperature + combine + scatter\u2011sum; stack a `sample_id` axis. **All clones within one draw share the SAME `p_ct` draw** (one coherent joint per `sample_id`) \u2014 never independent per\u2011clone draws.\n- **`covariate=None` computes the joint across ALL covariate values in one pass from a single `p_ct` draw** (the all\u2011timepoints path the sankey and multi\u2011covariate metrics need).\n- **Draw\u2011once invariant:** for `n_samples>0`, the number of Dirichlet draws is exactly `n_samples`, **independent of `#groups` and `#covariates`** \u2014 draws are reused across groups/covariates by cell\u2011masking, not re\u2011drawn. Enforced by a draw\u2011counter test (\u00a79).\n- **`groupby` is implemented by cell/clone RESTRICTION over the FULL `adata`** (positional masks into full\u2011space `uns` arrays + `clones=`), **never by passing a sliced `AnnData` to the engine** \u2014 this avoids the hard full\u2011space\u2011vs\u2011subset alignment guard that today's `tcri_boxplot` slicing would trip. `_metric_boxplot` is rewritten off the slice\u2011and\u2011call pattern.\n- **`groupby` requires the cell\u2011informed path** (or group keys that are clone\u2011nested / constant within `clone \u00d7 covariate`); it is **ill\u2011defined on the ct\u2011level table** for non\u2011clone\u2011determined columns. The design assumes **clones are disjoint across `groupby` groups** (a TCR clone does not span two patients) \u2014 now stated explicitly, with a validation that errors when a group split would bisect a clone's cells.\n- **Engine bug fixes folded in:** weighting keyed on the **`ct` index** (not the clone index); consistent normalization (row\u2011 vs whole\u2011table) across the two old functions; all\u2011zero\u2011clone reindex yields **NaN, not inflated uniform entropy**; **torch\u2011seeded** determinism.\n\n**6.2 Four metrics** (`tools/_entropy.py`, `_mutual_information.py`, `_flux.py`), uniform signatures; **none expose `use_logits`/`posterior=`/`point_estimate=`**:\n```\ntl.clonotypic_entropy(adata_or_jd, *, covariate=None, groupby=None, n_samples=0, temperature=1.0,\n clones=None, normalised=True, random_state=None)\ntl.phenotypic_entropy (adata_or_jd, *, covariate=None, groupby=None, n_samples=0, temperature=1.0,\n clones=None, normalised=True, random_state=None)\ntl.mutual_information (adata_or_jd, *, covariate=None, groupby=None, n_samples=0, temperature=1.0,\n clones=None, normalised=True, normalise_mode='min', random_state=None)\ntl.phenotypic_flux (adata, *, cov_from, cov_to, groupby=None, n_samples=0, temperature=1.0,\n clones=None, distance_metric='l1', random_state=None)\n```\n- **`mutual_information` default `normalise_mode='min'` (was `'average'`).** Under the settled uniform\u2011clone prior, `P(c)=1/C` pins the clone marginal entropy `H_c` to `log2(C)` (structural, uninformative), which throttles `'average'` normalization by `~1/log2(C)` and makes it non\u2011comparable across groups with different clone counts \u2014 breaking the whole `groupby` workflow. `'min'` (`I/H_p`, coefficient of constraint) is in `[0,1]`, reaches 1 when clone determines phenotype, and is `C`\u2011independent. `'average'` is documented as not\u2011recommended (or dropped).\n- **`clonotypic_entropy` normalizer comparability:** divide by `log2` of the number of clones with **genuine support**, dropping requested\u2011but\u2011absent (all\u2011zero, reindexed) clones from `C`; expose a **common\u2011denominator option** (`log2` of a fixed reference clone count) for cross\u2011group plots. Document that normalized clonotypic entropy uses group\u2011specific denominators by default.\n- **Dual input.** Each accepts `(adata + covariate)` [compute joint internally] **or** a precomputed joint `DataFrame` [fast path]. **A bare precomputed joint forces `n_samples=0` and `groupby=None`** (it carries no `p_ct`/`local_scale`/logits/cells to resample or re\u2011partition); `clones=` just re\u2011filters rows; **raise a clear error** for `n_samples>0`/`groupby` on a bare jd. `phenotypic_flux` correctly takes `adata` only (it needs two joints). Metrics **propagate the input joint's provenance** into their output.\n- **Return\u2011shape rule (uniform):** no `groupby` & `n_samples=0` \u2192 scalar/Series; `groupby` set \u2192 tidy `DataFrame` (row per group [\u00d7 phenotype/clone]); `n_samples>0` \u2192 adds a `sample_id` axis and, on reduction, summary columns **`mean`, `sd`, `hdi_low`, `hdi_high`** (**no `p_gt`** \u2014 R5). Metrics reduce the stack by **iterating the `sample_id` level** (per\u2011draw full\u2011joint metric, then summarize).\n- **`n_samples>0` intervals are partial posterior** (only `p_ct` uncertainty; classifier logits held at their posterior\u2011mean encoding) \u2014 documented so users don't read them as full posterior\u2011predictive uncertainty.\n- **h5ad\u2011serializable returns (build\u2011toward `@tl_result`):** flat columns, **no object\u2011array columns and no `df.attrs`\u2011only provenance**; per\u2011draw vectors go in a separate long frame or a `uns` sidecar. Define the cache key as a hash of `(covariate, groupby, n_samples, temperature, clones, normalised, normalise_mode, distance_metric, random_state)`.\n\n**6.3 Comparisons via `groupby` + `tl.compare_groups` (public).** The deleted `*_compare`/`*_delta` functions are replaced by a **public mid\u2011level helper**, not private primitives:\n```\ntl.compare_groups(df, *, value, by, reference=None, paired=False, hdi_prob=0.94)\n -> tidy DataFrame with per-pair: mean_a, mean_b, delta, U, p (Mann-Whitney),\n p_gt (Bayesian P(delta>0)), hdi_low, hdi_high\n```\nThis is where `p_gt`/HDI live (on the **signed** `delta`, R5). The docstring shows the recipe that recreates `mi_compare`'s per\u2011pair output, so \"groupby + stats subsumes the tables/deltas\" is actually reproducible.\n\n---\n\n## 7. GPU / optimization architecture\n\nGrounded 1:1 in grafiti's `_compute/` wins. The **one architecturally load\u2011bearing decision**: write the engine's numeric core as a **batched, device\u2011routable function returning a `[n_samples, n_clones, P]` array**, with pandas only at the boundary \u2014 so every win below is designed *in*, not retrofitted later.\n\n**7.1 Device seam \u2014 `tcri/_compute/_xp.py` (copy grafiti's 58\u2011line reference).** `resolve_device` / `get_xp` / `asnumpy`. Because **torch\u22652.4.1 is already a hard dep**, a **torch\u2011tensor core is the first backend** (torch.cuda when present \u2014 zero new deps), with **cupy optional** as a numpy\u2011style second backend. Every accelerated function returns a plain numpy array via `asnumpy`. Device ladder: `None/'cpu'\u2192cpu`; `'mps'\u2192cpu`; `'auto'/'gpu'/'cuda'\u2192GPU only if the lib imports **and** `getDeviceCount()>0`, else CPU (explicit `'cuda'` warns on fallback, `'auto'/'gpu'` silent).\n\n**7.2 Batched engine core \u2014 `_compute/_joint.py::_joint_draws(...) -> ndarray[n_samples, n_clones, P]`.** Precompute clone integer codes **once**; batch the Dirichlet draw and the softmax on the leading sample axis; scatter\u2011sum by clone. `tools/joint_distribution` is a thin single\u2011draw/summary DataFrame wrapper over this stack. Metrics consume the stack and reduce vectorized (`_compute/_reduce.py`).\n\n**7.3 The wins, priority\u2011ordered, with expected gains:**\n\n| # | operation | current hot path | fix | expected gain | prio |\n|---|---|---|---|---|---|\n| P0 | joint\u2011by\u2011clone reduction | `pd.DataFrame(...).groupby(level=0).sum()` rebuilt every draw (`_preprocessing.py:320-322`) over 1e5\u20131e6 cells | precompute clone codes once; **scatter\u2011add** (`np.add.at`/`np.bincount` CPU, `torch.index_add_`/`cupy.bincount` GPU) batched across all `n_samples` \u2014 grafiti `contingency.py`/`edge_tensor` verbatim | **10\u201350\u00d7 on CPU** (pandas groupby \u2192 integer\u2011keyed bincount), multiplied again on GPU; the single biggest win | P0 |\n| P0/P1 | Dirichlet sampling loop + softmax | `for i in range(n_samples): joint_distribution_posterior(...)` (`_metrics.py:296/363/548/741/987`), redrawing the **full** `p_ct` each call | draw all `n_samples` at once, **restricted to the covariate's ct rows first**; batch `softmax((logits+log b)/T)` over the leading axis; route via torch.cuda | removes the `n_samples`\u2011fold Python loop + per\u2011iteration `.uns`/DataFrame setup; collapses 200\u20131000 iterations into a few kernels | P0/P1 |\n| P1 | entropy / MI reductions | `scipy.stats.entropy` per draw + `jd.loc[cl]` per clone (`_metrics.py:315/562-569/744`) | batched `xlogx` reduction over `[S, n_clones, P]`; MI as joint\u2011vs\u2011outer\u2011product; `nanmean`/HDI over the sample axis \u2014 grafiti `joint.py::_entropy/_mi`, float64 accumulators | removes two nested Python loops + slow pandas `.loc`; medium\u2011high | P1 |\n| P1 | share the joint across metrics | each of the 4 metrics rebuilds the joint for the same `(covariate, n_samples)` | compute the `[S, n_clones, P]` stack **once per covariate** and reduce it in all four metrics; **build `groupby` groups in one batched pass** | divides the dominant cost by `#metrics`; realizes the draw\u2011once invariant (\u00a76.1) | P1 |\n| \u2014 | deterministic point estimate | `n_samples=0` still enters the sampler / returns one draw today | zero\u2011draw read of `uns[K.P_CT]` (R1) | correctness **and** cheapest path (no Monte\u2011Carlo at all) | \u2014 |\n| P2 | latent UMAP | `umap.UMAP` with `import umap` at module top (`_preprocessing.py:20`) | `cuml.manifold.UMAP` behind `_use_gpu`, umap\u2011learn CPU fallback, **lazy import moved inside the function**; return float64 ndarray | cuML UMAP commonly **10\u201350\u00d7** umap\u2011learn, but runs once per analysis \u2192 lower total impact; also fixes the module\u2011top heavy import | P2 |\n\n**7.4 Grafiti guardrails (replicate ALL):** (1) every GPU lib imported **lazily inside** the function \u2014 `import tcri` never touches cupy/cuml/torch.cuda (the current `import umap` at module top already violates this and is fixed); (2) GPU deps never in `install_requires`; (3) permissive device ladder with `getDeviceCount()>0` verification; (4) `asnumpy` at every return boundary; (5) GPU body wrapped in `try/except` that degrades to CPU and reports which backend ran; (6) **float64 accumulators** so GPU matches CPU; (7) **validate the joint** (finiteness, nonnegativity, per\u2011row sum\u22481 \u2014 per\u2011row, to catch cancelling defects) **on\u2011device before compute**; (8) **chunk the batched reduction** over cells/draws (grafiti Moran's\u2011I `chunk_size=256`, KDE `blk=8192`) to bound device memory on the large `[n_samples, n_cells, P]` tensor and avoid OOM.\n\n**7.5 Seeding (R7).** Draws move to a seeded `torch.Generator` (and cupy/torch.cuda RNG on GPU); `random_state` is threaded from every public sampling function. The old `np.random.seed` was a no\u2011op for the torch draws.\n\n---\n\n## 8. Ordered PR sequence\n\nEach PR is independently landable with green CI, lowest\u2011risk first.\n\n### Phase 0 \u2014 Contract freeze + CI scaffolding + caller census *(docs/tests only; zero code change)*\n- Frozen contract: `tcri/_contract.pyi` + `tests/test_contract_conformance.py`, ported from grafiti's `_pyi_gen.py`/`test_contract_conformance.py`. Seed `IMPLEMENTED = {}`; each target function flips to implemented as its phase lands. Markdown\u2192`.pyi`\u2192live\u2011signature drift fails CI.\n- **Caller census over `example/` + `docs/` notebooks** (not package source only). Reclassify the draft's \"12 confirmed dead\": `gene_entropy` (5 uses), `polar_plot` (1), `register_*_key` (34), `pl.flux` boxplot (5), `probability_ternary` (24) are **not** safe deletes. Land the corrected disposition map (\u00a73) into `docs/contract/`.\n- **Risk:** none. **Depends on:** nothing.\n\n### Phase 1 \u2014 Shared helpers + `_keys` adoption *(internal dedup; API unchanged)*\n- Create `_keys.py` (incl. `X_LOGITS/GATE_PROB/CLASSIFIER_TEMPERATURE/LOCAL_SCALE`), `_console.py`, `_stats.py` (true HDI, `prob_gt_zero`), `_distance.py` (bits/log2, JSD) (\u00a74).\n- **Adopt `_keys` at every read/write site.** Migrate `clone_size` to `K.CLONE_COL` in the **same** change that retires `tcri_clone_key` (no orphaned reader).\n- **Risk:** low (mechanical, no numerics). **Depends on:** Phase 0. **Hazard:** must precede Phases 4/5.\n\n### Phase 2 \u2014 Safe deletions *(only symbols unreferenced in package AND notebooks)*\n- Delete the genuinely dead: `get_latent_embedding`, `group_small_clones`, `register_probability_columns`, `remove_meaningless_genes`, `metrics._ent`, `clone_fraction`, module\u2011level `dkl`, **`utils.probabilities` (and the `_plotting.py:18` import in the SAME PR)**, `pl.probability_distribution`, `pl.bayesian_mutual_information`, `SankeyNode.hex_to_rgb`.\n- **Removed from the draft's delete list:** `gene_entropy` (\u2192 examples, Phase 10), `polar_plot` (\u2192 drop with notebook\u2011cell rewrite). **Not deleted here:** `compare_joint_distribution` (\u2192 diag, Phase 8), `pl.mutual_information`/`pl.phenotypic_entropy` (keep+fix, Phase 7), and the consolidated\u2011away `*_table`/`*_delta`/`clonality`/`clonotypic_entropy_base`/`ridge_delta_entropy` (removed **with** their replacement in Phase 6, so no caller is orphaned).\n- **Risk:** very low. **Depends on:** Phase 1.\n\n### Phase 3 \u2014 Model module split *(mechanical; no behavior change)*\n- Split `model/_model.py` (1074 ln) \u2192 `_model.py` + `_module.py` + `_priors.py` + `_classifier.py` + `_training.py`. Rename `c2p_mat \u2192 clone_phenotype_prior`.\n- **`build_archetypes` keeps returning `(centers, labels)`** (labels drive `diag.archetypes`' cluster ordering); persist labels on the model/`uns`. Reconcile the default\u2011`K` mismatch (`build_archetypes` default `K=4` vs model `K=10`).\n- Explicit `__all__` per module. **Risk:** low. **Depends on:** Phase 1. **Verify:** `test_model_setup`, `test_pyro_params` unchanged\u2011green.\n\n### Phase 4 \u2014 Model\u2192AnnData streamline *(behavior change; kills the manager hack)*\n- Implement \u00a75: `setup_anndata` registration\u2011only (keeps `obs['indices']`; **removes the `_model.py:697` manager stash**); `group_singletons` enforced to precede `setup_anndata`.\n- `register_model \u2192 model.to_anndata` writing the canonical set **including `X_LOGITS`, `GATE_PROB`, `CLASSIFIER_TEMPERATURE`, `LOCAL_SCALE`, and the raw `P_CT`** (R6/R8). `get_cell_phenotype_probs \u2192 predict` (order\u2011preserving loader, `obs_names` index).\n- `write_adata_safely \u2192 save_tcri_session` (plain h5ad); load rebuilds the registry via `setup_anndata` and sets the pyro param store before any `get_p_ct`.\n- **Risk:** HIGH. **Depends on:** Phase 1, Phase 3. **Gate:** `test_session_round_trip` rewritten \u2014 proves save/load reproduces `p_ct` + latent + `predict` probs with **no `tcri_manager` in `uns`**, `setup_anndata` leaves analysis `obs` untouched, and `to_anndata` writes **exactly** the canonical key set (asserts logits/gate/cls\u2011temp/local\u2011scale present).\n\n### Phase 5 \u2014 Engine consolidation *(the substrate; \u00a76.1, \u00a77.2)*\n- Create `tools/` (aliased `tl`) and `_compute/` (`_xp.py`, `_joint.py`, `_reduce.py`). Implement `_joint_draws` (batched, device\u2011routable, scatter\u2011add) and the thin `tools/joint_distribution` wrapper; re\u2011export as `tcri.joint_distribution`.\n- Signature per \u00a76.1 with **`use_logits`** (renamed from `posterior=`), `random_state`, `covariate=None` \u2192 all\u2011covariates one\u2011pass, draw\u2011once invariant, clamped\u2011Dirichlet draws, single\u2011knob temperature, groupby by full\u2011space restriction. Fold in the weighting/normalization/zero\u2011clone/seed bug fixes.\n- **Risk:** HIGH \u2014 every metric reads this. **Depends on:** Phase 4. **Verify:** `test_tools/test_joint`: `use_logits=False & n_samples=0 == tempered uns[K.P_CT]` **exactly**; `use_logits=True & n_samples=0 & T=1 == model.predict` aggregation; repeated `n_samples=0` **bit\u2011identical**; `n_samples>0` **torch\u2011seeded reproducible** and drawn from `Dirichlet(clamp(local_scale\u00b7m\u0303,1e-3))`; **draw\u2011counter == n_samples independent of #groups/#covariates**; weighting keyed on `ct`; serializable provenance.\n\n### Phase 6 \u2014 Metric\u2011API consolidation *(four metrics + `compare_groups`)*\n- Populate `tools/_entropy.py`, `_mutual_information.py`, `_flux.py` per \u00a76.2 (dual input, `normalise_mode='min'` default, support\u2011only clonotypic denominator + common\u2011denominator option, coherent per\u2011`sample_id` draws, `mean/sd/hdi` summary with **no `p_gt`**, serializable returns). Add public **`tl.compare_groups`** (\u00a76.3).\n- **Delete** `mi_compare`, `delta_entropy_table`, `flux_table`, `delta_clonotypic_entropy`, `phenotypic_entropy_delta`, `clonotypic_entropy_base`, `clonality`, `ridge_delta_entropy`, `compare_phenotypes`, the plural `*_entropies`, and `metrics.dkl`. Delete `metrics/` after migration.\n- **Risk:** HIGH. **Depends on:** Phase 5. **Verify:** ranges (`[0,1]` normalized), `n_samples=0` determinism, `groupby` tidy shape, **dual\u2011input equivalence at `n_samples=0` only**, all\u2011zero\u2011clone \u2192 NaN, `phenotypic_flux` `cov_from/cov_to` + seeded draws, `compare_groups` recreates `mi_compare`'s per\u2011pair output, **no test asserts `n_samples=0 == mean(n_samples>0)`** (R3).\n\n### Phase 7 \u2014 Plotting split + pl twins *(fix the broken core plots)*\n- Split `plotting/_plotting.py` (1437 ln) \u2192 `_base.py` (`_metric_boxplot` **with a `splitby=` axis**, `_finish`), `_colors.py` (`resolve_palette`), `_entropy.py`, `_mutual_information.py`, `_flux.py`, `_sankey.py`.\n- Ship the four tl\u2194pl twins: `pl.clonotypic_entropy` (was `_by_phenotype`), `pl.phenotypic_entropy` **[FIX]**, `pl.mutual_information` **[FIX]**, `pl.phenotypic_flux` (sankey). Keep `pl.probability_ternary`. pl functions are **cache renderers** (no metric math). **Retain `splitby`** (116 notebook uses) so two\u2011axis figures (dots=patient, boxes=response, x=phenotype) survive; document per\u2011figure the `groupby`+`splitby` recipe.\n- **Changelog behavior notes:** removing `pl.mutual_information`'s `weighted=True` default enshrines the uniform\u2011clonotype prior \u2192 user\u2011visible MI numbers change (cell\u2011weighted \u2192 per\u2011clonotype); `pl.flux` boxplot and `clonality` plot dropped.\n- **Move to `examples/`:** `top_clone_umap`, `clone_size_umap`. **Risk:** medium. **Depends on:** Phase 6, Phase 1. **Verify:** each twin returns a `Figure`/`Axes` from a tidy tl result; sankey renders.\n\n### Phase 8 \u2014 `diag/` seeding *(new; additive)*\n- `diag/_ppc.py`: the **fixed** `compare_joint_distribution` (no undefined\u2011global `NameError`; model `p(clone,phenotype)` vs empirical counts), phenotype\u2011probability calibration, reconstruction PPC, entropy/MI vs permutation null. **All return DataFrames.** Make the **model requirement explicit per function**: `joint_distribution_ppc`/calibration/permutation\u2011null run **adata\u2011only**; `reconstruction_ppc` **requires the live model** (ZINB decoder lives on the module). `diag/_training.py`: `plot_loss \u2192 loss curves`, `plot_archetypes \u2192 archetypes` (consumes `build_archetypes` labels).\n- **Risk:** low\u2011medium. **Depends on:** Phase 4 (finalized model), Phase 5 (engine). **Verify:** each PPC returns the expected columns on the `trained_model` fixture; permutation\u2011null seeded; single\u2011model\u2011per\u2011process param\u2011store scoping honored.\n\n### Phase 9 \u2014 PGM \u2192 docs; utils finalize\n- Move `build_nested_tcri_pgm`/`draw_tcri_pgm_nested` **out of the package into `docs/`**; drop `daft` from runtime deps \u2192 docs extras only. `utils/_utils.py \u2192 utils/_session.py` (session\u2011io + `_to_jsonable` only).\n- **Risk:** low. **Depends on:** Phase 1, Phase 8.\n\n### Phase 10 \u2014 Notebook + examples rewrite *(the real end\u2011to\u2011end acceptance test)*\n- Rewrite `example/` notebooks (`smith`, `renal`, `yost`, `zhang`, `synthetic`, `colitis`) to the new API: `setup_anndata \u2192 TCRIModel \u2192 train \u2192 model.to_anndata`; `tcri.joint_distribution`; four metrics with `groupby`/`n_samples`; `tl.compare_groups`; four pl twins; `diag` checks. Apply the **notebook\u2011rewrite mappings** for the retired `register_*_key` (34), `gene_entropy` (\u2192 examples), `polar_plot`, `pl.flux` boxplot, and the `X_tcri_phenotypes \u2192 K.X_PROBABILITIES` key fix in the moved `phenotype_probabilities` example. One canonical end\u2011to\u2011end tutorial notebook under `docs/`.\n- **Risk:** low (docs), high value. **Depends on:** Phases 4\u20138.\n\n### Phase 11 \u2014 Public API finalize + scverse ecosystem CI\n- `tcri/__init__.py`: **explicit** named re\u2011exports, `sys.modules` aliases for `tl/pp/pl/ml/diag/ut`, top\u2011level `tcri.joint_distribution`, **remove `import *`**, kill leaked aliases (`centropy`/`pentropy`/`*_tl`) via `__all__`. Flip **all** target functions to `IMPLEMENTED`; drift now hard\u2011fails CI. Turn on the full scverse\u2011ecosystem gate (\u00a79.2).\n- **Risk:** low\u2011medium. **Depends on:** all prior phases.\n\n---\n\n## 9. Testing + scverse\u2011CI strategy\n\n### 9.1 Tests per phase\n\n| phase | required tests |\n|---|---|\n| 0 | `test_contract_conformance` (markdown\u2192`.pyi`\u2192signature drift); import\u2011smoke py3.10/3.11; **notebook caller\u2011census fixture** feeding the disposition map |\n| 1 | `_stats` (`stars`, AUC/permutation, **true HDI vs equal\u2011tailed**, `prob_gt_zero` on a signed vector), `_distance` (`kl_divergence` bits/symmetry, JSD bound), `_console` silence flag, `_keys` \"no stray literal\" grep |\n| 2 | every deleted symbol absent from `__all__` **and** unreferenced (import\u2011graph test, package **+** notebooks); `import tcri` green after `utils.probabilities` + its `_plotting.py:18` import go together |\n| 3 | `test_model_setup`, `test_pyro_params` green through the split; `build_archetypes` returns `(centers, labels)`; submodule import smoke |\n| 4 | **`test_session_round_trip` rewritten** \u2014 save/load reproduces `p_ct`+latent+`predict` probs with **no `tcri_manager`**; `setup_anndata` leaves analysis `obs` untouched but keeps `obs['indices']`; `to_anndata` writes **exactly** the canonical set incl. `X_LOGITS/GATE_PROB/CLASSIFIER_TEMPERATURE/LOCAL_SCALE`; `predict` order\u2011preserving |\n| 5 | `test_tools/test_joint`: `use_logits=False,n=0 == tempered uns[P_CT]`; `use_logits=True,n=0,T=1 == predict` aggregation (**engine==predict** agreement test); repeat `n=0` bit\u2011identical; `n>0` torch\u2011seeded + clamped\u2011Dirichlet; **draw\u2011count == n_samples \u27c2 #groups/#covariates**; weighting on `ct`; `groupby` via full\u2011space restriction does not trip the alignment guard |\n| 6 | ranges, `n=0` determinism, `groupby` tidy shape, **dual\u2011input equivalence (n=0 only)**, all\u2011zero\u2011clone\u2192NaN, `normalise_mode='min'` C\u2011independence, `compare_groups` recreates `mi_compare`; **golden regression** computing MI/entropy from a fixed `uns[P_CT]` by hand; **no `n=0==mean(n>0)` assertion** |\n| 7 | each pl twin returns `Figure`/`Axes` from a tidy tl result without metric math; `splitby` two\u2011axis render; sankey renders; `probability_ternary` renders |\n| 8 | each PPC returns expected columns on `trained_model`; `reconstruction_ppc` requires model, others adata\u2011only; permutation\u2011null seeded |\n| 9 | no `daft`/PGM import in the installed package (import\u2011graph test) |\n| 10 | **notebook execution** (`pytest --nbmake`) on the synthetic tutorial end\u2011to\u2011end; retired\u2011symbol mappings applied |\n| 11 | contract test with **all** functions `IMPLEMENTED`; `__all__` completeness (public names \u2194 contract); \"no `import *`\" AST test |\n\nReuse `conftest.py` fixtures (`synthetic_adata`, `trained_model`, `mock_adata`); extend `mock_adata` to the unified `_keys` scheme in Phase 1 and to `to_anndata`'s canonical key set (incl. logits/gate/cls\u2011temp/local\u2011scale) in Phase 4.\n\n### 9.2 scverse\u2011ecosystem CI\n\nBring `.github/workflows/tests.yml` to the cookiecutter\u2011scverse bar, layered so each phase stays green:\n- **Matrix** py3.10 + py3.11 (extend to 3.12 before 1.0), `pip install -e \".[test]\"`, `pytest tests/ -v --cov=tcri`.\n- **Import\u2011smoke job** (grafiti pattern): `python -c \"import tcri; from tcri.model._model import TCRIModel; from tcri.tools._joint import joint_distribution\"` \u2014 **and assert no GPU lib (cupy/cuml/torch.cuda) was imported** (guardrail #1).\n- **Lint/format gate:** `pre-commit` (ruff + ruff\u2011format) \u2014 add in Phase 1, enforce from Phase 3.\n- **Contract\u2011conformance job** (Phase 0 on) \u2014 the markdown\u2192`.pyi`\u2192signature guardrail that lets views land independently.\n- **CPU\u2011only correctness job** \u2014 the full suite must pass with no GPU present; a separate optional GPU job (if a runner is available) asserts float64 GPU\u2248CPU parity.\n- **Notebook\u2011execution job** (`nbmake`) from Phase 10.\n- **Docs build** (sphinx + numpydoc; `.readthedocs.yaml` present) must pass; public API fully docstringed/typed.\n- **Ecosystem checklist** (final gate, Phase 11): AnnData\u2011native `setup_anndata`, no `import *`, typed public surface, tutorial notebook, `diag` returns data not plots, GPU strictly opt\u2011in.\n\n---\n\n## 10. Dependencies & ordering hazards\n\n```\n0 \u2500\u25b6 1 \u2500\u252c\u2500\u25b6 2\n \u251c\u2500\u25b6 3 \u2500\u25b6 4 \u2500\u25b6 5 \u2500\u25b6 6 \u2500\u25b6 7 \u2500\u25b6 8 \u2500\u25b6 9 \u2500\u25b6 10 \u2500\u25b6 11\n \u2514\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500(1 gates 4,5,6)\n```\nHard, load\u2011bearing constraints:\n1. **`_keys` (1) before `to_anndata` (4) and the engine (5).** Both rewrite key IO; centralizing strings first keeps diffs local and prevents two\u2011scheme drift. `clone_size`'s reader flip and `tcri_clone_key`'s writer retirement must land **together**.\n2. **Model split (3) before `to_anndata` (4).** `to_anndata` is a `TCRIModel` method; split first so the behavior\u2011change diff is readable.\n3. **`to_anndata` (4) before the engine (5).** The engine reads exactly the canonical keys/**logits/gate/cls\u2011temp/local\u2011scale** `to_anndata` writes; if the write\u2011set changes after the engine exists, `use_logits=True` breaks silently. The **manager\u2011in\u2011`uns` removal (at `setup_anndata:697`) is the single highest\u2011risk hazard** \u2014 it changes the session round\u2011trip; `test_session_round_trip` is rewritten in the same PR.\n4. **Engine (5) before metrics (6).** Metrics are thin readers; building them against a pre\u2011fix engine would bake in the weighting/normalization/seed/Jensen bugs.\n5. **Metrics (6) before pl twins (7).** pl are cache renderers.\n6. **Delete `*_table`/`*_delta`/`clonality`/`_base`/`ridge_delta_entropy`/`compare_phenotypes` WITH their replacement (6), never before** \u2014 they have live notebook callers; removing ahead of `groupby`/`compare_groups` orphans callers and reds the notebook\u2011execution job.\n7. **Rename freeze before Phase 5.** `from_this/to_that \u2192 cov_from/cov_to`, `flux \u2192 phenotypic_flux`, `metrics/ \u2192 tools/`, `posterior= \u2192 use_logits`, dropping `point_estimate=` are all breaking; batch them so notebooks are rewritten once (Phase 10).\n8. **`group_singletons` before `setup_anndata`** (clone\u2011relabel ordering) \u2014 enforced by `setup_anndata`.\n9. **Deferred, designed\u2011for:** `@tl_result` uns\u2011cache + `get.py`. Every tl function returns an **h5ad\u2011serializable** tidy result (flat columns, no object arrays, serializable provenance) so the decorator is a one\u2011line wrap; until then pl recomputes via the engine. Do **not** block the refactor on the cache.\n\n**Independent / parallelizable:** Phase 2 (safe deletions) and Phase 8 (`diag`, once 4\u20135 land) proceed alongside neighbors without contending for the same files.\n\n---\n\n## 11. Risks & open items\n\n- **Highest\u2011risk hazard: the manager\u2011stash removal + write\u2011set change (Phase 4).** It reshapes the train\u2192`to_anndata`\u2192session round\u2011trip and adds four persisted keys. Mitigation: rewrite `test_session_round_trip` in\u2011PR; assert the exact canonical set; keep the CPU\u2011only correctness job as the tripwire.\n- **`use_logits=True` parity depends on persisted model facts (R8).** If any of `X_LOGITS`/`GATE_PROB`/`CLASSIFIER_TEMPERATURE`/`LOCAL_SCALE` is dropped, parity degrades silently to the additive rule. Mitigation: the engine==predict agreement test (Phase 5) and the canonical\u2011set assertion (Phase 4). Note `gate_prob` defaults to `None` (models ungated by default), so the bug is latent until someone trains with a gate \u2014 the guarantee still must hold.\n- **Estimator semantics (R3) are a documentation risk, not a code risk.** `n_samples=0` (plug\u2011in) and the `n_samples>0` `mean` are *different estimators*; users will expect them to match. Mitigation: docstrings state the Jensen gap explicitly; no conformance test asserts equality.\n- **Global pyro param store.** `get_p_ct`/`to_anndata`/diag PPCs read the process\u2011global `q_p_ct_raw`; loading two models in one process clobbers it. **Open item:** namespace params per model or ship a documented single\u2011model\u2011per\u2011process contract with load setting the store immediately before use. Pre\u2011existing, not introduced by the refactor, but the multi\u2011model diag/round\u2011trip workflows expose it.\n- **`temperature\u22601` re\u2011centers the sampled distribution** away from the learned posterior (R6); intervals under non\u2011default `T` are not the model posterior. Documented; `T=1` is the parity point.\n- **Partial posterior.** `n_samples>0` intervals capture only `p_ct` uncertainty (classifier logits fixed). Documented so they are not read as full posterior\u2011predictive uncertainty.\n- **GPU memory on large `[n_samples, n_cells, P]`.** Without chunked reductions (guardrail #8) the batching win becomes an OOM risk on realistic `n_samples \u00d7 n_cells`. Mitigation: chunk over cells/draws; try/except degrade to CPU.\n- **`probability_ternary` / `splitby` disposition** are decided here (keep) but should be confirmed against the final notebook set in Phase 0's census before Phase 7 fixes the `pl` signatures.\n- **`@tl_result` cache is deferred.** Return shapes are constrained now to be h5ad\u2011serializable so the later wrap is a one\u2011liner; the cache\u2011key scheme is specified but not implemented.\n- **cuML/rapids GPU parity** for UMAP is layout\u2011different (both valid, as grafiti documents); not bit\u2011identical to CPU. No test asserts embedding bit\u2011identity.\n- **Open (deliberately deferred):** py3.12 matrix, the `get.py` accessor surface, and any KL/JSD default change for `phenotypic_flux` beyond offering JSD (l1 remains the bounded default).\n\n---\n\n## 12. Rollout checklist (per PR)\n\n- [ ] One behavior change (or zero, if a mechanical move).\n- [ ] `import tcri` green; `tl/pp/pl/ml/diag/ut` handles intact; **no GPU lib imported at import time**.\n- [ ] No new key\u2011string literal outside `_keys.py`; no new `import *`.\n- [ ] Deleted symbols removed from `__all__` and unreferenced **in package AND notebooks** (import\u2011site + call\u2011site grep).\n- [ ] Contract test updated (`IMPLEMENTED` flipped for functions this PR lands).\n- [ ] Phase\u2011specific tests from \u00a79.1 present and green on py3.10/3.11 (CPU\u2011only job passes).\n- [ ] Rename/disposition rows for this phase applied everywhere (code + fixtures; notebooks batched to Phase 10).\n- [ ] Any persisted\u2011key change reflected in `_keys`, `to_anndata`'s canonical set, and the round\u2011trip test.\n\n---\n\n## Appendix A \u2014 Audit findings by severity (traceability)\n\nEvery finding folded into the body above, cross\u2011referenced so nothing is lost. **Src:** PC = plan\u2011correctness, MS = math/stats, ML = missing\u2011links, GPU = GPU/optimization, PVM = prior\u2011vs\u2011mean resolution.\n\n### A.1 Blocking / High\n\n| # | src | finding | resolved in |\n|---|---|---|---|\n| B1 | PC/ML | `to_anndata` \"canonical minimum\" drops `obsm['X_tcri_logits']` that the default engine path hard\u2011requires | \u00a75.2, R8, Phase 4/5 |\n| B2 | PC/ML | `gate_prob` (and `classifier_temperature`) are model attributes never persisted \u2192 gate\u2011aware parity infeasible from `adata` | \u00a75.2, R8, Phase 4 |\n| B3 | PC | `setup_anndata` cannot be \"no obs mutation\" \u2014 must keep registered `obs['indices']` | \u00a71.1, \u00a75.1, Phase 4 |\n| B4 | PC | Deleting `utils.probabilities` (Phase 2) breaks `import tcri` via `_plotting.py:18` import | \u00a71.1, Phase 2 |\n| B5 | PC | `groupby` by slicing `adata` trips the full\u2011space\u2011vs\u2011subset alignment guard | \u00a76.1, Phase 5 |\n| B6 | MS | Jensen gap: `metric(E[p]) \u2260 E[metric(p)]`; `n=0` (plug\u2011in) vs `mean(n>0)` disagree \u2014 drop the equality test | R3, \u00a71.2, Phase 6 |\n| B7 | MS | `p_gt`/P(>0) vacuous per\u2011metric (all metrics \u22650) \u2014 only on signed contrasts | R5, \u00a76.2/6.3 |\n| B8 | MS | `n>0` draws must use the guide's **clamped** `Dirichlet(clamp(local_scale\u00b7m,1e-3))` | R4, \u00a76.1, Phase 5 |\n| B9 | MS | MI `normalise_mode='average'` breaks under uniform\u2011clone prior (`H_c=log2 C`) \u2014 default `'min'` | \u00a76.2, Phase 6 |\n| B10 | ML | `_stats` is private \u2014 no public comparison entry point; add `tl.compare_groups` | \u00a74, \u00a76.3, Phase 6 |\n| B11 | ML | \"safe deletion\" census run on package only \u2014 `gene_entropy`/`polar_plot`/`register_*_key` etc. have notebook callers | \u00a73, Phase 0/2/10 |\n| B12 | ML | `covariate \u00d7 groupby` semantics unspecified; requires cell\u2011informed path / clone\u2011disjoint groups | \u00a76.1, Phase 5/6 |\n| B13 | ML | No mechanism shares one `p_ct` draw across groups/covariates \u2192 O(groups\u00d7cov\u00d7n_samples) | \u00a76.1 draw\u2011once, \u00a77, Phase 5 |\n| B14 | GPU | Engine must be a batched device\u2011routable core returning `[n_samples,n_clones,P]`, DataFrame at boundary | \u00a77.2, Phase 5 |\n\n### A.2 Medium\n\n| # | src | finding | resolved in |\n|---|---|---|---|\n| M1 | PC | `n=0 == uns[P_CT]` holds only for `use_logits=False`; default folds logits \u2014 scope the test | R2, \u00a76.1, Phase 5 |\n| M2 | PC | `groupby` unrepresentable in ct\u2011level table for non\u2011clone\u2011determined columns | \u00a76.1 |\n| M3 | PC/MS | `adata_or_jd` fast path incompatible with `n>0`/`groupby` \u2014 restrict to `n=0`,`groupby=None` | \u00a76.2 |\n| M4 | PC | `clone_size` reads retired `tcri_clone_key` \u2014 migrate reader with writer | \u00a74, Phase 1 |\n| M5 | PC | `build_archetypes` must keep returning `labels`; default\u2011K mismatch (4 vs 10) | \u00a7Phase 3, \u00a7Phase 8 |\n| M6 | MS | `clonotypic_entropy` `log2(C)` denominator group\u2011dependent, inflated by absent clones | \u00a76.2, Phase 6 |\n| M7 | MS | \"HDI\" is actually equal\u2011tailed \u2014 implement true HDI or rename | \u00a74 (`_stats`), Phase 1/6 |\n| M8 | MS/ML | Dual\u2011input + `n>0` ill\u2011defined \u2014 force `n=0` on bare jd | \u00a76.2 |\n| M9 | MS/ML | Temperature re\u2011centers draws / double\u2011tempering with guide_temperature \u2014 single knob, store raw mean | R6, \u00a75.2, \u00a76.1 |\n| M10 | ML | `splitby` (116 uses) dropped \u2014 retain a `splitby` axis on pl twins | \u00a7Phase 7 |\n| M11 | ML | Undispositioned public symbols (`probability_ternary`\u00d724, plural `*_entropies`, `phenotypic_entropy_delta`, `pl.flux` boxplot) | \u00a73 |\n| M12 | ML | No `seed`/`random_state` on engine/metrics \u2014 add and thread to torch | R7, \u00a76, Phase 5/6 |\n| M13 | ML | Normalized entropy/MI non\u2011comparable across groups (group\u2011specific denominators) | \u00a76.2 common\u2011denominator option |\n| M14 | ML | `@tl_result` h5ad\u2011hostile returns (MultiIndex, object arrays, `.attrs`) | \u00a76.2, \u00a710(9) |\n| M15 | ML | `n>0` reduction must keep draws coherent (same `p_ct` per `sample_id`) | \u00a76.1/6.2, Phase 5/6 |\n| M16 | GPU | Seed torch (not `np.random`) for reproducible draws | R7, \u00a77.5 |\n\n### A.3 Low\n\n| # | src | finding | resolved in |\n|---|---|---|---|\n| L1 | PC | `predict()` DataFrame indexing assumes order\u2011preserving loader | \u00a75.2, Phase 4 |\n| L2 | PC | `pl.mutual_information` `weighted=True` default removal changes displayed numbers | \u00a7Phase 7 changelog |\n| L3 | PC | `tcri_manager` stash is in `setup_anndata:697`, not `register_model` | \u00a75.1, Phase 4 |\n| L4 | PC/ML | moved `plot_phenotype_probabilities` reads `X_tcri_phenotypes`, not `X_tcri_probabilities` | \u00a73, Phase 10 |\n| L5 | PC | `get_p_ct` reads the global pyro param store \u2014 multi\u2011model clobber | \u00a75.3, \u00a711, Phase 8 |\n| L6 | MS | KL flux asymmetric/unbounded, base mismatch \u2014 bits/log2, offer JSD, l1 default | \u00a74 (`_distance`) |\n| L7 | MS | `n>0` intervals are partial posterior (p_ct only) \u2014 document | \u00a76.2, \u00a711 |\n| L8 | MS | `local_scale` uns fallback `1.0` corrupts draw variance \u2014 `to_anndata` always writes it; engine raises if missing at `n>0` | \u00a75.2, \u00a76.1, Phase 4/5 |\n| L9 | ML | `group_singletons` ordering vs `setup_anndata` unspecified | \u00a75.1, Phase 4 |\n| L10 | ML | `diag` model\u2011required vs adata\u2011only inconsistent; missing golden + engine==predict agreement tests | \u00a7Phase 8, \u00a79.1 |\n| L11 | PVM | current `softmax(log p_ct/T)` path is already the posterior mean at `T=1`; \"prior/non\u2011posterior\" label is a misnomer \u2014 collapse the axis to `use_logits` | R1/R2, \u00a76.1 |\n| L12 | GPU | module\u2011top `import umap` violates lazy\u2011import guardrail; cuML UMAP behind `_use_gpu` | \u00a77.3/7.4, Phase 5+ |", "audit": {"plan": {"findings": [{"area": "model/to_anndata \u2194 tl.joint_distribution engine (obsm['X_tcri_logits'])", "issue": "The engine's DEFAULT path is posterior=True, defined as softmax(logits + log p_ct) aggregated per clone (API \u00a75.1; math \u00a71). That reads obsm['X_tcri_logits']. But to_anndata is specified to write only 'metadata + categories, X_tcri latent, predict() probs + hard labels, p_ct (+ ct index maps). Nothing else' (API \u00a73). predict() returns probabilities, NOT logits (API \u00a78). So logits are never persisted and the default metric path has nothing to read.", "severity": "high", "fix": "Add a canonical _write_logits step to to_anndata (or fold it into _write_predictions) that writes obsm[K.X_LOGITS]='X_tcri_logits'; keep X_LOGITS in the 'canonical minimum'. Reconcile the '_keys lists X_LOGITS' vs 'to_anndata writes nothing else' contradiction. Currently register_model writes it explicitly (_preprocessing.py:226) and joint_distribution_posterior hard-requires it (_preprocessing.py:309-311, raises RuntimeError if missing).", "why": "joint_distribution_posterior reads adata.obsm['X_tcri_logits'][idx_cov] (_preprocessing.py:311) and raises if absent (310); the plan's posterior=True engine inherits this dependency but to_anndata's stated write-set excludes it.", "blocking": true}, {"area": "model/to_anndata + engine (gate_prob not persisted)", "issue": "The plan promises the posterior=True joint be 'identical to predict() \u2014 gate-aware' (math \u00a70.3, \u00a71). predict()/training combine logits and log-prior via a GATE when gate_prob is not None: local = g*logits + (1-g)*log_prior, else additive (_model.py:961-964, 609-612). But gate_prob is a model attribute; the engine takes only `adata`. gate_prob is not written to uns by register_model and is absent from _keys.py. From adata alone the engine cannot apply the gate, so it can only reproduce the additive rule \u2014 exactly the disagreement the plan claims to fix.", "severity": "high", "fix": "Persist gate_prob (scalar, or None) to uns via to_anndata and add K.GATE_PROB to _keys.py; the engine must branch on it to match predict(). Note default gate_prob=None (_model.py:722) makes models ungated by default, so the bug is latent until a user trains with a gate \u2014 but the plan's gate-aware parity guarantee is unachievable without persisting it. Persisting the already-gated per-cell logposterior only fixes n_samples=0; n_samples>0 still needs raw logits + gate_prob to recombine with fresh p_ct draws.", "blocking": true}, {"area": "model/setup_anndata ('registration-only, no obs mutation')", "issue": "Plan says setup_anndata becomes 'registration only \u2014 no obs mutation' (API \u00a73; Plan Phase 4). But setup_anndata writes adata.obs['indices']=range(n) (_model.py:678) and registers it as CategoricalObsField('indices','indices') (684). Training consumes it: idx_diag=batch['indices'] (603) and idx=batch['indices'] (641) drive per-cell ct lookup. Removing the obs write breaks the registered 'indices' field and training.", "severity": "high", "fix": "Keep writing adata.obs['indices'] in setup_anndata (it is registration glue, not analysis output), or derive the per-cell index without an obs column. Restate the invariant as 'no analysis/label obs mutation' rather than 'no obs mutation'.", "why": "batch['indices'] is read in training_step (_model.py:603) and validation_step (641); the field is registered at 684 and populated at 678.", "blocking": true}, {"area": "deletions ordering (utils.probabilities) vs import graph", "issue": "Plan Phase 2 deletes utils.probabilities as '0-caller dead code', scheduled BEFORE the plotting split (Phase 7). But _plotting.py imports it at module top: `from ..utils._utils import probabilities` (_plotting.py:18). Deleting the symbol without editing that import makes `import tcri.plotting` (hence `import tcri`) raise ImportError, violating the 'import tcri stays green every PR' invariant.", "severity": "high", "fix": "In the same PR that deletes probabilities, remove the _plotting.py:18 import. Generally: run an import-site grep before every 'safe deletion'; the inventory's '0-caller' label counts call-sites, not import-sites.", "why": "probabilities is unused functionally but is imported at plotting module load (_plotting.py:18); it also references never-written uns keys (probability_columns, joint_distribution) at _utils.py:290,293.", "blocking": true}, {"area": "engine/groupby \u2194 AnnData alignment guard", "issue": "joint_distribution_posterior hard-raises if per-cell uns arrays don't match n_obs, because uns['tcri_*_array_for_cells'] stay in full-cell space while obs/obsm are subset on a view/slice (_preprocessing.py:281-291). The natural groupby implementation slices adata per group \u2014 exactly what the current tcri_boxplot does: function(adata[adata.obs[groupby]==group]) (_plotting.py:616,637,640). Any engine-backed metric called on a sliced adata therefore hits the guard (ValueError) or silent misalignment.", "severity": "high", "fix": "Implement groupby by cell/clone RESTRICTION over the FULL adata (positional masks into the full-space uns arrays + clones=), never by passing a sliced AnnData to the engine \u2014 mirror how mi_compare restricts via clones= per patient (_metrics.py:125-131). _metric_boxplot (renamed tcri_boxplot) must be rewritten off the slice-and-call pattern.", "why": "The guard at _preprocessing.py:282 fires whenever len(ct_per_cell)!=n_obs; current tcri_boxplot slicing (616) would trip it for entropy metrics that route through the posterior engine.", "blocking": true}, {"area": "engine semantics / Phase-5 acceptance test", "issue": "Plan \u00a79/\u00a70.3 and Phase-5 test assert 'n_samples=0 == uns[K.P_CT] exactly' and 'closed-form posterior mean'. That equality holds ONLY for posterior=False (ct-level p_ct table). The DEFAULT metric path is posterior=True, which folds per-cell classifier logits (softmax(logits+log p_ct) aggregated per clone) and is therefore NOT equal to uns[K.P_CT]. The determinism claim (reproducible, no draw) is fine for both, but the equality/closed-form framing conflates the two joint definitions.", "severity": "medium", "fix": "Scope the 'n_samples=0 == uns[K.P_CT]' test to posterior=False. For posterior=True, test determinism (bit-identical repeat calls) and equivalence-to-predict()-aggregation instead. Clarify in \u00a70.3/\u00a79 that the closed-form-mean identity is the posterior=False object, not the default posterior=True joint.", "why": "posterior=True combines obsm logits with log p_ct (plan \u00a71); only posterior=False returns the tempered p_ct rows (current joint_distribution, _preprocessing.py:397,415).", "blocking": false}, {"area": "engine/groupby \u2194 model covariate conditioning", "issue": "p_ct is indexed by ct=(clonotype, covariate) only (ct_to_c, ct_to_cov; _model.py:754-767). The engine's `covariate=` selects ct rows via ct_to_cov. A new `groupby` by an arbitrary obs column that is NOT functionally determined by clonotype-at-fixed-covariate (e.g. a tissue that cross-cuts a clone at one timepoint) is unrepresentable in the ct-level (posterior=False) table \u2014 there is no cell handle there to partition by. It is only well-defined on the cell-level posterior=True path (aggregate that group's cells' logits).", "severity": "medium", "fix": "Document/enforce that groupby requires the cell-informed (posterior=True) path, or that groupby columns must be clone-constant (like patient when clone ids are patient-unique, matching mi_compare's clones=patient_clones). Reject or warn on posterior=False + non-clone-determined groupby.", "why": "ct index carries only clone and covariate (_model.py:760-767); grouping by another cell attribute has no representation in uns['tcri_p_ct'] rows.", "blocking": false}, {"area": "metrics adata_or_jd fast path feasibility", "issue": "The precomputed-joint fast path (adata_or_jd) carries no p_ct, local_scale, logits, or cells. So it is incompatible with n_samples>0 (cannot draw more Dirichlet samples) and with groupby (cannot re-partition cells). The plan advertises a uniform signature (n_samples, groupby, clones) on the fast path without restricting these combinations.", "severity": "medium", "fix": "Specify that a precomputed-jd input is valid only for n_samples=0 and groupby=None (and that clones= just re-filters rows); raise a clear error otherwise. For n_samples>0 fast path, require the passed jd to already contain the sample axis and treat n_samples as validation, not resampling. Note phenotypic_flux correctly takes adata-only (it needs two joints), so it is consistent.", "why": "joint_distribution draws use uns['tcri_local_scale'] and p_ct (_preprocessing.py:445-448); a bare DataFrame lacks them.", "blocking": false}, {"area": "pp/_keys migration (clone_size depends on retired key)", "issue": "clone_size (kept in pp) reads adata.uns['tcri_clone_key'] (_preprocessing.py:549). The single-scheme migration retires tcri_clone_key/tcri_phenotype_key in favor of uns['tcri_metadata'] (API \u00a72 note). If the key is removed before clone_size is migrated, clone_size raises KeyError. Same key is also read by top_clone_umap (_plotting.py:550, moving to examples) and by deleted clonality/clone_fraction (_metrics.py:605-627).", "severity": "medium", "fix": "In the Phase-1 _keys adoption, migrate clone_size to read K.CLONE_COL from uns['tcri_metadata']['clone_col'] in the SAME change that stops writing tcri_clone_key (register_clonotype_key, _preprocessing.py:80). Ensure register_clonotype_key/register_phenotype_key retirement and every reader flip together.", "why": "tcri_clone_key is written only by register_clonotype_key (_preprocessing.py:80, called at register_model:258) and read at _preprocessing.py:549; retiring the writer orphans the reader.", "blocking": false}, {"area": "model/_training.build_archetypes return shape \u2194 diag.archetypes", "issue": "Current build_archetypes returns a tuple (centers, labels) (_model.py:45-51) and TCRIModel stores both (self.centers, self.labels; _model.py:751). plot_archetypes orders the heatmap by self.labels (argsort(self.labels), _model.py:1018). The plan's _training.build_archetypes is specified to return np.ndarray (centroids only) (API \u00a73 _training). Dropping labels removes the cluster ordering diag.archetypes needs.", "severity": "medium", "fix": "Have build_archetypes return both centers and labels (or persist labels on the model / in uns) so diag.archetypes can reproduce the cluster-ordered heatmap. Update the API \u00a73 signature accordingly. Also note the default K mismatch: build_archetypes defaults K=4 (_model.py:45) while the plan/model use K=10.", "why": "plot_archetypes(_model.py:1017-1019) consumes self.labels from build_archetypes; a centroids-only return breaks the relocated diag.archetypes.", "blocking": false}, {"area": "model/predict() DataFrame indexing", "issue": "predict() (renamed get_cell_phenotype_probs) is specced to return a DataFrame indexed by cell barcode (API \u00a78). The implementation walks the loader with sequential current_idx bookkeeping and maps cells to ct via ct_array in loader order (_model.py:943-969), assuming the data loader yields cells in obs/ct_array order. If the inference loader ever reorders/shuffles, both the ct assignment and the obs_names index would be mislabeled.", "severity": "low", "fix": "Either assert the inference DataLoader is order-preserving (shuffle=False, sequential sampler) and index by adata.obs_names, or carry an explicit per-batch index (the registered 'indices' field) and reindex, so ct lookup and barcode labels can't drift.", "why": "ct assignment uses current_idx slicing into ct_array (_model.py:951) which is only correct if the loader preserves setup order.", "blocking": false}, {"area": "metrics/plotting weighted-removal semantics", "issue": "pl.mutual_information currently defaults weighted=True (_plotting.py:1087). The plan removes `weighted` entirely and enshrines a uniform-clonotype prior (math \u00a70.6). This silently changes what existing MI plots display (cell-weighted \u2192 per-clonotype), not just the API.", "severity": "low", "fix": "Call out the uniform-clonotype convention as a behavior change in the changelog and each metric docstring; verify no notebook relies on weighted=True MI. (tl.mutual_information never accepted weighted, so only the pl default changes user-visible numbers.)", "why": "_plotting.py:1087 signature sets weighted=True by default; the plan's convention makes every clone count equally regardless of cell count.", "blocking": false}, {"area": "tcri_manager stash location (mislocated in plan)", "issue": "Plan attributes the AnnDataManager uns-stash to register_model/to_anndata and says removing it there deletes the write_adata_safely hack. The stash actually happens in setup_anndata: adata.uns['tcri_manager']=adata_manager (_model.py:697). register_model does NOT stash it. If only to_anndata is changed and _model.py:697 is left, tcri_manager is still written and write_adata_safely (_utils.py:546-561) is still required.", "severity": "low", "fix": "Remove the stash at setup_anndata (_model.py:697) as part of making setup_anndata registration-only; then write_adata_safely/_pop_nonserializables can be dropped. Confirm nothing reads uns['tcri_manager'] (grep shows only the save-time pop consumes it).", "why": "The only writer is _model.py:697; the only reader is the save-time pop in _utils.py:548-550.", "blocking": false}, {"area": "examples/keys (moved plot_phenotype_probabilities)", "issue": "plot_phenotype_probabilities reads obsm['X_tcri_phenotypes'] (_plotting.py:813-816), the slot written by the legacy classify_phenotypes (_preprocessing.py:95,129), not the register_model/to_anndata slot 'X_tcri_probabilities' (_preprocessing.py:237). classify_phenotypes is being dropped and to_anndata standardizes on X_tcri_probabilities, so the moved example reads a key nothing writes.", "severity": "low", "fix": "Update the moved example to read K.X_PROBABILITIES ('X_tcri_probabilities'); drop the X_tcri_phenotypes slot entirely.", "why": "Two divergent probability slots exist today (X_tcri_phenotypes vs X_tcri_probabilities); the retained writer is the latter.", "blocking": false}, {"area": "model/serialization global pyro param store", "issue": "get_p_ct() reads the GLOBAL pyro param store: pyro.get_param_store()['q_p_ct_raw'] (_model.py:512-515), not module-local state. to_anndata calls model.module.get_p_ct() and diag PPCs run on the finalized model; both depend on the global singleton. Loading two sessions/models in one process clobbers q_p_ct_raw, and load restores it into the same global (_utils.py:214-218, _ensure_pyro_posterior_params:87-135).", "severity": "low", "fix": "Not introduced by the refactor, but the plan's multi-model diag/round-trip workflows should scope the param store per model (namespaced params or restore-before-use), or document single-model-per-process. At minimum, session load must set the store immediately before any get_p_ct/to_anndata call.", "why": "get_p_ct pulls q_p_ct_raw from the process-global store (_model.py:514-515); there is no per-instance isolation.", "blocking": false}], "summary": "Audited the drafted refactor plan/API against the actual tcri source. Five blocking issues: (1) to_anndata's 'canonical minimum, nothing else' drops obsm['X_tcri_logits'], which the engine's DEFAULT posterior=True path hard-requires (_preprocessing.py:309-311); (2) gate_prob is a model attribute never persisted to adata, so the promised gate-aware parity between the joint and predict() is infeasible from adata alone; (3) setup_anndata cannot be 'registration-only, no obs mutation' because it must write the registered obs['indices'] column that training consumes (batch['indices'], _model.py:603/641/678/684); (4) deleting utils.probabilities in Phase 2 breaks `import tcri` since _plotting.py:18 imports it at module load (it is not truly 0-caller); (5) any groupby implemented by slicing adata trips joint_distribution_posterior's hard full-space-vs-subset alignment guard (_preprocessing.py:281-291), which is exactly what today's tcri_boxplot slicing does. Non-blocking but important: the 'n_samples=0 == uns[tcri_p_ct]' closed-form claim holds only for posterior=False (the default posterior=True folds logits); groupby is unrepresentable in the ct-level table for non-clone-determined columns; the adata_or_jd fast path is incompatible with n_samples>0/groupby; clone_size still reads the retired tcri_clone_key; build_archetypes must keep returning labels for diag.archetypes; the tcri_manager stash is in setup_anndata (line 697), not register_model as the plan states. The plan correctly identifies the genuine current bugs (compare_joint_distribution NameError on undefined `model` at _plotting.py:209, MI/flux posterior=False dead branches, flux torch-vs-numpy seeding, KL base mismatch)."}, "math": {"findings": [{"area": "n_samples=0 point estimate for entropy/MI/flux (Jensen gap)", "issue": "The resolution correctly proves E_q[p_ct] = m = uns['tcri_p_ct'] is the closed-form posterior MEAN of the DISTRIBUTION p_ct, and defines n_samples=0 as 'the deterministic point estimate.' It then computes entropy/MI/flux ON that mean distribution and presents the mean of the n_samples>0 draws as the same quantity's posterior mean. But entropy, MI, and KL/l1 flux are NONLINEAR functionals of p, so metric(E_q[p]) != E_q[metric(p)] by a Jensen gap that is NOT Monte-Carlo noise. Shannon entropy is concave, so the plug-in H(E[p]) OVER-estimates E[H(p)]; l1/KL flux is convex, so plug-in flux(mean,mean) UNDER-estimates E[flux] (a clone with no real change gives flux 0 at n_samples=0 but strictly positive mean flux at n_samples>0); for MI = H(phi) - H(phi|c) the two gaps partly cancel so the sign is indeterminate but the magnitudes differ. So n_samples=0 and the n_samples>0 'mean' column systematically disagree, and a conformance test asserting they match (or users expecting consistency) will be wrong.", "severity": "high", "why": "clonotypic_entropy/phenotypic_entropy (_metrics.py:315,569) apply scipy entropy() to a single distribution; mutual_information (_metrics.py:732) applies _mi_from_joint to one table; the settled design in the spec (sections 0.4/0.5, delta table rows 1-2) equates the closed-form mean point estimate with the posterior-mean metric. Verified E[Dirichlet(local_scale*m)] = m (local_scale cancels) but E[H(Dirichlet)] != H(m).", "fix": "Decide and document explicitly that n_samples=0 is the PLUG-IN estimator metric(E_q[p_ct]) = metric-at-posterior-mean, which is NOT an estimator of E_q[metric]; do not advertise the n_samples>0 'mean' as the same number. If E_q[metric] is wanted as the point estimate, it requires draws (or a bias correction) even for the 'deterministic' path. Keep the two clearly labeled as different estimators in docstrings and drop any test asserting n_samples=0 == mean(n_samples>0).", "blocking": true}, {"area": "P(>0) / p_gt summary column on single-metric draws", "issue": "Section 0.5's uniform return-shape rule adds p_gt (prob_greater_than_zero) as a summary column for the n_samples>0 output of EVERY metric. But clonotypic entropy, phenotypic entropy, mutual information, and phenotypic flux are all mathematically NONNEGATIVE (entropy >=0, MI >=0, l1 >=0, KL >=0). Every posterior draw is >= 0, so P(draw > 0) is approximately 1 for all of them and carries zero information. P(>0) is a directional/sign statistic that is only meaningful for a SIGNED quantity such as a between-group difference Delta = metric_B - metric_A.", "severity": "high", "why": "_mi_from_joint returns MI>=0 (_metrics.py:82-89); entropy() outputs >=0; flux l1/dkl are >=0 (_metrics.py:960-978). The spec's _stats.prob_greater_than_zero (section 7) is correctly scoped to a delta in prob_direction, but section 0.5 and the summarize() reducer attach p_gt to the raw single-metric draw vector, where it is vacuous.", "fix": "Remove p_gt from the single-metric n_samples>0 summary (summarize() should emit mean, sd/HDI only for a raw metric). Compute P(>0) ONLY on a paired between-group difference vector produced by groupby + the _stats comparison helper. Document that p_gt requires a signed contrast, not a nonnegative metric.", "blocking": true}, {"area": "Posterior draws must reproduce the guide's clamped Dirichlet concentration", "issue": "The HDI/mean/sd story for n_samples>0 is sold as summarizing the learned variational posterior q(p_ct). But the guide's actual posterior is Dirichlet(clamp(local_scale * m, min=1e-3)) with a hard 1e-3 concentration floor. There are three DIFFERENT concentration formulas in play: guide = clamp(local_scale*m, 1e-3); joint_distribution draws = local_scale*p_ct_for_cov (NO clamp, NO eps); joint_distribution_posterior = local_scale*p_ct_mean + 1e-8. For committed clones (m_k < 1e-3/local_scale, i.e. exactly the low-entropy/high-MI clones that dominate signal) the clamp raises tiny alpha_k, so E and variance of the clamped Dirichlet differ from the unclamped one, AND the closed-form mean m no longer equals the draw mean. Draws from the unclamped/eps-shifted forms therefore come from a distribution the model never learned, so the reported HDIs are miscalibrated and, worse, the n_samples=0 mean (= m) is inconsistent with the mean of these draws for committed clones.", "severity": "high", "why": "Guide: _model.py:490 conc_ct_guide = torch.clamp(self.local_scale * q_p_ct_sharp, min=1e-3). joint_distribution: _preprocessing.py:446 conc = local_scale * p_ct_for_cov (no clamp). joint_distribution_posterior: _preprocessing.py:305 Dirichlet(local_scale * p_ct_mean + 1e-8). get_p_ct returns unclamped normalized m (_model.py:520-525), so it equals the guide mean only when the 1e-3 clamp is inactive.", "fix": "Make the n_samples>0 engine draw from EXACTLY the guide's distribution: Dirichlet(clamp(local_scale * m, min=1e-3)), reusing the same floor as _model.py:490. Either accept that the point estimate is then the clamped-Dirichlet mean (recompute m as alpha/sum(alpha) after clamping) or document the clamp-induced discrepancy. Do not use the bare local_scale*m or +1e-8 variants for 'posterior' draws.", "blocking": true}, {"area": "mutual_information normalise_mode='average' default under the uniform-clone prior", "issue": "Weighting is removed, enshrining P(c)=1/C uniform over clones. Because each joint row is normalized to sum 1 and the whole table is then divided by its sum, px = P(c) = 1/C exactly, so the clone marginal entropy H_c is PINNED to log2(C) (maximal, structural, uninformative). The default normalise_mode='average' uses denom = 0.5(H_c + H_p) = 0.5(log2(C) + H_p). For realistic C (hundreds-thousands of clones) log2(C) dominates, so normalized MI is throttled by ~1/log2(C): it can never approach 1 and its ceiling shrinks as clone count grows, making it non-comparable across covariates/groups/patients with different C -- which breaks the entire groupby-comparison workflow the refactor is built around.", "severity": "high", "why": "_mi_from_joint (_metrics.py:80,86-88): px = pxy.sum(1); h_c = -sum(px*log2 px); with pxy=J/sum(J) and each J row summing to 1, px==1/C so h_c==log2(C). denom = 0.5*(h_c+h_p) if mode=='average'. mutual_information passes normalise_mode default 'average' (_metrics.py:641).", "fix": "Given the settled uniform-clone construction makes H_c a constant log2(C), change the default to 'min' (which yields I/H_p = coefficient of constraint, in [0,1], reaching 1 when clone determines phenotype and independent of C), or drop 'average' entirely. Document that H_c is structurally log2(C) and is not a meaningful normalizer here.", "blocking": true}, {"area": "clonotypic_entropy log2(n_clones) normalizer comparability", "issue": "clonotypic_entropy normalizes by log2(C) where C = number of clone rows in the table. Unlike phenotypic_entropy's log2(P) (P is a fixed constant across all clones/groups), C varies with the covariate, the clones= filter, and how many requested-but-absent clones get reindexed in as all-zero rows. Absent clones inflate C (hence log2(C)) while adding only negligible eps mass to the numerator, so normalized clonotypic entropy is systematically DEFLATED and is not comparable across groups with different clone counts -- again breaking the groupby comparison use case.", "severity": "medium", "why": "_metrics.py:316-317 H /= np.log2(len(vec)) with len(vec) = rows in the column; joint_distribution_posterior reindex(clones).fillna(0.0) (_preprocessing.py:326) adds zero rows for absent clones that are then eps-clipped and counted in len(vec).", "fix": "Define and document the normalizer's support precisely: divide by log2 of the number of clones with genuine support (or a fixed reference count) rather than the raw reindexed row count; drop requested-but-absent (all-zero) clones from C before normalizing. Warn that clonotypic_entropy across groups uses group-specific denominators and provide a common-denominator option for comparisons.", "blocking": false}, {"area": "HDI vs equal-tailed credible interval on skewed bounded posteriors", "issue": "Section 7 (and the current code) computes np.percentile(x,[2.5,97.5]) but labels it HDI. For committed clones the Dirichlet-induced entropy/MI/flux posteriors pile up against the boundary (entropy near 0, flux near 0), i.e. they are strongly skewed, so the equal-tailed interval is NOT the highest-density interval: it excludes high-density mass near the boundary and includes low-density mass in the tail. Reporting it as HDI is statistically incorrect, and near the [0,1]/[0,inf) boundary even a true HDI from a few hundred draws is unstable.", "severity": "medium", "why": "mutual_information reports np.percentile(mi_samples,[2.5,97.5]) as '95% CI' (_metrics.py:752-753); the spec's hdi()/summarize() (section 7) promise an HDI but the referenced implementation is equal-tailed. Dirichlet(local_scale*m) with committed m is right-skewed in derived entropy.", "fix": "Either implement a true highest-density interval (narrowest interval containing the mass) or rename the output ci_low/ci_high and stop calling it HDI. Note boundary instability and prefer HDI for the bounded, skewed entropy/flux posteriors.", "blocking": false}, {"area": "Dual-input (precomputed joint) combined with n_samples>0", "issue": "Every metric is specified to accept EITHER (adata + covariate) OR a precomputed joint DataFrame as a fast path, and the test plan asserts adata-path == precomputed-joint-path equivalence. But a precomputed joint is a single realized table; there is nothing to draw from, so n_samples>0 is ill-defined on the precomputed-joint input. The equivalence only holds at n_samples=0.", "severity": "medium", "why": "Spec sections 2-5 give each metric both an adata path (calls joint_distribution internally, which samples) and an 'adata_or_jd' precomputed path; n_samples>0 sampling lives in joint_distribution (_preprocessing.py:443-490) and cannot be reconstructed from one collapsed joint DataFrame.", "fix": "Specify that a precomputed-joint input forces n_samples=0 (or must itself carry a sample axis / MultiIndex with sample_id); raise a clear error if a bare joint is passed with n_samples>0. Scope the dual-input equivalence test to n_samples=0 only.", "blocking": false}, {"area": "Temperature placement and its effect on what distribution is sampled", "issue": "At n_samples>0 the engine draws Dirichlet(local_scale * tempered_mean) where tempered_mean = softmax(log m / T). For T != 1 the draws are centered on a re-tempered vector, NOT the learned posterior q(p_ct); the analysis-time temperature silently redefines the sampled distribution. Separately, for posterior=True the goal is to match predict(), but predict() already bakes in classifier_temperature and applies no extra analysis temperature, whereas the metric divides (logits + log prior) by T -- so a non-default T double-scales relative to the model's own cell probabilities.", "severity": "medium", "why": "joint_distribution: p_ct_for_cov = softmax(log(p_ct_for_cov+eps)/T) then conc = local_scale*p_ct_for_cov (_preprocessing.py:397,446). predict/get_cell_phenotype_probs combine logits+log prior with only classifier_temperature, no analysis T (_model.py:907-972). joint_distribution_posterior applies T in softmax((logits+log_prior)/T) (_preprocessing.py:314), a different stage than the prior path.", "fix": "Document that temperature != 1 makes n_samples>0 draws a re-tempered distribution, not the model posterior. Pick one consistent placement for T across both branches, and clarify the relationship between analysis temperature and the model's classifier_temperature so posterior=True at T=1 exactly reproduces predict().", "blocking": false}, {"area": "KL (dkl) flux: asymmetry, unboundedness, base", "issue": "phenotypic_flux offers distance_metric='dkl' = sum p*log(p/q), which is asymmetric (flux(from->to) != flux(to->from)) and unbounded in [0,inf); when a phenotype present at 'from' is near-absent at 'to' the term log(p/eps) explodes and dominates. Advertising it alongside l1 as a 'distance' invites misuse for a symmetric shift measure, and the current implementations disagree on log base (natural log in flux.dkl_func vs log2 elsewhere).", "severity": "low", "why": "flux dkl_func uses natural log with eps=1e-15 (_metrics.py:962-964); module-level dkl uses scipy entropy natural log eps=1e-10 (_metrics.py:159-163); entropy/MI use log2. KL is asymmetric and unbounded by definition.", "fix": "Standardize KL to bits (log2) with one eps in _distance.kl_divergence. Document dkl flux as directional and unbounded. Offer symmetric Jensen-Shannon divergence (bounded by 1 bit) as the recommended symmetric alternative; keep l1 as the safe bounded default.", "blocking": false}, {"area": "Partial posterior: only p_ct uncertainty, classifier logits treated as fixed", "issue": "For posterior=True with n_samples>0 the only randomness is the Dirichlet draw of the ct-level prior p_ct; the per-cell classifier logits are held fixed. The resulting HDI therefore reflects ONLY p_ct uncertainty, not encoder/classifier uncertainty, and understates total model uncertainty. This is defensible (it is the only cheaply available uncertainty) but must be stated so users do not read the intervals as full posterior predictive uncertainty.", "severity": "low", "why": "joint_distribution_posterior draws p_ct once and reuses fixed obsm['X_tcri_logits'] across the softmax (_preprocessing.py:305,311-314); logits carry no draw axis.", "fix": "Document that n_samples>0 intervals capture prior/p_ct uncertainty only, with classifier logits fixed at their posterior-mean encoding.", "blocking": false}, {"area": "local_scale default fallback 1.0 vs trained value", "issue": "The draw variance is set by local_scale, read as adata.uns.get('tcri_local_scale', 1.0). The model is constructed with local_scale=3.0 (TCRIModel) / 5.0 (TCRIModule). If to_anndata fails to write tcri_local_scale, all posterior draws silently use scale 1.0, i.e. much higher variance than the trained posterior, corrupting every HDI/sd without error.", "severity": "low", "why": "_preprocessing.py:300 and 445 both use .get('tcri_local_scale', 1.0); model defaults differ (TCRIModel/TCRIModule signatures). The mean path is unaffected (local_scale cancels) but draws are not.", "fix": "Make to_anndata always write K.LOCAL_SCALE from the trained module, and have the engine raise (not default to 1.0) if the key is missing when n_samples>0. Assert the written value equals the module's local_scale.", "blocking": false}], "summary": "The refactor's core resolution -- that n_samples=0 should return the closed-form variational posterior mean E_q[p_ct] = m = uns['tcri_p_ct'], with local_scale canceling in the mean, and that the prior/MAP paths are wrong -- is mathematically correct and well-grounded (verified against the guide at _model.py:487-491 and get_p_ct at :510-525). However, several proposed statistical conventions are wrong or will mislead before implementation. THREE are blocking. (1) Jensen gap: computing entropy/MI/flux on the mean distribution (metric(E[p])) is a PLUG-IN estimator that does NOT equal the posterior mean E[metric(p)] for these nonlinear functionals; entropy plug-in over-estimates, flux plug-in under-estimates, so n_samples=0 and the mean of n_samples>0 draws systematically disagree and must be documented as different estimators. (2) P(>0)/p_gt is vacuous as a per-metric summary because entropy, MI, and flux are all >= 0 (P(>0)~1 always); it is only meaningful on signed between-group differences. (3) The n_samples>0 draws must reproduce the guide's clamped concentration Dirichlet(clamp(local_scale*m, 1e-3)); the current engines use three inconsistent unclamped/eps-shifted variants, so the advertised HDIs summarize a distribution the model never learned and disagree with the mean for committed clones. Additionally, the uniform-clone prior pins the clone marginal entropy H_c to log2(C), which breaks normalise_mode='average' (default) for MI -- it becomes clone-count-dependent and non-comparable across groups; 'min' (I/H_p) is the correct choice. Also flagged: clonotypic entropy's log2(C) denominator is group-dependent and inflated by reindexed absent clones (hurts groupby comparisons); the 'HDI' is actually an equal-tailed interval and is mislabeled for skewed bounded posteriors; dual-input + n_samples>0 is ill-defined on a single precomputed joint; temperature != 1 silently re-centers the sampled distribution; KL flux is asymmetric/unbounded (prefer JSD); the posterior is partial (p_ct only); and the local_scale=1.0 uns fallback can corrupt draw variance if the key is missing."}, "priorMean": {"options_analyzed": [{"name": "(a) Closed-form posterior MEAN E_q[p_ct] under the learned Dirichlet q(p_ct)", "math": "The guide learns q(p_ct)=Dirichlet(alpha) with alpha = local_scale * m, where m = normalize(q_p_ct_raw) is on the simplex (guide L487-491). Because alpha is proportional to a simplex vector, Sum(alpha)=local_scale and the mean is closed-form: E_q[p_ct] = alpha/Sum(alpha) = m. local_scale cancels exactly. Numerically m == get_p_ct() (L510-525) == adata.uns['tcri_p_ct'].", "pros": "Exact, closed-form, deterministic, reproducible. Already computed and stored (uns['tcri_p_ct']) so n_samples=0 costs zero sampling. It is the Bayes point estimator under squared-error loss. Lies on the interior of the simplex (a valid probability vector). Invariant to local_scale (which only controls posterior spread, not location).", "cons": "Essentially none. Minor: get_p_ct does not reproduce the guide's clamp(min=1e-3) on the concentration, so for a handful of near-degenerate rows m differs from the guide's true Dirichlet mean at the 1e-3 scale (negligible)."}, {"name": "(b) The PRIOR clone_phen_prior (or the archetype mixture_concentration)", "math": "clone_phen_prior is the buffer prepare_two_level_params builds from c2p_mat = row-normalized empirical hard-label counts (+eps, optional prior_temperature). It enters training ONLY as guide init (init_mat = clone_phen_prior*guide_init_scale+1e-3, L444/L471) and as the generative anchor via the archetype MixtureDirichlet. It is not a function of the trained variational posterior at all.", "pros": "None as a point estimate of what the model learned.", "cons": "Wrong object: it is the initialization / generative anchor, discarding everything training learned; metrics built on it would be independent of the fit. It is derived from argmax hard labels (label leakage / circularity). There are also TWO distinct prior objects (clonotype-level clone_phen_prior vs K-archetype centroids), neither indexed at the ct level that metrics need."}, {"name": "(c) MAP / posterior mode", "math": "Dirichlet(alpha) mode = (alpha-1)/(Sum(alpha)-P), defined only when every alpha_k > 1. Here alpha_k = local_scale * m_k with local_scale ~ 3-5 and m_k < 1, so alpha_k is routinely < 1.", "pros": "None here.", "cons": "When any alpha_k < 1 the density is unbounded at the simplex boundary: the mode is ill-defined / sits on a face (zeros), is non-unique, and is non-representative. Unstable and discontinuous in local_scale. Not a usable point estimate for this posterior."}, {"name": "(d) Mean of many Dirichlet draws", "math": "(1/S) * Sum_s p_ct^(s), with p_ct^(s) ~ Dirichlet(local_scale * m). By LLN this converges to E_q[p_ct] = m, but with Monte-Carlo error O(1/sqrt(S)).", "pros": "Converges to the correct answer (a).", "cons": "A strictly dominated estimator of a quantity available in closed form (Rao-Blackwell): it only adds MC noise and cost. Non-deterministic / non-reproducible unless a generator is seeded. Never justified when the closed form exists."}, {"name": "(e) Current non-posterior softmax(log p_ct / T) in joint_distribution", "math": "joint_distribution n_samples=0 computes p_tilde = softmax(log(m + 1e-8)/T) = m^(1/T)/Sum(m^(1/T)), reading m = uns['tcri_p_ct']. At T=1 this is the identity up to the eps renormalization, so p_tilde == m == option (a). At T!=1 it is a temperature-tempered mean.", "pros": "At T=1 it already equals the correct estimator (a) and is already deterministic. Cheap.", "cons": "The 'non-posterior' / 'prior' label is a MISNOMER: it already uses the posterior mean m (never the generative prior). Its real deficiency vs the posterior=True path is that it does NOT fold in the per-cell classifier logits, so it disagrees with model.predict. The current 'prior vs posterior' framing conflates that classifier-mixing axis with a (nonexistent) prior/posterior axis."}], "recommendation": "Adopt option (a): n_samples=0 is the closed-form variational posterior MEAN E_q[p_ct] = normalize(q_p_ct_raw) = TCRIModel.get_p_ct() = adata.uns['tcri_p_ct'] (temperature-adjusted when temperature != 1). It is read directly and NEVER sampled. Reject the prior (b) \u2014 it is the guide-init/generative anchor, not what training learned, and is built from leaked hard labels; reject the MAP/mode (c) \u2014 alpha_k = local_scale*m_k is routinely < 1, so the mode is on the simplex boundary / undefined; reject mean-of-draws (d) \u2014 a strictly worse, noisier estimator of a quantity available in closed form; and recognize that the current softmax(log p_ct/T) path (e) already IS this posterior mean at T=1 (its 'prior/non-posterior' label is wrong). The prior path and the mutual_information posterior=False NotImplementedError branch are deleted, not implemented.", "arg_semantics": "n_samples is the single point-vs-draws knob (delete point_estimate=). n_samples=0 -> deterministic closed-form mean m; n_samples=N>0 -> N i.i.d. draws p_ct^(s) ~ Dirichlet(local_scale * m_tilde) via a seeded torch.Generator. temperature power-tempers the mean identically in both the mean and draw paths: m_tilde = softmax(log(m+1e-8)/T) (identity at T=1) \u2014 this must be applied one consistent way, unlike today where the two engines temper at different stages. uns['tcri_local_scale'] is consulted ONLY for n_samples>0 draws; it cancels in the mean and is irrelevant to n_samples=0.", "keep_posterior_arg": "Drop posterior= from ALL FOUR metrics and from phenotypic_flux: they always use the learned posterior (mean at n_samples=0, draws at n_samples>0) and, given an adata, compute the joint via joint_distribution with classifier mixing on. Keep exactly ONE flag on the engine joint_distribution, but REDEFINE and rename it to end the misnomer: it is a classifier-mixing switch, not a prior/posterior switch. Recommended name use_logits (alias cell_informed), replacing both posterior= and combine_with_logits=. use_logits=True (default) folds per-cell classifier logits into log(base) exactly like model.predict; use_logits=False returns the ct-level table directly. Both branches use the posterior mean/draws of p_ct; NEITHER ever touches the generative prior. So the 'prior vs posterior' axis is collapsed entirely (there was never a live prior branch), and what survives is only 'with logits vs without logits'.", "math_justification": "See math_justification field above (posterior-mean = Bayes estimator under squared loss; alpha proportional to simplex vector m gives closed-form mean m independent of local_scale; mode degenerate because alpha_k<1; mean-of-draws Rao-Blackwell dominated; prior is not a function of the fit). combine_with_logits is orthogonal to the point/draw choice and is applied to whichever base (mean m at n_samples=0, or a draw at n_samples>0) is used.", "implementation": "TENSORS/KEYS.\nDeterministic base (n_samples=0):\n m = adata.uns['tcri_p_ct'] # shape (ct_count, P); == TCRIModel.get_p_ct() == normalize(pyro param 'q_p_ct_raw')\nCovariate selection:\n cov_idx = list(uns['tcri_covariate_categories']).index(covariate)\n ct_mask = np.asarray(uns['tcri_ct_to_cov']) == cov_idx\n m_cov = m[ct_mask] # (n_ct_in_cov, P)\n clone_ids= np.asarray(uns['tcri_ct_to_c'])[ct_mask] # clone index per selected ct row\nTemperature (one consistent form, both branches):\n m_tilde = softmax(log(m_cov + 1e-8) / temperature, axis=-1) # == m_cov at T=1\nBranch use_logits=False (ct-level table, no classifier):\n return DataFrame(m_tilde, index=clone_ids, columns=phenotype_categories) # each row a simplex; deterministic\nBranch use_logits=True (default; fold classifier, match model.predict):\n per cell i in the covariate (selected via uns['tcri_cov_array_for_cells']==cov_idx),\n ct(i)=uns['tcri_ct_array_for_cells'][i], l_i = obsm['X_tcri_logits'][i]:\n logb = log(m_tilde[row_of ct(i)] + 1e-8)\n combine= gate_prob*l_i + (1-gate_prob)*logb if model.use_gate else l_i + logb\n P_phi_i= softmax(combine / temperature)\n J[clone,:] += P_phi_i ; then row-normalize J. # gate-aware, mean-prior -> identical to TCRIModel.predict except base is the mean m, not a draw. Fixes the standing bug where the metric used ONE Dirichlet draw and NEVER applied the gate.\nDraws (n_samples=N>0), SAME code path with base = draw:\n s = uns['tcri_local_scale']; g = torch.Generator().manual_seed(seed)\n p_ct_s ~ Dirichlet(s * m_tilde).sample((N,), generator=g) # torch, seeded (numpy seeding alone was a no-op bug)\n feed each p_ct_s as 'base' through the same temperature + combine + aggregate; stack a sample axis.\nNOTES: local_scale (uns['tcri_local_scale']) is used ONLY in the draw path; it cancels in the mean. Delete point_estimate= and the metric-level posterior=/combine_with_logits args; keep a single engine-level use_logits (replacing posterior=). Repeated n_samples=0 calls are bit-identical by construction."}, "gpu": {"findings": [{"area": "Engine core: write the joint as a single batched, device-routable function (not per-draw pandas)", "issue": "Today the substrate is per-draw pandas: joint_distribution_posterior returns ONE DataFrame per call (Dirichlet.sample() at _preprocessing.py:305 + groupby at :320-322), and every metric loops it (_metrics.py:296, :741). If the refactored tools/_joint.py preserves that per-draw/DataFrame shape, every GPU/vectorization win below is impossible to retrofit without a second engine rewrite.", "severity": "high", "fix": "Refactor the joint's numeric core into a device-routed `_joint_draws(adata, covariate, n_samples, device) -> array[n_samples, n_clones, P]` (precompute clone codes once; batch Dirichlet + softmax; scatter-sum by clone). Keep the public tl.joint_distribution DataFrame as a thin single-draw/summary wrapper over the stack. Grafiti precedent: _compute/joint.py loads q2d as [N,S,C] and reduces it as one array core, DataFrame only at the boundary.", "why": "The math spec already mandates a [n_samples, n_clones, P] stack for all metrics; building the engine around it from the start is the vehicle for P0-P2 and avoids baking pandas per-draw cost into the new layout.", "blocking": true}, {"area": "Reduction: replace pandas groupby-by-clone with precomputed-code scatter-add (P0, biggest single win)", "issue": "_preprocessing.py:320-322 does `pd.DataFrame(probs_cell, index=clone_labels).groupby(level=0).sum()` over the covariate's cells (1e5-1e6 rows) and rebuilds it from scratch on every one of the 50-1000 draws \u2014 pandas re-hashes and re-groups the SAME cell->clone key set each iteration even though that grouping is constant.", "severity": "high", "fix": "Compute clone integer codes ONCE outside the loop, then reduce with a constant-index scatter-add: np.add.at / np.bincount per phenotype column on CPU, torch.zeros(n_clones,P).index_add_(0, clone_codes, probs) or cupy.bincount on GPU, batched across all n_samples in one kernel. This is grafiti _compute/contingency.py::contingency and graph.py::edge_tensor verbatim (one O(N) bincount over cells, rest is small [n_clones,P] arithmetic).", "why": "The cell->clone map is a fixed integer map; a groupby recomputes it per call, a scatter-add reuses it. numpy index_add/bincount vs pandas groupby is typically 10-50x on wide integer-keyed sums, before any GPU; on GPU (torch scatter, torch already a hard dep) it multiplies again. Dominant per-draw cost.", "blocking": false}, {"area": "Sampling: draw all n_samples at once + batch the softmax; also stop re-sampling all covariates' ct rows", "issue": "The `for i in range(n_samples)` loops (_metrics.py:296 clonotypic, :741 MI, plus phenotypic/flux) re-enter joint_distribution_posterior each iteration, which re-does every .uns lookup, the alignment guard, one Dirichlet.sample() and one scipy softmax. Worse, _preprocessing.py:305 draws Dirichlet over the FULL p_ct (all ct_count rows across every covariate) each call, then indexes down to one covariate.", "severity": "high", "fix": "Draw once: Dirichlet(local_scale * p_ct_mean_for_cov).sample((n_samples,)) -> [n_samples, n_ct_cov, P] (restrict to the covariate's ct rows first), then compute softmax((logits + log_prior)/T) batched over the leading draw axis in one tensor op, then the scatter-sum from the prior finding. Eliminates the Python loop and the per-iteration setup. Route via torch.cuda when available (torch already required), mirroring grafiti embedding.py::_use_gpu. Grafiti precedent: joint.py computes PMI as one [N,S*C] matvec with no per-cell [N,S,C] intermediate.", "why": "The prior mean, logits, and cell->ct/clone maps are constant across draws; only the Dirichlet noise changes. Batching on the leading axis exposes that invariance to one kernel instead of re-deriving it n_samples times, and shrinks the Dirichlet from ct_count to n_ct_cov rows.", "blocking": false}, {"area": "Vectorize entropy/MI reductions over the whole draw stack", "issue": "Entropy is computed with scipy.stats.entropy once per draw (_metrics.py:315) and phenotypic_entropy adds an inner per-clone `jd.loc[cl]` pandas label lookup; MI calls _mi_from_joint per draw (_metrics.py:745). That is O(n_samples * n_clones) Python-level calls with pandas .loc overhead.", "severity": "medium", "fix": "Once the joint is a dense [n_samples, n_clones, P] array, compute entropy as a batched xlogx reduction along the phenotype axis (sum(p*log2 p)) and MI as joint-vs-outer-product-marginals in one shot; take posterior mean + HDI with one nanmean/percentile over the sample axis. Mirror grafiti _compute/joint.py::_entropy/_mi: float64 accumulators, on-device, asnumpy only at the return boundary. The per-clone jd.loc lookup disappears entirely.", "why": "Entropy/MI are elementwise-then-reduce over the phenotype axis; the per-draw scipy call and pandas label indexing are pure Python overhead that vanishes once the substrate is an array stack on the same device as the sampling.", "blocking": false}, {"area": "Algorithmic streamlining: compute the joint once and share it across all metrics; batch groupby groups", "issue": "clonotypic_entropy, phenotypic_entropy, mutual_information and flux each independently re-enter their own sampling loop and rebuild the joint for the same (covariate, n_samples) \u2014 the joint is recomputed per metric. With a groupby, each group is also handled by a separate Python-level recomputation.", "severity": "high", "fix": "Compute the [n_samples, n_clones, P] joint stack once per (covariate) and have every metric reduce that same stack (the dual-input `adata_or_jd` design already allows passing a precomputed joint \u2014 extend it internally to the sampled stack). For groupby, build the joint per group in one batched pass rather than a Python loop over groups. This directly implements the 'avoid recomputing the joint per metric, batch groups' streamlining.", "why": "The joint is the expensive object and is invariant across the four metrics at fixed covariate/n_samples; recomputing it per metric multiplies the dominant cost by the number of metrics requested.", "blocking": false}, {"area": "Deterministic point estimate must do ZERO draws (closed-form mean)", "issue": "n_samples=0 in the current metrics still enters the sampling machinery (clonotypic_entropy requires n_samples>=1 at _metrics.py:289; MI/flux at n_samples=0 return a single random Dirichlet draw, not a point estimate). Sampling-and-averaging only converges to a value available in closed form.", "severity": "medium", "fix": "Per the math spec, n_samples=0 reads adata.uns['tcri_p_ct'] (= E_q[p_ct], temperature-adjusted) directly with no Dirichlet draw. This is simultaneously the correctness fix and the cheapest possible path \u2014 the deterministic estimator does no Monte-Carlo at all. Only n_samples>0 touches local_scale / the sampler.", "why": "q(p_ct)=Dirichlet(local_scale\u00b7m) has mean exactly m=normalize(q_p_ct_raw); local_scale cancels. Drawing and averaging is a strictly worse, and far more expensive, estimator of a closed-form quantity.", "blocking": false}, {"area": "Add tcri/_compute/_xp.py device seam (torch-first, cupy optional, CPU default)", "issue": "There is no device abstraction today; the joint/entropy/MI math is written directly against pandas/scipy/torch-CPU, so there is nowhere to route a GPU backend. The target layout (API surface) also omits a _compute/ package.", "severity": "high", "fix": "Add tcri/_compute/_xp.py copying grafiti's resolve_device/get_xp/asnumpy 1:1 (None/'cpu'->cpu; 'auto'/'gpu'/'cuda'->GPU only if the lib imports AND a device is present, else CPU; explicit 'cuda' warns on fallback, 'auto'/'gpu' silent; 'mps'->cpu). Because torch>=2.4.1 is already a hard dep, make a torch-tensor core the FIRST backend (torch.cuda when available, zero new deps) with cupy as an optional numpy-style second backend. Every accelerated function returns a plain numpy array via asnumpy. This one seam lets the joint/entropy/MI core be written once and routed to numpy / torch-cpu / torch-cuda / cupy.", "why": "Grafiti routes all of _compute through this 58-line seam; it is the precedent and the single dispatch point that makes the acceleration additive and reversible without touching caller code.", "blocking": false}, {"area": "cuML UMAP behind a _use_gpu gate + lazy-import umap off module top", "issue": "The model-latent UMAP uses umap-learn (umap.UMAP at _preprocessing.py:249), and `import umap` sits at module top (_preprocessing.py:20) \u2014 a heavyweight import executed on every `import tcri.preprocessing`, and single-threaded/slow on large n_cells. scVI training already uses GPU via torch, but this step does not.", "severity": "low", "fix": "Swap to cuml.manifold.UMAP behind a _use_gpu(device) gate, keep umap-learn as the CPU fallback, return a plain float64 ndarray, and move the `import umap` (and the lazy cuml import) INSIDE the function so `import tcri` never touches either. Near-verbatim copy of grafiti _compute/embedding.py::umap/_use_gpu; document that GPU and CPU layouts differ but are both valid (as grafiti does).", "why": "cuML UMAP is commonly 10-50x umap-learn, but it runs once per analysis (not inside the sampling loop), so total impact is lower than P0/P1 \u2014 hence low priority. Moving the import also fixes a module-top heavy-import that the grafiti 'never required at import' guardrail forbids.", "blocking": false}, {"area": "Replicate ALL grafiti GPU guardrails, including bounded device memory", "issue": "The batched [n_samples, n_cells, P] tensor (e.g. n_samples=1000 x ~1e6 cells x P) can be enormous; a naive whole-batch GPU allocation will OOM. And nothing in tcri today enforces lazy GPU imports, optional-dep isolation, host-numpy boundaries, or degrade-on-failure.", "severity": "high", "fix": "Copy grafiti's guardrail set uniformly: (1) import every GPU lib lazily inside the function, never at module top \u2014 `import tcri` must never touch cupy/cuml/torch.cuda; (2) keep all GPU libs optional, never in install_requires, CPU path always fully functional; (3) permissive device ladder with getDeviceCount()>0 verification; (4) asnumpy at every return boundary; (5) wrap the GPU body in try/except that falls back to CPU and reports which backend ran; (6) float64 accumulators so GPU matches CPU (grafiti Moran's I); (7) validate the joint (finiteness, nonnegativity, per-row sum~1) on-device before compute; (8) CHUNK the batched reduction over cells or draws to bound device memory \u2014 grafiti Moran's I chunk_size=256 and KDE blk=8192 are the precedent.", "why": "These are exactly the invariants that let grafiti ship GPU as pure opt-in without ever breaking CPU users or crashing on a missing/OOM GPU; skipping the memory chunking specifically turns the batching win into an OOM risk on realistic n_samples x n_cells.", "blocking": false}, {"area": "Seed torch (and GPU) generators, not np.random, for reproducible draws", "issue": "Metric/flux draws use torch Dirichlet(...).sample(), but the seed argument calls np.random.seed (_metrics.py:350/417/789/921), which does not affect torch \u2014 so 'seed' is currently a silent no-op for posterior draws.", "severity": "medium", "fix": "When the sampling moves to a batched torch (and optionally cupy/torch.cuda) draw, seed a torch.Generator (device-aware) \u2014 and cupy/torch.cuda RNG on the GPU path \u2014 alongside numpy. Determinism of the batched draw is a stated test requirement in the refactor plan.", "why": "Directly relevant to moving sampling onto torch/GPU: the acceleration must preserve (and finally deliver) the reproducibility the current seed argument silently fails to provide.", "blocking": false}], "summary": "The tcri sampling stack is dominated by a Python for-loop over 50-1000 draws where each iteration re-enters joint_distribution_posterior, re-draws the FULL p_ct (all ct_count rows) via torch Dirichlet (_preprocessing.py:305), re-softmaxes, and rebuilds a pandas groupby-by-clone (_preprocessing.py:320-322) whose cell->clone grouping is constant across draws. The single biggest win (grafiti contingency.py/edge_tensor precedent) is precomputing clone integer codes once and replacing the pandas groupby with a bincount/index_add scatter-add batched across all n_samples \u2014 10-50x even on CPU, more on GPU. Batching the Dirichlet draw and the softmax on a leading sample axis eliminates the Python loop; vectorizing entropy/MI as xlogx reductions over one [n_samples, n_clones, P] stack (grafiti joint.py precedent, float64 accumulators, asnumpy at boundary) removes the per-draw scipy calls and per-clone pandas .loc lookups. The key algorithmic streamlining is to compute that joint stack ONCE per covariate and share it across all four metrics (they currently each rebuild it), and to make n_samples=0 a zero-draw closed-form read of uns['tcri_p_ct']. Acceleration should route through a torch-first tcri/_compute/_xp.py seam (torch>=2.4.1 is already a hard dep, so torch-CUDA is zero-new-dep; cupy optional) copied 1:1 from grafiti, with cuML UMAP behind a _use_gpu gate for the one-off latent embedding. All of it must replicate grafiti's full guardrail set \u2014 lazy GPU imports (never at module top; the current `import umap` at _preprocessing.py:20 already violates this), optional deps never in install_requires, permissive device ladder, asnumpy boundaries, try/except degrade, float64 parity, validate-before-compute, and chunked reductions to bound device memory on the large [n_samples, n_cells, P] tensor. The one architecturally load-bearing (blocking) recommendation: write the refactored engine core as a batched device-routable function returning the draw stack rather than per-draw DataFrames, so these wins are not designed out of the new layout. Also fold in the torch-generator seeding fix, since draws move to torch and the current np.random.seed is a no-op for them."}, "missing": {"findings": [{"area": "to_anndata write-set / metric<->model agreement", "issue": "The headline fix 'posterior=True joint matches model.predict()' is unimplementable from the adata alone as specified. predict() uses a gate-aware, classifier-temperature-scaled combination gate_prob*logits + (1-gate_prob)*log(p_ct) (model/_model.py:609,648,961). But gate_prob and classifier_temperature are MODEL attributes (self.gate_prob, self.classifier_temperature) never written to uns, and the to_anndata 'canonical minimum' write-set in the API surface (metadata/categories, X_tcri, predict() probs + hard labels, p_ct) drops X_tcri_logits even though _keys.py defines K.X_LOGITS and the engine's posterior path reads adata.obsm['X_tcri_logits'] (preprocessing/_preprocessing.py:311). register_model today additionally stores an ADDITIVE logits+prior (line 229) that already disagrees with the gate-aware predict(). So the engine has neither the logits, the gate, nor the temperature it needs to reproduce predict().", "severity": "high", "blocking": true, "why": "Grounded: gate_prob/classifier_temperature live only on the module; register_model writes X_tcri_logits but the streamlined to_anndata spec omits it; the round-trip test (test_session_round_trip.py) asserts X_tcri/probabilities/p_ct/q_p_ct_raw but never logits, so dropping logits would pass CI while silently breaking every posterior=True metric.", "fix": "Add K.GATE_PROB and K.CLASSIFIER_TEMPERATURE to _keys.py; make to_anndata persist gate_prob, classifier_temperature, AND X_tcri_logits as part of the canonical set. Specify two engine paths explicitly: (a) n_samples=0 & temperature=1 aggregates the stored per-cell probs (X_tcri_probabilities) for an exact, zero-recompute match to predict(); (b) tempering/draws recompute softmax(combine(logits, log b)/T) using the persisted gate_prob+cls_temp. Add a test asserting engine posterior=True per-cell probs == model.predict() before clone aggregation."}, {"area": "engine efficiency across groups/covariates", "issue": "Nothing in the plan makes a single p_ct posterior draw shared across covariates and groups. joint_distribution draws the FULL p_ct Dirichlet per call, and every consumer loops calls: groupby iterates groups, pl.phenotypic_flux calls tl.phenotypic_flux pairwise across an ordered covariate series (2*(M-1) joint builds), mi_compare-style per-patient analysis re-invokes the metric per (patient,covariate,sample), and diag.permutation_null recomputes the metric R times. Each rebuild re-draws all ct rows then filters. Result is O(groups x covariates x n_samples) redundant full-table Dirichlet draws.", "severity": "high", "blocking": true, "why": "Grounded in current code: clonotypic_entropy loops n_samples calls to joint_distribution_posterior (each redraws all p_ct), flux redraws per sample, mi_compare nests patients x covariates x n_samples. The refactor keeps the same per-call draw with no sharing spec; the prompt's stated concern #1 is not addressed anywhere in the drafts.", "fix": "Define covariate=None to compute the joint for ALL covariate values in one pass from a single p_ct draw (the natural all-timepoints path the sankey and multi-covariate metrics need). Specify a draw-once invariant: for n_samples>0, draw the full p_ct table once per sample and reuse it across all groups/covariates/clones, aggregating per group by cell-masking rather than re-drawing. Add a test/counter asserting the number of Dirichlet draws is n_samples, independent of #groups and #covariates."}, {"area": "comparison-stats public API", "issue": "The design deletes mi_compare/delta_entropy_table/flux_table/delta_clonotypic_entropy and asserts 'comparisons happen via groupby + a general stats helper', but _stats.py is a PRIVATE module absent from every __all__ (tcri/_stats.py). There is therefore NO public entry point to run the group comparison, and no mid-level helper that turns a tidy groupby DataFrame into the old outputs (per-pair Mann-Whitney p, Bayesian P(>0), HDI, per-group summary). Users would have to hand-split the tidy frame and call private primitives.", "severity": "high", "blocking": true, "why": "mi_compare (metrics/_metrics.py:93) and delta_entropy_table (:392) currently do this orchestration (pair enumeration, per-group MWU/summary, p_gt/p_lt). The settled design removes them but the replacement is private, so the claimed subsumption is non-functional for any programmatic user.", "fix": "Add a PUBLIC mid-level helper (e.g. tl.compare_groups(df, *, value, by, reference=None, paired=False, hdi_prob=0.94) -> tidy DataFrame with mean/delta/U/p/p_gt/hdi_low/hdi_high) or expose a public tcri.stats surface. Give its exact signature in the API doc so 'groupby + stats subsumes *_compare/*_delta' is actually reproducible, and show the recipe that recreates mi_compare's per-pair output."}, {"area": "safe-deletion caller census", "issue": "The Phase-2 'confirmed 0-caller / dead' deletion list is computed against package source only, not the example notebooks. A grep of example/ shows gene_entropy used 5x, polar_plot 1x, register_clonotype_key 23x, register_phenotype_key 11x, plus pl.flux (boxplot) 5x and pl.probability_ternary 24x. Deleting/privatizing these orphans live notebook cells.", "severity": "high", "blocking": true, "why": "Directly grounded: `grep -rhoE ... example/` returns gene_entropy=5, polar_plot=1, register_clonotype_key=23, register_phenotype_key=11. The plan calls gene_entropy a 'confirmed 0-caller' \u2014 false against the notebook corpus. Notebook execution is itself a Phase-10 acceptance gate, so this contradicts the plan's own CI.", "fix": "Run a caller census over example/ + docs notebooks BEFORE Phase 2 and reclassify. For genuinely-used-but-removed functions, provide either a kept public path (standalone key registration) or an explicit notebook-rewrite mapping in the same PR. Do not mark anything 'safe delete' on package-only evidence."}, {"area": "plotting axis coverage (splitby)", "issue": "splitby appears in 116 example call sites and is the most-used plotting argument, yet the new pl twins expose only groupby. Current plots carry TWO categorical axes simultaneously (tcri_boxplot groupby+splitby, plotting/_plotting.py:633; clonotypic_entropy_by_phenotype uses phenotype x covariate(hue) x patient(dots)). Collapsing to a single groupby loses a dimension that most published figures rely on (e.g. dots=patient, boxes=response, x=phenotype).", "severity": "high", "blocking": false, "why": "Grounded: 116 splitby usages; tcri_boxplot has a dedicated groupby!=None & splitby!=None branch; the entropy-by-phenotype plot enumerates itertools.product(batches, covariates). The new pl.* signatures list only groupby, no splitby.", "fix": "Either retain a splitby axis on the pl twins (groupby=dots, splitby=box hue) or document, per figure type, the exact groupby-only + _base._metric_boxplot recipe that reproduces each two-axis plot. Decide before Phase 7 since it changes the pl signatures and _metric_boxplot contract."}, {"area": "covariate x groupby semantics", "issue": "How groupby interacts with the model's per-(clone,covariate) conditioning is unspecified. p_ct is indexed by ct=(clone,covariate) with no per-cell/per-patient axis. groupby on an arbitrary obs column (patient, response) can only partition CELLS, so it is only well-defined on the posterior=True (cell-informed) path; with posterior=False (ct-level table) a groupby is ill-defined unless the group label is constant within each clone x covariate. The whole scheme silently assumes clones are disjoint across patients (a TCR clone never spans two patients), which is never stated.", "severity": "medium", "blocking": true, "why": "Grounded: joint_distribution/joint_distribution_posterior have no groupby today; the per-patient plots achieve grouping only by restricting the `clones` list to that patient's clones (plotting/_plotting.py:714, metrics mi_compare:125) \u2014 which works only because clones are patient-specific. The refactor adds a first-class groupby without stating this precondition or which posterior branch it requires.", "fix": "Specify that groupby requires the cell-informed path (or that group keys must be clone-nested / constant within clone x covariate), and document the clone-disjoint-across-groups assumption. Add a validation that errors when a group split would bisect a clone's cells under posterior=False."}, {"area": "API surface completeness", "issue": "Several public symbols actually used in notebooks are absent from every keep/drop/move list in the drafts: pl.probability_ternary (24 uses), tl.phenotypic_entropies and tl.clonotypic_entropies (plural batch forms, 15 uses combined), pl.phenotypic_entropy_delta (3), and pl.flux as a boxplot distinct from the sankey (5). The inventory claims 131 functions but the disposition tables don't cover these.", "severity": "medium", "blocking": false, "why": "Grounded in the example grep counts. The plural *_entropies and probability_ternary do not appear anywhere in the settled design, API surface, or deltas, so they'd silently vanish with no migration target.", "fix": "Reconcile the disposition tables against the set of public symbols actually imported/called in example/ and docs. Assign each an explicit keep/rename/move/drop, especially probability_ternary (ternary phenotype-simplex plot) and the plural convenience wrappers."}, {"area": "reproducibility / seed argument", "issue": "The settled design lists 'seed the torch Generator' as a bugfix, but the final signatures for joint_distribution and all four metrics have NO seed/random_state parameter. So n_samples>0 posterior draws are irreproducible through any public knob \u2014 the exact bug the design set out to fix (flux's seed only touched numpy while draws use torch) recurs by omission.", "severity": "medium", "blocking": false, "why": "Grounded: current flux has seed= (metrics/_metrics.py:864) but it only calls np.random.seed while draws use torch Dirichlet; the new engine signature joint_distribution(adata,*,covariate,groupby,n_samples,posterior,clones,temperature) drops seed entirely and none of the metric signatures add it.", "fix": "Add random_state (int | torch.Generator) to joint_distribution and to every sampling metric/diag function; thread it into the torch Dirichlet draw. Document that n_samples=0 is deterministic regardless."}, {"area": "temperature composition", "issue": "uns[tcri_p_ct] is already guide_temperature-adjusted at write time (get_p_ct applies ^(1/guide_temperature) then normalizes, model/_model.py:510). The engine/metrics then apply a SECOND `temperature` on top. The two temperings are never reconciled, so the effective exponent seen by a metric is 1/(guide_temperature*metric_T) or a re-softmax of an already-tempered vector, depending on path.", "severity": "medium", "blocking": false, "why": "Grounded: get_p_ct bakes guide_temperature into the stored p_ct; joint_distribution then does softmax(log(p_ct)/T). Double application is real and undocumented.", "fix": "State explicitly whether metric temperature composes with or is independent of guide_temperature, and whether the stored p_ct is the guide-tempered or raw mean. Recommend storing the raw posterior mean and letting metric temperature be the single knob, or document the composition."}, {"area": "normalised-entropy comparability", "issue": "clones/groupby change the entropy normalizer (H divided by log2(n_clones) for clonotypic, and the whole-table normalization imposes P(c)=1/C for MI). When groupby restricts to a group's clones, each group gets a DIFFERENT denominator, so normalized entropy/MI values are on different scales and are not directly comparable across groups \u2014 yet the pl twins plot exactly that cross-group comparison.", "severity": "medium", "blocking": false, "why": "Grounded: clonotypic_entropy divides by log2(len(vec)) (metrics/_metrics.py:317) where len(vec)=restricted clone count; mutual_information normalizes the whole table (:731) imposing uniform clone prior. Per-group clone counts differ, so denominators differ.", "fix": "Either offer a fixed global normalizer option (log2 of the full clone count) for cross-group plots, or document in each metric+pl twin that normalised values are within-group only and default the comparison plots to a comparable convention."}, {"area": "@tl_result / uns-cache forward-compat", "issue": "The plan says 'design every tl fn to return a clean tidy result so wrapping @tl_result is a one-liner later', but the chosen return shapes are h5ad-hostile: MultiIndex frames (group,clone,sample_id), per-draw object-array columns (delta_samples/flux_samples), and provenance in df.attrs['params']. pandas .attrs and object-dtype arrays do not survive h5ad round-trips, so a uns cache would drop provenance and shape.", "severity": "medium", "blocking": false, "why": "Grounded: delta_entropy_table/flux_table already stash full numpy vectors in object columns (metrics/_metrics.py:461, 840); the API surface puts provenance in df.attrs. AnnData write drops both. The deferred decorator can't cache these losslessly.", "fix": "Constrain tl return frames now to be h5ad-serializable: flat columns, no object arrays (store draws in a separate long frame or a uns sidecar), and put provenance in a serializable column or a companion uns key rather than .attrs. Define the cache-key scheme (hash of covariate/groupby/n_samples/temperature/clones/normalised/normalise_mode/distance_metric)."}, {"area": "n_samples>0 reduction correctness", "issue": "How a metric reduces the engine's large n_samples>0 tidy frame is unspecified, and correctness hinges on it: all clones within one draw must share the SAME p_ct draw (one coherent joint per sample), not independent per-clone draws. If the metric groups the tidy frame wrong (e.g., per-clone instead of per sample_id), MI/entropy posterior spread is corrupted.", "severity": "medium", "blocking": false, "why": "Grounded: current metrics preserve this by building one joint per draw inside the loop (mutual_information:742, clonotypic_entropy:296). The refactor externalizes draws into a sample_id level; the reduction from that frame back to per-draw metrics is not written down.", "fix": "Specify that metrics iterate the sample_id level and compute the full-joint metric per draw, then summarize. Add a test that n_samples>0 MI variance matches repeated coherent-joint draws (and differs from independent per-clone draws)."}, {"area": "group_singletons ordering", "issue": "group_singletons rewrites clone labels (collapsing small clones into Singleton_{group}), but the plan keeps it a separate pp step without stating it must run BEFORE setup_anndata/train. If run after, the model's learned clonotype categories (and p_ct's ct_to_c mapping) desync from the obs clone labels the engine reads.", "severity": "low", "blocking": false, "why": "Grounded: p_ct/ct_to_c are learned over the registered clone categories (register_model:186, module ct maps). group_singletons changes those labels; ordering is load-bearing but unspecified.", "fix": "Document and enforce that group_singletons (and any clone relabeling) precedes setup_anndata; optionally have setup_anndata refuse to register if a later relabel is detected."}, {"area": "key naming reconciliation", "issue": "The moved-to-examples plot_phenotype_probabilities reads obsm['X_tcri_phenotypes'] (plotting/_plotting.py:813 default slot), but the canonical key is K.X_PROBABILITIES='X_tcri_probabilities' and register currently writes a configurable phenotype_prob_slot. The example and the canonical key don't match, so the moved example breaks.", "severity": "low", "blocking": false, "why": "Grounded: default slot name mismatch between the plot function and the _keys registry.", "fix": "When moving the function to examples/, update it to read K.X_PROBABILITIES, or add a documented alias. Verify no other consumer references X_tcri_phenotypes."}, {"area": "diag signatures / provenance / test gaps", "issue": "Three loose ends the individual lenses each miss: (1) diag functions default model=None but reconstruction_ppc needs the live ZINB decoder (model-only) while joint_distribution_ppc can run from adata alone \u2014 the required-vs-optional model is inconsistent and unspecified. (2) metrics accept a precomputed joint df but the returned metric frame is not specified to carry through df.attrs provenance, breaking traceability of the fast path. (3) The test matrix lacks a golden regression tying the new deterministic metrics to a hand-computed value from uns[p_ct], and lacks the engine==predict() agreement test that validates the central fix.", "severity": "low", "blocking": false, "why": "Grounded: reconstruction_ppc simulates from the fitted decoder (px_r/decoder live on the module, not adata); joint_distribution_ppc only needs engine+empirical counts. Provenance-in-.attrs won't propagate unless stated. \u00a76.1 tests assert n_samples=0==uns[p_ct] but no numeric golden value and no model-agreement assertion.", "fix": "Make diag model-requirement explicit per function (adata-only vs model-required) and validate inputs. Specify that metrics propagate/merge the input joint's provenance into their output. Add (a) a golden test computing MI/entropy from a fixed uns[p_ct] by hand, and (b) a test asserting posterior=True joint per-cell == model.predict()."}], "summary": "The three drafts (API surface, math spec, implementation plan) are internally strong but leave load-bearing seams between the lenses unaddressed. The most serious gaps: (1) the headline 'metrics agree with model.predict()' fix is unimplementable as specified because gate_prob, classifier_temperature, and X_tcri_logits are not in the to_anndata write-set/_keys registry (gate_prob/cls_temp live only on the module; the streamlined to_anndata drops logits); (2) no mechanism makes a single p_ct draw shared across groups/covariates, so groupby, the flux sankey, and permutation-null are O(groups x covariates x n_samples) \u2014 the prompt's explicit efficiency concern; (3) the design deletes *_compare/*_delta in favor of 'groupby + _stats' but _stats is a private module with no public comparison entry point, so the subsumption is non-functional; (4) the 'safe deletion' 0-caller analysis was run against package source only \u2014 gene_entropy (5), polar_plot (1), register_*_key (34) and others have live example-notebook callers; and (5) splitby (116 call sites) and probability_ternary (24) are dropped/undispositioned, collapsing a whole plotting dimension. Secondary gaps: no seed on the engine/metrics, double guide/metric tempering, cross-group non-comparability of normalized entropy, h5ad-hostile return shapes for the deferred @tl_result cache, unspecified n_samples>0 reduction (draw-coherence), group_singletons ordering vs setup_anndata, a key-name mismatch, and missing golden/agreement tests. Findings are ordered most-blocking first; items 1-4 should be resolved before Phases 4-6 land."}}, "ground": {"gpu": {"grafiti_patterns": "GRAFITI GPU/ACCELERATION PATTERNS (the wins to copy). All live under grafiti/_compute/ plus grafiti/model/grafiti.py.\n\n1) THE CANONICAL FALLBACK \u2014 grafiti/_compute/_xp.py (58 lines, the reference implementation)\n Three functions: `resolve_device(device)`, `get_xp(device)`, `asnumpy(x)`.\n - resolve_device: None/'cpu'->'cpu'; 'mps'->'cpu' (no Metal backend yet); 'cuda'/'gpu'/'auto'->'cuda' ONLY if `import cupy` succeeds AND `cupy.cuda.runtime.getDeviceCount()>0`, else 'cpu'. Explicit 'cuda' warns on fallback; 'auto'/'gpu' fall back silently; unknown device warns.\n - get_xp: returns the array module \u2014 cupy on cuda, numpy otherwise. cupy imported lazily INSIDE the function.\n - asnumpy: `cupy.asnumpy(x)` if x is a cupy array else `np.asarray(x)` \u2014 the single host-boundary shim. Callers stay backend-agnostic.\n Consumers just do: `xp = get_xp(device); ...; return asnumpy(result)`.\n\n2) NEVER-DENSIFY-N\u00d7G + STREAMED/BLOCKED REDUCTION, DEVICE-ROUTED WHOLE \u2014 grafiti/_compute/stats.py :: morans_i_streaming (L296-383)\n Streams gene-column blocks of chunk_size=256; at most chunk_size columns are densified at once, so peak extra memory is O(n_cells*chunk), never the full n_cells\u00d7n_genes dense matrix. The SAME routine runs GPU-native: `import cupy as xp; import cupyx.scipy.sparse as xsp`, Wn->csr on device, X->CSC on device (or dense via xp.asarray), then the per-block densify + spatial lag (Wn @ block) + num/den reductions ALL run on GPU in float64, with `to_np = xp.asnumpy` only at the boundary. Explicitly float64 so the GPU result MATCHES the CPU path rather than approximating. b0389d9 shows they even patched a dense-input case on this GPU path.\n\n3) GPU UMAP / PCA (RAPIDS cuML) \u2014 grafiti/_compute/embedding.py\n `_use_gpu(device)` gate: True only if GPU requested AND `import cuml` succeeds (warns on explicit cuda/gpu fallback). umap(): GPU branch `from cuml.manifold import UMAP as cuUMAP; cuUMAP(output_type=\"numpy\", ...).fit_transform(X)`, CPU branch `umap-learn`. pca(): `cuml.decomposition.PCA` vs `sklearn.decomposition.PCA`. Both branches return a plain float64 ndarray so callers are backend-agnostic. Docstring is candid that GPU and CPU give different (both valid) layouts \u2014 parity claimed only where achievable.\n\n4) GPU LEIDEN (rapids-singlecell) \u2014 grafiti/model/grafiti.py :: _run_leiden_gpu (L1799) / _run_leiden_cpu (L1790)\n GPU: `import cupy as cp; import rapids_singlecell as rsc`; build AnnData on cp.asarray(z), `rsc.pp.neighbors(...)`, `rsc.tl.leiden(...)`. CPU mirror: scanpy `sc.pp.neighbors`+`sc.tl.leiden`. Contract: GPU returns `(labels, backend_str)` or `(None, reason_str)` on any import/runtime failure; caller (L1973) prints the backend and falls back to CPU. Extra SIZE/FEATURE GATES: `leiden_gpu_max_features=150`, `leiden_warmup_max_cells=50000` cap the working subset when the GPU path is unavailable, so the CPU fallback stays tractable.\n\n5) GPU kNN smoothing / label projection (cuML) \u2014 grafiti/model/grafiti.py :: _get_soft_labels_from_leiden_gpu (L2095), _predict_labels_from_subsample_gpu (L2150)\n `from cuml.neighbors import NearestNeighbors`; kNN on cp arrays; soft-label smoothing via `cp.add.at(probs,(row,col),weights)` scatter (GPU analogue of np.add.at); label projection via one-hot vote `cp.eye(K)[neighbor_labels].sum(1).argmax(1)`. Each returns `(result, backend)` / `(None, reason)`; the plain method (_get_soft_labels_from_leiden L2133) tries GPU then falls to sklearn CPU.\n\n6) GPU DELAUNAY SPATIAL GRAPH (cupyx) \u2014 grafiti/_compute/graph.py :: _build_cupy_delaunay_spatial_neighbors (L98) / build_delaunay_spatial_neighbors (L198)\n `from cupyx.scipy.spatial import Delaunay`; per-FOV triangulation on GPU with a PER-FOV scipy.spatial.Delaunay fallback caught inside the loop (records cpu_fallback_fovs), and a whole-op try/except that falls back to the squidpy CPU builder. Device dispatch via `getattr(device,'type',device)!='cuda'`.\n\n7) DEVICE-ROUTED BINCOUNT CROSSTAB (the direct analogue of tcri's joint-by-clone) \u2014 grafiti/_compute/contingency.py :: contingency (L12), and graph.py :: edge_tensor (L257)\n `xp = get_xp(device); a=xp.asarray(codes_a).astype(xp.int64); ...; counts = xp.bincount(a*n_b+b, minlength=n_a*n_b).reshape(n_a,n_b); return asnumpy(counts)`. One O(N) bincount over cells replaces a groupby; all the small [S,C] arithmetic sits on top. edge_tensor scatters a flat 3-D index with a single bincount the same way.\n\n8) DEVICE-ROUTED POPULATION JOINT / ENTROPY / MI \u2014 grafiti/_compute/joint.py\n Loads q2d as [N,S,C] onto device via get_xp; `_mean_joint` computes P=q.mean(axis=0, dtype=xp.float64) (float64 accumulator keeps the O(N) mean drift-free); per-cell PMI computed as ONE [N,S*C]@[S*C] matvec \"no [N,S,C] intermediate\" (L123). Validation (isfinite/min/row-sum) runs ON DEVICE before compute; asnumpy only at the return boundary.\n\n9) torch-CUDA BLOCKED KDE EVAL \u2014 grafiti/_compute/occupancy.py :: kde_eval_torch (L191), _torch_cuda (L171)\n scipy fits the KDE (keeps bandwidth+weights); torch evaluates on the grid in blocks of 8192 via torch.cdist on cuda; `_torch_cuda(device)` returns a torch.device or None; any exception falls back to scipy `kde(gp)`. Blocked to bound device memory \u2014 same shape as the Moran's I chunking.\n\n10) DEVICE-ROUTED connected_components \u2014 grafiti/_compute/patch.py :: _connected_components (L18): cupyx.scipy.sparse.csgraph on cuda, scipy on cpu, via `_xp.resolve_device`.", "guardrails": "GUARDRAILS grafiti applies uniformly (tcri must replicate ALL of these):\n\n1) NEVER required at import. Every GPU library (cupy, cupyx, cuml, rapids_singlecell, torch.cuda) is imported LAZILY inside the function that uses it, never at module top. Module docstrings state this explicitly (\"cupy imported lazily only when device='cuda'\", \"squidpy/cupy imported lazily so `import grafiti._compute` stays light\"). `import grafiti` never touches a GPU lib.\n\n2) Optional dependency, never in install_requires. Nothing GPU is a hard dep; the CPU path is always fully functional. GPU is pure opt-in via device=.\n\n3) Permissive device resolution with a graceful ladder. 'auto'/'gpu'/'cuda' all attempt GPU; if cupy/cuml/CUDA is missing they fall to CPU. Only an EXPLICIT 'cuda' warns on fallback (the user asked for it by name); 'auto'/'gpu' fall back silently. 'mps'->cpu. Unrecognized device -> warn + cpu. GPU presence is verified (getDeviceCount()>0), not assumed from a successful import.\n\n4) Host-numpy at the boundary, always. Every GPU function returns a plain numpy ndarray (via asnumpy, or output_type=\"numpy\" for cuML) so no device array ever escapes into caller code. Results are backend-identical in TYPE regardless of where they ran.\n\n5) try/except wraps the GPU body; failure degrades, never crashes. GPU helpers return `(None, reason_str)` or fall through to the CPU implementation on ANY import or runtime exception, and print/return which backend actually ran (Leiden, kNN, Delaunay all do this). A broken/OOM GPU path can never take down an analysis.\n\n6) Numerical parity where it must match, candor where it can't. Moran's I runs float64 on GPU so it MATCHES the CPU result and squidpy parity (documented). UMAP docstring states GPU vs CPU layouts differ (both valid) \u2014 they do NOT overpromise bit-identity for stochastic embeddings.\n\n7) Correctness gate before compute. joint.py validates q2d invariants (finiteness, nonnegativity, per-row sum\u22481 \u2014 per-row, not aggregate, to catch defects that cancel) on-device before doing the math.\n\n8) Bounded memory on both host and device. Chunking (Moran's I chunk_size=256; KDE blk=8192) bounds peak memory on whichever backend runs, and the size/feature caps (leiden_gpu_max_features=150, leiden_warmup_max_cells=50000) keep the CPU fallback tractable.", "opportunities": [{"operation": "Per-draw joint-table construction: softmax over selected cells then groupby-sum by clone", "gpu_approach": "Precompute clone integer codes ONCE (outside the draw loop) and replace pandas .groupby(level=0).sum() with a constant-index scatter-add: numpy np.add.at / np.bincount per phenotype column on CPU, or torch.zeros(n_clones,P).index_add_(0, clone_codes, probs) / cupy bincount on GPU \u2014 batched across all n_samples in one kernel. This is grafiti/_compute/contingency.py::contingency and graph.py::edge_tensor verbatim (one O(N) bincount over cells, rest is small [n_clones,P] arithmetic).", "tcri_hot_op": "tcri/preprocessing/_preprocessing.py::joint_distribution_posterior L320-322 \u2014 pd.DataFrame(probs_cell, index=clone_labels).groupby(level=0).sum(); rebuilt on every one of the 50-1000 draws inside mutual_information/clonotypic_entropy/phenotypic_entropy/delta_entropy loops (metrics L296, L363, L548, L741, L987).", "current_cpu": "pandas DataFrame build + groupby(level=0).sum() over the covariate's cells (often 1e5-1e6 rows), redone from scratch every draw even though the cell->clone grouping is CONSTANT across draws. pandas hashes and regroups the keys each time. Dominant per-draw cost.", "expected_gain": "Large even on CPU: numpy index_add/bincount replacing pandas groupby is typically 10-50x on wide integer-keyed sums, and the cell->clone code map is computed once instead of n_samples times. On GPU (torch scatter, torch already a hard dep) multiply again on large covariates. Overall the biggest single win.", "mechanism": "Cell->clone membership is a fixed integer map for a given covariate; a groupby recomputes it per call, a scatter-add reuses it. Moving the reduction to bincount/index_add turns O(draws) pandas passes into one precompute + O(draws) vectorized scatters.", "priority": "P0 - highest"}, {"operation": "Posterior Dirichlet sampling looped n_samples times, each redrawing p_ct and re-softmaxing the per-cell logits", "gpu_approach": "Draw ALL n_samples posterior samples at once: Dirichlet(local_scale*p_ct_mean+eps).sample((n_samples,)) -> tensor [n_samples, C, P] on GPU (torch, already required). Gather log_prior by ct_idx and compute softmax((logits+log_prior)/T) BATCHED over draws in one tensor op [n_samples, n_cells, P]. Eliminates the python draw loop entirely; run on cuda when available via an _xp/torch-device gate mirroring grafiti/_compute/embedding.py::_use_gpu.", "tcri_hot_op": "The `for i in range(n_samples): jd = joint_distribution_posterior(...)` loops in tcri/metrics/_metrics.py (L296 clonotypic_entropy, L363 delta, L548 phenotypic_entropy, L741 mutual_information, L987 flux); each iteration calls _preprocessing.py::joint_distribution_posterior L305 (Dirichlet.sample()) + L314 (scipy softmax).", "current_cpu": "torch Dirichlet.sample() for ONE draw then .numpy(), and scipy.special.softmax over [n_cells, n_pheno] \u2014 repeated 50-1000x in a pure-python loop with all the AnnData .uns lookups, validation, and DataFrame overhead re-executed each iteration.", "expected_gain": "Removes the python-loop and per-iteration setup overhead (n_samples-fold), and batches sampling+softmax into one GPU kernel. Combined with P0 this collapses a 200-1000-iteration python loop into a handful of batched tensor ops.", "mechanism": "The prior mean p_ct, the logits, and the cell->ct/clone maps are all constant across draws; only the Dirichlet noise changes. Batching draws on the leading axis exposes that invariance to a single vectorized kernel instead of re-deriving it per iteration.", "priority": "P0/P1 - high"}, {"operation": "Entropy / mutual-information reductions over posterior draws (and the per-clone entropy loop)", "gpu_approach": "Compute entropy/MI as vectorized xlogx reductions over the whole [n_samples, n_clones, P] (or [n_samples, S, C]) stack in one shot \u2014 sum(p*log2 p) along the phenotype axis, MI via joint vs outer-product marginals \u2014 instead of scipy.stats.entropy / _mi_from_joint per draw and jd.loc[cl] per clone. Mirror grafiti/_compute/joint.py::_entropy/_mi (float64 accumulators, on-device, asnumpy at the boundary).", "tcri_hot_op": "tcri/metrics/_metrics.py: the entropy() call per draw and the inner `for j, cl in enumerate(clones_list): p = jd.loc[cl]` loop in phenotypic_entropy (L562-569); _mi_from_joint per draw in mutual_information (L744-745); clonotypic_entropy L315.", "current_cpu": "scipy.stats.entropy / _mi_from_joint invoked once per draw in python, plus a per-clone pandas label-index (jd.loc[cl]) loop \u2014 O(n_samples * n_clones) python-level calls with pandas .loc overhead.", "expected_gain": "Medium-high: replaces two nested python loops with a single batched reduction; the per-clone jd.loc[cl] indexing (slow pandas label lookup) disappears once the joint is a dense [draws, clones, P] array. Also enables computing the full posterior mean+CI with one nanmean/percentile over the axis.", "mechanism": "Entropy and MI are elementwise-then-reduce over the phenotype axis; once the joint table is an array stack rather than per-draw DataFrames, the reduction vectorizes cleanly and moves to the same device as the sampling.", "priority": "P1 - high"}, {"operation": "Model latent UMAP embedding", "gpu_approach": "Swap umap-learn for cuml.manifold.UMAP behind a _use_gpu(device) gate, keeping umap-learn as the CPU fallback and returning a plain float64 ndarray \u2014 a near-verbatim copy of grafiti/_compute/embedding.py::umap/_use_gpu. Optionally cuml PCA for any pre-reduction.", "tcri_hot_op": "tcri/preprocessing/_preprocessing.py `import umap` (L20-21) \u2014 the CPU UMAP step used for the model latent embedding; scVI training itself already uses GPU via scvi-tools/torch.", "current_cpu": "umap-learn UMAP.fit_transform on CPU (single-threaded-ish, slow on large n_cells).", "expected_gain": "Large one-off speedup on big datasets (cuML UMAP is commonly 10-50x umap-learn), but it runs once per analysis, not inside the sampling loop \u2014 so lower total impact than P0/P1.", "mechanism": "UMAP is embarrassingly GPU-friendly (kNN + layout); cuML implements both on device. Not bit-identical to CPU (document the differing-but-valid layout, as grafiti does).", "priority": "P2 - medium"}, {"operation": "Structural recommendation: add tcri/_compute/_xp.py and refactor the joint into a device-routed core", "gpu_approach": "Create tcri/_compute/_xp.py copying grafiti's resolve_device/get_xp/asnumpy 1:1. Because torch>=2.4.1 is ALREADY a hard tcri dep, prefer a torch-tensor core (torch.cuda when available, zero new deps) as the first backend, with cupy as an optional second backend for numpy-style ops. Refactor joint_distribution_posterior into a device-routed `_joint_draws(adata, cov, n_samples, device)` returning a [n_samples, n_clones, P] array (precompute clone codes once; batch Dirichlet+softmax; scatter-sum by clone); keep the current pandas-DataFrame return only as a thin single-draw public wrapper. Metrics consume the stack and reduce vectorized.", "tcri_hot_op": "Consolidates the hot path shared by _preprocessing.py::joint_distribution_posterior (L264) and all _metrics.py sampling loops.", "current_cpu": "No device abstraction today; every metric re-enters pandas/scipy per draw. torch is imported but only for the model, not for the metric reductions.", "expected_gain": "Enables all of P0-P2 behind one dispatch and one CPU/GPU fallback contract; keeps CPU users unaffected (torch-CPU or numpy path) while GPU users get the batched kernels for free.", "mechanism": "A single _xp seam lets the joint/entropy/MI math be written once and routed to numpy, torch-cpu, torch-cuda, or cupy \u2014 exactly grafiti's structure \u2014 so the acceleration is additive and reversible.", "priority": "P0 - foundational (do first; it is the vehicle for P0-P2)"}]}, "mathX": {"functions": [{"name": "joint_distribution (tcri/preprocessing/_preprocessing.py:364-490)", "formula": "PRIOR-ONLY joint over (clone c, phenotype \u03c6) at covariate m, built from the guide posterior-mean p_ct (NOT the per-cell classifier). Select the ct rows with ct_to_cov==m: P0 = p_ct[chosen]. Temper: p\u0303_{ct} = softmax( log(P0+\u03b5)/T ) (\u03b5=1e-8) = P0^{1/T}/\u03a3_\u03c6 P0^{1/T}. n_samples=0 (point est): row per chosen ct, value = p\u0303_{ct}. n_samples>0: draw S_{ct} ~ Dirichlet(local_scale \u00b7 p\u0303_{ct}), shape (n_samples, num_ct, P). weighted: multiply each row by cell-count n_ct then divide by grand total \u03a3 so whole table sums to 1.", "operations": "1) tensors p_ct(ct_count\u00d7P), ct_to_cov, ct_to_c. 2) cov_value=index(label). 3) chosen_mask=ct_to_cov==cov_value; p_ct_for_cov=p_ct[mask]. 4) F.softmax(log(p_ct_for_cov+1e-8)/T, dim=-1). 5) clone_indices=ct_to_c[chosen_idx]. 6) Counter over ct_array_for_cells[cov_array_for_cells==cov_value] \u2192 ct_counts_dict (keyed by CT index). 7a) n_samples=0: DataFrame(p\u0303), attach clonotype_index/id, filter clones, optional weight, index=clonotype_id. 7b) n_samples>0: Dirichlet(local_scale\u00b7p\u0303).sample((n_samples,)); transpose(1,0,2).reshape(-1,P); repeat clone ids; index=f'{cid}_{sid}'.", "args_effects": "covariate_label\u2192selects ct rows via ct_to_cov. temperature T\u2192power-tempering p^{1/T}; T<1 sharpens, T>1 flattens; T=1 is \u2248identity (just renormalizes). n_samples: 0=deterministic posterior-mean rows, >0=Dirichlet draws with concentration local_scale\u00b7p\u0303 (variance \u2193 as local_scale \u2191). clones: filters rows by clonotype_id. weighted: True multiplies rows by clone cell-count then normalizes the ENTIRE table to sum 1 (a joint P(c,\u03c6)); False leaves rows as tempered p\u0303 (each already a per-clone simplex). local_scale read from uns['tcri_local_scale'] (default 1.0), only used when sampling.", "normalizations": "Temperature step renormalizes each row via softmax. weighted path divides by total_mass=\u03a3_{c,\u03c6} (whole-table L1 \u2192 sums to 1). Unweighted rows are already simplex from softmax. Dirichlet samples are simplex by construction.", "edge_cases": "ValueError if covariate_label not in categories. total_mass>0 guard before dividing. No eps other than 1e-8 in the log. Within one covariate each clone maps to exactly one ct, so there is exactly one row per clone (no groupby aggregation).", "bugs": "WEIGHTING BUG: ct_counts_dict is keyed by CT (clone\u00d7cov) indices from ct_array_for_cells, but df['clonotype_index']=ct_to_c[chosen_idx] holds CLONE indices; the weight lookup ct_counts_dict.get(clone_index) therefore reads the wrong bucket (usually 0/garbage). Should key on the ct index (chosen_idx), not the clone index. Consequence: weighted=True produces near-arbitrary weights. Also normalization convention (whole-table sum=1) is INCONSISTENT with joint_distribution_posterior's weighted path (which returns raw un-normalized counts). Temperature via softmax(log p) subtly differs from true p^{1/T}/\u03a3 only by the \u03b5 floor, negligible."}, {"name": "joint_distribution_posterior (tcri/preprocessing/_preprocessing.py:264-329)", "formula": "POSTERIOR per-cell joint. Draw ONE Dirichlet sample of the clone\u00d7cov prob: s_{ct} ~ Dirichlet(local_scale\u00b7p_ct + 1e-8). For each cell i in covariate m: if combine_with_logits: P(\u03c6|i) = softmax( (\u2113_i + log(s_{ct(i)}+1e-8)) / T ) where \u2113_i=obsm['X_tcri_logits'][i]; else P(\u03c6|i)=s_{ct(i)}. Aggregate by clone: J[c,\u03c6] = \u03a3_{i\u2208c} P(\u03c6|i) (groupby-sum). Unweighted: row-normalize J[c,\u00b7] /= \u03a3_\u03c6 J[c,\u03c6] \u2192 P(\u03c6|c). Weighted: leave J as summed masses (\u221d clone size).", "operations": "1) resolve cov_col/clone_col/ph_cats/cov_idx. 2) alignment guard: len(ct_per_cell)!=n_obs \u2192 ValueError (view/subset misalignment, Notion #4). 3) idx_cov=nonzero(cov_per_cell==cov_idx); optional clone filter via np.isin. 4) p_ct_mean=tensor(uns['tcri_p_ct']); non-finite\u21921/n_phen; s=Dirichlet(local_scale\u00b7p_ct_mean+1e-8).sample().numpy() (single draw). 5) combine: log_prior=log(s[ct_idx_sel]+1e-8); probs_cell=scipy.softmax((logits+log_prior)/T,axis=1). 6) DataFrame(probs_cell, index=clone_labels).groupby(level=0).sum(). 7) if not weighted: div by row-sum, fillna 0. 8) if clones: reindex(clones).fillna(0). 9) round(precision).", "args_effects": "covariate_label\u2192cell selection. temperature T\u2192softmax temperature on (logits+log prior); ONLY applied when combine_with_logits=True. clones\u2192isin filter on cells + final reindex to that exact list/order (missing clones become all-zero rows). weighted: False\u2192rows are P(\u03c6|c) (simplex); True\u2192rows are expected phenotype cell-counts (NOT normalized, NOT divided by any total). combine_with_logits: True\u2192full posterior softmax(logits+log sampled-prior); False\u2192ignores logits AND temperature, returns raw sampled prior per cell (so unweighted collapses to a single s_{ct} row per clone). precision\u2192rounding. silent\u2192suppresses prints. local_scale from uns (default 1.0).", "normalizations": "Per-cell softmax \u2192 simplex per cell. groupby-sum then optional row L1 normalization for unweighted. Weighted path performs NO normalization (unlike joint_distribution which normalizes the whole table).", "edge_cases": "Hard alignment guard raises if per-cell .uns arrays don't match n_obs (protects against sliced AnnData). Non-finite p_ct entries replaced elementwise by 1/n_phen (row may not re-sum to 1, but only used as Dirichlet concentration). reindex(clones) fills absent clones with all-zero rows \u2192 downstream entropy clips to \u03b5 \u2192 those rows become ~uniform (max entropy).", "bugs": "combine_with_logits=False silently IGNORES temperature. weighted=True returns un-normalized counts, inconsistent with joint_distribution's weighted (whole-table sum=1). Uses the SAMPLED prior draw s_{ct} (Dirichlet) whereas the model's get_cell_phenotype_probs uses the posterior-MEAN p_ct \u2014 so metric per-cell probs \u2260 stored model probs. It also ALWAYS uses additive logits+log_prior and NEVER applies gate_prob, so if the model was trained with a gate the metric's combination rule disagrees with training. Only one Dirichlet draw per call \u2192 callers must re-invoke in a loop for posterior spread (expensive, re-samples full p_ct each time)."}, {"name": "_mi_from_joint (tcri/metrics/_metrics.py:74-89)", "formula": "Given an already-normalized joint pxy (C\u00d7P, \u03a3=1): px=\u03a3_\u03c6 pxy (C\u00d71), py=\u03a3_c pxy (1\u00d7P). MI = \u03a3_{c,\u03c6} pxy\u00b7log2( (pxy+\u03b5)/(px\u00b7py\u1d40+\u03b5) ), \u03b5=1e-15, in BITS. If normalised: H_c=-\u03a3 px\u00b7log2(px+\u03b5), H_p=-\u03a3 py\u00b7log2(py+\u03b5), denom = 0.5(H_c+H_p) if mode=='average' else min(H_c,H_p); return MI/denom (0.0 if denom\u22640). This is symmetric-uncertainty / normalized MI \u2208[0,1].", "operations": "row marginal px, col marginal py, outer product px@py (C\u00d7P), elementwise pxy\u00b7log2(ratio), sum. Marginal entropies for normalization denominator.", "args_effects": "normalised: False\u2192raw MI in bits; True\u2192divide by average (mode='average') or min (any other mode) of the two marginal entropies. mode: 'average'\u2192symmetric uncertainty denom; else\u2192min(H_c,H_p) (yields the redundancy/IQR-style normalization, \u2265 average form).", "normalizations": "log base 2 (bits). Normalization denominators are the marginal Shannon entropies. Assumes pxy pre-normalized to sum 1 by caller.", "edge_cases": "\u03b5=1e-15 added inside both logs (asymmetric with the numerator's pxy). 0\u00b7log(\u00b7)=0 for zero cells. denom>0 guard returns 0.0 for degenerate (single-clone or single-phenotype) tables.", "bugs": "Correct math but the \u03b5 is added to BOTH pxy and px\u00b7py inside the log, giving a tiny systematic bias at the \u03b5 scale for exact-zero cells (negligible). No guard that pxy is actually normalized \u2014 relies entirely on caller doing pxy/=pxy.sum()."}, {"name": "clonotypic_entropy_base (tcri/metrics/_metrics.py:183-225)", "formula": "For phenotype \u03c6 at covariate m: get joint jd (rows=clones, cols=phenotypes). vec = jd[\u03c6] over clones; clip to \u03b5=1e-15; normalize v = vec/\u03a3vec \u2192 distribution over clones for that phenotype. H = Shannon entropy_base(v). If normalised and len>1: H /= log_base(n_clones) = log(n_clones)/log(base). Returns H \u2208 [0,1] (normalized).", "operations": "posterior=True\u2192one joint_distribution_posterior draw; posterior=False\u2192joint_distribution(n_samples=0). Extract column, clip, L1-normalize, scipy.stats.entropy(base). Divide by log_base(len(vec)).", "args_effects": "base\u2192log base for entropy AND the max-entropy divisor. normalised\u2192toggles /log_base(n_clones). temperature\u2192passed to joint builder. clones\u2192restrict/reindex clone set (changes n_clones=len(vec), the normalizer). weighted\u2192if True, jd[\u03c6] holds clone-size-weighted mass so v becomes size-weighted P(c|\u03c6); if False, equal per-clone weighting. posterior\u2192posterior draw vs prior point estimate. combine_with_logits\u2192forwarded to posterior builder.", "normalizations": "Max-entropy normalization by log_base(number of clones). Column L1-normalized to a simplex over clones first.", "edge_cases": "Returns 0.0 if jd is None/empty or \u03c6 not a column. len(vec)==1 \u2192 skip normalization (H already 0). All-zero column (absent clones) \u2192 clipped to \u03b5 \u2192 contributes ~uniform \u2192 slightly inflates entropy.", "bugs": "Single Dirichlet draw when posterior=True \u2192 high per-call variance (delta_clonotypic_entropy loops it n_samples times as its Monte-Carlo). Semantic subtlety: normalizing jd[\u03c6] across clones gives P(c|\u03c6) only under a uniform (unweighted) or size (weighted) clone prior \u2014 it is not P(c|\u03c6) from the true joint unless weighted matches clone sizes."}, {"name": "clonotypic_entropy (tcri/metrics/_metrics.py:227-322)", "formula": "For each phenotype \u03c6 and each of n_samples posterior draws: vec = jd_draw[\u03c6] over clones; clip \u03b5=1e-15; v=vec/\u03a3; H = entropy(v, base=2); if n_clones>1: H /= log2(n_clones). samples[i,\u03c6]=H. point_estimate\u2192 Series = nanmean over draws per phenotype; else return (n_samples \u00d7 n_phenotypes) matrix. Always normalized, always base-2 (bits, \u2208[0,1]).", "operations": "Loop i=0..n_samples-1: fresh joint_distribution_posterior draw; per phenotype column clip/normalize/entropy/normalize-by-log2(n_clones); np.nanmean(axis=0).", "args_effects": "covariate\u2192condition. point_estimate: True\u2192posterior-mean per phenotype (Series indexed by phenotype); False\u2192full draw matrix. n_samples: number of Dirichlet draws averaged (must be \u22651). temperature\u2192per-cell softmax temp inside the joint. combine_with_logits\u2192full posterior vs sampled-prior-only. _clones\u2192restrict clone set (changes n_clones normalizer).", "normalizations": "Fixed base=2; max-entropy divisor log2(n_clones). Column normalized to simplex over clones.", "edge_cases": "ValueError if n_samples<1. jd None/empty \u2192 row of NaN. \u03c6 absent \u2192 NaN entry. n_clones==1 \u2192 no normalization. nanmean tolerates NaN draws.", "bugs": "No `normalised` toggle (unlike _base) \u2014 always normalized. Base hard-coded to 2 (ignores any base). Re-samples the full p_ct Dirichlet every iteration \u2192 O(n_samples) redundant sampling. Inherits joint_distribution_posterior's additive/no-gate combination inconsistency with the trained model."}, {"name": "phenotypic_entropy (tcri/metrics/_metrics.py:470-573)", "formula": "For each clone c and each of n_samples draws: p = jd_draw.loc[c] over phenotypes (row = P(\u03c6|c)); clip \u03b5=1e-15; p/=\u03a3; H = entropy(p, base=2)/norm where norm=log2(n_phenotypes) (1.0 if n_phen\u22641). samples[i,c]=H. point_estimate\u2192Series=nanmean over draws per clone; else (n_samples \u00d7 n_clones). Bits, \u2208[0,1].", "operations": "Build clones_list=unique clones at covariate (from obs). Loop draws of joint_distribution_posterior(clones=clones_list); per clone row clip/normalize/entropy/\u00f7log2(n_phen). np.nanmean(axis=0).", "args_effects": "covariate\u2192condition + defines clones_list. point_estimate\u2192mean Series (indexed by clone) vs full matrix. n_samples\u2192draws averaged (\u22651). temperature\u2192per-cell softmax temp. combine_with_logits\u2192posterior vs sampled-prior-only.", "normalizations": "Fixed base=2; divisor log2(number of phenotypes). Row normalized to simplex over phenotypes.", "edge_cases": "ValueError if n_samples<1. Empty clones_list\u2192empty Series / (n_samples,0). jd None/empty\u2192NaN row. Clone not in jd.index\u2192NaN. n_phen\u22641\u2192norm=1.0.", "bugs": "An all-zero clone row (e.g. clone present in obs but zero posterior mass) is still in jd.index (reindex fills 0) \u2192 after \u03b5-clip it normalizes to UNIFORM \u2192 phenotypic entropy 1.0 instead of NaN (the `cl not in jd.index` guard never triggers). Same O(n_samples) redundant Dirichlet sampling and additive/no-gate inconsistency as clonotypic_entropy."}, {"name": "mutual_information (tcri/metrics/_metrics.py:633-760)", "formula": "Get joint df (rows=clones, cols=phenotypes). pxy = df.values; pxy /= \u03a3pxy (normalize whole table \u2192 P(c,\u03c6)). MI = _mi_from_joint(pxy, normalised, normalise_mode) = \u03a3 pxy\u00b7log2(pxy/(px\u00b7py)) optionally \u00f7 marginal-entropy denom. n_samples=0\u2192single value; n_samples>0\u2192array of per-draw MI. I(c;\u03c6|m).", "operations": "_get_df: posterior=True\u2192joint_distribution_posterior (unweighted, so rows are P(\u03c6|c)); posterior=False\u2192raise NotImplementedError. Normalize table, call _mi_from_joint. Loop for multiple draws; reports mean\u00b1sd and 2.5/97.5 percentiles.", "args_effects": "covariate\u2192condition. temperature\u2192forwarded per-cell softmax temp. n_samples: 0=point estimate, >0=array of draws. clones\u2192restrict clone set (rows). normalised\u2192raw vs [0,1] MI. normalise_mode\u2192'average'(0.5(H_c+H_p)) vs anything-else(min). posterior\u2192True only (False raises NotImplementedError \u2014 prior path not wired). combine_with_logits\u2192posterior sampled-prior+logits vs sampled-prior-only. verbose/graph\u2192reporting only.", "normalizations": "Whole-table L1 normalization (pxy/=\u03a3) BEFORE MI \u2192 because unweighted jd rows each sum to 1, dividing by C imposes a UNIFORM clone prior P(c)=1/C. MI in bits; optional symmetric-uncertainty or min normalization.", "edge_cases": "NotImplementedError for posterior=False. Relies on _mi_from_joint \u03b5 guards. Single Dirichlet draw per sample.", "bugs": "Never passes weighted\u2192 MI ALWAYS assumes equiprobable clones (a clone of 1 cell counts as much as a clone of 10\u2074), which is a strong, undocumented modeling choice. posterior=False is dead (raises). mi_compare() calls this per patient with clones=\u2026 and n_samples; inherits the uniform-clone assumption. Same additive/no-gate posterior inconsistency."}, {"name": "flux (tcri/metrics/_metrics.py:851-1008)", "formula": "Per clone c, distance between its phenotype distribution at 'from_this' (p=jd_from.loc[c]) and 'to_that' (q=jd_to.loc[c]). l1: d=\u03a3_\u03c6|p_\u03c6\u2212q_\u03c6| (\u2208[0,2]). dkl: d=\u03a3 p\u00b7ln(p/q) with clip \u03b5=1e-15 and renorm (NATURAL log \u2192 nats, asymmetric KL(from\u2016to)). callable: d=metric(p,q). n_samples=0\u2192Series over common clones; n_samples>0\u2192(n_samples \u00d7 n_common) redrawn each iteration.", "operations": "seed\u2192np.random.seed. Resolve clones (None\u2192all, str\u2192[str]). get = joint_distribution_posterior if posterior else joint_distribution. Build jd_from, jd_to; common=index intersection. Dispatch on distance_metric (callable / 'l1' / 'dkl'). n_samples>0: loop fresh draws, recompute distances over the FIRST-draw `common`.", "args_effects": "from_this/to_that\u2192the two covariates compared. clones\u2192restrict. temperature\u2192forwarded. distance_metric\u2192'l1' (abs-diff sum), 'dkl' (natural-log KL from\u2016to), or any callable f(p,q). n_samples: 0=point Series, >0=posterior draw matrix. weighted\u2192forwarded to joint builder (affects whether rows are P(\u03c6|c) or size-weighted). posterior\u2192posterior_joint vs prior joint. combine_with_logits\u2192forwarded only when posterior (else passed as None). graph\u2192ASCII histogram. seed\u2192np.random.seed only.", "normalizations": "l1 operates on jd rows (already per-clone simplex when unweighted). dkl re-clips/re-normalizes p,q internally. No max normalization; l1\u2208[0,2], dkl\u2208[0,\u221e).", "edge_cases": "ValueError if jd_from or jd_to empty (no clone overlap). common = intersection of the two indices. dkl clips to \u03b5=1e-15.", "bugs": "(1) posterior=False is BROKEN: it still passes silent=True and combine_with_logits=None to get(...), but joint_distribution(adata,covariate_label,temperature,n_samples,clones,weighted) accepts NEITHER kwarg \u2192 TypeError. (2) seed=np.random.seed only seeds NumPy, but posterior draws use torch Dirichlet(...).sample() which is unaffected \u2192 'seed' does NOT make posterior flux reproducible. (3) dkl uses natural log while all entropies/MI use log2 \u2014 mixed units across the library. (4) flux_table indexes the returned matrix by clones_g.index(cl) but flux columns are ordered by `common` (intersection), so if common\u2260clones_g (order or membership) the per-clone column lookup misaligns."}, {"name": "get_cell_phenotype_probs (tcri/model/_model.py:907-972)", "formula": "Per cell: z=encoder(x,b).loc; \u2113=classifier(z) (= mlp(z)/classifier_temperature); prior_log=log(p_ct[ct(cell)]+\u03b5), \u03b5=1e-8, p_ct=get_p_ct() (posterior-MEAN). If use_gate: local = gate_prob\u00b7\u2113 + (1\u2212gate_prob)\u00b7prior_log (convex combo in logit space). Else: local = \u2113 + prior_log (additive Bayesian product). probs = softmax(local, dim=-1). Returns (n_cells \u00d7 P).", "operations": "get_p_ct\u2192(ct_count\u00d7P); map cells\u2192ct via module.ct_array. Batch loop: encoder mean, classifier logits (already \u00f7classifier_temperature), log prior from mean p_ct, gate or additive combine, softmax, concat.", "args_effects": "adata\u2192data (default training adata). batch_size\u2192loader chunking (numerics unaffected). eps\u2192log-prior floor. gate_prob (module attr): if set, convex logit mix; if None, pure additive. classifier_temperature (module attr) divides logits inside classifier.forward.", "normalizations": "Final softmax \u2192 per-cell simplex over P phenotypes. No temperature at this stage other than the classifier's built-in classifier_temperature.", "edge_cases": "eps=1e-8 guards log(0). Uses posterior-MEAN p_ct (get_p_ct), not a Dirichlet draw. current_idx bookkeeping assumes loader yields cells in ct_array order (same order as setup).", "bugs": "This gate-aware, mean-prior combination is the 'ground truth' cell probability, but the metric layer's joint_distribution_posterior uses a Dirichlet DRAW of the prior AND never applies the gate \u2192 metrics can disagree with the model's own cell probabilities. register_model additionally stores softmax(logits+prior_log) with the mean prior but ALWAYS additive (ignoring gate), a third variant."}, {"name": "get_p_ct (tcri/model/_model.py:510-525)", "formula": "Read raw positive guide param q_p_ct_raw (ct_count\u00d7P). Non-finite\u21921/n_phen. If guide_temperature\u22601: sharp = q_p_ct_raw^(1/guide_temperature); p_ct = sharp/\u03a3_\u03c6 sharp. Else p_ct = q_p_ct_raw/\u03a3_\u03c6 q_p_ct_raw. Returns row-stochastic p_ct = the guide Dirichlet's MEAN direction (concentration normalized).", "operations": "param_store['q_p_ct_raw']; finite-fix; power-temper by 1/guide_temperature; row L1-normalize.", "args_effects": "guide_temperature: <1 sharpens each ct's phenotype distribution, >1 flattens, =1 plain normalization. n_phen inferred from shape.", "normalizations": "Row L1 normalization to a simplex over phenotypes. The tempering matches the guide's own q_p_ct_sharp computation.", "edge_cases": "Non-finite entries elementwise replaced by 1/n_phen then renormalized. No explicit clamp here (the guide itself clamps to 1e-8/1e-3).", "bugs": "Returns the normalized concentration DIRECTION (Dirichlet mean = conc/\u03a3conc), which equals the guide mean only because conc=local_scale\u00b7sharp is proportional to sharp (clamps aside) \u2014 the clamp min=1e-3 applied in the guide is NOT reproduced here, so extreme rows can differ slightly from the actual guide-sample mean. Silent 1/n_phen fill can mask a diverged parameter."}, {"name": "guide q_p_ct_raw / q_p_c_raw (tcri/model/_model.py:442-491)", "formula": "Amortized-free variational Dirichlet guide. ct level: init = clone_phen_prior[ct_to_c]\u00b7guide_init_scale + 1e-3; q_p_ct_raw = positive pyro.param(init). sharp = clamp(q_p_ct_raw^(1/guide_temperature), min=1e-8); sharp/=\u03a3; conc = clamp(local_scale\u00b7sharp, min=1e-3); p_ct ~ Dirichlet(conc). Clone level: init = clone_phen_prior\u00b7guide_init_scale+1e-3; q_p_c_raw positive param; sharp likewise; conc = clamp(global_scale\u00b7sharp, min=1e-3); p_c ~ Dirichlet(conc).", "operations": "Register/retrieve param; non-finite\u2192init fill; power-temper^(1/guide_temperature); clamp\u22651e-8; row-normalize; scale by local_scale (ct) or global_scale (c); clamp\u22651e-3; sample Dirichlet. z guide: latent ~ Normal(z_loc, clamp(z_scale,1e-3,10)) scaled by kl_weight.", "args_effects": "guide_init_scale\u2192initial concentration magnitude (confidence of prior init). guide_temperature\u2192sharpen(<1)/flatten(>1) the guide mean direction. local_scale\u2192ct-level Dirichlet concentration (posterior tightness of p_ct); global_scale\u2192clone-level concentration (tightness of p_c). Larger scale \u2192 lower-variance draws around the mean.", "normalizations": "Row L1-normalize sharp before scaling; clamps (1e-8 pre-norm, 1e-3 on concentration) keep Dirichlet valid.", "edge_cases": "Non-finite params replaced by scaled prior init (bad_ct/bad_c masks). Double clamp guards against zero/negative concentration.", "bugs": "Two-stage clamp+normalize means the realized Dirichlet mean \u2260 exactly the normalized sharp when any entry hits the 1e-3 concentration floor; get_p_ct ignores that floor, so the reported p_ct can diverge from the guide's true expectation for near-degenerate rows."}, {"name": "model prior Dirichlet(local_scale\u00b7p_ct) & clone level (tcri/model/_model.py:382-399)", "formula": "Hierarchical prior. Clone plate: p_c ~ MixtureDirichlet(uniform weights 1/B, concentration = mixture_concentration = K archetype centroids from KMeans on empirical c2p_mat, each row clamped\u22651e-3). ct plate: base_p = p_c[ct_to_c] + \u03b5 (\u03b5=1e-6); conc_ct = clamp(local_scale\u00b7base_p, min=1e-3); p_ct ~ Dirichlet(conc_ct). So E[p_ct | p_c] = p_c[clone]; local_scale sets concentration (variance \u221d 1/(1+\u03a3conc)).", "operations": "Uniform mixture weights over B archetypes; MixtureDirichlet.sample picks a component then its Dirichlet; map p_c to cts via ct_to_c; scale by local_scale; clamp; Dirichlet sample p_ct.", "args_effects": "local_scale\u2192tightness of p_ct around its clone mean p_c (higher=tighter). global_scale (guide side)\u2192clone-level tightness. mixture_concentration (=archetype centers, K=KMeans clusters)\u2192the shape family clones are drawn toward; because centers are row-normalized to sum 1, each Dirichlet concentration vector sums to \u22481 \u2192 LOW total concentration \u2192 near-one-hot, committed archetype draws. prior_temperature (applied earlier to clone_phen_prior)\u2192sharpens the empirical prior fed to guide init.", "normalizations": "base_p normalized implicitly via Dirichlet (mean = conc/\u03a3conc). Centers row-normalized in build_archetypes (clip\u22651e-8 then /rowsum).", "edge_cases": "\u03b5=1e-6 on base_p; clamp min=1e-3 on concentrations; MixtureDirichlet clamps concentration\u22651e-3. has_rsample=False (score-function/enumeration path).", "bugs": "mixture_concentration = normalized centers (\u03a3=1) gives Dirichlet concentration <1 per component \u2192 extremely sparse, high-variance clone draws (nearly one-hot); this is likely unintended for a 'prior mean' and makes p_c very peaked. Using normalized centroids as Dirichlet concentrations conflates a mean vector with a concentration vector (loses any notion of prior strength/count)."}, {"name": "clone_phen_prior (prepare_two_level_params + c2p_mat, tcri/model/_model.py:331-360, 745-751)", "formula": "Empirical clone\u00d7phenotype prior. c2p_mat[c,\u03c6] = count of cells with clone c and hard phenotype \u03c6; +1e-6; row-normalize \u2192 P\u0302(\u03c6|c). In prepare_two_level_params: prior_mat = (clone_phen_prior_mat + \u03b5)/rowsum, \u03b5=1e-6; if prior_temperature\u22601: prior_mat = prior_mat^(1/prior_temperature); renormalize rows. Stored as buffer clone_phen_prior. Archetypes: KMeans(K) on c2p_mat \u2192 centers clipped\u22651e-8, row-normalized \u2192 mixture_concentration.", "operations": "Count matrix from cat codes; Laplace-ish +1e-6; row L1-normalize; optional power-temper^(1/prior_temperature)+renorm; KMeans for archetype centroids.", "args_effects": "prior_temperature: <1 sharpens P\u0302(\u03c6|c) toward its mode, >1 flattens toward uniform, =1 leaves empirical estimate. K (KMeans clusters)\u2192number of archetype mixture components B. \u03b5=1e-6\u2192smoothing so no exact zeros. Used to initialize both the guide (init = prior\u00b7guide_init_scale+1e-3) and the mixture prior.", "normalizations": "Row L1 normalization to a per-clone phenotype simplex; re-normalized after tempering.", "edge_cases": "+1e-6 prevents zero rows/log issues. KMeans centers re-clipped\u22651e-8 and re-normalized.", "bugs": "Derived from HARD phenotype labels (argmax categories), so the 'prior' bakes in the very labels the model then refines \u2014 potential circularity/label leakage into the posterior. boost_phenotype_prior mutates c2p_mat, clone_phen_prior, and centers in place (multiply one column by boost_factor then renormalize), which can silently desync the external self.c2p_mat copy from the registered buffer if called inconsistently."}]}, "inference": "Both files are read. Here is a precise, concrete account of the PRIOR / POSTERIOR / MEAN / DRAW structure of `p_ct`.\n\nFiles:\n- `/Users/ceglian/Codebase/GitHub/tcri/tcri/model/_model.py`\n- `/Users/ceglian/Codebase/GitHub/tcri/tcri/preprocessing/_preprocessing.py`\n\nShapes/notation: `P` = #phenotypes, `c_count` = #clonotypes, `ct_count` = #(clonotype, covariate/timepoint) groups, `K` = #archetypes (the `K=10` KMeans param). `p_c` lives at clonotype level `(c_count, P)`; `p_ct` at clone\u00d7timepoint level `(ct_count, P)`.\n\n\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n(1) clone_phen_prior: c2p_mat, archetypes, mixture\n\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\nThere are two distinct prior objects, built in `TCRIModel.__init__` (_model.py L745\u2013751, L800):\n\n- `c2p_mat` (L745\u2013750): raw empirical phenotype counts per clonotype. For every cell, `c2p_mat[clone_code, phenotype_code] += 1`; then `+= 1e-6` and row-normalize. Result is the empirical clone\u2192phenotype distribution, `(c_count, P)`, each row on the simplex. Stored as `self.c2p_mat`.\n\n- `clone_phen_prior` (buffer): `prepare_two_level_params` (L346\u2013353) takes `c2p_mat`, re-adds eps, re-normalizes, optionally sharpens by `prior_temperature`, and registers it as buffer `clone_phen_prior`, `(c_count, P)`. This buffer is used ONLY to initialize the guide's variational params (not in the generative model at all).\n\n- `archetypes / centers` (L751, via `build_archetypes`, L45\u201351): KMeans with `K` clusters on `c2p_mat`'s rows; centroids clipped positive and row-normalized \u2192 `centers` `(K, P)`. Passed into the module as `mixture_concentration=torch.from_numpy(self.centers)` (L800).\n\nThe mixture is the actual generative PRIOR over `p_c`, built in `model()` (L382\u2013393):\n- `B = mixture_concentration.shape[0] = K` archetypes, uniform `mixture_weights = ones(K)/K`.\n- `expanded_conc` = archetype centers broadcast to `(c_count, K, P)`; `MixtureDirichlet(weights, conc)`.\n- `p_c ~ MixtureDirichlet` \u2014 each clonotype's prior over its phenotype distribution is a uniform mixture of `K` Dirichlets whose concentration vectors ARE the archetype centroids (note: centroids sum to 1, so these are very low-concentration/diffuse Dirichlets).\n\nThen the `p_ct` PRIOR (L396\u2013399), the hierarchical link:\n- `base_p = p_c[ct_to_c] + eps` (broadcast clonotype-level `p_c` down to each clone\u00d7timepoint group),\n- `conc_ct = clamp(local_scale * base_p, 1e-3)`,\n- `p_ct ~ Dirichlet(local_scale * p_c[ct_to_c])`.\nSo the prior on `p_ct` is a Dirichlet centered at the clonotype's `p_c`, with total concentration \u2248 `local_scale`.\n\n\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n(2) What the guide learns; what get_p_ct returns\n\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\nThe guide (L442\u2013491) introduces two `pyro.param` tensors, both positive-constrained, initialized from `clone_phen_prior * guide_init_scale + 1e-3`:\n- `q_p_c_raw` `(c_count, P)` (L449),\n- `q_p_ct_raw` `(ct_count, P)` (L475).\n\nFor `p_ct` the guide (L487\u2013491) does: sharpen by `guide_temperature` (identity at T=1), clamp, **row-normalize to the simplex** \u2192 `q_p_ct_sharp` (call it `m`, sums to 1), then `conc_ct_guide = clamp(local_scale * m, 1e-3)`, and samples `p_ct ~ Dirichlet(local_scale * m)`.\n\nKey fact: for `Dirichlet(\u03b1)` with `\u03b1 = local_scale * m` and `m` a simplex vector, `sum(\u03b1) = local_scale`, so the **mean is `\u03b1/sum(\u03b1) = m`** exactly. The normalization step throws away the magnitude of `q_p_ct_raw`; only its direction (the normalized vector `m`) sets the variational posterior mean.\n\n`get_p_ct()` (L510\u2013525): reads `q_p_ct_raw`, repairs non-finite rows, and returns `q_p_ct_raw / q_p_ct_raw.sum(dim=1, keepdim=True)` (with the same optional `guide_temperature` sharpening). That is:\n- It returns the **row-normalized raw variational parameter**, which is numerically **exactly the mean of the guide's Dirichlet posterior `q(p_ct)`** (`E_q[p_ct] = m`, modulo the negligible 1e-3 clamp).\n- It is NOT a sample, and NOT the mode. The Dirichlet mode `(\u03b1\u22121)/(\u03a3\u03b1 \u2212 P)` is different and, here, typically ill-defined: `\u03b1_k = local_scale\u00b7m_k` is usually < 1 (since `local_scale`\u22483\u20135 and `m_k`<1), pushing the mode onto the simplex boundary.\n\nSo `get_p_ct()` = posterior MEAN `E_q[p_ct]`, `(ct_count, P)`. This is what gets stored as `adata.uns[\"tcri_p_ct\"]` in `register_model` (_preprocessing.py L184).\n\n\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n(3) What joint_distribution_posterior samples; combine_with_logits\n\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n`joint_distribution_posterior` (_preprocessing.py L264\u2013329):\n- `p_ct_mean = adata.uns[\"tcri_p_ct\"]` \u2014 i.e. the posterior mean `m` from (2), `(ct_count, P)`.\n- `local_scale = adata.uns[\"tcri_local_scale\"]`.\n- `p_ct_sample = Dirichlet(local_scale * p_ct_mean + 1e-8).sample()` (L305) \u2014 it **reconstructs the guide's Dirichlet `q(p_ct)` and draws exactly ONE Monte-Carlo sample**. This is a single stochastic realization of the posterior; its expectation over repeated calls equals `p_ct_mean`, but any one call is noisy.\n\n`combine_with_logits=True` (L308\u2013315): mixes the classifier into the sampled prior additively:\n- `logits = adata.obsm[\"X_tcri_logits\"][cells]` (classifier logits from `z_loc`),\n- `log_prior = log(p_ct_sample[ct_idx] + 1e-8)`,\n- `probs_cell = softmax((logits + log_prior)/temperature)`.\nThis is the Bayesian product rule `cls_logits + log(prior)` (same additive form as the non-gated branch of `get_cell_phenotype_probs`, _model.py L964), but with the DRAW substituted for the mean, and notably it does NOT apply the `gate_prob` weighting. If `combine_with_logits=False`, `probs_cell = p_ct_sample[ct]` directly (classifier ignored). Then per-clonotype groupby-sum and optional normalization.\n\nNet: this function returns ONE posterior draw of the joint clone\u00d7phenotype table \u2014 stochastic across calls. A proper posterior would require many calls.\n\n\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n(4) What non-posterior joint_distribution computes (n_samples=0)\n\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n`joint_distribution` (_preprocessing.py L364\u2013441), `n_samples=0` branch:\n- `p_ct = adata.uns[\"tcri_p_ct\"]` \u2014 again the posterior MEAN `m`.\n- select rows for the covariate: `p_ct_for_cov = p_ct[ct_to_cov == cov]`.\n- `p_ct_for_cov = F.softmax(log(p_ct_for_cov + eps)/temperature)` (L397). At `temperature=1` this is the identity (up to eps + renormalization).\n- Return the resulting `(n_ct_in_cov, P)` table indexed by clonotype (optionally cell-count weighted).\n\nSo the point estimate is the **temperature-scaled variational posterior MEAN of `p_ct`** \u2014 it is NOT the prior (`clone_phen_prior`/archetypes are never touched here) and NOT the mode. At `T=1` it is essentially `adata.uns[\"tcri_p_ct\"]` itself. Note it uses ONLY `p_ct`; unlike `joint_distribution_posterior`, it never touches `X_tcri_logits`/the classifier. (Its `n_samples>0` branch draws from `Dirichlet(local_scale * p_ct_for_cov)`, i.e. the same `q(p_ct)`, so its draws' expectation is that same mean.)\n\nOne cross-function inconsistency worth flagging for the design decision: `joint_distribution` folds `temperature` into the Dirichlet concentration base (`softmax(log p_ct/T)` then `local_scale\u00b7\u00b7`), while `joint_distribution_posterior` leaves the concentration as `local_scale\u00b7p_ct_mean` and applies `temperature` only in the final `softmax(logits+log_prior)`. And only `joint_distribution_posterior` mixes in the classifier. So the two \"distribution\" functions do not agree even at `T=1`.\n\n\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n(5) What n_samples=0 SHOULD compute for a deterministic point estimate\n\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\nIt should compute the **posterior MEAN `E_q[p_ct]` under the learned variational Dirichlet** \u2014 not the prior, not the MAP, and not an empirical average of draws.\n\nRationale, concretely:\n- The learned posterior is `q(p_ct) = Dirichlet(\u03b1)`, `\u03b1 = local_scale \u00b7 normalized(q_p_ct_raw)`. Its mean has a **closed form**: `E_q[p_ct] = \u03b1/\u03a3\u03b1 = normalized(q_p_ct_raw)`. This is exactly `get_p_ct()`'s return and exactly `adata.uns[\"tcri_p_ct\"]`. No sampling is needed or desirable \u2014 drawing many Dirichlet samples and averaging (the \"mean of many draws\") converges to this same vector but only up to Monte-Carlo noise, so it is a strictly worse estimator of a quantity available in closed form.\n- The prior (`clone_phen_prior` buffer, or the archetype `mixture_concentration`) is the wrong object \u2014 it is only the guide's initialization/generative anchor, not what training learned.\n- The MAP/mode `(\u03b1\u22121)/(\u03a3\u03b1\u2212P)` is wrong here: with `\u03b1_k = local_scale\u00b7m_k` frequently < 1 (small `local_scale`, `m_k`<1), the mode sits on the simplex boundary / is undefined, so it is unstable and non-representative.\n\nConcrete tensors/keys for the deterministic estimate:\n- Source of truth: guide param `q_p_ct_raw` `(ct_count, P)` \u2192 `E_q[p_ct] = q_p_ct_raw / q_p_ct_raw.sum(dim=1, keepdim=True)` = `module.get_p_ct()` = `adata.uns[\"tcri_p_ct\"]`.\n- Index maps for selection/aggregation: `adata.uns[\"tcri_ct_to_cov\"]`, `adata.uns[\"tcri_ct_to_c\"]`, `adata.uns[\"tcri_ct_array_for_cells\"]`, `adata.uns[\"tcri_cov_array_for_cells\"]`.\n- `adata.uns[\"tcri_local_scale\"]` is needed ONLY to reconstruct `q(p_ct)` for DRAWS (n_samples>0); it is irrelevant to the mean, since it cancels in `\u03b1/\u03a3\u03b1`.\n\nIn short: `n_samples=0` = closed-form posterior mean = `adata.uns[\"tcri_p_ct\"]` (temperature-adjusted if desired); `n_samples>0` = i.i.d. draws from `Dirichlet(local_scale \u00b7 p_ct_mean)`; and `joint_distribution_posterior` is that same draw machinery but with exactly one sample and an extra additive classifier-logit mixing step."}, "drafts": {"api": "I have everything I need. Producing the definitive API surface document.\n\n# TCRI \u2014 Final Functional API Surface\n\n_The definitive public + private surface for the refactored, grafiti-mirrored `tcri` package (Door A: standalone scverse package). Every module, every symbol, its exact final signature, and its one-line responsibility. Layout mirrors `grafiti`: one file per topic, private cross-cutting helper modules, explicit `__all__` re-export (never `import *`)._\n\n> **Directory-name note.** The `tl` view lives in `tcri/tools/` (grafiti-mirrored `tools\u2194tl`, superseding the earlier working name `metrics/`). All five view directories map to short handles exactly as grafiti does: `model\u2192ml`, `tools\u2192tl`, `preprocessing\u2192pp`, `plotting\u2192pl`, `diagnostics\u2192diag`, plus `utils\u2192ut` (tcri-specific, session I/O).\n>\n> **Convention baked into every signature below.** `n_samples=0` \u2192 deterministic **closed-form posterior mean** point estimate (no draw, reproducible); `n_samples=N>0` \u2192 `N` i.i.d. posterior Dirichlet draws. There is **no** `point_estimate=` argument anywhere. Keyword-only (`*`) marks the boundary after the primary positional data argument.\n\n---\n\n## 0. Package tree\n\n```\ntcri/\n __init__.py # explicit re-export + sys.modules aliases (tl/pp/pl/ml/ut/diag); top-level joint_distribution; NO import *\n _keys.py # single source of every uns/obsm/obs key string (constants only)\n _console.py # leveled, silenceable logging over scanpy.logging (no raw ANSI, no _ascii_hist)\n _stats.py # stars, AUROC+permutation, bootstrap, + posterior-comparison stats (MWU, P(>0), HDI, summarize)\n _distance.py # kl_divergence, l1_distance, phenotype_distance dispatcher\n model/ # ml\n _model.py # TCRIModel (setup_anndata, train, get_latent_representation, predict, get_p_ct, to_anndata)\n _module.py # TCRIModule (pyro model/guide, get_latent, get_p_ct)\n _priors.py # MixtureDirichlet, VampPrior\n _classifier.py # PhenotypeClassifier\n _training.py # UnifiedTrainingPlan, build_archetypes\n preprocessing/ # pp\n _register.py # registration writers behind TCRIModel.to_anndata (all private)\n _clones.py # group_singletons, clone_size\n tools/ # tl (mirrors pl by filename)\n _joint.py # joint_distribution (THE ENGINE; unifies the two current joint fns)\n _entropy.py # clonotypic_entropy, phenotypic_entropy\n _mutual_information.py # mutual_information (+ private _mi_from_joint)\n _flux.py # phenotypic_flux\n plotting/ # pl (twins mirror tl by filename)\n _base.py # _metric_boxplot, _finish\n _colors.py # tcri_colors, resolve_palette\n _entropy.py # clonotypic_entropy, phenotypic_entropy\n _mutual_information.py # mutual_information\n _flux.py # phenotypic_flux (sankey)\n _sankey.py # SankeyNode, _phenotype_mass_per_clone\n diagnostics/ # diag (NEW \u2014 PPCs + model validation; returns DataFrames, plus two relocated model plots)\n _ppc.py # joint_distribution_ppc, phenotype_calibration, reconstruction_ppc, permutation_null\n _training.py # loss, archetypes\n utils/ # ut\n _session.py # save_tcri_session, load_tcri_session (+ private load/serialization helpers)\n```\n\n`examples/` (outside the package): `top_clone_umap`, `clone_size_umap`, `phenotype_probabilities`, `compare_phenotypes`, rewritten notebooks. `docs/`: the model PGM diagram (`build_nested_tcri_pgm`).\n\n---\n\n## 1. Top-level `__init__.py` and the `__all__` story\n\n**Pattern (copied from grafiti):** `__all__` is declared at **both** levels. Every impl module (`_joint.py`, `_entropy.py`, \u2026) declares its own `__all__`. Every view `__init__.py` imports each symbol by name and re-declares an aggregate `__all__` grouped by view. The root imports the six view packages, aliases them into `sys.modules`, and additionally re-exports `joint_distribution` for prominence. **No `import *` anywhere** \u2014 numpy/pandas/torch and every `_helper` stay unexported.\n\n```python\n# tcri/__init__.py\nfrom importlib.metadata import PackageNotFoundError, version as _version\ntry:\n __version__ = _version(\"tcri\")\nexcept PackageNotFoundError:\n __version__ = \"0.0.0+unknown\"\n\nfrom . import tools as tl\nfrom . import preprocessing as pp\nfrom . import plotting as pl\nfrom . import model as ml\nfrom . import diagnostics as diag\nfrom . import utils as ut\nfrom .tools import joint_distribution # top-level prominence: tcri.joint_distribution\n\nimport sys\nsys.modules.update({f\"{__name__}.{m}\": globals()[m]\n for m in (\"tl\", \"pp\", \"pl\", \"ml\", \"diag\", \"ut\")})\n\n__all__ = [\"tl\", \"pp\", \"pl\", \"ml\", \"diag\", \"ut\", \"joint_distribution\", \"__version__\"]\n```\n\n| View `__init__` | `__all__` (public surface it re-exports) |\n|---|---|\n| `tools/__init__.py` (`tl`) | `joint_distribution`, `clonotypic_entropy`, `phenotypic_entropy`, `mutual_information`, `phenotypic_flux` |\n| `preprocessing/__init__.py` (`pp`) | `group_singletons`, `clone_size` |\n| `plotting/__init__.py` (`pl`) | `clonotypic_entropy`, `phenotypic_entropy`, `mutual_information`, `phenotypic_flux`, `tcri_colors`, `resolve_palette` |\n| `model/__init__.py` (`ml`) | `TCRIModel` |\n| `diagnostics/__init__.py` (`diag`) | `joint_distribution_ppc`, `phenotype_calibration`, `reconstruction_ppc`, `permutation_null`, `loss`, `archetypes` |\n| `utils/__init__.py` (`ut`) | `save_tcri_session`, `load_tcri_session` |\n\nPrivate modules (`_keys`, `_console`, `_stats`, `_distance`) and every private module symbol (`_mi_from_joint`, `_metric_boxplot`, `_finish`, `SankeyNode`, `_phenotype_mass_per_clone`, all `_register.py` writers) are **absent** from every `__all__` \u2014 reachable only via their fully-qualified private path.\n\n---\n\n## 2. Shared private helper modules\n\n### `tcri/_keys.py` \u2014 canonical AnnData key registry (constants only, no functions)\n\nSingle source of every string. Import as `from tcri import _keys as K`; no key literal lives anywhere else.\n\n| Constant | Value | Slot | Meaning |\n|---|---|---|---|\n| `K.METADATA` | `\"tcri_metadata\"` | `uns` | dict: `covariate_col`, `clone_col`, `phenotype_col`, `batch_col` (the **single** metadata scheme \u2014 replaces the old dual `tcri_clone_key`/`tcri_phenotype_key`) |\n| `K.PHENOTYPE_CATEGORIES` | `\"tcri_phenotype_categories\"` | `uns` | ordered phenotype category list |\n| `K.CLONOTYPE_CATEGORIES` | `\"tcri_clonotype_categories\"` | `uns` | ordered clonotype category list |\n| `K.COVARIATE_CATEGORIES` | `\"tcri_covariate_categories\"` | `uns` | ordered covariate category list |\n| `K.P_CT` | `\"tcri_p_ct\"` | `uns` | `(ct_count, P)` learned posterior-mean `E_q[p_ct]` |\n| `K.CT_TO_COV` | `\"tcri_ct_to_cov\"` | `uns` | `(ct_count,)` ct\u2192covariate index map |\n| `K.CT_TO_C` | `\"tcri_ct_to_c\"` | `uns` | `(ct_count,)` ct\u2192clonotype index map |\n| `K.CT_ARRAY` | `\"tcri_ct_array_for_cells\"` | `uns` | `(n_obs,)` per-cell ct index |\n| `K.COV_ARRAY` | `\"tcri_cov_array_for_cells\"` | `uns` | `(n_obs,)` per-cell covariate index |\n| `K.LOCAL_SCALE` | `\"tcri_local_scale\"` | `uns` | scalar Dirichlet concentration (draws only) |\n| `K.X_LATENT` | `\"X_tcri\"` | `obsm` | `(n_obs, n_latent)` encoder posterior-mean latent |\n| `K.X_LOGITS` | `\"X_tcri_logits\"` | `obsm` | `(n_obs, P)` classifier logits |\n| `K.X_PROBABILITIES` | `\"X_tcri_probabilities\"` | `obsm` | `(n_obs, P)` per-cell phenotype probabilities |\n| `K.PHENOTYPE_LABEL` | `\"tcri_phenotype\"` | `obs` | per-cell argmax hard phenotype label |\n| `K.CLONE_SIZE` | `\"clone_size\"` | `obs` | per-cell clone cell-count |\n\n> The AnnDataManager is **no longer** stashed in `uns` (`tcri_manager` retired) \u2014 this deletes the `write_adata_safely` / `_pop_nonserializables` hack.\n\n### `tcri/_console.py` \u2014 leveled, silenceable logging (private)\n\nThin wrappers over `scanpy.logging` for ecosystem consistency; respects scanpy verbosity. Raw ANSI prints and `_ascii_hist` (and every `graph=`/ASCII-histogram code path) are **deleted**.\n\n| Signature | Responsibility |\n|---|---|\n| `info(msg, *, deep=None)` | Emit an info line (`scanpy.logging.info`); silenced by scanpy verbosity. |\n| `warning(msg)` | Emit a warning line (`scanpy.logging.warning`). |\n| `success(msg)` | Emit a hint/success line (`scanpy.logging.hint`). |\n| `done(msg=\"done\")` | Emit the terminal completion line. |\n\n### `tcri/_stats.py` \u2014 significance + posterior-comparison statistics (private)\n\n| Signature | Responsibility |\n|---|---|\n| `stars(p)` | Map a p-value to a significance-star string (`****`/`***`/`**`/`*`/`ns`). |\n| `auc_and_label_permutation(scores, labels, *, pos_label=None, n_perm=200_000, seed=42, max_exact=200_000)` | Observed ROC-AUC plus two-sided permutation p-value (exact combinations when `C(n,k) \u2264 max_exact`, else Monte-Carlo); returns `(auc, p, perm_stats, mode)`. |\n| `bootstrap_auc(scores, labels, *, pos_label=None, n_boot=5000, seed=42)` | Bootstrap 95% CI (2.5/97.5 quantiles) of ROC-AUC, resampling until both classes present. |\n| `mann_whitney(a, b, *, alternative=\"two-sided\")` | Mann\u2013Whitney U test between two sample vectors; returns `(U, p)`. |\n| `prob_greater_than_zero(samples)` | Bayesian `P(x > 0)` = fraction of posterior draws above zero. |\n| `hdi(samples, *, hdi_prob=0.94)` | Highest-density interval `(low, high)` of a posterior sample vector. |\n| `summarize(samples, *, hdi_prob=0.94)` | Collapse a draw vector to the canonical summary dict: `mean`, `sd`, `hdi_low`, `hdi_high`, `p_gt` (the shared \"n_samples>0 \u2192 summary columns\" reducer). |\n\n### `tcri/_distance.py` \u2014 phenotype-distribution distances (private)\n\nDedupes the old module-level `dkl` and `flux.dkl_func`; unifies units to **bits (log\u2082)** across the library.\n\n| Signature | Responsibility |\n|---|---|\n| `kl_divergence(p, q, *, base=2.0, eps=1e-15)` | KL(p\u2016q) with clip+renormalize; base-2 by default (bits), matching entropy/MI. |\n| `l1_distance(p, q)` | L1 / total-variation-style distance `\u03a3|p\u1d62\u2212q\u1d62|` (\u2208 [0, 2]). |\n| `phenotype_distance(p, q, *, metric=\"l1\")` | Dispatcher: `\"l1\"` \u2192 `l1_distance`, `\"dkl\"` \u2192 `kl_divergence`, or a user callable `f(p, q)`. |\n\n---\n\n## 3. `tcri.ml` \u2014 model (`model/`)\n\n### `model/_model.py`\n\n**`class TCRIModel(BaseModelClass)`** \u2014 high-level scvi model: register \u2192 build \u2192 train \u2192 extract \u2192 write.\n\n| Method (signature) | Responsibility |\n|---|---|\n| `@classmethod setup_anndata(cls, adata, *, layer=None, clonotype_key=\"unique_clone_id\", phenotype_key=\"phenotype_col\", covariate_key=\"timepoint\", batch_key=\"patient\", **kwargs)` | **Registration only** \u2014 register clonotype/phenotype/covariate/batch/count fields with scvi and store the layer. No `obs` mutation, no `tcri_manager` stash in `uns`. |\n| `__init__(self, adata, *, n_latent=128, n_hidden=128, n_layers=3, classifier_n_layers=3, global_scale=5.0, local_scale=3.0, prior_temperature=1.0, guide_temperature=1.0, use_enumeration=False, patience=300, classifier_hidden=128, classifier_dropout=0.1, n_pseudo_obs=10, K=10, phenotype_weights=None, gate_prob=None, kl_weight_max=1.0, guide_init_scale=10.0, classifier_temperature=1.0, **kwargs)` | Build the empirical clone\u2192phenotype prior + KMeans archetypes + clonotype/covariate index maps + class weights, then construct and prime `TCRIModule`. |\n| `train(self, *, max_epochs=1000, batch_size=1000, lr=1e-3, reconstruction_loss_scale=1e-3, n_steps_kl_warmup=2000, **kwargs)` | 0.9/0.1 split, build `UnifiedTrainingPlan`, run `TrainRunner` with `elbo_validation` early stopping. |\n| `get_latent_representation(self, adata=None, *, indices=None, batch_size=None) -> np.ndarray` | Batched encode to the `(n_cells, n_latent)` posterior-mean latent. |\n| `predict(self, adata=None, *, batch_size=256, eps=1e-8) -> pd.DataFrame` | **(renamed from `get_cell_phenotype_probs`)** Per-cell phenotype-probability DataFrame (cells \u00d7 phenotypes) combining classifier logits with `log p_ct` (gate or additive), matching training \u2014 the scvi/CellAssign idiom. |\n| `get_p_ct(self) -> np.ndarray` | Return the learned `(ct_count, P)` closed-form posterior mean `E_q[p_ct]` = `normalize(q_p_ct_raw)`. |\n| `to_anndata(self, adata=None, *, latent_key=\"X_tcri\", predictions_key=\"X_tcri_probabilities\", label_key=\"tcri_phenotype\") -> AnnData` | **(replaces the heavy `register_model`)** Thin writer: write the canonical minimum \u2014 metadata + categories (from registry), `X_tcri` latent, `predict()` probs + argmax hard labels, and `p_ct` (+ ct index maps). Nothing else; no manager stash. |\n\n> Relocated off the model: `plot_archetypes` \u2192 `diag.archetypes`; `plot_loss` \u2192 `diag.loss`. `boost_phenotype_prior`, `use_gate` remain internal.\n\n### `model/_module.py`\n\n**`class TCRIModule(PyroBaseModuleClass)`** \u2014 Pyro CVAE with hierarchical clonotype\u2192(clonotype\u00d7covariate) Dirichlet priors and a phenotype classifier. *(internal)*\n\n| Member (signature) | Responsibility |\n|---|---|\n| `__init__(self, n_input, n_latent, P, n_batch, *, global_scale=10.0, local_scale=5.0, prior_temperature=1.0, guide_temperature=1.0, gate_prob=0.5, mixture_concentration=None, n_pseudo_obs=10, use_enumeration=False, classifier_hidden=128, classifier_dropout=0.1, classifier_n_layers=3, n_hidden=128, n_layers=3, class_weights=None, kl_weight_max=1.0, guide_init_scale=10.0, classifier_temperature=1.0)` | Construct encoder/decoder/classifier/VampPrior, `px_r`, and register empty two-level buffers + class weights (`mixture_concentration` required). |\n| `prepare_two_level_params(self, clone_phen_prior_mat, ct_to_c, ct_to_cov, ct_array_for_cells, cov_array_for_cells, *, eps=1e-6)` | Normalize/temperature the clone-phenotype prior and register all two-level index buffers. |\n| `model(self, x, batch_idx, log_library, ...)` | Generative model: sample `p_c` (MixtureDirichlet), `p_ct` (Dirichlet centered at `p_c`), latent `z` (VampPrior), ZINB gene obs. |\n| `guide(self, x, batch_idx, log_library, ...)` | Variational guide: learnable Dirichlet params `q(p_c)`, `q(p_ct)` and Normal `q(z)` from the encoder. |\n| `get_latent(self, tensor_dict) -> torch.Tensor` | Encode a batch to the posterior-mean latent `z_loc`. |\n| `get_p_ct(self) -> torch.Tensor` | Read `q_p_ct_raw` from the param store; return the row-normalized (guide-temperature-adjusted) `(ct_count, P)` posterior mean. |\n| `use_gate(self) -> bool` (property) | `True` when `gate_prob is not None` (convex-gate vs additive combination). |\n| `@staticmethod _get_fn_args_from_batch(tensor_dict) -> tuple` | Extract `(x, batch_idx, log_library)` from a scvi batch dict. |\n\n### `model/_priors.py`\n\n**`class MixtureDirichlet(dist.TorchDistribution)`** \u2014 mixture of Dirichlets over the phenotype simplex; the clonotype prior `p_c`. *(internal)*\n\n| Member (signature) | Responsibility |\n|---|---|\n| `__init__(self, mixture_weights, concentration, validate_args=None)` | Clamp concentrations (\u22651e-3), infer `batch_shape`/`B`/`K`, init the TorchDistribution. |\n| `sample(self, sample_shape=torch.Size())` | Sample a component per batch element, then draw from the selected Dirichlet. |\n| `log_prob(self, value)` | Log-sum-exp of component Dirichlet log-probs weighted by (log) mixture weights. |\n| `score_parts(self, value)` | Return `(log_prob, 0, 0)` so Pyro treats it as reparam-free. |\n| `__call__(self, *args, **kwargs)` | Alias for `sample()`. |\n\n**`class VampPrior(torch.nn.Module)`** \u2014 VampPrior over latent `z`: uniform mixture of encoder posteriors at learnable pseudo-inputs. *(internal)*\n\n| Member (signature) | Responsibility |\n|---|---|\n| `__init__(self, pseudo_inputs, encoder)` | Register learnable pseudo-inputs as a `Parameter`; hold the shared encoder. |\n| `get_mixture(self)` | Encode pseudo-inputs \u2192 uniform `MixtureSameFamily` of `Independent` Normals as `p(z)`. |\n| `log_prob(self, z)` | Log density of `z` under the mixture. |\n| `sample(self, sample_shape=torch.Size())` | Draw from the mixture. |\n\n### `model/_classifier.py`\n\n**`class PhenotypeClassifier(nn.Module)`** \u2014 temperature-scaled MLP head `z \u2192 phenotype logits`. *(internal)*\n\n| Member (signature) | Responsibility |\n|---|---|\n| `__init__(self, n_latent, classifier_hidden, P, *, num_layers=3, dropout_rate=0.1, temperature=1.0)` | Build stacked Linear/ReLU/Dropout MLP; store softmax temperature. |\n| `forward(self, x)` | Return MLP logits divided by `temperature`. |\n\n### `model/_training.py`\n\n**`class UnifiedTrainingPlan(PyroTrainingPlan)`** \u2014 training plan with KL warmup, diagnostics logging, and `elbo_validation` early stopping. *(internal)*\n\n| Member (signature) | Responsibility |\n|---|---|\n| `__init__(self, module, *, n_steps_kl_warmup=1000, reconstruction_loss_scale=1e-2, num_particles=5, optimizer_config=None, class_weights=None, **kwargs)` | Choose `TraceEnum_ELBO` vs `Trace_ELBO` by `module.use_enumeration`; store optimizer/warmup config. |\n| `loss(self)` (property) | Expose the configured ELBO loss object. |\n| `configure_optimizers(self)` | Build an Adam optimizer over module parameters from `optimizer_config`. |\n| `training_step(self, batch, batch_idx)` | Apply KL warmup, run the Pyro ELBO step, log KL/entropy/confidence diagnostics. |\n| `validation_step(self, batch, batch_idx)` | Compute validation ELBO + prior-KL, log `elbo_validation` for early stopping. |\n\n**Module function:** `build_archetypes(c2p_mat, K=10) -> np.ndarray` \u2014 KMeans-cluster clone\u2192phenotype rows into `K` normalized archetype centroids (the Dirichlet-mixture prior concentration).\n\n---\n\n## 4. `tcri.pp` \u2014 preprocessing (`preprocessing/`)\n\nThe engine has moved **up to `tl`**. `pp` shrinks to clone-utility functions plus the private registration writers behind `TCRIModel.to_anndata`.\n\n### `preprocessing/_clones.py` \u2014 public\n\n| Signature | Responsibility |\n|---|---|\n| `group_singletons(adata, *, clonotype_key=\"trb\", groupby=\"patient\", target_col=\"trb_unique\", min_clone_size=10) -> AnnData` | Collapse clones smaller than `min_clone_size` (per `groupby`) into `\"Singleton_{group}\"` labels in `target_col`. Kept separate from `setup_anndata` (which stays registration-only). |\n| `clone_size(adata, *, key_added=\"clone_size\", return_counts=False)` | Per-clone cell counts from the registered clone key, written per cell into `obs[key_added]`. |\n\n**Private inner:** `group_singletons.collapse_singleton(row)` \u2014 map a row to `\"Singleton_{group}\"` when its clone count `< min_clone_size`, else keep the candidate.\n\n### `preprocessing/_register.py` \u2014 all private (the `to_anndata` writer, one topic per file)\n\nNot re-exported; called only by `TCRIModel.to_anndata`. Folds in the old `register_phenotype_key` / `register_clonotype_key` / `_compute_logits_and_prior`.\n\n| Signature | Responsibility |\n|---|---|\n| `_write_metadata(adata, model)` | Write `uns[K.METADATA]` (single scheme) + the three category lists from the scvi registry. |\n| `_register_clonotype_key(adata, clonotype_key, *, order=None)` | Register the clonotype `obs` column and its ordered categories into `uns`. |\n| `_register_phenotype_key(adata, phenotype_key, *, order=None)` | Register the phenotype `obs` column and its ordered categories into `uns`. |\n| `_write_latent(adata, model, *, latent_key=\"X_tcri\", batch_size=256)` | Write the encoder posterior-mean latent to `obsm[latent_key]`. |\n| `_write_predictions(adata, model, *, predictions_key=\"X_tcri_probabilities\", label_key=\"tcri_phenotype\", batch_size=256)` | Write `predict()` probs to `obsm` + argmax hard labels to `obs`. |\n| `_write_p_ct(adata, model)` | Write `p_ct`, `ct_to_cov`, `ct_to_c`, per-cell ct/cov arrays, and `local_scale` to `uns`. |\n| `_compute_logits_and_prior(model, adata, *, batch_size=256, eps=1e-8) -> tuple[np.ndarray, np.ndarray]` | Run encoder+classifier over a loader to extract per-cell logits and the `log`-prior from `get_p_ct()`. |\n\n---\n\n## 5. `tcri.tl` \u2014 tools / metrics (`tools/`)\n\n### `tools/_joint.py` \u2014 THE ENGINE\n\n```python\njoint_distribution(\n adata, *,\n covariate=None,\n groupby=None,\n n_samples=0,\n posterior=True,\n clones=None,\n temperature=1.0,\n) -> pandas.DataFrame\n```\n\n**Responsibility.** Single Bayesian engine unifying today's `joint_distribution` + `joint_distribution_posterior`. Returns a tidy clone \u00d7 phenotype `DataFrame` (with a covariate/sample axis when relevant). `n_samples=0` \u2192 **closed-form posterior mean** `E_q[p_ct]` = `uns[K.P_CT]` (temperature-adjusted; deterministic). `n_samples>0` \u2192 i.i.d. draws from `Dirichlet(local_scale \u00b7 p_ct_mean)`. `posterior=True` folds in the per-cell classifier logits \u2014 `softmax(logits + log p_ct)` \u2014 for the full per-cell posterior joint; `posterior=False` returns the ct-level `p_ct` table directly (no logits). Provenance recorded in `df.attrs[\"params\"]`.\n\n- Weighting is **fixed and explicit** (repairs today's `ct`-vs-`clone` index mislookup): per-clone rows are cell-count-weighted consistently, one normalization convention.\n- Both `posterior` branches draw from the **same** learned variational posterior `q(p_ct)`; the flag toggles classifier mixing, **not** prior-vs-posterior. The prior path is removed (see \u00a79).\n\n`__all__ = [\"joint_distribution\"]`\n\n### `tools/_entropy.py`\n\n```python\nclonotypic_entropy(\n adata_or_jd, *,\n covariate=None, groupby=None,\n n_samples=0, temperature=1.0,\n clones=None, normalised=True,\n) -> float | pandas.Series | pandas.DataFrame\n```\nPer-phenotype clonotypic entropy `H(clone | phenotype)` at a covariate. Accepts **either** `(adata + covariate)` (computes the joint internally via the engine, `posterior=True`) **or** a precomputed joint `DataFrame` (fast path). `normalised=True` divides by `log\u2082(n_clones)`. `base` is fixed to 2 (bits) \u2014 the old `base=` / `_base` special-case is gone.\n\n```python\nphenotypic_entropy(\n adata_or_jd, *,\n covariate=None, groupby=None,\n n_samples=0, temperature=1.0,\n clones=None, normalised=True,\n) -> float | pandas.Series | pandas.DataFrame\n```\nPer-clone phenotypic entropy `H(phenotype | clone)` at a covariate; `normalised=True` divides by `log\u2082(n_phenotypes)`. Fixes the all-zero-clone-row \u2192 spurious-uniform bug (absent clones yield `NaN`, not max entropy).\n\n**Return-shape rule (both, uniform):** no `groupby` & `n_samples=0` \u2192 scalar/`Series`; `groupby` set \u2192 tidy `DataFrame` (row per group [\u00d7 phenotype/clone]); `n_samples>0` \u2192 adds a sample axis / summary columns (`mean`, `sd`, `hdi_low`, `hdi_high`, `p_gt` via `_stats.summarize`).\n\n`__all__ = [\"clonotypic_entropy\", \"phenotypic_entropy\"]`\n\n### `tools/_mutual_information.py`\n\n```python\nmutual_information(\n adata_or_jd, *,\n covariate=None, groupby=None,\n n_samples=0, temperature=1.0,\n clones=None, normalised=True,\n normalise_mode=\"average\",\n) -> float | pandas.Series | pandas.DataFrame\n```\nClone \u00d7 phenotype mutual information `I(clone; phenotype)` at a covariate (bits). `normalised=True` divides by the symmetric-uncertainty denominator (`normalise_mode=\"average\"` \u2192 `\u00bd(H_c+H_p)`, else `min(H_c,H_p)`). `n_samples=0` now returns a **true deterministic point estimate** (fixes today's \"one random draw at n_samples=0\" bug); `posterior=False`'s `NotImplementedError` path is gone.\n\n**Private helper (same file):**\n`_mi_from_joint(pxy, normalised, mode=\"average\") -> float` \u2014 MI kernel over an already-normalized `C\u00d7P` joint (`px=\u03a3_\u03c6`, `py=\u03a3_c`, `\u03a3 pxy\u00b7log\u2082(pxy/(px\u00b7py))`, optional marginal-entropy normalization). Not re-exported.\n\n`__all__ = [\"mutual_information\"]`\n\n### `tools/_flux.py`\n\n```python\nphenotypic_flux(\n adata, *,\n cov_from, cov_to,\n groupby=None,\n n_samples=0, temperature=1.0,\n clones=None,\n distance_metric=\"l1\",\n) -> pandas.Series | pandas.DataFrame\n```\n**(renamed from `flux`)** Per-clone distance between a clone's phenotype distribution at `cov_from` vs `cov_to`, over the clone intersection. `distance_metric` dispatches through `_distance.phenotype_distance` (`\"l1\"`, `\"dkl\"` in **bits**, or a callable). `n_samples=0` \u2192 deterministic `Series` over common clones; `n_samples>0` \u2192 draw matrix / summary. Fixes today's `flux` bugs: the `posterior=False` `TypeError`, torch-vs-numpy seeding non-reproducibility, mixed KL units, and the `flux_table` column-misalignment.\n\n**Return-shape rule:** no `groupby` & `n_samples=0` \u2192 `Series`; `groupby` set \u2192 tidy `DataFrame` (row per group \u00d7 clone); `n_samples>0` \u2192 sample axis / summary columns.\n\n`__all__ = [\"phenotypic_flux\"]`\n\n> **No `*_table` and no `*_delta` functions.** `groupby` subsumes the old `mi_compare` / `delta_entropy_table` / `flux_table`; comparisons (\u0394, P(>0), HDI) are produced by `groupby` + `_stats` reducers, not bespoke functions. `delta_clonotypic_entropy`, `delta_entropy_table`, `mi_compare`, `clonality`, `clonotypic_entropy_base` are all removed.\n\n---\n\n## 6. `tcri.pl` \u2014 plotting (`plotting/`)\n\nTwins mirror `tl` by filename and function name. Each renderer accepts the same metric arguments as its `tl` twin (computing the joint internally when needed) plus rendering args; comparison across groups is driven by `groupby`, not a separate `*_compare` function.\n\n### `plotting/_entropy.py`\n\n```python\nclonotypic_entropy(\n adata, *,\n covariate=None, groupby=None,\n n_samples=0, temperature=1.0, clones=None, normalised=True,\n palette=None, hue_order=None, ax=None, figsize=(6, 3),\n rotation=90, legend_fontsize=6, bbox_to_anchor=(1.15, 1.0),\n save=None, return_df=False,\n)\n```\n**(renamed from `clonotypic_entropy_by_phenotype`)** Box-and-dot plot of clonotypic entropy per phenotype across covariate values, with per-group dots and significance brackets (via `_stats`).\n\n```python\nphenotypic_entropy(\n adata, *,\n covariate=None, groupby=None,\n n_samples=0, temperature=1.0, clones=None, normalised=True,\n palette=None, ax=None, figsize=(8, 4),\n rotation=90, legend_fontsize=6, bbox_to_anchor=(1.15, 1.0),\n save=None, return_df=False,\n)\n```\n**[FIXED]** Box/strip plot of phenotypic entropy per covariate/group.\n\n`__all__ = [\"clonotypic_entropy\", \"phenotypic_entropy\"]`\n\n### `plotting/_mutual_information.py`\n\n```python\nmutual_information(\n adata, *,\n covariate=None, groupby=None,\n n_samples=0, temperature=1.0, clones=None,\n normalised=True, normalise_mode=\"average\",\n palette=None, ax=None, figsize=(8, 4), rotation=90,\n legend_fontsize=6, bbox_to_anchor=(1.15, 1.0),\n save=None, return_df=False,\n)\n```\n**[FIXED; subsumes `mi_compare`]** Box/strip plot of clone\u00d7phenotype MI per covariate; `groupby` (e.g. `\"patient\"`) supplies the per-group points and drives AUROC/MWU/label-permutation stats from `_stats`.\n\n`__all__ = [\"mutual_information\"]`\n\n### `plotting/_flux.py`\n\n```python\nphenotypic_flux(\n adata, *,\n order,\n groupby=None, clones=None,\n normalize=True, temperature=1.0,\n phenotype_colors=None, ax=None, figsize=(6, 3),\n show_legend=True, title=None, save=None, return_axes=False,\n)\n```\nThe Sankey. Draws phenotype-distribution flow across the ordered `order` sequence of covariate values (calling `tl.phenotypic_flux` pairwise between consecutive values under the hood). `order` replaces the pairwise `cov_from`/`cov_to` because a Sankey spans the full ordered series.\n\n`__all__ = [\"phenotypic_flux\"]`\n\n### `plotting/_sankey.py` \u2014 private drawing primitives\n\n**`class SankeyNode`** *(internal)* \u2014 one rectangular node + its outgoing ribbons.\n\n| Member (signature) | Responsibility |\n|---|---|\n| `__init__(self, x, y, val, *, dx=0.2, color=None, **kwargs)` | Compute the node bounding box and build the `mpatches.Rectangle`. |\n| `plot(self, ax)` | Render the rectangle onto `ax`. |\n| `plot_node_connection(self, destination_node, ax, **kwargs)` | Draw the curved, color-interpolated ribbon (sigmoid-edged `fill_between`) to a destination node. |\n\n`_phenotype_mass_per_clone(adata, covariate, clones, normalize) -> dict[str, np.ndarray]` \u2014 `{clone_id \u2192 phenotype-mass vector}` at one covariate, summing `joint_distribution` rows (optionally weighted). *(internal)*\n\n> `SankeyNode.hex_to_rgb` is **deleted** (0 callers; ribbons use `mcolors.to_rgb`).\n\n### `plotting/_colors.py`\n\n| Symbol (signature) | Responsibility |\n|---|---|\n| `tcri_colors` (module constant: `list[str]`) | Canonical categorical hex palette. |\n| `resolve_palette(adata, columns, *, palette=None) -> dict` | **(renamed from `set_color_palette`)** Assign `tcri_colors` to each `obs` column's categories, store in `uns[\"_colors\"]`, return the category\u2192color map. Fixes the old \"writes on `adata.copy()`\" bug (mutates in place). |\n\n`__all__ = [\"tcri_colors\", \"resolve_palette\"]`\n\n### `plotting/_base.py` \u2014 private plotting engine\n\n| Signature | Responsibility |\n|---|---|\n| `_metric_boxplot(adata, function, *, groupby=None, splitby=None, ylabel=\"\", order=None, palette=None, s=20, ax=None, figsize=(8, 4)) -> (fig, ax)` | **(renamed from `tcri_boxplot`)** Generic per-phenotype metric box/strip engine applying a metric `function` across `groupby`/`splitby` strata. |\n| `_finish(fig, ax, *, save=None, show=None, return_axes=False)` | scanpy-style show/save/return finalizer shared by every `pl` function. |\n\n---\n\n## 7. `tcri.diag` \u2014 diagnostics (`diagnostics/`) \u2014 NEW\n\n`gf.diag`-style: read-only checks on the finalized model. The PPCs return `DataFrame`s; the two relocated model plots render figures.\n\n### `diagnostics/_ppc.py`\n\n| Signature | Responsibility |\n|---|---|\n| `joint_distribution_ppc(adata, model=None, *, covariate=None) -> pandas.DataFrame` | **(the fixed `compare_joint_distribution`)** Model `p(clone, phenotype)` vs empirical hard-label counts, tidy per covariate; the undefined-`model` NameError is repaired (model passed explicitly). |\n| `phenotype_calibration(adata, model=None, *, n_bins=10) -> pandas.DataFrame` | Reliability table: predicted phenotype-probability bins vs observed frequency (ECE-style). |\n| `reconstruction_ppc(adata, model=None, *, n_samples=100, seed=0) -> pandas.DataFrame` | ZINB reconstruction PPC: simulate from the fitted likelihood and compare library size / dropout / mean-variance against observed. |\n| `permutation_null(adata, *, metric=\"mutual_information\", covariate=None, groupby=None, n_permutations=1000, seed=0) -> pandas.DataFrame` | Observed entropy/MI vs a label-permutation null; returns statistic, null summary, and empirical p. |\n\n`__all__ = [\"joint_distribution_ppc\", \"phenotype_calibration\", \"reconstruction_ppc\", \"permutation_null\"]`\n\n### `diagnostics/_training.py`\n\n| Signature | Responsibility |\n|---|---|\n| `loss(model, *, log_scale=False, ax=None, save=None)` | **(relocated `plot_loss`)** Plot training/validation ELBO and prior-KL curves from `model.history_`. |\n| `archetypes(model, *, ax=None, save=None)` | **(relocated `plot_archetypes`)** Heatmap the cluster-ordered clone-phenotype matrix and the archetype centroids. |\n\n`__all__ = [\"joint_distribution_ppc\", \"phenotype_calibration\", \"reconstruction_ppc\", \"permutation_null\", \"loss\", \"archetypes\"]`\n\n> The model PGM (`build_nested_tcri_pgm` / `draw_tcri_pgm_nested`) is moved **out of the package** to `docs/` \u2014 not part of the surface.\n\n---\n\n## 8. `tcri.ut` \u2014 utilities (`utils/`)\n\n### `utils/_session.py` \u2014 public\n\n| Signature | Responsibility |\n|---|---|\n| `save_tcri_session(model, adata, out_dir, *, save_adata=True, compression=\"gzip\") -> dict` | Persist a trained session: scvi model (weights + registry, no embedded adata), Pyro param store, `setup.json`, and the h5ad. |\n| `load_tcri_session(run_dir, *, adata_path=None, map_location=None, layer=None) -> (TCRIModel, AnnData)` | Reconstruct a trained `TCRIModel` + `AnnData` from a saved run dir: read h5ad, restore setup/category order, re-run `setup_anndata`, load model + Pyro params. |\n\n### `utils/_session.py` \u2014 private load/serialization helpers (not re-exported)\n\n| Signature | Responsibility |\n|---|---|\n| `_to_jsonable(x) -> Any` | Recursively coerce numpy/torch/nested values into JSON-serializable primitives. |\n| `_collect_setup_from_adata_or_model(adata, model) -> dict` | Assemble the setup dict (phenotype/clone/covariate/batch cols, category lists, layer) from `uns[K.METADATA]`/categories and, if present, the model registry. |\n| `_restore_category_order(adata, setup) -> None` | Re-impose saved categorical ordering on the `obs` phenotype/clone/covariate columns. |\n| `_resolve_TCRIModel()` | Dynamically locate/import `TCRIModel` across common paths (editable installs). |\n| `_disable_scvi_onload_train()` | Context manager that no-ops scvi's `on_load` warmup train during load (inner `_noop(self, *args, **kwargs)` clears the Pyro store). |\n| `_ensure_pyro_posterior_params(model, adata) -> None` | After load, guarantee `q_p_ct_raw` exists; if missing, warn and re-init to a uniform `1/P` simplex. |\n| `_pyro_load(path, *, map_location=None)` | `torch.load` a Pyro param-store state dict and set it into the global store. |\n| `_ensure_dir(path) -> None` | `os.makedirs(path, exist_ok=True)` wrapper. |\n\n> **Removed from utils:** `write_adata_safely`, `_pop_nonserializables` (manager no longer stashed \u2192 hack unneeded); `probabilities` (dead: read a never-written `uns` key); `build_nested_tcri_pgm` / `draw_tcri_pgm_nested` (\u2192 `docs/`); `stars` / `auc_and_label_permutation` / `bootstrap_auc` (\u2192 `_stats.py`).\n\n---\n\n## 9. Resolution of the OPEN question \u2014 prior vs. mean\n\n**Resolved.** The deterministic point estimate (`n_samples=0`) is the **closed-form variational posterior mean** `E_q[p_ct] = normalize(q_p_ct_raw)` \u2014 exactly `TCRIModel.get_p_ct()` = `adata.uns[K.P_CT]`. Grounding:\n\n- The learned posterior is `q(p_ct) = Dirichlet(\u03b1)` with `\u03b1 = local_scale \u00b7 normalize(q_p_ct_raw)`. Its mean has the closed form `\u03b1/\u03a3\u03b1 = normalize(q_p_ct_raw)`; `local_scale` cancels. Sampling-and-averaging converges to this same vector but only up to Monte-Carlo noise \u2014 a strictly worse estimator of a quantity available in closed form. So `n_samples=0` reads `uns[K.P_CT]` directly and never draws.\n- The **prior is the wrong object** and is **dropped from the public surface**: `clone_phen_prior` / the archetype `mixture_concentration` are guide-initialization / generative anchors, not what training learned. There is **no prior point estimate** and **no prior/`NotImplementedError` branch** in any metric.\n- The **MAP/mode** `(\u03b1\u22121)/(\u03a3\u03b1\u2212P)` is rejected: with `\u03b1_k = local_scale\u00b7m_k` frequently `< 1`, the mode sits on the simplex boundary / is undefined \u2014 unstable and non-representative.\n\n**Consequences for the surface:**\n- **Metrics carry no `posterior=` argument.** They always use the learned posterior \u2014 `n_samples=0` (closed-form mean) or `n_samples>0` (Dirichlet draws) \u2014 and, when handed an `adata`, compute the joint via `joint_distribution(posterior=True)` internally.\n- **`posterior=` survives only on `tl.joint_distribution`, redefined** to mean \"fold in per-cell classifier logits (`True`) vs use the ct-level `p_ct` table directly (`False`).\" Both are posterior objects; the flag is a classifier-mixing switch, not a prior/posterior switch. `local_scale` is consulted **only** for `n_samples>0` draws.\n\n---\n\n## 10. Surface deltas (removed / moved, for completeness)\n\n- **Deleted (dead/broken/out-of-scope):** `clonality` (both tl + pl), `clonotypic_entropy_base`, `delta_clonotypic_entropy`, `delta_entropy_table`, `mi_compare` (tl + pl), `flux_table`, `bayesian_mutual_information`, `probability_distribution`, `polar_plot`, `classify_phenotypes`, `gene_entropy`, `get_latent_embedding`, `group_small_clones`, `register_probability_columns`, `remove_meaningless_genes`, `_ent`, `clone_fraction`, `dkl` (\u2192 `_distance.kl_divergence`), `probabilities`, `SankeyNode.hex_to_rgb`, `_ascii_hist` (+ all `graph=`/ASCII paths), `write_adata_safely`, `_pop_nonserializables`.\n- **Renamed:** `flux`\u2192`tl.phenotypic_flux`; `get_cell_phenotype_probs`\u2192`TCRIModel.predict`; `register_model`\u2192`TCRIModel.to_anndata`; `clonotypic_entropy_by_phenotype`\u2192`pl.clonotypic_entropy`; `tcri_boxplot`\u2192`_base._metric_boxplot`; `set_color_palette`\u2192`resolve_palette`; params `from_this`/`to_that`\u2192`cov_from`/`cov_to`, `point_estimate=`\u2192removed (use `n_samples`).\n- **Moved to `examples/`:** `top_clone_umap`, `clone_size_umap`, `plot_phenotype_probabilities`, `compare_phenotypes`. **Moved to `docs/`:** `build_nested_tcri_pgm`, `draw_tcri_pgm_nested`. **Moved to `diag/`:** `compare_joint_distribution`(\u2192`joint_distribution_ppc`), `plot_loss`(\u2192`loss`), `plot_archetypes`(\u2192`archetypes`).\n\n---\n\n_Document home (if persisted): `/Users/ceglian/Codebase/GitHub/tcri/docs/contract/tcri_api_surface.md`. Source of truth cross-checked against `tcri/model/_model.py`, `tcri/preprocessing/_preprocessing.py`, `tcri/metrics/_metrics.py`, `tcri/plotting/_plotting.py`, `tcri/plotting/_sankey.py`, `tcri/utils/_utils.py`, and the grafiti reference at `/Users/ceglian/Codebase/GitHub/grafiti/grafiti`._", "mathspec": "I have enough grounding from the code, the extracted math, the inference-path analysis, and the inventory. Producing the spec.\n\n# TCRI \u2014 Per-Function Responsibility & Math/Stats Specification\n\n*Validate-the-math-ahead-of-time artifact for the refactor. Every function that performs a mathematical or statistical operation in the target API is specified below: (a) responsibility, (b) exact math, (c) how each argument dictates the operation, (d) exact return shape per argument combination. Formulas are grounded in the current code (`tcri/metrics/_metrics.py`, `tcri/preprocessing/_preprocessing.py`, `tcri/model/_model.py`, `tcri/utils/_utils.py`, `tcri/plotting/_plotting.py`); where the target design changes the math, the NEW intended operation is stated precisely and marked **[NEW]**.*\n\n---\n\n## 0. Shared substrate, notation, and resolved conventions\n\n### 0.1 Notation\n\n| Symbol | Meaning |\n|---|---|\n| $P$ | number of phenotypes (columns of every joint) |\n| $C$ | number of clonotypes (rows of a joint at one covariate) |\n| $M$ | number of covariate values (e.g. timepoints) |\n| $\\mathrm{ct}$ | a (clonotype, covariate) group; `ct_count` = number of them |\n| $p_{ct}\\in\\Delta^{P}$ | learned per-`ct` phenotype distribution, `adata.uns[\"tcri_p_ct\"]`, shape `(ct_count, P)` |\n| $\\ell_i\\in\\mathbb{R}^P$ | per-cell classifier logits, `adata.obsm[\"X_tcri_logits\"]` |\n| $s$ | local scale, `adata.uns[\"tcri_local_scale\"]` (Dirichlet total concentration; default 1.0) |\n| $T$ | `temperature` |\n| $\\varepsilon$ | numerical floor (values below in situ) |\n\nThe index maps used for selection/aggregation: `uns[\"tcri_ct_to_cov\"]`, `uns[\"tcri_ct_to_c\"]`, `uns[\"tcri_ct_array_for_cells\"]`, `uns[\"tcri_cov_array_for_cells\"]`.\n\n### 0.2 The one substrate every metric reads\n\nTraining learns a variational Dirichlet posterior on each `ct` row:\n\n$$q(p_{ct})=\\mathrm{Dirichlet}(\\alpha),\\qquad \\alpha = s\\cdot m,\\quad m=\\text{normalize}\\big(q\\_p\\_ct\\_raw\\big)\\in\\Delta^{P}.$$\n\nBecause $\\alpha=s\\cdot m$ with $m$ on the simplex, $\\sum_k\\alpha_k=s$ and the **mean is closed-form**:\n\n$$\\mathbb{E}_q[p_{ct}] \\;=\\; \\frac{\\alpha}{\\sum_k\\alpha_k} \\;=\\; m \\;=\\; \\texttt{get\\_p\\_ct()} \\;=\\; \\texttt{adata.uns[\"tcri\\_p\\_ct\"]}.$$\n\n`get_p_ct()` returns $m$ exactly (it row-normalizes `q_p_ct_raw`); the local scale $s$ cancels and is irrelevant to the mean. $s$ matters **only** for the spread of draws.\n\n### 0.3 RESOLUTION of the open question (prior vs. mean vs. MAP vs. draw)\n\nThis is the audit's required decision. Grounded in the inference path:\n\n1. **The deterministic point estimate (`n_samples=0`) is the closed-form posterior mean $\\mathbb{E}_q[p_{ct}]=m=$ `uns[\"tcri_p_ct\"]`** (temperature-adjusted, \u00a71). It is available in closed form; no sampling is needed or desirable.\n2. **Not the generative prior.** The archetype `mixture_concentration` (KMeans centroids) and the `clone_phen_prior` buffer are the model's *initialization / generative anchor*, never what training learned. **No metric may read them.** Consequently the current \"prior-only\" language is misleading and the `mutual_information(posterior=False) \u2192 NotImplementedError` dead branch is deleted, not implemented. *Note:* both current \"distribution\" functions already use the posterior mean of $p_{ct}$ \u2014 neither ever touched the generative prior \u2014 so the word \"posterior\" in today's flags is doubly misleading.\n3. **Not the MAP/mode.** The Dirichlet mode $(\\alpha-1)/(\\sum\\alpha-P)$ is degenerate here: $\\alpha_k=s\\,m_k$ is frequently $<1$ (small $s$, $m_k<1$), pushing the mode onto the simplex boundary. Unstable and non-representative.\n4. **Not the mean-of-draws.** Averaging many Dirichlet draws converges to $m$ only up to Monte-Carlo noise \u2014 a strictly worse estimator of a quantity available exactly.\n\n**What remains of the `posterior=` axis.** The only *real* difference between today's two engines is not prior-vs-posterior but **whether per-cell classifier logits $\\ell_i$ are folded in**. The resolved meaning of the engine's `posterior=` flag is therefore:\n\n- `posterior=True` (default): the joint reflects the model's full per-cell phenotype posterior \u2014 classifier logits combined with $\\log p_{ct}$, aggregated to clone level (matches `predict()`).\n- `posterior=False`: the joint is the clone-level $p_{ct}$ table alone (classifier evidence ignored; faster).\n\nBoth branches use the *posterior mean/draws* of $p_{ct}$; neither is the generative prior. **Recommendation:** document `posterior=` with this meaning (a rename such as `cell_informed=`/`use_logits=` would remove the residual ambiguity), and make the `posterior=True` combination rule *identical* to `predict()` \u2014 gate-aware when the model was trained with a gate \u2014 to fix the standing bug that metrics silently disagree with the model's own cell probabilities (today they use a Dirichlet *draw* and *never* apply the gate).\n\n### 0.4 Uniform sampling convention (all metrics + engine)\n\n| `n_samples` | Operation |\n|---|---|\n| `0` | **Deterministic point estimate.** Use posterior mean $m$; no Dirichlet draw; reproducible. **[NEW]** \u2014 fixes today's latent bug where `mutual_information`/`flux` at `n_samples=0` return *one random draw*. |\n| `N>0` | **$N$ i.i.d. posterior draws** $p_{ct}^{(s)}\\sim\\mathrm{Dirichlet}(s\\cdot \\tilde p_{ct})$. Adds a sample axis; mean \u00b1 credible interval fall out. |\n\nThe old `point_estimate=` argument is **deleted**; `n_samples` is the only knob.\n\n### 0.5 Uniform return-shape rule\n\n| `groupby` | `n_samples` | Return |\n|---|---|---|\n| unset | `0` | scalar (MI, flux-of-one-clone) or `Series` (entropies over phenotypes/clones) |\n| unset | `N>0` | draw array with a sample axis **+** summary columns `mean, hdi_low, hdi_high` (and `p_gt` for signed/paired quantities) |\n| set | `0` | tidy `DataFrame`, one row per group (\u00d7 phenotype / \u00d7 clone) |\n| set | `N>0` | tidy `DataFrame`, one row per group (\u00d7 phenotype / \u00d7 clone) with summary columns `mean, sd, hdi_low, hdi_high, p_gt` |\n\nComparisons between groups are done by `groupby` **plus** the general stats helper (\u00a77), not by `*_delta` / `*_table` functions (all deleted).\n\n### 0.6 The `weighted` axis \u2014 resolved by removal\n\n`weighted` is **dropped from every public signature** (engine and all four metrics). This enshrines a **uniform prior over clonotypes**: each clonotype is one unit on the simplex regardless of cell count (a 1-cell clone counts as much as a $10^4$-cell clone). Every entropy/MI/flux is therefore a **repertoire-level (per-clonotype) statistic**, not a per-cell statistic. This convention must be stated in each metric's docstring. (It also removes the current weighting bug in `joint_distribution`, where the weight lookup keyed a CT-indexed `Counter` with clone indices and produced near-arbitrary weights, and the inconsistency where `joint_distribution` normalized the whole table to sum 1 while `joint_distribution_posterior` returned un-normalized counts.)\n\n---\n\n## 1. Engine \u2014 `tcri.tl.joint_distribution`\n\n```\njoint_distribution(adata, *, covariate=None, groupby=None,\n n_samples=0, posterior=True, clones=None,\n temperature=1.0) -> pandas.DataFrame\n```\nUnifies today's `joint_distribution` + `joint_distribution_posterior`. Re-exported top-level as `tcri.joint_distribution`.\n\n### (a) Responsibility\nProduce the clone\u00d7phenotype distribution (the substrate every metric consumes) at a covariate value, from the learned variational posterior of $p_{ct}$, either as a deterministic posterior-mean table or as posterior draws. Provenance recorded in `.attrs`.\n\n### (b) Math\n\nLet $\\mathcal{I}_m=\\{\\mathrm{ct}: \\texttt{ct\\_to\\_cov}[\\mathrm{ct}]=m\\}$ select the covariate. Within one covariate each clone maps to exactly one `ct`, so rows are indexed by clonotype.\n\n**Temperature-tempered base** (applied to the posterior mean $m$):\n$$\\tilde p_{ct} \\;=\\; \\mathrm{softmax}\\!\\Big(\\tfrac{1}{T}\\log(m_{ct}+\\varepsilon)\\Big) \\;=\\; \\frac{m_{ct}^{1/T}}{\\sum_\\phi m_{ct,\\phi}^{1/T}},\\qquad \\varepsilon=10^{-8}.$$\nAt $T=1$ this is the identity (renormalization only); $T<1$ sharpens, $T>1$ flattens.\n\n**`posterior=False` (clone-level $p_{ct}$ only):**\n- `n_samples=0`: row $c$ = $\\tilde p_{ct}$ (the covariate-restricted posterior-mean table). Deterministic.\n- `n_samples=N`: draw $p_{ct}^{(s)}\\sim\\mathrm{Dirichlet}(s\\cdot\\tilde p_{ct})$, $s=1,\\dots,N$. (Grounded in current `joint_distribution` n_samples>0 branch: `Dirichlet(local_scale * p_ct_for_cov)`.)\n\n**`posterior=True` (fold per-cell classifier evidence; matches `predict()`) [NEW: deterministic at `n_samples=0`, gate-aware]:**\nPer cell $i$ at covariate $m$ with clonotype $c(i)$ and group $\\mathrm{ct}(i)$, using base $b_{ct}$ ($=m_{ct}$ for the mean path, or a draw for `n_samples>0`):\n$$P(\\phi\\mid i)=\\mathrm{softmax}\\!\\Big(\\tfrac1T\\,\\mathrm{combine}\\big(\\ell_i,\\ \\log(b_{\\mathrm{ct}(i)}+\\varepsilon)\\big)\\Big),\\quad\n\\mathrm{combine}=\\begin{cases}g\\,\\ell_i+(1-g)\\log b & \\text{if gated }(g=\\texttt{gate\\_prob})\\\\ \\ell_i+\\log b & \\text{additive}\\end{cases}$$\nthen aggregate to clone level and normalize:\n$$J[c,\\phi]=\\sum_{i\\in c}P(\\phi\\mid i),\\qquad J[c,\\cdot]\\leftarrow J[c,\\cdot]\\big/\\textstyle\\sum_\\phi J[c,\\phi].$$\n- `n_samples=0`: $b=m$ (posterior mean). **Deterministic** \u2014 fixes today's `joint_distribution_posterior`, which always drew exactly one Dirichlet sample and always used the additive rule with no gate.\n- `n_samples=N`: $N$ draws $b=p_{ct}^{(s)}\\sim\\mathrm{Dirichlet}(s\\cdot m)$.\n\n**`groupby=g`:** partition cells by `obs[g]` and compute the above independently per group value, stacking a group axis into the tidy frame.\n\n### (c) Arguments \u2192 math\n\n| Argument | Effect on the math |\n|---|---|\n| `covariate` | Selects the `ct` rows via `ct_to_cov`; defines the condition $m$. Required unless a precomputed joint is supplied downstream. |\n| `groupby` | Computes a separate joint per group value (adds a group axis). Unset \u2192 single joint. |\n| `n_samples` | `0` \u2192 posterior-mean table (deterministic); `N` \u2192 $N$ Dirichlet draws from $q(p_{ct})$ with total concentration $s$. Only place $s=$`local_scale` enters. |\n| `posterior` | `True` \u2192 fold per-cell logits $\\ell_i$ with $\\log b$ (gate-aware), aggregate per clone, row-normalize; `False` \u2192 clone-level $\\tilde p_{ct}$ rows directly (no logits). Neither is the generative prior. |\n| `clones` | Filters rows to the listed clonotypes; with `posterior=True` also restricts the cells aggregated; final reindex to the exact list (absent clones \u2192 all-zero rows \u2192 should be dropped/NaN, see \u00a72\u2013\u00a73 bug note). |\n| `temperature` | $T$ divides the logit before softmax: power-tempers $m$ (`posterior=False`) or the combined per-cell logit (`posterior=True`). $T=1$ identity. **[NEW]** applied consistently in both branches (today the two functions temper at different stages). |\n\n### (d) Return shape\n\n`pandas.DataFrame`, columns = phenotype categories.\n\n| `groupby` | `n_samples` | Index / axes |\n|---|---|---|\n| unset | `0` | rows = clonotype id; shape `(C_m, P)` |\n| unset | `N>0` | rows = `f\"{clone}_{s}\"` for $s=0..N-1$; shape `(C_m\u00b7N, P)` (a `sample_id` level) |\n| set | `0` | MultiIndex (group, clonotype); one block per group |\n| set | `N>0` | MultiIndex (group, clonotype, sample_id) |\n\n---\n\n## 2. Metric \u2014 `tcri.tl.clonotypic_entropy`\n\n```\nclonotypic_entropy(adata_or_jd, *, covariate=None, groupby=None,\n n_samples=0, temperature=1.0, clones=None,\n normalised=True) -> Series | DataFrame\n```\n\n### (a) Responsibility\nFor each phenotype $\\phi$ (at covariate $m$), the normalized Shannon entropy of the distribution over clonotypes carrying that phenotype, $H[P(c\\mid\\phi,m)]$. High \u2192 the phenotype is spread across many clones; low \u2192 a few clones dominate it.\n\n### (b) Math\nFrom the joint $J$ (rows=clones, cols=phenotypes), take column $\\phi$ and normalize over clones:\n$$v_c=\\frac{\\max(J[c,\\phi],\\varepsilon)}{\\sum_{c'}\\max(J[c',\\phi],\\varepsilon)},\\quad \\varepsilon=10^{-15},\\qquad\nH_\\phi=-\\sum_{c} v_c\\log_2 v_c\\ \\text{(bits)}.$$\nIf `normalised` and $C>1$:\n$$H_\\phi \\leftarrow H_\\phi/\\log_2 C,\\qquad C=\\text{n\\_clones}\\in[0,1].$$\n(Grounded in current `clonotypic_entropy` / `clonotypic_entropy_base`.) $v$ equals $P(c\\mid\\phi)$ under the uniform-clonotype prior (\u00a70.6).\n\n**[NEW]** `n_samples=0` returns the deterministic estimate from the posterior mean (today the function has no `n_samples=0` path and always samples $\\ge1$ draw). **[NEW]** `normalised` is exposed (today `clonotypic_entropy` is hard-normalized; only the merged-away `_base` had the toggle). Base is fixed to 2 (bits).\n\n### (c) Arguments \u2192 math\n\n| Argument | Effect |\n|---|---|\n| `adata_or_jd` | If AnnData \u2192 compute $J$ internally via \u00a71 (`covariate` required); if a precomputed joint DataFrame \u2192 skip straight to the entropy (fast path). |\n| `covariate` | Condition $m$; selects `ct` rows. |\n| `groupby` | Entropy computed per group value \u2192 tidy rows (group \u00d7 phenotype). |\n| `n_samples` | `0` \u2192 posterior-mean entropy per phenotype; `N` \u2192 $N$ per-draw entropies (sample axis + summary). |\n| `temperature` | Tempers the joint before the column is extracted (sharpen `<1` / flatten `>1`). |\n| `clones` | Restricts the clone set \u2192 changes $C$ (the $\\log_2 C$ normalizer) and which clones enter $v$. |\n| `normalised` | `True` \u2192 divide by $\\log_2 C$ (range $[0,1]$); `False` \u2192 raw bits. $C\\le1$ \u2192 no division (H already 0). |\n| `weighted` | *Removed.* Uniform-clonotype prior enforced (\u00a70.6). |\n\n### (d) Return shape\n\n| `groupby` | `n_samples` | Return |\n|---|---|---|\n| unset | `0` | `Series` indexed by phenotype (bits or $[0,1]$) |\n| unset | `N>0` | `(N, P)` array **or** long DataFrame `(sample_id, phenotype, value)` + summary `mean, hdi_low, hdi_high` per phenotype |\n| set | `0` | tidy DataFrame, row per (group, phenotype) |\n| set | `N>0` | tidy DataFrame, row per (group, phenotype) + `mean, sd, hdi_low, hdi_high` |\n\n**Edge cases / bug fixes:** absent phenotype \u2192 NaN (not 0). An absent/zero-mass clone column that clips to $\\varepsilon$ slightly inflates entropy toward uniform \u2014 target should drop zero-support clones before normalizing rather than let $\\varepsilon$-clipping fabricate mass.\n\n---\n\n## 3. Metric \u2014 `tcri.tl.phenotypic_entropy`\n\n```\nphenotypic_entropy(adata_or_jd, *, covariate=None, groupby=None,\n n_samples=0, temperature=1.0, clones=None,\n normalised=True) -> Series | DataFrame\n```\n\n### (a) Responsibility\nFor each clonotype $c$ (at covariate $m$), the normalized Shannon entropy of its phenotype distribution $H[P(\\phi\\mid c,m)]$. High \u2192 phenotypically plastic clone; low \u2192 committed to one phenotype.\n\n### (b) Math\nRow $c$ of the joint is $P(\\phi\\mid c)$. Normalize and take entropy over phenotypes:\n$$p_\\phi=\\frac{\\max(J[c,\\phi],\\varepsilon)}{\\sum_{\\phi'}\\max(J[c,\\phi'],\\varepsilon)},\\quad\nH_c=-\\sum_\\phi p_\\phi\\log_2 p_\\phi,\\qquad\nH_c\\leftarrow H_c/\\log_2 P\\ \\text{if normalised and }P>1.$$\n(Grounded in current `phenotypic_entropy`.) **[NEW]** `n_samples=0` deterministic; `normalised` exposed; base fixed to 2.\n\n### (c) Arguments \u2192 math\n\n| Argument | Effect |\n|---|---|\n| `adata_or_jd` / `covariate` | Compute $J$ internally, or consume a precomputed joint. |\n| `groupby` | Per-group entropy \u2192 tidy rows (group \u00d7 clone). |\n| `n_samples` | `0` \u2192 posterior-mean $H$ per clone; `N` \u2192 per-draw matrix + summary. |\n| `temperature` | Tempers the joint before rows are read. |\n| `clones` | Restricts which clones are scored (the normalizer $\\log_2 P$ is unaffected \u2014 it depends on $P$, not clone count). |\n| `normalised` | `True` \u2192 divide by $\\log_2 P$; `False` \u2192 raw bits. $P\\le1$ \u2192 divisor 1. |\n| `weighted` | *Removed* (uniform-clonotype prior). |\n\n### (d) Return shape\n\n| `groupby` | `n_samples` | Return |\n|---|---|---|\n| unset | `0` | `Series` indexed by clonotype |\n| unset | `N>0` | `(N, C)` array + per-clone `mean, hdi_low, hdi_high` |\n| set | `0` | tidy DataFrame, row per (group, clone) |\n| set | `N>0` | tidy DataFrame + `mean, sd, hdi_low, hdi_high` |\n\n**Bug fix (critical):** today an all-zero clone row (clone present in `obs` but zero posterior mass) is reindexed to zeros, $\\varepsilon$-clipped, and normalizes to **uniform \u2192 entropy 1.0** instead of NaN (the `cl not in jd.index` guard never fires). Target: zero-support clones return NaN / are excluded, not assigned maximal entropy.\n\n---\n\n## 4. Metric \u2014 `tcri.tl.mutual_information` (+ kernel `_mi_from_joint`)\n\n```\nmutual_information(adata_or_jd, *, covariate=None, groupby=None,\n n_samples=0, temperature=1.0, clones=None,\n normalised=True, normalise_mode='average') -> float | ndarray | DataFrame\n```\n\n### (a) Responsibility\n$I(c;\\phi\\mid m)$ \u2014 how much a cell's clonotype tells you about its phenotype at covariate $m$; the strength of clone\u2013phenotype coupling. Optionally normalized to $[0,1]$ (symmetric uncertainty).\n\n### (b) Math \u2014 kernel `_mi_from_joint(pxy, normalised, mode)`\nNormalize the whole joint to $p(c,\\phi)$: $\\ \\text{pxy}=J/\\sum J$. Marginals $p_x=\\sum_\\phi\\text{pxy}$ (=$P(c)$), $p_y=\\sum_c\\text{pxy}$ (=$P(\\phi)$).\n$$I=\\sum_{c,\\phi}\\text{pxy}\\,\\log_2\\frac{\\text{pxy}+\\varepsilon}{p_x p_y^\\top+\\varepsilon}\\ \\text{(bits)},\\qquad \\varepsilon=10^{-15}.$$\nIf `normalised`, with marginal entropies $H_c=-\\sum p_x\\log_2 p_x$, $H_p=-\\sum p_y\\log_2 p_y$:\n$$I_{\\text{norm}}=\\frac{I}{D},\\quad D=\\begin{cases}\\tfrac12(H_c+H_p)&\\texttt{mode='average'}\\ \\text{(symmetric uncertainty)}\\\\[2pt]\\min(H_c,H_p)&\\text{otherwise (coefficient of constraint; }\\ge\\text{average form)}\\end{cases}$$\nreturning $0$ if $D\\le0$ (degenerate single-clone / single-phenotype table). (Grounded exactly in `_mi_from_joint`.)\n\n**Uniform-clonotype prior:** because each joint row sums to 1 and the whole table is then divided by its sum, $P(c)=1/C$ is imposed \u2014 the resolved convention (\u00a70.6), to be documented in the docstring. **[NEW]** `n_samples=0` = deterministic point estimate (today it returns one random draw). The `posterior=False \u2192 NotImplementedError` dead branch is removed (\u00a70.3).\n\n### (c) Arguments \u2192 math\n\n| Argument | Effect |\n|---|---|\n| `adata_or_jd` / `covariate` | Compute $J$ internally, or consume a precomputed joint. |\n| `groupby` | MI per group value \u2192 tidy rows (one MI per group). Group comparisons via \u00a77. |\n| `n_samples` | `0` \u2192 single deterministic MI (from posterior mean); `N` \u2192 array of per-draw MI + summary. |\n| `temperature` | Tempers the joint before marginalization. |\n| `clones` | Restricts rows (the clone support of the table). |\n| `normalised` | `True` \u2192 divide by $D$ (range $[0,1]$); `False` \u2192 raw bits. |\n| `normalise_mode` | `'average'` \u2192 $D=\\tfrac12(H_c+H_p)$; any other value \u2192 $D=\\min(H_c,H_p)$. |\n| `weighted` | *Removed* \u2014 MI is per-clonotype, not per-cell. |\n\n### (d) Return shape\n\n| `groupby` | `n_samples` | Return |\n|---|---|---|\n| unset | `0` | scalar `float` |\n| unset | `N>0` | `(N,)` array + `mean, sd, hdi_low, hdi_high` |\n| set | `0` | tidy DataFrame, one row per group, column `MI` |\n| set | `N>0` | tidy DataFrame per group + `mean, sd, hdi_low, hdi_high` |\n\n---\n\n## 5. Metric \u2014 `tcri.tl.phenotypic_flux` (renamed from `flux`)\n\n```\nphenotypic_flux(adata, *, cov_from, cov_to, groupby=None,\n n_samples=0, temperature=1.0, clones=None,\n distance_metric='l1') -> Series | ndarray | DataFrame\n```\n\n### (a) Responsibility\nPer-clonotype distance between its phenotype distribution at `cov_from` and at `cov_to`; how much a clone's phenotype mix shifts across two conditions.\n\n### (b) Math\nBuild $J_{\\text{from}}=$ joint at `cov_from`, $J_{\\text{to}}=$ joint at `cov_to` (rows $P(\\phi\\mid c)$). Over the common clones $c\\in J_{\\text{from}}\\cap J_{\\text{to}}$, with $p=J_{\\text{from}}[c]$, $q=J_{\\text{to}}[c]$:\n$$d_c = \\text{distance\\_metric}(p,q),\\qquad\n\\text{l1: } \\sum_\\phi|p_\\phi-q_\\phi|\\in[0,2];\\quad\n\\text{kl: } \\sum_\\phi p_\\phi\\log\\frac{p_\\phi}{q_\\phi}\\ (\\text{clip }\\varepsilon,\\text{ asymmetric }\\mathrm{KL}(p\\Vert q)).$$\n(Grounded in current `flux` `l1` / `dkl_func` / callable dispatch.) Distance kernels live in `_distance.py` (\u00a76).\n\n### (c) Arguments \u2192 math\n\n| Argument | Effect |\n|---|---|\n| `cov_from`, `cov_to` | The two conditions compared (was `from_this`/`to_that`). |\n| `groupby` | Per-group flux \u2192 tidy rows (group \u00d7 clone), with `clone_size` column (replaces `flux_table`). |\n| `n_samples` | `0` \u2192 deterministic per-clone Series (from posterior mean); `N` \u2192 $N$ redrawn distance vectors + summary. |\n| `temperature` | Tempers both joints identically before differencing. |\n| `clones` | Restricts the clone set on both sides; distances taken over the intersection. |\n| `distance_metric` | `'l1'` \u2192 $\\sum|p-q|$; `'kl'` \u2192 asymmetric KL (specify base, see \u00a76); callable $f(p,q)\\to\\mathbb{R}$. |\n| `weighted` | *Removed* (uniform-clonotype prior). |\n\n### (d) Return shape\n\n| `groupby` | `n_samples` | Return |\n|---|---|---|\n| unset | `0` | `Series` indexed by clonotype (distance per clone) |\n| unset | `N>0` | `(N, |common|)` array + per-clone `mean, sd, hdi_low, hdi_high` |\n| set | `0` | tidy DataFrame, row per (group, clone) + `clone_size` |\n| set | `N>0` | tidy DataFrame + `flux_mean, flux_sd, hdi_low, hdi_high` + `clone_size` |\n\n**Bug fixes carried by the rewrite:** (1) the `posterior=False` path is gone, so it no longer passes unsupported `silent=`/`combine_with_logits=` kwargs into `joint_distribution` \u2192 today's `TypeError` is eliminated. (2) **Reproducibility:** posterior draws use `torch` Dirichlet, but the old `seed=` only seeded NumPy \u2192 the seed did nothing; **[NEW]** seed a `torch.Generator` (and NumPy) so flux draws are reproducible. (3) **Unit consistency:** the KL kernel used natural log while entropies/MI use $\\log_2$ \u2014 **[NEW]** pick one base (recommend bits) library-wide, or name the units explicitly. (4) the `flux_table` column-misalignment (`clones_g.index(cl)` vs columns ordered by `common`) is removed by returning a tidy per-(group,clone) frame keyed by clone id.\n\n---\n\n## 6. Distance kernels \u2014 `tcri/_distance.py`\n\nDedupes the dead module-level `dkl` and `flux.dkl_func`.\n\n### `l1_distance(p, q)`\n- **Responsibility / math:** $\\displaystyle \\sum_i |p_i-q_i|$. Symmetric; range $[0,2]$ on simplex inputs. Assumes normalized inputs; defensively renormalize.\n- **Return:** scalar float.\n\n### `kl_divergence(p, q)`\n- **Responsibility / math:** $\\mathrm{KL}(p\\Vert q)=\\sum_i p_i\\log\\frac{p_i}{q_i}$, clip $p,q\\in[\\varepsilon,1]$ then renormalize. **Asymmetric.** Range $[0,\\infty)$.\n- **[NEW] base:** current copies disagree \u2014 the dead module `dkl` uses `scipy.stats.entropy` (natural log, $\\varepsilon=10^{-10}$); `flux.dkl_func` uses natural log with $\\varepsilon=10^{-15}$. Fix one base (recommend $\\log_2$, matching entropies/MI) and one $\\varepsilon$.\n- **Return:** scalar float.\n\n### `phenotype_distance(p, q, metric='l1')`\n- **Responsibility:** dispatcher. `'l1'` \u2192 `l1_distance`; `'kl'` (`'dkl'` alias) \u2192 `kl_divergence`; callable \u2192 `metric(p, q)`; else `ValueError`.\n- **Arguments \u2192 math:** `metric` selects the kernel above.\n- **Return:** scalar float.\n\n---\n\n## 7. Comparison statistics helper \u2014 `tcri/_stats.py`\n\nReplaces every `*_delta` / `*_table` function. Given metric values across a `groupby` (per-patient point estimates) **or** paired posterior-draw vectors, it quantifies group differences.\n\n### `mannwhitney(a, b, alternative='two-sided')`\n- **Responsibility / math:** Mann\u2013Whitney $U$ rank-sum test on two independent samples (e.g. responder vs non-responder patient-level metric values). $U=\\sum_{i,j}\\mathbb{1}[a_i>b_j]+\\tfrac12\\mathbb{1}[a_i=b_j]$; two-sided $p$ from the rank-sum null. Non-parametric; no distributional assumption. (Grounded in `scipy.stats.mannwhitneyu`, currently inline in plotting.)\n- **Arguments:** `alternative` selects one/two-sided.\n- **Return:** `(U: float, p: float)`.\n\n### `prob_direction(delta)` \u2014 Bayesian $P(\\Delta>0)$\n- **Responsibility / math:** given a posterior-draw vector of a signed difference $\\Delta=\\text{metric}_B-\\text{metric}_A$ (paired across draws), the posterior probability of direction:\n$$p_{gt}=\\tfrac1N\\sum_s\\mathbb{1}[\\Delta^{(s)}>0],\\qquad p_{lt}=\\tfrac1N\\sum_s\\mathbb{1}[\\Delta^{(s)}<0].$$\n(Grounded in `delta_entropy_table`: `p_greater=(delta>0).mean()`.)\n- **Return:** `(p_gt: float, p_lt: float)`.\n\n### `credible_interval(x, prob=0.95)` \u2014 HDI/CI\n- **Responsibility / math:** the current code computes the **equal-tailed interval** `np.percentile(x, [2.5, 97.5])` but labels it \"HDI.\" For a rigorous artifact: **[NEW]** either (i) keep equal-tailed and name it `ci_low/ci_high`, or (ii) implement the true **highest-density interval** (narrowest interval containing `prob` mass) \u2014 these differ for skewed posteriors. Return `mean`, `sd`, `low`, `high`.\n- **Arguments:** `prob` sets the mass (default 0.95).\n- **Return:** `(mean, sd, low, high)`.\n\n### `stars(p)`\n- **Math:** $p<10^{-4}\\to$`****`; $<10^{-3}\\to$`***`; $<10^{-2}\\to$`**`; $<0.05\\to$`*`; else `ns`. (Exact from `utils.stars`.)\n\n### `auc_and_label_permutation(scores, labels, pos_label=None, n_perm=200000, max_exact=200000)`\n- **Responsibility / math:** observed $\\mathrm{AUROC}$ of `scores` vs binarized `labels`, plus a two-sided label-permutation $p$. If $\\binom{n}{k}\\le$`max_exact` enumerate all label assignments (exact null), else Monte-Carlo `n_perm` permutations:\n$$p_{\\text{perm}}=\\text{mean}\\big(|\\mathrm{AUC}_{\\text{perm}}-0.5|\\ \\ge\\ |\\mathrm{AUC}_{\\text{obs}}-0.5|\\big).$$\n- **Arguments:** `pos_label` sets the positive class (default max label); `max_exact` switches exact\u2194MC; `n_perm`/`seed` control the MC null.\n- **Return:** `(obs_auc, p_perm, perm_stats, perm_mode)`.\n\n### `bootstrap_auc(scores, labels, pos_label=None, n_boot=5000, seed=42)`\n- **Responsibility / math:** resample cells with replacement (rejecting draws missing a class), recompute AUROC, return the $[2.5\\%, 97.5\\%]$ quantiles as a 95% CI.\n- **Arguments:** `n_boot` resamples; `pos_label`/`seed` as above.\n- **Return:** `np.array([lo, hi])`.\n\n---\n\n## 8. Model \u2014 `TCRIModel.predict` (renamed from `get_cell_phenotype_probs`) and `get_p_ct`\n\n### `predict(adata=None, batch_size=256, eps=1e-8) -> DataFrame`\n\n#### (a) Responsibility\nThe model's own per-cell phenotype-probability matrix (the ground-truth cell posterior), computed exactly as in training. scvi/CellAssign idiom. **[NEW]** returns a `DataFrame` (index = cell barcode, columns = phenotypes) rather than a bare array.\n\n#### (b) Math\nPer cell with input $x$, batch $b$, group $\\mathrm{ct}$:\n$$z=\\text{encoder}(x,b).\\text{loc},\\quad \\ell=\\text{classifier}(z)=\\frac{\\text{mlp}(z)}{\\tau_{\\text{cls}}},\\quad \\pi=\\log(p_{ct}+\\varepsilon),\\ p_{ct}=\\texttt{get\\_p\\_ct()}\\ (\\text{mean}),$$\n$$\\text{local}=\\begin{cases}g\\,\\ell+(1-g)\\,\\pi & \\text{if } \\texttt{use\\_gate}\\ (g=\\texttt{gate\\_prob})\\\\ \\ell+\\pi & \\text{additive}\\end{cases},\\qquad \\text{probs}=\\mathrm{softmax}(\\text{local}).$$\n(Exact from `get_cell_phenotype_probs`.) This gate-aware, mean-prior rule is the reference the metric layer's `posterior=True` joint must match (\u00a70.3).\n\n#### (c) Arguments \u2192 math\n\n| Argument | Effect |\n|---|---|\n| `adata` | Which cells to score (default = training AnnData). |\n| `batch_size` | Loader chunking only; no numeric effect (assumes loader yields cells in `ct_array` order). |\n| `eps` | Floor inside $\\log(p_{ct}+\\varepsilon)$. |\n| `gate_prob` (module attr) | Set \u2192 convex logit mix; `None` \u2192 additive product. |\n| `classifier_temperature` $\\tau_{\\text{cls}}$ (module attr) | Divides classifier logits inside `classifier.forward`. |\n\n#### (d) Return shape\n`DataFrame` `(n_cells, P)`; rows = cell barcodes, columns = phenotype categories; each row a simplex (softmax output).\n\n### `get_p_ct() -> ndarray`\n- **Responsibility / math:** returns the learned posterior mean $m$. Reads `q_p_ct_raw`, replaces non-finite rows with $1/P$, optionally sharpens by $1/\\text{guide\\_temperature}$, row-normalizes: $\\ p_{ct}=q\\_p\\_ct\\_raw^{1/T_g}/\\sum_\\phi(\\cdot)$. Numerically equals $\\mathbb{E}_q[p_{ct}]$ (\u00a70.2).\n- **Arguments:** `guide_temperature` $T_g<1$ sharpens / $>1$ flattens / $=1$ plain normalize.\n- **Return:** `(ct_count, P)` row-stochastic array. (Stored as `uns[\"tcri_p_ct\"]`.)\n- **Caveat:** the guide's `min=1e-3` concentration clamp is not reproduced here, so near-degenerate rows can diverge slightly from the guide's true expectation.\n\n---\n\n## 9. Diagnostics \u2014 `tcri.diag` (PPCs & model validation, return DataFrames)\n\nRead-only checks on the finalized model; return DataFrames (grafiti `gf.diag` style), outside the tl-writes / pl-reads loop.\n\n### 9.1 `joint_distribution_ppc` \u2014 fixed `compare_joint_distribution`\n\n#### (a) Responsibility\nPosterior predictive check: does the model's clone\u00d7phenotype distribution match the empirical per-clone phenotype frequencies, per covariate?\n\n#### (b) Math\nPer covariate $m$ and clone $c$:\n$$P_{\\text{model}}(\\phi\\mid c,m)=\\texttt{joint\\_distribution}(adata,\\ covariate{=}m)[c],\\qquad\nP_{\\text{emp}}(\\phi\\mid c,m)=\\frac{\\#\\{i\\in c,m:\\ \\text{pheno}_i=\\phi\\}}{\\#\\{i\\in c,m\\}}$$\n(empirical from `obs.groupby(clone)[phenotype].value_counts(normalize=True).unstack`). Per-clone discrepancy:\n$$\\delta_c=\\ \\mathrm{L1}\\big(P_{\\text{model}}[c],P_{\\text{emp}}[c]\\big)\\ \\text{or}\\ \\mathrm{KL}\\big(P_{\\text{emp}}[c]\\Vert P_{\\text{model}}[c]\\big),$$\nplus an aggregate (mean/median $\\delta$, correlation) per covariate.\n**Bug fix:** the current plot reads an undefined global `model` (\u2192 `NameError`); target reads `clonotype_col`/`phenotype_col` from `adata.uns[\"tcri_metadata\"]`.\n\n#### (c) Arguments \u2192 math\n`temperature` tempers $P_{\\text{model}}$; `covariate`/`groupby` scope which strata are compared; `distance_metric` selects L1/KL for $\\delta$.\n\n#### (d) Return\nTidy DataFrame, row per (covariate, clone): `p_model_*`, `p_emp_*`, `divergence`; plus a per-covariate summary frame.\n\n### 9.2 `phenotype_calibration`\n- **Responsibility / math:** reliability of `predict()` probabilities. Bin cells by predicted max-prob; per bin compare mean predicted probability to empirical accuracy vs the hard label; expected calibration error $\\text{ECE}=\\sum_b \\frac{n_b}{N}\\,|\\text{acc}_b-\\text{conf}_b|$.\n- **Return:** DataFrame (`bin, mean_pred, emp_freq, count`) + scalar `ECE`.\n\n### 9.3 `reconstruction_ppc`\n- **Responsibility / math:** simulate counts from the fitted ZINB decoder ($\\mu,\\theta,\\pi_{\\text{dropout}}$), compare summary statistics \u2014 library size, per-gene dropout rate, mean\u2013variance relation \u2014 between simulated and observed data.\n- **Return:** DataFrame of statistic \u00d7 {observed, simulated, discrepancy}.\n\n### 9.4 `entropy_permutation_null` / `mi_permutation_null`\n- **Responsibility / math:** significance of an observed entropy/MI against a label-permutation null. Permute phenotype labels within each covariate $R$ times, recompute the metric to form a null distribution; report\n$$p=\\text{mean}\\big(\\text{null}\\ \\ge\\ \\text{obs}\\big),\\qquad z=\\frac{\\text{obs}-\\overline{\\text{null}}}{\\text{sd(null)}}.$$\n- **Arguments:** `n_perm` (null size), `seed`, `metric` (entropy/MI), `covariate`/`groupby` scope.\n- **Return:** DataFrame per (stratum): `observed, null_mean, null_sd, z, p`.\n\n*(Also relocated into `diag` as data/diagnostic renderers, not metric math: training-loss curves `plot_loss`, archetype heatmap `plot_archetypes`, model PGM `build_nested_tcri_pgm`.)*\n\n---\n\n## Appendix \u2014 Current \u2192 target math deltas (what changes and why)\n\n| # | Function | Current math | Target math | Rationale |\n|---|---|---|---|---|\n| 1 | engine `n_samples=0` | `joint_distribution_posterior` always draws 1 Dirichlet sample (stochastic) | posterior **mean** $m$, no draw | closed-form, reproducible (\u00a70.3) |\n| 2 | `mutual_information`/`flux` `n_samples=0` | returns one random draw | deterministic point estimate | fixes latent bug (\u00a70.4) |\n| 3 | `posterior=` semantics | conflates draw-vs-mean **and** logit-folding; `posterior=False` in MI raises `NotImplementedError` | axis means *fold per-cell logits* only; both branches use the posterior (never generative prior) | \u00a70.3 |\n| 4 | metric \u2194 model agreement | metrics use a Dirichlet **draw** of $p_{ct}$ and **never** apply the gate | `posterior=True` uses the same gate-aware, mean-prior rule as `predict()` | removes silent disagreement |\n| 5 | `weighted` | buggy weight key; whole-table-sum-1 vs un-normalized inconsistency between the two engines | removed; uniform-clonotype prior documented | \u00a70.6 |\n| 6 | `temperature` | applied at different stages in the two engines | one consistent power-temper of the (combined) logit | \u00a71 |\n| 7 | phenotypic entropy zero clone | zero-mass clone \u2192 uniform \u2192 $H=1.0$ | zero-support clone \u2192 NaN/excluded | \u00a73 |\n| 8 | flux `seed` | seeds NumPy only; torch draws unaffected | seed torch Generator (+NumPy) | reproducibility |\n| 9 | KL base | natural log in flux; $\\log_2$ in entropies/MI | one base library-wide (recommend bits) | unit consistency \u00a76 |\n| 10 | \"HDI\" | equal-tailed percentiles labeled HDI | rename to CI, or implement true HDI | \u00a77 |\n| 11 | joint-distribution PPC | references undefined global `model` \u2192 `NameError` | read cols from `uns[\"tcri_metadata\"]` | \u00a79.1 |", "plan": "I have everything I need. Here is the implementation plan.\n\n---\n\n# TCRI Refactor \u2014 Implementation Plan\n\n**From** the current five-monolith package **to** the settled Door\u2011A, scverse\u2011ecosystem target (grafiti\u2011mirrored, one\u2011file\u2011per\u2011topic).\n\n**Source of truth.** The *SETTLED DESIGN* block is authoritative. Where it diverges from `docs/contract/tcri_function_inventory.md` \u00a73/\u00a79 (engine lives in `tl` not `pp`; `groupby` subsumes all `*_table`/`*_compare`/`*_delta`; `flux \u2192 phenotypic_flux`; `register_model \u2192 model.to_anndata`; `tl` package mirrors grafiti's `tools/`), this plan follows the settled design and flags the divergence inline. The inventory remains the canonical per\u2011function disposition list.\n\n**How to read.** Section 1 fixes the invariants and resolves the one open statistical question so the engine/metric phases are unblocked. Section 2 is the target tree. Sections 3\u20134 are the rename map and shared\u2011helper extraction (they cut across every later phase). Section 5 is the ordered PR sequence \u2014 each PR is independently landable, lowest\u2011risk first, green CI at every step. Sections 6\u20137 are the testing strategy and the ordering\u2011hazard graph.\n\n---\n\n## 1. Invariants and the resolved open question\n\n### 1.1 Invariants held at every PR boundary\n- **`import tcri` stays green** and the public handles `tl / pp / pl / ml / diag / ut` remain importable. No PR is allowed to leave `main` with a broken import or red CI.\n- **One behavior change per PR.** Mechanical moves (file splits, helper extraction) never change numerics; numeric changes (engine, metrics) never also move files. This keeps every diff reviewable and every regression bisectable.\n- **`_keys.py` is the only place a key string is written**, from Phase 1 onward. No new string literal for a `uns/obsm/obs` key may be introduced after that.\n- **No `import *`** is added; each module split lands with an explicit `__all__` and named re\u2011exports (the top\u2011level `__init__` flip is the last PR).\n\n### 1.2 RESOLVED \u2014 prior vs. posterior\u2011mean (the parked open question)\n\nThe inference\u2011path grounding settles this decisively; the plan adopts it and **supersedes the \"PARKED\" note in inventory \u00a710**:\n\n- **`n_samples=0` is the closed\u2011form variational posterior mean** `E_q[p_ct] = q_p_ct_raw / \u03a3 = adata.uns[\"tcri_p_ct\"]` (temperature\u2011adjusted if `temperature\u22601`). It is exact, reproducible, and requires **no sampling** \u2014 drawing Dirichlet samples and averaging converges to the *same* vector but only up to Monte\u2011Carlo noise, so it is a strictly worse estimator of a quantity available in closed form.\n- **`n_samples=N>0` is `N` i.i.d. draws** from `Dirichlet(local_scale \u00b7 p_ct_mean)` (torch), the reconstructed guide posterior `q(p_ct)`.\n- **The prior path is dropped from the public metric/engine API.** `clone_phen_prior` / the archetype mixture are the guide's *initialization / generative anchor*, not what training learned \u2014 the wrong object for a metric. The **MAP/mode is rejected**: `\u03b1_k = local_scale\u00b7m_k` is routinely `< 1`, so the mode sits on the simplex boundary / is undefined.\n- **Consequence for the API:** **delete `point_estimate=` and delete the public `posterior=` axis.** `n_samples` is the single point/draws knob. Any prior\u2011predictive comparison (\"does the learned posterior move off the prior?\") is a **`diag/` PPC**, not a metric argument. The current `posterior=False` branch (which raises `NotImplementedError` in MI and `TypeError` in flux) is removed rather than finished.\n\nThis also fixes the two cross\u2011function inconsistencies noted in grounding: the single unified engine applies `temperature` **and** the classifier\u2011logit combination **and** `gate_prob` identically to `TCRIModel.predict`, so metric per\u2011cell probabilities can no longer disagree with the model's own.\n\n---\n\n## 2. Target layout (grafiti\u2011mirrored, authoritative)\n\n```\ntcri/\n __init__.py # explicit re-export; sys.modules aliases tl/pp/pl/ml/diag/ut;\n # top-level tcri.joint_distribution; NO import *\n _keys.py # single source of every uns/obsm/obs key string (adopt in Phase 1)\n _console.py # leveled, silenceable logging via scanpy logging; NO ANSI, NO _ascii_hist\n _stats.py # stars, auc_and_label_permutation, bootstrap_auc,\n # + posterior-comparison: mann_whitney, prob_gt_zero, hdi\n _distance.py # kl_divergence, l1_distance, phenotype_distance dispatcher\n model/ # ml\n _model.py # TCRIModel: setup_anndata, train, get_latent_representation,\n # predict (was get_cell_phenotype_probs), get_p_ct, to_anndata\n _module.py # TCRIModule (pyro model/guide)\n _priors.py # MixtureDirichlet, VampPrior\n _classifier.py # PhenotypeClassifier\n _training.py # UnifiedTrainingPlan, build_archetypes\n preprocessing/ # pp (shrinks to registration-glue + clone utilities)\n _clones.py # group_singletons, clone_size\n tools/ # tl (metrics + the engine; mirrors grafiti tools/)\n _joint.py # joint_distribution (the engine; re-exported as tcri.joint_distribution)\n _entropy.py # clonotypic_entropy, phenotypic_entropy\n _mutual_information.py # mutual_information (+ private _mi_from_joint)\n _flux.py # phenotypic_flux (was flux; cov_from / cov_to)\n plotting/ # pl (twins mirror tl by filename)\n _base.py # _metric_boxplot (was tcri_boxplot), _finish\n _colors.py # tcri_colors, resolve_palette (was set_color_palette)\n _entropy.py # clonotypic_entropy [FIX], phenotypic_entropy [FIX]\n _mutual_information.py # mutual_information [FIX]\n _flux.py # phenotypic_flux (sankey)\n _sankey.py # SankeyNode, _phenotype_mass_per_clone\n diagnostics/ # diag (NEW \u2014 PPCs + model validation; returns DataFrames)\n _ppc.py # joint-distribution PPC (fixed compare_joint_distribution),\n # calibration, reconstruction PPC, permutation-null\n _training.py # loss curves (was plot_loss), archetypes (was plot_archetypes)\n utils/ # ut\n _session.py # save/load_tcri_session, _to_jsonable (plain h5ad; no manager hack)\nexamples/ # top_clone_umap, clone_size_umap, phenotype_probabilities,\n # compare_phenotypes + rewritten notebooks\ndocs/ # model PGM (build_nested_tcri_pgm) lives here, out of the package\n```\n\nDivergences from inventory \u00a73, called out: **(a)** the `tl` package is `tools/` (grafiti parity), not `metrics/`; **(b)** the engine is `tools/_joint.py`, not `preprocessing/_engine.py`; **(c)** `preprocessing/` loses `_register.py`/`_engine.py` \u2014 registration collapses into `model.to_anndata`, leaving `pp` = `_clones.py`; **(d)** no `_tables.py` \u2014 `groupby` subsumes it.\n\n---\n\n## 3. Rename\u2011map application\n\nFreeze the map **before** Phase 5 (the first breaking PR). Renames are breaking; pre\u20111.0 we pay once. Full table lives in inventory \u00a79; the deltas this plan enforces:\n\n| current | \u2192 target | lands in |\n|---|---|---|\n| `metrics/` package | `tools/` package (aliased `tl`) | Phase 5 |\n| `joint_distribution` + `joint_distribution_posterior` | `tl.joint_distribution` (engine, function) \u2192 top\u2011level `tcri.joint_distribution` | Phase 5 |\n| `flux` / `from_this` / `to_that` | `phenotypic_flux` / `cov_from` / `cov_to` | Phase 6 |\n| `clonotypic_entropy_base`, `clonality` | merged into `clonotypic_entropy` | Phase 6 |\n| `mi_compare`, `delta_entropy_table`, `flux_table`, `delta_clonotypic_entropy` | **deleted** \u2014 expressed via `groupby=` + `_stats` | Phase 6 |\n| `point_estimate=`, public `posterior=` | **removed** \u2014 `n_samples` is the only knob (see \u00a71.2) | Phase 5\u20136 |\n| `get_cell_phenotype_probs` | `predict` (scvi/CellAssign idiom) | Phase 4 |\n| `register_model` (+ `classify_phenotypes`, `register_*_key`) | `model.to_anndata` (thin) | Phase 4 |\n| `pl.clonotypic_entropy_by_phenotype` | `pl.clonotypic_entropy` | Phase 7 |\n| `plot_pheno_sankey` / `plot_phenotype_probabilities` | `phenotypic_flux` (sankey) / `phenotype_probabilities` (\u2192 examples) | Phase 7 |\n| `tcri_boxplot` | `_metric_boxplot` (private) | Phase 7 |\n| `set_color_palette` | `resolve_palette` | Phase 7 |\n| `dkl`, `flux.dkl_func` | `_distance.kl_divergence` | Phase 1 |\n| `\u0394` (unicode) | `delta` (ASCII, greppable) | Phase 6 |\n| `c2p_mat` | `clone_phenotype_prior` | Phase 3 |\n| `centropy` / `pentropy` / `*_tl` leaked aliases | removed via `__all__` | Phase 11 |\n| `uns[\"tcri_clone_key\"]` / `[\"tcri_phenotype_key\"]` **and** `uns[\"tcri_metadata\"][...]` (two schemes) | one scheme via `_keys.py` (single `tcri_metadata`) | Phase 1 |\n\n---\n\n## 4. Shared\u2011helper extraction (`_keys` / `_console` / `_stats` / `_distance`)\n\nLands as **Phase 1**, before any file move, because every later phase imports these. Pure internal dedup; public API unchanged.\n\n| new module | absorbs | notes |\n|---|---|---|\n| `tcri/_keys.py` | every `uns/obsm/obs` key literal, both current schemes | constants only; unifies `tcri_clone_key`/`tcri_metadata['clone_col']` to one `tcri_metadata` scheme; standard scvi `REGISTRY_KEYS` for registry fields |\n| `tcri/_console.py` | triplicated `_ok/_info/_warn/_fin` (metrics/pp/pl copies) | reimplement over **scanpy's `logging`/verbosity** for ecosystem consistency; drop raw ANSI prints; **drop `_ascii_hist` and every `graph=`/ASCII\u2011histogram code path** |\n| `tcri/_stats.py` | `stars`, `auc_and_label_permutation`, `bootstrap_auc` (from `utils/_utils.py`) | **add** the general posterior\u2011comparison stats used to replace `*_delta`: `mann_whitney`, `prob_gt_zero` (Bayesian P(>0)), `hdi` |\n| `tcri/_distance.py` | module\u2011level dead `dkl` + `flux.dkl_func` | one `kl_divergence` + `l1_distance` + `phenotype_distance(metric=...)` dispatcher; **fix the mixed\u2011units bug** \u2014 standardize on `log2` (bits) across entropy/MI/KL, or document nats explicitly for KL |\n\n`_mi_from_joint` stays module\u2011private in `tools/_mutual_information.py` (single source, not shared).\n\n---\n\n## 5. Ordered PR sequence\n\nEach PR is independently landable with green CI. Risk and dependencies are explicit.\n\n### Phase 0 \u2014 Contract freeze + CI scaffolding *(docs/tests only; zero code change)*\n- Write the target signatures as a frozen contract: a `tcri/_contract.pyi` + a `tests/test_contract_conformance.py`, ported from grafiti's `_pyi_gen.py` / `test_contract_conformance.py`. Seed `IMPLEMENTED = {}`; every target function is on the worklist and flips to implemented as its phase lands. Drift at the markdown\u2192`.pyi`\u2192live\u2011signature seam fails CI.\n- Land the rename map (\u00a73) and this document into `docs/contract/`.\n- **Risk:** none. **Depends on:** nothing.\n\n### Phase 1 \u2014 Shared helpers + `_keys` adoption *(internal dedup; API unchanged)*\n- Create `_keys.py`, `_console.py`, `_stats.py`, `_distance.py` (\u00a74). Replace all triplicated printers, dead `dkl`/`dkl_func`, and relocate stats out of the `utils` monolith.\n- **Adopt `_keys` at every read/write site** in `model/`, `preprocessing/`, `metrics/`, `plotting/`, `utils/`. This is the single most valuable pre\u2011move safety step: the subsequent `to_anndata`/engine rewrites touch key IO, and centralizing the strings first makes those diffs local.\n- **Risk:** low (mechanical, no numerics). **Depends on:** Phase 0. **Hazard:** must precede Phase 4/5.\n\n### Phase 2 \u2014 Safe deletions *(0\u2011caller dead code)*\n- Delete the 12 confirmed 0\u2011caller / dead\u2011broken functions: `gene_entropy`, `get_latent_embedding`, `group_small_clones`, `register_probability_columns`, `remove_meaningless_genes`, `metrics._ent`, `clone_fraction`, module\u2011level `dkl`, `utils.probabilities`, `pl.polar_plot`, `pl.probability_distribution`, `pl.bayesian_mutual_information`, `SankeyNode.hex_to_rgb`.\n- **Not deleted here:** `compare_joint_distribution` (\u2192 diag, Phase 8), `pl.mutual_information` / `pl.phenotypic_entropy` / `ridge_delta_entropy` (keep+fix, Phase 7), and the consolidated\u2011away `*_table`/`*_delta`/`clonality`/`_base` (removed *with* their replacement in Phase 6, so no caller is ever orphaned).\n- **Risk:** very low. **Depends on:** Phase 1 (so deletions don't fight helper moves).\n\n### Phase 3 \u2014 Model module split *(mechanical; no behavior change)*\n- Split `model/_model.py` (1074 ln) \u2192 `_model.py` (`TCRIModel`) + `_module.py` (`TCRIModule`) + `_priors.py` (`MixtureDirichlet`, `VampPrior`) + `_classifier.py` (`PhenotypeClassifier`) + `_training.py` (`UnifiedTrainingPlan`, `build_archetypes`).\n- Rename `c2p_mat \u2192 clone_phenotype_prior`. Explicit `__all__` per module; `model/__init__` re\u2011exports named symbols.\n- **Risk:** low (import wiring). **Depends on:** Phase 1. **Verify:** `test_model_setup.py`, `test_pyro_params.py` unchanged\u2011green.\n\n### Phase 4 \u2014 Model\u2192AnnData streamline *(behavior change; kills the manager hack)*\n- `setup_anndata` becomes **registration\u2011only** \u2014 no `obs` mutation, no manager stashed in `uns`. `group_singletons` stays a separate `pp` step.\n- `register_model` \u2192 **`model.to_anndata(adata)`** (thin): writes only the canonical minimum via `_keys` \u2014 metadata/categories from the registry, `X_tcri` latent, `predict()` phenotype probs + hard labels, `p_ct`. **Stops writing `uns['tcri_manager']`**, which **deletes the `write_adata_safely` / `_pop_nonserializables` hack**.\n- Rename `get_cell_phenotype_probs \u2192 predict` (returns per\u2011cell phenotype\u2011prob DataFrame). Keep the gate\u2011aware/additive combination as the single ground\u2011truth cell probability.\n- Rework `utils/_session.py`: `write_adata_safely \u2192 save_tcri_session` writes a plain h5ad (nothing non\u2011picklable in `uns` anymore); load rebuilds the registry by re\u2011running `setup_anndata`.\n- **Risk:** HIGH \u2014 touches train\u2192register\u2192session round\u2011trip. **Depends on:** Phase 1 (`_keys`), Phase 3 (split). **Hazard:** `test_session_round_trip.py` is the gate \u2014 update it in\u2011PR; it must prove a saved/loaded session reproduces `p_ct`, latent, and probs without the manager.\n\n### Phase 5 \u2014 Engine consolidation *(the substrate; resolves \u00a71.2)*\n- Create `tools/` (aliased `tl`). Add **`tools/_joint.py::joint_distribution`** unifying `joint_distribution` + `joint_distribution_posterior` into one function; re\u2011export as **`tcri.joint_distribution`** (top\u2011level prominence).\n- Signature: `joint_distribution(adata, *, covariate=None, groupby=None, n_samples=0, clones=None, temperature=1.0)` \u2192 tidy `DataFrame` (clone \u00d7 phenotype; covariate/sample axis when relevant), provenance in `.attrs`.\n- Implement \u00a71.2: `n_samples=0` \u2192 closed\u2011form posterior mean `uns[K.P_CT]` (temperature\u2011adjusted); `n_samples>0` \u2192 torch `Dirichlet(local_scale\u00b7p_ct_mean)` draws. **Drop `posterior=`/`point_estimate=`/`combine_with_logits` from the public surface.** Apply `temperature`, the classifier\u2011logit combine, and `gate_prob` **identically to `model.predict`** so metrics and model agree.\n- **Bug fixes folded in:** the weighting bug (weights keyed by clone index instead of the `ct` index \u2192 arbitrary weights); the inconsistent whole\u2011table\u2011vs\u2011row normalization between the two old functions; the all\u2011zero\u2011clone reindex that inflates entropy to uniform; **seed reproducibility** (seed torch, not just numpy \u2014 `np.random.seed` never touched the Dirichlet draws).\n- **Risk:** HIGH \u2014 every metric reads this. **Depends on:** Phase 4 (reads the keys/logits `to_anndata` writes). **Verify:** rewrite `test_preprocessing/test_joint_distribution_posterior.py` \u2192 `test_tools/test_joint.py` asserting `n_samples=0` == `uns[K.P_CT]` exactly, and that repeated `n_samples=0` calls are bit\u2011identical.\n\n### Phase 6 \u2014 Metric\u2011API consolidation *(four tl metrics; `groupby` subsumes tables/compare/delta)*\n- Populate `tools/_entropy.py`, `tools/_mutual_information.py`, `tools/_flux.py`. Uniform signature across all four:\n ```\n tl.clonotypic_entropy (adata_or_jd, *, covariate=None, groupby=None, n_samples=0, temperature=1.0, clones=None, normalised=True)\n tl.phenotypic_entropy (adata_or_jd, *, covariate=None, groupby=None, n_samples=0, temperature=1.0, clones=None, normalised=True)\n tl.mutual_information (adata_or_jd, *, covariate=None, groupby=None, n_samples=0, temperature=1.0, clones=None, normalised=True, normalise_mode='average')\n tl.phenotypic_flux (adata, *, cov_from, cov_to, groupby=None, n_samples=0, temperature=1.0, clones=None, distance_metric='l1')\n ```\n- **Dual input:** each accepts `(adata + covariate)` [compute joint internally] **or** a precomputed joint `DataFrame` [fast path].\n- **Return\u2011shape rule (uniform):** no `groupby` & `n_samples=0` \u2192 scalar/Series; `groupby` set \u2192 tidy `DataFrame` (row per group [\u00d7 phenotype/clone]); `n_samples>0` \u2192 adds a sample axis / summary columns (`mean`, `hdi_low`, `hdi_high`, `p_gt`).\n- **Comparisons via `groupby` + `_stats`, not delta functions.** **Delete** `mi_compare`, `delta_entropy_table`, `flux_table`, `delta_clonotypic_entropy`, `clonotypic_entropy_base`, `clonality`, `metrics.dkl`. `delta`\u2011style analyses become `groupby` results fed to `_stats.mann_whitney`/`prob_gt_zero`/`hdi`.\n- Delete `metrics/` after migration. Design each function to return a clean tidy result so the deferred `@tl_result` cache decorator is a one\u2011line wrap later (build\u2011toward, not implemented).\n- **Risk:** HIGH. **Depends on:** Phase 5. **Verify:** rewrite `test_metrics/` \u2192 `test_tools/`; assert normalized MI \u2208 [0,1], entropies \u2208 [0,1], `n_samples=0` deterministic, `groupby` tidy shape, dual\u2011input equivalence (adata path == precomputed\u2011joint path).\n\n### Phase 7 \u2014 Plotting split + pl twins *(fix the broken core plots)*\n- Split `plotting/_plotting.py` (1437 ln) \u2192 `_base.py` (`_metric_boxplot`, `_finish`), `_colors.py` (`resolve_palette`), `_entropy.py`, `_mutual_information.py`, `_flux.py`, `_sankey.py`.\n- Ship the four tl\u2194pl twins (filename + function name match): `pl.clonotypic_entropy` (was `_by_phenotype`), `pl.phenotypic_entropy` **[FIX]**, `pl.mutual_information` **[FIX]**, `pl.phenotypic_flux` (sankey). Consolidate `SankeyNode` + `_phenotype_mass_per_clone` behind `pl.phenotypic_flux`.\n- pl functions are **cache renderers**: read the tidy tl result (or recompute via the engine on demand until `@tl_result`/`get.py` lands). No metric math in `pl`.\n- **Move to `examples/`:** `top_clone_umap`, `clone_size_umap`, `plot_phenotype_probabilities`, `compare_phenotypes`. **Drop:** `clonality` plot, `flux` boxplot.\n- **Risk:** medium. **Depends on:** Phase 6 (tidy results), Phase 1 (`_colors`/`_console`). **Verify:** `test_plotting/` returns a `Figure`/`Axes` without error for each twin (smoke\u2011level).\n\n### Phase 8 \u2014 `diag/` seeding *(new; additive)*\n- `diag/_ppc.py`: joint\u2011distribution PPC (the **fixed** `compare_joint_distribution` \u2014 model `p(clone,phenotype)` vs empirical counts, no undefined\u2011global `NameError`); phenotype\u2011probability calibration; reconstruction PPC (ZINB simulate \u2192 compare library/dropout/mean\u2011var); entropy/MI vs permutation null. **All return DataFrames** (read\u2011only on the finalized model, outside the tl\u2011writes/pl\u2011reads loop).\n- `diag/_training.py`: `plot_loss \u2192 loss curves`, `plot_archetypes \u2192 archetypes` relocated here (off `TCRIModel`).\n- **Risk:** low\u2011medium (new surface). **Depends on:** Phase 4 (finalized model), Phase 5 (engine).\n\n### Phase 9 \u2014 PGM \u2192 docs; utils finalize\n- Move `build_nested_tcri_pgm` / `draw_tcri_pgm_nested` **out of the package into `docs/`** (a docs script/notebook). Drop `daft` from runtime deps \u2192 docs extras only.\n- `utils/_utils.py` \u2192 `utils/_session.py` (session\u2011io + `_to_jsonable` only; stats/console/distance already extracted in Phase 1).\n- **Risk:** low. **Depends on:** Phase 1, Phase 8.\n\n### Phase 10 \u2014 Notebook + examples rewrite\n- Rewrite the `example/` notebooks (`smith`, `renal`, `yost`, `zhang`, `synthetic`, `colitis`) to the new API: `setup_anndata \u2192 TCRIModel \u2192 train \u2192 model.to_anndata`; `tcri.joint_distribution`; the four metrics with `groupby`/`n_samples`; the four pl twins; `diag` checks. One canonical end\u2011to\u2011end tutorial notebook under `docs/`.\n- Move bespoke one\u2011off plots into `examples/`.\n- **Risk:** low (docs), but this is the real end\u2011to\u2011end acceptance test of the whole API. **Depends on:** Phases 4\u20138.\n\n### Phase 11 \u2014 Public API finalize + scverse ecosystem CI\n- `tcri/__init__.py`: **explicit** named re\u2011exports, `sys.modules` aliases for `tl/pp/pl/ml/diag/ut`, top\u2011level `tcri.joint_distribution`, **remove `import *`**, kill leaked aliases (`centropy`/`pentropy`/`*_tl`) via `__all__`.\n- Flip **all** target functions to `IMPLEMENTED` in the contract test; drift now hard\u2011fails CI.\n- Turn on the full scverse\u2011ecosystem gate (\u00a76.2).\n- **Risk:** low\u2011medium (surface finalization). **Depends on:** all prior phases.\n\n---\n\n## 6. Testing strategy\n\n### 6.1 Tests that must exist / pass, per phase\n\n| phase | required tests |\n|---|---|\n| 0 | `test_contract_conformance` (markdown\u2192`.pyi`\u2192signature drift); import\u2011smoke on py3.10/3.11 |\n| 1 | unit tests for `_stats` (`stars`, AUC/permutation, `hdi`, `prob_gt_zero`), `_distance` (`kl_divergence` symmetry/units), `_console` (silence flag); a `_keys` \"no stray literal\" grep test |\n| 2 | regression: every deleted symbol is absent from `__all__` **and** unreferenced (import\u2011graph test) |\n| 3 | `test_model_setup`, `test_pyro_params` unchanged\u2011green through the split; `import tcri.model._{priors,classifier,module,training}` smoke |\n| 4 | **`test_session_round_trip` rewritten** \u2014 save/load reproduces `p_ct` + latent + `predict` probs with **no `tcri_manager` in `uns`**; `setup_anndata` leaves `obs` untouched; `to_anndata` writes exactly the canonical key set |\n| 5 | `test_tools/test_joint`: `n_samples=0 == uns[K.P_CT]` exactly; repeated `n_samples=0` bit\u2011identical (determinism); `n_samples>0` seeded\u2011reproducible (torch); weighting keyed on `ct` index; tidy `.attrs` provenance |\n| 6 | `test_tools/test_{entropy,mutual_information,flux}`: ranges (`[0,1]` normalized), `n_samples=0` determinism, `groupby` tidy shape, **dual\u2011input equivalence** (adata path == precomputed\u2011joint path), all\u2011zero\u2011clone \u2192 NaN (not uniform), `phenotypic_flux` `cov_from/cov_to` + seeded draws |\n| 7 | `test_plotting`: each pl twin returns `Figure`/`Axes` from a tidy tl result without touching metric math; sankey renders |\n| 8 | `test_diag`: each PPC returns a DataFrame of expected columns on the `trained_model` fixture; permutation\u2011null is seeded |\n| 9 | no `daft`/PGM import remains in the installed package (import\u2011graph test) |\n| 10 | **notebook execution** (`nbmake`/`pytest --nbmake`) on the synthetic tutorial notebook end\u2011to\u2011end |\n| 11 | contract test with **all** functions `IMPLEMENTED`; `__all__` completeness (public names \u2194 contract); \"no `import *`\" AST test |\n\nReuse the existing `conftest.py` fixtures (`synthetic_adata`, `trained_model`, `mock_adata`); extend `mock_adata` to the unified `_keys` scheme in Phase 1 and to `to_anndata`'s canonical key set in Phase 4.\n\n### 6.2 scverse\u2011ecosystem CI requirement\n\nBring `.github/workflows/tests.yml` up to the cookiecutter\u2011scverse bar, layered so each phase keeps CI green:\n- **Matrix** py3.10 + py3.11 (extend to 3.12 before 1.0), `pip install -e \".[test]\"`, `pytest tests/ -v --cov=tcri`.\n- **Import\u2011smoke job** (grafiti pattern): `python -c \"import tcri; from tcri.model._model import TCRIModel; from tcri.tools._joint import joint_distribution\"`.\n- **Lint/format gate:** `pre-commit` (ruff + ruff\u2011format) \u2014 add in Phase 1, enforce from Phase 3.\n- **Contract\u2011conformance job** (Phase 0 onward) \u2014 the markdown\u2192`.pyi`\u2192signature guardrail that lets views land independently.\n- **Notebook\u2011execution job** (`nbmake`) from Phase 10.\n- **Docs build** (sphinx + numpydoc, readthedocs `.readthedocs.yaml` already present) must pass; public API is fully docstringed/typed.\n- **Ecosystem checklist** (final gate, Phase 11): AnnData\u2011native `setup_anndata`, no `import *`, typed public surface, tutorial notebook, `diag` returns data not plots \u2014 the scverse ecosystem\u2011package acceptance criteria.\n\n---\n\n## 7. Dependencies & ordering hazards\n\n```\n0 \u2500\u25b6 1 \u2500\u252c\u2500\u25b6 2\n \u251c\u2500\u25b6 3 \u2500\u25b6 4 \u2500\u25b6 5 \u2500\u25b6 6 \u2500\u25b6 7 \u2500\u25b6 8 \u2500\u25b6 9 \u2500\u25b6 10 \u2500\u25b6 11\n \u2514\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500(1 gates 4,5,6)\n```\n\nHard ordering constraints and the reason each is load\u2011bearing:\n\n1. **`_keys` (1) before `to_anndata` (4) and the engine (5).** Both rewrite key IO; centralizing the strings first keeps those diffs local and prevents a two\u2011scheme drift (`tcri_clone_key` vs `tcri_metadata['clone_col']`) from leaking into the new code.\n2. **Model split (3) before `to_anndata` (4).** `to_anndata` is a `TCRIModel` method; splitting the 1074\u2011line module first keeps the behavior\u2011change diff readable.\n3. **`to_anndata` (4) before the engine (5).** The engine reads exactly the canonical keys/logits `to_anndata` writes; if the write set changes after the engine is built, the engine breaks silently. The **manager\u2011in\u2011`uns` removal is the single highest\u2011risk hazard** \u2014 it changes the session round\u2011trip; `test_session_round_trip` must be rewritten in the same PR.\n4. **Engine (5) before metrics (6).** Metrics are thin readers of the joint; consolidating them against a pre\u2011fix engine would bake in the weighting/normalization/seed bugs.\n5. **Metrics (6) before pl twins (7).** pl are cache renderers of tidy tl results.\n6. **Delete `*_table`/`*_delta`/`clonality`/`_base` *with* their replacement (6), never before.** These have live callers (notebooks, `pl.mi_compare`); removing them ahead of the `groupby`/`_stats` replacement would orphan callers and red the notebook\u2011execution job.\n7. **Rename freeze before Phase 5.** `from_this/to_that \u2192 cov_from/cov_to`, `flux \u2192 phenotypic_flux`, `metrics/ \u2192 tools/`, dropping `posterior=`/`point_estimate=` are all breaking; batch them at the first breaking PR so downstream notebooks are rewritten once (Phase 10), not repeatedly.\n8. **Deferred, but designed\u2011for:** `@tl_result` uns\u2011cache + `get.py`. Every tl function in Phase 5\u20136 must return a clean tidy result so the decorator is a one\u2011line wrap later; until then pl recomputes via the engine. Do **not** block this refactor on the cache.\n\n**Independent / parallelizable:** Phase 2 (safe deletions) and Phase 8 (`diag` seeding, once 4\u20135 land) can proceed alongside their neighbors without contending for the same files.\n\n---\n\n## 8. Rollout checklist (per PR)\n\n- [ ] One behavior change (or zero, if a mechanical move).\n- [ ] `import tcri` green; `tl/pp/pl/ml/diag/ut` handles intact.\n- [ ] No new key\u2011string literal outside `_keys.py`; no new `import *`.\n- [ ] Deleted symbols removed from `__all__` and unreferenced.\n- [ ] Contract test updated (`IMPLEMENTED` flipped for functions this PR lands).\n- [ ] Phase\u2011specific tests from \u00a76.1 present and green on py3.10/3.11.\n- [ ] Rename\u2011map rows for this phase applied everywhere (code + fixtures; notebooks batched to Phase 10)."}} \ No newline at end of file diff --git a/docs/contract/build_tcri_contract.py b/docs/contract/build_tcri_contract.py new file mode 100644 index 0000000..2704d5f --- /dev/null +++ b/docs/contract/build_tcri_contract.py @@ -0,0 +1,551 @@ +#!/usr/bin/env python3 +"""Build the TCRI target API contract — emits .md (source) + color HTML. + +Single source of truth for the *ideal* post-refactor API. Renders a compact +summary table, per-function detail cards (color-coded by namespace), the +adata-state schema (the data contract everything couples through), the shared +primitives index, and the dropped/merged list. Work the refactor toward this; +do not drift. + +Companion docs: + - tcri_api_refactor.md (prose design spec + sequencing) [optional] + - tcri_dependency_map.* (call + producer/consumer graph) [build_tcri_depgraph.py] +""" +import html as _html +import os + +_HERE = os.path.dirname(os.path.abspath(__file__)) +OUT_MD = os.path.join(_HERE, "tcri_api_contract.md") +OUT_HTML = os.path.join(_HERE, "tcri_api_contract.html") + +# ---- namespaces (color-coded) ---------------------------------------------- +VIEWS = { + "ml": ("#dc2626", "Model (ml)"), + "pp": ("#0f766e", "Preprocess (pp)"), + "tl": ("#7c3aed", "Metrics (tl)"), + "pl": ("#2563eb", "Plotting (pl)"), + "ut": ("#475569", "Session / IO (ut)"), +} + +# ---- function inventory ----------------------------------------------------- +# fields: ns, group, name, status, signature, returns, writes, reads, calls, +# upstream (run-first), invariants, edges, mirror, source (provenance), notes +F = [] +def fn(**k): F.append(k) + +# ============================ ml — model ==================================== +fn(ns="ml", group="model", name="TCRIModel.setup_anndata", status="keep", + signature="(adata, *, layer=None, clonotype_key='unique_clone_id', " + "phenotype_key='phenotype_col', covariate_key='timepoint', batch_key='patient')", + returns="adata (registers scvi AnnDataManager)", + writes="uns[tcri_manager]; registry{clonotype/phenotype/covariate/batch_col}", + reads="obs[clonotype/phenotype/covariate/batch_key]", + calls="—", upstream="—", + invariants="all four obs columns must exist → ValueError otherwise", + edges="layer=None → counts from X", mirror="—", + source="keep as-is", notes="scvi classmethod entry point") +fn(ns="ml", group="model", name="TCRIModel.train", status="keep", + signature="(max_epochs=1000, batch_size=1000, lr=1e-3, " + "reconstruction_loss_scale=1e-3, n_steps_kl_warmup=2000, **kw)", + returns="None (fits in place; populates history_)", + writes="pyro param store (q_p_c_raw, q_p_ct_raw); model weights", + reads="adata_manager", calls="UnifiedTrainingPlan, TrainRunner, DataSplitter", + upstream="setup_anndata", + invariants="early-stops on elbo_validation", edges="—", mirror="pl.model_loss", + source="keep as-is", notes="") +fn(ns="ml", group="model", name="TCRIModel.get_latent_representation", status="keep", + signature="(adata=None, indices=None, batch_size=None)", + returns="ndarray[N, n_latent] — posterior-mean z", + writes="—", reads="encoder", calls="module.get_latent", upstream="train", + invariants="deterministic (mean, no sampling)", edges="—", mirror="—", + source="keep as-is", notes="") +fn(ns="ml", group="model", name="TCRIModel.get_cell_phenotype_probs", status="keep — now the ONLY prob path", + signature="(adata=None, batch_size=256, eps=1e-8)", + returns="ndarray[N, P] — per-cell phenotype probabilities", + writes="—", reads="encoder, classifier, get_p_ct(), module.ct_array", + calls="module.get_p_ct", upstream="train", + invariants="softmax(gate·logits + (1-gate)·log prior) OR additive when no gate", + edges="use_gate toggles the two combination rules", + mirror="—", + source="canonical; absorbs pp._compute_logits_and_prior + pp.classify_phenotypes", + notes="register_model now calls THIS instead of reimplementing the softmax") +fn(ns="ml", group="model", name="TCRIModel.get_p_ct", status="keep", + signature="()", returns="ndarray[CT, P] — learned clone×covariate phenotype prior", + writes="—", reads="pyro param store q_p_ct_raw", calls="module.get_p_ct", + upstream="train", invariants="rows sum to 1; NaN→uniform guard", edges="—", + mirror="—", source="keep as-is", notes="") +fn(ns="ml", group="model", name="TCRIModel.boost_phenotype_prior", status="keep — advanced", + signature="(phenotype_name, boost_factor=5.0, *, affect_mixture=True)", + returns="None (mutates clone_phen_prior / mixture in place)", + writes="module.clone_phen_prior, module.mixture_concentration", + reads="c2p_mat, centers", calls="—", upstream="(before) train", + invariants="rows renormalized to 1 after boost", edges="unknown phenotype → ValueError", + mirror="—", source="keep; drop its inline _ok (use _console)", + notes="niche manual-prior knob; kept but flagged advanced") + +# ============================ pp — preprocess =============================== +fn(ns="pp", group="register", name="register_model", status="keep", + signature="(adata, model, *, latent_slot='X_tcri', batch_size=256, " + "store_logits=True, store_logposterior=True, compute_umap=False, " + "clonotype_key='trb_unique', ...)", + returns="adata (all tcri_* state written)", + writes="uns[tcri_metadata, tcri_p_ct, tcri_ct_to_cov, tcri_ct_to_c, tcri_local_scale, " + "tcri_{covariate,clonotype,phenotype}_categories, tcri_{ct,cov}_array_for_cells]; " + "obsm[X_tcri, X_tcri_logits, X_tcri_logposterior, X_tcri_probabilities]; obs[tcri_phenotype]", + reads="model.module.*, adata_manager.registry", + calls="model.get_latent_representation, model.get_cell_phenotype_probs, model.get_p_ct", + upstream="model.train", + invariants="per-cell arrays length == n_obs (guarded downstream)", + edges="compute_umap optional; prob slot only written if absent", + mirror="—", + source="keep; folds in register_phenotype_key + register_clonotype_key (metadata only)", + notes="THE bridge: model outputs → canonical adata state") +fn(ns="pp", group="engine", name="joint_distribution", status="merge — unified engine", + signature="(adata, covariate, *, posterior=True, n_samples=0, temperature=1.0, " + "clones=None, weighted=False, combine_with_logits=True, seed=None, silent=True)", + returns="DataFrame[clone × phenotype] (rows sum to 1 unless weighted); " + "n_samples>0 → stacked draws", + writes="— (pure)", reads="uns[tcri_p_ct, tcri_local_scale, tcri_*_categories, " + "tcri_metadata, tcri_{ct,cov}_array_for_cells, tcri_ct_to_*]; obsm[X_tcri_logits]", + calls="—", upstream="register_model", + invariants="posterior=True draws Dirichlet(local_scale·p_ct)+logit combine; " + "posterior=False = point-estimate prior; FAILS LOUDLY on filtered-view length mismatch", + edges="clones filter + reindex; weighted → mass-weighted, no renorm", + mirror="—", + source="MERGES joint_distribution_posterior (posterior=True) + joint_distribution (posterior=False)", + notes="the single computational core under every metric") +fn(ns="pp", group="bookkeeping", name="group_singletons", status="keep", + signature="(adata, *, clonotype_key='trb', groupby='patient', " + "target_col='trb_unique', min_clone_size=10)", + returns="None (writes obs)", writes="obs[target_col], obs[trb_candidate]", + reads="obs[clonotype_key], obs[groupby]", calls="—", upstream="—", + invariants="clones < min_clone_size → 'Singleton_'", edges="—", + mirror="—", source="keep; subsumes group_small_clones (hardcoded dup)", + notes="canonical small-clone collapse") +fn(ns="pp", group="bookkeeping", name="clone_size", status="keep", + signature="(adata, *, key_added='clone_size', return_counts=False)", + returns="None | dict{clone: size}", writes="obs[clone_size]", + reads="uns[tcri_metadata][clone_col] (was tcri_clone_key)", calls="—", upstream="register_model", + invariants="size == cells per clone", edges="—", mirror="—", + source="keep; retarget onto tcri_metadata (drop tcri_clone_key)", notes="") +fn(ns="pp", group="bookkeeping", name="filter_genes", status="rename + fix", + signature="(adata, *, mt=True, rp=True, tcr=True, hsp=True, mtrn=True, ribo=True)", + returns="adata (subset copy)", writes="— (returns new view)", + reads="var_names", calls="—", upstream="—", + invariants="each flag composes (AND), never resets the running mask", + edges="HLA-* kept despite '-'/'.'", mirror="—", + source="rename of remove_meaningless_genes + FIX flag-reset bug", + notes="TCR-gene removal pre-embedding is methodologically in-scope; generic flags fixed") + +# ============================ tl — metrics ================================== +fn(ns="tl", group="per-covariate", name="clonotypic_entropy", status="keep (+ single-phenotype mode)", + signature="(adata, covariate, *, n_samples=0, temperature=1.0, clones=None, " + "combine_with_logits=True, normalised=True, phenotype=None)", + returns="Series[phenotype] (point) | ndarray[n_samples, P] (draws) | float (phenotype=)", + writes="—", reads="uns[tcri_phenotype_categories]; (joint_distribution)", + calls="pp.joint_distribution", upstream="register_model", + invariants="H normalised by log2(n_clones); n_samples<1 → ValueError", + edges="empty joint → NaN row", mirror="pl.clonotypic_entropy", + source="keep; absorbs clonotypic_entropy_base via phenotype= arg", + notes="H[P(c|phi,m)] — phenotype spread across clones") +fn(ns="tl", group="per-covariate", name="phenotypic_entropy", status="keep", + signature="(adata, covariate, *, n_samples=0, temperature=1.0, clones=None, " + "combine_with_logits=True, normalised=True)", + returns="Series[clone] (point) | ndarray[n_samples, n_clones] (draws)", + writes="—", reads="uns[tcri_metadata]; (joint_distribution)", + calls="pp.joint_distribution", upstream="register_model", + invariants="H normalised by log2(P); n_samples<1 → ValueError", + edges="no clones at covariate → empty", mirror="pl.phenotypic_entropy", + source="keep as-is", notes="H[P(phi|c,m)] — clone phenotypic plasticity") +fn(ns="tl", group="per-covariate", name="mutual_information", status="keep", + signature="(adata, covariate, *, n_samples=0, temperature=1.0, clones=None, " + "normalised=True, normalise_mode='average', posterior=True, " + "combine_with_logits=True, verbose=True)", + returns="float (point) | ndarray[n_samples] (draws)", + writes="—", reads="(joint_distribution)", calls="pp.joint_distribution, _joint_to_mi", + upstream="register_model", + invariants="I = ΣΣ p·log2(p/(px·py)); normalised by mean/min marginal H", + edges="posterior=False path implemented (no NotImplementedError)", + mirror="pl.mutual_information", + source="keep; wire the prior path that currently raises", + notes="clone↔phenotype coupling") +fn(ns="tl", group="per-covariate", name="clonality", status="keep", + signature="(adata)", returns="dict{phenotype: clonality∈[0,1]}", + writes="—", reads="obs[tcri_phenotype], uns[tcri_metadata]", calls="—", + upstream="register_model", + invariants="1 - H(clone sizes)/log2(K); hard labels (no posterior)", + edges="single clone → 1; nan→0", mirror="pl.clonality", + source="keep; retarget onto tcri_metadata", notes="") +fn(ns="tl", group="between-covariate", name="flux", status="keep", + signature="(adata, *, cov_from, cov_to, clones=None, distance_metric='l1', " + "n_samples=0, temperature=1.0, weighted=False, posterior=True, " + "combine_with_logits=True, seed=42)", + returns="Series[clone] (point) | ndarray[n_samples, n_clones] (draws)", + writes="—", reads="(joint_distribution ×2)", calls="pp.joint_distribution, _stats.distance", + upstream="register_model", + invariants="dist over common clones at both covariates", + edges="no overlap → ValueError; metric ∈ {l1, dkl, callable}", + mirror="pl.flux", + source="keep; from_this/to_that → cov_from/cov_to; dkl via _stats registry", + notes="phenotype-distribution shift per clone") +fn(ns="tl", group="between-covariate", name="delta_clonotypic_entropy", status="keep", + signature="(adata, phenotype, *, cov_pre, cov_post, n_samples=1000, temperature=1.0, " + "clones=None, weighted=False, normalised=True, posterior=True, " + "combine_with_logits=True, seed=None)", + returns="ndarray[n_samples] — H_post − H_pre", + writes="—", reads="(clonotypic_entropy)", calls="tl.clonotypic_entropy", + upstream="register_model", + invariants="positive ⇒ entropy rose pre→post", edges="—", + mirror="pl.ridge_delta_entropy (via delta_entropy_table)", + source="keep; calls clonotypic_entropy(phenotype=)", notes="") +fn(ns="tl", group="tables", name="mi_compare", status="keep", + signature="(adata, groupby, *, groups=None, treatment=None, n_samples=50, " + "patient_col=None, clone_col=None, covariate_col=None, verbose=True)", + returns="dict{samples, summary, pairs, params}", + writes="—", reads="uns[tcri_metadata]; obs[groupby, patient_col]", + calls="tl.mutual_information (per patient×covariate)", upstream="register_model", + invariants="patient-level samples → group summary; pairs from groups", + edges="missing group/cov skipped", mirror="pl.mi_compare", + source="keep; uses shared _group_table loop", notes="patient-level MI comparison") +fn(ns="tl", group="tables", name="delta_entropy_table", status="keep", + signature="(adata, *, cov_pre, cov_post, splitby='response', n_samples=1000, " + "temperature=1.0, weighted=False, normalised=True, posterior=True, " + "combine_with_logits=True, seed=42)", + returns="DataFrame[phenotype × splitby] — delta_samples + summary stats", + writes="—", reads="uns[tcri_metadata]; obs[splitby]", + calls="tl.delta_clonotypic_entropy", upstream="register_model", + invariants="keeps full delta vector per row", edges="—", + mirror="pl.ridge_delta_entropy", + source="keep; uses shared _group_table loop", notes="") +fn(ns="tl", group="tables", name="flux_table", status="keep", + signature="(adata, *, cov_pre, cov_post, splitby='response', n_samples=0, " + "temperature=1.0, weighted=False, posterior=True, " + "combine_with_logits=True, distance_metric='l1', seed=42)", + returns="DataFrame[clone × splitby] — flux_samples, flux_mean/sd, clone_size", + writes="—", reads="uns[tcri_metadata]; obs[splitby]", calls="tl.flux", + upstream="register_model", invariants="per-group clone scoping", edges="—", + mirror="—", source="keep; uses shared _group_table loop", notes="") + +# ============================ pl — plotting ================================= +def plf(group, name, reads, viewtype, calls, status="keep", upstream="matching tl tool", + source="keep", mirror="—"): + fn(ns="pl", group=group, name=name, status=status, + signature="(adata, *, ..., palette=None, figsize=..., ax=None, save=None)", + returns="Axes | Figure | None", writes="—", reads=reads, calls=calls, + upstream=upstream, invariants="reads computed state / calls tl; never owns model math", + edges="—", mirror=mirror, source=source, notes=viewtype) + +plf("metric", "mutual_information", "(via tl.mutual_information)", "box/strip across covariate ±splitby", + "tl.mutual_information") +plf("metric", "clonotypic_entropy", "(via tl.clonotypic_entropy)", "per-phenotype box/dot ±covariate", + "tl.clonotypic_entropy", status="rename", source="rename of clonotypic_entropy_by_phenotype") +plf("metric", "phenotypic_entropy", "(via tl.phenotypic_entropy)", "box/strip per covariate", + "tl.phenotypic_entropy", status="fix", source="FIX broken tl call signature") +plf("metric", "clonality", "(via tl.clonality)", "stripplot per phenotype ±group", + "tl.clonality, pl._metric_boxplot") +plf("metric", "flux", "(via tl.flux)", "box of flux distance by group", + "tl.flux", status="fix", source="FIX broken key= passed to tl.flux") +plf("compare", "mi_compare", "uns/(via tl.mi_compare)", "patient MI box + AUROC/permutation stats", + "tl.mi_compare, _stats.auc_and_label_permutation, _stats.bootstrap_auc") +plf("compare", "bayesian_mutual_information", "(via tl.mutual_information ×2)", + "ΔMI KDE / posterior / bar across two covariates", "tl.mutual_information") +plf("compare", "ridge_delta_entropy", "DataFrame from tl.delta_entropy_table", + "ridge plot of Δ-entropy posteriors", "—", status="fix", + upstream="tl.delta_entropy_table", source="FIX undefined cm/st imports") +plf("distribution", "phenotypic_flux", "(via pp.joint_distribution)", "phenotype-flow sankey across covariates", + "pp.joint_distribution, _build_sankey", source="keep; absorbs plot_pheno_sankey as private _build_sankey") +plf("distribution", "polar_plot", "(via pp.joint_distribution / tl.clonotypic_entropy)", + "radar of phenotype distribution or entropy", "pp.joint_distribution, tl.clonotypic_entropy", + status="fix", source="FIX undefined clonotypic_entropy ref + string phenotypes") +plf("umap", "clone_size_umap", "obs[clone_size], obsm[X_umap]", "UMAP colored by log clone size", + "pp.clone_size") +plf("umap", "top_clone_umap", "obs[clone_col], obsm[X_umap]", "UMAP highlighting top-N clones", "—") +plf("umap", "phenotype_probabilities_umap", "obsm[X_tcri_probabilities]", "per-phenotype probability UMAP grid", + "—", status="rename", source="rename of plot_phenotype_probabilities") +plf("diagnostic", "model_loss", "model.history_", "ELBO + dKL training curves", "—", + status="move", upstream="model.train", source="moved from TCRIModel.plot_loss") +plf("diagnostic", "archetypes", "model.c2p_mat, model.centers", "archetype / clone-prior heatmaps", "—", + status="move", upstream="model.train", source="moved from TCRIModel.plot_archetypes") +plf("diagnostic", "model_pgm", "— (static daft diagram)", "TCRI plate-diagram (PGM)", "—", + status="move+merge", upstream="—", + source="moved from ut.build_nested_tcri_pgm + ut.draw_tcri_pgm_nested (merged)") + +# ============================ ut — session / IO ============================= +fn(ns="ut", group="session", name="save_tcri_session", status="keep", + signature="(model, adata, out_dir, *, save_adata=True, compression='gzip')", + returns="dict{paths}", writes="run_dir/{model, pyro_params.pt, setup.json, adata.h5ad, meta.json}", + reads="model.save, pyro store, adata", calls="write_adata_safely, _collect_setup_from_adata_or_model", + upstream="train", invariants="adata written without tcri_manager (non-picklable)", + edges="—", mirror="—", source="keep as-is", notes="") +fn(ns="ut", group="session", name="load_tcri_session", status="keep", + signature="(run_dir, *, adata_path=None, map_location=None, layer=None)", + returns="(model, adata)", writes="pyro param store (restored)", + reads="run_dir artifacts", calls="TCRIModel.setup_anndata/.load, _pyro_load, " + "_restore_category_order, _ensure_pyro_posterior_params, _disable_scvi_onload_train", + upstream="save_tcri_session", invariants="category order restored; posterior params ensured", + edges="missing pyro store → warn + uniform-prior fallback", mirror="—", + source="keep as-is", notes="") +fn(ns="ut", group="session", name="write_adata_safely", status="keep", + signature="(adata, path, *, compression='gzip')", returns="None", + writes="path (h5ad without tcri_manager)", reads="adata", calls="_pop_nonserializables", + upstream="—", invariants="strips non-serializable manager before write", edges="—", + mirror="—", source="keep as-is", notes="") + +# ---- shared primitives (grouped) ------------------------------------------- +PRIMS = { + "Console (tcri/_console.py) — replaces 3 duplicated copies": [ + ("_ok / _info / _warn / _fin", "every verbose tl/pp/pl/ml function"), + ("_ascii_hist(samples)", "mutual_information, flux, delta_clonotypic_entropy (graph=)"), + ("ANSI constants (RESET/BOLD/GRN/...)", "all of the above"), + ], + "Stats (tcri/_stats.py)": [ + ("_joint_to_mi(pxy, normalised, mode)", "tl.mutual_information"), + ("distance(metric) -> f(p,q) [l1 | dkl | callable]", "tl.flux, tl.flux_table"), + ("auc_and_label_permutation(scores, labels)", "pl.mi_compare"), + ("bootstrap_auc(scores, labels)", "pl.mi_compare"), + ("_norm_entropy(p, base, n)", "clonotypic_entropy, phenotypic_entropy"), + ], + "Table builder (tcri/metrics/_tables.py)": [ + ("_group_table(adata, splitby, per_group_fn)", "mi_compare, delta_entropy_table, flux_table"), + ], + "Plot helpers (tcri/plotting/_base.py)": [ + ("_metric_boxplot(adata, fn, ...)", "clonality (+ any group×split metric box)"), + ("_resolve_palette / tcri_colors", "every pl function"), + ("_build_sankey / SankeyNode / _phenotype_mass_per_clone", "phenotypic_flux"), + ], + "Session internals (tcri/utils/_io.py)": [ + ("_ensure_dir, _to_jsonable, _pop_nonserializables", "save_tcri_session, write_adata_safely"), + ("_collect_setup_from_adata_or_model, _restore_category_order", "save/load_tcri_session"), + ("_pyro_load, _resolve_TCRIModel, _disable_scvi_onload_train, _ensure_pyro_posterior_params", + "load_tcri_session"), + ], +} + +# ---- adata-state schema (the data contract; root of the DAG) --------------- +STATE = [ + ("uns[tcri_metadata]", "register_model", "{covariate,clone,phenotype,batch}_col", + "≈ every tl/pp/pl function"), + ("uns[tcri_p_ct]", "register_model", "learned clone×cov phenotype prior [CT,P]", + "pp.joint_distribution"), + ("uns[tcri_local_scale]", "register_model", "Dirichlet concentration scale", + "pp.joint_distribution (posterior draw)"), + ("uns[tcri_{covariate,clonotype,phenotype}_categories]", "register_model", "category orders", + "joint_distribution, all metrics"), + ("uns[tcri_ct_to_cov], uns[tcri_ct_to_c]", "register_model", "ct→cov / ct→clone maps", + "pp.joint_distribution (posterior=False)"), + ("uns[tcri_{ct,cov}_array_for_cells]", "register_model", "per-cell ct/cov indices", + "pp.joint_distribution (cell selection)"), + ("obsm[X_tcri]", "register_model", "latent posterior mean z", "pl.*_umap (if X_umap absent)"), + ("obsm[X_tcri_logits]", "register_model", "classifier logits per cell", + "pp.joint_distribution (combine_with_logits)"), + ("obsm[X_tcri_probabilities]", "register_model", "softmax phenotype probs [N,P]", + "pl.phenotype_probabilities_umap"), + ("obs[tcri_phenotype]", "register_model", "hard phenotype label", "tl.clonality"), + ("obs[clone_size]", "pp.clone_size", "cells per clone", "pl.clone_size_umap"), + ("obs[trb_unique]", "pp.group_singletons", "collapsed clone id", "setup_anndata(clonotype_key)"), + ("pyro param store", "ml.TCRIModel.train", "q_p_c_raw / q_p_ct_raw", "get_p_ct, save/load_session"), +] + +DROPPED = [ + ("pp.classify_phenotypes", "deleted", "redundant cosine phenotype assignment; 0 callers"), + ("pp.get_latent_embedding", "deleted", "trivial gaussian sampler; 0 callers"), + ("pp.register_probability_columns", "deleted", "feeds only the dead probabilities()"), + ("pp.gene_entropy", "deleted", "generic gene QC, out of scope; 0 callers"), + ("pp.group_small_clones", "subsumed → group_singletons", "hardcoded inferior dup"), + ("pp.register_phenotype_key / register_clonotype_key", "subsumed → register_model", + "kill the tcri_*_key shadow convention"), + ("pp._compute_logits_and_prior", "subsumed → ml.get_cell_phenotype_probs", "one prob path"), + ("pp.joint_distribution_posterior", "subsumed → joint_distribution(posterior=True)", "unified engine"), + ("tl.clonotypic_entropy_base", "subsumed → clonotypic_entropy(phenotype=)", "one entropy fn"), + ("tl.clone_fraction", "deleted", "one-line value-counts; 0 callers"), + ("tl.dkl", "subsumed → _stats.distance", "dead; flux had its own copy"), + ("pl.compare_phenotypes", "deleted", "trivial crosstab heatmap; 0 callers"), + ("pl.compare_joint_distribution", "deleted", "broken (undefined `model`)"), + ("pl.probability_distribution", "deleted", "broken (self-recursion); covered by polar_plot"), + ("pl.set_color_palette", "deleted", "buggy palette helper; 0 callers"), + ("pl.plot_pheno_sankey", "subsumed → phenotypic_flux (_build_sankey)", "one public sankey"), + ("pl.tcri_boxplot", "internalized → _metric_boxplot", "shared plot helper"), + ("ut.probabilities", "deleted", "reads uns[joint_distribution] nothing ever writes"), + ("ut.stars", "deleted", "imported but never called"), + ("ut.build_nested_tcri_pgm / draw_tcri_pgm_nested", "moved+merged → pl.model_pgm", + "describes the model, not IO"), + ("ut.auc_and_label_permutation / bootstrap_auc", "internalized → _stats", "used only by pl.mi_compare"), + ("SankeyNode.hex_to_rgb", "deleted", "unused method"), + ("_ok/_info/_warn/_fin (×3), _ascii_hist (×2)", "subsumed → _console", "dedup"), +] + +CONVENTIONS = ( + "ml: scvi/pyro model; the single source of phenotype probabilities " + "(get_cell_phenotype_probs). " + "pp: writes the canonical tcri_* state and owns the one " + "joint_distribution engine (posterior= flag). " + "tl: keyword-only; pure (reads state, returns arrays/frames, writes nothing); " + "n_samples=0 point estimate else posterior draws. " + "pl: ax/save/palette; calls the matching tl/pp " + "(never owns model math). " + "ut: session save/load only. " + "Every module declares __all__; shared logic lives in " + "_console/_stats/_base, never copied." +) + +# ============================================================================= +# render +# ============================================================================= +def mc(s): + """Escape a value for use inside a Markdown table cell.""" + return str(s).replace("|", "\\|").replace("\n", " ") + + +def md(): + L = ["# TCRI — Target API Contract (`ml`/`pp`/`tl`/`pl`/`ut`)", + "", + "_The ideal post-refactor surface. Work toward this; do not drift. " + "Generated from `build_tcri_contract.py` (single source of truth)._", + "", + "**Conventions.** " + CONVENTIONS.replace("", "**").replace("", "**") + .replace("", "`").replace("", "`"), + "", + "## Namespaces", ""] + for v, (c, label) in VIEWS.items(): + L.append(f"- **{label}** (`{v}`)") + L += ["", "## Summary (target inventory)", "", + "| ns | group | function | status | returns | writes |", + "|---|---|---|---|---|---|"] + for f in F: + L.append(f"| `{f['ns']}` | {mc(f['group'])} | `{f['name']}` | {mc(f['status'])} | " + f"{mc(f['returns'])} | {mc(f['writes'])} |") + L += ["", "## adata-state schema (the data contract — root of the DAG)", "", + "| key | produced by | meaning | consumed by |", "|---|---|---|---|"] + for k, prod, mean, cons in STATE: + L.append(f"| `{k}` | `{prod}` | {mc(mean)} | {mc(cons)} |") + L += ["", "## Detail cards", ""] + for f in F: + L.append(f"### `{f['ns']}.{f['name']}` — _{f['group']}_ · {f['status']}") + L.append("") + sig = f["signature"] if f["signature"].startswith("(") else "(" + f["signature"] + ")" + L.append(f"```python\n{f['ns']}.{f['name']}{sig}\n```") + L.append(f"- **returns:** {f['returns']}") + L.append(f"- **writes:** {f['writes']}") + L.append(f"- **reads:** {f['reads']}") + L.append(f"- **calls:** {f['calls']}") + L.append(f"- **upstream:** {f['upstream']}") + L.append(f"- **invariants:** {f['invariants']}") + L.append(f"- **edge cases:** {f['edges']}") + if f.get("mirror") and f["mirror"] != "—": + L.append(f"- **plot mirror:** {f['mirror']}") + L.append(f"- **provenance:** {f['source']}") + if f.get("notes"): + L.append(f"- **note:** {f['notes']}") + L.append("") + L += ["## Shared primitives", ""] + for grp, items in PRIMS.items(): + L.append(f"**{grp}**") + L.append("") + for sig, used in items: + L.append(f"- `{sig}` — {used}") + L.append("") + L += ["## Dropped / merged from the current code", "", + "| current | disposition | why |", "|---|---|---|"] + for cur, disp, why in DROPPED: + L.append(f"| `{cur}` | {mc(disp)} | {mc(why)} |") + L += ["", "_See `tcri_dependency_map.md` for the full call + producer/consumer graph._"] + return "\n".join(L) + + +def esc(s): + return _html.escape(str(s)) + + +def render_html(): + css = """ + body{font:14px/1.5 -apple-system,Segoe UI,Roboto,sans-serif;margin:0;background:#fafafa;color:#1a1a1a} + .wrap{max-width:1180px;margin:0 auto;padding:28px} + h1{font-size:26px;margin:0 0 4px} h2{font-size:20px;margin:34px 0 12px;border-bottom:2px solid #eee;padding-bottom:6px} + .sub{color:#555;margin:0 0 14px} + .conv{background:#f1f5f9;border-left:4px solid #64748b;padding:10px 14px;border-radius:6px;margin:14px 0} + .legend span{display:inline-block;margin:2px 8px 2px 0;padding:3px 10px;border-radius:12px;color:#fff;font-weight:600;font-size:12px} + table{border-collapse:collapse;width:100%;font-size:12.5px;margin:8px 0 20px} + th,td{border:1px solid #e5e7eb;padding:6px 9px;text-align:left;vertical-align:top} + th{background:#f3f4f6;position:sticky;top:0} + code{background:#eef2ff;padding:1px 5px;border-radius:4px;font-size:92%} + .badge{display:inline-block;padding:2px 9px;border-radius:11px;color:#fff;font-weight:600;font-size:11px;white-space:nowrap} + .card{border:1px solid #e5e7eb;border-left:6px solid #999;border-radius:8px;padding:12px 16px;margin:12px 0;background:#fff} + .card h3{margin:0 0 8px;font-size:15px} + .card pre{background:#0f172a;color:#e2e8f0;padding:10px 12px;border-radius:6px;overflow-x:auto;font-size:12px;white-space:pre-wrap} + .card .row{margin:3px 0} .card .k{font-weight:700;color:#334155} + .grp{font-weight:700;margin:12px 0 4px;color:#334155} + .drop{color:#b91c1c} .merge{color:#b45309} .new{color:#15803d} + ul{margin:4px 0 4px 18px} + """ + def badge(v): + c, label = VIEWS[v] + return f'{esc(label)}' + + H = ['TCRI target API contract', + f'
'] + H.append("

TCRI — Target API Contract

") + H.append('

Ideal post-refactor surface · ml/pp/' + 'tl/pl/ut · work toward this, do not drift.

') + H.append(f'
Conventions. {CONVENTIONS}
') + H.append('
' + "".join(badge(v) for v in VIEWS) + "
") + + H.append("

Summary (target inventory)

") + H.append("" + "") + for f in F: + H.append(f"" + f"" + f"") + H.append("
nsgroupfunctionstatusreturnswrites
{badge(f['ns'])}{esc(f['group'])}{esc(f['name'])}{esc(f['status'])}{esc(f['returns'])}{esc(f['writes'])}
") + + H.append("

adata-state schema (the data contract — root of the DAG)

") + H.append("") + for k, prod, mean, cons in STATE: + H.append(f"" + f"") + H.append("
keyproduced bymeaningconsumed by
{esc(k)}{esc(prod)}{esc(mean)}{esc(cons)}
") + + H.append("

Detail cards

") + for f in F: + c = VIEWS[f["ns"]][0] + H.append(f'
') + H.append(f'

{badge(f["ns"])}   {esc(f["ns"])}.{esc(f["name"])} ' + f'· {esc(f["group"])} · {esc(f["status"])}

') + sig = f["signature"] if f["signature"].startswith("(") else "(" + f["signature"] + ")" + H.append(f'
{esc(f["ns"])}.{esc(f["name"])}{esc(sig)}
') + for k in ("returns", "writes", "reads", "calls", "upstream", "invariants", "edges"): + label = {"edges": "edge cases"}.get(k, k) + H.append(f'
{label}: {esc(f[k])}
') + if f.get("mirror") and f["mirror"] != "—": + H.append(f'
plot mirror: {esc(f["mirror"])}
') + H.append(f'
provenance: {esc(f["source"])}
') + if f.get("notes"): + H.append(f'
note: {esc(f["notes"])}
') + H.append("
") + + H.append("

Shared primitives

") + for grp, items in PRIMS.items(): + H.append(f'
{esc(grp)}
    ') + for sig, used in items: + H.append(f"
  • {esc(sig)} — {esc(used)}
    • ") + H.append("
") + + H.append("

Dropped / merged from the current code

") + H.append("") + for cur, disp, why in DROPPED: + cls = "drop" if disp.startswith("deleted") else ("merge" if ("subsumed" in disp or "internal" in disp or "moved" in disp) else "") + H.append(f'') + H.append("
currentdispositionwhy
{esc(cur)}{esc(disp)}{esc(why)}
") + H.append('

See tcri_dependency_map.html for the full call + ' + 'producer/consumer graph.

') + H.append("
") + return "".join(H) + + +with open(OUT_MD, "w") as f: + f.write(md()) +with open(OUT_HTML, "w") as f: + f.write(render_html()) +print("wrote", OUT_MD) +print("wrote", OUT_HTML) +print("functions:", len(F), "| state keys:", len(STATE), "| dropped/merged:", len(DROPPED)) diff --git a/docs/contract/build_tcri_depgraph.py b/docs/contract/build_tcri_depgraph.py new file mode 100644 index 0000000..30db060 --- /dev/null +++ b/docs/contract/build_tcri_depgraph.py @@ -0,0 +1,376 @@ +#!/usr/bin/env python3 +"""Build the TCRI dependency map for the *target* API. + +Two graphs over the same nodes: + (1) call graph — function -> function (who calls whom) + (2) dataflow graph — function -> adata-state key -> function (producers/consumers) + +Renders: producer/consumer + call-adjacency tables (Markdown + HTML), a Mermaid +DAG for each graph (renders on GitHub / in the HTML via CDN), and a Graphviz DOT +combining both (render offline: `dot -Tsvg tcri_dependency_map.dot -o dep.svg`). + +This is curated from the target contract (build_tcri_contract.py), so it stays a +*contract* we refactor toward — not an after-the-fact extraction. See the +"How this is built / standard tooling" section for the auto-extraction path. +""" +import html as _html +import os +from collections import defaultdict + +_HERE = os.path.dirname(os.path.abspath(__file__)) +OUT_MD = os.path.join(_HERE, "tcri_dependency_map.md") +OUT_HTML = os.path.join(_HERE, "tcri_dependency_map.html") +OUT_DOT = os.path.join(_HERE, "tcri_dependency_map.dot") + +NS_COLOR = { + "ml": "#dc2626", "pp": "#0f766e", "tl": "#7c3aed", + "pl": "#2563eb", "ut": "#475569", "core": "#b45309", "state": "#334155", +} +NS_FILL = { + "ml": "#fde2e2", "pp": "#d6f0ec", "tl": "#ece3fb", + "pl": "#dbe8fd", "ut": "#e5e9ef", "core": "#fef3c7", "state": "#ffffff", +} + +# ---- function nodes (target API + shared prims) ----------------------------- +FUNCS = [ + # ml + ("ml", "setup_anndata"), ("ml", "train"), ("ml", "get_latent_representation"), + ("ml", "get_cell_phenotype_probs"), ("ml", "get_p_ct"), ("ml", "boost_phenotype_prior"), + # pp + ("pp", "register_model"), ("pp", "joint_distribution"), ("pp", "group_singletons"), + ("pp", "clone_size"), ("pp", "filter_genes"), + # tl + ("tl", "clonotypic_entropy"), ("tl", "phenotypic_entropy"), ("tl", "mutual_information"), + ("tl", "clonality"), ("tl", "flux"), ("tl", "delta_clonotypic_entropy"), + ("tl", "mi_compare"), ("tl", "delta_entropy_table"), ("tl", "flux_table"), + # pl + ("pl", "mutual_information"), ("pl", "clonotypic_entropy"), ("pl", "phenotypic_entropy"), + ("pl", "clonality"), ("pl", "flux"), ("pl", "mi_compare"), + ("pl", "bayesian_mutual_information"), ("pl", "ridge_delta_entropy"), + ("pl", "phenotypic_flux"), ("pl", "polar_plot"), ("pl", "clone_size_umap"), + ("pl", "top_clone_umap"), ("pl", "phenotype_probabilities_umap"), + ("pl", "model_loss"), ("pl", "archetypes"), ("pl", "model_pgm"), + # ut + ("ut", "save_tcri_session"), ("ut", "load_tcri_session"), ("ut", "write_adata_safely"), + # core (shared primitives that participate in call edges) + ("core", "_joint_to_mi"), ("core", "_stats.distance"), ("core", "_stats.auc_perm"), + ("core", "_stats.bootstrap_auc"), ("core", "_group_table"), ("core", "_metric_boxplot"), + ("core", "_build_sankey"), +] +NS = {f"{ns}.{nm}": ns for ns, nm in FUNCS} + +# ---- call edges: caller -> callee ------------------------------------------ +CALLS = [ + # pl -> tl / pp / core + ("pl.mutual_information", "tl.mutual_information"), + ("pl.clonotypic_entropy", "tl.clonotypic_entropy"), + ("pl.phenotypic_entropy", "tl.phenotypic_entropy"), + ("pl.clonality", "tl.clonality"), + ("pl.clonality", "core._metric_boxplot"), + ("pl.flux", "tl.flux"), + ("pl.mi_compare", "tl.mi_compare"), + ("pl.mi_compare", "core._stats.auc_perm"), + ("pl.mi_compare", "core._stats.bootstrap_auc"), + ("pl.bayesian_mutual_information", "tl.mutual_information"), + ("pl.ridge_delta_entropy", "tl.delta_entropy_table"), + ("pl.phenotypic_flux", "pp.joint_distribution"), + ("pl.phenotypic_flux", "core._build_sankey"), + ("pl.polar_plot", "pp.joint_distribution"), + ("pl.polar_plot", "tl.clonotypic_entropy"), + ("pl.clone_size_umap", "pp.clone_size"), + ("pl.model_loss", "ml.train"), + ("pl.archetypes", "ml.train"), + # tl -> pp / tl / core + ("tl.clonotypic_entropy", "pp.joint_distribution"), + ("tl.phenotypic_entropy", "pp.joint_distribution"), + ("tl.mutual_information", "pp.joint_distribution"), + ("tl.mutual_information", "core._joint_to_mi"), + ("tl.flux", "pp.joint_distribution"), + ("tl.flux", "core._stats.distance"), + ("tl.delta_clonotypic_entropy", "tl.clonotypic_entropy"), + ("tl.mi_compare", "tl.mutual_information"), + ("tl.mi_compare", "core._group_table"), + ("tl.delta_entropy_table", "tl.delta_clonotypic_entropy"), + ("tl.delta_entropy_table", "core._group_table"), + ("tl.flux_table", "tl.flux"), + ("tl.flux_table", "core._group_table"), + # pp -> ml + ("pp.register_model", "ml.get_cell_phenotype_probs"), + ("pp.register_model", "ml.get_latent_representation"), + ("pp.register_model", "ml.get_p_ct"), + # ml internal + ("ml.get_cell_phenotype_probs", "ml.get_p_ct"), + # ut + ("ut.save_tcri_session", "ut.write_adata_safely"), + ("ut.load_tcri_session", "ml.setup_anndata"), +] + +# ---- dataflow: adata-state keys with producer + concrete consumers ---------- +# (consumers listed are DIRECT readers; metrics that read only via +# pp.joint_distribution are attributed to joint_distribution, not re-listed) +STATE = [ + ("uns[tcri_metadata]", "pp.register_model", + ["pp.joint_distribution", "tl.clonality", "tl.mi_compare", + "tl.delta_entropy_table", "tl.flux_table", "pp.clone_size"]), + ("uns[tcri_p_ct]", "pp.register_model", ["pp.joint_distribution"]), + ("uns[tcri_local_scale]", "pp.register_model", ["pp.joint_distribution"]), + ("uns[tcri_*_categories]", "pp.register_model", + ["pp.joint_distribution", "tl.clonotypic_entropy", "tl.phenotypic_entropy"]), + ("uns[tcri_ct_to_cov/ct_to_c]", "pp.register_model", ["pp.joint_distribution"]), + ("uns[tcri_{ct,cov}_array_for_cells]", "pp.register_model", ["pp.joint_distribution"]), + ("obsm[X_tcri_logits]", "pp.register_model", ["pp.joint_distribution"]), + ("obsm[X_tcri_probabilities]", "pp.register_model", ["pl.phenotype_probabilities_umap"]), + ("obsm[X_tcri]", "pp.register_model", ["pl.clone_size_umap", "pl.top_clone_umap"]), + ("obs[tcri_phenotype]", "pp.register_model", ["tl.clonality", "pl.top_clone_umap"]), + ("obs[clone_size]", "pp.clone_size", ["pl.clone_size_umap"]), + ("obs[trb_unique]", "pp.group_singletons", ["ml.setup_anndata"]), + ("pyro_param_store", "ml.train", + ["ml.get_p_ct", "ut.save_tcri_session", "ut.load_tcri_session"]), +] + +# ============================================================================= +# derive adjacency +# ============================================================================= +calls_out = defaultdict(list) +calls_in = defaultdict(list) +for a, b in CALLS: + calls_out[a].append(b) + calls_in[b].append(a) + +ALL = [f"{ns}.{nm}" for ns, nm in FUNCS] +entry_points = [f for f in ALL if f not in calls_in and NS.get(f) in ("pl", "ut", "pp")] +leaves = [f for f in ALL if f not in calls_out and NS.get(f) in ("ml", "core", "pp")] + + +def sid(s): + out = [] + for ch in s: + out.append(ch if ch.isalnum() else "_") + return "n_" + "".join(out) + + +# ---- Mermaid ---------------------------------------------------------------- +def mermaid_calls(): + L = ["flowchart TB"] + for ns in ("pl", "tl", "pp", "ml", "core", "ut"): + members = [f for f in ALL if NS[f] == ns and (f in calls_out or f in calls_in)] + if not members: + continue + L.append(f' subgraph {ns.upper()}') + for f in members: + L.append(f' {sid(f)}["{f}"]') + L.append(" end") + for a, b in CALLS: + L.append(f" {sid(a)} --> {sid(b)}") + for ns, color in NS_FILL.items(): + L.append(f" classDef {ns} fill:{color},stroke:{NS_COLOR[ns]},color:#111;") + for f in ALL: + if f in calls_out or f in calls_in: + L.append(f" class {sid(f)} {NS[f]};") + return "\n".join(L) + + +def mermaid_dataflow(): + L = ["flowchart LR"] + seen = set() + for key, prod, cons in STATE: + kid = sid("KEY_" + key) + L.append(f' {kid}[("{key}")]') + L.append(f' class {kid} state;') + if prod not in seen: + L.append(f' {sid(prod)}["{prod}"]'); L.append(f' class {sid(prod)} {NS[prod]};'); seen.add(prod) + L.append(f" {sid(prod)} ==> {kid}") + for c in cons: + if c not in seen: + L.append(f' {sid(c)}["{c}"]'); L.append(f' class {sid(c)} {NS[c]};'); seen.add(c) + L.append(f" {kid} -.-> {sid(c)}") + for ns, color in NS_FILL.items(): + L.append(f" classDef {ns} fill:{color},stroke:{NS_COLOR[ns]},color:#111;") + return "\n".join(L) + + +# ---- Graphviz DOT (combined) ------------------------------------------------ +def dot(): + L = ['digraph tcri {', ' rankdir=TB; node [fontname="Helvetica",fontsize=10];', + ' edge [fontname="Helvetica",fontsize=8];'] + for ns in ("pl", "tl", "pp", "ml", "ut", "core"): + L.append(f' subgraph cluster_{ns} {{ label="{ns}"; style=rounded; color="{NS_COLOR[ns]}";') + for f in ALL: + if NS[f] == ns: + L.append(f' "{f}" [shape=box,style=filled,fillcolor="{NS_FILL[ns]}",' + f'color="{NS_COLOR[ns]}"];') + L.append(" }") + L.append(" // call edges (solid)") + for a, b in CALLS: + L.append(f' "{a}" -> "{b}" [color="#475569"];') + L.append(" // state keys + dataflow (write=bold, read=dashed)") + for key, prod, cons in STATE: + L.append(f' "{key}" [shape=cylinder,style=filled,fillcolor="#f8fafc",color="#334155"];') + L.append(f' "{prod}" -> "{key}" [color="#0f766e",penwidth=1.6];') + for c in cons: + L.append(f' "{key}" -> "{c}" [style=dashed,color="#94a3b8"];') + L.append("}") + return "\n".join(L) + + +# ============================================================================= +# render markdown + html +# ============================================================================= +STANDARD = [ + ("Call graph (who-calls-whom)", + "The de-facto standard is a directed graph rendered with Graphviz/DOT. " + "Auto-extract from source with the stdlib ast module, or tools like " + "pyan3, code2flow, pydeps (module-level), or " + "griffe (the engine behind mkdocstrings). For Markdown-native rendering use " + "Mermaid flowchart (GitHub renders it inline); for interactive web use " + "cytoscape.js or d3."), + ("Dataflow / producer-consumer", + "Because TCRI couples through adata.uns/obsm/obs keys (not just direct calls), " + "the precise model is a bipartite graph of functions ↔ state keys. This is exactly a " + "build-system DAG: state keys are the artifacts/targets, functions are the rules. The " + "standard tools for that shape are Make, Snakemake, or dbt " + "(dbt docs renders an interactive lineage graph) — and the same idea is what " + "scverse calls a data-flow/provenance graph. Here it is curated by hand from the " + "contract so it can lead the refactor rather than trail it."), + ("Preventing drift after the refactor", + "Once the target lands, point an ast walker at tcri/ to extract the " + "actual call edges + uns/obsm/obs read/writes and diff them against this " + "curated graph in CI. Drift = a function that calls or reads/writes something the contract " + "doesn't list."), +] + + +def md(): + L = ["# TCRI — Dependency Map (target API)", + "", + "_Call graph + `adata`-state producer/consumer links for the post-refactor API. " + "Curated from `build_tcri_contract.py`; regenerate with `build_tcri_depgraph.py`._", + "", + "**Two graphs, same nodes.** (1) **Call graph** — `pl` → `tl` → `pp` → `ml`, bottoming " + "out at the model + shared `core` primitives. (2) **Dataflow** — everything funnels " + "through the `tcri_*` state that `pp.register_model` writes; `pp.joint_distribution` is " + "the universal consumer hub.", + "", + "## Call graph", "", + "```mermaid", mermaid_calls(), "```", "", + "### Call adjacency", "", + "| function | calls | called by |", "|---|---|---|"] + for f in ALL: + if f in calls_out or f in calls_in: + out = ", ".join(f"`{x}`" for x in calls_out.get(f, [])) or "—" + inn = ", ".join(f"`{x}`" for x in calls_in.get(f, [])) or "—" + L.append(f"| `{f}` | {out} | {inn} |") + L += ["", f"**Entry points** (no caller): " + ", ".join(f"`{f}`" for f in entry_points), + "", f"**Leaves** (call nothing in-package): " + ", ".join(f"`{f}`" for f in leaves), + "", + "## Dataflow (producers / consumers)", "", + "```mermaid", mermaid_dataflow(), "```", "", + "### State producer/consumer table", "", + "| adata key | produced by | consumed by |", "|---|---|---|"] + for key, prod, cons in STATE: + L.append(f"| `{key}` | `{prod}` | " + ", ".join(f"`{c}`" for c in cons) + " |") + L += ["", + "> Metrics that read only `uns[tcri_metadata]` + categories **via** " + "`pp.joint_distribution` are attributed to that hub, not re-listed per key.", + "", + "## How this is built / standard tooling", ""] + for title, body in STANDARD: + b = body.replace("", "**").replace("", "**").replace("", "_").replace("", "_") + b = b.replace("", "`").replace("", "`") + L.append(f"**{title}.** {b}") + L.append("") + L += ["## Graphviz", "", + "A combined DOT file is emitted to `tcri_dependency_map.dot` " + "(call edges solid, writes bold-teal, reads dashed). Render:", + "", "```bash", "dot -Tsvg tcri_dependency_map.dot -o tcri_dependency_map.svg", "```"] + return "\n".join(L) + + +def esc(s): + return _html.escape(str(s)) + + +def render_html(): + css = """ + body{font:14px/1.5 -apple-system,Segoe UI,Roboto,sans-serif;margin:0;background:#fafafa;color:#1a1a1a} + .wrap{max-width:1180px;margin:0 auto;padding:28px} + h1{font-size:26px;margin:0 0 4px} h2{font-size:20px;margin:34px 0 12px;border-bottom:2px solid #eee;padding-bottom:6px} + h3{font-size:15px;margin:18px 0 8px} + .sub{color:#555;margin:0 0 14px} + .conv{background:#f1f5f9;border-left:4px solid #64748b;padding:10px 14px;border-radius:6px;margin:14px 0} + table{border-collapse:collapse;width:100%;font-size:12.5px;margin:8px 0 20px} + th,td{border:1px solid #e5e7eb;padding:6px 9px;text-align:left;vertical-align:top} + th{background:#f3f4f6;position:sticky;top:0} + code{background:#eef2ff;padding:1px 5px;border-radius:4px;font-size:92%} + .mermaid{background:#fff;border:1px solid #e5e7eb;border-radius:8px;padding:14px;margin:10px 0;overflow-x:auto} + .note{background:#fffbeb;border-left:4px solid #d97706;padding:8px 12px;border-radius:5px;margin:10px 0} + .legend span{display:inline-block;margin:2px 8px 2px 0;padding:3px 10px;border-radius:12px;color:#fff;font-weight:600;font-size:12px} + """ + legend = "".join( + f'{ns}' + for ns in ("ml", "pp", "tl", "pl", "ut", "core")) + H = ['TCRI dependency map', + f'', + '', + '
'] + H.append("

TCRI — Dependency Map (target API)

") + H.append('

Call graph + adata-state producer/consumer links · ' + 'curated from the target contract.

') + H.append(f'
{legend}
') + H.append('
Two graphs, same nodes. (1) Call graph: ' + 'pl → tl → pp → ml, bottoming out at the model + shared core ' + 'primitives. (2) Dataflow: everything funnels through the tcri_* ' + 'state that pp.register_model writes; pp.joint_distribution is ' + 'the universal consumer hub.
') + + H.append("

Call graph

") + H.append(f'
{esc(mermaid_calls())}
') + H.append("

Call adjacency

") + H.append("") + for f in ALL: + if f in calls_out or f in calls_in: + out = ", ".join(f"{esc(x)}" for x in calls_out.get(f, [])) or "—" + inn = ", ".join(f"{esc(x)}" for x in calls_in.get(f, [])) or "—" + color = NS_COLOR[NS[f]] + H.append(f'' + f"") + H.append("
functioncallscalled by
{esc(f)}{out}{inn}
") + H.append('
Entry points (no caller): ' + + ", ".join(f"{esc(f)}" for f in entry_points) + + "
Leaves (call nothing in-package): " + + ", ".join(f"{esc(f)}" for f in leaves) + "
") + + H.append("

Dataflow (producers / consumers)

") + H.append(f'
{esc(mermaid_dataflow())}
') + H.append("") + for key, prod, cons in STATE: + H.append(f"" + f"") + H.append("
adata keyproduced byconsumed by
{esc(key)}{esc(prod)}" + ", ".join(f"{esc(c)}" for c in cons) + "
") + H.append('
Metrics that read only uns[tcri_metadata] + categories ' + 'via pp.joint_distribution are attributed to that hub, not re-listed ' + 'per key.
') + + H.append("

How this is built / standard tooling

") + for title, body in STANDARD: + H.append(f"

{esc(title)}. {body}

") + H.append("

Graphviz

") + H.append('

Combined DOT emitted to tcri_dependency_map.dot ' + '(call edges solid, writes bold-teal, reads dashed). Render: ' + 'dot -Tsvg tcri_dependency_map.dot -o tcri_dependency_map.svg.

') + H.append("
") + return "".join(H) + + +with open(OUT_MD, "w") as f: + f.write(md()) +with open(OUT_HTML, "w") as f: + f.write(render_html()) +with open(OUT_DOT, "w") as f: + f.write(dot()) +print("wrote", OUT_MD) +print("wrote", OUT_HTML) +print("wrote", OUT_DOT) +print(f"nodes: {len(ALL)} | call edges: {len(CALLS)} | state keys: {len(STATE)} | " + f"entry points: {len(entry_points)} | leaves: {len(leaves)}") diff --git a/docs/contract/build_tcri_inventory.py b/docs/contract/build_tcri_inventory.py new file mode 100644 index 0000000..d15e63b --- /dev/null +++ b/docs/contract/build_tcri_inventory.py @@ -0,0 +1,316 @@ +#!/usr/bin/env python3 +"""Render the full TCRI function inventory + consolidation plan from the +inventory-workflow output (tcri_inventory_data.json). + +Data-driven sections (inventory table, consolidation groups, deletions, helper +extraction, plotting triage) come straight from the workflow result. The +authored overlay sections (grafiti target shape, target tree, rename map, metric +conventions, diagnostics scoping) encode the five design decisions layered on +top. Regenerate: python3 docs/contract/build_tcri_inventory.py +""" +import json, os, textwrap + +HERE = os.path.dirname(os.path.abspath(__file__)) +DATA = os.path.join(HERE, "tcri_inventory_data.json") +OUT_MD = os.path.join(HERE, "tcri_function_inventory.md") + +res = json.load(open(DATA)) +merged = res["merged"] +crit = res["critic"] +syn = res["synthesis"] + +# ---- critic corrections (label overrides) -------------------------------- +CRITIC_FIX = { + ("pl", "clonality"): "plotting-beyond-core", + ("pl", "tcri_boxplot"): "helper", + ("ml", "TCRIModule.get_latent"): "model-construction", + ("ml", "TCRIModule.get_p_ct"): "model-construction", +} +for r in merged: + key = (r["namespace"], r["name"]) + if key in CRITIC_FIX: + r["label"] = CRITIC_FIX[key] + +# ---- disposition per row ------------------------------------------------- +DISPO_DEFAULT = { + "core": "keep", + "redundant": "merge", + "helper": "extract → shared", + "model-construction": "keep (internal)", + "session-io": "keep", + "plotting-beyond-core": "move→examples / drop", + "dead-broken": "delete", +} +# name-based overrides (apply our 5 overlays + the KEEP+FIX pl twins) +DISPO_OVERRIDE = { + "mutual_information": None, # ambiguous (tl vs pl) — handled by (ns,name) below +} +NAME_OVERRIDE = { + ("pl", "mutual_information"): "keep + FIX", + ("pl", "phenotypic_entropy"): "keep + FIX", + ("pl", "ridge_delta_entropy"): "keep + FIX", + ("pl", "compare_joint_distribution"): "→ diagnostics (PPC)", # OVERLAY: was 'delete' + ("ml", "TCRIModel.plot_loss"): "→ diagnostics", # OVERLAY + ("ml", "TCRIModel.plot_archetypes"): "→ diagnostics", # OVERLAY + ("ut", "build_nested_tcri_pgm"): "→ diagnostics", # OVERLAY + ("ut", "draw_tcri_pgm_nested"): "→ diagnostics", # OVERLAY +} + +def disposition(r): + key = (r["namespace"], r["name"]) + if key in NAME_OVERRIDE: + return NAME_OVERRIDE[key] + d = DISPO_DEFAULT.get(r["label"], "?") + if r["label"] == "redundant" and r.get("consolidate_into"): + return "merge → " + r["consolidate_into"] + return d + +def mc(s): + return str(s).replace("|", "\\|").replace("\n", " ").strip() + +def short(s, n=150): + return textwrap.shorten(str(s), width=n, placeholder=" …") + +# ===================================================================== +CORE_DEF = """\ +**Core = five things; everything else must justify itself against them.** +1. **Model (`ml`)** — `TCRIModel`: build / train / evaluate / register outputs onto an AnnData. +2. **Engine (`pp`)** — `joint_distribution`: Bayesian posterior sampling of the clone×phenotype distribution (the substrate every metric reads). +3. **Metrics (`tl`)** — `clonotypic_entropy`, `phenotypic_entropy`, `mutual_information` (+ `flux`, `delta_clonotypic_entropy`, and the tidy-table builders). +4. **Plotting (`pl`)** — plots that *directly* visualize those metrics and the joint distribution / flux. +5. **Utils (`ut`) + shared helpers** — session save/load; deduplicated console / stats / distance / color helpers. +""" + +GRAFITI = """\ +The reference layout (`../grafiti`) is a flat package with scanpy-style sub-packages aliased to short handles +(`model→ml`, `tools→tl`, `plotting→pl`, `preprocessing→pp`, `diagnostics→diag`, `datasets→ds`, `get.py→get`). +Six patterns tcri should copy: + +1. **One file per topic, never a monolith.** `tools/_joint.py`, `tools/_motif.py`, … and `plotting/` mirrors them 1:1 by filename. (tcri's 1008-line `_metrics.py` and 1437-line `_plotting.py` are the anti-pattern.) +2. **Private cross-cutting sub-packages.** `_state/` (keys, resolve, storage, schemas) + `_compute/` (device-routed math) hold everything shared, so impl files stay thin. tcri analog: `_keys.py`, `_console.py`, `_stats.py`, `_distance.py` (and later a `_compute`). +3. **`__all__` at BOTH levels; NO `import *`.** Each impl module declares `__all__`; each `__init__` names every symbol explicitly and re-declares an aggregate `__all__` grouped by view. **This corrects the earlier `import *` instinct** — the mature pattern is explicit re-export, which keeps numpy/pandas/helpers *unexported*. +4. **`get.py` + `@tl_result` cache convention.** tl writes a versioned uns blob and returns a tidy result; `pl` functions are pure *cache renderers* (`load_result(adata, key)` → draw, never compute). *(tcri: adopt `_keys.py` now; defer the cache decorator — see open questions.)* +5. **`diagnostics/` returns DATA, not plots.** `gf.diag` runs read-only concordance/quality checks on the *finalized* model and returns a DataFrame ("did the model fit?"), deliberately outside the tl-writes / pl-reads loop. **This is exactly where PPCs + model-validation live.** +6. **Naming.** public package + private `_topic.py` impl modules; helper *packages* underscore-prefixed dirs, helper *files* underscore-prefixed, helper *functions* underscore-prefixed; tl↔pl twins share filename + function name. +""" + +TARGET_TREE = """\ +``` +tcri/ + __init__.py # explicit re-export + sys.modules aliases (tl/pp/pl/ml/ut/diag); NO import * + _keys.py # single source of every uns/obsm/obs key string + _console.py # _ok/_info/_warn/_fin/_ascii_hist (was triplicated across 3 files) + _stats.py # stars, auc_and_label_permutation, bootstrap_auc + _distance.py # kl_divergence, l1_distance, phenotype_distance dispatch (was dkl + flux.dkl_func) + model/ # ml + _model.py # TCRIModel: setup_anndata, train, get_latent_representation, get_cell_phenotype_probs, get_p_ct + _module.py # TCRIModule (pyro model/guide, get_latent, get_p_ct) + _priors.py # MixtureDirichlet, VampPrior + _classifier.py # PhenotypeClassifier + _training.py # UnifiedTrainingPlan, build_archetypes + preprocessing/ # pp + _register.py # register_model (+ folded register_*_key, _compute_logits_and_prior) + _engine.py # joint_distribution(posterior=, n_samples=) (unifies the two current fns) + _clones.py # group_singletons, clone_size + metrics/ # tl + _entropy.py # clonotypic_entropy, phenotypic_entropy, delta_clonotypic_entropy + _mutual_information.py # mutual_information (+ private _mi_from_joint) + _flux.py # flux + _tables.py # mi_compare, delta_entropy_table, flux_table + plotting/ # pl (twins mirror tl by filename) + _base.py # _metric_boxplot (was tcri_boxplot), _finish + _colors.py # tcri_colors, resolve_palette (was set_color_palette) + _entropy.py # clonotypic_entropy (was _by_phenotype), phenotypic_entropy [FIX], ridge_delta_entropy [FIX] + _mutual_information.py # mutual_information [FIX], mi_compare + _flux.py # phenotypic_flux (sankey) + _sankey.py # SankeyNode, _phenotype_mass_per_clone + diagnostics/ # diag (NEW — PPCs + model validation, returns DataFrames) + _ppc.py # joint-distribution PPC (was compare_joint_distribution, fixed) + calibration + reconstruction PPC + _training.py # loss curves (was plot_loss), archetypes (was plot_archetypes) + _pgm.py # model PGM (was build_nested_tcri_pgm) + utils/ # ut + _session.py # save/load_tcri_session, write_adata_safely, _to_jsonable +examples/ # bespoke one-offs move here: top_clone_umap, clone_size_umap, + # phenotype_probabilities UMAP, compare_phenotypes + the rewritten notebooks +``` +""" + +RENAME_MAP = """\ +Freeze this **before** touching code or notebooks — renames are breaking and we only pay once (pre-1.0, Alpha). + +### Modules / files +| current | → target | +|---|---| +| `metrics/_metrics.py` (1008 ln) | `tl/_entropy.py` + `_mutual_information.py` + `_flux.py` + `_tables.py` | +| `preprocessing/_preprocessing.py` (559) | `pp/_register.py` + `_engine.py` + `_clones.py` | +| `plotting/_plotting.py` (1437) | `pl/_entropy.py` + `_mutual_information.py` + `_flux.py` + `_base.py` + `_colors.py` | +| `model/_model.py` (1074) | `model/_model.py` + `_module.py` + `_priors.py` + `_classifier.py` + `_training.py` | +| `utils/_utils.py` (665) | `utils/_session.py` + new `_console.py` / `_stats.py` / `_distance.py` / `_keys.py` + `diagnostics/` | + +### Functions +| current | → target | why | +|---|---|---| +| `joint_distribution` + `joint_distribution_posterior` | `joint_distribution(posterior=, n_samples=)` | one engine, one point/draws knob | +| `clonotypic_entropy_base` | merged into `clonotypic_entropy` | `_base` means nothing; it's the single-phenotype case | +| `pl.clonotypic_entropy_by_phenotype` | `pl.clonotypic_entropy` | tl↔pl twin name; matches notebook expectation | +| `plot_phenotype_probabilities` / `plot_pheno_sankey` | drop `plot_` prefix (`phenotype_probabilities`, internal `_sankey`) | scanpy/scvi pl fns are unprefixed | +| `tcri_boxplot` | `_metric_boxplot` (private helper) | generic engine, not public API | +| `dkl` / `flux.dkl_func` | `_distance.kl_divergence` | dedupe the KL kernel | +| `centropy` / `pentropy` / `*_tl` aliases | **removed** (via `__all__`) | leaked import aliases | +| `classify_phenotypes` | **removed** → `register_model` | duplicate phenotype-assignment path | +| `remove_meaningless_genes` | `filter_genes` (if kept) *or* delete | broken flag logic + 0 callers | + +### Parameters +| current | → target | +|---|---| +| `from_this` / `to_that` (flux) | `cov_from` / `cov_to` | +| `point_estimate=True` (entropies) | **removed** — use `n_samples=0` | +| ad-hoc arg orders | standardize `covariate` / `splitby` / `groupby` / `clones` / `n_samples` / `temperature` / `posterior` everywhere | + +### State keys +| current | → target | +|---|---| +| `uns["tcri_clone_key"]` / `["tcri_phenotype_key"]` **and** `uns["tcri_metadata"][...]` (two conventions) | one scheme via `_keys.py` constants (single `tcri_metadata`) | +| ad-hoc registry keys (`"clonotype_col_in_registry"`) | standard scvi `REGISTRY_KEYS` | + +### Internal variables +| current | → target | +|---|---| +| `Δ` (unicode) in `delta_clonotypic_entropy` | `delta` (ASCII, greppable) | +| `c2p_mat` | `clone_phenotype_prior` | +| `p_ct` / `ct_to_c` / `ct_to_cov` (terse) | keep, but document (`ct` = (clone, covariate) index) | +""" + +CONVENTIONS = """\ +### Uniform point-estimate / sampling convention (applies to every metric) +- **`n_samples=0` → deterministic point estimate.** Posterior-*mean* `p_ct` (× logits), softmax, **no draw**, reproducible. +- **`n_samples=N>0` → N posterior draws** (adds a sampling axis; mean ± CI fall out). +- **Delete the `point_estimate=` argument** — `n_samples` is the only knob. This also fixes a latent bug: today `mutual_information(n_samples=0)` / `flux(n_samples=0)` return **one random draw**, not a deterministic estimate. + +### Prior-vs-posterior — PARKED (open, do not collapse in this pass) +The `{prior, posterior} × {point, draws}` 2×2 is deferred. Until resolved the plan assumes `posterior=True` with +`n_samples` as the point/draws knob, and keeps `posterior=` as a documented-but-unfinalized argument +(the current prior-only branch raises `NotImplementedError`). + +### diagnostics/ = PPCs + model validation +`gf.diag`-style, returns DataFrames, read-only on the finalized model. Seeded by: +- **joint-distribution PPC** — model p(clone,phenotype) vs empirical counts (the *fixed* `compare_joint_distribution`) +- phenotype-probability **calibration**; **reconstruction PPC** (ZINB simulate → compare library/dropout/mean-var); entropy/MI vs permutation null +- relocated: training curves (`plot_loss`), archetypes (`plot_archetypes`), model PGM (`build_nested_tcri_pgm`) +""" + +# ===================================================================== +def md(): + L = [] + A = L.append + A("# TCRI — Full Function Inventory & Consolidation Plan") + A("") + A("_Generated from the inventory workflow (9 agents, 131 functions, completeness-verified) via " + "`build_tcri_inventory.py`. Data: `tcri_inventory_data.json`. This is the working list we reduce " + "the repo against._") + A("") + A("## 0. Core definition") + A("") + A(CORE_DEF) + # counts + from collections import Counter + lab = Counter(r["label"] for r in merged) + A("## 1. Label counts (critic-corrected)") + A("") + A("| label | count | disposition |") + A("|---|---|---|") + order = ["core","redundant","plotting-beyond-core","model-construction","session-io","helper","dead-broken"] + disp = { + "core":"keep (22 survive)","redundant":"merge into core (5 groups)", + "plotting-beyond-core":"move to examples / drop","model-construction":"keep, split across model/_module,_priors,_classifier,_training", + "session-io":"keep as utils/_session","helper":"dedupe → _console/_stats/_distance/_base/_colors", + "dead-broken":"12 delete · 2 merge · 3 keep+fix", + } + for k in order: + A(f"| {k} | {lab.get(k,0)} | {disp.get(k,'')} |") + A(f"| **total** | **{len(merged)}** | |") + A("") + A("## 2. Grafiti reference layout (the target shape)") + A("") + A(GRAFITI) + A("## 3. Target tcri layout (grafiti-mirrored)") + A("") + A(TARGET_TREE) + # full inventory grouped by namespace + A("## 4. Full inventory — every function") + A("") + A("Label is critic-corrected. Disposition folds in the five overlays " + "(diagnostics reclass, keep+fix pl twins, n_samples convention).") + NS_TITLE = {"ml":"ml — model","pp":"pp — preprocessing","tl":"tl — metrics","pl":"pl — plotting","ut":"ut — utils"} + labrank = {l:i for i,l in enumerate(order)} + for ns in ["ml","pp","tl","pl","ut"]: + rows = [r for r in merged if r["namespace"]==ns] + rows.sort(key=lambda r:(labrank.get(r["label"],9), r["name"])) + A(f"### `tcri.{ns}` — {NS_TITLE[ns].split(' — ')[1]} ({len(rows)} records)") + A("") + A("| name | kind | label | disposition | purpose |") + A("|---|---|---|---|---|") + for r in rows: + A(f"| `{mc(r['name'])}` | {r['kind']} | {r['label']} | {mc(disposition(r))} | {mc(short(r['purpose'],140))} |") + A("") + # synthesis sections + A("## 5. Consolidation groups (redundant → core)") + A("") + A("| into | members | rationale |") + A("|---|---|---|") + for g in syn.get("consolidation_groups",[]): + A(f"| `{mc(g.get('into'))}` | {mc(', '.join(g.get('members',[])))} | {mc(short(g.get('rationale'),160))} |") + A("") + A("## 6. Deletions") + A("") + A("| function | reason |") + A("|---|---|") + for d in syn.get("deletions",[]): + nm = d.get("name") + note = mc(short(d.get("reason"),150)) + if "compare_joint_distribution" in nm: + note = "**OVERLAY OVERRIDE → keep as diagnostics PPC** (was: " + note + ")" + A(f"| `{mc(nm)}` | {note} |") + A("") + A("## 7. Helper extraction (dedupe → shared modules)") + A("") + A("| helper | → module | current copies |") + A("|---|---|---|") + for h in syn.get("helper_extraction",[]): + A(f"| {mc(h.get('helper'))} | `{mc(h.get('into_module'))}` | {mc(short(h.get('current_copies'),90))} |") + A("") + A("## 8. Plotting triage") + A("") + pt = syn.get("plotting_triage",{}) + A("**Core (keep):** " + ", ".join(f"`{mc(x.split(' ')[0])}`" for x in pt.get("core_keep",[]))) + A("") + A("**Beyond core (move→examples / drop):**") + A("") + for x in pt.get("beyond_core_drop_or_move",[]): + A(f"- {mc(x)}") + A("") + A("## 9. Rename / readability map") + A("") + A(RENAME_MAP) + A("## 10. Metric conventions & scoping") + A("") + A(CONVENTIONS) + A("## 11. Open questions (decide before executing)") + A("") + oq = syn.get("open_questions") or "" + if isinstance(oq, str) and oq.strip(): + A(oq) + else: + A("_(see synthesis output)_") + A("") + A("---") + A("_Overlays applied on top of the workflow synthesis: (a) `diagnostics/` = PPCs + model validation " + "(`compare_joint_distribution` reclassified from delete → diagnostics PPC seed; `plot_loss`/`plot_archetypes`/PGM relocated); " + "(b) uniform `n_samples=0` point-estimate convention, drop `point_estimate=`; (c) prior-vs-posterior parked; " + "(d) explicit `__all__` re-export (NOT `import *`) per grafiti; (e) full rename map._") + return "\n".join(L) + +open(OUT_MD,"w").write(md()) +print("wrote", OUT_MD, f"({os.path.getsize(OUT_MD)} bytes)") +print("functions:", len(merged), "| consolidation:", len(syn.get('consolidation_groups',[])), + "| deletions:", len(syn.get('deletions',[])), "| helpers:", len(syn.get('helper_extraction',[]))) diff --git a/docs/contract/tcri_api_and_responsibilities.md b/docs/contract/tcri_api_and_responsibilities.md new file mode 100644 index 0000000..1146bab --- /dev/null +++ b/docs/contract/tcri_api_and_responsibilities.md @@ -0,0 +1,791 @@ +# TCRI — API & Function Responsibilities (Final) + +*The definitive, publishable specification for the refactored, grafiti-mirrored `tcri` package (Door A: standalone scverse package). It folds the API-surface draft and the math/stats draft into a single artifact and applies every fix from the plan-correctness, math/stats, prior-vs-mean, missing-links, and GPU/optimization audits **inline** — no known-wrong math survives below. For each function you get: exact final signature and module home; responsibility; the precise math/stats it performs; how every argument dictates that math; and the exact return shape per argument combination.* + +--- + +## 0. Conventions, notation, and resolved decisions + +### 0.1 Layout principle + +Mirrors `grafiti`: one file per topic (never a monolith), private cross-cutting helper modules, explicit `__all__` re-export, **never `import *`**. Five view directories map to short handles exactly as grafiti does — `model→ml`, `tools→tl`, `preprocessing→pp`, `plotting→pl`, `diagnostics→diag` — plus `utils→ut` (tcri-specific session I/O) and a new private `_compute/` numeric+device seam. The `tl` view lives in `tcri/tools/` (grafiti `tools↔tl`), superseding the earlier working name `metrics/`. + +### 0.2 Notation + +| Symbol | Meaning | +|---|---| +| $P$ | number of phenotypes (columns of every joint) | +| $C$ | number of clonotypes (rows of a joint at one covariate) | +| $M$ | number of covariate values (e.g. timepoints) | +| $\mathrm{ct}$ | a $(\text{clonotype},\text{covariate})$ group; `ct_count` = number of them | +| $p_{ct}\in\Delta^{P}$ | learned per-`ct` phenotype distribution; `uns[K.P_CT]`, shape `(ct_count, P)` | +| $m$ | $=\text{normalize}(q\_p\_ct\_raw)=\mathbb{E}_q[p_{ct}]=$ `uns[K.P_CT]` | +| $\ell_i\in\mathbb{R}^P$ | per-cell classifier logits, `obsm[K.X_LOGITS]` (already scaled by classifier temperature) | +| $s$ | local scale, `uns[K.LOCAL_SCALE]` (Dirichlet total concentration; controls draw spread only) | +| $g$ | `gate_prob`, `uns[K.GATE_PROB]` (scalar $\in[0,1]$, or `None` → additive) | +| $\tau_{\text{cls}}$ | classifier temperature, `uns[K.CLASSIFIER_TEMPERATURE]` | +| $T$ | analysis-time `temperature` argument | +| $\varepsilon$ | numerical floor; values stated in situ | + +Index maps: `uns[K.CT_TO_COV]`, `uns[K.CT_TO_C]`, `uns[K.CT_ARRAY]` (per-cell ct index), `uns[K.COV_ARRAY]` (per-cell covariate index). + +### 0.3 The one substrate every metric reads + +Training learns a variational Dirichlet posterior per `ct` row: + +$$q(p_{ct})=\mathrm{Dirichlet}(\alpha),\qquad \alpha=\mathrm{clamp}(s\cdot m,\ \min=10^{-3}),\quad m=\text{normalize}(q\_p\_ct\_raw)\in\Delta^{P}.$$ + +Where the clamp is **inactive** (the common case), $\alpha=s\cdot m$, so $\sum_k\alpha_k=s$ and the mean is closed-form: + +$$\mathbb{E}_q[p_{ct}]=\frac{\alpha}{\sum_k\alpha_k}=m=\texttt{get\_p\_ct()}=\texttt{uns[K.P\_CT]}.$$ + +The local scale $s$ **cancels in the mean** and matters only for the spread of draws. + +> **Persisted-substrate decision (fixes the double-tempering bug).** `uns[K.P_CT]` stores the **raw** normalized posterior mean $m=\text{normalize}(q\_p\_ct\_raw)$ with **no** guide-temperature and **no** analysis-temperature baked in. `TCRIModel.get_p_ct()` returns exactly this at its default `guide_temperature=1.0`. The analysis-time `temperature` argument (§0.9) is therefore the **single** tempering knob; it is never composed on top of a pre-tempered vector. + +### 0.4 RESOLVED — the point estimate is the closed-form posterior MEAN (prior vs mean vs MAP vs draw) + +This is the audit's required decision. `n_samples=0` returns the **closed-form variational posterior mean** $\mathbb{E}_q[p_{ct}]=m=\texttt{uns[K.P\_CT]}=\texttt{get\_p\_ct()}$ (temperature-adjusted per §0.9). It is read directly and **never sampled**. + +Options analyzed and their disposition: + +| Option | Verdict | Reason | +|---|---|---| +| **(a) Closed-form posterior mean** $\mathbb{E}_q[p_{ct}]=m$ | **ADOPTED** | Exact, closed-form, deterministic, reproducible; already stored in `uns[K.P_CT]` (zero sampling cost). It is the Bayes point estimator under squared-error loss and lies in the simplex interior. Invariant to $s$ (which only sets spread). | +| (b) Generative prior `clone_phen_prior` / archetype `mixture_concentration` | **REJECTED** | It is guide **initialization** / generative anchor, not what training learned; built from argmax hard labels (leakage/circularity); indexed at clone/archetype level, not the `ct` level metrics need. No metric may read it. | +| (c) MAP / posterior mode $(\alpha-1)/(\sum\alpha-P)$ | **REJECTED** | $\alpha_k=s\,m_k$ is routinely $<1$ (small $s$, $m_k<1$), so the mode sits on the simplex boundary / is undefined — unstable, discontinuous in $s$. | +| (d) Mean of many Dirichlet draws | **REJECTED** | Converges to (a) only up to Monte-Carlo noise; a strictly Rao-Blackwell-dominated, non-reproducible estimator of a quantity available in closed form. | + +**Fate of the `posterior=` argument.** The *only* real difference between today's two engines was never prior-vs-posterior — both already used the posterior mean of $p_{ct}$; neither ever touched the generative prior. The real axis is **whether per-cell classifier logits $\ell_i$ are folded in**. Therefore: + +- **`posterior=` is DELETED from all four metrics and from `phenotypic_flux`.** They always use the learned posterior (mean at `n_samples=0`, draws at `n_samples>0`) and, given an `adata`, compute the joint with logits folded in. +- **On the engine `joint_distribution`, the flag survives but is REDEFINED and RENAMED to `use_logits`** (alias `cell_informed=`), replacing both `posterior=` and the old `combine_with_logits=`. It is a **classifier-mixing switch**, not a prior/posterior switch: `use_logits=True` folds per-cell logits into $\log(\text{base})$ exactly like `predict()`; `use_logits=False` returns the `ct`-level table directly. **Both branches use the posterior mean/draws of $p_{ct}$; neither ever touches the generative prior.** The dead `mutual_information(posterior=False) → NotImplementedError` branch and the "prior path" are removed, not implemented. + +### 0.5 Uniform sampling convention (`n_samples`) + +| `n_samples` | Operation | +|---|---| +| `0` | **Deterministic point estimate.** Use the posterior mean $m$ (temperature-adjusted); **no Dirichlet draw**; bit-reproducible on repeat calls. Fixes today's latent bug where `mutual_information`/`flux` at `n_samples=0` returned *one random draw*. | +| `N>0` | **$N$ i.i.d. posterior draws** from the **exact guide posterior** $p_{ct}^{(s)}\sim\mathrm{Dirichlet}\big(\mathrm{clamp}(s\cdot\tilde p_{ct},\ \min=10^{-3})\big)$, seeded (§0.11). Adds a sample axis; posterior mean/HDI of the functional fall out. | + +The old `point_estimate=` argument is **deleted**; `n_samples` is the only point-vs-draws knob. + +> **Default `n_samples = 250`** on the engine, all four metrics, and the sampling `diag` functions — every default call **samples** and reports mean + interval (the honest posterior). `n_samples=0` (the deterministic point estimate) is **opt-in** for speed. (HDI is stabler at `n_samples ≳ 500` near boundaries — tunable.) Signatures ship this default; the `n_samples=0` shown in per-function code blocks below denotes the point-estimate *identity*, not the default. + +> **Clamp fix (blocking).** Draws use the guide's clamped concentration `clamp(local_scale·m̃, 1e-3)` — **not** the bare `local_scale·m̃` or `local_scale·m̃+1e-8` variants that appear in today's three inconsistent engines and summarize a distribution the model never learned. See §0.10 for the induced (documented, intentional) mean discrepancy on committed clones. + +### 0.6 Estimator honesty — plug-in vs posterior-mean (Jensen gap) + +Entropy, MI, and KL/L1 flux are **nonlinear** functionals of $p$, so $\text{metric}(\mathbb{E}_q[p])\neq\mathbb{E}_q[\text{metric}(p)]$; the difference is a **Jensen gap**, *not* Monte-Carlo noise. Consequently the two numbers below are **different estimators by design** and must be documented and tested as such: + +- **`n_samples=0`** computes the **plug-in-at-posterior-mean** estimator $\text{metric}(m̃)$. +- The **`mean` summary column of `n_samples>0`** estimates the **posterior mean of the functional** $\mathbb{E}_q[\text{metric}(p)]$. + +Directions of the gap (Shannon entropy concave; L1/KL flux convex): + +| Metric | Relationship | +|---|---| +| clonotypic / phenotypic entropy | plug-in $\ge$ posterior-mean (over-estimates) | +| phenotypic flux (L1, KL) | plug-in $\le$ posterior-mean (under-estimates); a clone with no real shift reads exactly $0$ at `n_samples=0` but strictly $>0$ in the `n_samples>0` mean | +| mutual information $=H(\phi)-H(\phi\mid c)$ | gaps partly cancel; **sign indeterminate**, magnitudes differ | + +Additionally, because draws use the **clamped** concentration while the `n_samples=0` base is the **unclamped** mean $m$, the two also differ on **committed clones** (where $s\,m_k<10^{-3}$) by a second, clamp-induced term. **No conformance test may assert `n_samples=0 == mean(n_samples>0)`.** Docstrings label the two estimators distinctly. + +### 0.7 Uniform return-shape rule + +| `groupby` | `n_samples` | Return | +|---|---|---| +| unset | `0` | scalar (MI, single-clone flux) or `Series` (entropies over phenotypes/clones) | +| unset | `N>0` | draw array with a sample axis **+** summary columns `mean, sd, hdi_low, hdi_high` | +| set | `0` | tidy `DataFrame`, one row per group [× phenotype / × clone] | +| set | `N>0` | tidy `DataFrame`, one row per group [× phenotype / × clone] + `mean, sd, hdi_low, hdi_high` | + +> **`p_gt` fix (blocking).** `p_gt`/`P(>0)` is **removed from every single-metric summary**. Entropy, MI, L1 and KL flux are all **$\ge 0$**, so `P(draw>0) ≈ 1` always and is vacuous. A signed-direction probability is emitted **only** by `tl.compare_groups` on a between-group **difference** $\Delta$ (§7.6), where it is meaningful. + +> **HDI fix.** The interval columns are a **true highest-density interval** (`hdi_low/hdi_high`), i.e. the narrowest interval containing `hdi_prob` mass — **not** the equal-tailed `np.percentile(x,[2.5,97.5])` mislabeled "HDI" in today's code. For the bounded, right-skewed entropy/flux posteriors (mass piled against the boundary for committed clones) the equal-tailed interval is materially wrong. HDIs from few hundred draws near a boundary are documented as unstable. + +> **Draw-coherence rule (correctness).** For `n_samples>0`, **all clones within one sample share the same $p_{ct}$ draw** (one coherent joint per sample). Metrics iterate the `sample_id` level and compute the full-joint metric per draw, then summarize — never independent per-clone draws. + +### 0.8 The `weighted` axis — KEPT as a dial (default `False`) + +`weighted` is **retained** on the engine and all four metrics, **default `False`**: each clonotype is one unit on the simplex (a **repertoire-level / per-clonotype** statistic). `weighted=True` recovers the **cell-weighted** statistic (large clones dominate — a cell-level statistic). The two answer different biological questions (repertoire structure vs clonal expansion), can reorder samples, and are both valid — the choice is the user's, stated in each docstring. The refactor **fixes the current weight-lookup bug** (a `ct`-indexed `Counter` keyed with clone indices) and unifies the two engines so weighting is applied consistently (removing the inconsistency where `joint_distribution` normalized to sum 1 while `joint_distribution_posterior` returned un-normalized counts). `min`/uncertainty-coefficient normalization (§7.4) is robust to either mode (`MI ≤ min(H_c, H_p)` holds regardless). + +> **Behavior-change note (changelog).** The default is now `weighted=False`, flipping `pl.mutual_information`'s displayed MI from the old cell-weighted default to per-clonotype. `weighted=True` restores the old behavior. + +### 0.9 Temperature — single knob, one consistent placement + +$T$ power-tempers the base **once**, identically in the mean and draw paths: + +$$\tilde p_{ct}=\mathrm{softmax}\!\Big(\tfrac1T\log(m+\varepsilon)\Big)=\frac{m^{1/T}}{\sum_\phi m_\phi^{1/T}},\qquad \varepsilon=10^{-8}.$$ + +$T=1$ is the identity (renormalization only); $T<1$ sharpens, $T>1$ flattens. This fixes today's split where the two engines tempered at different stages. + +- For `use_logits=True`, the combined per-cell logit is divided by $T$ **once**: $P(\phi\mid i)=\mathrm{softmax}\big(\text{combine}(\ell_i,\log\tilde b)/T\big)$. At **$T=1$** this reproduces `predict()` **bit-for-bit** (classifier temperature $\tau_{\text{cls}}$ is already baked into $\ell_i$; no second division). $T\neq1$ is an analysis-time temper that intentionally diverges from `predict()`; documented. +- For `n_samples>0`, draws are centered on the **re-tempered** $\tilde p_{ct}$; docstrings state that $T\neq1$ makes the sampled distribution a re-tempered object, not the raw learned posterior. + +### 0.10 Reproducibility / seeding + +`random_state` (`int | numpy.Generator | torch.Generator | None`) is added to `joint_distribution`, all four metrics, `compare_groups`'s bootstrap, and the sampling `diag` functions. It seeds a **`torch.Generator`** (and, on the GPU path, the CUDA RNG) because draws are `torch` Dirichlet — fixing the standing no-op where `seed` only touched `np.random`. `n_samples=0` is deterministic regardless of `random_state`. + +--- + +## 1. Package tree + +``` +tcri/ + __init__.py # explicit re-export + sys.modules aliases (tl/pp/pl/ml/ut/diag); top-level joint_distribution; NO import * + _keys.py # single source of every uns/obsm/obs key string (constants only) + _console.py # leveled, silenceable logging over scanpy.logging (no raw ANSI, no _ascii_hist) + _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) + _reduce.py # batched entropy / mutual-information / distance reductions over the stack + model/ # ml + _model.py # TCRIModel + _module.py # TCRIModule (pyro model/guide, get_latent, get_p_ct) + _priors.py # MixtureDirichlet, VampPrior + _classifier.py # PhenotypeClassifier + _training.py # UnifiedTrainingPlan, build_archetypes + preprocessing/ # pp + _register.py # registration writers behind TCRIModel.to_anndata (all private) + _clones.py # group_singletons, clone_size + tools/ # tl (mirrors pl by filename) + _joint.py # joint_distribution (THE ENGINE) + _entropy.py # clonotypic_entropy, phenotypic_entropy + _mutual_information.py # mutual_information (+ private _mi_from_joint) + _flux.py # phenotypic_flux + _compare.py # compare_groups (public group-comparison orchestrator) + plotting/ # pl + _base.py # _metric_boxplot, _finish + _colors.py # tcri_colors, resolve_palette + _entropy.py # clonotypic_entropy, phenotypic_entropy + _mutual_information.py # mutual_information + _flux.py # phenotypic_flux (sankey) + _sankey.py # SankeyNode, _phenotype_mass_per_clone (private) + diagnostics/ # diag (NEW) + _ppc.py # joint_distribution_ppc, phenotype_calibration, reconstruction_ppc, permutation_null + _training.py # loss, archetypes + utils/ # ut + _session.py # save_tcri_session, load_tcri_session (+ private helpers) +``` + +**Dropped entirely** (deleted — not relocated anywhere): `top_clone_umap`, `clone_size_umap`, `plot_phenotype_probabilities`, `gene_entropy`, `polar_plot`, `probability_ternary`. Out of the package to `docs/` (a figure script only): the model PGM (`build_nested_tcri_pgm` / `draw_tcri_pgm_nested`). + +--- + +## 2. Top-level `__init__.py` and the `__all__` story + +`__all__` is declared at **both** levels (grafiti pattern). Every impl module declares its own `__all__`; every view `__init__` imports symbols by name and re-declares an aggregate `__all__`. The root imports the six view packages, aliases them into `sys.modules`, and re-exports `joint_distribution` for prominence. **No `import *` anywhere** — numpy/pandas/torch and every `_helper` stay unexported; GPU libs are never imported at module top (§4.3). + +```python +# tcri/__init__.py +from importlib.metadata import PackageNotFoundError, version as _version +try: + __version__ = _version("tcri") +except PackageNotFoundError: + __version__ = "0.0.0+unknown" + +from . import tools as tl +from . import preprocessing as pp +from . import plotting as pl +from . import model as ml +from . import diagnostics as diag +from . import utils as ut +from .tools import joint_distribution # tcri.joint_distribution + +import sys +sys.modules.update({f"{__name__}.{m}": globals()[m] + for m in ("tl", "pp", "pl", "ml", "diag", "ut")}) + +__all__ = ["tl", "pp", "pl", "ml", "diag", "ut", "joint_distribution", "__version__"] +``` + +| View `__init__` | `__all__` | +|---|---| +| `tools/__init__.py` (`tl`) | `joint_distribution`, `clonotypic_entropy`, `phenotypic_entropy`, `mutual_information`, `phenotypic_flux`, `compare_groups` | +| `preprocessing/__init__.py` (`pp`) | `group_singletons`, `clone_size` | +| `plotting/__init__.py` (`pl`) | `clonotypic_entropy`, `phenotypic_entropy`, `mutual_information`, `phenotypic_flux`, `tcri_colors`, `resolve_palette` | +| `model/__init__.py` (`ml`) | `TCRIModel` | +| `diagnostics/__init__.py` (`diag`) | `joint_distribution_ppc`, `phenotype_calibration`, `reconstruction_ppc`, `permutation_null`, `loss`, `archetypes` | +| `utils/__init__.py` (`ut`) | `save_tcri_session`, `load_tcri_session` | + +Private modules (`_keys`, `_console`, `_stats`, `_distance`, `_compute/*`) and every private symbol (`_mi_from_joint`, `_metric_boxplot`, `_finish`, `SankeyNode`, `_phenotype_mass_per_clone`, all `_register.py` writers) are **absent** from every `__all__`. + +--- + +## 3. Shared private helper modules + +### 3.1 `tcri/_keys.py` — canonical AnnData key registry (constants only) + +Import as `from tcri import _keys as K`; no key literal lives anywhere else. + +| Constant | Value | Slot | Meaning | +|---|---|---|---| +| `K.METADATA` | `"tcri_metadata"` | `uns` | dict: `covariate_col`, `clone_col`, `phenotype_col`, `batch_col` (single scheme; replaces dual `tcri_clone_key`/`tcri_phenotype_key`) | +| `K.PHENOTYPE_CATEGORIES` | `"tcri_phenotype_categories"` | `uns` | ordered phenotype categories | +| `K.CLONOTYPE_CATEGORIES` | `"tcri_clonotype_categories"` | `uns` | ordered clonotype categories | +| `K.COVARIATE_CATEGORIES` | `"tcri_covariate_categories"` | `uns` | ordered covariate categories | +| `K.P_CT` | `"tcri_p_ct"` | `uns` | `(ct_count, P)` **raw** posterior mean $m=\mathbb{E}_q[p_{ct}]$ (no temper baked in) | +| `K.CT_TO_COV` | `"tcri_ct_to_cov"` | `uns` | `(ct_count,)` ct→covariate index | +| `K.CT_TO_C` | `"tcri_ct_to_c"` | `uns` | `(ct_count,)` ct→clonotype index | +| `K.CT_ARRAY` | `"tcri_ct_array_for_cells"` | `uns` | `(n_obs,)` per-cell ct index | +| `K.COV_ARRAY` | `"tcri_cov_array_for_cells"` | `uns` | `(n_obs,)` per-cell covariate index | +| `K.LOCAL_SCALE` | `"tcri_local_scale"` | `uns` | scalar Dirichlet total concentration $s$ (draws only) | +| **`K.GATE_PROB`** | `"tcri_gate_prob"` | `uns` | **NEW** — scalar gate $g\in[0,1]$ or `None`; required for gate-aware `use_logits=True` parity with `predict()` | +| **`K.CLASSIFIER_TEMPERATURE`** | `"tcri_classifier_temperature"` | `uns` | **NEW** — $\tau_{\text{cls}}$; persisted for reproducibility/auditing (already baked into stored logits) | +| `K.X_LATENT` | `"X_tcri"` | `obsm` | `(n_obs, n_latent)` encoder posterior-mean latent | +| **`K.X_LOGITS`** | `"X_tcri_logits"` | `obsm` | `(n_obs, P)` classifier logits — **part of the canonical write-set** (§5.1); the `use_logits=True` engine path requires it | +| `K.X_PROBABILITIES` | `"X_tcri_probabilities"` | `obsm` | `(n_obs, P)` per-cell phenotype probabilities (`predict()`) | +| `K.PHENOTYPE_LABEL` | `"tcri_phenotype"` | `obs` | per-cell argmax hard label | +| `K.CLONE_SIZE` | `"clone_size"` | `obs` | per-cell clone cell-count | +| `K.OBS_INDICES` | `"indices"` | `obs` | per-cell integer index — **registration glue** written by `setup_anndata` (§5.1); not analysis output | + +> The AnnDataManager is **no longer** stashed in `uns` (`tcri_manager` retired) — deleting the `write_adata_safely` / `_pop_nonserializables` hack. The stash lived in `setup_anndata`, so it is removed there (§5.1). + +### 3.2 `tcri/_console.py` — leveled, silenceable logging (private) + +Thin wrappers over `scanpy.logging`; respects scanpy verbosity. Raw ANSI prints and `_ascii_hist` (and every `graph=`/ASCII-histogram code path) are **deleted**. + +| Signature | Responsibility | +|---|---| +| `info(msg, *, deep=None)` | `scanpy.logging.info`; silenced by scanpy verbosity. | +| `warning(msg)` | `scanpy.logging.warning`. | +| `success(msg)` | `scanpy.logging.hint`. | +| `done(msg="done")` | terminal completion line. | + +### 3.3 `tcri/_stats.py` — significance + posterior-comparison statistics (private) + +| Signature | Responsibility / math | +|---|---| +| `stars(p)` | $p<10^{-4}\to$`****`; $<10^{-3}\to$`***`; $<10^{-2}\to$`**`; $<0.05\to$`*`; else `ns`. | +| `mann_whitney(a, b, *, alternative="two-sided")` | Mann–Whitney $U=\sum_{i,j}\mathbb1[a_i>b_j]+\tfrac12\mathbb1[a_i=b_j]$; two-sided $p$ from the rank-sum null (`scipy.stats.mannwhitneyu`). Returns `(U, p)`. | +| `prob_direction(delta)` | **Signed-contrast only.** Given a paired difference-draw vector $\Delta$: $p_{gt}=\frac1N\sum_s\mathbb1[\Delta^{(s)}>0]$, $p_{lt}=\frac1N\sum_s\mathbb1[\Delta^{(s)}<0]$. Returns `(p_gt, p_lt)`. | +| `hdi(samples, *, hdi_prob=0.94)` | **True** highest-density interval: over sorted samples, the **narrowest** window containing $\lceil hdi\_prob\cdot N\rceil$ points. Returns `(low, high)`. Documented unstable near a bounded posterior's boundary for small $N$. | +| `summarize(samples, *, hdi_prob=0.94)` | Reducer for a **raw metric** draw vector → `{mean, sd, hdi_low, hdi_high}`. **No `p_gt`** (vacuous for $\ge0$ metrics, §0.7). | +| `auc_and_label_permutation(scores, labels, *, pos_label=None, n_perm=200_000, seed=42, max_exact=200_000)` | Observed ROC-AUC + two-sided permutation $p$: exact enumeration when $\binom{n}{k}\le$`max_exact`, else Monte-Carlo; $p_{\text{perm}}=\text{mean}(|\mathrm{AUC}_{\text{perm}}-0.5|\ge|\mathrm{AUC}_{\text{obs}}-0.5|)$. Returns `(auc, p, perm_stats, mode)`. | +| `bootstrap_auc(scores, labels, *, pos_label=None, n_boot=5000, seed=42)` | Resample cells with replacement (reject draws missing a class), recompute AUROC, return the 2.5/97.5 quantiles. Returns `np.array([lo, hi])`. | + +### 3.4 `tcri/_distance.py` — phenotype-distribution distances (private) + +Dedupes the old module-level `dkl` and `flux.dkl_func`; **one base (bits, $\log_2$) and one $\varepsilon=10^{-12}$ library-wide**, matching entropy/MI. + +| Signature | Responsibility / math | +|---|---| +| `l1_distance(p, q)` | $\sum_i|p_i-q_i|$; symmetric; range $[0,2]$ on the simplex. Defensively renormalizes inputs. Returns float. | +| `kl_divergence(p, q, *, base=2.0, eps=1e-12)` | $\mathrm{KL}(p\Vert q)=\sum_i p_i\log_2\frac{p_i}{q_i}$; clip to $[\varepsilon,1]$ then renormalize. **Asymmetric**, range $[0,\infty)$. **Single base fixed to $\log_2$ / single $\varepsilon$** (repairs the natural-log-vs-$\log_2$ and $10^{-10}$-vs-$10^{-15}$ divergence between the two dead copies). Returns float. | +| `js_divergence(p, q, *, base=2.0, eps=1e-12)` | **NEW** Jensen–Shannon $=\tfrac12\mathrm{KL}(p\Vert r)+\tfrac12\mathrm{KL}(q\Vert r)$, $r=\tfrac12(p+q)$; **symmetric, bounded $[0,1]$ bit** — the recommended symmetric shift measure. Returns float. | +| `phenotype_distance(p, q, *, metric="l1")` | Dispatcher: `"l1"`→`l1_distance`; `"kl"`/`"dkl"`→`kl_divergence` (directional, unbounded, bits); `"jsd"`→`js_divergence`; a callable `f(p,q)`; else `ValueError`. | + +--- + +## 4. `tcri/_compute/` — numeric + device seam (NEW, private) + +The engine's numeric core is written **once** as a batched, device-routable function so the acceleration is additive and reversible, and so the four metrics share one joint-draw stack. + +### 4.1 `_xp.py` — the device seam (grafiti reference, copied 1:1) + +| 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. | +| `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. +- **`_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) + +Lazy GPU imports (never at module top — `import tcri` never touches a GPU lib; the old module-top `import umap` is moved inside its function); GPU deps never in `install_requires` (CPU path always fully functional); permissive device ladder with device-count verification; `asnumpy`/`output_type="numpy"` at every boundary; `try/except` degrade-to-CPU reporting which backend ran; `float64` where parity must hold; validate-before-compute on **per-row** invariants; chunked reductions to bound host+device memory on the large `[n_samples, n_cells, P]` tensor. `n_samples=0` performs **zero** draws (closed-form read of `uns[K.P_CT]`). cuML UMAP for the one-off latent embedding sits behind a `_use_gpu(device)` gate with the `umap-learn` CPU fallback; layouts differ (both valid), documented, not claimed bit-identical. + +--- + +## 5. `tcri.ml` — model (`model/`) + +### 5.1 `model/_model.py` + +**`class TCRIModel(BaseModelClass)`** — register → build → train → extract → write. + +| Method (signature) | Responsibility | +|---|---| +| `@classmethod setup_anndata(cls, adata, *, layer=None, clonotype_key="unique_clone_id", phenotype_key="phenotype_col", covariate_key="timepoint", batch_key="patient", **kwargs)` | **Registration only** — register clonotype/phenotype/covariate/batch/count fields with scvi and store the layer. **Writes `obs["indices"]=range(n)` and registers it** (`CategoricalObsField`) — this is registration glue that `training_step`/`validation_step` consume via `batch["indices"]`; it is **not** analysis output. Invariant is **"no analysis/label `obs` mutation"** (labels/probabilities are written only by `to_anndata`). **Removes the `uns["tcri_manager"]` stash** (was here), deleting the need for `write_adata_safely`. | +| `__init__(self, adata, *, n_latent=128, n_hidden=128, n_layers=3, classifier_n_layers=3, global_scale=5.0, local_scale=3.0, prior_temperature=1.0, guide_temperature=1.0, use_enumeration=False, patience=300, classifier_hidden=128, classifier_dropout=0.1, n_pseudo_obs=10, K=10, phenotype_weights=None, gate_prob=None, kl_weight_max=1.0, guide_init_scale=10.0, classifier_temperature=1.0, **kwargs)` | Build the empirical clone→phenotype prior + KMeans archetypes + clonotype/covariate index maps + class weights, then construct/prime `TCRIModule`. Note `gate_prob=None` default ⇒ ungated model; the gate-parity guarantee is only exercised when a gate is trained. | +| `train(self, *, max_epochs=1000, batch_size=1000, lr=1e-3, reconstruction_loss_scale=1e-3, n_steps_kl_warmup=2000, **kwargs)` | 0.9/0.1 split, `UnifiedTrainingPlan`, `TrainRunner` with `elbo_validation` early stopping. | +| `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. | + +> Relocated off the model: `plot_archetypes`→`diag.archetypes`; `plot_loss`→`diag.loss`. `boost_phenotype_prior`, `use_gate` remain internal. + +### 5.2 `model/_module.py` + +**`class TCRIModule(PyroBaseModuleClass)`** — Pyro CVAE with hierarchical clonotype→(clonotype×covariate) Dirichlet priors and a phenotype classifier. *(internal)* + +| Member (signature) | Responsibility | +|---|---| +| `__init__(self, n_input, n_latent, P, n_batch, *, global_scale=10.0, local_scale=5.0, prior_temperature=1.0, guide_temperature=1.0, gate_prob=0.5, mixture_concentration=None, n_pseudo_obs=10, use_enumeration=False, classifier_hidden=128, classifier_dropout=0.1, classifier_n_layers=3, n_hidden=128, n_layers=3, class_weights=None, kl_weight_max=1.0, guide_init_scale=10.0, classifier_temperature=1.0)` | Construct encoder/decoder/classifier/VampPrior, `px_r`; register empty two-level buffers + class weights (`mixture_concentration` required). | +| `prepare_two_level_params(self, clone_phen_prior_mat, ct_to_c, ct_to_cov, ct_array_for_cells, cov_array_for_cells, *, eps=1e-6)` | Normalize/temperature the clone-phenotype prior; register two-level index buffers. | +| `model(self, x, batch_idx, log_library, ...)` | Generative: sample $p_c$ (MixtureDirichlet), $p_{ct}$ (Dirichlet centered at $p_c$), latent $z$ (VampPrior), ZINB gene obs. | +| `guide(self, x, batch_idx, log_library, ...)` | Guide: learnable Dirichlet params $q(p_c)$, $q(p_{ct})$ with **`clamp(min=1e-3)`** on the concentration (the floor the draw path must reproduce, §0.5); Normal $q(z)$ from the encoder. | +| `get_latent(self, tensor_dict) -> torch.Tensor` | Encode a batch to posterior-mean latent $z_{\text{loc}}$. | +| `get_p_ct(self) -> torch.Tensor` | Read `q_p_ct_raw` from the (process-global) param store; return the row-normalized `(ct_count, P)` posterior mean. | +| `use_gate(self) -> bool` (property) | `True` when `gate_prob is not None`. | +| `@staticmethod _get_fn_args_from_batch(tensor_dict) -> tuple` | Extract `(x, batch_idx, log_library)` from a scvi batch dict. | + +> **Param-store caveat (documented).** `get_p_ct` reads the **process-global** Pyro param store (`q_p_ct_raw`); loading two sessions in one process clobbers it. `to_anndata` and every `diag` PPC must be called immediately after the intended model's params are set; `load_tcri_session` sets the store before any `get_p_ct`/`to_anndata` call. Single-model-per-process otherwise. + +### 5.3 `model/_priors.py` + +**`class MixtureDirichlet(dist.TorchDistribution)`** *(internal)* — clonotype prior $p_c$. Members: `__init__(self, mixture_weights, concentration, validate_args=None)`; `sample`; `log_prob` (log-sum-exp of component Dirichlet log-probs); `score_parts` (returns `(log_prob, 0, 0)` → reparam-free); `__call__` (alias for `sample`). + +**`class VampPrior(torch.nn.Module)`** *(internal)* — VampPrior over $z$. Members: `__init__(self, pseudo_inputs, encoder)`; `get_mixture` (uniform `MixtureSameFamily` of `Independent` Normals); `log_prob(self, z)`; `sample`. + +### 5.4 `model/_classifier.py` + +**`class PhenotypeClassifier(nn.Module)`** *(internal)* — `__init__(self, n_latent, classifier_hidden, P, *, num_layers=3, dropout_rate=0.1, temperature=1.0)`; `forward(self, x)` returns MLP logits divided by `temperature` ($\tau_{\text{cls}}$, baked into the stored logits). + +### 5.5 `model/_training.py` + +**`class UnifiedTrainingPlan(PyroTrainingPlan)`** *(internal)* — `__init__(self, module, *, n_steps_kl_warmup=1000, reconstruction_loss_scale=1e-2, num_particles=5, optimizer_config=None, class_weights=None, **kwargs)`; `loss` (property); `configure_optimizers`; `training_step`; `validation_step` (logs `elbo_validation`). + +**Module function:** `build_archetypes(c2p_mat, *, K=10) -> tuple[np.ndarray, np.ndarray]` — KMeans-cluster clone→phenotype rows into `K` normalized archetype centroids. **Returns `(centers, labels)`** — labels are retained so `diag.archetypes` can reproduce the cluster-ordered heatmap. `K` default is **10** (aligned to `TCRIModel`, repairing the former `K=4` default mismatch). + +--- + +## 6. `tcri.pp` — preprocessing (`preprocessing/`) + +### 6.1 `preprocessing/_clones.py` — public + +| Signature | Responsibility | +|---|---| +| `group_singletons(adata, *, clonotype_key="trb", groupby="patient", target_col="trb_unique", min_clone_size=10) -> AnnData` | Collapse clones smaller than `min_clone_size` (per `groupby`) into `"Singleton_{group}"` labels in `target_col`. **Ordering invariant (documented + enforced):** any clone relabeling must run **before** `setup_anndata`/`train`, else the learned clonotype categories and `p_ct`'s `ct_to_c` map desync from `obs`; `setup_anndata` refuses registration if a later relabel is detected. | +| `clone_size(adata, *, key_added="clone_size", return_counts=False)` | Per-clone cell counts, written per cell into `obs[key_added]`. **Reads `uns[K.METADATA]["clone_col"]`** (migrated off the retired `tcri_clone_key` in the same change that stops writing it). | + +**Private inner:** `group_singletons.collapse_singleton(row)`. + +### 6.2 `preprocessing/_register.py` — private (the `to_anndata` writers) + +Called only by `TCRIModel.to_anndata`; folds in the old `register_phenotype_key` / `register_clonotype_key` / `_compute_logits_and_prior`. + +| Signature | Responsibility | +|---|---| +| `_write_metadata(adata, model)` | `uns[K.METADATA]` (single scheme) + the three category lists from the registry. | +| `_register_clonotype_key(adata, clonotype_key, *, order=None)` | Register the clonotype `obs` column + ordered categories. | +| `_register_phenotype_key(adata, phenotype_key, *, order=None)` | Register the phenotype `obs` column + ordered categories. | +| `_write_latent(adata, model, *, latent_key="X_tcri", batch_size=256)` | Encoder posterior-mean latent → `obsm`. | +| `_write_logits(adata, model, *, logits_key="X_tcri_logits", batch_size=256)` | **Per-cell classifier logits → `obsm[K.X_LOGITS]`** (canonical; required by the default engine path). | +| `_write_predictions(adata, model, *, predictions_key="X_tcri_probabilities", label_key="tcri_phenotype", batch_size=256)` | `predict()` probs → `obsm`; argmax hard labels → `obs`. | +| `_write_p_ct(adata, model)` | `p_ct`, `ct_to_cov`, `ct_to_c`, per-cell ct/cov arrays, `local_scale`, **`gate_prob`**, **`classifier_temperature`** → `uns`. | +| `_compute_logits_and_prior(model, adata, *, batch_size=256, eps=1e-8) -> tuple[np.ndarray, np.ndarray]` | Run encoder+classifier to extract per-cell logits and $\log p_{ct}$ from `get_p_ct()`. | + +--- + +## 7. `tcri.tl` — tools / metrics (`tools/`) + +### 7.1 Engine — `tools/_joint.py` + +```python +joint_distribution( + adata, *, + covariate=None, # None → ALL covariate values in one pass (shared draw) + groupby=None, + n_samples=0, + use_logits=True, # was posterior=; alias cell_informed=; classifier-mixing switch + clones=None, + temperature=1.0, + random_state=None, + device=None, +) -> pandas.DataFrame +``` +Re-exported top-level as `tcri.joint_distribution`. Unifies today's `joint_distribution` + `joint_distribution_posterior`. + +**(a) Responsibility.** Produce the clone×phenotype distribution (the substrate every metric consumes) at a covariate value from the learned variational posterior of $p_{ct}$ — a deterministic posterior-mean table or posterior draws. Provenance in `df.attrs["params"]` (and, for cache-friendliness, in a flat `_provenance` sidecar column, §7.7). + +**(b) Math.** Select `ct` rows for covariate $m$ via `ct_to_cov`; each clone maps to exactly one `ct`, so rows index by clonotype. Temperature-temper the base **once** (§0.9): $\tilde p_{ct}=\mathrm{softmax}(\tfrac1T\log(m+10^{-8}))$. + +*Base per point/draw:* +- `n_samples=0`: base $b=\tilde p_{ct}$ (posterior mean; deterministic). +- `n_samples=N`: bases $b^{(s)}\sim\mathrm{Dirichlet}\big(\mathrm{clamp}(s\cdot\tilde p_{ct},\ \min=10^{-3})\big)$, $s=$`local_scale`, seeded (§0.10). One coherent joint per sample (§0.7). + +*`use_logits=False` (ct-level table):* row $c=b$. For `n_samples=0` this is exactly $\tilde p_{ct}$; at $T=1$ it equals `uns[K.P_CT]` restricted to the covariate — the clean closed-form identity used in tests. + +*`use_logits=True` (fold per-cell logits; matches `predict()`):* per cell $i$ with clonotype $c(i)$, group $\mathrm{ct}(i)$, base $b_{\mathrm{ct}(i)}$: +$$P(\phi\mid i)=\mathrm{softmax}\!\Big(\tfrac1T\,\mathrm{combine}\big(\ell_i,\ \log(b_{\mathrm{ct}(i)}+\varepsilon)\big)\Big),\quad +\mathrm{combine}=\begin{cases}g\,\ell_i+(1-g)\log b & g=\texttt{gate\_prob}\neq\text{None}\\ \ell_i+\log b & \text{additive}\end{cases}$$ +then $J[c,\phi]=\sum_{i\in c}P(\phi\mid i)$, row-normalize. At $T=1$, `n_samples=0`, this equals `predict()` aggregated per clone — **bit-for-bit**, gate-aware (fixes the standing disagreement where metrics used one Dirichlet *draw* and *never* applied the gate). + +*`covariate=None`:* compute the joint for **all** covariate values from a **single shared draw** per sample (the draw-once invariant, §7.8), stacking a covariate axis. + +*`groupby=g`:* restrict by cell/clone masks into the **full-space** `uns` arrays (never by slicing the AnnData — see the guard note below), computing per group value and stacking a group axis. + +**(c) Arguments → math.** + +| Argument | Effect | +|---|---| +| `covariate` | Selects `ct` rows via `ct_to_cov`. `None` → all covariates in one shared-draw pass. | +| `groupby` | Separate joint per group value (adds a group axis), implemented by **restriction over full adata**, not slicing. Requires the cell-informed path or a clone-constant key (see semantics note). | +| `n_samples` | `0` → posterior-mean table (deterministic); `N` → $N$ clamped-Dirichlet draws. Only place `local_scale` enters. | +| `use_logits` | `True` → fold logits with $\log b$ (gate-aware), aggregate per clone, row-normalize; `False` → `ct`-level $\tilde p_{ct}$ rows. Neither is the generative prior. | +| `clones` | Filters rows to the listed clonotypes; with `use_logits=True` also restricts aggregated cells; final reindex to the exact list (absent clones → dropped, **not** all-zero rows — see §7.2 fix). | +| `temperature` | $T$ tempers the base once (§0.9). $T=1$ identity; at $T=1$, `use_logits=True` reproduces `predict()`. | +| `random_state` | Seeds the torch (and CUDA) Dirichlet generator for `n_samples>0`. Ignored at `n_samples=0`. | +| `device` | Routes the numeric core through `_compute/_xp` (CPU / torch-CUDA / cupy); result is always host numpy. | + +**(d) Return shape** — `pandas.DataFrame`, columns = phenotype categories. + +| `covariate` | `groupby` | `n_samples` | Index / axes | +|---|---|---|---| +| set | unset | `0` | rows = clonotype id; `(C_m, P)` | +| set | unset | `N>0` | MultiIndex (clonotype, `sample_id`); `(C_m·N, P)` | +| set | set | `0` | MultiIndex (group, clonotype) | +| set | set | `N>0` | MultiIndex (group, clonotype, `sample_id`) | +| `None` | — | — | adds a leading covariate level to any of the above | + +> **groupby ↔ alignment guard (blocking fix).** `joint_distribution_posterior` hard-raises if per-cell `uns[...array_for_cells]` lengths $\neq$ `n_obs`. Passing a **sliced** AnnData (today's `tcri_boxplot` pattern) trips this. groupby is therefore implemented by **positional cell/clone masks into the full-space `uns` arrays** + `clones=`, never by handing a slice to the engine. `_metric_boxplot` (§8.5) is rewritten off the slice-and-call pattern. + +> **groupby ↔ covariate semantics.** `p_ct` is indexed by `ct=(clonotype, covariate)` only. A `groupby` key that is **not** functionally determined by clonotype-at-fixed-covariate (e.g. a tissue cross-cutting one clone) is unrepresentable in the `use_logits=False` table and is only well-defined on the cell-informed `use_logits=True` path. The engine **requires the cell-informed path for such keys**, or requires the key be clone-nested / constant within a clone×covariate; it errors/warns on `use_logits=False` + a non-clone-determined groupby. The whole per-group scheme assumes **clones are disjoint across groups** (a TCR clone never spans two patients) — stated explicitly. + +`__all__ = ["joint_distribution"]` + +### 7.2 `tools/_entropy.py` — `clonotypic_entropy` + +```python +clonotypic_entropy( + adata_or_jd, *, + covariate=None, groupby=None, + n_samples=0, temperature=1.0, + clones=None, normalized=True, n_clones_ref=None, + random_state=None, device=None, +) -> float | pandas.Series | pandas.DataFrame +``` + +**(a) Responsibility.** For each phenotype $\phi$ (at covariate $m$), the normalized Shannon entropy of the distribution over clonotypes carrying that phenotype, $H[P(c\mid\phi,m)]$ — spread of a phenotype across clones. **Repertoire-level (uniform-clonotype prior; §0.8).** + +**(b) Math.** From joint $J$, take column $\phi$ over the **supported** clones only (absent/all-zero clones excluded — see fix), normalize, take entropy in bits: +$$v_c=\frac{J[c,\phi]}{\sum_{c'\in\text{supp}}J[c',\phi]},\qquad H_\phi=-\sum_{c\in\text{supp}} v_c\log_2 v_c.$$ +If `normalized`: divide by $\log_2 C_{\text{den}}$ where $C_{\text{den}}$ = number of **supported** clones (default), or `n_clones_ref` if given (a fixed reference for cross-group comparability). No division when $C_{\text{den}}\le1$. Base fixed to 2. + +- **Estimator (§0.6):** `n_samples=0` = plug-in $H_\phi(m̃)$; `n_samples>0` `mean` = $\mathbb{E}_q[H_\phi]$ (plug-in $\ge$ posterior-mean for entropy). Documented as distinct. +- **Fixes:** deterministic `n_samples=0` (no forced draw); `normalized` exposed (was hard-normalized); base fixed to 2; **absent/zero-support clones excluded before normalizing** (no $\varepsilon$-clip fabricating uniform mass or inflating $C$). + +**(c) Arguments → math.** + +| Argument | Effect | +|---|---| +| `adata_or_jd` | AnnData → compute $J$ internally via §7.1 (`use_logits=True`; `covariate` required); precomputed joint DataFrame → skip to entropy (fast path — valid only at `n_samples=0`, `groupby=None`; `clones` just re-filters; else `ValueError`, §7.9). | +| `covariate` / `groupby` | Condition $m$; per-group entropy → tidy rows (group × phenotype). | +| `n_samples` | `0` → plug-in per phenotype; `N` → per-draw + summary. | +| `temperature` | Tempers $J$ before the column is read. | +| `clones` | Restricts the clone set → changes support and the default $\log_2 C_{\text{den}}$. | +| `normalized` / `n_clones_ref` | `True` → divide by $\log_2 C_{\text{den}}$ (range $[0,1]$); `n_clones_ref` fixes the denominator for comparability; `False` → raw bits. | +| `random_state` / `device` | Seeding / backend routing for `n_samples>0`. | + +**(d) Return shape.** `n_samples=0`, no `groupby` → `Series` over phenotypes; `n_samples>0`, no `groupby` → per-phenotype `mean, sd, hdi_low, hdi_high`; `groupby` → tidy DataFrame row per (group, phenotype) [+ summary]. **Absent phenotype → `NaN`, not 0.** + +> **Comparability note.** Because the default denominator is group-specific, normalized clonotypic entropy is **within-group** unless `n_clones_ref` (a common denominator) is supplied. The `pl` twin defaults cross-group plots to a common `n_clones_ref`. + +### 7.3 `tools/_entropy.py` — `phenotypic_entropy` + +```python +phenotypic_entropy( + adata_or_jd, *, + covariate=None, groupby=None, + n_samples=0, temperature=1.0, + clones=None, normalized=True, + random_state=None, device=None, +) -> float | pandas.Series | pandas.DataFrame +``` + +**(a) Responsibility.** For each clonotype $c$, the normalized Shannon entropy of its phenotype distribution $H[P(\phi\mid c,m)]$ — plasticity vs commitment. + +**(b) Math.** Row $c$ = $P(\phi\mid c)$; normalize over phenotypes, entropy in bits, divide by $\log_2 P$ if `normalized` and $P>1$: +$$p_\phi=\frac{J[c,\phi]}{\sum_{\phi'}J[c,\phi']},\quad H_c=-\sum_\phi p_\phi\log_2 p_\phi.$$ +Estimator convention as §0.6 (plug-in at `n_samples=0`). **Critical bug fix:** a clone present in `obs` but with **zero posterior mass** returns **`NaN`** and is excluded — it is **not** reindexed to zeros, $\varepsilon$-clipped, and normalized to uniform → spurious $H=1.0$. Base fixed to 2; `normalized` exposed. + +**(c) Arguments → math.** As §7.2, except the normalizer is $\log_2 P$ (depends on $P$, not clone count), so `clones` does not change the divisor; $P\le1\Rightarrow$ divisor 1. + +**(d) Return shape.** `n_samples=0`, no `groupby` → `Series` over clonotypes; `n_samples>0` → per-clone `mean, sd, hdi_low, hdi_high`; `groupby` → tidy DataFrame row per (group, clone) [+ summary]. + +### 7.4 `tools/_mutual_information.py` — `mutual_information` (+ kernel) + +```python +mutual_information( + adata_or_jd, *, + covariate=None, groupby=None, + n_samples=0, temperature=1.0, + clones=None, normalized=True, + normalize_mode="min", # CHANGED default (was "average") + random_state=None, device=None, +) -> float | numpy.ndarray | pandas.DataFrame +``` + +**(a) Responsibility.** $I(c;\phi\mid m)$ in bits — strength of clone–phenotype coupling. Optionally normalized to $[0,1]$. + +**(b) Math — kernel `_mi_from_joint(pxy, *, normalized, mode="min") -> float`.** Normalize the whole joint: $\text{pxy}=J/\sum J$; marginals $p_x=\sum_\phi\text{pxy}$, $p_y=\sum_c\text{pxy}$. +$$I=\sum_{c,\phi}\text{pxy}\,\log_2\frac{\text{pxy}+\varepsilon}{p_x p_y^\top+\varepsilon}\ \text{(bits)},\quad \varepsilon=10^{-15}.$$ +With $H_c=-\sum p_x\log_2 p_x$, $H_p=-\sum p_y\log_2 p_y$: +$$I_{\text{norm}}=\frac{I}{D},\quad D=\begin{cases}\min(H_c,H_p) & \texttt{mode="min"}\ \text{(coefficient of constraint; default)}\\[2pt]\tfrac12(H_c+H_p) & \texttt{mode="average"}\end{cases}$$ +returning 0 if $D\le0$. + +> **Default `normalize_mode="min"` (blocking fix).** Under the uniform-clonotype prior (§0.8), each row sums to 1 and the table is divided by its sum, so $p_x=P(c)=1/C$ **exactly** and $H_c=\log_2 C$ is **structural and maximal**. `"average"` denom $=\tfrac12(\log_2 C+H_p)$ throttles normalized MI by $\sim1/\log_2 C$ and shrinks its ceiling as $C$ grows — non-comparable across groups/covariates with different $C$, breaking the groupby-comparison workflow. `"min"` gives $I/H_p$ (reaches 1 when clone determines phenotype, $C$-independent) and is the default. Docstring states $H_c=\log_2 C$ is not a meaningful normalizer here. + +**Fixes:** `n_samples=0` = deterministic plug-in $I(m̃)$ (was one random draw); `posterior=False → NotImplementedError` deleted (§0.4). Estimator honesty per §0.6 (MI Jensen-gap sign indeterminate). + +**(c) Arguments → math.** As the shared table; additionally `normalize_mode` selects $D$. `clones` restricts rows; `normalized` toggles $I$ vs $I/D$. + +**(d) Return shape.** `n_samples=0`, no `groupby` → scalar `float`; `n_samples>0`, no `groupby` → `(N,)` array + `mean, sd, hdi_low, hdi_high`; `groupby` → tidy DataFrame, one row per group, column `MI` [+ summary]. Fast path (precomputed jd) valid only at `n_samples=0`, `groupby=None`. + +`__all__ = ["mutual_information"]` + +### 7.5 `tools/_flux.py` — `phenotypic_flux` (renamed from `flux`) + +```python +phenotypic_flux( + adata, *, + cov_from, cov_to, + groupby=None, + n_samples=0, temperature=1.0, + clones=None, + distance_metric="l1", + random_state=None, device=None, +) -> pandas.Series | numpy.ndarray | pandas.DataFrame +``` + +**(a) Responsibility.** Per-clonotype distance between a clone's phenotype distribution at `cov_from` vs `cov_to`, over the clone intersection. + +**(b) Math.** $J_{\text{from}}$, $J_{\text{to}}$ (rows $P(\phi\mid c)$); over common clones $c$, $p=J_{\text{from}}[c]$, $q=J_{\text{to}}[c]$: +$$d_c=\text{phenotype\_distance}(p,q,\ \text{metric}=\texttt{distance\_metric}),$$ +dispatched through `_distance` (§3.4): `"l1"` (default, bounded $[0,2]$), `"kl"`/`"dkl"` (directional, unbounded, **bits**), `"jsd"` (symmetric, bounded $[0,1]$ bit), or callable. + +**Fixes carried by the rewrite:** (1) the `posterior=False` dead branch is gone (no more `TypeError` from unsupported `silent=`/`combine_with_logits=` kwargs); (2) **reproducibility** — draws use a seeded **torch** generator (old `seed=` only touched NumPy → no-op); (3) **unit consistency** — KL is $\log_2$/bits, single $\varepsilon$; (4) the `flux_table` column-misalignment (`clones_g.index(cl)` vs `common`-ordered columns) is removed by returning a tidy per-(group,clone) frame keyed by clone id. **Estimator honesty (§0.6):** `n_samples=0` = plug-in $d_c(m̃)$ (convex → under-estimates $\mathbb{E}_q[d_c]$; a clone with no real shift reads exactly 0 at `n_samples=0` but $>0$ in the `n_samples>0` mean). + +**(c) Arguments → math.** + +| Argument | Effect | +|---|---| +| `cov_from`, `cov_to` | The two conditions compared (was `from_this`/`to_that`). | +| `groupby` | Per-group flux → tidy rows (group × clone) with a `clone_size` column (replaces `flux_table`), via full-space restriction. | +| `n_samples` | `0` → deterministic per-clone plug-in `Series`; `N` → $N$ redrawn distance vectors + summary. | +| `temperature` | Tempers both joints identically before differencing. | +| `clones` | Restricts both sides; distances over the intersection. | +| `distance_metric` | `"l1"` / `"kl"` / `"jsd"` / callable. | +| `random_state` / `device` | Seeding / backend. | + +**(d) Return shape.** `n_samples=0`, no `groupby` → `Series` over common clones; `n_samples>0` → per-clone `mean, sd, hdi_low, hdi_high`; `groupby` → tidy DataFrame row per (group, clone) + `clone_size` [+ summary]. + +`__all__ = ["phenotypic_flux"]` + +### 7.6 `tools/_compare.py` — `compare_groups` (public group-comparison orchestrator) + +```python +compare_groups( + df, *, + value, # column holding the per-unit metric value or draw vector + by, # grouping column (e.g. "response") + reference=None, # baseline level; None → all pairwise + paired=False, # True → paired posterior-draw contrast (uses prob_direction) + hdi_prob=0.94, + alternative="two-sided", +) -> pandas.DataFrame +``` + +**Responsibility.** The **public** replacement for the deleted `mi_compare` / `delta_entropy_table` / `flux_table`: turn a tidy `groupby` result (per-unit point estimates, e.g. per patient) or paired posterior-draw vectors into group contrasts. This closes the audit gap where "`groupby` + `_stats` subsumes `*_compare`/`*_delta`" was non-functional because `_stats` is private. + +**Math.** For each contrast (`reference` vs other, or all pairs): +- **Unpaired point estimates:** Mann–Whitney $U$ + two-sided $p$ (`_stats.mann_whitney`), group means, and $\Delta=\text{mean}_B-\text{mean}_A$. +- **Paired posterior draws** (`paired=True`, one draw vector per group per unit, aligned by `sample_id`): the signed difference $\Delta^{(s)}=\text{metric}_B^{(s)}-\text{metric}_A^{(s)}$, then `mean(Δ)`, `hdi(Δ)`, and **`p_gt`/`p_lt` via `prob_direction`** — the **only** place a direction probability is emitted (§0.7). + +**Return.** Tidy DataFrame, one row per contrast: `group_a, group_b, mean_a, mean_b, delta, U, p, p_gt, hdi_low, hdi_high, stars`. Recreates `mi_compare`'s per-pair output exactly. + +`tools/__init__.py __all__ = ["joint_distribution", "clonotypic_entropy", "phenotypic_entropy", "mutual_information", "phenotypic_flux", "compare_groups"]` + +### 7.7 h5ad-serializable return shapes (forward-compat with the deferred `@tl_result` uns-cache) + +Every `tl` return frame is constrained **now** to survive an h5ad round-trip so the deferred `@tl_result` cache is a one-line wrapper later: **flat columns only, no object-dtype "samples" columns**, and provenance in a serializable `_provenance` column (JSON string) **plus** `df.attrs["params"]` (attrs are convenience-only; the column is the durable copy). Per-draw values live in a **separate long frame** carrying an explicit `sample_id` level, never as numpy vectors embedded in object columns (which `AnnData.write` drops). Cache key = hash of `(covariate, groupby, n_samples, temperature, clones, normalized, normalize_mode, distance_metric, use_logits, random_state)`. + +### 7.8 Draw-once efficiency invariant + +For `n_samples>0`, the engine draws the `p_ct` table **once per sample** and **reuses that draw across all covariates, groups, and clones**; groups are formed by cell/clone masking, not re-drawing. `covariate=None`, the flux sankey's pairwise series, per-patient analyses, and `diag.permutation_null` all consume one shared draw stack. A test/counter asserts the number of Dirichlet draws equals `n_samples`, independent of `#groups` and `#covariates`. + +### 7.9 Precomputed-joint fast path constraints + +A bare precomputed joint carries no `p_ct`/`local_scale`/logits/cells, so it is valid **only for `n_samples=0` and `groupby=None`** (`clones=` merely re-filters rows). `n_samples>0` or `groupby` on a bare joint raises a clear `ValueError`; a jd that already contains a `sample_id` axis is accepted, with `n_samples` treated as validation, not resampling. The adata-path == precomputed-jd equivalence test is scoped to `n_samples=0`. (`phenotypic_flux` takes adata only — it needs two joints — so it is unaffected.) Metrics propagate the input jd's `_provenance` into their output. + +--- + +## 8. `tcri.pl` — plotting (`plotting/`) + +Twins mirror `tl` by filename and function name. Each renderer accepts its `tl` twin's metric arguments (computing the joint internally when needed) plus rendering args. Cross-group comparison is driven by **`groupby`** (dots = aggregation unit) and **`splitby`** (box hue = comparison cohort) — **both axes are retained by design** because most figures carry two categorical axes simultaneously (e.g. dots = patient, boxes = response, x = phenotype). Statistics come from `_stats` / `compare_groups`. + +### 8.1 `plotting/_entropy.py` + +```python +clonotypic_entropy( + adata, *, + covariate=None, groupby=None, splitby=None, + n_samples=0, temperature=1.0, clones=None, normalized=True, n_clones_ref=None, + palette=None, hue_order=None, ax=None, figsize=(6, 3), + rotation=90, legend_fontsize=6, bbox_to_anchor=(1.15, 1.0), + random_state=None, save=None, return_df=False, +) +``` +**(renamed from `clonotypic_entropy_by_phenotype`)** Box-and-dot plot of clonotypic entropy per phenotype across covariate values, per-group dots, significance brackets. Cross-group plots default to a common `n_clones_ref` for comparability (§7.2). + +```python +phenotypic_entropy( + adata, *, + covariate=None, groupby=None, splitby=None, + n_samples=0, temperature=1.0, clones=None, normalized=True, + palette=None, ax=None, figsize=(8, 4), + rotation=90, legend_fontsize=6, bbox_to_anchor=(1.15, 1.0), + random_state=None, save=None, return_df=False, +) +``` +**[FIXED]** Box/strip plot of phenotypic entropy per covariate/group. + +`__all__ = ["clonotypic_entropy", "phenotypic_entropy"]` + +### 8.2 `plotting/_mutual_information.py` + +```python +mutual_information( + adata, *, + covariate=None, groupby=None, splitby=None, + n_samples=0, temperature=1.0, clones=None, + normalized=True, normalize_mode="min", + palette=None, ax=None, figsize=(8, 4), rotation=90, + legend_fontsize=6, bbox_to_anchor=(1.15, 1.0), + random_state=None, save=None, return_df=False, +) +``` +**[FIXED; subsumes `mi_compare`]** Box/strip plot of clone×phenotype MI per covariate; `groupby` (e.g. `"patient"`) supplies per-group points and drives AUROC/MWU/label-permutation stats. `weighted` retained (default `False`; §0.8); default `normalize_mode="min"` (§7.4). + +`__all__ = ["mutual_information"]` + +### 8.3 `plotting/_flux.py` + +```python +phenotypic_flux( + adata, *, + order, + groupby=None, clones=None, + normalize=True, temperature=1.0, + distance_metric="l1", + phenotype_colors=None, ax=None, figsize=(6, 3), + show_legend=True, title=None, random_state=None, + save=None, return_axes=False, +) +``` +The Sankey. Draws phenotype-distribution flow across the ordered `order` sequence of covariate values, calling `tl.phenotypic_flux` pairwise between consecutive values under one shared draw stack (§7.8). `order` replaces pairwise `cov_from`/`cov_to` because a Sankey spans the full ordered series. + +`__all__ = ["phenotypic_flux"]` + +### 8.4 `plotting/_sankey.py` — private drawing primitives + +**`class SankeyNode`** *(internal)* — `__init__(self, x, y, val, *, dx=0.2, color=None, **kwargs)`; `plot(self, ax)`; `plot_node_connection(self, destination_node, ax, **kwargs)` (curved, color-interpolated ribbon). `_phenotype_mass_per_clone(adata, covariate, clones, normalize) -> dict[str, np.ndarray]` — `{clone → phenotype-mass vector}` at one covariate. `SankeyNode.hex_to_rgb` is **deleted** (0 callers; ribbons use `mcolors.to_rgb`). + +### 8.5 `plotting/_base.py` — private plotting engine + +| Signature | Responsibility | +|---|---| +| `_metric_boxplot(adata, function, *, groupby=None, splitby=None, ylabel="", order=None, palette=None, s=20, ax=None, figsize=(8, 4)) -> (fig, ax)` | **(renamed from `tcri_boxplot`)** Generic per-phenotype metric box/strip engine across `groupby`/`splitby` strata. **Rewritten to compute each stratum by full-space restriction (`clones=`/masks), never `function(adata[mask])`** — so engine-backed metrics never trip the alignment guard (§7.1). | +| `_finish(fig, ax, *, save=None, show=None, return_axes=False)` | scanpy-style show/save/return finalizer. | + +### 8.6 `plotting/_colors.py` + +| Symbol | Responsibility | +|---|---| +| `tcri_colors` (`list[str]`) | Canonical categorical hex palette. | +| `resolve_palette(adata, columns, *, palette=None) -> dict` | **(renamed from `set_color_palette`)** Assign `tcri_colors` to each `obs` column's categories, store in `uns["_colors"]`, return the map. **Fixes the "writes on `adata.copy()`" bug (mutates in place).** | + +`__all__ = ["tcri_colors", "resolve_palette"]` + +--- + +## 9. `tcri.diag` — diagnostics (`diagnostics/`) — NEW + +Read-only checks on the finalized model. PPCs return `DataFrame`s; the two relocated model plots render figures. **`model` is required exactly where the live decoder/param store is needed, optional where `adata` suffices** (stated per function). + +### 9.1 `diagnostics/_ppc.py` + +| Signature | Responsibility / math | +|---|---| +| `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`. | + +`__all__ = ["joint_distribution_ppc", "phenotype_calibration", "reconstruction_ppc", "permutation_null"]` + +### 9.2 `diagnostics/_training.py` + +| Signature | Responsibility | +|---|---| +| `loss(model, *, log_scale=False, ax=None, save=None)` | **(relocated `plot_loss`)** Plot training/validation ELBO and prior-KL from `model.history_`. | +| `archetypes(model, *, ax=None, save=None)` | **(relocated `plot_archetypes`)** Cluster-ordered clone-phenotype heatmap + archetype centroids, ordered by the `labels` from `build_archetypes` (retained, §5.5). | + +`diagnostics/__init__.py __all__ = ["joint_distribution_ppc", "phenotype_calibration", "reconstruction_ppc", "permutation_null", "loss", "archetypes"]` + +> The model PGM (`build_nested_tcri_pgm` / `draw_tcri_pgm_nested`) is moved **out of the package** to `docs/`. + +--- + +## 10. `tcri.ut` — utilities (`utils/`) + +### 10.1 `utils/_session.py` — public + +| Signature | Responsibility | +|---|---| +| `save_tcri_session(model, adata, out_dir, *, save_adata=True, compression="gzip") -> dict` | Persist a trained session: scvi model (weights + registry, no embedded adata), Pyro param store, `setup.json`, the h5ad. | +| `load_tcri_session(run_dir, *, adata_path=None, map_location=None, layer=None) -> (TCRIModel, AnnData)` | Reconstruct `TCRIModel` + `AnnData`: read h5ad, restore setup/category order, re-run `setup_anndata`, load model + Pyro params. **Sets the global Pyro store before any `get_p_ct`/`to_anndata` call** (param-store caveat, §5.2). | + +### 10.2 `utils/_session.py` — private helpers (not re-exported) + +`_to_jsonable(x)`; `_collect_setup_from_adata_or_model(adata, model)`; `_restore_category_order(adata, setup)`; `_resolve_TCRIModel()`; `_disable_scvi_onload_train()`; `_ensure_pyro_posterior_params(model, adata)` (guarantees `q_p_ct_raw`; if missing, warn + re-init to uniform $1/P$); `_pyro_load(path, *, map_location=None)`; `_ensure_dir(path)`. + +> **Removed from utils:** `write_adata_safely`, `_pop_nonserializables` (manager stash retired at `setup_anndata`); `probabilities` (dead: read a never-written `uns` key — **and its module-top import in `_plotting.py` is removed in the same PR**, §11); `build_nested_tcri_pgm`/`draw_tcri_pgm_nested` (→ `docs/`); `stars`/`auc_and_label_permutation`/`bootstrap_auc` (→ `_stats.py`). + +--- + +## 11. Surface deltas (removed / renamed / moved) + +**Disposition rule.** Every function is kept or dropped by ONE test — *is it core?* (the model, the joint engine, the four metrics + their plots, session I/O, PPC diagnostics, shared helpers). Non-core = **dropped (deleted)**; nothing is relocated to `examples/`, and the disposable notebooks are never consulted for disposition. Deletion PRs grep import-sites as well as call-sites so a top-level import (e.g. `utils.probabilities` at `_plotting.py:18`) is removed in the same PR as its symbol. + +- **Deleted (not core — dropped, never moved to examples):** `clonality` (tl + pl), `probability_ternary`, `gene_entropy`, `top_clone_umap`, `clone_size_umap`, `plot_phenotype_probabilities`, `polar_plot`, `compare_phenotypes`, `clonotypic_entropy_base`, `delta_clonotypic_entropy`, `delta_entropy_table`, `mi_compare` (tl + pl), `flux_table`, `bayesian_mutual_information`, `probability_distribution`, `classify_phenotypes`, `get_latent_embedding`, `group_small_clones`, `register_probability_columns`, `remove_meaningless_genes`, `clone_fraction`, `_ent`, `ridge_delta_entropy`, `dkl` (→ `_distance.kl_divergence`), `probabilities` (**and its `_plotting.py` import**), `SankeyNode.hex_to_rgb`, `_ascii_hist` (+ all `graph=`/ASCII paths), `write_adata_safely`, `_pop_nonserializables`, and the retired `uns` keys `tcri_manager`, `tcri_clone_key`, `tcri_phenotype_key`, the `X_tcri_phenotypes` obsm slot. +- **Renamed:** `flux`→`tl.phenotypic_flux`; `get_cell_phenotype_probs`→`TCRIModel.predict`; `register_model`→`TCRIModel.to_anndata`; `clonotypic_entropy_by_phenotype`→`pl.clonotypic_entropy`; `tcri_boxplot`→`_base._metric_boxplot`; `set_color_palette`→`resolve_palette`; params `from_this`/`to_that`→`cov_from`/`cov_to`; engine `posterior=`→`use_logits=` (alias `cell_informed=`); `point_estimate=`→removed (use `n_samples`). `weighted=` is **retained** (default `False`; §0.8), not removed. +- **Made private (folded in):** `register_clonotype_key` / `register_phenotype_key` / `_compute_logits_and_prior` → `preprocessing/_register` internals, folded into `TCRIModel.to_anndata`. +- **Relocated (kept, NOT to examples):** `compare_joint_distribution`→`diag.joint_distribution_ppc`, `plot_loss`→`diag.loss`, `plot_archetypes`→`diag.archetypes`; `build_nested_tcri_pgm` / `draw_tcri_pgm_nested` → **`docs/`** (out of the package, a figure script only). +- **Subsumed by `groupby` + `compare_groups` (removed, with migration recipe):** the plural batch wrappers `clonotypic_entropies` / `phenotypic_entropies`, `pl.phenotypic_entropy_delta`, and every `*_compare` / `*_delta` / `*_table` variant. + +--- + +## 12. Appendix — current → target math/stats deltas (what changed and why) + +| # | Site | Current | Target | Rationale | +|---|---|---|---|---| +| 1 | engine `n_samples=0` | `joint_distribution_posterior` always draws 1 Dirichlet sample | closed-form posterior **mean** $m$, no draw | reproducible, Rao-Blackwell (§0.4) | +| 2 | `mutual_information`/`flux` `n_samples=0` | returns one random draw | deterministic plug-in point estimate | latent bug (§0.5) | +| 3 | `posterior=` semantics | conflates draw-vs-mean **and** logit-folding; MI `posterior=False` raises | axis renamed `use_logits`, means *fold per-cell logits* only; both branches use the posterior, never the generative prior | §0.4 | +| 4 | metric ↔ model agreement | metrics use a Dirichlet **draw** and **never** apply the gate | `use_logits=True` at $T=1$ uses the same gate-aware, mean-prior rule as `predict()` (needs persisted `X_tcri_logits`, `gate_prob`, `classifier_temperature`) | removes silent disagreement (§0.9, §5.1) | +| 5 | **plug-in vs posterior-mean** | drafts equate `n_samples=0` with `mean(n_samples>0)` | documented as **different estimators** (Jensen gap: entropy plug-in $\ge$ mean; flux plug-in $\le$ mean; MI indeterminate); **no equality test** | §0.6 | +| 6 | **`p_gt` summary** | attached to every `n_samples>0` metric | **removed** from single-metric summaries; emitted only by `compare_groups` on a signed $\Delta$ | metrics are $\ge0$ ⇒ $P(>0)\approx1$ (§0.7) | +| 7 | **posterior draw concentration** | three inconsistent variants: `clamp(s·m,1e-3)` (guide) vs `s·p_ct` vs `s·p_ct+1e-8` | draw from the **exact guide** `Dirichlet(clamp(s·m̃, 1e-3))` | HDIs must summarize the learned posterior (§0.5) | +| 8 | **MI `normalize_mode` default** | `"average"` ⇒ denom $\tfrac12(\log_2C+H_p)$, $C$-dependent | **`"min"`** ⇒ $I/H_p$, $C$-independent; document $H_c=\log_2C$ structural | cross-group comparability (§7.4) | +| 9 | clonotypic-entropy denominator | $\log_2$ of raw reindexed row count (inflated by absent clones) | $\log_2$ of **supported** clones; optional fixed `n_clones_ref` | comparability (§7.2) | +| 10 | phenotypic-entropy zero clone | zero-mass clone → uniform → $H=1.0$ | zero-support clone → **`NaN`/excluded** | §7.3 | +| 11 | flux `seed` | seeds NumPy only; torch draws unaffected | seed a **torch (+CUDA) Generator**; `random_state` on engine/metrics | reproducibility (§0.10) | +| 12 | KL base/$\varepsilon$ | natural log in flux; $\log_2$ elsewhere; mixed $\varepsilon$ | one base ($\log_2$/bits), one $\varepsilon=10^{-12}$; add bounded symmetric `jsd` | unit consistency (§3.4) | +| 13 | "HDI" | equal-tailed percentiles labeled HDI | **true** highest-density interval `hdi_low/hdi_high` | correct for skewed bounded posteriors (§0.7) | +| 14 | temperature | applied at different stages in the two engines; double-tempered with `guide_temperature` | **single** power-temper of the base; `uns[K.P_CT]` stores the **raw** mean; $T=1$ reproduces `predict()` | §0.9 | +| 15 | joint-distribution PPC | references undefined global `model` → `NameError` | reads cols from `uns[K.METADATA]` | §9.1 | +| 16 | groupby via slicing | `function(adata[mask])` trips the full-space alignment guard | full-space cell/clone **restriction**; `_metric_boxplot` rewritten | §7.1, §8.5 | +| 17 | `local_scale` fallback | `uns.get("tcri_local_scale", 1.0)` silently corrupts draw variance if unwritten | `to_anndata` always writes `K.LOCAL_SCALE`; engine **raises** (no `1.0` default) when missing at `n_samples>0` | draw-variance integrity (§5.1) | +| 18 | partial posterior | intervals silently read as full predictive uncertainty | documented: `n_samples>0` captures **`p_ct` uncertainty only** (classifier logits fixed at their posterior-mean encoding) | §0.6 | + +--- + +*Source of truth cross-checked against `tcri/model/_model.py`, `tcri/model/_module.py`, `tcri/preprocessing/_preprocessing.py`, `tcri/metrics/_metrics.py`, `tcri/plotting/_plotting.py`, `tcri/plotting/_sankey.py`, `tcri/utils/_utils.py`, and the grafiti reference at `/Users/ceglian/Codebase/GitHub/grafiti/grafiti`. Intended document home: `/Users/ceglian/Codebase/GitHub/tcri/docs/contract/tcri_api_and_responsibilities.md`.* \ No newline at end of file diff --git a/docs/contract/tcri_api_contract.html b/docs/contract/tcri_api_contract.html new file mode 100644 index 0000000..42958ea --- /dev/null +++ b/docs/contract/tcri_api_contract.html @@ -0,0 +1,20 @@ +TCRI target API contract

TCRI — Target API Contract

Ideal post-refactor surface · ml/pp/tl/pl/ut · work toward this, do not drift.

Conventions. ml: scvi/pyro model; the single source of phenotype probabilities (get_cell_phenotype_probs). pp: writes the canonical tcri_* state and owns the one joint_distribution engine (posterior= flag). tl: keyword-only; pure (reads state, returns arrays/frames, writes nothing); n_samples=0 point estimate else posterior draws. pl: ax/save/palette; calls the matching tl/pp (never owns model math). ut: session save/load only. Every module declares __all__; shared logic lives in _console/_stats/_base, never copied.
Model (ml)Preprocess (pp)Metrics (tl)Plotting (pl)Session / IO (ut)

Summary (target inventory)

nsgroupfunctionstatusreturnswrites
Model (ml)modelTCRIModel.setup_anndatakeepadata (registers scvi AnnDataManager)uns[tcri_manager]; registry{clonotype/phenotype/covariate/batch_col}
Model (ml)modelTCRIModel.trainkeepNone (fits in place; populates history_)pyro param store (q_p_c_raw, q_p_ct_raw); model weights
Model (ml)modelTCRIModel.get_latent_representationkeepndarray[N, n_latent] — posterior-mean z—
Model (ml)modelTCRIModel.get_cell_phenotype_probskeep — now the ONLY prob pathndarray[N, P] — per-cell phenotype probabilities—
Model (ml)modelTCRIModel.get_p_ctkeepndarray[CT, P] — learned clone×covariate phenotype prior—
Model (ml)modelTCRIModel.boost_phenotype_priorkeep — advancedNone (mutates clone_phen_prior / mixture in place)module.clone_phen_prior, module.mixture_concentration
Preprocess (pp)registerregister_modelkeepadata (all tcri_* state written)uns[tcri_metadata, tcri_p_ct, tcri_ct_to_cov, tcri_ct_to_c, tcri_local_scale, tcri_{covariate,clonotype,phenotype}_categories, tcri_{ct,cov}_array_for_cells]; obsm[X_tcri, X_tcri_logits, X_tcri_logposterior, X_tcri_probabilities]; obs[tcri_phenotype]
Preprocess (pp)enginejoint_distributionmerge — unified engineDataFrame[clone × phenotype] (rows sum to 1 unless weighted); n_samples>0 → stacked draws— (pure)
Preprocess (pp)bookkeepinggroup_singletonskeepNone (writes obs)obs[target_col], obs[trb_candidate]
Preprocess (pp)bookkeepingclone_sizekeepNone | dict{clone: size}obs[clone_size]
Preprocess (pp)bookkeepingfilter_genesrename + fixadata (subset copy)— (returns new view)
Metrics (tl)per-covariateclonotypic_entropykeep (+ single-phenotype mode)Series[phenotype] (point) | ndarray[n_samples, P] (draws) | float (phenotype=)—
Metrics (tl)per-covariatephenotypic_entropykeepSeries[clone] (point) | ndarray[n_samples, n_clones] (draws)—
Metrics (tl)per-covariatemutual_informationkeepfloat (point) | ndarray[n_samples] (draws)—
Metrics (tl)per-covariateclonalitykeepdict{phenotype: clonality∈[0,1]}—
Metrics (tl)between-covariatefluxkeepSeries[clone] (point) | ndarray[n_samples, n_clones] (draws)—
Metrics (tl)between-covariatedelta_clonotypic_entropykeepndarray[n_samples] — H_post − H_pre—
Metrics (tl)tablesmi_comparekeepdict{samples, summary, pairs, params}—
Metrics (tl)tablesdelta_entropy_tablekeepDataFrame[phenotype × splitby] — delta_samples + summary stats—
Metrics (tl)tablesflux_tablekeepDataFrame[clone × splitby] — flux_samples, flux_mean/sd, clone_size—
Plotting (pl)metricmutual_informationkeepAxes | Figure | None—
Plotting (pl)metricclonotypic_entropyrenameAxes | Figure | None—
Plotting (pl)metricphenotypic_entropyfixAxes | Figure | None—
Plotting (pl)metricclonalitykeepAxes | Figure | None—
Plotting (pl)metricfluxfixAxes | Figure | None—
Plotting (pl)comparemi_comparekeepAxes | Figure | None—
Plotting (pl)comparebayesian_mutual_informationkeepAxes | Figure | None—
Plotting (pl)compareridge_delta_entropyfixAxes | Figure | None—
Plotting (pl)distributionphenotypic_fluxkeepAxes | Figure | None—
Plotting (pl)distributionpolar_plotfixAxes | Figure | None—
Plotting (pl)umapclone_size_umapkeepAxes | Figure | None—
Plotting (pl)umaptop_clone_umapkeepAxes | Figure | None—
Plotting (pl)umapphenotype_probabilities_umaprenameAxes | Figure | None—
Plotting (pl)diagnosticmodel_lossmoveAxes | Figure | None—
Plotting (pl)diagnosticarchetypesmoveAxes | Figure | None—
Plotting (pl)diagnosticmodel_pgmmove+mergeAxes | Figure | None—
Session / IO (ut)sessionsave_tcri_sessionkeepdict{paths}run_dir/{model, pyro_params.pt, setup.json, adata.h5ad, meta.json}
Session / IO (ut)sessionload_tcri_sessionkeep(model, adata)pyro param store (restored)
Session / IO (ut)sessionwrite_adata_safelykeepNonepath (h5ad without tcri_manager)

adata-state schema (the data contract — root of the DAG)

keyproduced bymeaningconsumed by
uns[tcri_metadata]register_model{covariate,clone,phenotype,batch}_col≈ every tl/pp/pl function
uns[tcri_p_ct]register_modellearned clone×cov phenotype prior [CT,P]pp.joint_distribution
uns[tcri_local_scale]register_modelDirichlet concentration scalepp.joint_distribution (posterior draw)
uns[tcri_{covariate,clonotype,phenotype}_categories]register_modelcategory ordersjoint_distribution, all metrics
uns[tcri_ct_to_cov], uns[tcri_ct_to_c]register_modelct→cov / ct→clone mapspp.joint_distribution (posterior=False)
uns[tcri_{ct,cov}_array_for_cells]register_modelper-cell ct/cov indicespp.joint_distribution (cell selection)
obsm[X_tcri]register_modellatent posterior mean zpl.*_umap (if X_umap absent)
obsm[X_tcri_logits]register_modelclassifier logits per cellpp.joint_distribution (combine_with_logits)
obsm[X_tcri_probabilities]register_modelsoftmax phenotype probs [N,P]pl.phenotype_probabilities_umap
obs[tcri_phenotype]register_modelhard phenotype labeltl.clonality
obs[clone_size]pp.clone_sizecells per clonepl.clone_size_umap
obs[trb_unique]pp.group_singletonscollapsed clone idsetup_anndata(clonotype_key)
pyro param storeml.TCRIModel.trainq_p_c_raw / q_p_ct_rawget_p_ct, save/load_session

Detail cards

Model (ml)   ml.TCRIModel.setup_anndata · model · keep

ml.TCRIModel.setup_anndata(adata, *, layer=None, clonotype_key='unique_clone_id', phenotype_key='phenotype_col', covariate_key='timepoint', batch_key='patient')
returns: adata (registers scvi AnnDataManager)
writes: uns[tcri_manager]; registry{clonotype/phenotype/covariate/batch_col}
reads: obs[clonotype/phenotype/covariate/batch_key]
calls: —
upstream: —
invariants: all four obs columns must exist → ValueError otherwise
edge cases: layer=None → counts from X
provenance: keep as-is
note: scvi classmethod entry point

Model (ml)   ml.TCRIModel.train · model · keep

ml.TCRIModel.train(max_epochs=1000, batch_size=1000, lr=1e-3, reconstruction_loss_scale=1e-3, n_steps_kl_warmup=2000, **kw)
returns: None (fits in place; populates history_)
writes: pyro param store (q_p_c_raw, q_p_ct_raw); model weights
reads: adata_manager
calls: UnifiedTrainingPlan, TrainRunner, DataSplitter
upstream: setup_anndata
invariants: early-stops on elbo_validation
edge cases: —
plot mirror: pl.model_loss
provenance: keep as-is

Model (ml)   ml.TCRIModel.get_latent_representation · model · keep

ml.TCRIModel.get_latent_representation(adata=None, indices=None, batch_size=None)
returns: ndarray[N, n_latent] — posterior-mean z
writes: —
reads: encoder
calls: module.get_latent
upstream: train
invariants: deterministic (mean, no sampling)
edge cases: —
provenance: keep as-is

Model (ml)   ml.TCRIModel.get_cell_phenotype_probs · model · keep — now the ONLY prob path

ml.TCRIModel.get_cell_phenotype_probs(adata=None, batch_size=256, eps=1e-8)
returns: ndarray[N, P] — per-cell phenotype probabilities
writes: —
reads: encoder, classifier, get_p_ct(), module.ct_array
calls: module.get_p_ct
upstream: train
invariants: softmax(gate·logits + (1-gate)·log prior) OR additive when no gate
edge cases: use_gate toggles the two combination rules
provenance: canonical; absorbs pp._compute_logits_and_prior + pp.classify_phenotypes
note: register_model now calls THIS instead of reimplementing the softmax

Model (ml)   ml.TCRIModel.get_p_ct · model · keep

ml.TCRIModel.get_p_ct()
returns: ndarray[CT, P] — learned clone×covariate phenotype prior
writes: —
reads: pyro param store q_p_ct_raw
calls: module.get_p_ct
upstream: train
invariants: rows sum to 1; NaN→uniform guard
edge cases: —
provenance: keep as-is

Model (ml)   ml.TCRIModel.boost_phenotype_prior · model · keep — advanced

ml.TCRIModel.boost_phenotype_prior(phenotype_name, boost_factor=5.0, *, affect_mixture=True)
returns: None (mutates clone_phen_prior / mixture in place)
writes: module.clone_phen_prior, module.mixture_concentration
reads: c2p_mat, centers
calls: —
upstream: (before) train
invariants: rows renormalized to 1 after boost
edge cases: unknown phenotype → ValueError
provenance: keep; drop its inline _ok (use _console)
note: niche manual-prior knob; kept but flagged advanced

Preprocess (pp)   pp.register_model · register · keep

pp.register_model(adata, model, *, latent_slot='X_tcri', batch_size=256, store_logits=True, store_logposterior=True, compute_umap=False, clonotype_key='trb_unique', ...)
returns: adata (all tcri_* state written)
writes: uns[tcri_metadata, tcri_p_ct, tcri_ct_to_cov, tcri_ct_to_c, tcri_local_scale, tcri_{covariate,clonotype,phenotype}_categories, tcri_{ct,cov}_array_for_cells]; obsm[X_tcri, X_tcri_logits, X_tcri_logposterior, X_tcri_probabilities]; obs[tcri_phenotype]
reads: model.module.*, adata_manager.registry
calls: model.get_latent_representation, model.get_cell_phenotype_probs, model.get_p_ct
upstream: model.train
invariants: per-cell arrays length == n_obs (guarded downstream)
edge cases: compute_umap optional; prob slot only written if absent
provenance: keep; folds in register_phenotype_key + register_clonotype_key (metadata only)
note: THE bridge: model outputs → canonical adata state

Preprocess (pp)   pp.joint_distribution · engine · merge — unified engine

pp.joint_distribution(adata, covariate, *, posterior=True, n_samples=0, temperature=1.0, clones=None, weighted=False, combine_with_logits=True, seed=None, silent=True)
returns: DataFrame[clone × phenotype] (rows sum to 1 unless weighted); n_samples>0 → stacked draws
writes: — (pure)
reads: uns[tcri_p_ct, tcri_local_scale, tcri_*_categories, tcri_metadata, tcri_{ct,cov}_array_for_cells, tcri_ct_to_*]; obsm[X_tcri_logits]
calls: —
upstream: register_model
invariants: posterior=True draws Dirichlet(local_scale·p_ct)+logit combine; posterior=False = point-estimate prior; FAILS LOUDLY on filtered-view length mismatch
edge cases: clones filter + reindex; weighted → mass-weighted, no renorm
provenance: MERGES joint_distribution_posterior (posterior=True) + joint_distribution (posterior=False)
note: the single computational core under every metric

Preprocess (pp)   pp.group_singletons · bookkeeping · keep

pp.group_singletons(adata, *, clonotype_key='trb', groupby='patient', target_col='trb_unique', min_clone_size=10)
returns: None (writes obs)
writes: obs[target_col], obs[trb_candidate]
reads: obs[clonotype_key], obs[groupby]
calls: —
upstream: —
invariants: clones < min_clone_size → 'Singleton_<group>'
edge cases: —
provenance: keep; subsumes group_small_clones (hardcoded dup)
note: canonical small-clone collapse

Preprocess (pp)   pp.clone_size · bookkeeping · keep

pp.clone_size(adata, *, key_added='clone_size', return_counts=False)
returns: None | dict{clone: size}
writes: obs[clone_size]
reads: uns[tcri_metadata][clone_col] (was tcri_clone_key)
calls: —
upstream: register_model
invariants: size == cells per clone
edge cases: —
provenance: keep; retarget onto tcri_metadata (drop tcri_clone_key)

Preprocess (pp)   pp.filter_genes · bookkeeping · rename + fix

pp.filter_genes(adata, *, mt=True, rp=True, tcr=True, hsp=True, mtrn=True, ribo=True)
returns: adata (subset copy)
writes: — (returns new view)
reads: var_names
calls: —
upstream: —
invariants: each flag composes (AND), never resets the running mask
edge cases: HLA-* kept despite '-'/'.'
provenance: rename of remove_meaningless_genes + FIX flag-reset bug
note: TCR-gene removal pre-embedding is methodologically in-scope; generic flags fixed

Metrics (tl)   tl.clonotypic_entropy · per-covariate · keep (+ single-phenotype mode)

tl.clonotypic_entropy(adata, covariate, *, n_samples=0, temperature=1.0, clones=None, combine_with_logits=True, normalised=True, phenotype=None)
returns: Series[phenotype] (point) | ndarray[n_samples, P] (draws) | float (phenotype=)
writes: —
reads: uns[tcri_phenotype_categories]; (joint_distribution)
calls: pp.joint_distribution
upstream: register_model
invariants: H normalised by log2(n_clones); n_samples<1 → ValueError
edge cases: empty joint → NaN row
plot mirror: pl.clonotypic_entropy
provenance: keep; absorbs clonotypic_entropy_base via phenotype= arg
note: H[P(c|phi,m)] — phenotype spread across clones

Metrics (tl)   tl.phenotypic_entropy · per-covariate · keep

tl.phenotypic_entropy(adata, covariate, *, n_samples=0, temperature=1.0, clones=None, combine_with_logits=True, normalised=True)
returns: Series[clone] (point) | ndarray[n_samples, n_clones] (draws)
writes: —
reads: uns[tcri_metadata]; (joint_distribution)
calls: pp.joint_distribution
upstream: register_model
invariants: H normalised by log2(P); n_samples<1 → ValueError
edge cases: no clones at covariate → empty
plot mirror: pl.phenotypic_entropy
provenance: keep as-is
note: H[P(phi|c,m)] — clone phenotypic plasticity

Metrics (tl)   tl.mutual_information · per-covariate · keep

tl.mutual_information(adata, covariate, *, n_samples=0, temperature=1.0, clones=None, normalised=True, normalise_mode='average', posterior=True, combine_with_logits=True, verbose=True)
returns: float (point) | ndarray[n_samples] (draws)
writes: —
reads: (joint_distribution)
calls: pp.joint_distribution, _joint_to_mi
upstream: register_model
invariants: I = ΣΣ p·log2(p/(px·py)); normalised by mean/min marginal H
edge cases: posterior=False path implemented (no NotImplementedError)
plot mirror: pl.mutual_information
provenance: keep; wire the prior path that currently raises
note: clone↔phenotype coupling

Metrics (tl)   tl.clonality · per-covariate · keep

tl.clonality(adata)
returns: dict{phenotype: clonality∈[0,1]}
writes: —
reads: obs[tcri_phenotype], uns[tcri_metadata]
calls: —
upstream: register_model
invariants: 1 - H(clone sizes)/log2(K); hard labels (no posterior)
edge cases: single clone → 1; nan→0
plot mirror: pl.clonality
provenance: keep; retarget onto tcri_metadata

Metrics (tl)   tl.flux · between-covariate · keep

tl.flux(adata, *, cov_from, cov_to, clones=None, distance_metric='l1', n_samples=0, temperature=1.0, weighted=False, posterior=True, combine_with_logits=True, seed=42)
returns: Series[clone] (point) | ndarray[n_samples, n_clones] (draws)
writes: —
reads: (joint_distribution ×2)
calls: pp.joint_distribution, _stats.distance
upstream: register_model
invariants: dist over common clones at both covariates
edge cases: no overlap → ValueError; metric ∈ {l1, dkl, callable}
plot mirror: pl.flux
provenance: keep; from_this/to_that → cov_from/cov_to; dkl via _stats registry
note: phenotype-distribution shift per clone

Metrics (tl)   tl.delta_clonotypic_entropy · between-covariate · keep

tl.delta_clonotypic_entropy(adata, phenotype, *, cov_pre, cov_post, n_samples=1000, temperature=1.0, clones=None, weighted=False, normalised=True, posterior=True, combine_with_logits=True, seed=None)
returns: ndarray[n_samples] — H_post − H_pre
writes: —
reads: (clonotypic_entropy)
calls: tl.clonotypic_entropy
upstream: register_model
invariants: positive ⇒ entropy rose pre→post
edge cases: —
plot mirror: pl.ridge_delta_entropy (via delta_entropy_table)
provenance: keep; calls clonotypic_entropy(phenotype=)

Metrics (tl)   tl.mi_compare · tables · keep

tl.mi_compare(adata, groupby, *, groups=None, treatment=None, n_samples=50, patient_col=None, clone_col=None, covariate_col=None, verbose=True)
returns: dict{samples, summary, pairs, params}
writes: —
reads: uns[tcri_metadata]; obs[groupby, patient_col]
calls: tl.mutual_information (per patient×covariate)
upstream: register_model
invariants: patient-level samples → group summary; pairs from groups
edge cases: missing group/cov skipped
plot mirror: pl.mi_compare
provenance: keep; uses shared _group_table loop
note: patient-level MI comparison

Metrics (tl)   tl.delta_entropy_table · tables · keep

tl.delta_entropy_table(adata, *, cov_pre, cov_post, splitby='response', n_samples=1000, temperature=1.0, weighted=False, normalised=True, posterior=True, combine_with_logits=True, seed=42)
returns: DataFrame[phenotype × splitby] — delta_samples + summary stats
writes: —
reads: uns[tcri_metadata]; obs[splitby]
calls: tl.delta_clonotypic_entropy
upstream: register_model
invariants: keeps full delta vector per row
edge cases: —
plot mirror: pl.ridge_delta_entropy
provenance: keep; uses shared _group_table loop

Metrics (tl)   tl.flux_table · tables · keep

tl.flux_table(adata, *, cov_pre, cov_post, splitby='response', n_samples=0, temperature=1.0, weighted=False, posterior=True, combine_with_logits=True, distance_metric='l1', seed=42)
returns: DataFrame[clone × splitby] — flux_samples, flux_mean/sd, clone_size
writes: —
reads: uns[tcri_metadata]; obs[splitby]
calls: tl.flux
upstream: register_model
invariants: per-group clone scoping
edge cases: —
provenance: keep; uses shared _group_table loop

Plotting (pl)   pl.mutual_information · metric · keep

pl.mutual_information(adata, *, ..., palette=None, figsize=..., ax=None, save=None)
returns: Axes | Figure | None
writes: —
reads: (via tl.mutual_information)
calls: tl.mutual_information
upstream: matching tl tool
invariants: reads computed state / calls tl; never owns model math
edge cases: —
provenance: keep
note: box/strip across covariate ±splitby

Plotting (pl)   pl.clonotypic_entropy · metric · rename

pl.clonotypic_entropy(adata, *, ..., palette=None, figsize=..., ax=None, save=None)
returns: Axes | Figure | None
writes: —
reads: (via tl.clonotypic_entropy)
calls: tl.clonotypic_entropy
upstream: matching tl tool
invariants: reads computed state / calls tl; never owns model math
edge cases: —
provenance: rename of clonotypic_entropy_by_phenotype
note: per-phenotype box/dot ±covariate

Plotting (pl)   pl.phenotypic_entropy · metric · fix

pl.phenotypic_entropy(adata, *, ..., palette=None, figsize=..., ax=None, save=None)
returns: Axes | Figure | None
writes: —
reads: (via tl.phenotypic_entropy)
calls: tl.phenotypic_entropy
upstream: matching tl tool
invariants: reads computed state / calls tl; never owns model math
edge cases: —
provenance: FIX broken tl call signature
note: box/strip per covariate

Plotting (pl)   pl.clonality · metric · keep

pl.clonality(adata, *, ..., palette=None, figsize=..., ax=None, save=None)
returns: Axes | Figure | None
writes: —
reads: (via tl.clonality)
calls: tl.clonality, pl._metric_boxplot
upstream: matching tl tool
invariants: reads computed state / calls tl; never owns model math
edge cases: —
provenance: keep
note: stripplot per phenotype ±group

Plotting (pl)   pl.flux · metric · fix

pl.flux(adata, *, ..., palette=None, figsize=..., ax=None, save=None)
returns: Axes | Figure | None
writes: —
reads: (via tl.flux)
calls: tl.flux
upstream: matching tl tool
invariants: reads computed state / calls tl; never owns model math
edge cases: —
provenance: FIX broken key= passed to tl.flux
note: box of flux distance by group

Plotting (pl)   pl.mi_compare · compare · keep

pl.mi_compare(adata, *, ..., palette=None, figsize=..., ax=None, save=None)
returns: Axes | Figure | None
writes: —
reads: uns/(via tl.mi_compare)
calls: tl.mi_compare, _stats.auc_and_label_permutation, _stats.bootstrap_auc
upstream: matching tl tool
invariants: reads computed state / calls tl; never owns model math
edge cases: —
provenance: keep
note: patient MI box + AUROC/permutation stats

Plotting (pl)   pl.bayesian_mutual_information · compare · keep

pl.bayesian_mutual_information(adata, *, ..., palette=None, figsize=..., ax=None, save=None)
returns: Axes | Figure | None
writes: —
reads: (via tl.mutual_information ×2)
calls: tl.mutual_information
upstream: matching tl tool
invariants: reads computed state / calls tl; never owns model math
edge cases: —
provenance: keep
note: ΔMI KDE / posterior / bar across two covariates

Plotting (pl)   pl.ridge_delta_entropy · compare · fix

pl.ridge_delta_entropy(adata, *, ..., palette=None, figsize=..., ax=None, save=None)
returns: Axes | Figure | None
writes: —
reads: DataFrame from tl.delta_entropy_table
calls: —
upstream: tl.delta_entropy_table
invariants: reads computed state / calls tl; never owns model math
edge cases: —
provenance: FIX undefined cm/st imports
note: ridge plot of Δ-entropy posteriors

Plotting (pl)   pl.phenotypic_flux · distribution · keep

pl.phenotypic_flux(adata, *, ..., palette=None, figsize=..., ax=None, save=None)
returns: Axes | Figure | None
writes: —
reads: (via pp.joint_distribution)
calls: pp.joint_distribution, _build_sankey
upstream: matching tl tool
invariants: reads computed state / calls tl; never owns model math
edge cases: —
provenance: keep; absorbs plot_pheno_sankey as private _build_sankey
note: phenotype-flow sankey across covariates

Plotting (pl)   pl.polar_plot · distribution · fix

pl.polar_plot(adata, *, ..., palette=None, figsize=..., ax=None, save=None)
returns: Axes | Figure | None
writes: —
reads: (via pp.joint_distribution / tl.clonotypic_entropy)
calls: pp.joint_distribution, tl.clonotypic_entropy
upstream: matching tl tool
invariants: reads computed state / calls tl; never owns model math
edge cases: —
provenance: FIX undefined clonotypic_entropy ref + string phenotypes
note: radar of phenotype distribution or entropy

Plotting (pl)   pl.clone_size_umap · umap · keep

pl.clone_size_umap(adata, *, ..., palette=None, figsize=..., ax=None, save=None)
returns: Axes | Figure | None
writes: —
reads: obs[clone_size], obsm[X_umap]
calls: pp.clone_size
upstream: matching tl tool
invariants: reads computed state / calls tl; never owns model math
edge cases: —
provenance: keep
note: UMAP colored by log clone size

Plotting (pl)   pl.top_clone_umap · umap · keep

pl.top_clone_umap(adata, *, ..., palette=None, figsize=..., ax=None, save=None)
returns: Axes | Figure | None
writes: —
reads: obs[clone_col], obsm[X_umap]
calls: —
upstream: matching tl tool
invariants: reads computed state / calls tl; never owns model math
edge cases: —
provenance: keep
note: UMAP highlighting top-N clones

Plotting (pl)   pl.phenotype_probabilities_umap · umap · rename

pl.phenotype_probabilities_umap(adata, *, ..., palette=None, figsize=..., ax=None, save=None)
returns: Axes | Figure | None
writes: —
reads: obsm[X_tcri_probabilities]
calls: —
upstream: matching tl tool
invariants: reads computed state / calls tl; never owns model math
edge cases: —
provenance: rename of plot_phenotype_probabilities
note: per-phenotype probability UMAP grid

Plotting (pl)   pl.model_loss · diagnostic · move

pl.model_loss(adata, *, ..., palette=None, figsize=..., ax=None, save=None)
returns: Axes | Figure | None
writes: —
reads: model.history_
calls: —
upstream: model.train
invariants: reads computed state / calls tl; never owns model math
edge cases: —
provenance: moved from TCRIModel.plot_loss
note: ELBO + dKL training curves

Plotting (pl)   pl.archetypes · diagnostic · move

pl.archetypes(adata, *, ..., palette=None, figsize=..., ax=None, save=None)
returns: Axes | Figure | None
writes: —
reads: model.c2p_mat, model.centers
calls: —
upstream: model.train
invariants: reads computed state / calls tl; never owns model math
edge cases: —
provenance: moved from TCRIModel.plot_archetypes
note: archetype / clone-prior heatmaps

Plotting (pl)   pl.model_pgm · diagnostic · move+merge

pl.model_pgm(adata, *, ..., palette=None, figsize=..., ax=None, save=None)
returns: Axes | Figure | None
writes: —
reads: — (static daft diagram)
calls: —
upstream: —
invariants: reads computed state / calls tl; never owns model math
edge cases: —
provenance: moved from ut.build_nested_tcri_pgm + ut.draw_tcri_pgm_nested (merged)
note: TCRI plate-diagram (PGM)

Session / IO (ut)   ut.save_tcri_session · session · keep

ut.save_tcri_session(model, adata, out_dir, *, save_adata=True, compression='gzip')
returns: dict{paths}
writes: run_dir/{model, pyro_params.pt, setup.json, adata.h5ad, meta.json}
reads: model.save, pyro store, adata
calls: write_adata_safely, _collect_setup_from_adata_or_model
upstream: train
invariants: adata written without tcri_manager (non-picklable)
edge cases: —
provenance: keep as-is

Session / IO (ut)   ut.load_tcri_session · session · keep

ut.load_tcri_session(run_dir, *, adata_path=None, map_location=None, layer=None)
returns: (model, adata)
writes: pyro param store (restored)
reads: run_dir artifacts
calls: TCRIModel.setup_anndata/.load, _pyro_load, _restore_category_order, _ensure_pyro_posterior_params, _disable_scvi_onload_train
upstream: save_tcri_session
invariants: category order restored; posterior params ensured
edge cases: missing pyro store → warn + uniform-prior fallback
provenance: keep as-is

Session / IO (ut)   ut.write_adata_safely · session · keep

ut.write_adata_safely(adata, path, *, compression='gzip')
returns: None
writes: path (h5ad without tcri_manager)
reads: adata
calls: _pop_nonserializables
upstream: —
invariants: strips non-serializable manager before write
edge cases: —
provenance: keep as-is

Shared primitives

Console (tcri/_console.py) — replaces 3 duplicated copies
  • _ok / _info / _warn / _fin — every verbose tl/pp/pl/ml function
  • _ascii_hist(samples) — mutual_information, flux, delta_clonotypic_entropy (graph=)
  • ANSI constants (RESET/BOLD/GRN/...) — all of the above
Stats (tcri/_stats.py)
  • _joint_to_mi(pxy, normalised, mode) — tl.mutual_information
  • distance(metric) -> f(p,q) [l1 | dkl | callable] — tl.flux, tl.flux_table
  • auc_and_label_permutation(scores, labels) — pl.mi_compare
  • bootstrap_auc(scores, labels) — pl.mi_compare
  • _norm_entropy(p, base, n) — clonotypic_entropy, phenotypic_entropy
Table builder (tcri/metrics/_tables.py)
  • _group_table(adata, splitby, per_group_fn) — mi_compare, delta_entropy_table, flux_table
Plot helpers (tcri/plotting/_base.py)
  • _metric_boxplot(adata, fn, ...) — clonality (+ any group×split metric box)
  • _resolve_palette / tcri_colors — every pl function
  • _build_sankey / SankeyNode / _phenotype_mass_per_clone — phenotypic_flux
Session internals (tcri/utils/_io.py)
  • _ensure_dir, _to_jsonable, _pop_nonserializables — save_tcri_session, write_adata_safely
  • _collect_setup_from_adata_or_model, _restore_category_order — save/load_tcri_session
  • _pyro_load, _resolve_TCRIModel, _disable_scvi_onload_train, _ensure_pyro_posterior_params — load_tcri_session

Dropped / merged from the current code

currentdispositionwhy
pp.classify_phenotypesdeletedredundant cosine phenotype assignment; 0 callers
pp.get_latent_embeddingdeletedtrivial gaussian sampler; 0 callers
pp.register_probability_columnsdeletedfeeds only the dead probabilities()
pp.gene_entropydeletedgeneric gene QC, out of scope; 0 callers
pp.group_small_clonessubsumed → group_singletonshardcoded inferior dup
pp.register_phenotype_key / register_clonotype_keysubsumed → register_modelkill the tcri_*_key shadow convention
pp._compute_logits_and_priorsubsumed → ml.get_cell_phenotype_probsone prob path
pp.joint_distribution_posteriorsubsumed → joint_distribution(posterior=True)unified engine
tl.clonotypic_entropy_basesubsumed → clonotypic_entropy(phenotype=)one entropy fn
tl.clone_fractiondeletedone-line value-counts; 0 callers
tl.dklsubsumed → _stats.distancedead; flux had its own copy
pl.compare_phenotypesdeletedtrivial crosstab heatmap; 0 callers
pl.compare_joint_distributiondeletedbroken (undefined `model`)
pl.probability_distributiondeletedbroken (self-recursion); covered by polar_plot
pl.set_color_palettedeletedbuggy palette helper; 0 callers
pl.plot_pheno_sankeysubsumed → phenotypic_flux (_build_sankey)one public sankey
pl.tcri_boxplotinternalized → _metric_boxplotshared plot helper
ut.probabilitiesdeletedreads uns[joint_distribution] nothing ever writes
ut.starsdeletedimported but never called
ut.build_nested_tcri_pgm / draw_tcri_pgm_nestedmoved+merged → pl.model_pgmdescribes the model, not IO
ut.auc_and_label_permutation / bootstrap_aucinternalized → _statsused only by pl.mi_compare
SankeyNode.hex_to_rgbdeletedunused method
_ok/_info/_warn/_fin (×3), _ascii_hist (×2)subsumed → _consolededup

See tcri_dependency_map.html for the full call + producer/consumer graph.

\ No newline at end of file diff --git a/docs/contract/tcri_api_contract.md b/docs/contract/tcri_api_contract.md new file mode 100644 index 0000000..896c2d2 --- /dev/null +++ b/docs/contract/tcri_api_contract.md @@ -0,0 +1,722 @@ +# TCRI — Target API Contract (`ml`/`pp`/`tl`/`pl`/`ut`) + +_The ideal post-refactor surface. Work toward this; do not drift. Generated from `build_tcri_contract.py` (single source of truth)._ + +**Conventions.** **ml**: scvi/pyro model; the single source of phenotype probabilities (`get_cell_phenotype_probs`). **pp**: writes the canonical `tcri_*` state and owns the one `joint_distribution` engine (`posterior=` flag). **tl**: keyword-only; pure (reads state, returns arrays/frames, writes nothing); `n_samples=0` point estimate else posterior draws. **pl**: `ax/save/palette`; calls the matching `tl`/`pp` (never owns model math). **ut**: session save/load only. Every module declares `__all__`; shared logic lives in `_console`/`_stats`/`_base`, never copied. + +## Namespaces + +- **Model (ml)** (`ml`) +- **Preprocess (pp)** (`pp`) +- **Metrics (tl)** (`tl`) +- **Plotting (pl)** (`pl`) +- **Session / IO (ut)** (`ut`) + +## Summary (target inventory) + +| ns | group | function | status | returns | writes | +|---|---|---|---|---|---| +| `ml` | model | `TCRIModel.setup_anndata` | keep | adata (registers scvi AnnDataManager) | uns[tcri_manager]; registry{clonotype/phenotype/covariate/batch_col} | +| `ml` | model | `TCRIModel.train` | keep | None (fits in place; populates history_) | pyro param store (q_p_c_raw, q_p_ct_raw); model weights | +| `ml` | model | `TCRIModel.get_latent_representation` | keep | ndarray[N, n_latent] — posterior-mean z | — | +| `ml` | model | `TCRIModel.get_cell_phenotype_probs` | keep — now the ONLY prob path | ndarray[N, P] — per-cell phenotype probabilities | — | +| `ml` | model | `TCRIModel.get_p_ct` | keep | ndarray[CT, P] — learned clone×covariate phenotype prior | — | +| `ml` | model | `TCRIModel.boost_phenotype_prior` | keep — advanced | None (mutates clone_phen_prior / mixture in place) | module.clone_phen_prior, module.mixture_concentration | +| `pp` | register | `register_model` | keep | adata (all tcri_* state written) | uns[tcri_metadata, tcri_p_ct, tcri_ct_to_cov, tcri_ct_to_c, tcri_local_scale, tcri_{covariate,clonotype,phenotype}_categories, tcri_{ct,cov}_array_for_cells]; obsm[X_tcri, X_tcri_logits, X_tcri_logposterior, X_tcri_probabilities]; obs[tcri_phenotype] | +| `pp` | engine | `joint_distribution` | merge — unified engine | DataFrame[clone × phenotype] (rows sum to 1 unless weighted); n_samples>0 → stacked draws | — (pure) | +| `pp` | bookkeeping | `group_singletons` | keep | None (writes obs) | obs[target_col], obs[trb_candidate] | +| `pp` | bookkeeping | `clone_size` | keep | None \| dict{clone: size} | obs[clone_size] | +| `pp` | bookkeeping | `filter_genes` | rename + fix | adata (subset copy) | — (returns new view) | +| `tl` | per-covariate | `clonotypic_entropy` | keep (+ single-phenotype mode) | Series[phenotype] (point) \| ndarray[n_samples, P] (draws) \| float (phenotype=) | — | +| `tl` | per-covariate | `phenotypic_entropy` | keep | Series[clone] (point) \| ndarray[n_samples, n_clones] (draws) | — | +| `tl` | per-covariate | `mutual_information` | keep | float (point) \| ndarray[n_samples] (draws) | — | +| `tl` | per-covariate | `clonality` | keep | dict{phenotype: clonality∈[0,1]} | — | +| `tl` | between-covariate | `flux` | keep | Series[clone] (point) \| ndarray[n_samples, n_clones] (draws) | — | +| `tl` | between-covariate | `delta_clonotypic_entropy` | keep | ndarray[n_samples] — H_post − H_pre | — | +| `tl` | tables | `mi_compare` | keep | dict{samples, summary, pairs, params} | — | +| `tl` | tables | `delta_entropy_table` | keep | DataFrame[phenotype × splitby] — delta_samples + summary stats | — | +| `tl` | tables | `flux_table` | keep | DataFrame[clone × splitby] — flux_samples, flux_mean/sd, clone_size | — | +| `pl` | metric | `mutual_information` | keep | Axes \| Figure \| None | — | +| `pl` | metric | `clonotypic_entropy` | rename | Axes \| Figure \| None | — | +| `pl` | metric | `phenotypic_entropy` | fix | Axes \| Figure \| None | — | +| `pl` | metric | `clonality` | keep | Axes \| Figure \| None | — | +| `pl` | metric | `flux` | fix | Axes \| Figure \| None | — | +| `pl` | compare | `mi_compare` | keep | Axes \| Figure \| None | — | +| `pl` | compare | `bayesian_mutual_information` | keep | Axes \| Figure \| None | — | +| `pl` | compare | `ridge_delta_entropy` | fix | Axes \| Figure \| None | — | +| `pl` | distribution | `phenotypic_flux` | keep | Axes \| Figure \| None | — | +| `pl` | distribution | `polar_plot` | fix | Axes \| Figure \| None | — | +| `pl` | umap | `clone_size_umap` | keep | Axes \| Figure \| None | — | +| `pl` | umap | `top_clone_umap` | keep | Axes \| Figure \| None | — | +| `pl` | umap | `phenotype_probabilities_umap` | rename | Axes \| Figure \| None | — | +| `pl` | diagnostic | `model_loss` | move | Axes \| Figure \| None | — | +| `pl` | diagnostic | `archetypes` | move | Axes \| Figure \| None | — | +| `pl` | diagnostic | `model_pgm` | move+merge | Axes \| Figure \| None | — | +| `ut` | session | `save_tcri_session` | keep | dict{paths} | run_dir/{model, pyro_params.pt, setup.json, adata.h5ad, meta.json} | +| `ut` | session | `load_tcri_session` | keep | (model, adata) | pyro param store (restored) | +| `ut` | session | `write_adata_safely` | keep | None | path (h5ad without tcri_manager) | + +## adata-state schema (the data contract — root of the DAG) + +| key | produced by | meaning | consumed by | +|---|---|---|---| +| `uns[tcri_metadata]` | `register_model` | {covariate,clone,phenotype,batch}_col | ≈ every tl/pp/pl function | +| `uns[tcri_p_ct]` | `register_model` | learned clone×cov phenotype prior [CT,P] | pp.joint_distribution | +| `uns[tcri_local_scale]` | `register_model` | Dirichlet concentration scale | pp.joint_distribution (posterior draw) | +| `uns[tcri_{covariate,clonotype,phenotype}_categories]` | `register_model` | category orders | joint_distribution, all metrics | +| `uns[tcri_ct_to_cov], uns[tcri_ct_to_c]` | `register_model` | ct→cov / ct→clone maps | pp.joint_distribution (posterior=False) | +| `uns[tcri_{ct,cov}_array_for_cells]` | `register_model` | per-cell ct/cov indices | pp.joint_distribution (cell selection) | +| `obsm[X_tcri]` | `register_model` | latent posterior mean z | pl.*_umap (if X_umap absent) | +| `obsm[X_tcri_logits]` | `register_model` | classifier logits per cell | pp.joint_distribution (combine_with_logits) | +| `obsm[X_tcri_probabilities]` | `register_model` | softmax phenotype probs [N,P] | pl.phenotype_probabilities_umap | +| `obs[tcri_phenotype]` | `register_model` | hard phenotype label | tl.clonality | +| `obs[clone_size]` | `pp.clone_size` | cells per clone | pl.clone_size_umap | +| `obs[trb_unique]` | `pp.group_singletons` | collapsed clone id | setup_anndata(clonotype_key) | +| `pyro param store` | `ml.TCRIModel.train` | q_p_c_raw / q_p_ct_raw | get_p_ct, save/load_session | + +## Detail cards + +### `ml.TCRIModel.setup_anndata` — _model_ · keep + +```python +ml.TCRIModel.setup_anndata(adata, *, layer=None, clonotype_key='unique_clone_id', phenotype_key='phenotype_col', covariate_key='timepoint', batch_key='patient') +``` +- **returns:** adata (registers scvi AnnDataManager) +- **writes:** uns[tcri_manager]; registry{clonotype/phenotype/covariate/batch_col} +- **reads:** obs[clonotype/phenotype/covariate/batch_key] +- **calls:** — +- **upstream:** — +- **invariants:** all four obs columns must exist → ValueError otherwise +- **edge cases:** layer=None → counts from X +- **provenance:** keep as-is +- **note:** scvi classmethod entry point + +### `ml.TCRIModel.train` — _model_ · keep + +```python +ml.TCRIModel.train(max_epochs=1000, batch_size=1000, lr=1e-3, reconstruction_loss_scale=1e-3, n_steps_kl_warmup=2000, **kw) +``` +- **returns:** None (fits in place; populates history_) +- **writes:** pyro param store (q_p_c_raw, q_p_ct_raw); model weights +- **reads:** adata_manager +- **calls:** UnifiedTrainingPlan, TrainRunner, DataSplitter +- **upstream:** setup_anndata +- **invariants:** early-stops on elbo_validation +- **edge cases:** — +- **plot mirror:** pl.model_loss +- **provenance:** keep as-is + +### `ml.TCRIModel.get_latent_representation` — _model_ · keep + +```python +ml.TCRIModel.get_latent_representation(adata=None, indices=None, batch_size=None) +``` +- **returns:** ndarray[N, n_latent] — posterior-mean z +- **writes:** — +- **reads:** encoder +- **calls:** module.get_latent +- **upstream:** train +- **invariants:** deterministic (mean, no sampling) +- **edge cases:** — +- **provenance:** keep as-is + +### `ml.TCRIModel.get_cell_phenotype_probs` — _model_ · keep — now the ONLY prob path + +```python +ml.TCRIModel.get_cell_phenotype_probs(adata=None, batch_size=256, eps=1e-8) +``` +- **returns:** ndarray[N, P] — per-cell phenotype probabilities +- **writes:** — +- **reads:** encoder, classifier, get_p_ct(), module.ct_array +- **calls:** module.get_p_ct +- **upstream:** train +- **invariants:** softmax(gate·logits + (1-gate)·log prior) OR additive when no gate +- **edge cases:** use_gate toggles the two combination rules +- **provenance:** canonical; absorbs pp._compute_logits_and_prior + pp.classify_phenotypes +- **note:** register_model now calls THIS instead of reimplementing the softmax + +### `ml.TCRIModel.get_p_ct` — _model_ · keep + +```python +ml.TCRIModel.get_p_ct() +``` +- **returns:** ndarray[CT, P] — learned clone×covariate phenotype prior +- **writes:** — +- **reads:** pyro param store q_p_ct_raw +- **calls:** module.get_p_ct +- **upstream:** train +- **invariants:** rows sum to 1; NaN→uniform guard +- **edge cases:** — +- **provenance:** keep as-is + +### `ml.TCRIModel.boost_phenotype_prior` — _model_ · keep — advanced + +```python +ml.TCRIModel.boost_phenotype_prior(phenotype_name, boost_factor=5.0, *, affect_mixture=True) +``` +- **returns:** None (mutates clone_phen_prior / mixture in place) +- **writes:** module.clone_phen_prior, module.mixture_concentration +- **reads:** c2p_mat, centers +- **calls:** — +- **upstream:** (before) train +- **invariants:** rows renormalized to 1 after boost +- **edge cases:** unknown phenotype → ValueError +- **provenance:** keep; drop its inline _ok (use _console) +- **note:** niche manual-prior knob; kept but flagged advanced + +### `pp.register_model` — _register_ · keep + +```python +pp.register_model(adata, model, *, latent_slot='X_tcri', batch_size=256, store_logits=True, store_logposterior=True, compute_umap=False, clonotype_key='trb_unique', ...) +``` +- **returns:** adata (all tcri_* state written) +- **writes:** uns[tcri_metadata, tcri_p_ct, tcri_ct_to_cov, tcri_ct_to_c, tcri_local_scale, tcri_{covariate,clonotype,phenotype}_categories, tcri_{ct,cov}_array_for_cells]; obsm[X_tcri, X_tcri_logits, X_tcri_logposterior, X_tcri_probabilities]; obs[tcri_phenotype] +- **reads:** model.module.*, adata_manager.registry +- **calls:** model.get_latent_representation, model.get_cell_phenotype_probs, model.get_p_ct +- **upstream:** model.train +- **invariants:** per-cell arrays length == n_obs (guarded downstream) +- **edge cases:** compute_umap optional; prob slot only written if absent +- **provenance:** keep; folds in register_phenotype_key + register_clonotype_key (metadata only) +- **note:** THE bridge: model outputs → canonical adata state + +### `pp.joint_distribution` — _engine_ · merge — unified engine + +```python +pp.joint_distribution(adata, covariate, *, posterior=True, n_samples=0, temperature=1.0, clones=None, weighted=False, combine_with_logits=True, seed=None, silent=True) +``` +- **returns:** DataFrame[clone × phenotype] (rows sum to 1 unless weighted); n_samples>0 → stacked draws +- **writes:** — (pure) +- **reads:** uns[tcri_p_ct, tcri_local_scale, tcri_*_categories, tcri_metadata, tcri_{ct,cov}_array_for_cells, tcri_ct_to_*]; obsm[X_tcri_logits] +- **calls:** — +- **upstream:** register_model +- **invariants:** posterior=True draws Dirichlet(local_scale·p_ct)+logit combine; posterior=False = point-estimate prior; FAILS LOUDLY on filtered-view length mismatch +- **edge cases:** clones filter + reindex; weighted → mass-weighted, no renorm +- **provenance:** MERGES joint_distribution_posterior (posterior=True) + joint_distribution (posterior=False) +- **note:** the single computational core under every metric + +### `pp.group_singletons` — _bookkeeping_ · keep + +```python +pp.group_singletons(adata, *, clonotype_key='trb', groupby='patient', target_col='trb_unique', min_clone_size=10) +``` +- **returns:** None (writes obs) +- **writes:** obs[target_col], obs[trb_candidate] +- **reads:** obs[clonotype_key], obs[groupby] +- **calls:** — +- **upstream:** — +- **invariants:** clones < min_clone_size → 'Singleton_' +- **edge cases:** — +- **provenance:** keep; subsumes group_small_clones (hardcoded dup) +- **note:** canonical small-clone collapse + +### `pp.clone_size` — _bookkeeping_ · keep + +```python +pp.clone_size(adata, *, key_added='clone_size', return_counts=False) +``` +- **returns:** None | dict{clone: size} +- **writes:** obs[clone_size] +- **reads:** uns[tcri_metadata][clone_col] (was tcri_clone_key) +- **calls:** — +- **upstream:** register_model +- **invariants:** size == cells per clone +- **edge cases:** — +- **provenance:** keep; retarget onto tcri_metadata (drop tcri_clone_key) + +### `pp.filter_genes` — _bookkeeping_ · rename + fix + +```python +pp.filter_genes(adata, *, mt=True, rp=True, tcr=True, hsp=True, mtrn=True, ribo=True) +``` +- **returns:** adata (subset copy) +- **writes:** — (returns new view) +- **reads:** var_names +- **calls:** — +- **upstream:** — +- **invariants:** each flag composes (AND), never resets the running mask +- **edge cases:** HLA-* kept despite '-'/'.' +- **provenance:** rename of remove_meaningless_genes + FIX flag-reset bug +- **note:** TCR-gene removal pre-embedding is methodologically in-scope; generic flags fixed + +### `tl.clonotypic_entropy` — _per-covariate_ · keep (+ single-phenotype mode) + +```python +tl.clonotypic_entropy(adata, covariate, *, n_samples=0, temperature=1.0, clones=None, combine_with_logits=True, normalised=True, phenotype=None) +``` +- **returns:** Series[phenotype] (point) | ndarray[n_samples, P] (draws) | float (phenotype=) +- **writes:** — +- **reads:** uns[tcri_phenotype_categories]; (joint_distribution) +- **calls:** pp.joint_distribution +- **upstream:** register_model +- **invariants:** H normalised by log2(n_clones); n_samples<1 → ValueError +- **edge cases:** empty joint → NaN row +- **plot mirror:** pl.clonotypic_entropy +- **provenance:** keep; absorbs clonotypic_entropy_base via phenotype= arg +- **note:** H[P(c|phi,m)] — phenotype spread across clones + +### `tl.phenotypic_entropy` — _per-covariate_ · keep + +```python +tl.phenotypic_entropy(adata, covariate, *, n_samples=0, temperature=1.0, clones=None, combine_with_logits=True, normalised=True) +``` +- **returns:** Series[clone] (point) | ndarray[n_samples, n_clones] (draws) +- **writes:** — +- **reads:** uns[tcri_metadata]; (joint_distribution) +- **calls:** pp.joint_distribution +- **upstream:** register_model +- **invariants:** H normalised by log2(P); n_samples<1 → ValueError +- **edge cases:** no clones at covariate → empty +- **plot mirror:** pl.phenotypic_entropy +- **provenance:** keep as-is +- **note:** H[P(phi|c,m)] — clone phenotypic plasticity + +### `tl.mutual_information` — _per-covariate_ · keep + +```python +tl.mutual_information(adata, covariate, *, n_samples=0, temperature=1.0, clones=None, normalised=True, normalise_mode='average', posterior=True, combine_with_logits=True, verbose=True) +``` +- **returns:** float (point) | ndarray[n_samples] (draws) +- **writes:** — +- **reads:** (joint_distribution) +- **calls:** pp.joint_distribution, _joint_to_mi +- **upstream:** register_model +- **invariants:** I = ΣΣ p·log2(p/(px·py)); normalised by mean/min marginal H +- **edge cases:** posterior=False path implemented (no NotImplementedError) +- **plot mirror:** pl.mutual_information +- **provenance:** keep; wire the prior path that currently raises +- **note:** clone↔phenotype coupling + +### `tl.clonality` — _per-covariate_ · keep + +```python +tl.clonality(adata) +``` +- **returns:** dict{phenotype: clonality∈[0,1]} +- **writes:** — +- **reads:** obs[tcri_phenotype], uns[tcri_metadata] +- **calls:** — +- **upstream:** register_model +- **invariants:** 1 - H(clone sizes)/log2(K); hard labels (no posterior) +- **edge cases:** single clone → 1; nan→0 +- **plot mirror:** pl.clonality +- **provenance:** keep; retarget onto tcri_metadata + +### `tl.flux` — _between-covariate_ · keep + +```python +tl.flux(adata, *, cov_from, cov_to, clones=None, distance_metric='l1', n_samples=0, temperature=1.0, weighted=False, posterior=True, combine_with_logits=True, seed=42) +``` +- **returns:** Series[clone] (point) | ndarray[n_samples, n_clones] (draws) +- **writes:** — +- **reads:** (joint_distribution ×2) +- **calls:** pp.joint_distribution, _stats.distance +- **upstream:** register_model +- **invariants:** dist over common clones at both covariates +- **edge cases:** no overlap → ValueError; metric ∈ {l1, dkl, callable} +- **plot mirror:** pl.flux +- **provenance:** keep; from_this/to_that → cov_from/cov_to; dkl via _stats registry +- **note:** phenotype-distribution shift per clone + +### `tl.delta_clonotypic_entropy` — _between-covariate_ · keep + +```python +tl.delta_clonotypic_entropy(adata, phenotype, *, cov_pre, cov_post, n_samples=1000, temperature=1.0, clones=None, weighted=False, normalised=True, posterior=True, combine_with_logits=True, seed=None) +``` +- **returns:** ndarray[n_samples] — H_post − H_pre +- **writes:** — +- **reads:** (clonotypic_entropy) +- **calls:** tl.clonotypic_entropy +- **upstream:** register_model +- **invariants:** positive ⇒ entropy rose pre→post +- **edge cases:** — +- **plot mirror:** pl.ridge_delta_entropy (via delta_entropy_table) +- **provenance:** keep; calls clonotypic_entropy(phenotype=) + +### `tl.mi_compare` — _tables_ · keep + +```python +tl.mi_compare(adata, groupby, *, groups=None, treatment=None, n_samples=50, patient_col=None, clone_col=None, covariate_col=None, verbose=True) +``` +- **returns:** dict{samples, summary, pairs, params} +- **writes:** — +- **reads:** uns[tcri_metadata]; obs[groupby, patient_col] +- **calls:** tl.mutual_information (per patient×covariate) +- **upstream:** register_model +- **invariants:** patient-level samples → group summary; pairs from groups +- **edge cases:** missing group/cov skipped +- **plot mirror:** pl.mi_compare +- **provenance:** keep; uses shared _group_table loop +- **note:** patient-level MI comparison + +### `tl.delta_entropy_table` — _tables_ · keep + +```python +tl.delta_entropy_table(adata, *, cov_pre, cov_post, splitby='response', n_samples=1000, temperature=1.0, weighted=False, normalised=True, posterior=True, combine_with_logits=True, seed=42) +``` +- **returns:** DataFrame[phenotype × splitby] — delta_samples + summary stats +- **writes:** — +- **reads:** uns[tcri_metadata]; obs[splitby] +- **calls:** tl.delta_clonotypic_entropy +- **upstream:** register_model +- **invariants:** keeps full delta vector per row +- **edge cases:** — +- **plot mirror:** pl.ridge_delta_entropy +- **provenance:** keep; uses shared _group_table loop + +### `tl.flux_table` — _tables_ · keep + +```python +tl.flux_table(adata, *, cov_pre, cov_post, splitby='response', n_samples=0, temperature=1.0, weighted=False, posterior=True, combine_with_logits=True, distance_metric='l1', seed=42) +``` +- **returns:** DataFrame[clone × splitby] — flux_samples, flux_mean/sd, clone_size +- **writes:** — +- **reads:** uns[tcri_metadata]; obs[splitby] +- **calls:** tl.flux +- **upstream:** register_model +- **invariants:** per-group clone scoping +- **edge cases:** — +- **provenance:** keep; uses shared _group_table loop + +### `pl.mutual_information` — _metric_ · keep + +```python +pl.mutual_information(adata, *, ..., palette=None, figsize=..., ax=None, save=None) +``` +- **returns:** Axes | Figure | None +- **writes:** — +- **reads:** (via tl.mutual_information) +- **calls:** tl.mutual_information +- **upstream:** matching tl tool +- **invariants:** reads computed state / calls tl; never owns model math +- **edge cases:** — +- **provenance:** keep +- **note:** box/strip across covariate ±splitby + +### `pl.clonotypic_entropy` — _metric_ · rename + +```python +pl.clonotypic_entropy(adata, *, ..., palette=None, figsize=..., ax=None, save=None) +``` +- **returns:** Axes | Figure | None +- **writes:** — +- **reads:** (via tl.clonotypic_entropy) +- **calls:** tl.clonotypic_entropy +- **upstream:** matching tl tool +- **invariants:** reads computed state / calls tl; never owns model math +- **edge cases:** — +- **provenance:** rename of clonotypic_entropy_by_phenotype +- **note:** per-phenotype box/dot ±covariate + +### `pl.phenotypic_entropy` — _metric_ · fix + +```python +pl.phenotypic_entropy(adata, *, ..., palette=None, figsize=..., ax=None, save=None) +``` +- **returns:** Axes | Figure | None +- **writes:** — +- **reads:** (via tl.phenotypic_entropy) +- **calls:** tl.phenotypic_entropy +- **upstream:** matching tl tool +- **invariants:** reads computed state / calls tl; never owns model math +- **edge cases:** — +- **provenance:** FIX broken tl call signature +- **note:** box/strip per covariate + +### `pl.clonality` — _metric_ · keep + +```python +pl.clonality(adata, *, ..., palette=None, figsize=..., ax=None, save=None) +``` +- **returns:** Axes | Figure | None +- **writes:** — +- **reads:** (via tl.clonality) +- **calls:** tl.clonality, pl._metric_boxplot +- **upstream:** matching tl tool +- **invariants:** reads computed state / calls tl; never owns model math +- **edge cases:** — +- **provenance:** keep +- **note:** stripplot per phenotype ±group + +### `pl.flux` — _metric_ · fix + +```python +pl.flux(adata, *, ..., palette=None, figsize=..., ax=None, save=None) +``` +- **returns:** Axes | Figure | None +- **writes:** — +- **reads:** (via tl.flux) +- **calls:** tl.flux +- **upstream:** matching tl tool +- **invariants:** reads computed state / calls tl; never owns model math +- **edge cases:** — +- **provenance:** FIX broken key= passed to tl.flux +- **note:** box of flux distance by group + +### `pl.mi_compare` — _compare_ · keep + +```python +pl.mi_compare(adata, *, ..., palette=None, figsize=..., ax=None, save=None) +``` +- **returns:** Axes | Figure | None +- **writes:** — +- **reads:** uns/(via tl.mi_compare) +- **calls:** tl.mi_compare, _stats.auc_and_label_permutation, _stats.bootstrap_auc +- **upstream:** matching tl tool +- **invariants:** reads computed state / calls tl; never owns model math +- **edge cases:** — +- **provenance:** keep +- **note:** patient MI box + AUROC/permutation stats + +### `pl.bayesian_mutual_information` — _compare_ · keep + +```python +pl.bayesian_mutual_information(adata, *, ..., palette=None, figsize=..., ax=None, save=None) +``` +- **returns:** Axes | Figure | None +- **writes:** — +- **reads:** (via tl.mutual_information ×2) +- **calls:** tl.mutual_information +- **upstream:** matching tl tool +- **invariants:** reads computed state / calls tl; never owns model math +- **edge cases:** — +- **provenance:** keep +- **note:** ΔMI KDE / posterior / bar across two covariates + +### `pl.ridge_delta_entropy` — _compare_ · fix + +```python +pl.ridge_delta_entropy(adata, *, ..., palette=None, figsize=..., ax=None, save=None) +``` +- **returns:** Axes | Figure | None +- **writes:** — +- **reads:** DataFrame from tl.delta_entropy_table +- **calls:** — +- **upstream:** tl.delta_entropy_table +- **invariants:** reads computed state / calls tl; never owns model math +- **edge cases:** — +- **provenance:** FIX undefined cm/st imports +- **note:** ridge plot of Δ-entropy posteriors + +### `pl.phenotypic_flux` — _distribution_ · keep + +```python +pl.phenotypic_flux(adata, *, ..., palette=None, figsize=..., ax=None, save=None) +``` +- **returns:** Axes | Figure | None +- **writes:** — +- **reads:** (via pp.joint_distribution) +- **calls:** pp.joint_distribution, _build_sankey +- **upstream:** matching tl tool +- **invariants:** reads computed state / calls tl; never owns model math +- **edge cases:** — +- **provenance:** keep; absorbs plot_pheno_sankey as private _build_sankey +- **note:** phenotype-flow sankey across covariates + +### `pl.polar_plot` — _distribution_ · fix + +```python +pl.polar_plot(adata, *, ..., palette=None, figsize=..., ax=None, save=None) +``` +- **returns:** Axes | Figure | None +- **writes:** — +- **reads:** (via pp.joint_distribution / tl.clonotypic_entropy) +- **calls:** pp.joint_distribution, tl.clonotypic_entropy +- **upstream:** matching tl tool +- **invariants:** reads computed state / calls tl; never owns model math +- **edge cases:** — +- **provenance:** FIX undefined clonotypic_entropy ref + string phenotypes +- **note:** radar of phenotype distribution or entropy + +### `pl.clone_size_umap` — _umap_ · keep + +```python +pl.clone_size_umap(adata, *, ..., palette=None, figsize=..., ax=None, save=None) +``` +- **returns:** Axes | Figure | None +- **writes:** — +- **reads:** obs[clone_size], obsm[X_umap] +- **calls:** pp.clone_size +- **upstream:** matching tl tool +- **invariants:** reads computed state / calls tl; never owns model math +- **edge cases:** — +- **provenance:** keep +- **note:** UMAP colored by log clone size + +### `pl.top_clone_umap` — _umap_ · keep + +```python +pl.top_clone_umap(adata, *, ..., palette=None, figsize=..., ax=None, save=None) +``` +- **returns:** Axes | Figure | None +- **writes:** — +- **reads:** obs[clone_col], obsm[X_umap] +- **calls:** — +- **upstream:** matching tl tool +- **invariants:** reads computed state / calls tl; never owns model math +- **edge cases:** — +- **provenance:** keep +- **note:** UMAP highlighting top-N clones + +### `pl.phenotype_probabilities_umap` — _umap_ · rename + +```python +pl.phenotype_probabilities_umap(adata, *, ..., palette=None, figsize=..., ax=None, save=None) +``` +- **returns:** Axes | Figure | None +- **writes:** — +- **reads:** obsm[X_tcri_probabilities] +- **calls:** — +- **upstream:** matching tl tool +- **invariants:** reads computed state / calls tl; never owns model math +- **edge cases:** — +- **provenance:** rename of plot_phenotype_probabilities +- **note:** per-phenotype probability UMAP grid + +### `pl.model_loss` — _diagnostic_ · move + +```python +pl.model_loss(adata, *, ..., palette=None, figsize=..., ax=None, save=None) +``` +- **returns:** Axes | Figure | None +- **writes:** — +- **reads:** model.history_ +- **calls:** — +- **upstream:** model.train +- **invariants:** reads computed state / calls tl; never owns model math +- **edge cases:** — +- **provenance:** moved from TCRIModel.plot_loss +- **note:** ELBO + dKL training curves + +### `pl.archetypes` — _diagnostic_ · move + +```python +pl.archetypes(adata, *, ..., palette=None, figsize=..., ax=None, save=None) +``` +- **returns:** Axes | Figure | None +- **writes:** — +- **reads:** model.c2p_mat, model.centers +- **calls:** — +- **upstream:** model.train +- **invariants:** reads computed state / calls tl; never owns model math +- **edge cases:** — +- **provenance:** moved from TCRIModel.plot_archetypes +- **note:** archetype / clone-prior heatmaps + +### `pl.model_pgm` — _diagnostic_ · move+merge + +```python +pl.model_pgm(adata, *, ..., palette=None, figsize=..., ax=None, save=None) +``` +- **returns:** Axes | Figure | None +- **writes:** — +- **reads:** — (static daft diagram) +- **calls:** — +- **upstream:** — +- **invariants:** reads computed state / calls tl; never owns model math +- **edge cases:** — +- **provenance:** moved from ut.build_nested_tcri_pgm + ut.draw_tcri_pgm_nested (merged) +- **note:** TCRI plate-diagram (PGM) + +### `ut.save_tcri_session` — _session_ · keep + +```python +ut.save_tcri_session(model, adata, out_dir, *, save_adata=True, compression='gzip') +``` +- **returns:** dict{paths} +- **writes:** run_dir/{model, pyro_params.pt, setup.json, adata.h5ad, meta.json} +- **reads:** model.save, pyro store, adata +- **calls:** write_adata_safely, _collect_setup_from_adata_or_model +- **upstream:** train +- **invariants:** adata written without tcri_manager (non-picklable) +- **edge cases:** — +- **provenance:** keep as-is + +### `ut.load_tcri_session` — _session_ · keep + +```python +ut.load_tcri_session(run_dir, *, adata_path=None, map_location=None, layer=None) +``` +- **returns:** (model, adata) +- **writes:** pyro param store (restored) +- **reads:** run_dir artifacts +- **calls:** TCRIModel.setup_anndata/.load, _pyro_load, _restore_category_order, _ensure_pyro_posterior_params, _disable_scvi_onload_train +- **upstream:** save_tcri_session +- **invariants:** category order restored; posterior params ensured +- **edge cases:** missing pyro store → warn + uniform-prior fallback +- **provenance:** keep as-is + +### `ut.write_adata_safely` — _session_ · keep + +```python +ut.write_adata_safely(adata, path, *, compression='gzip') +``` +- **returns:** None +- **writes:** path (h5ad without tcri_manager) +- **reads:** adata +- **calls:** _pop_nonserializables +- **upstream:** — +- **invariants:** strips non-serializable manager before write +- **edge cases:** — +- **provenance:** keep as-is + +## Shared primitives + +**Console (tcri/_console.py) — replaces 3 duplicated copies** + +- `_ok / _info / _warn / _fin` — every verbose tl/pp/pl/ml function +- `_ascii_hist(samples)` — mutual_information, flux, delta_clonotypic_entropy (graph=) +- `ANSI constants (RESET/BOLD/GRN/...)` — all of the above + +**Stats (tcri/_stats.py)** + +- `_joint_to_mi(pxy, normalised, mode)` — tl.mutual_information +- `distance(metric) -> f(p,q) [l1 | dkl | callable]` — tl.flux, tl.flux_table +- `auc_and_label_permutation(scores, labels)` — pl.mi_compare +- `bootstrap_auc(scores, labels)` — pl.mi_compare +- `_norm_entropy(p, base, n)` — clonotypic_entropy, phenotypic_entropy + +**Table builder (tcri/metrics/_tables.py)** + +- `_group_table(adata, splitby, per_group_fn)` — mi_compare, delta_entropy_table, flux_table + +**Plot helpers (tcri/plotting/_base.py)** + +- `_metric_boxplot(adata, fn, ...)` — clonality (+ any group×split metric box) +- `_resolve_palette / tcri_colors` — every pl function +- `_build_sankey / SankeyNode / _phenotype_mass_per_clone` — phenotypic_flux + +**Session internals (tcri/utils/_io.py)** + +- `_ensure_dir, _to_jsonable, _pop_nonserializables` — save_tcri_session, write_adata_safely +- `_collect_setup_from_adata_or_model, _restore_category_order` — save/load_tcri_session +- `_pyro_load, _resolve_TCRIModel, _disable_scvi_onload_train, _ensure_pyro_posterior_params` — load_tcri_session + +## Dropped / merged from the current code + +| current | disposition | why | +|---|---|---| +| `pp.classify_phenotypes` | deleted | redundant cosine phenotype assignment; 0 callers | +| `pp.get_latent_embedding` | deleted | trivial gaussian sampler; 0 callers | +| `pp.register_probability_columns` | deleted | feeds only the dead probabilities() | +| `pp.gene_entropy` | deleted | generic gene QC, out of scope; 0 callers | +| `pp.group_small_clones` | subsumed → group_singletons | hardcoded inferior dup | +| `pp.register_phenotype_key / register_clonotype_key` | subsumed → register_model | kill the tcri_*_key shadow convention | +| `pp._compute_logits_and_prior` | subsumed → ml.get_cell_phenotype_probs | one prob path | +| `pp.joint_distribution_posterior` | subsumed → joint_distribution(posterior=True) | unified engine | +| `tl.clonotypic_entropy_base` | subsumed → clonotypic_entropy(phenotype=) | one entropy fn | +| `tl.clone_fraction` | deleted | one-line value-counts; 0 callers | +| `tl.dkl` | subsumed → _stats.distance | dead; flux had its own copy | +| `pl.compare_phenotypes` | deleted | trivial crosstab heatmap; 0 callers | +| `pl.compare_joint_distribution` | deleted | broken (undefined `model`) | +| `pl.probability_distribution` | deleted | broken (self-recursion); covered by polar_plot | +| `pl.set_color_palette` | deleted | buggy palette helper; 0 callers | +| `pl.plot_pheno_sankey` | subsumed → phenotypic_flux (_build_sankey) | one public sankey | +| `pl.tcri_boxplot` | internalized → _metric_boxplot | shared plot helper | +| `ut.probabilities` | deleted | reads uns[joint_distribution] nothing ever writes | +| `ut.stars` | deleted | imported but never called | +| `ut.build_nested_tcri_pgm / draw_tcri_pgm_nested` | moved+merged → pl.model_pgm | describes the model, not IO | +| `ut.auc_and_label_permutation / bootstrap_auc` | internalized → _stats | used only by pl.mi_compare | +| `SankeyNode.hex_to_rgb` | deleted | unused method | +| `_ok/_info/_warn/_fin (×3), _ascii_hist (×2)` | subsumed → _console | dedup | + +_See `tcri_dependency_map.md` for the full call + producer/consumer graph._ \ No newline at end of file diff --git a/docs/contract/tcri_consistency_sufficiency.md b/docs/contract/tcri_consistency_sufficiency.md new file mode 100644 index 0000000..b3cf36d --- /dev/null +++ b/docs/contract/tcri_consistency_sufficiency.md @@ -0,0 +1,45 @@ +# TCRI — Argument Consistency & Sufficiency (clean re-derivation) + +Replaces the quarantined `tcri_arg_consistency_sufficiency.md`. **Consistency** is checked over the API surface +only. **Sufficiency** is derived from each kept function's **responsibility** — the disposable notebooks are never +consulted (Hard Rules, `REFACTOR_HISTORY.md`). Findings here are deltas to fold into +`tcri_api_and_responsibilities.md` signatures. + +## A. Consistency — canonical decisions (fold into the signatures) + +| # | concept | drift found | CANONICAL | +|---|---|---|---| +| C1 | RNG arg | `seed=0` (diags) vs `random_state` (tl/pl) vs `n_perm` | **`random_state=None`** everywhere, typed `int \| numpy.Generator \| torch.Generator \| None`, keyword-only | +| C2 | arg order | `clones`/`temperature` order varies | canonical: `covariate`/`cov_*` · `groupby` · `splitby` · `n_samples` · `temperature` · `clones` · `use_logits` · `normalized` · `normalize_mode` · `distance_metric` · `random_state` · `device` | +| C3 | distance selector | `metric=` (`phenotype_distance`) vs `distance_metric=` (flux) | **`distance_metric`** everywhere | +| C4 | palette arg | `phenotype_colors=` (flux) vs `palette=` (rest) | **`palette`** | +| C5 | normalize flag | `normalize=` (flux) vs `normalized=` (metrics) | **`normalized`** (adjective), American spelling | +| C6 | category ordering | `order=`/`hue_order=` present on some pl plots, absent on siblings | expose **both** `order=None` and `hue_order=None` on all metric box plots | +| C7 | figsize | diverges across sibling box plots | unify `figsize=(8,4)` for the metric-plot family | +| C8 | keyword-only | `columns`/`function`/`adata` positional in places | only the primary data object is positional; everything else after `*` | +| C9 | scanpy `show=` | only on private `_finish` | thread `show=None` through every public `pl` function | +| C10 | PPC sample knob | `reconstruction_ppc(n_samples=100)` collides with the `n_samples=0` convention | rename to **`n_sims=100`**; reserve `n_samples` (default 0) for the point/draws convention | +| C11 | `compare_groups` axis | `by=` for the contrast column | rename **`by=`→`splitby=`** to match the two-axis vocab | +| C12 | `group_singletons` defaults | `clonotype_key="trb"`, `groupby="patient"` hardcoded | require or read the registered defaults (no dataset-specific literals) | +| C13 | eps floor | four different probability-clip epsilons | one `eps=1e-12` for the distance/MI/normalization paths (incl. the MI kernel) | +| C14 | spelling | British `normalise*` in the draft | **American** (`normalized`, `normalize_mode`, `color`) — already applied | +| C15 | `covariate=None` | one line said "required" on the `adata_or_jd` fast path | `covariate=None` = **all covariate values** everywhere; the jd fast path ignores it — already reconciled | + +## B. Sufficiency — from responsibility (kept functions only) + +| function | knob | why (its own responsibility) | +|---|---|---| +| `joint_distribution`, 4 metrics, sampling `diag` | `random_state=` | reproducible sampling is intrinsic to a posterior-draw function | +| per-clone metrics + `joint_distribution_ppc` | `clones=` | the function computes per-clone quantities → subsetting is intrinsic | +| the metric box plots | `order=` / `hue_order=` | a categorical plot must control axis / hue order | +| `compare_groups` | `pair_on=` | a **paired** covariate contrast (Pre→Post per unit) needs the unit-alignment column — a real responsibility, not notebook parity | + +## C. REJECTED — notebook-driven "sufficiency", NOT re-added + +These were flagged only because a disposable notebook used them. They correspond to **dropped** functions or +**deliberately removed** behavior, and are not re-added: + +- `gene_entropy`, `probability_ternary`, `top_clone_umap`, `clone_size_umap`, `plot_phenotype_probabilities` — **dropped** (not core). No args to "restore." +- `weighted=` — deliberately removed (uniform-clonotype prior; `REFACTOR_HISTORY` §2). +- `minimum_clone_size=`, `base=`, `decimals=` — notebook conveniences on dropped/reworked functions; not core responsibilities. +- `covariate_key=` — **invalid**: the covariate column is fixed by `setup_anndata` and baked into `p_ct` via `ct_to_cov`; a call-time override would silently mismatch the model. Correctly absent. diff --git a/docs/contract/tcri_dependency_map.dot b/docs/contract/tcri_dependency_map.dot new file mode 100644 index 0000000..46e17aa --- /dev/null +++ b/docs/contract/tcri_dependency_map.dot @@ -0,0 +1,151 @@ +digraph tcri { + rankdir=TB; node [fontname="Helvetica",fontsize=10]; + edge [fontname="Helvetica",fontsize=8]; + subgraph cluster_pl { label="pl"; style=rounded; color="#2563eb"; + "pl.mutual_information" [shape=box,style=filled,fillcolor="#dbe8fd",color="#2563eb"]; + "pl.clonotypic_entropy" [shape=box,style=filled,fillcolor="#dbe8fd",color="#2563eb"]; + "pl.phenotypic_entropy" [shape=box,style=filled,fillcolor="#dbe8fd",color="#2563eb"]; + "pl.clonality" [shape=box,style=filled,fillcolor="#dbe8fd",color="#2563eb"]; + "pl.flux" [shape=box,style=filled,fillcolor="#dbe8fd",color="#2563eb"]; + "pl.mi_compare" [shape=box,style=filled,fillcolor="#dbe8fd",color="#2563eb"]; + "pl.bayesian_mutual_information" [shape=box,style=filled,fillcolor="#dbe8fd",color="#2563eb"]; + "pl.ridge_delta_entropy" [shape=box,style=filled,fillcolor="#dbe8fd",color="#2563eb"]; + "pl.phenotypic_flux" [shape=box,style=filled,fillcolor="#dbe8fd",color="#2563eb"]; + "pl.polar_plot" [shape=box,style=filled,fillcolor="#dbe8fd",color="#2563eb"]; + "pl.clone_size_umap" [shape=box,style=filled,fillcolor="#dbe8fd",color="#2563eb"]; + "pl.top_clone_umap" [shape=box,style=filled,fillcolor="#dbe8fd",color="#2563eb"]; + "pl.phenotype_probabilities_umap" [shape=box,style=filled,fillcolor="#dbe8fd",color="#2563eb"]; + "pl.model_loss" [shape=box,style=filled,fillcolor="#dbe8fd",color="#2563eb"]; + "pl.archetypes" [shape=box,style=filled,fillcolor="#dbe8fd",color="#2563eb"]; + "pl.model_pgm" [shape=box,style=filled,fillcolor="#dbe8fd",color="#2563eb"]; + } + subgraph cluster_tl { label="tl"; style=rounded; color="#7c3aed"; + "tl.clonotypic_entropy" [shape=box,style=filled,fillcolor="#ece3fb",color="#7c3aed"]; + "tl.phenotypic_entropy" [shape=box,style=filled,fillcolor="#ece3fb",color="#7c3aed"]; + "tl.mutual_information" [shape=box,style=filled,fillcolor="#ece3fb",color="#7c3aed"]; + "tl.clonality" [shape=box,style=filled,fillcolor="#ece3fb",color="#7c3aed"]; + "tl.flux" [shape=box,style=filled,fillcolor="#ece3fb",color="#7c3aed"]; + "tl.delta_clonotypic_entropy" [shape=box,style=filled,fillcolor="#ece3fb",color="#7c3aed"]; + "tl.mi_compare" [shape=box,style=filled,fillcolor="#ece3fb",color="#7c3aed"]; + "tl.delta_entropy_table" [shape=box,style=filled,fillcolor="#ece3fb",color="#7c3aed"]; + "tl.flux_table" [shape=box,style=filled,fillcolor="#ece3fb",color="#7c3aed"]; + } + subgraph cluster_pp { label="pp"; style=rounded; color="#0f766e"; + "pp.register_model" [shape=box,style=filled,fillcolor="#d6f0ec",color="#0f766e"]; + "pp.joint_distribution" [shape=box,style=filled,fillcolor="#d6f0ec",color="#0f766e"]; + "pp.group_singletons" [shape=box,style=filled,fillcolor="#d6f0ec",color="#0f766e"]; + "pp.clone_size" [shape=box,style=filled,fillcolor="#d6f0ec",color="#0f766e"]; + "pp.filter_genes" [shape=box,style=filled,fillcolor="#d6f0ec",color="#0f766e"]; + } + subgraph cluster_ml { label="ml"; style=rounded; color="#dc2626"; + "ml.setup_anndata" [shape=box,style=filled,fillcolor="#fde2e2",color="#dc2626"]; + "ml.train" [shape=box,style=filled,fillcolor="#fde2e2",color="#dc2626"]; + "ml.get_latent_representation" [shape=box,style=filled,fillcolor="#fde2e2",color="#dc2626"]; + "ml.get_cell_phenotype_probs" [shape=box,style=filled,fillcolor="#fde2e2",color="#dc2626"]; + "ml.get_p_ct" [shape=box,style=filled,fillcolor="#fde2e2",color="#dc2626"]; + "ml.boost_phenotype_prior" [shape=box,style=filled,fillcolor="#fde2e2",color="#dc2626"]; + } + subgraph cluster_ut { label="ut"; style=rounded; color="#475569"; + "ut.save_tcri_session" [shape=box,style=filled,fillcolor="#e5e9ef",color="#475569"]; + "ut.load_tcri_session" [shape=box,style=filled,fillcolor="#e5e9ef",color="#475569"]; + "ut.write_adata_safely" [shape=box,style=filled,fillcolor="#e5e9ef",color="#475569"]; + } + subgraph cluster_core { label="core"; style=rounded; color="#b45309"; + "core._joint_to_mi" [shape=box,style=filled,fillcolor="#fef3c7",color="#b45309"]; + "core._stats.distance" [shape=box,style=filled,fillcolor="#fef3c7",color="#b45309"]; + "core._stats.auc_perm" [shape=box,style=filled,fillcolor="#fef3c7",color="#b45309"]; + "core._stats.bootstrap_auc" [shape=box,style=filled,fillcolor="#fef3c7",color="#b45309"]; + "core._group_table" [shape=box,style=filled,fillcolor="#fef3c7",color="#b45309"]; + "core._metric_boxplot" [shape=box,style=filled,fillcolor="#fef3c7",color="#b45309"]; + "core._build_sankey" [shape=box,style=filled,fillcolor="#fef3c7",color="#b45309"]; + } + // call edges (solid) + "pl.mutual_information" -> "tl.mutual_information" [color="#475569"]; + "pl.clonotypic_entropy" -> "tl.clonotypic_entropy" [color="#475569"]; + "pl.phenotypic_entropy" -> "tl.phenotypic_entropy" [color="#475569"]; + "pl.clonality" -> "tl.clonality" [color="#475569"]; + "pl.clonality" -> "core._metric_boxplot" [color="#475569"]; + "pl.flux" -> "tl.flux" [color="#475569"]; + "pl.mi_compare" -> "tl.mi_compare" [color="#475569"]; + "pl.mi_compare" -> "core._stats.auc_perm" [color="#475569"]; + "pl.mi_compare" -> "core._stats.bootstrap_auc" [color="#475569"]; + "pl.bayesian_mutual_information" -> "tl.mutual_information" [color="#475569"]; + "pl.ridge_delta_entropy" -> "tl.delta_entropy_table" [color="#475569"]; + "pl.phenotypic_flux" -> "pp.joint_distribution" [color="#475569"]; + "pl.phenotypic_flux" -> "core._build_sankey" [color="#475569"]; + "pl.polar_plot" -> "pp.joint_distribution" [color="#475569"]; + "pl.polar_plot" -> "tl.clonotypic_entropy" [color="#475569"]; + "pl.clone_size_umap" -> "pp.clone_size" [color="#475569"]; + "pl.model_loss" -> "ml.train" [color="#475569"]; + "pl.archetypes" -> "ml.train" [color="#475569"]; + "tl.clonotypic_entropy" -> "pp.joint_distribution" [color="#475569"]; + "tl.phenotypic_entropy" -> "pp.joint_distribution" [color="#475569"]; + "tl.mutual_information" -> "pp.joint_distribution" [color="#475569"]; + "tl.mutual_information" -> "core._joint_to_mi" [color="#475569"]; + "tl.flux" -> "pp.joint_distribution" [color="#475569"]; + "tl.flux" -> "core._stats.distance" [color="#475569"]; + "tl.delta_clonotypic_entropy" -> "tl.clonotypic_entropy" [color="#475569"]; + "tl.mi_compare" -> "tl.mutual_information" [color="#475569"]; + "tl.mi_compare" -> "core._group_table" [color="#475569"]; + "tl.delta_entropy_table" -> "tl.delta_clonotypic_entropy" [color="#475569"]; + "tl.delta_entropy_table" -> "core._group_table" [color="#475569"]; + "tl.flux_table" -> "tl.flux" [color="#475569"]; + "tl.flux_table" -> "core._group_table" [color="#475569"]; + "pp.register_model" -> "ml.get_cell_phenotype_probs" [color="#475569"]; + "pp.register_model" -> "ml.get_latent_representation" [color="#475569"]; + "pp.register_model" -> "ml.get_p_ct" [color="#475569"]; + "ml.get_cell_phenotype_probs" -> "ml.get_p_ct" [color="#475569"]; + "ut.save_tcri_session" -> "ut.write_adata_safely" [color="#475569"]; + "ut.load_tcri_session" -> "ml.setup_anndata" [color="#475569"]; + // state keys + dataflow (write=bold, read=dashed) + "uns[tcri_metadata]" [shape=cylinder,style=filled,fillcolor="#f8fafc",color="#334155"]; + "pp.register_model" -> "uns[tcri_metadata]" [color="#0f766e",penwidth=1.6]; + "uns[tcri_metadata]" -> "pp.joint_distribution" [style=dashed,color="#94a3b8"]; + "uns[tcri_metadata]" -> "tl.clonality" [style=dashed,color="#94a3b8"]; + "uns[tcri_metadata]" -> "tl.mi_compare" [style=dashed,color="#94a3b8"]; + "uns[tcri_metadata]" -> "tl.delta_entropy_table" [style=dashed,color="#94a3b8"]; + "uns[tcri_metadata]" -> "tl.flux_table" [style=dashed,color="#94a3b8"]; + "uns[tcri_metadata]" -> "pp.clone_size" [style=dashed,color="#94a3b8"]; + "uns[tcri_p_ct]" [shape=cylinder,style=filled,fillcolor="#f8fafc",color="#334155"]; + "pp.register_model" -> "uns[tcri_p_ct]" [color="#0f766e",penwidth=1.6]; + "uns[tcri_p_ct]" -> "pp.joint_distribution" [style=dashed,color="#94a3b8"]; + "uns[tcri_local_scale]" [shape=cylinder,style=filled,fillcolor="#f8fafc",color="#334155"]; + "pp.register_model" -> "uns[tcri_local_scale]" [color="#0f766e",penwidth=1.6]; + "uns[tcri_local_scale]" -> "pp.joint_distribution" [style=dashed,color="#94a3b8"]; + "uns[tcri_*_categories]" [shape=cylinder,style=filled,fillcolor="#f8fafc",color="#334155"]; + "pp.register_model" -> "uns[tcri_*_categories]" [color="#0f766e",penwidth=1.6]; + "uns[tcri_*_categories]" -> "pp.joint_distribution" [style=dashed,color="#94a3b8"]; + "uns[tcri_*_categories]" -> "tl.clonotypic_entropy" [style=dashed,color="#94a3b8"]; + "uns[tcri_*_categories]" -> "tl.phenotypic_entropy" [style=dashed,color="#94a3b8"]; + "uns[tcri_ct_to_cov/ct_to_c]" [shape=cylinder,style=filled,fillcolor="#f8fafc",color="#334155"]; + "pp.register_model" -> "uns[tcri_ct_to_cov/ct_to_c]" [color="#0f766e",penwidth=1.6]; + "uns[tcri_ct_to_cov/ct_to_c]" -> "pp.joint_distribution" [style=dashed,color="#94a3b8"]; + "uns[tcri_{ct,cov}_array_for_cells]" [shape=cylinder,style=filled,fillcolor="#f8fafc",color="#334155"]; + "pp.register_model" -> "uns[tcri_{ct,cov}_array_for_cells]" [color="#0f766e",penwidth=1.6]; + "uns[tcri_{ct,cov}_array_for_cells]" -> "pp.joint_distribution" [style=dashed,color="#94a3b8"]; + "obsm[X_tcri_logits]" [shape=cylinder,style=filled,fillcolor="#f8fafc",color="#334155"]; + "pp.register_model" -> "obsm[X_tcri_logits]" [color="#0f766e",penwidth=1.6]; + "obsm[X_tcri_logits]" -> "pp.joint_distribution" [style=dashed,color="#94a3b8"]; + "obsm[X_tcri_probabilities]" [shape=cylinder,style=filled,fillcolor="#f8fafc",color="#334155"]; + "pp.register_model" -> "obsm[X_tcri_probabilities]" [color="#0f766e",penwidth=1.6]; + "obsm[X_tcri_probabilities]" -> "pl.phenotype_probabilities_umap" [style=dashed,color="#94a3b8"]; + "obsm[X_tcri]" [shape=cylinder,style=filled,fillcolor="#f8fafc",color="#334155"]; + "pp.register_model" -> "obsm[X_tcri]" [color="#0f766e",penwidth=1.6]; + "obsm[X_tcri]" -> "pl.clone_size_umap" [style=dashed,color="#94a3b8"]; + "obsm[X_tcri]" -> "pl.top_clone_umap" [style=dashed,color="#94a3b8"]; + "obs[tcri_phenotype]" [shape=cylinder,style=filled,fillcolor="#f8fafc",color="#334155"]; + "pp.register_model" -> "obs[tcri_phenotype]" [color="#0f766e",penwidth=1.6]; + "obs[tcri_phenotype]" -> "tl.clonality" [style=dashed,color="#94a3b8"]; + "obs[tcri_phenotype]" -> "pl.top_clone_umap" [style=dashed,color="#94a3b8"]; + "obs[clone_size]" [shape=cylinder,style=filled,fillcolor="#f8fafc",color="#334155"]; + "pp.clone_size" -> "obs[clone_size]" [color="#0f766e",penwidth=1.6]; + "obs[clone_size]" -> "pl.clone_size_umap" [style=dashed,color="#94a3b8"]; + "obs[trb_unique]" [shape=cylinder,style=filled,fillcolor="#f8fafc",color="#334155"]; + "pp.group_singletons" -> "obs[trb_unique]" [color="#0f766e",penwidth=1.6]; + "obs[trb_unique]" -> "ml.setup_anndata" [style=dashed,color="#94a3b8"]; + "pyro_param_store" [shape=cylinder,style=filled,fillcolor="#f8fafc",color="#334155"]; + "ml.train" -> "pyro_param_store" [color="#0f766e",penwidth=1.6]; + "pyro_param_store" -> "ml.get_p_ct" [style=dashed,color="#94a3b8"]; + "pyro_param_store" -> "ut.save_tcri_session" [style=dashed,color="#94a3b8"]; + "pyro_param_store" -> "ut.load_tcri_session" [style=dashed,color="#94a3b8"]; +} \ No newline at end of file diff --git a/docs/contract/tcri_dependency_map.html b/docs/contract/tcri_dependency_map.html new file mode 100644 index 0000000..6f2fc48 --- /dev/null +++ b/docs/contract/tcri_dependency_map.html @@ -0,0 +1,257 @@ +TCRI dependency map

TCRI — Dependency Map (target API)

Call graph + adata-state producer/consumer links · curated from the target contract.

mlpptlplutcore
Two graphs, same nodes. (1) Call graph: pl → tl → pp → ml, bottoming out at the model + shared core primitives. (2) Dataflow: everything funnels through the tcri_* state that pp.register_model writes; pp.joint_distribution is the universal consumer hub.

Call graph

flowchart TB + subgraph PL + n_pl_mutual_information["pl.mutual_information"] + n_pl_clonotypic_entropy["pl.clonotypic_entropy"] + n_pl_phenotypic_entropy["pl.phenotypic_entropy"] + n_pl_clonality["pl.clonality"] + n_pl_flux["pl.flux"] + n_pl_mi_compare["pl.mi_compare"] + n_pl_bayesian_mutual_information["pl.bayesian_mutual_information"] + n_pl_ridge_delta_entropy["pl.ridge_delta_entropy"] + n_pl_phenotypic_flux["pl.phenotypic_flux"] + n_pl_polar_plot["pl.polar_plot"] + n_pl_clone_size_umap["pl.clone_size_umap"] + n_pl_model_loss["pl.model_loss"] + n_pl_archetypes["pl.archetypes"] + end + subgraph TL + n_tl_clonotypic_entropy["tl.clonotypic_entropy"] + n_tl_phenotypic_entropy["tl.phenotypic_entropy"] + n_tl_mutual_information["tl.mutual_information"] + n_tl_clonality["tl.clonality"] + n_tl_flux["tl.flux"] + n_tl_delta_clonotypic_entropy["tl.delta_clonotypic_entropy"] + n_tl_mi_compare["tl.mi_compare"] + n_tl_delta_entropy_table["tl.delta_entropy_table"] + n_tl_flux_table["tl.flux_table"] + end + subgraph PP + n_pp_register_model["pp.register_model"] + n_pp_joint_distribution["pp.joint_distribution"] + n_pp_clone_size["pp.clone_size"] + end + subgraph ML + n_ml_setup_anndata["ml.setup_anndata"] + n_ml_train["ml.train"] + n_ml_get_latent_representation["ml.get_latent_representation"] + n_ml_get_cell_phenotype_probs["ml.get_cell_phenotype_probs"] + n_ml_get_p_ct["ml.get_p_ct"] + end + subgraph CORE + n_core__joint_to_mi["core._joint_to_mi"] + n_core__stats_distance["core._stats.distance"] + n_core__stats_auc_perm["core._stats.auc_perm"] + n_core__stats_bootstrap_auc["core._stats.bootstrap_auc"] + n_core__group_table["core._group_table"] + n_core__metric_boxplot["core._metric_boxplot"] + n_core__build_sankey["core._build_sankey"] + end + subgraph UT + n_ut_save_tcri_session["ut.save_tcri_session"] + n_ut_load_tcri_session["ut.load_tcri_session"] + n_ut_write_adata_safely["ut.write_adata_safely"] + end + n_pl_mutual_information --> n_tl_mutual_information + n_pl_clonotypic_entropy --> n_tl_clonotypic_entropy + n_pl_phenotypic_entropy --> n_tl_phenotypic_entropy + n_pl_clonality --> n_tl_clonality + n_pl_clonality --> n_core__metric_boxplot + n_pl_flux --> n_tl_flux + n_pl_mi_compare --> n_tl_mi_compare + n_pl_mi_compare --> n_core__stats_auc_perm + n_pl_mi_compare --> n_core__stats_bootstrap_auc + n_pl_bayesian_mutual_information --> n_tl_mutual_information + n_pl_ridge_delta_entropy --> n_tl_delta_entropy_table + n_pl_phenotypic_flux --> n_pp_joint_distribution + n_pl_phenotypic_flux --> n_core__build_sankey + n_pl_polar_plot --> n_pp_joint_distribution + n_pl_polar_plot --> n_tl_clonotypic_entropy + n_pl_clone_size_umap --> n_pp_clone_size + n_pl_model_loss --> n_ml_train + n_pl_archetypes --> n_ml_train + n_tl_clonotypic_entropy --> n_pp_joint_distribution + n_tl_phenotypic_entropy --> n_pp_joint_distribution + n_tl_mutual_information --> n_pp_joint_distribution + n_tl_mutual_information --> n_core__joint_to_mi + n_tl_flux --> n_pp_joint_distribution + n_tl_flux --> n_core__stats_distance + n_tl_delta_clonotypic_entropy --> n_tl_clonotypic_entropy + n_tl_mi_compare --> n_tl_mutual_information + n_tl_mi_compare --> n_core__group_table + n_tl_delta_entropy_table --> n_tl_delta_clonotypic_entropy + n_tl_delta_entropy_table --> n_core__group_table + n_tl_flux_table --> n_tl_flux + n_tl_flux_table --> n_core__group_table + n_pp_register_model --> n_ml_get_cell_phenotype_probs + n_pp_register_model --> n_ml_get_latent_representation + n_pp_register_model --> n_ml_get_p_ct + n_ml_get_cell_phenotype_probs --> n_ml_get_p_ct + n_ut_save_tcri_session --> n_ut_write_adata_safely + n_ut_load_tcri_session --> n_ml_setup_anndata + classDef ml fill:#fde2e2,stroke:#dc2626,color:#111; + classDef pp fill:#d6f0ec,stroke:#0f766e,color:#111; + classDef tl fill:#ece3fb,stroke:#7c3aed,color:#111; + classDef pl fill:#dbe8fd,stroke:#2563eb,color:#111; + classDef ut fill:#e5e9ef,stroke:#475569,color:#111; + classDef core fill:#fef3c7,stroke:#b45309,color:#111; + classDef state fill:#ffffff,stroke:#334155,color:#111; + class n_ml_setup_anndata ml; + class n_ml_train ml; + class n_ml_get_latent_representation ml; + class n_ml_get_cell_phenotype_probs ml; + class n_ml_get_p_ct ml; + class n_pp_register_model pp; + class n_pp_joint_distribution pp; + class n_pp_clone_size pp; + class n_tl_clonotypic_entropy tl; + class n_tl_phenotypic_entropy tl; + class n_tl_mutual_information tl; + class n_tl_clonality tl; + class n_tl_flux tl; + class n_tl_delta_clonotypic_entropy tl; + class n_tl_mi_compare tl; + class n_tl_delta_entropy_table tl; + class n_tl_flux_table tl; + class n_pl_mutual_information pl; + class n_pl_clonotypic_entropy pl; + class n_pl_phenotypic_entropy pl; + class n_pl_clonality pl; + class n_pl_flux pl; + class n_pl_mi_compare pl; + class n_pl_bayesian_mutual_information pl; + class n_pl_ridge_delta_entropy pl; + class n_pl_phenotypic_flux pl; + class n_pl_polar_plot pl; + class n_pl_clone_size_umap pl; + class n_pl_model_loss pl; + class n_pl_archetypes pl; + class n_ut_save_tcri_session ut; + class n_ut_load_tcri_session ut; + class n_ut_write_adata_safely ut; + class n_core__joint_to_mi core; + class n_core__stats_distance core; + class n_core__stats_auc_perm core; + class n_core__stats_bootstrap_auc core; + class n_core__group_table core; + class n_core__metric_boxplot core; + class n_core__build_sankey core;

Call adjacency

functioncallscalled by
ml.setup_anndata—ut.load_tcri_session
ml.train—pl.model_loss, pl.archetypes
ml.get_latent_representation—pp.register_model
ml.get_cell_phenotype_probsml.get_p_ctpp.register_model
ml.get_p_ct—pp.register_model, ml.get_cell_phenotype_probs
pp.register_modelml.get_cell_phenotype_probs, ml.get_latent_representation, ml.get_p_ct—
pp.joint_distribution—pl.phenotypic_flux, pl.polar_plot, tl.clonotypic_entropy, tl.phenotypic_entropy, tl.mutual_information, tl.flux
pp.clone_size—pl.clone_size_umap
tl.clonotypic_entropypp.joint_distributionpl.clonotypic_entropy, pl.polar_plot, tl.delta_clonotypic_entropy
tl.phenotypic_entropypp.joint_distributionpl.phenotypic_entropy
tl.mutual_informationpp.joint_distribution, core._joint_to_mipl.mutual_information, pl.bayesian_mutual_information, tl.mi_compare
tl.clonality—pl.clonality
tl.fluxpp.joint_distribution, core._stats.distancepl.flux, tl.flux_table
tl.delta_clonotypic_entropytl.clonotypic_entropytl.delta_entropy_table
tl.mi_comparetl.mutual_information, core._group_tablepl.mi_compare
tl.delta_entropy_tabletl.delta_clonotypic_entropy, core._group_tablepl.ridge_delta_entropy
tl.flux_tabletl.flux, core._group_table—
pl.mutual_informationtl.mutual_information—
pl.clonotypic_entropytl.clonotypic_entropy—
pl.phenotypic_entropytl.phenotypic_entropy—
pl.clonalitytl.clonality, core._metric_boxplot—
pl.fluxtl.flux—
pl.mi_comparetl.mi_compare, core._stats.auc_perm, core._stats.bootstrap_auc—
pl.bayesian_mutual_informationtl.mutual_information—
pl.ridge_delta_entropytl.delta_entropy_table—
pl.phenotypic_fluxpp.joint_distribution, core._build_sankey—
pl.polar_plotpp.joint_distribution, tl.clonotypic_entropy—
pl.clone_size_umappp.clone_size—
pl.model_lossml.train—
pl.archetypesml.train—
ut.save_tcri_sessionut.write_adata_safely—
ut.load_tcri_sessionml.setup_anndata—
ut.write_adata_safely—ut.save_tcri_session
core._joint_to_mi—tl.mutual_information
core._stats.distance—tl.flux
core._stats.auc_perm—pl.mi_compare
core._stats.bootstrap_auc—pl.mi_compare
core._group_table—tl.mi_compare, tl.delta_entropy_table, tl.flux_table
core._metric_boxplot—pl.clonality
core._build_sankey—pl.phenotypic_flux
Entry points (no caller): pp.register_model, pp.group_singletons, pp.filter_genes, pl.mutual_information, pl.clonotypic_entropy, pl.phenotypic_entropy, pl.clonality, pl.flux, pl.mi_compare, pl.bayesian_mutual_information, pl.ridge_delta_entropy, pl.phenotypic_flux, pl.polar_plot, pl.clone_size_umap, pl.top_clone_umap, pl.phenotype_probabilities_umap, pl.model_loss, pl.archetypes, pl.model_pgm, ut.save_tcri_session, ut.load_tcri_session
Leaves (call nothing in-package): ml.setup_anndata, ml.train, ml.get_latent_representation, ml.get_p_ct, ml.boost_phenotype_prior, pp.joint_distribution, pp.group_singletons, pp.clone_size, pp.filter_genes, core._joint_to_mi, core._stats.distance, core._stats.auc_perm, core._stats.bootstrap_auc, core._group_table, core._metric_boxplot, core._build_sankey

Dataflow (producers / consumers)

flowchart LR + n_KEY_uns_tcri_metadata_[("uns[tcri_metadata]")] + class n_KEY_uns_tcri_metadata_ state; + n_pp_register_model["pp.register_model"] + class n_pp_register_model pp; + n_pp_register_model ==> n_KEY_uns_tcri_metadata_ + n_pp_joint_distribution["pp.joint_distribution"] + class n_pp_joint_distribution pp; + n_KEY_uns_tcri_metadata_ -.-> n_pp_joint_distribution + n_tl_clonality["tl.clonality"] + class n_tl_clonality tl; + n_KEY_uns_tcri_metadata_ -.-> n_tl_clonality + n_tl_mi_compare["tl.mi_compare"] + class n_tl_mi_compare tl; + n_KEY_uns_tcri_metadata_ -.-> n_tl_mi_compare + n_tl_delta_entropy_table["tl.delta_entropy_table"] + class n_tl_delta_entropy_table tl; + n_KEY_uns_tcri_metadata_ -.-> n_tl_delta_entropy_table + n_tl_flux_table["tl.flux_table"] + class n_tl_flux_table tl; + n_KEY_uns_tcri_metadata_ -.-> n_tl_flux_table + n_pp_clone_size["pp.clone_size"] + class n_pp_clone_size pp; + n_KEY_uns_tcri_metadata_ -.-> n_pp_clone_size + n_KEY_uns_tcri_p_ct_[("uns[tcri_p_ct]")] + class n_KEY_uns_tcri_p_ct_ state; + n_pp_register_model ==> n_KEY_uns_tcri_p_ct_ + n_KEY_uns_tcri_p_ct_ -.-> n_pp_joint_distribution + n_KEY_uns_tcri_local_scale_[("uns[tcri_local_scale]")] + class n_KEY_uns_tcri_local_scale_ state; + n_pp_register_model ==> n_KEY_uns_tcri_local_scale_ + n_KEY_uns_tcri_local_scale_ -.-> n_pp_joint_distribution + n_KEY_uns_tcri___categories_[("uns[tcri_*_categories]")] + class n_KEY_uns_tcri___categories_ state; + n_pp_register_model ==> n_KEY_uns_tcri___categories_ + n_KEY_uns_tcri___categories_ -.-> n_pp_joint_distribution + n_tl_clonotypic_entropy["tl.clonotypic_entropy"] + class n_tl_clonotypic_entropy tl; + n_KEY_uns_tcri___categories_ -.-> n_tl_clonotypic_entropy + n_tl_phenotypic_entropy["tl.phenotypic_entropy"] + class n_tl_phenotypic_entropy tl; + n_KEY_uns_tcri___categories_ -.-> n_tl_phenotypic_entropy + n_KEY_uns_tcri_ct_to_cov_ct_to_c_[("uns[tcri_ct_to_cov/ct_to_c]")] + class n_KEY_uns_tcri_ct_to_cov_ct_to_c_ state; + n_pp_register_model ==> n_KEY_uns_tcri_ct_to_cov_ct_to_c_ + n_KEY_uns_tcri_ct_to_cov_ct_to_c_ -.-> n_pp_joint_distribution + n_KEY_uns_tcri__ct_cov__array_for_cells_[("uns[tcri_{ct,cov}_array_for_cells]")] + class n_KEY_uns_tcri__ct_cov__array_for_cells_ state; + n_pp_register_model ==> n_KEY_uns_tcri__ct_cov__array_for_cells_ + n_KEY_uns_tcri__ct_cov__array_for_cells_ -.-> n_pp_joint_distribution + n_KEY_obsm_X_tcri_logits_[("obsm[X_tcri_logits]")] + class n_KEY_obsm_X_tcri_logits_ state; + n_pp_register_model ==> n_KEY_obsm_X_tcri_logits_ + n_KEY_obsm_X_tcri_logits_ -.-> n_pp_joint_distribution + n_KEY_obsm_X_tcri_probabilities_[("obsm[X_tcri_probabilities]")] + class n_KEY_obsm_X_tcri_probabilities_ state; + n_pp_register_model ==> n_KEY_obsm_X_tcri_probabilities_ + n_pl_phenotype_probabilities_umap["pl.phenotype_probabilities_umap"] + class n_pl_phenotype_probabilities_umap pl; + n_KEY_obsm_X_tcri_probabilities_ -.-> n_pl_phenotype_probabilities_umap + n_KEY_obsm_X_tcri_[("obsm[X_tcri]")] + class n_KEY_obsm_X_tcri_ state; + n_pp_register_model ==> n_KEY_obsm_X_tcri_ + n_pl_clone_size_umap["pl.clone_size_umap"] + class n_pl_clone_size_umap pl; + n_KEY_obsm_X_tcri_ -.-> n_pl_clone_size_umap + n_pl_top_clone_umap["pl.top_clone_umap"] + class n_pl_top_clone_umap pl; + n_KEY_obsm_X_tcri_ -.-> n_pl_top_clone_umap + n_KEY_obs_tcri_phenotype_[("obs[tcri_phenotype]")] + class n_KEY_obs_tcri_phenotype_ state; + n_pp_register_model ==> n_KEY_obs_tcri_phenotype_ + n_KEY_obs_tcri_phenotype_ -.-> n_tl_clonality + n_KEY_obs_tcri_phenotype_ -.-> n_pl_top_clone_umap + n_KEY_obs_clone_size_[("obs[clone_size]")] + class n_KEY_obs_clone_size_ state; + n_pp_clone_size ==> n_KEY_obs_clone_size_ + n_KEY_obs_clone_size_ -.-> n_pl_clone_size_umap + n_KEY_obs_trb_unique_[("obs[trb_unique]")] + class n_KEY_obs_trb_unique_ state; + n_pp_group_singletons["pp.group_singletons"] + class n_pp_group_singletons pp; + n_pp_group_singletons ==> n_KEY_obs_trb_unique_ + n_ml_setup_anndata["ml.setup_anndata"] + class n_ml_setup_anndata ml; + n_KEY_obs_trb_unique_ -.-> n_ml_setup_anndata + n_KEY_pyro_param_store[("pyro_param_store")] + class n_KEY_pyro_param_store state; + n_ml_train["ml.train"] + class n_ml_train ml; + n_ml_train ==> n_KEY_pyro_param_store + n_ml_get_p_ct["ml.get_p_ct"] + class n_ml_get_p_ct ml; + n_KEY_pyro_param_store -.-> n_ml_get_p_ct + n_ut_save_tcri_session["ut.save_tcri_session"] + class n_ut_save_tcri_session ut; + n_KEY_pyro_param_store -.-> n_ut_save_tcri_session + n_ut_load_tcri_session["ut.load_tcri_session"] + class n_ut_load_tcri_session ut; + n_KEY_pyro_param_store -.-> n_ut_load_tcri_session + classDef ml fill:#fde2e2,stroke:#dc2626,color:#111; + classDef pp fill:#d6f0ec,stroke:#0f766e,color:#111; + classDef tl fill:#ece3fb,stroke:#7c3aed,color:#111; + classDef pl fill:#dbe8fd,stroke:#2563eb,color:#111; + classDef ut fill:#e5e9ef,stroke:#475569,color:#111; + classDef core fill:#fef3c7,stroke:#b45309,color:#111; + classDef state fill:#ffffff,stroke:#334155,color:#111;
adata keyproduced byconsumed by
uns[tcri_metadata]pp.register_modelpp.joint_distribution, tl.clonality, tl.mi_compare, tl.delta_entropy_table, tl.flux_table, pp.clone_size
uns[tcri_p_ct]pp.register_modelpp.joint_distribution
uns[tcri_local_scale]pp.register_modelpp.joint_distribution
uns[tcri_*_categories]pp.register_modelpp.joint_distribution, tl.clonotypic_entropy, tl.phenotypic_entropy
uns[tcri_ct_to_cov/ct_to_c]pp.register_modelpp.joint_distribution
uns[tcri_{ct,cov}_array_for_cells]pp.register_modelpp.joint_distribution
obsm[X_tcri_logits]pp.register_modelpp.joint_distribution
obsm[X_tcri_probabilities]pp.register_modelpl.phenotype_probabilities_umap
obsm[X_tcri]pp.register_modelpl.clone_size_umap, pl.top_clone_umap
obs[tcri_phenotype]pp.register_modeltl.clonality, pl.top_clone_umap
obs[clone_size]pp.clone_sizepl.clone_size_umap
obs[trb_unique]pp.group_singletonsml.setup_anndata
pyro_param_storeml.trainml.get_p_ct, ut.save_tcri_session, ut.load_tcri_session
Metrics that read only uns[tcri_metadata] + categories via pp.joint_distribution are attributed to that hub, not re-listed per key.

How this is built / standard tooling

Call graph (who-calls-whom). The de-facto standard is a directed graph rendered with Graphviz/DOT. Auto-extract from source with the stdlib ast module, or tools like pyan3, code2flow, pydeps (module-level), or griffe (the engine behind mkdocstrings). For Markdown-native rendering use Mermaid flowchart (GitHub renders it inline); for interactive web use cytoscape.js or d3.

Dataflow / producer-consumer. Because TCRI couples through adata.uns/obsm/obs keys (not just direct calls), the precise model is a bipartite graph of functions ↔ state keys. This is exactly a build-system DAG: state keys are the artifacts/targets, functions are the rules. The standard tools for that shape are Make, Snakemake, or dbt (dbt docs renders an interactive lineage graph) — and the same idea is what scverse calls a data-flow/provenance graph. Here it is curated by hand from the contract so it can lead the refactor rather than trail it.

Preventing drift after the refactor. Once the target lands, point an ast walker at tcri/ to extract the actual call edges + uns/obsm/obs read/writes and diff them against this curated graph in CI. Drift = a function that calls or reads/writes something the contract doesn't list.

Graphviz

Combined DOT emitted to tcri_dependency_map.dot (call edges solid, writes bold-teal, reads dashed). Render: dot -Tsvg tcri_dependency_map.dot -o tcri_dependency_map.svg.

\ No newline at end of file diff --git a/docs/contract/tcri_dependency_map.md b/docs/contract/tcri_dependency_map.md new file mode 100644 index 0000000..9488324 --- /dev/null +++ b/docs/contract/tcri_dependency_map.md @@ -0,0 +1,344 @@ +# TCRI — Dependency Map (target API) + +_Call graph + `adata`-state producer/consumer links for the post-refactor API. Curated from `build_tcri_contract.py`; regenerate with `build_tcri_depgraph.py`._ + +**Two graphs, same nodes.** (1) **Call graph** — `pl` → `tl` → `pp` → `ml`, bottoming out at the model + shared `core` primitives. (2) **Dataflow** — everything funnels through the `tcri_*` state that `pp.register_model` writes; `pp.joint_distribution` is the universal consumer hub. + +## Call graph + +```mermaid +flowchart TB + subgraph PL + n_pl_mutual_information["pl.mutual_information"] + n_pl_clonotypic_entropy["pl.clonotypic_entropy"] + n_pl_phenotypic_entropy["pl.phenotypic_entropy"] + n_pl_clonality["pl.clonality"] + n_pl_flux["pl.flux"] + n_pl_mi_compare["pl.mi_compare"] + n_pl_bayesian_mutual_information["pl.bayesian_mutual_information"] + n_pl_ridge_delta_entropy["pl.ridge_delta_entropy"] + n_pl_phenotypic_flux["pl.phenotypic_flux"] + n_pl_polar_plot["pl.polar_plot"] + n_pl_clone_size_umap["pl.clone_size_umap"] + n_pl_model_loss["pl.model_loss"] + n_pl_archetypes["pl.archetypes"] + end + subgraph TL + n_tl_clonotypic_entropy["tl.clonotypic_entropy"] + n_tl_phenotypic_entropy["tl.phenotypic_entropy"] + n_tl_mutual_information["tl.mutual_information"] + n_tl_clonality["tl.clonality"] + n_tl_flux["tl.flux"] + n_tl_delta_clonotypic_entropy["tl.delta_clonotypic_entropy"] + n_tl_mi_compare["tl.mi_compare"] + n_tl_delta_entropy_table["tl.delta_entropy_table"] + n_tl_flux_table["tl.flux_table"] + end + subgraph PP + n_pp_register_model["pp.register_model"] + n_pp_joint_distribution["pp.joint_distribution"] + n_pp_clone_size["pp.clone_size"] + end + subgraph ML + n_ml_setup_anndata["ml.setup_anndata"] + n_ml_train["ml.train"] + n_ml_get_latent_representation["ml.get_latent_representation"] + n_ml_get_cell_phenotype_probs["ml.get_cell_phenotype_probs"] + n_ml_get_p_ct["ml.get_p_ct"] + end + subgraph CORE + n_core__joint_to_mi["core._joint_to_mi"] + n_core__stats_distance["core._stats.distance"] + n_core__stats_auc_perm["core._stats.auc_perm"] + n_core__stats_bootstrap_auc["core._stats.bootstrap_auc"] + n_core__group_table["core._group_table"] + n_core__metric_boxplot["core._metric_boxplot"] + n_core__build_sankey["core._build_sankey"] + end + subgraph UT + n_ut_save_tcri_session["ut.save_tcri_session"] + n_ut_load_tcri_session["ut.load_tcri_session"] + n_ut_write_adata_safely["ut.write_adata_safely"] + end + n_pl_mutual_information --> n_tl_mutual_information + n_pl_clonotypic_entropy --> n_tl_clonotypic_entropy + n_pl_phenotypic_entropy --> n_tl_phenotypic_entropy + n_pl_clonality --> n_tl_clonality + n_pl_clonality --> n_core__metric_boxplot + n_pl_flux --> n_tl_flux + n_pl_mi_compare --> n_tl_mi_compare + n_pl_mi_compare --> n_core__stats_auc_perm + n_pl_mi_compare --> n_core__stats_bootstrap_auc + n_pl_bayesian_mutual_information --> n_tl_mutual_information + n_pl_ridge_delta_entropy --> n_tl_delta_entropy_table + n_pl_phenotypic_flux --> n_pp_joint_distribution + n_pl_phenotypic_flux --> n_core__build_sankey + n_pl_polar_plot --> n_pp_joint_distribution + n_pl_polar_plot --> n_tl_clonotypic_entropy + n_pl_clone_size_umap --> n_pp_clone_size + n_pl_model_loss --> n_ml_train + n_pl_archetypes --> n_ml_train + n_tl_clonotypic_entropy --> n_pp_joint_distribution + n_tl_phenotypic_entropy --> n_pp_joint_distribution + n_tl_mutual_information --> n_pp_joint_distribution + n_tl_mutual_information --> n_core__joint_to_mi + n_tl_flux --> n_pp_joint_distribution + n_tl_flux --> n_core__stats_distance + n_tl_delta_clonotypic_entropy --> n_tl_clonotypic_entropy + n_tl_mi_compare --> n_tl_mutual_information + n_tl_mi_compare --> n_core__group_table + n_tl_delta_entropy_table --> n_tl_delta_clonotypic_entropy + n_tl_delta_entropy_table --> n_core__group_table + n_tl_flux_table --> n_tl_flux + n_tl_flux_table --> n_core__group_table + n_pp_register_model --> n_ml_get_cell_phenotype_probs + n_pp_register_model --> n_ml_get_latent_representation + n_pp_register_model --> n_ml_get_p_ct + n_ml_get_cell_phenotype_probs --> n_ml_get_p_ct + n_ut_save_tcri_session --> n_ut_write_adata_safely + n_ut_load_tcri_session --> n_ml_setup_anndata + classDef ml fill:#fde2e2,stroke:#dc2626,color:#111; + classDef pp fill:#d6f0ec,stroke:#0f766e,color:#111; + classDef tl fill:#ece3fb,stroke:#7c3aed,color:#111; + classDef pl fill:#dbe8fd,stroke:#2563eb,color:#111; + classDef ut fill:#e5e9ef,stroke:#475569,color:#111; + classDef core fill:#fef3c7,stroke:#b45309,color:#111; + classDef state fill:#ffffff,stroke:#334155,color:#111; + class n_ml_setup_anndata ml; + class n_ml_train ml; + class n_ml_get_latent_representation ml; + class n_ml_get_cell_phenotype_probs ml; + class n_ml_get_p_ct ml; + class n_pp_register_model pp; + class n_pp_joint_distribution pp; + class n_pp_clone_size pp; + class n_tl_clonotypic_entropy tl; + class n_tl_phenotypic_entropy tl; + class n_tl_mutual_information tl; + class n_tl_clonality tl; + class n_tl_flux tl; + class n_tl_delta_clonotypic_entropy tl; + class n_tl_mi_compare tl; + class n_tl_delta_entropy_table tl; + class n_tl_flux_table tl; + class n_pl_mutual_information pl; + class n_pl_clonotypic_entropy pl; + class n_pl_phenotypic_entropy pl; + class n_pl_clonality pl; + class n_pl_flux pl; + class n_pl_mi_compare pl; + class n_pl_bayesian_mutual_information pl; + class n_pl_ridge_delta_entropy pl; + class n_pl_phenotypic_flux pl; + class n_pl_polar_plot pl; + class n_pl_clone_size_umap pl; + class n_pl_model_loss pl; + class n_pl_archetypes pl; + class n_ut_save_tcri_session ut; + class n_ut_load_tcri_session ut; + class n_ut_write_adata_safely ut; + class n_core__joint_to_mi core; + class n_core__stats_distance core; + class n_core__stats_auc_perm core; + class n_core__stats_bootstrap_auc core; + class n_core__group_table core; + class n_core__metric_boxplot core; + class n_core__build_sankey core; +``` + +### Call adjacency + +| function | calls | called by | +|---|---|---| +| `ml.setup_anndata` | — | `ut.load_tcri_session` | +| `ml.train` | — | `pl.model_loss`, `pl.archetypes` | +| `ml.get_latent_representation` | — | `pp.register_model` | +| `ml.get_cell_phenotype_probs` | `ml.get_p_ct` | `pp.register_model` | +| `ml.get_p_ct` | — | `pp.register_model`, `ml.get_cell_phenotype_probs` | +| `pp.register_model` | `ml.get_cell_phenotype_probs`, `ml.get_latent_representation`, `ml.get_p_ct` | — | +| `pp.joint_distribution` | — | `pl.phenotypic_flux`, `pl.polar_plot`, `tl.clonotypic_entropy`, `tl.phenotypic_entropy`, `tl.mutual_information`, `tl.flux` | +| `pp.clone_size` | — | `pl.clone_size_umap` | +| `tl.clonotypic_entropy` | `pp.joint_distribution` | `pl.clonotypic_entropy`, `pl.polar_plot`, `tl.delta_clonotypic_entropy` | +| `tl.phenotypic_entropy` | `pp.joint_distribution` | `pl.phenotypic_entropy` | +| `tl.mutual_information` | `pp.joint_distribution`, `core._joint_to_mi` | `pl.mutual_information`, `pl.bayesian_mutual_information`, `tl.mi_compare` | +| `tl.clonality` | — | `pl.clonality` | +| `tl.flux` | `pp.joint_distribution`, `core._stats.distance` | `pl.flux`, `tl.flux_table` | +| `tl.delta_clonotypic_entropy` | `tl.clonotypic_entropy` | `tl.delta_entropy_table` | +| `tl.mi_compare` | `tl.mutual_information`, `core._group_table` | `pl.mi_compare` | +| `tl.delta_entropy_table` | `tl.delta_clonotypic_entropy`, `core._group_table` | `pl.ridge_delta_entropy` | +| `tl.flux_table` | `tl.flux`, `core._group_table` | — | +| `pl.mutual_information` | `tl.mutual_information` | — | +| `pl.clonotypic_entropy` | `tl.clonotypic_entropy` | — | +| `pl.phenotypic_entropy` | `tl.phenotypic_entropy` | — | +| `pl.clonality` | `tl.clonality`, `core._metric_boxplot` | — | +| `pl.flux` | `tl.flux` | — | +| `pl.mi_compare` | `tl.mi_compare`, `core._stats.auc_perm`, `core._stats.bootstrap_auc` | — | +| `pl.bayesian_mutual_information` | `tl.mutual_information` | — | +| `pl.ridge_delta_entropy` | `tl.delta_entropy_table` | — | +| `pl.phenotypic_flux` | `pp.joint_distribution`, `core._build_sankey` | — | +| `pl.polar_plot` | `pp.joint_distribution`, `tl.clonotypic_entropy` | — | +| `pl.clone_size_umap` | `pp.clone_size` | — | +| `pl.model_loss` | `ml.train` | — | +| `pl.archetypes` | `ml.train` | — | +| `ut.save_tcri_session` | `ut.write_adata_safely` | — | +| `ut.load_tcri_session` | `ml.setup_anndata` | — | +| `ut.write_adata_safely` | — | `ut.save_tcri_session` | +| `core._joint_to_mi` | — | `tl.mutual_information` | +| `core._stats.distance` | — | `tl.flux` | +| `core._stats.auc_perm` | — | `pl.mi_compare` | +| `core._stats.bootstrap_auc` | — | `pl.mi_compare` | +| `core._group_table` | — | `tl.mi_compare`, `tl.delta_entropy_table`, `tl.flux_table` | +| `core._metric_boxplot` | — | `pl.clonality` | +| `core._build_sankey` | — | `pl.phenotypic_flux` | + +**Entry points** (no caller): `pp.register_model`, `pp.group_singletons`, `pp.filter_genes`, `pl.mutual_information`, `pl.clonotypic_entropy`, `pl.phenotypic_entropy`, `pl.clonality`, `pl.flux`, `pl.mi_compare`, `pl.bayesian_mutual_information`, `pl.ridge_delta_entropy`, `pl.phenotypic_flux`, `pl.polar_plot`, `pl.clone_size_umap`, `pl.top_clone_umap`, `pl.phenotype_probabilities_umap`, `pl.model_loss`, `pl.archetypes`, `pl.model_pgm`, `ut.save_tcri_session`, `ut.load_tcri_session` + +**Leaves** (call nothing in-package): `ml.setup_anndata`, `ml.train`, `ml.get_latent_representation`, `ml.get_p_ct`, `ml.boost_phenotype_prior`, `pp.joint_distribution`, `pp.group_singletons`, `pp.clone_size`, `pp.filter_genes`, `core._joint_to_mi`, `core._stats.distance`, `core._stats.auc_perm`, `core._stats.bootstrap_auc`, `core._group_table`, `core._metric_boxplot`, `core._build_sankey` + +## Dataflow (producers / consumers) + +```mermaid +flowchart LR + n_KEY_uns_tcri_metadata_[("uns[tcri_metadata]")] + class n_KEY_uns_tcri_metadata_ state; + n_pp_register_model["pp.register_model"] + class n_pp_register_model pp; + n_pp_register_model ==> n_KEY_uns_tcri_metadata_ + n_pp_joint_distribution["pp.joint_distribution"] + class n_pp_joint_distribution pp; + n_KEY_uns_tcri_metadata_ -.-> n_pp_joint_distribution + n_tl_clonality["tl.clonality"] + class n_tl_clonality tl; + n_KEY_uns_tcri_metadata_ -.-> n_tl_clonality + n_tl_mi_compare["tl.mi_compare"] + class n_tl_mi_compare tl; + n_KEY_uns_tcri_metadata_ -.-> n_tl_mi_compare + n_tl_delta_entropy_table["tl.delta_entropy_table"] + class n_tl_delta_entropy_table tl; + n_KEY_uns_tcri_metadata_ -.-> n_tl_delta_entropy_table + n_tl_flux_table["tl.flux_table"] + class n_tl_flux_table tl; + n_KEY_uns_tcri_metadata_ -.-> n_tl_flux_table + n_pp_clone_size["pp.clone_size"] + class n_pp_clone_size pp; + n_KEY_uns_tcri_metadata_ -.-> n_pp_clone_size + n_KEY_uns_tcri_p_ct_[("uns[tcri_p_ct]")] + class n_KEY_uns_tcri_p_ct_ state; + n_pp_register_model ==> n_KEY_uns_tcri_p_ct_ + n_KEY_uns_tcri_p_ct_ -.-> n_pp_joint_distribution + n_KEY_uns_tcri_local_scale_[("uns[tcri_local_scale]")] + class n_KEY_uns_tcri_local_scale_ state; + n_pp_register_model ==> n_KEY_uns_tcri_local_scale_ + n_KEY_uns_tcri_local_scale_ -.-> n_pp_joint_distribution + n_KEY_uns_tcri___categories_[("uns[tcri_*_categories]")] + class n_KEY_uns_tcri___categories_ state; + n_pp_register_model ==> n_KEY_uns_tcri___categories_ + n_KEY_uns_tcri___categories_ -.-> n_pp_joint_distribution + n_tl_clonotypic_entropy["tl.clonotypic_entropy"] + class n_tl_clonotypic_entropy tl; + n_KEY_uns_tcri___categories_ -.-> n_tl_clonotypic_entropy + n_tl_phenotypic_entropy["tl.phenotypic_entropy"] + class n_tl_phenotypic_entropy tl; + n_KEY_uns_tcri___categories_ -.-> n_tl_phenotypic_entropy + n_KEY_uns_tcri_ct_to_cov_ct_to_c_[("uns[tcri_ct_to_cov/ct_to_c]")] + class n_KEY_uns_tcri_ct_to_cov_ct_to_c_ state; + n_pp_register_model ==> n_KEY_uns_tcri_ct_to_cov_ct_to_c_ + n_KEY_uns_tcri_ct_to_cov_ct_to_c_ -.-> n_pp_joint_distribution + n_KEY_uns_tcri__ct_cov__array_for_cells_[("uns[tcri_{ct,cov}_array_for_cells]")] + class n_KEY_uns_tcri__ct_cov__array_for_cells_ state; + n_pp_register_model ==> n_KEY_uns_tcri__ct_cov__array_for_cells_ + n_KEY_uns_tcri__ct_cov__array_for_cells_ -.-> n_pp_joint_distribution + n_KEY_obsm_X_tcri_logits_[("obsm[X_tcri_logits]")] + class n_KEY_obsm_X_tcri_logits_ state; + n_pp_register_model ==> n_KEY_obsm_X_tcri_logits_ + n_KEY_obsm_X_tcri_logits_ -.-> n_pp_joint_distribution + n_KEY_obsm_X_tcri_probabilities_[("obsm[X_tcri_probabilities]")] + class n_KEY_obsm_X_tcri_probabilities_ state; + n_pp_register_model ==> n_KEY_obsm_X_tcri_probabilities_ + n_pl_phenotype_probabilities_umap["pl.phenotype_probabilities_umap"] + class n_pl_phenotype_probabilities_umap pl; + n_KEY_obsm_X_tcri_probabilities_ -.-> n_pl_phenotype_probabilities_umap + n_KEY_obsm_X_tcri_[("obsm[X_tcri]")] + class n_KEY_obsm_X_tcri_ state; + n_pp_register_model ==> n_KEY_obsm_X_tcri_ + n_pl_clone_size_umap["pl.clone_size_umap"] + class n_pl_clone_size_umap pl; + n_KEY_obsm_X_tcri_ -.-> n_pl_clone_size_umap + n_pl_top_clone_umap["pl.top_clone_umap"] + class n_pl_top_clone_umap pl; + n_KEY_obsm_X_tcri_ -.-> n_pl_top_clone_umap + n_KEY_obs_tcri_phenotype_[("obs[tcri_phenotype]")] + class n_KEY_obs_tcri_phenotype_ state; + n_pp_register_model ==> n_KEY_obs_tcri_phenotype_ + n_KEY_obs_tcri_phenotype_ -.-> n_tl_clonality + n_KEY_obs_tcri_phenotype_ -.-> n_pl_top_clone_umap + n_KEY_obs_clone_size_[("obs[clone_size]")] + class n_KEY_obs_clone_size_ state; + n_pp_clone_size ==> n_KEY_obs_clone_size_ + n_KEY_obs_clone_size_ -.-> n_pl_clone_size_umap + n_KEY_obs_trb_unique_[("obs[trb_unique]")] + class n_KEY_obs_trb_unique_ state; + n_pp_group_singletons["pp.group_singletons"] + class n_pp_group_singletons pp; + n_pp_group_singletons ==> n_KEY_obs_trb_unique_ + n_ml_setup_anndata["ml.setup_anndata"] + class n_ml_setup_anndata ml; + n_KEY_obs_trb_unique_ -.-> n_ml_setup_anndata + n_KEY_pyro_param_store[("pyro_param_store")] + class n_KEY_pyro_param_store state; + n_ml_train["ml.train"] + class n_ml_train ml; + n_ml_train ==> n_KEY_pyro_param_store + n_ml_get_p_ct["ml.get_p_ct"] + class n_ml_get_p_ct ml; + n_KEY_pyro_param_store -.-> n_ml_get_p_ct + n_ut_save_tcri_session["ut.save_tcri_session"] + class n_ut_save_tcri_session ut; + n_KEY_pyro_param_store -.-> n_ut_save_tcri_session + n_ut_load_tcri_session["ut.load_tcri_session"] + class n_ut_load_tcri_session ut; + n_KEY_pyro_param_store -.-> n_ut_load_tcri_session + classDef ml fill:#fde2e2,stroke:#dc2626,color:#111; + classDef pp fill:#d6f0ec,stroke:#0f766e,color:#111; + classDef tl fill:#ece3fb,stroke:#7c3aed,color:#111; + classDef pl fill:#dbe8fd,stroke:#2563eb,color:#111; + classDef ut fill:#e5e9ef,stroke:#475569,color:#111; + classDef core fill:#fef3c7,stroke:#b45309,color:#111; + classDef state fill:#ffffff,stroke:#334155,color:#111; +``` + +### State producer/consumer table + +| adata key | produced by | consumed by | +|---|---|---| +| `uns[tcri_metadata]` | `pp.register_model` | `pp.joint_distribution`, `tl.clonality`, `tl.mi_compare`, `tl.delta_entropy_table`, `tl.flux_table`, `pp.clone_size` | +| `uns[tcri_p_ct]` | `pp.register_model` | `pp.joint_distribution` | +| `uns[tcri_local_scale]` | `pp.register_model` | `pp.joint_distribution` | +| `uns[tcri_*_categories]` | `pp.register_model` | `pp.joint_distribution`, `tl.clonotypic_entropy`, `tl.phenotypic_entropy` | +| `uns[tcri_ct_to_cov/ct_to_c]` | `pp.register_model` | `pp.joint_distribution` | +| `uns[tcri_{ct,cov}_array_for_cells]` | `pp.register_model` | `pp.joint_distribution` | +| `obsm[X_tcri_logits]` | `pp.register_model` | `pp.joint_distribution` | +| `obsm[X_tcri_probabilities]` | `pp.register_model` | `pl.phenotype_probabilities_umap` | +| `obsm[X_tcri]` | `pp.register_model` | `pl.clone_size_umap`, `pl.top_clone_umap` | +| `obs[tcri_phenotype]` | `pp.register_model` | `tl.clonality`, `pl.top_clone_umap` | +| `obs[clone_size]` | `pp.clone_size` | `pl.clone_size_umap` | +| `obs[trb_unique]` | `pp.group_singletons` | `ml.setup_anndata` | +| `pyro_param_store` | `ml.train` | `ml.get_p_ct`, `ut.save_tcri_session`, `ut.load_tcri_session` | + +> Metrics that read only `uns[tcri_metadata]` + categories **via** `pp.joint_distribution` are attributed to that hub, not re-listed per key. + +## How this is built / standard tooling + +**Call graph (who-calls-whom).** The de-facto standard is a directed graph rendered with **Graphviz/DOT**. Auto-extract from source with the stdlib `ast` module, or tools like `pyan3`, `code2flow`, `pydeps` (module-level), or `griffe` (the engine behind mkdocstrings). For Markdown-native rendering use **Mermaid** `flowchart` (GitHub renders it inline); for interactive web use `cytoscape.js` or `d3`. + +**Dataflow / producer-consumer.** Because TCRI couples through `adata.uns/obsm/obs` keys (not just direct calls), the precise model is a **bipartite graph** of functions ↔ state keys. This is exactly a **build-system DAG**: state keys are the artifacts/targets, functions are the rules. The standard tools for that shape are **Make**, **Snakemake**, or **dbt** (`dbt docs` renders an interactive lineage graph) — and the same idea is what scverse calls a _data-flow_/provenance graph. Here it is curated by hand from the contract so it can lead the refactor rather than trail it. + +**Preventing drift after the refactor.** Once the target lands, point an `ast` walker at `tcri/` to extract the _actual_ call edges + `uns/obsm/obs` read/writes and diff them against this curated graph in CI. Drift = a function that calls or reads/writes something the contract doesn't list. + +## Graphviz + +A combined DOT file is emitted to `tcri_dependency_map.dot` (call edges solid, writes bold-teal, reads dashed). Render: + +```bash +dot -Tsvg tcri_dependency_map.dot -o tcri_dependency_map.svg +``` \ No newline at end of file diff --git a/docs/contract/tcri_function_inventory.md b/docs/contract/tcri_function_inventory.md new file mode 100644 index 0000000..626bde5 --- /dev/null +++ b/docs/contract/tcri_function_inventory.md @@ -0,0 +1,369 @@ +# TCRI — Full Function Inventory & Consolidation Plan + +_Generated from the inventory workflow (9 agents, 131 functions, completeness-verified) via `build_tcri_inventory.py`. Data: `tcri_inventory_data.json`. This is the working list we reduce the repo against._ + +## 0. Core definition + +**Core = five things; everything else must justify itself against them.** +1. **Model (`ml`)** — `TCRIModel`: build / train / evaluate / register outputs onto an AnnData. +2. **Engine (`pp`)** — `joint_distribution`: Bayesian posterior sampling of the clone×phenotype distribution (the substrate every metric reads). +3. **Metrics (`tl`)** — `clonotypic_entropy`, `phenotypic_entropy`, `mutual_information` (+ `flux`, `delta_clonotypic_entropy`, and the tidy-table builders). +4. **Plotting (`pl`)** — plots that *directly* visualize those metrics and the joint distribution / flux. +5. **Utils (`ut`) + shared helpers** — session save/load; deduplicated console / stats / distance / color helpers. + +## 1. Label counts (critic-corrected) + +| label | count | disposition | +|---|---|---| +| core | 22 | keep (22 survive) | +| redundant | 4 | merge into core (5 groups) | +| plotting-beyond-core | 10 | move to examples / drop | +| model-construction | 31 | keep, split across model/_module,_priors,_classifier,_training | +| session-io | 12 | keep as utils/_session | +| helper | 35 | dedupe → _console/_stats/_distance/_base/_colors | +| dead-broken | 17 | 12 delete · 2 merge · 3 keep+fix | +| **total** | **131** | | + +## 2. Grafiti reference layout (the target shape) + +The reference layout (`../grafiti`) is a flat package with scanpy-style sub-packages aliased to short handles +(`model→ml`, `tools→tl`, `plotting→pl`, `preprocessing→pp`, `diagnostics→diag`, `datasets→ds`, `get.py→get`). +Six patterns tcri should copy: + +1. **One file per topic, never a monolith.** `tools/_joint.py`, `tools/_motif.py`, … and `plotting/` mirrors them 1:1 by filename. (tcri's 1008-line `_metrics.py` and 1437-line `_plotting.py` are the anti-pattern.) +2. **Private cross-cutting sub-packages.** `_state/` (keys, resolve, storage, schemas) + `_compute/` (device-routed math) hold everything shared, so impl files stay thin. tcri analog: `_keys.py`, `_console.py`, `_stats.py`, `_distance.py` (and later a `_compute`). +3. **`__all__` at BOTH levels; NO `import *`.** Each impl module declares `__all__`; each `__init__` names every symbol explicitly and re-declares an aggregate `__all__` grouped by view. **This corrects the earlier `import *` instinct** — the mature pattern is explicit re-export, which keeps numpy/pandas/helpers *unexported*. +4. **`get.py` + `@tl_result` cache convention.** tl writes a versioned uns blob and returns a tidy result; `pl` functions are pure *cache renderers* (`load_result(adata, key)` → draw, never compute). *(tcri: adopt `_keys.py` now; defer the cache decorator — see open questions.)* +5. **`diagnostics/` returns DATA, not plots.** `gf.diag` runs read-only concordance/quality checks on the *finalized* model and returns a DataFrame ("did the model fit?"), deliberately outside the tl-writes / pl-reads loop. **This is exactly where PPCs + model-validation live.** +6. **Naming.** public package + private `_topic.py` impl modules; helper *packages* underscore-prefixed dirs, helper *files* underscore-prefixed, helper *functions* underscore-prefixed; tl↔pl twins share filename + function name. + +## 3. Target tcri layout (grafiti-mirrored) + +``` +tcri/ + __init__.py # explicit re-export + sys.modules aliases (tl/pp/pl/ml/ut/diag); NO import * + _keys.py # single source of every uns/obsm/obs key string + _console.py # _ok/_info/_warn/_fin/_ascii_hist (was triplicated across 3 files) + _stats.py # stars, auc_and_label_permutation, bootstrap_auc + _distance.py # kl_divergence, l1_distance, phenotype_distance dispatch (was dkl + flux.dkl_func) + model/ # ml + _model.py # TCRIModel: setup_anndata, train, get_latent_representation, get_cell_phenotype_probs, get_p_ct + _module.py # TCRIModule (pyro model/guide, get_latent, get_p_ct) + _priors.py # MixtureDirichlet, VampPrior + _classifier.py # PhenotypeClassifier + _training.py # UnifiedTrainingPlan, build_archetypes + preprocessing/ # pp + _register.py # register_model (+ folded register_*_key, _compute_logits_and_prior) + _engine.py # joint_distribution(posterior=, n_samples=) (unifies the two current fns) + _clones.py # group_singletons, clone_size + metrics/ # tl + _entropy.py # clonotypic_entropy, phenotypic_entropy, delta_clonotypic_entropy + _mutual_information.py # mutual_information (+ private _mi_from_joint) + _flux.py # flux + _tables.py # mi_compare, delta_entropy_table, flux_table + plotting/ # pl (twins mirror tl by filename) + _base.py # _metric_boxplot (was tcri_boxplot), _finish + _colors.py # tcri_colors, resolve_palette (was set_color_palette) + _entropy.py # clonotypic_entropy (was _by_phenotype), phenotypic_entropy [FIX], ridge_delta_entropy [FIX] + _mutual_information.py # mutual_information [FIX], mi_compare + _flux.py # phenotypic_flux (sankey) + _sankey.py # SankeyNode, _phenotype_mass_per_clone + diagnostics/ # diag (NEW — PPCs + model validation, returns DataFrames) + _ppc.py # joint-distribution PPC (was compare_joint_distribution, fixed) + calibration + reconstruction PPC + _training.py # loss curves (was plot_loss), archetypes (was plot_archetypes) + _pgm.py # model PGM (was build_nested_tcri_pgm) + utils/ # ut + _session.py # save/load_tcri_session, write_adata_safely, _to_jsonable +examples/ # bespoke one-offs move here: top_clone_umap, clone_size_umap, + # phenotype_probabilities UMAP, compare_phenotypes + the rewritten notebooks +``` + +## 4. Full inventory — every function + +Label is critic-corrected. Disposition folds in the five overlays (diagnostics reclass, keep+fix pl twins, n_samples convention). +### `tcri.ml` — model (42 records) + +| name | kind | label | disposition | purpose | +|---|---|---|---|---| +| `TCRIModel` | class | core | keep | High-level scvi model API: setup, build, train, and extract latent/phenotype/p_ct outputs. | +| `TCRIModel.__init__` | method | core | keep | Build clone->phenotype matrix, archetypes, clonotype-covariate index maps and class weights, then construct+prime TCRIModule. | +| `TCRIModel.get_cell_phenotype_probs` | method | core | keep | Per-cell phenotype probabilities by combining classifier logits with log p_ct prior (gate or additive), matching training. | +| `TCRIModel.get_latent_representation` | method | core | keep | Batched encode of an AnnData to a (n_cells, n_latent) numpy latent matrix. | +| `TCRIModel.get_p_ct` | method | core | keep | Return the learned clone-covariate x phenotype posterior p_ct as a numpy array. | +| `TCRIModel.setup_anndata` | method | core | keep | Register clonotype/phenotype/covariate/batch/count fields with scvi and stash the manager/layer on the AnnData. | +| `TCRIModel.train` | method | core | keep | Split data, build UnifiedTrainingPlan, and run TrainRunner with elbo_validation early stopping. | +| `TCRIModel.plot_archetypes` | method | plotting-beyond-core | → diagnostics | Heatmap the cluster-ordered clone-phenotype matrix and the archetype centroids. | +| `TCRIModel.plot_loss` | method | plotting-beyond-core | → diagnostics | Plot training/validation ELBO and prior-KL curves from self.history_. | +| `MixtureDirichlet` | class | model-construction | keep (internal) | Custom Pyro distribution: a mixture of Dirichlets over the phenotype simplex, used as the clonotype prior p_c. | +| `MixtureDirichlet.__call__` | method | model-construction | keep (internal) | Make the distribution callable as an alias for sample(). | +| `MixtureDirichlet.__init__` | method | model-construction | keep (internal) | Clamp concentrations, infer batch/B/K shapes, and init the TorchDistribution. | +| `MixtureDirichlet.log_prob` | method | model-construction | keep (internal) | Log-sum-exp of component Dirichlet log-probs weighted by (log) mixture weights. | +| `MixtureDirichlet.sample` | method | model-construction | keep (internal) | Sample a mixture component per batch element then draw from the selected Dirichlet. | +| `MixtureDirichlet.score_parts` | method | model-construction | keep (internal) | Return (log_prob, zero score-fn, zero entropy) so Pyro treats it as reparam-free. | +| `PhenotypeClassifier` | class | model-construction | keep (internal) | Temperature-scaled MLP head mapping latent z to phenotype logits. | +| `PhenotypeClassifier.__init__` | method | model-construction | keep (internal) | Build the stacked Linear/ReLU/Dropout MLP and store the softmax temperature. | +| `PhenotypeClassifier.forward` | method | model-construction | keep (internal) | Return MLP logits divided by temperature. | +| `TCRIModel.boost_phenotype_prior` | method | model-construction | keep (internal) | Multiply one phenotype's column in the clone prior (and optionally mixture centers) by a factor, renormalize, and overwrite module buffers. | +| `TCRIModule` | class | model-construction | keep (internal) | Pyro CVAE module with hierarchical clonotype->clonotype-covariate Dirichlet priors and a phenotype classifier. | +| `TCRIModule.__init__` | method | model-construction | keep (internal) | Construct encoder/decoder/classifier/VampPrior, px_r param, and register empty two-level buffers + class weights. | +| `TCRIModule._get_fn_args_from_batch` | method | model-construction | keep (internal) | Extract (x, batch_idx, log_library) tuple from a scvi batch dict for model()/guide(). | +| `TCRIModule.get_latent` | method | model-construction | keep (internal) | Encode a batch to the posterior-mean latent z_loc (collapses MC dim if present). | +| `TCRIModule.get_p_ct` | method | model-construction | keep (internal) | Read learned q_p_ct_raw from the Pyro param store and normalize (with guide-temperature) to the clone-covariate x phenotype posterior. | +| `TCRIModule.guide` | method | model-construction | keep (internal) | Variational guide: learnable Dirichlet params q(p_c), q(p_ct) and Normal q(z) from the encoder. | +| `TCRIModule.model` | method | model-construction | keep (internal) | Generative model: sample p_c (MixtureDirichlet), p_ct (Dirichlet), latent z (VampPrior), and ZINB gene obs. | +| `TCRIModule.prepare_two_level_params` | method | model-construction | keep (internal) | Normalize/temperature the clone-phenotype prior and register all two-level index buffers onto the module. | +| `TCRIModule.use_gate` | method | model-construction | keep (internal) | True when gate_prob is not None (selects convex-gate vs additive phenotype combination). | +| `UnifiedTrainingPlan` | class | model-construction | keep (internal) | Training plan adding KL warmup, diagnostics logging, and a validation_step emitting elbo_validation for early stopping. | +| `UnifiedTrainingPlan.__init__` | method | model-construction | keep (internal) | Choose TraceEnum_ELBO vs Trace_ELBO by module.use_enumeration and store optimizer/warmup config. | +| `UnifiedTrainingPlan.configure_optimizers` | method | model-construction | keep (internal) | Build an Adam optimizer over module parameters from optimizer_config. | +| `UnifiedTrainingPlan.loss` | method | model-construction | keep (internal) | Expose the configured ELBO loss object. | +| `UnifiedTrainingPlan.training_step` | method | model-construction | keep (internal) | Apply KL warmup, run the Pyro ELBO step, and log KL/entropy/confidence diagnostics. | +| `UnifiedTrainingPlan.validation_step` | method | model-construction | keep (internal) | Compute validation ELBO and prior-KL, logging elbo_validation for scvi early stopping. | +| `VampPrior` | class | model-construction | keep (internal) | VampPrior over latent z: a uniform mixture of encoder posteriors evaluated at learnable pseudo-inputs. | +| `VampPrior.__init__` | method | model-construction | keep (internal) | Register learnable pseudo-inputs as a Parameter and hold the shared encoder. | +| `VampPrior.get_mixture` | method | model-construction | keep (internal) | Encode pseudo-inputs and assemble a uniform MixtureSameFamily of Independent Normals as p(z). | +| `VampPrior.log_prob` | method | model-construction | keep (internal) | Log density of z under the VampPrior mixture. | +| `VampPrior.sample` | method | model-construction | keep (internal) | Draw samples from the VampPrior mixture. | +| `build_archetypes` | function | model-construction | keep (internal) | KMeans-cluster clone->phenotype rows into K normalized archetype centroids used as the Dirichlet-mixture prior concentration. | +| `TCRIModel.boost_phenotype_prior._ok` | inner-function | helper | extract → shared | Print a green check-marked status line. | +| `TCRIModel.use_gate` | method | helper | extract → shared | Public passthrough to module.use_gate. | + +### `tcri.pp` — preprocessing (20 records) + +| name | kind | label | disposition | purpose | +|---|---|---|---|---| +| `joint_distribution_posterior` | function | core | keep | Draw one posterior Dirichlet sample of p_ct, combine with per-cell classifier logits, and aggregate to a clone x phenotype distribution … | +| `register_model` | function | core | keep | Register all TCRIModel outputs (priors, metadata, categories, per-cell ct/cov arrays, latent means, logits, log-posterior, phenotype … | +| `classify_phenotypes` | function | redundant | merge → register_model | Alternate per-cell phenotype assignment via cosine similarity of latent to per-phenotype archetypes, reweighted by posterior p_ct. | +| `joint_distribution` | function | redundant | merge → joint_distribution_posterior | Build a clone x phenotype distribution for a covariate from ct-level p_ct point estimates (n_samples=0) or Dirichlet draws (n_samples>0), … | +| `_ascii_hist` | function | helper | extract → shared | Render a numpy histogram of samples as an ASCII bar chart string. | +| `_compute_logits_and_prior` | function | helper | extract → shared | Run the model encoder+classifier over a data loader to extract per-cell classifier logits and the log-prior from get_p_ct. | +| `_fin` | function | helper | extract → shared | Print a magenta 'Done!' flourish unless quiet. | +| `_info` | function | helper | extract → shared | Print a dim key-value info line unless quiet. | +| `_ok` | function | helper | extract → shared | Print a green success line unless quiet. | +| `_warn` | function | helper | extract → shared | Print a yellow warning line unless quiet. | +| `clone_size` | function | helper | extract → shared | Compute per-clone cell counts from the registered clone key and write them per cell into obs. | +| `group_singletons` | function | helper | extract → shared | Collapse clones smaller than min_clone_size (per groupby) into 'Singleton_{group}' labels. | +| `group_singletons.collapse_singleton` | inner-function | helper | extract → shared | Map a row to 'Singleton_{group}' when its candidate clone count < min_clone_size, else keep candidate. | +| `register_clonotype_key` | function | helper | extract → shared | Register the clonotype obs column and its unique categories into uns. | +| `register_phenotype_key` | function | helper | extract → shared | Register the phenotype obs column and its unique categories into uns. | +| `gene_entropy` | function | dead-broken | delete | Compute per-gene Shannon entropy of expression-value counts (optionally per batch, aggregated) into var. | +| `get_latent_embedding` | function | dead-broken | delete | Draw Gaussian samples around the stored latent means with a scalar std. | +| `group_small_clones` | function | dead-broken | delete | Collapse clones with clone_size<4 into 'Singleton_{patient}', else '{trb}_{patient}', into obs['trb_unique']. | +| `register_probability_columns` | function | dead-broken | delete | Store a probability_columns list into uns. | +| `remove_meaningless_genes` | function | dead-broken | delete | Intend to filter out MT/RP/HSP/MTRN/TCR/RIK/GM/LINC/ambiguous genes and return a sliced copy. | + +### `tcri.tl` — metrics (21 records) + +| name | kind | label | disposition | purpose | +|---|---|---|---|---| +| `_mi_from_joint` | function | core | keep | Compute (optionally normalised) mutual information from an already-normalised C×P joint table. | +| `clonotypic_entropy` | function | core | keep | Posterior mean (or per-draw matrix) normalised clonotypic entropy per phenotype at one covariate. | +| `delta_clonotypic_entropy` | function | core | keep | Monte-Carlo posterior samples of H_post−H_pre clonotypic entropy for one phenotype. | +| `delta_entropy_table` | function | core | keep | Tidy Δ-clonotypic-entropy table: one row per phenotype × splitby group with mean/sd/HDI/p and raw sample vector. | +| `flux` | function | core | keep | Per-clone phenotypic-distribution distance (l1/dkl/callable) between two covariates; point estimate or per-draw. | +| `flux_table` | function | core | keep | Tidy per-clone flux table: flux mean/sd + sample vector + clone size, per splitby group. | +| `mi_compare` | function | core | keep | Build a tidy per-patient MI samples+summary table across covariates and group pairs. | +| `mutual_information` | function | core | keep | Clone×phenotype mutual information at one covariate: point estimate (n_samples=0) or per-draw array. | +| `phenotypic_entropy` | function | core | keep | Posterior mean (or per-draw) normalised phenotypic entropy per clonotype at one covariate. | +| `clonality` | function | redundant | merge → clonotypic_entropy | Per-phenotype clonality 1 − H/log2K from observed hard clone-size counts. | +| `clonotypic_entropy_base` | function | redundant | merge → clonotypic_entropy | Single-draw normalised clonotypic entropy for ONE phenotype at ONE covariate. | +| `_ascii_hist` | function | helper | extract → shared | Build a text ASCII histogram of a sample vector for notebook/SSH display. | +| `_fin` | function | helper | extract → shared | Print a magenta 'Done!' flourish unless quiet. | +| `_info` | function | helper | extract → shared | Print a cyan key-value info line unless quiet. | +| `_ok` | function | helper | extract → shared | Print a green success/checkmark line unless quiet. | +| `_warn` | function | helper | extract → shared | Print a yellow warning line unless quiet. | +| `flux.dkl_func` | inner-function | helper | extract → shared | KL divergence kernel used when distance_metric=='dkl' inside flux. | +| `mutual_information._get_df` | inner-function | helper | extract → shared | Return one joint-distribution table for the MI computation (posterior draw). | +| `_ent` | function | dead-broken | delete | Normalise a vector and return its Shannon entropy in given base. | +| `clone_fraction` | function | dead-broken | delete | Per-group nested dict of each clone's frequency (count/total) within the group. | +| `dkl` | function | dead-broken | delete | KL divergence of p‖q via scipy.stats.entropy after clipping. | + +### `tcri.pl` — plotting (29 records) + +| name | kind | label | disposition | purpose | +|---|---|---|---|---| +| `clonotypic_entropy_by_phenotype` | function | core | keep | Box-and-dot plot of clonotypic entropy per phenotype and covariate, with bootstrap/MWU significance brackets and per-patient dots. | +| `mi_compare` | function | core | keep | Plot per-patient TCRi normalized MI across covariate group pairs as boxplots with jittered points and AUROC/MWU/label-permutation stats. | +| `phenotypic_flux` | function | core | keep | Public flux Sankey across `order` values of `splitby`; thin convenience wrapper adding x-ticks and save. | +| `plot_pheno_sankey` | function | core | keep | Draw a Sankey of phenotype-distribution flow across ordered covariate values using per-clone outer-product flow geometry. | +| `clonality` | function | plotting-beyond-core | move→examples / drop | Boxplot of the clonality metric via tcri_boxplot. | +| `clone_size_umap` | function | plotting-beyond-core | move→examples / drop | UMAP scatter colored by log10 clone size. | +| `compare_phenotypes` | function | plotting-beyond-core | move→examples / drop | Heatmap of the row-normalized crosstab between two arbitrary obs columns. | +| `flux` | function | plotting-beyond-core | move→examples / drop | Boxplot of per-clone flux (flux_tl distance) between consecutive `order` values, grouped by `groupby` and colored by a `paint` obs category. | +| `plot_phenotype_probabilities` | function | plotting-beyond-core | move→examples / drop | UMAP panels colored by each per-cell phenotype probability. | +| `top_clone_umap` | function | plotting-beyond-core | move→examples / drop | UMAP scatter highlighting the top-N largest clones over a grey background. | +| `SankeyNode` | class | helper | extract → shared | Drawing primitive representing one rectangular sankey node and the ribbons flowing out of it, used to render the phenotypic-flux sankey. | +| `SankeyNode.__init__` | method | helper | extract → shared | Compute node bounding box (min/max x,y from center x, base y, width dx, height val) and build the mpatches.Rectangle patch. | +| `SankeyNode.plot` | method | helper | extract → shared | Render the node's rectangle patch onto the given matplotlib axis. | +| `SankeyNode.plot_node_connection` | method | helper | extract → shared | Draw the curved, color-interpolated ribbon (500 fill_between segments, sigmoid-shaped top/bottom edges) from this node to a destination … | +| `_fin` | function | helper | extract → shared | Print a magenta 'Done!' final flourish to stdout. | +| `_info` | function | helper | extract → shared | Print a cyan key-value info line to stdout. | +| `_ok` | function | helper | extract → shared | Print a green checkmark success line to stdout. | +| `_phenotype_mass_per_clone` | function | helper | extract → shared | Return {clone_id -> phenotype-mass vector} at one covariate by summing joint_distribution rows, optionally weighting each row by the … | +| `_warn` | function | helper | extract → shared | Print a yellow warning line to stdout. | +| `set_color_palette` | function | helper | extract → shared | Assign tcri_colors to the categories of each obs column and store them in uns['_colors']; return category->color map. | +| `tcri_boxplot` | function | helper | extract → shared | Generic per-phenotype metric boxplot/stripplot engine that applies a metric `function` across groupby/splitby strata. | +| `SankeyNode.hex_to_rgb` | method | dead-broken | delete | Parse a #RRGGBB hex string into a normalized (r,g,b) float tuple in [0,1]. | +| `bayesian_mutual_information` | function | dead-broken | delete | 3-panel Bayesian Δ-MI (post-pre) analysis across splitby strata: Δ-MI KDEs, per-condition pre/post KDEs, and Δ bar summary with HDI/P(>0). | +| `compare_joint_distribution` | function | dead-broken | → diagnostics (PPC) | Side-by-side clustered heatmaps/dendrograms of model-inferred vs empirical joint distributions per covariate. | +| `mutual_information` | function | dead-broken | keep + FIX | Box/strip plot of clonotype<->phenotype mutual information per covariate and batch (intended core MI plot). | +| `phenotypic_entropy` | function | dead-broken | keep + FIX | Box/strip plot of phenotypic entropy per covariate and batch (intended core entropy plot). | +| `polar_plot` | function | dead-broken | delete | Radar/polar plot of per-phenotype distribution or entropy across split values. | +| `probability_distribution` | function | dead-broken | delete | Intended: barplots of the phenotype probability distribution per split value. | +| `ridge_delta_entropy` | function | dead-broken | keep + FIX | Ridge/joyplot of Δ-clonotypic-entropy posteriors per phenotype with significance brackets on the first two groups. | + +### `tcri.ut` — utils (19 records) + +| name | kind | label | disposition | purpose | +|---|---|---|---|---| +| `build_nested_tcri_pgm` | function | plotting-beyond-core | → diagnostics | Construct a daft probabilistic-graphical-model diagram of the nested TCRI generative model (plates for batch/clonotype/ct/data, nodes … | +| `draw_tcri_pgm_nested` | function | plotting-beyond-core | → diagnostics | Render build_nested_tcri_pgm and save it to a hardcoded PDF, then plt.show(). | +| `_collect_setup_from_adata_or_model` | function | session-io | keep | Assemble the setup dict (phenotype/clone/covariate/batch cols, category lists, layer) from adata.uns['tcri_metadata']/categories and, if … | +| `_disable_scvi_onload_train` | function | session-io | keep | Context manager that monkey-patches scvi PyroBaseModuleClass.on_load to a no-op during model load, avoiding the one-step warmup train that … | +| `_disable_scvi_onload_train._noop` | inner-function | session-io | keep | Replacement on_load that swallows all args and just clears the Pyro param store (own params loaded afterward). | +| `_ensure_dir` | function | session-io | keep | os.makedirs(path, exist_ok=True) wrapper. | +| `_ensure_pyro_posterior_params` | function | session-io | keep | After load, guarantee the Pyro param 'q_p_ct_raw' exists; if missing, warn and re-init it to a uniform 1/P simplex so posterior metrics … | +| `_pop_nonserializables` | function | session-io | keep | Remove the non-picklable AnnDataManager (uns['tcri_manager']) before writing h5ad, returning a sidecar note. | +| `_pyro_load` | function | session-io | keep | torch.load a Pyro param-store state dict (weights_only=False for constraint objects) and set it into the global param store. | +| `_resolve_TCRIModel` | function | session-io | keep | Dynamically locate and import the TCRIModel class from common module paths or sibling files (editable installs). | +| `_restore_category_order` | function | session-io | keep | Re-impose saved categorical ordering on adata.obs phenotype/clone/covariate columns from the setup dict. | +| `load_tcri_session` | function | session-io | keep | Reconstruct a trained TCRIModel + AnnData from a saved run dir: read h5ad, restore setup/category order, re-run setup_anndata, load model … | +| `save_tcri_session` | function | session-io | keep | Persist a trained session: scvi model (weights+registry, no embedded adata), Pyro param store, setup.json, sanitized adata.h5ad, and … | +| `write_adata_safely` | function | session-io | keep | Write an AnnData to h5ad after stripping the non-serializable tcri_manager (not restored; rebuilt on load). | +| `_to_jsonable` | function | helper | extract → shared | Recursively coerce arbitrary values (numpy/torch/nested) into JSON-serializable primitives. | +| `auc_and_label_permutation` | function | helper | extract → shared | Compute observed ROC-AUC plus a two-sided permutation p-value (exact combinations if feasible, else Monte-Carlo). | +| `bootstrap_auc` | function | helper | extract → shared | Bootstrap 95% CI (2.5/97.5 quantiles) of ROC-AUC, resampling until both classes present. | +| `stars` | function | helper | extract → shared | Map a p-value to a significance-star string (****/***/**/*/ns). | +| `probabilities` | function | dead-broken | delete | Build a per-cell {barcode: {phenotype: prob}} dict from the probability columns and the joint_distribution index. | + +## 5. Consolidation groups (redundant → core) + +| into | members | rationale | +|---|---|---| +| `metrics.clonotypic_entropy` | metrics.clonotypic_entropy_base, metrics.clonality | clonotypic_entropy_base is the single-phenotype/single-draw special case (only extra is posterior=False + weighted) and its sole caller … | +| `preprocessing.joint_distribution_posterior` | preprocessing.joint_distribution | joint_distribution is the point-estimate/ct-level path producing the same covariate->clone x phenotype table WITHOUT per-cell logits. Unify into one engine … | +| `preprocessing.register_model` | preprocessing.classify_phenotypes | classify_phenotypes writes the same phenotype-probability + hard-label slots register_model produces, via a different archetype-cosine algorithm instead of … | +| `preprocessing.group_singletons` | preprocessing.group_small_clones | Both collapse sub-threshold clones into 'Singleton_{group}' labels written to obs['trb_unique']. group_small_clones is dead-broken/dataset-specific (0 … | +| `tcri/_distance.py (kl_divergence)` | metrics.dkl, metrics.flux.dkl_func | Module-level dkl is dead (0 callers) and flux's inner dkl_func reimplements the same KL kernel for the distance_metric=='dkl' branch. Fold both into one … | + +## 6. Deletions + +| function | reason | +|---|---| +| `metrics._ent` | 0 callers; would-be entropy helper never wired in — the entropy metrics inline their own clip/normalise/entropy. | +| `metrics.clone_fraction` | 0 callers; slow reimplementation of obs.groupby(groupby)[clone].value_counts(normalize=True); also uses legacy uns['tcri_clone_key']. | +| `preprocessing.remove_meaningless_genes` | 0 callers AND broken: the include_mtrn/include_hsp branches reassign `genes` from the FULL adata.var.index, silently discarding all prior filters. | +| `preprocessing.get_latent_embedding` | 0 callers; fabricates a spherical Gaussian with a constant posterior_scale (not the model's true latent variance), and default n_samples=0 yields an … | +| `preprocessing.register_probability_columns` | 0 callers; trivial setter whose uns['probability_columns'] key is never read anywhere in the package. | +| `preprocessing.gene_entropy` | 0 callers; per-GENE expression entropy unrelated to core clonotypic/phenotypic entropy; also drops the first count bin on an assumption. | +| `plotting.compare_joint_distribution` | **OVERLAY OVERRIDE → keep as diagnostics PPC** (was: dead-broken: references an undefined global `model` (model.adata_manager.registry) -> NameError at runtime.) | +| `plotting.probability_distribution` | dead-broken: infinite self-recursion (calls probability_distribution instead of utils.probabilities); also ax indexing fails when splitby=None. | +| `plotting.bayesian_mutual_information` | dead-broken (passes unsupported weighted= to tl.mutual_information -> TypeError) AND redundant with the functional core pl.mi_compare MI-comparison … | +| `plotting.polar_plot` | dead-broken: phenotypes defaults to a column-NAME string (iterated char-by-char) and the entropy branch calls unimported clonotypic_entropy -> … | +| `plotting._sankey.SankeyNode.hex_to_rgb` | 0 callers; plot_node_connection uses matplotlib mcolors.to_rgb instead; contract doc already marks it a deleted/unused method. | +| `utils.probabilities` | dead-broken: reads uns['joint_distribution'] which is never written anywhere in the package -> KeyError. | + +## 7. Helper extraction (dedupe → shared modules) + +| helper | → module | current copies | +|---|---|---| +| _ok, _info, _warn, _fin (ANSI console printers) | `tcri/_console.py` | metrics/_metrics.py:47-56; preprocessing/_preprocessing.py:50-59; … | +| _ascii_hist (numpy-histogram ASCII bar builder) | `tcri/_console.py` | metrics/_metrics.py:62 (the used copy); preprocessing/_preprocessing.py:62 (identical, … | +| stars, auc_and_label_permutation, bootstrap_auc (significance + AUROC stats) | `tcri/_stats.py` | utils/_utils.py (single copies today; relocate out of the utils monolith into a named … | +| kl_divergence + l1_distance + phenotype_distance dispatcher | `tcri/_distance.py` | metrics/_metrics.py:159 (dead module-level dkl) and metrics/_metrics.py flux.dkl_func … | +| _mi_from_joint (single-source MI kernel) | `tcri/metrics/_mutual_information.py (module-private)` | metrics/_metrics.py:_mi_from_joint (single copy; keep private, called by … | +| tcri_boxplot (generic per-phenotype box/strip engine) + _finish (scanpy show/save) | `tcri/plotting/_base.py` | plotting/_plotting.py:598 tcri_boxplot; _finish is new (formalize the fig/ax show/save … | +| SankeyNode drawing primitive + _phenotype_mass_per_clone data-prep | `tcri/plotting/_sankey.py` | plotting/_sankey.py (SankeyNode already; drop hex_to_rgb); … | +| set_color_palette + tcri_colors palette constants | `tcri/plotting/_colors.py` | plotting/_plotting.py set_color_palette (also fix: it writes uns on adata.copy() so … | +| register_phenotype_key, register_clonotype_key, _compute_logits_and_prior | `tcri/preprocessing/_register.py` | preprocessing/_preprocessing.py (relocate alongside register_model; migrate the legacy … | +| group_singletons (+ collapse_singleton inner), clone_size | `tcri/preprocessing/_clones.py` | preprocessing/_preprocessing.py:83 group_singletons; preprocessing/_preprocessing.py … | +| _to_jsonable (recursive JSON coercion of numpy/torch/nested) | `tcri/utils/_session.py` | utils/_utils.py:_to_jsonable (0 external callers today; wire it into save_tcri_session's … | + +## 8. Plotting triage + +**Core (keep):** `clonotypic_entropy_by_phenotype`, `mi_compare`, `phenotypic_flux`, `plot_pheno_sankey`, `mutual_information`, `phenotypic_entropy`, `ridge_delta_entropy` + +**Beyond core (move→examples / drop):** + +- compare_phenotypes (move->examples: generic categorical-crosstab heatmap, visualizes no core metric) +- top_clone_umap (move->examples: bespoke top-N-clone UMAP overlay, hardcoded title) +- clone_size_umap (move->examples: bespoke clone-size UMAP overlay; mutates adata.obs as a side effect) +- plot_phenotype_probabilities (move->examples: per-cell phenotype-probability UMAP panels) +- clonality (drop: only plots the merged-away redundant clonality metric) +- flux boxplot (drop: broken on default paint=None path — pcat used before assignment -> NameError — plus dataset-specific `paint` overlay; the Sankey is the flux visualization) +- plot_archetypes (drop or keep as optional TCRIModel model-diagnostic method: model prior heatmap, not a core-metric plot) +- plot_loss (drop or keep as optional TCRIModel model-diagnostic method: training ELBO/KL curves) +- build_nested_tcri_pgm (move->docs/examples: daft PGM diagram of the model architecture) +- draw_tcri_pgm_nested (drop: one-off PGM export with a hardcoded output PDF filename) + +## 9. Rename / readability map + +Freeze this **before** touching code or notebooks — renames are breaking and we only pay once (pre-1.0, Alpha). + +### Modules / files +| current | → target | +|---|---| +| `metrics/_metrics.py` (1008 ln) | `tl/_entropy.py` + `_mutual_information.py` + `_flux.py` + `_tables.py` | +| `preprocessing/_preprocessing.py` (559) | `pp/_register.py` + `_engine.py` + `_clones.py` | +| `plotting/_plotting.py` (1437) | `pl/_entropy.py` + `_mutual_information.py` + `_flux.py` + `_base.py` + `_colors.py` | +| `model/_model.py` (1074) | `model/_model.py` + `_module.py` + `_priors.py` + `_classifier.py` + `_training.py` | +| `utils/_utils.py` (665) | `utils/_session.py` + new `_console.py` / `_stats.py` / `_distance.py` / `_keys.py` + `diagnostics/` | + +### Functions +| current | → target | why | +|---|---|---| +| `joint_distribution` + `joint_distribution_posterior` | `joint_distribution(posterior=, n_samples=)` | one engine, one point/draws knob | +| `clonotypic_entropy_base` | merged into `clonotypic_entropy` | `_base` means nothing; it's the single-phenotype case | +| `pl.clonotypic_entropy_by_phenotype` | `pl.clonotypic_entropy` | tl↔pl twin name; matches notebook expectation | +| `plot_phenotype_probabilities` / `plot_pheno_sankey` | drop `plot_` prefix (`phenotype_probabilities`, internal `_sankey`) | scanpy/scvi pl fns are unprefixed | +| `tcri_boxplot` | `_metric_boxplot` (private helper) | generic engine, not public API | +| `dkl` / `flux.dkl_func` | `_distance.kl_divergence` | dedupe the KL kernel | +| `centropy` / `pentropy` / `*_tl` aliases | **removed** (via `__all__`) | leaked import aliases | +| `classify_phenotypes` | **removed** → `register_model` | duplicate phenotype-assignment path | +| `remove_meaningless_genes` | `filter_genes` (if kept) *or* delete | broken flag logic + 0 callers | + +### Parameters +| current | → target | +|---|---| +| `from_this` / `to_that` (flux) | `cov_from` / `cov_to` | +| `point_estimate=True` (entropies) | **removed** — use `n_samples=0` | +| ad-hoc arg orders | standardize `covariate` / `splitby` / `groupby` / `clones` / `n_samples` / `temperature` / `posterior` everywhere | + +### State keys +| current | → target | +|---|---| +| `uns["tcri_clone_key"]` / `["tcri_phenotype_key"]` **and** `uns["tcri_metadata"][...]` (two conventions) | one scheme via `_keys.py` constants (single `tcri_metadata`) | +| ad-hoc registry keys (`"clonotype_col_in_registry"`) | standard scvi `REGISTRY_KEYS` | + +### Internal variables +| current | → target | +|---|---| +| `Δ` (unicode) in `delta_clonotypic_entropy` | `delta` (ASCII, greppable) | +| `c2p_mat` | `clone_phenotype_prior` | +| `p_ct` / `ct_to_c` / `ct_to_cov` (terse) | keep, but document (`ct` = (clone, covariate) index) | + +## 10. Metric conventions & scoping + +### Uniform point-estimate / sampling convention (applies to every metric) +- **`n_samples=0` → deterministic point estimate.** Posterior-*mean* `p_ct` (× logits), softmax, **no draw**, reproducible. +- **`n_samples=N>0` → N posterior draws** (adds a sampling axis; mean ± CI fall out). +- **Delete the `point_estimate=` argument** — `n_samples` is the only knob. This also fixes a latent bug: today `mutual_information(n_samples=0)` / `flux(n_samples=0)` return **one random draw**, not a deterministic estimate. + +### Prior-vs-posterior — PARKED (open, do not collapse in this pass) +The `{prior, posterior} × {point, draws}` 2×2 is deferred. Until resolved the plan assumes `posterior=True` with +`n_samples` as the point/draws knob, and keeps `posterior=` as a documented-but-unfinalized argument +(the current prior-only branch raises `NotImplementedError`). + +### diagnostics/ = PPCs + model validation +`gf.diag`-style, returns DataFrames, read-only on the finalized model. Seeded by: +- **joint-distribution PPC** — model p(clone,phenotype) vs empirical counts (the *fixed* `compare_joint_distribution`) +- phenotype-probability **calibration**; **reconstruction PPC** (ZINB simulate → compare library/dropout/mean-var); entropy/MI vs permutation null +- relocated: training curves (`plot_loss`), archetypes (`plot_archetypes`), model PGM (`build_nested_tcri_pgm`) + +## 11. Open questions (decide before executing) + +_(see synthesis output)_ + +--- +_Overlays applied on top of the workflow synthesis: (a) `diagnostics/` = PPCs + model validation (`compare_joint_distribution` reclassified from delete → diagnostics PPC seed; `plot_loss`/`plot_archetypes`/PGM relocated); (b) uniform `n_samples=0` point-estimate convention, drop `point_estimate=`; (c) prior-vs-posterior parked; (d) explicit `__all__` re-export (NOT `import *`) per grafiti; (e) full rename map._ \ No newline at end of file diff --git a/docs/contract/tcri_implementation_plan.md b/docs/contract/tcri_implementation_plan.md new file mode 100644 index 0000000..cfc7747 --- /dev/null +++ b/docs/contract/tcri_implementation_plan.md @@ -0,0 +1,463 @@ +# TCRI Refactor — Implementation Plan (Final) + +**From** the current five-monolith package **to** the settled Door‑A, scverse‑ecosystem target (grafiti‑mirrored, one‑file‑per‑topic), with every audit fix folded in. + +**Source of truth.** The *SETTLED DESIGN* block is authoritative; the *prior‑vs‑mean RESOLUTION* settles the one parked statistical question and is adopted verbatim. Where the settled design diverges from `docs/contract/tcri_function_inventory.md` §3/§9 (engine lives in `tl` not `pp`; `groupby` subsumes all `*_table`/`*_compare`/`*_delta`; `flux → phenotypic_flux`; `register_model → model.to_anndata`; `tl` mirrors grafiti's `tools/`), this plan follows the settled design and flags the divergence inline. This document supersedes the earlier draft wherever the audit corrected it; the corrections are integrated at the point they bite, and every finding is cross‑referenced in **Appendix A** so nothing is lost. + +**How to read.** §1 fixes the invariants and records the resolved decisions that unblock the engine/metric phases. §2 is the target tree. §3–4 are the rename/disposition map and the shared‑helper extraction. §5 is the model→AnnData streamline. §6 is the engine + metric **numeric contract** (all math/stats fixes live here). §7 is the GPU/optimization architecture, grounded in grafiti. §8 is the ordered, independently‑landable PR sequence. §9 is testing + scverse‑CI. §10 is the ordering‑hazard graph. §11 is risks & open items. §12 is the per‑PR checklist. + +--- + +## 1. Invariants and resolved decisions + +### 1.1 Invariants held at every PR boundary +- **`import tcri` stays green** and the public handles `tl / pp / pl / ml / diag / ut` remain importable. No PR leaves `main` with a broken import or red CI. **Corollary:** "0‑caller" means *no call‑sites **and** no import‑sites in the package* — a symbol imported at module top (e.g. `utils.probabilities` at `_plotting.py:18`) is a live dependency even with zero calls, so it and its import are removed in the same PR. (The `example/` notebooks are disposable and never counted.) +- **One behavior change per PR.** Mechanical moves (splits, helper extraction) never change numerics; numeric changes (engine, metrics) never also move files. Every diff stays reviewable and every regression bisectable. +- **`_keys.py` is the only place a key string is written**, from Phase 1 on. No new `uns/obsm/obs` key literal may appear outside it. +- **No `import *`.** Each split lands with an explicit `__all__` and named re‑exports; the top‑level `__init__` flip is the last PR. +- **`setup_anndata` performs no *analysis/label* obs mutation.** *(Corrected from the draft's "no obs mutation.")* It must still write and register the `obs['indices']` field the training step consumes (`batch['indices']`, `_model.py:603/641/678/684`). The invariant forbids writing *results* (probabilities, hard labels, latent) into `obs`, not the registration glue column. +- **GPU is never on the import path and never in `install_requires`.** Every accelerated path is pure opt‑in behind a device seam with a fully‑functional CPU fallback (§7). + +### 1.2 Resolved decisions (adopted from the prior‑vs‑mean RESOLUTION + math/stats audit) + +**(R1) Point estimate = closed‑form variational posterior mean.** `n_samples=0` returns `E_q[p_ct] = normalize(q_p_ct_raw) = TCRIModel.get_p_ct() = adata.uns[K.P_CT]`, read directly, **never sampled**. Because the guide is `Dirichlet(local_scale · m)` with `m` on the simplex, `Σα = local_scale` and the mean is exactly `m` — `local_scale` cancels. The **prior/archetype path is dropped** (it is guide‑init/generative anchor built from leaked hard labels, not what training learned); the **MAP/mode is rejected** (`α_k = local_scale·m_k` is routinely `< 1`, so the mode sits on the simplex boundary / is undefined); **mean‑of‑draws is rejected** (Rao‑Blackwell‑dominated, adds only MC noise to a closed‑form quantity). + +**(R2) Rename the mis‑named `posterior=` axis to `use_logits` — a classifier‑mixing switch, not a prior/posterior switch.** There was never a live prior branch; what actually differs is *with logits vs without*. The **engine** `joint_distribution` keeps exactly one such flag, `use_logits` (alias `cell_informed`), replacing both `posterior=` and `combine_with_logits=`: + - `use_logits=True` (default) folds per‑cell classifier logits into `log(base)` exactly like `model.predict` (gate‑aware, §5). + - `use_logits=False` returns the ct‑level table directly. + Both branches use the posterior **mean** (`n_samples=0`) or **draws** (`n_samples>0`) of `p_ct`; neither ever touches the generative prior. **The four metrics do NOT expose `use_logits`, `posterior=`, or `point_estimate=`** — given an `adata` they always compute the cell‑informed joint. `point_estimate=` is deleted. + +**(R3) `n_samples=0` is a *plug‑in* estimator, not the posterior mean of the metric.** Entropy, MI, and l1/KL flux are **nonlinear** functionals, so `metric(E_q[p]) ≠ E_q[metric(p)]` by a **Jensen gap that is not Monte‑Carlo noise**: entropy (concave) plug‑in **over‑estimates**, flux (convex) plug‑in **under‑estimates**, MI is sign‑indeterminate. Therefore: + - Document `n_samples=0` as **`metric‑at‑posterior‑mean`** and the `n_samples>0` `mean` column as **`E_q[metric]`** — two different, clearly‑labeled estimators. + - **No test may assert `n_samples=0 == mean(n_samples>0)`.** The determinism/reproducibility guarantee (bit‑identical repeated `n_samples=0` calls) still holds and is tested; the *equality across estimators* claim is dropped. + +**(R4) `n_samples>0` draws must reproduce the guide's clamped concentration.** Draw from `Dirichlet(clamp(local_scale · m̃, min=1e-3))`, reusing the guide's floor (`_model.py:490`). The three inconsistent current variants (`local_scale·m` unclamped; `local_scale·m + 1e-8`; the clamped guide form) are unified onto the guide form so reported HDIs summarize the distribution the model actually learned. Note that for committed clones (`m_k < 1e-3/local_scale`) the clamp makes the draw‑mean differ slightly from `m`; this is documented, and `n_samples=0` remains the closed‑form `m` by definition. + +**(R5) `p_gt` (P(>0)) is only meaningful on a signed contrast.** Entropy, MI, and flux are all `≥ 0`, so `P(draw > 0) ≈ 1` and is vacuous. **Remove `p_gt` from the single‑metric `n_samples>0` summary** (emit `mean`, `sd`, `hdi_low`, `hdi_high` only). `p_gt` is computed **only** on a paired between‑group difference vector by the comparison helper (§4, §6). + +**(R6) Temperature is a single analysis knob applied one consistent way.** `m̃ = softmax(log(m + 1e-8) / T)` (identity at `T=1`), applied identically in the mean and draw branches. To avoid double‑tempering, **`to_anndata` persists the *raw* posterior mean** `m = normalize(q_p_ct_raw)` (not the guide‑temperature‑adjusted vector), so analysis `temperature` is the sole tempering knob. At `T=1`, `use_logits=True` reproduces `model.predict` exactly. `temperature≠1` re‑centers the sampled distribution away from the learned posterior — documented in the metric docstrings. + +**(R7) Reproducibility is via a seeded torch generator.** All draws move to a seeded `torch.Generator` (device‑aware; also seed cupy/torch.cuda RNG on GPU). Add `random_state: int | torch.Generator | None` to `joint_distribution` and every sampling metric/diag function; thread it into the Dirichlet draw. `np.random.seed` alone was a **no‑op** for the torch draws and is retired. `n_samples=0` is deterministic regardless. + +**(R8) `use_logits=True` parity with `predict` requires three model facts persisted to `adata`.** The gate‑aware combination `gate_prob·logits + (1-gate_prob)·log(base)` and the classifier temperature are **model attributes**. `to_anndata` must persist `obsm[K.X_LOGITS]`, `uns[K.GATE_PROB]`, and `uns[K.CLASSIFIER_TEMPERATURE]` (§5). Without them the engine can only reproduce the additive rule — the exact disagreement the refactor set out to remove. + +--- + +## 2. Target layout (grafiti‑mirrored, authoritative) + +``` +tcri/ + __init__.py # explicit re-export; sys.modules aliases tl/pp/pl/ml/diag/ut; + # top-level tcri.joint_distribution; NO import * + _keys.py # single source of every uns/obsm/obs key string (adopt in Phase 1) + # NOW INCLUDES: X_LOGITS, GATE_PROB, CLASSIFIER_TEMPERATURE, + # LOCAL_SCALE, P_CT, X_TCRI, X_PROBABILITIES, CLONE_COL/PHENO_COL + _console.py # leveled, silenceable logging over scanpy logging; NO ANSI, NO _ascii_hist + _stats.py # stars, auc_and_label_permutation, bootstrap_auc, + # + posterior-comparison primitives: mann_whitney, prob_gt_zero, hdi (TRUE HDI) + _distance.py # kl_divergence (log2/bits), l1_distance, jensen_shannon, phenotype_distance dispatcher + _compute/ # NEW — device seam + batched numeric core (grafiti/_compute parity) + _xp.py # resolve_device / get_xp / asnumpy (torch-first, cupy optional, cpu default) + _joint.py # _joint_draws(adata, covariate, n_samples, *, use_logits, temperature, device, + # random_state) -> ndarray[n_samples, n_clones, P] (scatter-add reduction) + _reduce.py # batched entropy / mutual_information / flux over the [S, n_clones, P] stack + _embedding.py # umap() behind _use_gpu gate (cuML on GPU, umap-learn CPU); lazy import + model/ # ml + _model.py # TCRIModel: setup_anndata, train, get_latent_representation, + # predict (was get_cell_phenotype_probs), get_p_ct, to_anndata + _module.py # TCRIModule (pyro model/guide) + _priors.py # MixtureDirichlet, VampPrior + _classifier.py # PhenotypeClassifier + _training.py # UnifiedTrainingPlan, build_archetypes (returns centers AND labels) + preprocessing/ # pp (shrinks to clone utilities; engine moved to tools/) + _clones.py # group_singletons (must precede setup_anndata), clone_size + tools/ # tl (metrics + engine; mirrors grafiti tools/) + _joint.py # joint_distribution (thin DataFrame wrapper over _compute._joint; + # re-exported as tcri.joint_distribution) + _entropy.py # clonotypic_entropy, phenotypic_entropy + _mutual_information.py # mutual_information (+ private _mi_from_joint) + _flux.py # phenotypic_flux (was flux; cov_from / cov_to) + _compare.py # compare_groups (PUBLIC mid-level stats helper; replaces *_compare/*_delta) + plotting/ # pl (twins mirror tl by filename) + _base.py # _metric_boxplot (was tcri_boxplot; groupby + splitby), _finish + _colors.py # tcri_colors, resolve_palette (was set_color_palette) + _entropy.py # clonotypic_entropy [FIX], phenotypic_entropy [FIX] + _mutual_information.py # mutual_information [FIX] + _flux.py # phenotypic_flux (sankey) + _sankey.py # SankeyNode, _phenotype_mass_per_clone + diagnostics/ # diag (NEW — PPCs + model validation; returns DataFrames) + _ppc.py # joint-distribution PPC (fixed compare_joint_distribution), + # calibration, reconstruction PPC (model-required), permutation-null + _training.py # loss curves (was plot_loss), archetypes (was plot_archetypes) + utils/ # ut + _session.py # save/load_tcri_session, _to_jsonable (plain h5ad; no manager hack) +docs/ # model PGM (build_nested_tcri_pgm) lives here, out of the package; + # the rewritten tutorial notebook lives here too +``` + +Divergences from inventory §3, called out: **(a)** `tl` package is `tools/` (grafiti parity), not `metrics/`; **(b)** the engine's *numeric core* is `_compute/_joint.py` with a thin `tools/_joint.py` DataFrame wrapper — **not** `preprocessing/_engine.py`; **(c)** `preprocessing/` loses `_register.py`/`_engine.py` (registration collapses into `model.to_anndata`), leaving `pp` = `_clones.py`; **(d)** no `_tables.py` — `groupby` + `tl.compare_groups` subsume it; **(e)** a **new `_compute/` package** (not in the original inventory) is the device seam that makes the GPU wins additive. + +--- + +## 3. Rename & disposition map + +Freeze the map **before** Phase 5 (the first breaking PR). Renames are breaking; pre‑1.0 we pay once. The deltas this plan enforces — **including the four settled‑design corrections the draft had wrong** (marked ⚠): + +| current | → target | lands in | +|---|---|---| +| `metrics/` package | `tools/` package (aliased `tl`) | Phase 5 | +| `joint_distribution` + `joint_distribution_posterior` | `tl.joint_distribution` (engine) → top‑level `tcri.joint_distribution` | Phase 5 | +| `posterior=` / `combine_with_logits=` (engine) | **`use_logits=`** (alias `cell_informed`), default `True` (R2) | Phase 5 | +| `point_estimate=`; public `posterior=` on metrics | **removed** — `n_samples` is the only point/draws knob (R1–R2) | Phase 5–6 | +| `flux` / `from_this` / `to_that` | `phenotypic_flux` / `cov_from` / `cov_to` | Phase 6 | +| ⚠ `clonality` | **DROP entirely — do NOT merge into `clonotypic_entropy`** (generic repertoire stat) | Phase 6 | +| ⚠ `clonotypic_entropy_base` | **DROP** (not merged); log base standardized via `_distance` (bits) | Phase 6 | +| ⚠ `ridge_delta_entropy` | **DROP** (not keep‑and‑fix) | Phase 7 | +| ⚠ `compare_phenotypes` | **DROP** | Phase 6 | +| `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_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 | +| `plot_phenotype_probabilities` | **DROP** (not core) | Phase 6 | +| `probability_ternary` | **DROP** (not core — a bespoke simplex plot of no core metric) | Phase 6 | +| `gene_entropy` | **DROP** (out of scope — generic gene QC) | Phase 2 | +| `polar_plot` | **DROP** | Phase 2 | +| `pl.flux` boxplot | **DROP** (sankey is the flux plot) | Phase 7 | +| `tcri_boxplot` | `_metric_boxplot` (private; keeps a **`splitby=` axis**, §Phase 7) | Phase 7 | +| `set_color_palette` | `resolve_palette` | Phase 7 | +| `dkl`, `flux.dkl_func` | `_distance.kl_divergence` (log2/bits, single eps) + `_distance.jensen_shannon` | Phase 1 | +| `Δ` (unicode) | `delta` (ASCII, greppable) | Phase 6 | +| `c2p_mat` | `clone_phenotype_prior` | Phase 3 | +| `centropy` / `pentropy` / `*_tl` leaked aliases | removed via `__all__` | Phase 11 | +| `uns["tcri_clone_key"]`/`["tcri_phenotype_key"]` **and** `uns["tcri_metadata"][...]` (two schemes) | one scheme via `_keys.py` (single `tcri_metadata`) | Phase 1 | + +--- + +## 4. Shared‑helper extraction (`_keys` / `_console` / `_stats` / `_distance`) + +Lands as **Phase 1**, before any file move, because every later phase imports these. Pure internal dedup; public API unchanged except for the new public `tl.compare_groups` (Phase 6, built on `_stats`). + +| new module | absorbs | notes / fixes folded in | +|---|---|---| +| `tcri/_keys.py` | every `uns/obsm/obs` key literal, both current schemes | constants only; unifies `tcri_clone_key`/`tcri_metadata['clone_col']` to one `tcri_metadata` scheme; **adds `X_LOGITS`, `GATE_PROB`, `CLASSIFIER_TEMPERATURE`, `LOCAL_SCALE`** (R4/R6/R8); `clone_size` and any other reader flips to `K.CLONE_COL` in the **same** change that retires the writer (no orphaned reader) | +| `tcri/_console.py` | triplicated `_ok/_info/_warn/_fin` (metrics/pp/pl copies) | reimplement over **scanpy's `logging`/verbosity**; drop raw ANSI; **drop `_ascii_hist` and every `graph=`/ASCII‑histogram code path** | +| `tcri/_stats.py` | `stars`, `auc_and_label_permutation`, `bootstrap_auc` (from `utils/_utils.py`) | **add** `mann_whitney`, `prob_gt_zero` (Bayesian P(>0), for signed contrasts only — R5), **`hdi` implemented as a TRUE highest‑density interval** (not the mislabeled equal‑tailed `percentile[2.5,97.5]`); document boundary instability for bounded skewed posteriors | +| `tcri/_distance.py` | module‑level dead `dkl` + `flux.dkl_func` | one `kl_divergence` + `l1_distance` + **`jensen_shannon`** + `phenotype_distance(metric=...)`; **fix the mixed‑units bug — standardize on `log2` (bits) across entropy/MI/KL with one eps**; document `dkl` as directional/unbounded and recommend JSD (bounded ≤1 bit) for symmetric shift; `l1` stays the safe bounded default | + +`_mi_from_joint` stays module‑private in `tools/_mutual_information.py`. The **public comparison surface** is `tl.compare_groups` (§6) — not the private `_stats` primitives — so "comparisons via `groupby` + stats" is programmatically reproducible. + +--- + +## 5. Model → AnnData streamline + +The single highest‑risk behavior change, and the one that makes the metric↔model agreement guarantee (R8) achievable. It kills the `AnnDataManager`‑in‑`uns` hack and fixes the write‑set the draft under‑specified. + +**5.1 `setup_anndata` — registration only (no *analysis* obs mutation).** +- Registers fields via scvi `REGISTRY_KEYS`; **keeps writing/registering `obs['indices']`** (training glue, `_model.py:678/684`) — this is not an analysis mutation and must not be removed. +- **Removes the manager stash `adata.uns['tcri_manager'] = adata_manager` (`_model.py:697`)** — *this is where the stash actually lives, not in `register_model`.* Removing it here is what lets `write_adata_safely`/`_pop_nonserializables` be deleted. +- `group_singletons` stays a **separate `pp` step that must run BEFORE `setup_anndata`** (it relabels clones; running it after desyncs `ct_to_c`/`p_ct` from `obs`). Enforced: `setup_anndata` errors if a later relabel is detected. + +**5.2 `model.to_anndata(adata)` — thin, canonical write‑set.** Writes **only** the canonical minimum via `_keys`, and the canonical minimum now **explicitly includes the three items the engine needs for `predict` parity** (correcting the draft's "nothing else"): + +| slot | key | why it is canonical | +|---|---|---| +| metadata/categories | `K.META`, covariate/phenotype/ct category maps | registry provenance | +| latent | `obsm[K.X_TCRI]` | embedding | +| phenotype probs + hard labels | `obsm[K.X_PROBABILITIES]`, `obs[...]` | `predict()` output; standard slot (retire `X_tcri_phenotypes`) | +| ct‑level prior mean | `uns[K.P_CT]` = **raw** `normalize(q_p_ct_raw)` (R6) | `n_samples=0` closed‑form mean | +| **per‑cell logits** | **`obsm[K.X_LOGITS]`** | `use_logits=True` engine path (R8) — hard‑required | +| **gate probability** | **`uns[K.GATE_PROB]`** (scalar or `None`) | gate‑aware combine parity with `predict` (R8) | +| **classifier temperature** | **`uns[K.CLASSIFIER_TEMPERATURE]`** | matches `predict`'s logit scaling (R8) | +| **local scale** | **`uns[K.LOCAL_SCALE]`** | draw variance for `n_samples>0`; engine **raises** (never defaults to 1.0) if missing when `n_samples>0` (R4/R8) | + +- **Stops writing `uns['tcri_manager']`** (already removed in 5.1) → deletes the `write_adata_safely`/`_pop_nonserializables` hack. +- `predict` (renamed from `get_cell_phenotype_probs`): returns a per‑cell phenotype‑prob `DataFrame`; **asserts the inference `DataLoader` is order‑preserving (`shuffle=False`) and indexes by `adata.obs_names`** (or carries the registered `indices` field and reindexes) so ct assignment and barcode labels cannot drift. + +**5.3 Session IO.** `write_adata_safely → save_tcri_session` writes a **plain h5ad** (nothing non‑picklable in `uns` anymore); `load_tcri_session` rebuilds the registry by re‑running `setup_anndata`. `get_p_ct` reads the **process‑global** pyro param store (`q_p_ct_raw`); load must set the store immediately before any `get_p_ct`/`to_anndata` call, and multi‑model/round‑trip diag workflows are documented as single‑model‑per‑process unless params are namespaced (§11). + +--- + +## 6. Engine + metric numeric contract + +This section is the substrate all metrics read; it folds in every math/stats and missing‑link correction. + +**6.1 Engine `tools/_joint.py::joint_distribution`** (thin DataFrame wrapper over `_compute/_joint.py::_joint_draws`, §7): + +``` +joint_distribution(adata, *, covariate=None, groupby=None, n_samples=0, + use_logits=True, clones=None, temperature=1.0, + random_state=None) -> pandas.DataFrame +``` +- Unifies `joint_distribution` + `joint_distribution_posterior`; re‑exported as `tcri.joint_distribution`. Provenance in a **serializable form** (a `params` column or a companion `uns` sidecar), **not** only `df.attrs` (R‑forward: h5ad round‑trips must not silently drop it, §11). +- `n_samples=0`: closed‑form path. `use_logits=False` → tempered `m̃` rows (`== uns[K.P_CT]` at `T=1`); `use_logits=True` (default) → per‑cell `softmax((logits + gate‑combine(log m̃))/T)` aggregated per clone, **identical to `model.predict`** (R6/R8). +- `n_samples>0`: draw `p_ct ~ Dirichlet(clamp(local_scale·m̃, 1e-3))` **once per sample via a seeded torch generator** (R4/R7), feed each draw through the same temperature + combine + scatter‑sum; stack a `sample_id` axis. **All clones within one draw share the SAME `p_ct` draw** (one coherent joint per `sample_id`) — never independent per‑clone draws. +- **`covariate=None` computes the joint across ALL covariate values in one pass from a single `p_ct` draw** (the all‑timepoints path the sankey and multi‑covariate metrics need). +- **Draw‑once invariant:** for `n_samples>0`, the number of Dirichlet draws is exactly `n_samples`, **independent of `#groups` and `#covariates`** — draws are reused across groups/covariates by cell‑masking, not re‑drawn. Enforced by a draw‑counter test (§9). +- **`groupby` is implemented by cell/clone RESTRICTION over the FULL `adata`** (positional masks into full‑space `uns` arrays + `clones=`), **never by passing a sliced `AnnData` to the engine** — this avoids the hard full‑space‑vs‑subset alignment guard that today's `tcri_boxplot` slicing would trip. `_metric_boxplot` is rewritten off the slice‑and‑call pattern. +- **`groupby` requires the cell‑informed path** (or group keys that are clone‑nested / constant within `clone × covariate`); it is **ill‑defined on the ct‑level table** for non‑clone‑determined columns. The design assumes **clones are disjoint across `groupby` groups** (a TCR clone does not span two patients) — now stated explicitly, with a validation that errors when a group split would bisect a clone's cells. +- **Engine bug fixes folded in:** weighting keyed on the **`ct` index** (not the clone index); consistent normalization (row‑ vs whole‑table) across the two old functions; all‑zero‑clone reindex yields **NaN, not inflated uniform entropy**; **torch‑seeded** determinism. + +**6.2 Four metrics** (`tools/_entropy.py`, `_mutual_information.py`, `_flux.py`), uniform signatures; **none expose `use_logits`/`posterior=`/`point_estimate=`**: +``` +tl.clonotypic_entropy(adata_or_jd, *, covariate=None, groupby=None, n_samples=0, temperature=1.0, + clones=None, normalized=True, random_state=None) +tl.phenotypic_entropy (adata_or_jd, *, covariate=None, groupby=None, n_samples=0, temperature=1.0, + clones=None, normalized=True, random_state=None) +tl.mutual_information (adata_or_jd, *, covariate=None, groupby=None, n_samples=0, temperature=1.0, + clones=None, normalized=True, normalize_mode='min', random_state=None) +tl.phenotypic_flux (adata, *, cov_from, cov_to, groupby=None, n_samples=0, temperature=1.0, + clones=None, distance_metric='l1', random_state=None) +``` +- **`mutual_information` default `normalize_mode='min'` (was `'average'`).** Under the settled uniform‑clone prior, `P(c)=1/C` pins the clone marginal entropy `H_c` to `log2(C)` (structural, uninformative), which throttles `'average'` normalization by `~1/log2(C)` and makes it non‑comparable across groups with different clone counts — breaking the whole `groupby` workflow. `'min'` (`I/H_p`, coefficient of constraint) is in `[0,1]`, reaches 1 when clone determines phenotype, and is `C`‑independent. `'average'` is documented as not‑recommended (or dropped). +- **`clonotypic_entropy` normalizer comparability:** divide by `log2` of the number of clones with **genuine support**, dropping requested‑but‑absent (all‑zero, reindexed) clones from `C`; expose a **common‑denominator option** (`log2` of a fixed reference clone count) for cross‑group plots. Document that normalized clonotypic entropy uses group‑specific denominators by default. +- **Dual input.** Each accepts `(adata + covariate)` [compute joint internally] **or** a precomputed joint `DataFrame` [fast path]. **A bare precomputed joint forces `n_samples=0` and `groupby=None`** (it carries no `p_ct`/`local_scale`/logits/cells to resample or re‑partition); `clones=` just re‑filters rows; **raise a clear error** for `n_samples>0`/`groupby` on a bare jd. `phenotypic_flux` correctly takes `adata` only (it needs two joints). Metrics **propagate the input joint's provenance** into their output. +- **Return‑shape rule (uniform):** no `groupby` & `n_samples=0` → scalar/Series; `groupby` set → tidy `DataFrame` (row per group [× phenotype/clone]); `n_samples>0` → adds a `sample_id` axis and, on reduction, summary columns **`mean`, `sd`, `hdi_low`, `hdi_high`** (**no `p_gt`** — R5). Metrics reduce the stack by **iterating the `sample_id` level** (per‑draw full‑joint metric, then summarize). +- **`n_samples>0` intervals are partial posterior** (only `p_ct` uncertainty; classifier logits held at their posterior‑mean encoding) — documented so users don't read them as full posterior‑predictive uncertainty. +- **h5ad‑serializable returns (build‑toward `@tl_result`):** flat columns, **no object‑array columns and no `df.attrs`‑only provenance**; per‑draw vectors go in a separate long frame or a `uns` sidecar. Define the cache key as a hash of `(covariate, groupby, n_samples, temperature, clones, normalized, normalize_mode, distance_metric, random_state)`. + +**6.3 Comparisons via `groupby` + `tl.compare_groups` (public).** The deleted `*_compare`/`*_delta` functions are replaced by a **public mid‑level helper**, not private primitives: +``` +tl.compare_groups(df, *, value, by, reference=None, paired=False, hdi_prob=0.94) + -> tidy DataFrame with per-pair: mean_a, mean_b, delta, U, p (Mann-Whitney), + p_gt (Bayesian P(delta>0)), hdi_low, hdi_high +``` +This is where `p_gt`/HDI live (on the **signed** `delta`, R5). The docstring shows the recipe that recreates `mi_compare`'s per‑pair output, so "groupby + stats subsumes the tables/deltas" is actually reproducible. + +--- + +## 7. GPU / optimization architecture + +Grounded 1:1 in grafiti's `_compute/` wins. The **one architecturally load‑bearing decision**: write the engine's numeric core as a **batched, device‑routable function returning a `[n_samples, n_clones, P]` array**, with pandas only at the boundary — so every win below is designed *in*, not retrofitted later. + +**7.1 Device seam — `tcri/_compute/_xp.py` (copy grafiti's 58‑line reference).** `resolve_device` / `get_xp` / `asnumpy`. Because **torch≥2.4.1 is already a hard dep**, a **torch‑tensor core is the first backend** (torch.cuda when present — zero new deps), with **cupy optional** as a numpy‑style second backend. Every accelerated function returns a plain numpy array via `asnumpy`. Device ladder: `None/'cpu'→cpu`; `'mps'→cpu`; `'auto'/'gpu'/'cuda'→GPU only if the lib imports **and** `getDeviceCount()>0`, else CPU (explicit `'cuda'` warns on fallback, `'auto'/'gpu'` silent). + +**7.2 Batched engine core — `_compute/_joint.py::_joint_draws(...) -> ndarray[n_samples, n_clones, P]`.** Precompute clone integer codes **once**; batch the Dirichlet draw and the softmax on the leading sample axis; scatter‑sum by clone. `tools/joint_distribution` is a thin single‑draw/summary DataFrame wrapper over this stack. Metrics consume the stack and reduce vectorized (`_compute/_reduce.py`). + +**7.3 The wins, priority‑ordered, with expected gains:** + +| # | operation | current hot path | fix | expected gain | prio | +|---|---|---|---|---|---| +| P0 | joint‑by‑clone reduction | `pd.DataFrame(...).groupby(level=0).sum()` rebuilt every draw (`_preprocessing.py:320-322`) over 1e5–1e6 cells | precompute clone codes once; **scatter‑add** (`np.add.at`/`np.bincount` CPU, `torch.index_add_`/`cupy.bincount` GPU) batched across all `n_samples` — grafiti `contingency.py`/`edge_tensor` verbatim | **10–50× on CPU** (pandas groupby → integer‑keyed bincount), multiplied again on GPU; the single biggest win | P0 | +| P0/P1 | Dirichlet sampling loop + softmax | `for i in range(n_samples): joint_distribution_posterior(...)` (`_metrics.py:296/363/548/741/987`), redrawing the **full** `p_ct` each call | draw all `n_samples` at once, **restricted to the covariate's ct rows first**; batch `softmax((logits+log b)/T)` over the leading axis; route via torch.cuda | removes the `n_samples`‑fold Python loop + per‑iteration `.uns`/DataFrame setup; collapses 200–1000 iterations into a few kernels | P0/P1 | +| P1 | entropy / MI reductions | `scipy.stats.entropy` per draw + `jd.loc[cl]` per clone (`_metrics.py:315/562-569/744`) | batched `xlogx` reduction over `[S, n_clones, P]`; MI as joint‑vs‑outer‑product; `nanmean`/HDI over the sample axis — grafiti `joint.py::_entropy/_mi`, float64 accumulators | removes two nested Python loops + slow pandas `.loc`; medium‑high | P1 | +| P1 | share the joint across metrics | each of the 4 metrics rebuilds the joint for the same `(covariate, n_samples)` | compute the `[S, n_clones, P]` stack **once per covariate** and reduce it in all four metrics; **build `groupby` groups in one batched pass** | divides the dominant cost by `#metrics`; realizes the draw‑once invariant (§6.1) | P1 | +| — | deterministic point estimate | `n_samples=0` still enters the sampler / returns one draw today | zero‑draw read of `uns[K.P_CT]` (R1) | correctness **and** cheapest path (no Monte‑Carlo at all) | — | +| P2 | latent UMAP | `umap.UMAP` with `import umap` at module top (`_preprocessing.py:20`) | `cuml.manifold.UMAP` behind `_use_gpu`, umap‑learn CPU fallback, **lazy import moved inside the function**; return float64 ndarray | cuML UMAP commonly **10–50×** umap‑learn, but runs once per analysis → lower total impact; also fixes the module‑top heavy import | P2 | + +**7.4 Grafiti guardrails (replicate ALL):** (1) every GPU lib imported **lazily inside** the function — `import tcri` never touches cupy/cuml/torch.cuda (the current `import umap` at module top already violates this and is fixed); (2) GPU deps never in `install_requires`; (3) permissive device ladder with `getDeviceCount()>0` verification; (4) `asnumpy` at every return boundary; (5) GPU body wrapped in `try/except` that degrades to CPU and reports which backend ran; (6) **float64 accumulators** so GPU matches CPU; (7) **validate the joint** (finiteness, nonnegativity, per‑row sum≈1 — per‑row, to catch cancelling defects) **on‑device before compute**; (8) **chunk the batched reduction** over cells/draws (grafiti Moran's‑I `chunk_size=256`, KDE `blk=8192`) to bound device memory on the large `[n_samples, n_cells, P]` tensor and avoid OOM. + +**7.5 Seeding (R7).** Draws move to a seeded `torch.Generator` (and cupy/torch.cuda RNG on GPU); `random_state` is threaded from every public sampling function. The old `np.random.seed` was a no‑op for the torch draws. + +--- + +## 8. Ordered PR sequence + +Each PR is independently landable with green CI, lowest‑risk first. + +### Phase 0 — Contract freeze + CI scaffolding *(docs/tests only; zero code change)* +- Frozen contract: `tcri/_contract.pyi` + `tests/test_contract_conformance.py`, ported from grafiti's `_pyi_gen.py`/`test_contract_conformance.py`. Seed `IMPLEMENTED = {}`; each target function flips to implemented as its phase lands. Markdown→`.pyi`→live‑signature drift fails CI. +- Land the corrected disposition map (§3) into `docs/contract/`. **Disposition is decided by one test — *is it core?* The disposable `example/` notebooks are never consulted for what to keep or drop.** +- **Risk:** none. **Depends on:** nothing. + +### Phase 1 — Shared helpers + `_keys` adoption *(internal dedup; API unchanged)* +- Create `_keys.py` (incl. `X_LOGITS/GATE_PROB/CLASSIFIER_TEMPERATURE/LOCAL_SCALE`), `_console.py`, `_stats.py` (true HDI, `prob_gt_zero`), `_distance.py` (bits/log2, JSD) (§4). +- **Adopt `_keys` at every read/write site.** Migrate `clone_size` to `K.CLONE_COL` in the **same** change that retires `tcri_clone_key` (no orphaned reader). +- **Risk:** low (mechanical, no numerics). **Depends on:** Phase 0. **Hazard:** must precede Phases 4/5. + +### Phase 2 — Safe deletions *(non-core symbols, unreferenced in the package)* +- Delete the non-core / dead: `get_latent_embedding`, `group_small_clones`, `register_probability_columns`, `remove_meaningless_genes`, `gene_entropy`, `polar_plot`, `metrics._ent`, `clone_fraction`, module‑level `dkl`, **`utils.probabilities` (and the `_plotting.py:18` import in the SAME PR)**, `pl.probability_distribution`, `pl.bayesian_mutual_information`, `SankeyNode.hex_to_rgb`. +- **Not deleted here** (each dropped in its own phase, *with* its replacement so no in‑package caller is orphaned): `compare_joint_distribution` (→ diag, Phase 8); `pl.mutual_information`/`pl.phenotypic_entropy` (keep+fix, Phase 7); the consolidated‑away `*_table`/`*_delta`/`clonality`/`clonotypic_entropy_base`/`ridge_delta_entropy` (Phase 6); `probability_ternary`/`plot_phenotype_probabilities` (Phase 6/7). +- **Risk:** very low. **Depends on:** Phase 1. + +### Phase 3 — Model module split *(mechanical; no behavior change)* +- Split `model/_model.py` (1074 ln) → `_model.py` + `_module.py` + `_priors.py` + `_classifier.py` + `_training.py`. Rename `c2p_mat → clone_phenotype_prior`. +- **`build_archetypes` keeps returning `(centers, labels)`** (labels drive `diag.archetypes`' cluster ordering); persist labels on the model/`uns`. Reconcile the default‑`K` mismatch (`build_archetypes` default `K=4` vs model `K=10`). +- Explicit `__all__` per module. **Risk:** low. **Depends on:** Phase 1. **Verify:** `test_model_setup`, `test_pyro_params` unchanged‑green. + +### Phase 4 — Model→AnnData streamline *(behavior change; kills the manager hack)* +- Implement §5: `setup_anndata` registration‑only (keeps `obs['indices']`; **removes the `_model.py:697` manager stash**); `group_singletons` enforced to precede `setup_anndata`. +- `register_model → model.to_anndata` writing the canonical set **including `X_LOGITS`, `GATE_PROB`, `CLASSIFIER_TEMPERATURE`, `LOCAL_SCALE`, and the raw `P_CT`** (R6/R8). `get_cell_phenotype_probs → predict` (order‑preserving loader, `obs_names` index). +- `write_adata_safely → save_tcri_session` (plain h5ad); load rebuilds the registry via `setup_anndata` and sets the pyro param store before any `get_p_ct`. +- **Risk:** HIGH. **Depends on:** Phase 1, Phase 3. **Gate:** `test_session_round_trip` rewritten — proves save/load reproduces `p_ct` + latent + `predict` probs with **no `tcri_manager` in `uns`**, `setup_anndata` leaves analysis `obs` untouched, and `to_anndata` writes **exactly** the canonical key set (asserts logits/gate/cls‑temp/local‑scale present). + +### Phase 5 — Engine consolidation *(the substrate; §6.1, §7.2)* +- Create `tools/` (aliased `tl`) and `_compute/` (`_xp.py`, `_joint.py`, `_reduce.py`). Implement `_joint_draws` (batched, device‑routable, scatter‑add) and the thin `tools/joint_distribution` wrapper; re‑export as `tcri.joint_distribution`. +- Signature per §6.1 with **`use_logits`** (renamed from `posterior=`), `random_state`, `covariate=None` → all‑covariates one‑pass, draw‑once invariant, clamped‑Dirichlet draws, single‑knob temperature, groupby by full‑space restriction. Fold in the weighting/normalization/zero‑clone/seed bug fixes. +- **Risk:** HIGH — every metric reads this. **Depends on:** Phase 4. **Verify:** `test_tools/test_joint`: `use_logits=False & n_samples=0 == tempered uns[K.P_CT]` **exactly**; `use_logits=True & n_samples=0 & T=1 == model.predict` aggregation; repeated `n_samples=0` **bit‑identical**; `n_samples>0` **torch‑seeded reproducible** and drawn from `Dirichlet(clamp(local_scale·m̃,1e-3))`; **draw‑counter == n_samples independent of #groups/#covariates**; weighting keyed on `ct`; serializable provenance. + +### Phase 6 — Metric‑API consolidation *(four metrics + `compare_groups`)* +- Populate `tools/_entropy.py`, `_mutual_information.py`, `_flux.py` per §6.2 (dual input, `normalize_mode='min'` default, support‑only clonotypic denominator + common‑denominator option, coherent per‑`sample_id` draws, `mean/sd/hdi` summary with **no `p_gt`**, serializable returns). Add public **`tl.compare_groups`** (§6.3). +- **Delete** `mi_compare`, `delta_entropy_table`, `flux_table`, `delta_clonotypic_entropy`, `phenotypic_entropy_delta`, `clonotypic_entropy_base`, `clonality`, `ridge_delta_entropy`, `compare_phenotypes`, the plural `*_entropies`, and `metrics.dkl`. Delete `metrics/` after migration. +- **Risk:** HIGH. **Depends on:** Phase 5. **Verify:** ranges (`[0,1]` normalized), `n_samples=0` determinism, `groupby` tidy shape, **dual‑input equivalence at `n_samples=0` only**, all‑zero‑clone → NaN, `phenotypic_flux` `cov_from/cov_to` + seeded draws, `compare_groups` recreates `mi_compare`'s per‑pair output, **no test asserts `n_samples=0 == mean(n_samples>0)`** (R3). + +### Phase 7 — Plotting split + pl twins *(fix the broken core plots)* +- Split `plotting/_plotting.py` (1437 ln) → `_base.py` (`_metric_boxplot` **with a `splitby=` axis**, `_finish`), `_colors.py` (`resolve_palette`), `_entropy.py`, `_mutual_information.py`, `_flux.py`, `_sankey.py`. +- Ship the four tl↔pl twins: `pl.clonotypic_entropy` (was `_by_phenotype`), `pl.phenotypic_entropy` **[FIX]**, `pl.mutual_information` **[FIX]**, `pl.phenotypic_flux` (sankey). pl functions are **cache renderers** (no metric math). **Retain `splitby`** as a distinct box‑hue axis (design decision — `groupby`=aggregation unit, `splitby`=comparison cohort) so two‑axis figures (dots=patient, boxes=response, x=phenotype) work; document per‑figure the `groupby`+`splitby` recipe. +- **Changelog behavior notes:** the default flips to `weighted=False` → `pl.mutual_information` displayed MI changes (cell‑weighted → per‑clonotype); `weighted=True` restores the old behavior. `pl.flux` boxplot and `clonality` plot dropped. +- **Drop (non-core):** `top_clone_umap`, `clone_size_umap`, `plot_phenotype_probabilities`, `probability_ternary`. **Risk:** medium. **Depends on:** Phase 6, Phase 1. **Verify:** each twin returns a `Figure`/`Axes` from a tidy tl result; sankey renders. + +### Phase 8 — `diag/` seeding *(new; additive)* +- `diag/_ppc.py`: the **fixed** `compare_joint_distribution` (no undefined‑global `NameError`; model `p(clone,phenotype)` vs empirical counts), phenotype‑probability calibration, reconstruction PPC, entropy/MI vs permutation null. **All return DataFrames.** Make the **model requirement explicit per function**: `joint_distribution_ppc`/calibration/permutation‑null run **adata‑only**; `reconstruction_ppc` **requires the live model** (ZINB decoder lives on the module). `diag/_training.py`: `plot_loss → loss curves`, `plot_archetypes → archetypes` (consumes `build_archetypes` labels). +- **Risk:** low‑medium. **Depends on:** Phase 4 (finalized model), Phase 5 (engine). **Verify:** each PPC returns the expected columns on the `trained_model` fixture; permutation‑null seeded; single‑model‑per‑process param‑store scoping honored. + +### Phase 9 — PGM → docs; utils finalize +- Move `build_nested_tcri_pgm`/`draw_tcri_pgm_nested` **out of the package into `docs/`**; drop `daft` from runtime deps → docs extras only. `utils/_utils.py → utils/_session.py` (session‑io + `_to_jsonable` only). +- **Risk:** low. **Depends on:** Phase 1, Phase 8. + +### Phase 10 — Notebook rewrite *(fresh tutorials against the new API)* +- Rewrite the `example/` notebooks **fresh** against the new API: `setup_anndata → TCRIModel → train → model.to_anndata`; `tcri.joint_distribution`; four metrics with `groupby`/`n_samples`; `tl.compare_groups`; four pl twins; `diag` checks. Dropped functions (`gene_entropy`, `polar_plot`, `probability_ternary`, `pl.flux` boxplot, `plot_phenotype_probabilities`, the `register_*_key` writers) have **no successor** — the rewrite simply does not call them. One canonical end‑to‑end tutorial notebook under `docs/`. **The notebooks are an OUTPUT of the refactor, never an input to it.** +- **Risk:** low (docs), high value. **Depends on:** Phases 4–8. + +### Phase 11 — Public API finalize + scverse ecosystem CI +- `tcri/__init__.py`: **explicit** named re‑exports, `sys.modules` aliases for `tl/pp/pl/ml/diag/ut`, top‑level `tcri.joint_distribution`, **remove `import *`**, kill leaked aliases (`centropy`/`pentropy`/`*_tl`) via `__all__`. Flip **all** target functions to `IMPLEMENTED`; drift now hard‑fails CI. Turn on the full scverse‑ecosystem gate (§9.2). +- **Risk:** low‑medium. **Depends on:** all prior phases. + +--- + +## 9. Testing + scverse‑CI strategy + +### 9.1 Tests per phase + +| phase | required tests | +|---|---| +| 0 | `test_contract_conformance` (markdown→`.pyi`→signature drift); import‑smoke py3.10/3.11 | +| 1 | `_stats` (`stars`, AUC/permutation, **true HDI vs equal‑tailed**, `prob_gt_zero` on a signed vector), `_distance` (`kl_divergence` bits/symmetry, JSD bound), `_console` silence flag, `_keys` "no stray literal" grep | +| 2 | every deleted symbol absent from `__all__` **and** unreferenced in the package (import‑graph test); `import tcri` green after `utils.probabilities` + its `_plotting.py:18` import go together | +| 3 | `test_model_setup`, `test_pyro_params` green through the split; `build_archetypes` returns `(centers, labels)`; submodule import smoke | +| 4 | **`test_session_round_trip` rewritten** — save/load reproduces `p_ct`+latent+`predict` probs with **no `tcri_manager`**; `setup_anndata` leaves analysis `obs` untouched but keeps `obs['indices']`; `to_anndata` writes **exactly** the canonical set incl. `X_LOGITS/GATE_PROB/CLASSIFIER_TEMPERATURE/LOCAL_SCALE`; `predict` order‑preserving | +| 5 | `test_tools/test_joint`: `use_logits=False,n=0 == tempered uns[P_CT]`; `use_logits=True,n=0,T=1 == predict` aggregation (**engine==predict** agreement test); repeat `n=0` bit‑identical; `n>0` torch‑seeded + clamped‑Dirichlet; **draw‑count == n_samples ⟂ #groups/#covariates**; weighting on `ct`; `groupby` via full‑space restriction does not trip the alignment guard | +| 6 | ranges, `n=0` determinism, `groupby` tidy shape, **dual‑input equivalence (n=0 only)**, all‑zero‑clone→NaN, `normalize_mode='min'` C‑independence, `compare_groups` recreates `mi_compare`; **golden regression** computing MI/entropy from a fixed `uns[P_CT]` by hand; **no `n=0==mean(n>0)` assertion** | +| 7 | each pl twin returns `Figure`/`Axes` from a tidy tl result without metric math; `splitby` two‑axis render; sankey renders | +| 8 | each PPC returns expected columns on `trained_model`; `reconstruction_ppc` requires model, others adata‑only; permutation‑null seeded | +| 9 | no `daft`/PGM import in the installed package (import‑graph test) | +| 10 | **tutorial execution** (`pytest --nbmake`) on the rewritten synthetic tutorial, end‑to‑end | +| 11 | contract test with **all** functions `IMPLEMENTED`; `__all__` completeness (public names ↔ contract); "no `import *`" AST test | + +Reuse `conftest.py` fixtures (`synthetic_adata`, `trained_model`, `mock_adata`); extend `mock_adata` to the unified `_keys` scheme in Phase 1 and to `to_anndata`'s canonical key set (incl. logits/gate/cls‑temp/local‑scale) in Phase 4. + +### 9.2 scverse‑ecosystem CI + +Bring `.github/workflows/tests.yml` to the cookiecutter‑scverse bar, layered so each phase stays green: +- **Matrix** py3.10 + py3.11 (extend to 3.12 before 1.0), `pip install -e ".[test]"`, `pytest tests/ -v --cov=tcri`. +- **Import‑smoke job** (grafiti pattern): `python -c "import tcri; from tcri.model._model import TCRIModel; from tcri.tools._joint import joint_distribution"` — **and assert no GPU lib (cupy/cuml/torch.cuda) was imported** (guardrail #1). +- **Lint/format gate:** `pre-commit` (ruff + ruff‑format) — add in Phase 1, enforce from Phase 3. +- **Contract‑conformance job** (Phase 0 on) — the markdown→`.pyi`→signature guardrail that lets views land independently. +- **CPU‑only correctness job** — the full suite must pass with no GPU present; a separate optional GPU job (if a runner is available) asserts float64 GPU≈CPU parity. +- **Notebook‑execution job** (`nbmake`) from Phase 10. +- **Docs build** (sphinx + numpydoc; `.readthedocs.yaml` present) must pass; public API fully docstringed/typed. +- **Ecosystem checklist** (final gate, Phase 11): AnnData‑native `setup_anndata`, no `import *`, typed public surface, tutorial notebook, `diag` returns data not plots, GPU strictly opt‑in. + +--- + +## 10. Dependencies & ordering hazards + +``` +0 ─▶ 1 ─┬─▶ 2 + ├─▶ 3 ─▶ 4 ─▶ 5 ─▶ 6 ─▶ 7 ─▶ 8 ─▶ 9 ─▶ 10 ─▶ 11 + └─────────────────────(1 gates 4,5,6) +``` +Hard, load‑bearing constraints: +1. **`_keys` (1) before `to_anndata` (4) and the engine (5).** Both rewrite key IO; centralizing strings first keeps diffs local and prevents two‑scheme drift. `clone_size`'s reader flip and `tcri_clone_key`'s writer retirement must land **together**. +2. **Model split (3) before `to_anndata` (4).** `to_anndata` is a `TCRIModel` method; split first so the behavior‑change diff is readable. +3. **`to_anndata` (4) before the engine (5).** The engine reads exactly the canonical keys/**logits/gate/cls‑temp/local‑scale** `to_anndata` writes; if the write‑set changes after the engine exists, `use_logits=True` breaks silently. The **manager‑in‑`uns` removal (at `setup_anndata:697`) is the single highest‑risk hazard** — it changes the session round‑trip; `test_session_round_trip` is rewritten in the same PR. +4. **Engine (5) before metrics (6).** Metrics are thin readers; building them against a pre‑fix engine would bake in the weighting/normalization/seed/Jensen bugs. +5. **Metrics (6) before pl twins (7).** pl are cache renderers. +6. **Delete `*_table`/`*_delta`/`clonality`/`_base`/`ridge_delta_entropy`/`compare_phenotypes` WITH their replacement (Phase 6), never before** — landing the `groupby`/`compare_groups` replacement in the same PR keeps CI green with no half‑migrated in‑package state. +7. **Rename freeze before Phase 5.** `from_this/to_that → cov_from/cov_to`, `flux → phenotypic_flux`, `metrics/ → tools/`, `posterior= → use_logits`, dropping `point_estimate=` are all breaking; batch them into one breaking window. +8. **`group_singletons` before `setup_anndata`** (clone‑relabel ordering) — enforced by `setup_anndata`. +9. **Deferred, designed‑for:** `@tl_result` uns‑cache + `get.py`. Every tl function returns an **h5ad‑serializable** tidy result (flat columns, no object arrays, serializable provenance) so the decorator is a one‑line wrap; until then pl recomputes via the engine. Do **not** block the refactor on the cache. + +**Independent / parallelizable:** Phase 2 (safe deletions) and Phase 8 (`diag`, once 4–5 land) proceed alongside neighbors without contending for the same files. + +--- + +## 11. Risks & open items + +- **Highest‑risk hazard: the manager‑stash removal + write‑set change (Phase 4).** It reshapes the train→`to_anndata`→session round‑trip and adds four persisted keys. Mitigation: rewrite `test_session_round_trip` in‑PR; assert the exact canonical set; keep the CPU‑only correctness job as the tripwire. +- **`use_logits=True` parity depends on persisted model facts (R8).** If any of `X_LOGITS`/`GATE_PROB`/`CLASSIFIER_TEMPERATURE`/`LOCAL_SCALE` is dropped, parity degrades silently to the additive rule. Mitigation: the engine==predict agreement test (Phase 5) and the canonical‑set assertion (Phase 4). Note `gate_prob` defaults to `None` (models ungated by default), so the bug is latent until someone trains with a gate — the guarantee still must hold. +- **Estimator semantics (R3) are a documentation risk, not a code risk.** `n_samples=0` (plug‑in) and the `n_samples>0` `mean` are *different estimators*; users will expect them to match. Mitigation: docstrings state the Jensen gap explicitly; no conformance test asserts equality. +- **Global pyro param store.** `get_p_ct`/`to_anndata`/diag PPCs read the process‑global `q_p_ct_raw`; loading two models in one process clobbers it. **Open item:** namespace params per model or ship a documented single‑model‑per‑process contract with load setting the store immediately before use. Pre‑existing, not introduced by the refactor, but the multi‑model diag/round‑trip workflows expose it. +- **`temperature≠1` re‑centers the sampled distribution** away from the learned posterior (R6); intervals under non‑default `T` are not the model posterior. Documented; `T=1` is the parity point. +- **Partial posterior.** `n_samples>0` intervals capture only `p_ct` uncertainty (classifier logits fixed). Documented so they are not read as full posterior‑predictive uncertainty. +- **GPU memory on large `[n_samples, n_cells, P]`.** Without chunked reductions (guardrail #8) the batching win becomes an OOM risk on realistic `n_samples × n_cells`. Mitigation: chunk over cells/draws; try/except degrade to CPU. +- **`splitby`** is retained as a distinct box‑hue axis by design (≠ `groupby`, not a notebook artifact); its `pl` signatures are fixed in Phase 7. +- **`@tl_result` cache is deferred.** Return shapes are constrained now to be h5ad‑serializable so the later wrap is a one‑liner; the cache‑key scheme is specified but not implemented. +- **cuML/rapids GPU parity** for UMAP is layout‑different (both valid, as grafiti documents); not bit‑identical to CPU. No test asserts embedding bit‑identity. +- **Open (deliberately deferred):** py3.12 matrix, the `get.py` accessor surface, and any KL/JSD default change for `phenotypic_flux` beyond offering JSD (l1 remains the bounded default). + +--- + +## 12. Rollout checklist (per PR) + +- [ ] One behavior change (or zero, if a mechanical move). +- [ ] `import tcri` green; `tl/pp/pl/ml/diag/ut` handles intact; **no GPU lib imported at import time**. +- [ ] No new key‑string literal outside `_keys.py`; no new `import *`. +- [ ] Deleted symbols removed from `__all__` and unreferenced **in the package** (import‑site + call‑site grep). +- [ ] Contract test updated (`IMPLEMENTED` flipped for functions this PR lands). +- [ ] Phase‑specific tests from §9.1 present and green on py3.10/3.11 (CPU‑only job passes). +- [ ] Rename/disposition rows for this phase applied everywhere (code + fixtures). +- [ ] Any persisted‑key change reflected in `_keys`, `to_anndata`'s canonical set, and the round‑trip test. + +--- + +## Appendix A — Audit findings by severity (traceability) + +Every finding folded into the body above, cross‑referenced so nothing is lost. **Src:** PC = plan‑correctness, MS = math/stats, ML = missing‑links, GPU = GPU/optimization, PVM = prior‑vs‑mean resolution. + +### A.1 Blocking / High + +| # | src | finding | resolved in | +|---|---|---|---| +| B1 | PC/ML | `to_anndata` "canonical minimum" drops `obsm['X_tcri_logits']` that the default engine path hard‑requires | §5.2, R8, Phase 4/5 | +| B2 | PC/ML | `gate_prob` (and `classifier_temperature`) are model attributes never persisted → gate‑aware parity infeasible from `adata` | §5.2, R8, Phase 4 | +| B3 | PC | `setup_anndata` cannot be "no obs mutation" — must keep registered `obs['indices']` | §1.1, §5.1, Phase 4 | +| B4 | PC | Deleting `utils.probabilities` (Phase 2) breaks `import tcri` via `_plotting.py:18` import | §1.1, Phase 2 | +| B5 | PC | `groupby` by slicing `adata` trips the full‑space‑vs‑subset alignment guard | §6.1, Phase 5 | +| B6 | MS | Jensen gap: `metric(E[p]) ≠ E[metric(p)]`; `n=0` (plug‑in) vs `mean(n>0)` disagree — drop the equality test | R3, §1.2, Phase 6 | +| B7 | MS | `p_gt`/P(>0) vacuous per‑metric (all metrics ≥0) — only on signed contrasts | R5, §6.2/6.3 | +| B8 | MS | `n>0` draws must use the guide's **clamped** `Dirichlet(clamp(local_scale·m,1e-3))` | R4, §6.1, Phase 5 | +| B9 | MS | MI `normalize_mode='average'` breaks under uniform‑clone prior (`H_c=log2 C`) — default `'min'` | §6.2, Phase 6 | +| B10 | ML | `_stats` is private — no public comparison entry point; add `tl.compare_groups` | §4, §6.3, Phase 6 | +| B12 | ML | `covariate × groupby` semantics unspecified; requires cell‑informed path / clone‑disjoint groups | §6.1, Phase 5/6 | +| B13 | ML | No mechanism shares one `p_ct` draw across groups/covariates → O(groups×cov×n_samples) | §6.1 draw‑once, §7, Phase 5 | +| B14 | GPU | Engine must be a batched device‑routable core returning `[n_samples,n_clones,P]`, DataFrame at boundary | §7.2, Phase 5 | + +### A.2 Medium + +| # | src | finding | resolved in | +|---|---|---|---| +| M1 | PC | `n=0 == uns[P_CT]` holds only for `use_logits=False`; default folds logits — scope the test | R2, §6.1, Phase 5 | +| M2 | PC | `groupby` unrepresentable in ct‑level table for non‑clone‑determined columns | §6.1 | +| M3 | PC/MS | `adata_or_jd` fast path incompatible with `n>0`/`groupby` — restrict to `n=0`,`groupby=None` | §6.2 | +| M4 | PC | `clone_size` reads retired `tcri_clone_key` — migrate reader with writer | §4, Phase 1 | +| M5 | PC | `build_archetypes` must keep returning `labels`; default‑K mismatch (4 vs 10) | §Phase 3, §Phase 8 | +| M6 | MS | `clonotypic_entropy` `log2(C)` denominator group‑dependent, inflated by absent clones | §6.2, Phase 6 | +| M7 | MS | "HDI" is actually equal‑tailed — implement true HDI or rename | §4 (`_stats`), Phase 1/6 | +| M8 | MS/ML | Dual‑input + `n>0` ill‑defined — force `n=0` on bare jd | §6.2 | +| M9 | MS/ML | Temperature re‑centers draws / double‑tempering with guide_temperature — single knob, store raw mean | R6, §5.2, §6.1 | +| M10 | ML | `splitby` is a distinct box‑hue axis (≠ `groupby`) — retain it on the pl twins | §Phase 7 | +| M12 | ML | No `seed`/`random_state` on engine/metrics — add and thread to torch | R7, §6, Phase 5/6 | +| M13 | ML | Normalized entropy/MI non‑comparable across groups (group‑specific denominators) | §6.2 common‑denominator option | +| M14 | ML | `@tl_result` h5ad‑hostile returns (MultiIndex, object arrays, `.attrs`) | §6.2, §10(9) | +| M15 | ML | `n>0` reduction must keep draws coherent (same `p_ct` per `sample_id`) | §6.1/6.2, Phase 5/6 | +| M16 | GPU | Seed torch (not `np.random`) for reproducible draws | R7, §7.5 | + +### A.3 Low + +| # | src | finding | resolved in | +|---|---|---|---| +| L1 | PC | `predict()` DataFrame indexing assumes order‑preserving loader | §5.2, Phase 4 | +| L2 | PC | `pl.mutual_information` default flip to `weighted=False` changes displayed numbers (`weighted=True` restores) | §Phase 7 changelog | +| L3 | PC | `tcri_manager` stash is in `setup_anndata:697`, not `register_model` | §5.1, Phase 4 | +| L4 | PC/ML | moved `plot_phenotype_probabilities` reads `X_tcri_phenotypes`, not `X_tcri_probabilities` | §3, Phase 10 | +| L5 | PC | `get_p_ct` reads the global pyro param store — multi‑model clobber | §5.3, §11, Phase 8 | +| L6 | MS | KL flux asymmetric/unbounded, base mismatch — bits/log2, offer JSD, l1 default | §4 (`_distance`) | +| L7 | MS | `n>0` intervals are partial posterior (p_ct only) — document | §6.2, §11 | +| L8 | MS | `local_scale` uns fallback `1.0` corrupts draw variance — `to_anndata` always writes it; engine raises if missing at `n>0` | §5.2, §6.1, Phase 4/5 | +| L9 | ML | `group_singletons` ordering vs `setup_anndata` unspecified | §5.1, Phase 4 | +| L10 | ML | `diag` model‑required vs adata‑only inconsistent; missing golden + engine==predict agreement tests | §Phase 8, §9.1 | +| L11 | PVM | current `softmax(log p_ct/T)` path is already the posterior mean at `T=1`; "prior/non‑posterior" label is a misnomer — collapse the axis to `use_logits` | R1/R2, §6.1 | +| L12 | GPU | module‑top `import umap` violates lazy‑import guardrail; cuML UMAP behind `_use_gpu` | §7.3/7.4, Phase 5+ | \ No newline at end of file diff --git a/docs/contract/tcri_inventory_data.json b/docs/contract/tcri_inventory_data.json new file mode 100644 index 0000000..d93d0f1 --- /dev/null +++ b/docs/contract/tcri_inventory_data.json @@ -0,0 +1 @@ +{"grafiti": {"diagnostics_vs_pl": "The split is by RETURN TYPE and lifecycle, not by topic. `gf.pl` (plotting/) returns Matplotlib Figure/Axes/None and NEVER computes \u2014 every pl function is a \"cache renderer\" that calls `load_result(adata, key)` and draws the already-stored tl blob (see plotting/_affinity.py: `res = load_result(adata, key); M = res[\"coupling\"]`). `gf.diag` (diagnostics/) returns tidy DATA (a pandas DataFrame of quality/sanity numbers), runs against FINALIZED model outputs, and is deliberately NOT part of the tl-writes-cache / pl-reads-cache loop. Key signals from diagnostics/_affinity.py: (1) it reads raw model state directly (`adata.uns[grafiti_affinity_matrix]`, `adata.obsm[grafiti_motif_2d]`) and recomputes an empirical target on the fly, rather than reading a tl result blob; (2) it does NOT use the `@tl_result` decorator and writes nothing back to uns \u2014 it is a pure read-only computation returning a DataFrame (cosine, signed_slope, residual_norm, coupling_norm, gamma_empirical, mi, n_cells per group); (3) it is validation/concordance \u2014 \"does the learned coupling agree with the data?\" \u2014 i.e. did the model fit, not a scientific result the user plots. So the rule: an analysis result the user wants to VISUALIZE and that should be cached + reproducible \u2192 tl + pl twin. A cheap numeric sanity/concordance check on the finalized fit that returns a small table and is never plotted \u2192 diag. diagnostics is intentionally thin (one module, _affinity.py, plus a 5-line __init__ exposing `affinity_concordance` + its `affinity_diagnostic` alias). It also carries its OWN private validators (_read_coupling/_read_q2d/_summarize/_mi) rather than sharing tl's helpers, keeping it decoupled from the contract loop.", "helper_pattern": "Helpers are factored by LAYER into private modules so no math or house-style is duplicated, and public modules import them rather than re-implementing. Three distinct helper tiers:\n\n1. `_state/` \u2014 the shared substrate every pp/tl/pl/get sits on. `storage.py` realizes the entire \"write uns as dict-of-arrays + params + version, and return the tidy result\" convention ONCE as the `@tl_result(key=, version=, schema=)` decorator (~16 tl functions physically cannot drift from it); it also owns `load_result` / `load_result_params` / `decode_blob` used by pl and get to read blobs back, plus `_encode`/`_decode` for h5ad-safe serialization. `resolve.py` is the ONE place key-resolution lives (`resolve_fov_key`, `resolve_motif_key`, `resolve_celltype_key`, `resolve_connectivity_key`, `grouping_values`) \u2014 the docstring notes this replaced a \"literal-guessing bug that had spread to five functions.\" `keys.py` is the single source of truth for every uns/obs/obsm string constant (imported everywhere as `from .._state import keys as K`). `schemas.py` holds TypedDicts that the decorator runtime-checks. `experimental.py` holds the `@experimental`/`ExperimentalWarning` marker.\n\n2. `_compute/` \u2014 view-agnostic, device-routed math shared across tools. `_xp.py` is the device-dispatch fallback (`resolve_device`/`get_xp`/`asnumpy` \u2192 numpy or cupy). `stats.py` (550 lines) + `joint.py`, `graph.py`, `contingency.py`, `composition.py`, etc. hold the heavy reductions (population-joint math, Moran's I, permutation nulls, connected components) so a tl file like tools/_affinity.py just calls `grouped_coupling_strength(...)` from `_compute.joint` and stays thin.\n\n3. `plotting/_base.py` \u2014 the pl \"house style\" layer: `_finish(fig, ret, show, save)` implements the scanpy show/save return convention (Figure|Axes|None) once, and `_labeled_heatmap` / `_clustered_heatmap` / `_heatmap_figsize` are the shared visual primitives every heatmap view reuses. `plotting/_colors.py` is the shared palette/color-resolution helper (`resolve_colors`, `MOTIF_PALETTE`, etc.). Naming: cross-cutting helper SUB-PACKAGES are underscore-prefixed dirs (`_state`, `_compute`); within-package shared helpers are underscore-prefixed FILES (`_base.py`, `_colors.py`); helper FUNCTIONS are underscore-prefixed (`_finish`, `_encode`, `_summarize`).", "public_api_pattern": "Two-level explicit re-export with `__all__` declared at BOTH levels \u2014 deliberately NO `import *` from monoliths. (1) Each private per-view impl module declares its own `__all__` (e.g. tools/_affinity.py: `__all__ = [\"affinity_matrix\", \"affinity_gamma\"]`). (2) The sub-package `__init__.py` does NOT `from ._x import *`; it names each symbol explicitly (`from ._affinity import affinity_matrix, affinity_gamma`) and re-declares an aggregate `__all__` that lists every public name, grouped by view with `# joint`/`# motif`/`# affinity` comment banners (see plotting/__init__.py and tools/__init__.py). This keeps the namespace clean: numpy/pandas/matplotlib and every `_helper` stay unexported even though the impl files import them. (3) The top-level grafiti/__init__.py binds the short aliases (`from . import tools as tl`, `... diagnostics as diag`, etc.), then registers them in `sys.modules` (`sys.modules[f\"{__name__}.tl\"] = tl`) so `import grafiti.tl` and `gf.tl.x` both resolve; heavy namespaces (pl/ml/ds) are wrapped in a `_LazyModule` proxy so plain `import grafiti` stays light and doesn't eagerly import matplotlib/scanpy/torch. A generated `_contract.pyi` stub plus `test_signature_matches_contract` machine-checks that every tl signature matches the documented contract.", "layout": "grafiti is a flat single package `grafiti/` where the public API is a set of scanpy-style sub-packages, each aliased to a two-letter handle in the top-level __init__: `model/`\u2192`gf.ml`, `tools/`\u2192`gf.tl`, `plotting/`\u2192`gf.pl`, `preprocessing/`\u2192`gf.pp`, `datasets/`\u2192`gf.ds`, `diagnostics/`\u2192`gf.diag`, plus top-level `get.py`\u2192`gf.get` and `migrate.py`. Cross-cutting internals live in underscore-prefixed private sub-packages that are NOT part of the public surface: `_state/` (the shared AnnData-state substrate: keys, resolve, storage, schemas, experimental) and `_compute/` (view-agnostic device-routed math). A generated `_contract.pyi` conformance stub sits at package root, and `scratch/` holds dev throwaways. MODEL: the neural model lives in `model/` as a multi-file package \u2014 `model/grafiti.py` holds the `Grafiti` class (and `TensorizedDEC` head), with siblings `_config.py`, `_data_loader.py`, `_decoders.py`, `_encoder.py`, `_losses.py`, `_early_stopping.py`, `_tiling.py`, `_regularizations.py`, `_recursive_normalized_cut.py`, `_utils.py`, `attention.py`; `model/__init__.py` is just two lines re-exporting `Grafiti` and `AffinityConfig`. The model owns training and `finalize_anndata`, which WRITES the canonical state (uns[grafiti_config], uns[grafiti_affinity_matrix], obsm[grafiti_motif_2d]) that tl reads. tl/pl/pp are each strictly ONE FILE PER CONTRACT VIEW (joint / motif / celltype / affinity / graph / patch), never a monolith: tools/_joint.py, tools/_motif.py, tools/_affinity.py, ... and plotting/ mirrors them 1:1 by filename.", "naming_conventions": "Public sub-packages are plain names (`tools`, `plotting`, `diagnostics`, `preprocessing`, `model`, `datasets`) exposed via short two-letter aliases (tl/pl/diag/pp/ml/ds). Within each, the actual code lives in PRIVATE underscore-prefixed impl modules named by contract view: `_joint.py`, `_motif.py`, `_celltype.py`, `_affinity.py`, `_graph.py`, `_patch.py`, `_spatial.py`. The pattern is \"public package, private per-topic _impl modules; __init__ is the only public face.\" Cross-cutting helper collections are underscore-prefixed PACKAGES (`_state/`, `_compute/`); within-package shared helpers are underscore-prefixed FILES (`_base.py`, `_colors.py`, `_spatial_core.py`); private functions/vars are underscore-prefixed (`_finish`, `_encode`, `_read_affinity_bundle`, `_EPS`). tl\u2194pl TWINS share the same view filename AND the same public function names where they pair (tl.affinity_matrix writes, pl.affinity_matrix reads); pl adds extra display-only derivations with descriptive suffixes (affinity_rankings, motif_markers_matrixplot, motif_markers_specificity, graph_centrality_heatmap vs graph_centrality_scatter). Every string key is a NAMED CONSTANT in `_state/keys.py` (K.AFFINITY, K.MOTIF_2D) with a nested `Config` class for config-field names \u2014 no string literal for an AnnData slot appears anywhere else. Result blobs are versioned (`version=1` in the decorator) and namespaced `grafiti_*`. The `.pyi` contract stub and diagnostics-alias pattern (`affinity_diagnostic = alias of affinity_concordance`) round it out.", "lessons_for_tcri": "tcri today is the anti-pattern grafiti fixed: four monoliths (metrics/_metrics.py 1008 lines, plotting/_plotting.py 1437, model/_model.py 1074, utils/_utils.py 665), each `__init__.py` is a bare `from .._x import *` with NO `__all__`, so numpy/pandas/every helper leaks into the public namespace, and there is no state/keys/storage layer or diagnostics category. Concrete steps to mirror grafiti:\n\n1. SPLIT THE MONOLITHS ONE-FILE-PER-TOPIC. Break metrics/_metrics.py into metrics/_.py (e.g. _clonality.py, _phenotypic.py, _flux.py per metric family) and plotting/_plotting.py + _sankey.py into plotting/_.py that mirror the metric files 1:1 by name. Keep each file focused; grafiti's largest view file is ~650 lines.\n\n2. ADD A `_state/` PACKAGE. Create tcri/_state/ with keys.py (every obs/obsm/uns string as a named constant, imported as `from .._state import keys as K`), resolve.py (one home for all key/groupby resolution \u2014 kills literal-guessing bugs), storage.py (a `@tl_result(key=, version=, schema=)` decorator that writes uns[key] = dict-of-arrays + params + version AND returns the tidy result, so every metric caches identically and is reproducible), and schemas.py (TypedDict per return, runtime-checked by the decorator).\n\n3. ADD `gf.get`-STYLE ACCESSORS. A tcri/get.py with `result(adata, name)` / `params(...)` / per-metric `_df` accessors that call load_result, so a user who ran a metric inplace can reconstruct the table.\n\n4. MAKE PLOTTING PURE CACHE-RENDERERS. Every tcri pl function should `load_result(adata, key)` and only draw \u2014 never recompute a metric. Add plotting/_base.py with a `_finish(fig, ret, show, save)` scanpy show/save helper and shared heatmap/palette primitives (tcri already has _sankey; generalize the shared parts into _base.py + _colors.py).\n\n5. DECLARE `__all__` AT BOTH LEVELS AND DROP `import *` FROM MONOLITHS. Each _topic.py declares its own `__all__`; each package __init__ names symbols explicitly and re-declares an aggregate `__all__` grouped by topic. This is the single biggest namespace-cleanliness win.\n\n6. ADD A `diagnostics/` PACKAGE (tcri's missing category). Put read-only sanity/concordance checks that return small DataFrames and are never plotted here (e.g. model-fit quality, posterior/joint-distribution concordance vs empirical \u2014 tcri already has test_joint_distribution_posterior, that concordance logic is a natural diag.*). Expose as `tcri.diag`, keep it thin, and DON'T route it through the tl-cache/pl loop \u2014 it reads model state directly and returns numbers.\n\n7. SPLIT THE MODEL PACKAGE. model/_model.py (1074 lines) should become a model/ package: model/.py for the class, plus _config.py, _losses.py, _encoder.py, etc.; model/__init__.py re-exports only the public class(es).\n\n8. KEEP THE TOP-LEVEL __init__ ALIAS+sys.modules PATTERN tcri already has (tl/pp/pl/ut/ml), but add `diag`, consider lazy-loading heavy namespaces (pl/ml) via a _LazyModule proxy, and add a `get` alias.\n\n9. MOVE VIEW-AGNOSTIC MATH OUT OF metrics into a `_compute/` package with an `_xp.py` device helper if any GPU/heavy reductions exist, so metric files stay thin wrappers.", "modules": [{"name": "model/ (gf.ml)", "purpose": "The neural model + training. Multi-file package; the class trains and its finalize_anndata WRITES the canonical AnnData state (uns config/affinity bundle, obsm soft assignments) that every tl reads. model/__init__.py re-exports only the public class(es).", "key_files": "model/grafiti.py (Grafiti class + TensorizedDEC head), _config.py (AffinityConfig), _data_loader.py, _encoder.py, _decoders.py, _losses.py, _early_stopping.py, _tiling.py, _regularizations.py, attention.py; __init__.py = 2-line re-export of Grafiti + AffinityConfig"}, {"name": "tools/ (gf.tl)", "purpose": "Analysis tools run after training. Strictly ONE FILE PER CONTRACT VIEW (not a monolith). Each tool is decorated with @tl_result: it computes, the decorator caches to uns[key] as dict-of-arrays+params+version, and it returns the tidy result. Bodies stay thin by calling _compute for heavy math.", "key_files": "_joint.py, _motif.py (471L, largest), _celltype.py, _affinity.py, _graph.py, _patch.py (652L); __init__.py explicitly re-exports every symbol grouped by view with an aggregate __all__"}, {"name": "plotting/ (gf.pl)", "purpose": "Cache renderers that MIRROR tl 1:1 by filename and, where paired, by function name. Every pl function calls load_result(adata,key) and only draws \u2014 NEVER recomputes. Returns Figure|Axes|None via the shared _finish convention. Adds display-only derivations (rankings, matrixplots) beyond the tl twins.", "key_files": "_joint.py, _motif.py, _celltype.py, _affinity.py, _graph.py, _patch.py, _spatial.py + _spatial_core.py; _base.py (shared _finish + heatmap primitives), _colors.py (palettes/resolve_colors)"}, {"name": "diagnostics/ (gf.diag)", "purpose": "The category tcri LACKS. Read-only sanity/concordance checks on FINALIZED model outputs that return small tidy DataFrames of quality numbers and are never plotted. Reads raw model state directly, does NOT use @tl_result, writes nothing to uns, carries its own private validators. Deliberately thin and decoupled from the tl-cache/pl loop.", "key_files": "_affinity.py (affinity_concordance + affinity_diagnostic alias; private _read_coupling/_read_q2d/_summarize/_mi); __init__.py = 5 lines"}, {"name": "_state/ (private)", "purpose": "The shared substrate every pp/tl/pl/get sits on \u2014 the layer that makes the whole toolkit consistent. Not part of the public surface. Realizes the storage convention, key resolution, key registry, and return schemas each exactly ONCE so ~16 functions cannot drift.", "key_files": "keys.py (single source of truth for every uns/obs/obsm string constant + nested Config class), resolve.py (all key/groupby resolution), storage.py (@tl_result decorator + load_result/decode_blob + h5ad-safe encode/decode), schemas.py (TypedDicts, runtime-checked), experimental.py (@experimental marker)"}, {"name": "_compute/ (private)", "purpose": "View-agnostic, device-routed heavy math shared across tools, behind a device= arg. Keeps tl files thin wrappers. Not public.", "key_files": "_xp.py (resolve_device/get_xp/asnumpy \u2192 numpy or cupy), stats.py (550L: Moran's I, permutation nulls), joint.py, graph.py, contingency.py, composition.py, occupancy.py, patch.py, ncut.py, fields.py, embedding.py"}, {"name": "preprocessing/ (gf.pp)", "purpose": "Steps run BEFORE training (setup_anndata pipeline, spatial graph, SVG, scVI, UMAP, link pruning). Same one-file-per-view + explicit __all__ pattern as tl/pl; supporting primitives live in _compute.", "key_files": "_preprocessing.py; __init__.py explicitly re-exports 8 functions with a grouped __all__"}, {"name": "get.py (gf.get) + _contract.pyi", "purpose": "get.py = accessors that read cached tl results back out of uns (result/params generic + per-tool _df tidy-table accessors), via a name\u2192(uns key, tidy sub-key) table. _contract.pyi = generated conformance stub whose signatures are machine-checked against every tl by the test suite.", "key_files": "get.py (_RESULTS map, result(), params(), *_df accessors, own __all__); _contract.pyi (generated, do-not-edit); tests/_pyi_gen.py"}]}, "merged": [{"name": "_ok", "kind": "function", "signature": "_ok(msg:str, quiet=False)", "purpose": "Print a green success/checkmark line unless quiet.", "io": "prints ANSI line to stdout", "calls": "-", "label": "helper", "consolidate_into": "", "notes": "Console-print family (ANSI); used by mutual_information/delta_clonotypic_entropy. Belongs in a shared console-helper module (duplicated in pl/pp).", "module": "tcri/metrics/_metrics.py", "namespace": "tl"}, {"name": "_info", "kind": "function", "signature": "_info(key:str, txt:str, quiet=False)", "purpose": "Print a cyan key-value info line unless quiet.", "io": "prints ANSI line to stdout", "calls": "-", "label": "helper", "consolidate_into": "", "notes": "Console-print family; used across mutual_information/delta_clonotypic_entropy. Extract to shared helper module (also defined in pl/pp).", "module": "tcri/metrics/_metrics.py", "namespace": "tl"}, {"name": "_warn", "kind": "function", "signature": "_warn(msg:str, quiet=False)", "purpose": "Print a yellow warning line unless quiet.", "io": "prints ANSI line to stdout", "calls": "-", "label": "helper", "consolidate_into": "", "notes": "Console-print family but 0 callers in this file (live copies are in pl/_plotting.py:43 and pp/_preprocessing.py:56). Consolidate the triplicated console helpers into one shared module.", "module": "tcri/metrics/_metrics.py", "namespace": "tl"}, {"name": "_fin", "kind": "function", "signature": "_fin(quiet=False)", "purpose": "Print a magenta 'Done!' flourish unless quiet.", "io": "prints ANSI line to stdout", "calls": "-", "label": "helper", "consolidate_into": "", "notes": "Console-print family but 0 callers in this file (the used copy is pp/_preprocessing.py:59). Triplicated across modules; move to shared helper module.", "module": "tcri/metrics/_metrics.py", "namespace": "tl"}, {"name": "_ascii_hist", "kind": "function", "signature": "_ascii_hist(samples, bins=25, width=40) -> str", "purpose": "Build a text ASCII histogram of a sample vector for notebook/SSH display.", "io": "-", "calls": "np.histogram", "label": "helper", "consolidate_into": "", "notes": "Pure string builder; used by delta_clonotypic_entropy, mutual_information, flux (graph=True). Shared display helper.", "module": "tcri/metrics/_metrics.py", "namespace": "tl"}, {"name": "_mi_from_joint", "kind": "function", "signature": "_mi_from_joint(pxy:np.ndarray, normalised:bool, mode:str='average') -> float", "purpose": "Compute (optionally normalised) mutual information from an already-normalised C\u00d7P joint table.", "io": "-", "calls": "-", "label": "core", "consolidate_into": "", "notes": "The single-source-of-truth MI kernel behind mutual_information (called at lines 732,745). Pure numpy; essential MI substrate.", "module": "tcri/metrics/_metrics.py", "namespace": "tl"}, {"name": "mi_compare", "kind": "function", "signature": "mi_compare(adata, groupby, groups=None, treatment=None, n_samples=50, patient_col=None, clone_col=None, covariate_col=None, verbose=True, **mi_kwargs)", "purpose": "Build a tidy per-patient MI samples+summary table across covariates and group pairs.", "io": "reads uns['tcri_metadata'] (batch_col/clone_col/covariate_col) and obs[groupby,patient_col,clone_col,covariate_col]; returns dict{samples,summary,pairs,params}", "calls": "mutual_information (posterior sampling per patient/covariate); tqdm", "label": "core", "consolidate_into": "", "notes": "Explicit core tidy-table builder for MI; wrapped by pl.mi_compare (_plotting.py:95).", "module": "tcri/metrics/_metrics.py", "namespace": "tl"}, {"name": "dkl", "kind": "function", "signature": "dkl(p, q)", "purpose": "KL divergence of p\u2016q via scipy.stats.entropy after clipping.", "io": "-", "calls": "scipy.stats.entropy", "label": "dead-broken", "consolidate_into": "", "notes": "0 callers anywhere in tcri. Superseded by flux's inner dkl_func which reimplements KL for the 'dkl' distance branch. Remove or fold into flux.", "module": "tcri/metrics/_metrics.py", "namespace": "tl"}, {"name": "_ent", "kind": "function", "signature": "_ent(p, base=2)", "purpose": "Normalise a vector and return its Shannon entropy in given base.", "io": "-", "calls": "scipy.stats.entropy", "label": "dead-broken", "consolidate_into": "", "notes": "0 callers. Would-be stats helper never wired in; the entropy metrics inline their own clip/normalise/entropy instead. Dead.", "module": "tcri/metrics/_metrics.py", "namespace": "tl"}, {"name": "clonotypic_entropy_base", "kind": "function", "signature": "clonotypic_entropy_base(adata, covariate_label, phenotype, *, base=2, normalised=True, temperature=1.0, clones=None, weighted=False, posterior=True, combine_with_logits=True) -> float", "purpose": "Single-draw normalised clonotypic entropy for ONE phenotype at ONE covariate.", "io": "-", "calls": "joint_distribution_posterior / joint_distribution (posterior sampling engine)", "label": "redundant", "consolidate_into": "clonotypic_entropy", "notes": "Computes the same normalised H[P(c|\u03c6)] as clonotypic_entropy but per-single-phenotype/single-draw; only extra is posterior=False + weighted support. Sole caller is delta_clonotypic_entropy, which could be rebuilt on clonotypic_entropy(point_estimate=False).", "module": "tcri/metrics/_metrics.py", "namespace": "tl"}, {"name": "clonotypic_entropy", "kind": "function", "signature": "clonotypic_entropy(adata, covariate, *, point_estimate=True, n_samples=200, temperature=1.0, combine_with_logits=True, _clones=None) -> pd.Series|np.ndarray", "purpose": "Posterior mean (or per-draw matrix) normalised clonotypic entropy per phenotype at one covariate.", "io": "reads uns['tcri_phenotype_categories']; returns Series[phenotype] or (n_samples,n_pheno) array", "calls": "joint_distribution_posterior (posterior sampling, one draw per sample)", "label": "core", "consolidate_into": "", "notes": "Listed core metric H[P(c|\u03c6,m)]; raises ValueError if n_samples<1.", "module": "tcri/metrics/_metrics.py", "namespace": "tl"}, {"name": "delta_clonotypic_entropy", "kind": "function", "signature": "delta_clonotypic_entropy(adata, phenotype, *, cov_pre='Pre-treatment', cov_post='Post-treatment', n_samples=1000, temperature=1.0, clones=None, weighted=False, normalised=True, base=2, posterior=True, combine_with_logits=True, verbose=True, graph=False, seed=None) -> np.ndarray", "purpose": "Monte-Carlo posterior samples of H_post\u2212H_pre clonotypic entropy for one phenotype.", "io": "prints progress/stats when verbose; returns (n_samples,) array", "calls": "clonotypic_entropy_base (2\u00d7 per draw, posterior sampling); _info/_ok/_ascii_hist", "label": "core", "consolidate_into": "", "notes": "Explicit core between-covariate change metric.", "module": "tcri/metrics/_metrics.py", "namespace": "tl"}, {"name": "delta_entropy_table", "kind": "function", "signature": "delta_entropy_table(adata, *, cov_pre, cov_post, splitby='response', n_samples=1000, temperature=1.0, weighted=False, normalised=True, base=2, posterior=True, combine_with_logits=True, seed=42, show_progress=True) -> pd.DataFrame", "purpose": "Tidy \u0394-clonotypic-entropy table: one row per phenotype \u00d7 splitby group with mean/sd/HDI/p and raw sample vector.", "io": "reads uns['tcri_metadata'](clone_col,phenotype_col) and obs[splitby,phenotype_col]; returns DataFrame", "calls": "delta_clonotypic_entropy (posterior sampling per phenotype/group); tqdm", "label": "core", "consolidate_into": "", "notes": "Explicit core tidy-table builder.", "module": "tcri/metrics/_metrics.py", "namespace": "tl"}, {"name": "phenotypic_entropy", "kind": "function", "signature": "phenotypic_entropy(adata, covariate, *, point_estimate=True, n_samples=200, temperature=1.0, combine_with_logits=True) -> pd.Series|np.ndarray", "purpose": "Posterior mean (or per-draw) normalised phenotypic entropy per clonotype at one covariate.", "io": "reads uns['tcri_metadata'](clone_col,covariate_col) and obs; returns Series[clone] or (n_samples,n_clones) array", "calls": "joint_distribution_posterior (posterior sampling)", "label": "core", "consolidate_into": "", "notes": "Listed core metric H[P(\u03c6|c,m)]; raises ValueError if n_samples<1; empty-clone guard returns empty Series/array.", "module": "tcri/metrics/_metrics.py", "namespace": "tl"}, {"name": "clonality", "kind": "function", "signature": "clonality(adata) -> dict", "purpose": "Per-phenotype clonality 1 \u2212 H/log2K from observed hard clone-size counts.", "io": "reads uns['tcri_phenotype_key'],uns['tcri_clone_key'] and obs (hard labels); returns {phenotype: clonality}", "calls": "scipy.stats.entropy", "label": "redundant", "consolidate_into": "clonotypic_entropy", "notes": "Hard-count complement of clonotypic entropy (overlaps clonotypic_entropy but on observed counts, no posterior). Consumed by pl.clonality (_plotting.py:661). Uses legacy uns keys tcri_phenotype_key/tcri_clone_key (written by pp:76,80) rather than the tcri_metadata convention used elsewhere \u2014 inconsistent.", "module": "tcri/metrics/_metrics.py", "namespace": "tl"}, {"name": "clone_fraction", "kind": "function", "signature": "clone_fraction(adata, groupby) -> dict", "purpose": "Per-group nested dict of each clone's frequency (count/total) within the group.", "io": "reads uns['tcri_clone_key'] and obs[groupby]; returns {group:{clone:freq}}", "calls": "-", "label": "dead-broken", "consolidate_into": "", "notes": "0 callers anywhere in tcri. Thin, slow reimplementation of obs.groupby(groupby)[clone].value_counts(normalize=True); also uses legacy uns['tcri_clone_key']. Unused with dubious value.", "module": "tcri/metrics/_metrics.py", "namespace": "tl"}, {"name": "mutual_information", "kind": "function", "signature": "mutual_information(adata, covariate, *, temperature=1.0, n_samples=0, clones=None, normalised=True, normalise_mode='average', posterior=True, combine_with_logits=True, verbose=True, graph=False) -> float|np.ndarray", "purpose": "Clone\u00d7phenotype mutual information at one covariate: point estimate (n_samples=0) or per-draw array.", "io": "prints progress/stats when verbose; returns float or (n_samples,) array", "calls": "joint_distribution_posterior via inner _get_df (posterior sampling); _mi_from_joint; _ok/_info/_ascii_hist", "label": "core", "consolidate_into": "", "notes": "Listed core metric I(c;\u03c6|m). NOTE: posterior=False path is broken \u2014 inner _get_df raises NotImplementedError('prior-only joint_distribution not included'), so mutual_information(posterior=False) always errors.", "module": "tcri/metrics/_metrics.py", "namespace": "tl"}, {"name": "mutual_information._get_df", "kind": "inner-function", "signature": "_get_df(silent_flag:bool)", "purpose": "Return one joint-distribution table for the MI computation (posterior draw).", "io": "-", "calls": "joint_distribution_posterior (posterior sampling); raises NotImplementedError when posterior is False", "calls_note": "", "label": "helper", "consolidate_into": "", "notes": "Inner plumbing of core mutual_information; wraps a single posterior draw. Its else-branch raises NotImplementedError, making the parent's posterior=False mode non-functional.", "module": "tcri/metrics/_metrics.py", "namespace": "tl"}, {"name": "flux_table", "kind": "function", "signature": "flux_table(adata, *, cov_pre, cov_post, splitby='response', n_samples=0, temperature=1.0, weighted=False, posterior=True, combine_with_logits=True, distance_metric='l1', seed=42, show_progress=True) -> pd.DataFrame", "purpose": "Tidy per-clone flux table: flux mean/sd + sample vector + clone size, per splitby group.", "io": "reads uns['tcri_metadata'](clone_col) and obs[splitby,clone_col]; returns DataFrame", "calls": "flux (posterior sampling per group); tqdm (nested bars)", "label": "core", "consolidate_into": "", "notes": "Explicit core tidy-table builder.", "module": "tcri/metrics/_metrics.py", "namespace": "tl"}, {"name": "flux", "kind": "function", "signature": "flux(adata, *, from_this, to_that, clones=None, temperature=1.0, distance_metric='l1', n_samples=0, weighted=False, posterior=True, combine_with_logits=True, graph=False, seed=42) -> pd.Series|np.ndarray", "purpose": "Per-clone phenotypic-distribution distance (l1/dkl/callable) between two covariates; point estimate or per-draw.", "io": "reads uns['tcri_metadata'](clone_col) and obs; returns Series[clone] or (n_samples,n_clones) array", "calls": "joint_distribution_posterior / joint_distribution (posterior sampling); inner dkl_func; _ascii_hist", "label": "core", "consolidate_into": "", "notes": "Listed core metric (phenotypic flux). Raises ValueError if no clone overlap at both covariates.", "module": "tcri/metrics/_metrics.py", "namespace": "tl"}, {"name": "flux.dkl_func", "kind": "inner-function", "signature": "dkl_func(p, q)", "purpose": "KL divergence kernel used when distance_metric=='dkl' inside flux.", "io": "-", "calls": "-", "label": "helper", "consolidate_into": "", "notes": "Inner distance kernel of core flux; duplicates the dead module-level dkl. Could be replaced by a shared distance helper.", "module": "tcri/metrics/_metrics.py", "namespace": "tl"}, {"name": "_ok", "kind": "function", "signature": "_ok(msg:str, quiet=False)", "purpose": "Print a green success line unless quiet.", "label": "helper", "io": "side effect: stdout (ANSI)", "calls": "-", "consolidate_into": "", "notes": "Low-level ANSI console printing; belongs in shared console-helper module. Used by register_model/joint_distribution_posterior.", "module": "tcri/preprocessing/_preprocessing.py", "namespace": "pp"}, {"name": "_info", "kind": "function", "signature": "_info(key:str, txt:str, quiet=False)", "purpose": "Print a dim key-value info line unless quiet.", "label": "helper", "io": "side effect: stdout (ANSI)", "calls": "-", "consolidate_into": "", "notes": "ANSI console printing helper.", "module": "tcri/preprocessing/_preprocessing.py", "namespace": "pp"}, {"name": "_warn", "kind": "function", "signature": "_warn(msg:str, quiet=False)", "purpose": "Print a yellow warning line unless quiet.", "label": "helper", "io": "side effect: stdout (ANSI)", "calls": "-", "consolidate_into": "", "notes": "ANSI console printing helper; no in-module caller but part of the console helper set.", "module": "tcri/preprocessing/_preprocessing.py", "namespace": "pp"}, {"name": "_fin", "kind": "function", "signature": "_fin(quiet=False)", "purpose": "Print a magenta 'Done!' flourish unless quiet.", "label": "helper", "io": "side effect: stdout (ANSI)", "calls": "-", "consolidate_into": "", "notes": "ANSI console printing helper.", "module": "tcri/preprocessing/_preprocessing.py", "namespace": "pp"}, {"name": "_ascii_hist", "kind": "function", "signature": "_ascii_hist(samples, bins=25, width=40) -> str", "purpose": "Render a numpy histogram of samples as an ASCII bar chart string.", "label": "helper", "io": "-", "calls": "np.histogram", "consolidate_into": "", "notes": "Exact duplicate of tcri/metrics/_metrics.py:62 _ascii_hist (that copy is the one actually used); 0 callers within this module. Consolidate the two copies into one shared helper.", "module": "tcri/preprocessing/_preprocessing.py", "namespace": "pp"}, {"name": "register_phenotype_key", "kind": "function", "signature": "register_phenotype_key(adata, phenotype_key, order=None)", "purpose": "Register the phenotype obs column and its unique categories into uns.", "label": "helper", "io": "reads obs[phenotype_key]; writes uns['tcri_unique_phenotypes'], uns['tcri_phenotype_key']", "calls": "-", "consolidate_into": "", "notes": "Low-level uns registration setter called at the end of register_model. Minor quirk: when order!=None, tcri_unique_phenotypes is never set (order param otherwise unused).", "module": "tcri/preprocessing/_preprocessing.py", "namespace": "pp"}, {"name": "register_clonotype_key", "kind": "function", "signature": "register_clonotype_key(adata, tcr_key)", "purpose": "Register the clonotype obs column and its unique categories into uns.", "label": "helper", "io": "reads obs[tcr_key]; writes uns['tcri_clone_key'], uns['tcri_unique_clonotypes']", "calls": "-", "consolidate_into": "", "notes": "Low-level uns registration setter called by register_model; tcri_clone_key is consumed by clone_size.", "module": "tcri/preprocessing/_preprocessing.py", "namespace": "pp"}, {"name": "group_singletons", "kind": "function", "signature": "group_singletons(adata, clonotype_key='trb', groupby='patient', target_col='trb_unique', min_clone_size=10)", "purpose": "Collapse clones smaller than min_clone_size (per groupby) into 'Singleton_{group}' labels.", "label": "helper", "io": "reads obs[clonotype_key], obs[groupby]; writes obs['trb_candidate'], obs[target_col]", "calls": "collapse_singleton (inner)", "consolidate_into": "", "notes": "Clone-collapse preprocessing util; 0 external callers; overlaps group_small_clones (same purpose). Hardcodes intermediate col 'trb_candidate'. Consolidate with group_small_clones into one parametrized util.", "module": "tcri/preprocessing/_preprocessing.py", "namespace": "pp"}, {"name": "group_singletons.collapse_singleton", "kind": "inner-function", "signature": "collapse_singleton(row)", "purpose": "Map a row to 'Singleton_{group}' when its candidate clone count < min_clone_size, else keep candidate.", "label": "helper", "io": "reads row['trb_candidate'], row[groupby]; closes over clone_counts, min_clone_size, groupby", "calls": "-", "consolidate_into": "", "notes": "Row-apply closure inside group_singletons.", "module": "tcri/preprocessing/_preprocessing.py", "namespace": "pp"}, {"name": "classify_phenotypes", "kind": "function", "signature": "classify_phenotypes(adata, phenotype_prob_slot='X_tcri_phenotypes', phenotype_assignment_obs='tcri_phenotype')", "purpose": "Alternate per-cell phenotype assignment via cosine similarity of latent to per-phenotype archetypes, reweighted by posterior p_ct.", "label": "redundant", "io": "reads uns['tcri_metadata'], uns['tcri_ct_array_for_cells'], uns['tcri_p_ct'], uns['tcri_phenotype_categories'], obsm['X_tcri'], obs[phenotype_col]; writes obsm[phenotype_prob_slot], obs[phenotype_assignment_obs]", "calls": "sklearn cosine_similarity; uses stored p_ct (no live posterior sampling)", "consolidate_into": "register_model", "notes": "0 external callers. Writes the same phenotype-probability/hard-label slots register_model produces, via a different archetype-cosine algorithm instead of softmax(logits+prior). Overlaps register_model's phenotype assignment.", "module": "tcri/preprocessing/_preprocessing.py", "namespace": "pp"}, {"name": "_compute_logits_and_prior", "kind": "function", "signature": "@torch.no_grad _compute_logits_and_prior(model, adata, batch_size=256, eps=1e-8) -> (np.float32, np.float32)", "purpose": "Run the model encoder+classifier over a data loader to extract per-cell classifier logits and the log-prior from get_p_ct.", "label": "helper", "io": "reads model data loader / module state; returns arrays (no adata writes)", "calls": "HITS MODEL: model._make_data_loader, model.module.encoder/classifier/get_p_ct, model.module.ct_array; REGISTRY_KEYS", "consolidate_into": "", "notes": "Private extraction helper called only by register_model; inference plumbing (not model construction).", "module": "tcri/preprocessing/_preprocessing.py", "namespace": "pp"}, {"name": "register_model", "kind": "function", "signature": "@torch.no_grad register_model(adata, model, phenotype_prob_slot='X_tcri_probabilities', phenotype_assignment_obs='tcri_phenotype', latent_slot='X_tcri', batch_size=256, store_logits=True, store_logposterior=True, compute_umap=False, umap_*=..., clonotype_key='trb_unique')", "purpose": "Register all TCRIModel outputs (priors, metadata, categories, per-cell ct/cov arrays, latent means, logits, log-posterior, phenotype probs+labels, optional UMAP) onto the AnnData.", "label": "core", "io": "writes uns[tcri_p_ct, tcri_ct_to_cov, tcri_ct_to_c, tcri_local_scale, tcri_metadata, tcri_{covariate,clonotype,phenotype}_categories, tcri_ct_array_for_cells, tcri_cov_array_for_cells]; obsm[X_tcri, X_tcri_logits, X_tcri_logposterior, phenotype_prob_slot, X_umap]; obs[phenotype_assignment_obs]", "calls": "HITS MODEL: model.module.get_p_ct/ct_to_cov/ct_to_c/local_scale/ct_array, model.adata_manager.registry, model.get_latent_representation; _compute_logits_and_prior; register_phenotype_key; register_clonotype_key; umap.UMAP; scipy softmax", "consolidate_into": "", "notes": "CORE MODEL step: the documented pp.register_model entrypoint that materializes the substrate every metric/engine reads. Referenced throughout metrics docstrings.", "module": "tcri/preprocessing/_preprocessing.py", "namespace": "pp"}, {"name": "joint_distribution_posterior", "kind": "function", "signature": "joint_distribution_posterior(adata, covariate_label, *, temperature=1.0, clones=None, weighted=False, combine_with_logits=True, precision=3, silent=False) -> pd.DataFrame", "purpose": "Draw one posterior Dirichlet sample of p_ct, combine with per-cell classifier logits, and aggregate to a clone x phenotype distribution for a covariate.", "label": "core", "io": "reads uns[tcri_metadata, tcri_phenotype_categories, tcri_covariate_categories, tcri_ct_array_for_cells, tcri_cov_array_for_cells, tcri_p_ct, tcri_local_scale], obsm['X_tcri_logits'], obs[clone_col]; returns DataFrame (no writes)", "calls": "POSTERIOR SAMPLING: Dirichlet(local_scale*p_ct).sample(); scipy softmax; _ok/_info/_fin; guards n_obs vs per-cell arrays", "consolidate_into": "", "notes": "CORE ENGINE: the Bayesian posterior-sampling clone x phenotype substrate that essentially every metric in _metrics.py calls (posterior=True path). Includes the Notion #4 filtered-AnnData misalignment guard.", "module": "tcri/preprocessing/_preprocessing.py", "namespace": "pp"}, {"name": "remove_meaningless_genes", "kind": "function", "signature": "remove_meaningless_genes(adata, include_mt=True, include_rp=True, include_mtrn=True, include_hsp=True, include_tcr=True) -> AnnData", "purpose": "Intend to filter out MT/RP/HSP/MTRN/TCR/RIK/GM/LINC/ambiguous genes and return a sliced copy.", "label": "dead-broken", "io": "reads var.index; returns adata[:,genes].copy()", "calls": "-", "consolidate_into": "", "notes": "BROKEN: the `if include_mtrn:` and `if include_hsp:` branches reassign `genes` from the FULL adata.var.index (not the running filtered list), silently discarding all prior RIK/GM/dash/dot/LINC filters (and each other). With defaults, only HSP/MT-/RP/TCR filtering survives. 0 external callers.", "module": "tcri/preprocessing/_preprocessing.py", "namespace": "pp"}, {"name": "joint_distribution", "kind": "function", "signature": "joint_distribution(adata, covariate_label:str, temperature=1.0, n_samples=0, clones=None, weighted=False) -> pd.DataFrame", "purpose": "Build a clone x phenotype distribution for a covariate from ct-level p_ct point estimates (n_samples=0) or Dirichlet draws (n_samples>0), with optional clone-size weighting.", "label": "redundant", "io": "reads uns[tcri_p_ct, tcri_ct_to_cov, tcri_ct_to_c, tcri_covariate_categories, tcri_phenotype_categories, tcri_clonotype_categories, tcri_metadata, tcri_ct_array_for_cells, tcri_cov_array_for_cells, tcri_local_scale]; returns DataFrame", "calls": "Dirichlet(conc).sample (when n_samples>0); torch softmax", "consolidate_into": "joint_distribution_posterior", "notes": "Point-estimate/ct-level alternative to the posterior engine producing the same covariate->clone x phenotype output but WITHOUT per-cell logits. Still live (metrics posterior=False branch + plotting compare_joint_distribution/polar_plot), but overlaps joint_distribution_posterior; unify into one engine with a posterior/n_samples switch.", "module": "tcri/preprocessing/_preprocessing.py", "namespace": "pp"}, {"name": "get_latent_embedding", "kind": "function", "signature": "get_latent_embedding(adata, latent_slot='X_tcri', n_samples=0, posterior_scale=1.0) -> np.ndarray", "purpose": "Draw Gaussian samples around the stored latent means with a scalar std.", "label": "dead-broken", "io": "reads obsm[latent_slot]; returns array", "calls": "np.random.normal", "consolidate_into": "", "notes": "0 callers; dubious value: fabricates a spherical Gaussian with a constant posterior_scale rather than the model's true per-cell latent variance, and the default n_samples=0 yields an empty (0,n_cells,dim) array.", "module": "tcri/preprocessing/_preprocessing.py", "namespace": "pp"}, {"name": "group_small_clones", "kind": "function", "signature": "group_small_clones(adata, patient_key='')", "purpose": "Collapse clones with clone_size<4 into 'Singleton_{patient}', else '{trb}_{patient}', into obs['trb_unique'].", "label": "dead-broken", "io": "reads obs['trb'], obs['clone_size'], obs[patient_key]; writes obs['trb_unique']", "calls": "-", "consolidate_into": "", "notes": "0 callers; dataset-specific: hardcodes obs['trb'] and requires a precomputed obs['clone_size']. Duplicates group_singletons (same collapse-small-clones purpose). Remove or fold into group_singletons.", "module": "tcri/preprocessing/_preprocessing.py", "namespace": "pp"}, {"name": "register_probability_columns", "kind": "function", "signature": "register_probability_columns(adata, probability_columns)", "purpose": "Store a probability_columns list into uns.", "label": "dead-broken", "io": "writes uns['probability_columns']", "calls": "-", "consolidate_into": "", "notes": "0 callers; trivial one-line setter whose uns['probability_columns'] key is never read anywhere in the package.", "module": "tcri/preprocessing/_preprocessing.py", "namespace": "pp"}, {"name": "gene_entropy", "kind": "function", "signature": "gene_entropy(adata, key_added='entropy', batch_key=None, agg_function=None)", "purpose": "Compute per-gene Shannon entropy of expression-value counts (optionally per batch, aggregated) into var.", "label": "dead-broken", "io": "reads adata.X (dense), obs[batch_key]; writes var[key_added]", "calls": "scipy.stats.entropy", "consolidate_into": "", "notes": "0 callers; orphaned feature unrelated to core clonotypic/phenotypic entropy (this is per-GENE expression entropy). Drops the first count bin (counts[1][1:]), assuming index 0 is the zero-expression bin.", "module": "tcri/preprocessing/_preprocessing.py", "namespace": "pp"}, {"name": "clone_size", "kind": "function", "signature": "clone_size(adata, key_added='clone_size', return_counts=False)", "purpose": "Compute per-clone cell counts from the registered clone key and write them per cell into obs.", "label": "helper", "io": "reads uns['tcri_clone_key'], obs[tcr_key]; writes obs[key_added]; optionally returns dict", "calls": "-", "consolidate_into": "", "notes": "Low-level obs utility; used by tcri/plotting/_plotting.py clone_size_umap. Depends on register_clonotype_key having set tcri_clone_key.", "module": "tcri/preprocessing/_preprocessing.py", "namespace": "pp"}, {"name": "_ok", "kind": "function", "signature": "(msg:str, quiet=False)", "purpose": "Print a green checkmark success line to stdout.", "label": "helper", "io": "stdout print", "calls": "-", "consolidate_into": "", "notes": "Console-printing helper; exact duplicate of tcri.utils._utils._ok (redefined here). Should import from utils instead of re-declaring.", "module": "tcri/plotting/_plotting.py", "namespace": "pl"}, {"name": "_info", "kind": "function", "signature": "(key:str, txt:str, quiet=False)", "purpose": "Print a cyan key-value info line to stdout.", "label": "helper", "io": "stdout print", "calls": "-", "consolidate_into": "", "notes": "Console-printing helper; duplicate of tcri.utils._utils._info. Used by bayesian_mutual_information.", "module": "tcri/plotting/_plotting.py", "namespace": "pl"}, {"name": "_warn", "kind": "function", "signature": "(msg:str, quiet=False)", "purpose": "Print a yellow warning line to stdout.", "label": "helper", "io": "stdout print", "calls": "-", "consolidate_into": "", "notes": "Console-printing helper; duplicate of tcri.utils._utils._warn. Used by clonotypic_entropy_by_phenotype.", "module": "tcri/plotting/_plotting.py", "namespace": "pl"}, {"name": "_fin", "kind": "function", "signature": "(quiet=False)", "purpose": "Print a magenta 'Done!' final flourish to stdout.", "label": "helper", "io": "stdout print", "calls": "-", "consolidate_into": "", "notes": "Console-printing helper; duplicate of tcri.utils._utils._fin. 0 callers in this file (unused here).", "module": "tcri/plotting/_plotting.py", "namespace": "pl"}, {"name": "mi_compare", "kind": "function", "signature": "(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)", "purpose": "Plot per-patient TCRi normalized MI across covariate group pairs as boxplots with jittered points and AUROC/MWU/label-permutation stats.", "label": "core", "io": "reads via tl (uns['tcri_metadata']); fig.savefig(save)", "calls": "mi_compare_tl (tl, posterior MI sampling); auc_and_label_permutation, bootstrap_auc (utils)", "consolidate_into": "", "notes": "pl visualization of the core MI-comparison tidy table; call signature to mi_compare_tl and utils stats verified correct and functional.", "module": "tcri/plotting/_plotting.py", "namespace": "pl"}, {"name": "compare_phenotypes", "kind": "function", "signature": "(adata, variable1, variable2)", "purpose": "Heatmap of the row-normalized crosstab between two arbitrary obs columns.", "label": "plotting-beyond-core", "io": "reads adata.obs[variable1,variable2]", "calls": "-", "consolidate_into": "", "notes": "Generic categorical-crosstab heatmap; does not visualize any core metric or the joint distribution.", "module": "tcri/plotting/_plotting.py", "namespace": "pl"}, {"name": "compare_joint_distribution", "kind": "function", "signature": "(adata, temperature=1)", "purpose": "Side-by-side clustered heatmaps/dendrograms of model-inferred vs empirical joint distributions per covariate.", "label": "dead-broken", "io": "reads uns['tcri_metadata'], obs; plt.show()", "calls": "joint_distribution (engine)", "consolidate_into": "", "notes": "References undefined global `model` at L209-210 (`model.adata_manager.registry[...]`) \u2014 never defined or imported \u2192 NameError at runtime.", "module": "tcri/plotting/_plotting.py", "namespace": "pl"}, {"name": "_phenotype_mass_per_clone", "kind": "function", "signature": "(adata, covariate, clones, normalize)", "purpose": "Return {clone_id -> phenotype-mass vector} at one covariate by summing joint_distribution rows, optionally weighting each row by the clone's cell count.", "label": "helper", "io": "reads uns['tcri_phenotype_categories'], uns['tcri_metadata'], obs value_counts", "calls": "joint_distribution (engine / posterior substrate)", "consolidate_into": "", "notes": "Private extraction helper feeding the sankey/flux plots; low-level data-prep utility.", "module": "tcri/plotting/_plotting.py", "namespace": "pl"}, {"name": "plot_pheno_sankey", "kind": "function", "signature": "(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)", "purpose": "Draw a Sankey of phenotype-distribution flow across ordered covariate values using per-clone outer-product flow geometry.", "label": "core", "io": "reads uns['tcri_phenotype_categories']; draws/returns fig,ax", "calls": "_phenotype_mass_per_clone -> joint_distribution (engine); SankeyNode", "consolidate_into": "", "notes": "Core flux/joint-distribution visualization engine; SankeyNode attributes (.patch/.max_y/.plot/.plot_node_connection) all present, functional.", "module": "tcri/plotting/_plotting.py", "namespace": "pl"}, {"name": "phenotypic_flux", "kind": "function", "signature": "(adata, splitby, order, clones=None, normalize=True, phenotype_colors=None, save=None, figsize=(6,3), show_legend=True, title=None)", "purpose": "Public flux Sankey across `order` values of `splitby`; thin convenience wrapper adding x-ticks and save.", "label": "core", "io": "fig.savefig(save)", "calls": "plot_pheno_sankey", "consolidate_into": "", "notes": "Thin wrapper over plot_pheno_sankey exposing the flux-sankey as the public pl API; visualizes core flux.", "module": "tcri/plotting/_plotting.py", "namespace": "pl"}, {"name": "probability_distribution", "kind": "function", "signature": "(adata, phenotype_order=None, color='#000000', rotation=90, splitby=None, order=None, figsize=(7,5), save=None)", "purpose": "Intended: barplots of the phenotype probability distribution per split value.", "label": "dead-broken", "io": "fig.savefig(save)", "calls": "itself (recursion)", "consolidate_into": "", "notes": "Infinite self-recursion: L538 calls `probability_distribution(zdata)` (almost certainly meant utils.probabilities). Also `ax[i]` fails when splitby=None (ncols==1 -> single Axes) and `order=sorted(adata.obs[None])` raises. Non-functional.", "module": "tcri/plotting/_plotting.py", "namespace": "pl"}, {"name": "top_clone_umap", "kind": "function", "signature": "(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)", "purpose": "UMAP scatter highlighting the top-N largest clones over a grey background.", "label": "plotting-beyond-core", "io": "reads uns['tcri_clone_key'], obsm['X_umap']; plt.savefig; optional return df", "calls": "-", "consolidate_into": "", "notes": "Bespoke UMAP overlay with hardcoded title 'Top 10 TCR Clone by Size'; not a core metric/joint-distribution plot.", "module": "tcri/plotting/_plotting.py", "namespace": "pl"}, {"name": "tcri_boxplot", "kind": "function", "signature": "(adata, function, groupby=None, ylabel='', splitby=None, figsize=(8,4), s=20, order=None, palette=None)", "purpose": "Generic per-phenotype metric boxplot/stripplot engine that applies a metric `function` across groupby/splitby strata.", "label": "core", "io": "draws fig", "calls": "the passed metric fn (e.g. clonality_tl)", "consolidate_into": "", "notes": "Shared plotting primitive that directly renders core per-phenotype metrics; currently the engine behind clonality.", "module": "tcri/plotting/_plotting.py", "namespace": "pl"}, {"name": "clonality", "kind": "function", "signature": "(adata, groupby=None, splitby=None, s=10, order=None, figsize=(12,5), palette=None)", "purpose": "Boxplot of the clonality metric via tcri_boxplot.", "label": "core", "io": "draws fig", "calls": "tcri_boxplot; clonality_tl (tl)", "consolidate_into": "", "notes": "Directly visualizes the core clonality metric.", "module": "tcri/plotting/_plotting.py", "namespace": "pl"}, {"name": "clonotypic_entropy_by_phenotype", "kind": "function", "signature": "(adata, *, temperature=1.0, n_samples=200, combine_with_logits=True, bayesian=True, bayes_samples=1000, palette=None, group_colors=None, hue_order=None, legend_fontsize=6, bbox_to_anchor=(1.15,1.), figsize=(6,3), rotation=90, save=None, return_df=False, progress=True)", "purpose": "Box-and-dot plot of clonotypic entropy per phenotype and covariate, with bootstrap/MWU significance brackets and per-patient dots.", "label": "core", "io": "reads uns['tcri_metadata']; fig.savefig(save); optional return df", "calls": "centropy (tl clonotypic_entropy, posterior sampling); _ok, _warn", "consolidate_into": "", "notes": "centropy call (covariate=, point_estimate=, n_samples=, temperature=, combine_with_logits=, _clones=) matches metrics signature exactly; functional core plot of clonotypic entropy.", "module": "tcri/plotting/_plotting.py", "namespace": "pl"}, {"name": "plot_phenotype_probabilities", "kind": "function", "signature": "(adata, phenotype_prob_slot='X_tcri_phenotypes', add_outline=False, save=None, ncols=2, cmap='magma')", "purpose": "UMAP panels colored by each per-cell phenotype probability.", "label": "plotting-beyond-core", "io": "reads obsm['X_tcri_phenotypes'], uns['tcri_phenotype_categories']; copies adata, writes temp obs '_probability' cols; plt.savefig", "calls": "sc.pl.umap", "consolidate_into": "", "notes": "UMAP overlay of model phenotype probabilities; not a core metric or the joint distribution itself.", "module": "tcri/plotting/_plotting.py", "namespace": "pl"}, {"name": "clone_size_umap", "kind": "function", "signature": "(adata, reduction='umap', figsize=(10,8), size=1, alpha=0.7, palette='coolwarm', save=None)", "purpose": "UMAP scatter colored by log10 clone size.", "label": "plotting-beyond-core", "io": "clone_size(adata) writes obs['clone_size']; writes obs 'UMAP1'/'UMAP2'/'log(Clone Size)'; reads obsm['X_umap']; fig.savefig", "calls": "clone_size (preprocessing)", "consolidate_into": "", "notes": "Bespoke UMAP overlay; mutates adata.obs as a side effect; not a core metric plot.", "module": "tcri/plotting/_plotting.py", "namespace": "pl"}, {"name": "ridge_delta_entropy", "kind": "function", "signature": "(df_delta:pd.DataFrame, *, splitby='complete_response', order_group=None, order_phen=None, palette=None, bw_adjust=0.8, jitter=0.15, density_scale=0.9, significance=True, sig_test='mannwhitney', bayes_iters=5000, bracket_pad=0.15, star_size=16, figsize=(10,6), ax=None)", "purpose": "Ridge/joyplot of \u0394-clonotypic-entropy posteriors per phenotype with significance brackets on the first two groups.", "label": "dead-broken", "io": "consumes precomputed df_delta; draws/returns fig,ax", "calls": "-", "consolidate_into": "", "notes": "References undefined names: `cm` (L888 cm.tab10.colors) and `st` (L915 st.gaussian_kde, L936 st.mannwhitneyu). Neither matplotlib.cm nor scipy.stats is imported under those aliases \u2192 NameError.", "module": "tcri/plotting/_plotting.py", "namespace": "pl"}, {"name": "phenotypic_entropy", "kind": "function", "signature": "(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)", "purpose": "Box/strip plot of phenotypic entropy per covariate and batch (intended core entropy plot).", "label": "dead-broken", "io": "reads uns['tcri_metadata']; fig.savefig(save)", "calls": "pentropy (tl phenotypic_entropy, posterior)", "consolidate_into": "", "notes": "Wrong call signature: calls pentropy(subt, t, ph, temperature=, clones=vclones, n_samples=, normalized=) but metrics.phenotypic_entropy is (adata, covariate, *, point_estimate, n_samples, temperature, combine_with_logits) \u2014 `ph` is an illegal 3rd positional (rest are keyword-only) and `clones=`/`normalized=` are unknown kwargs \u2192 TypeError. Also splitby branch's `subts = subt[subt.obs[cov_col]==t]` ignores the split value.", "module": "tcri/plotting/_plotting.py", "namespace": "pl"}, {"name": "set_color_palette", "kind": "function", "signature": "(adata, columns)", "purpose": "Assign tcri_colors to the categories of each obs column and store them in uns['_colors']; return category->color map.", "label": "helper", "io": "writes uns['_colors'] on adata.copy() (lost); returns dict", "calls": "-", "consolidate_into": "", "notes": "Palette utility. Operates on adata.copy() so the uns writes never reach the caller; `colors.pop(i)` indexed by enumerate position is fragile/order-dependent.", "module": "tcri/plotting/_plotting.py", "namespace": "pl"}, {"name": "flux", "kind": "function", "signature": "(adata, key, order, groupby, paint_dict=None, method='probabilistic', paint=None, distance_metric='l1', figsize=(12,5), paint_order=None, palette=None)", "purpose": "Boxplot of per-clone flux (flux_tl distance) between consecutive `order` values, grouped by `groupby` and colored by a `paint` obs category.", "label": "plotting-beyond-core", "io": "reads obs[groupby], obs[paint], uns['_colors']; leftover debug print(palette)", "calls": "flux_tl (tl)", "consolidate_into": "", "notes": "Dataset-specific `paint`-overlay boxplot. Bug: on the default `paint=None` path, `pcat` at L1075 (`df['Comparison']=pcat`) is referenced before assignment \u2192 NameError; only works when paint is supplied. `method` and `paint_order` params are unused; leftover print(palette).", "module": "tcri/plotting/_plotting.py", "namespace": "pl"}, {"name": "mutual_information", "kind": "function", "signature": "(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)", "purpose": "Box/strip plot of clonotype<->phenotype mutual information per covariate and batch (intended core MI plot).", "label": "dead-broken", "io": "reads uns['tcri_metadata']; plt.savefig(save)", "calls": "mutual_information_tl (tl, posterior)", "consolidate_into": "", "notes": "Passes `weighted=weighted` (L1170, L1186) to metrics.mutual_information, whose signature is (adata, covariate, *, temperature, n_samples, clones, normalised, normalise_mode, posterior, combine_with_logits, verbose, graph) \u2014 no `weighted` param \u2192 TypeError. Intended core MI plot but non-functional as written.", "module": "tcri/plotting/_plotting.py", "namespace": "pl"}, {"name": "bayesian_mutual_information", "kind": "function", "signature": "(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)", "purpose": "3-panel Bayesian \u0394-MI (post-pre) analysis across splitby strata: \u0394-MI KDEs, per-condition pre/post KDEs, and \u0394 bar summary with HDI/P(>0).", "label": "dead-broken", "io": "prints ANSI summary; draws 3-panel fig; returns results dict", "calls": "mutual_information_tl (tl, posterior sampling); _info", "consolidate_into": "", "notes": "Passes `weighted=weighted` (L1279, L1284) to metrics.mutual_information which has no `weighted` param \u2192 TypeError. (Other kwargs normalised/normalise_mode/posterior/combine_with_logits/verbose are valid; only `weighted` is unsupported.)", "module": "tcri/plotting/_plotting.py", "namespace": "pl"}, {"name": "polar_plot", "kind": "function", "signature": "(adata, phenotypes=None, statistic='distribution', method='joint_distribution', splitby=None, color_dict=None, temperature=1.0)", "purpose": "Radar/polar plot of per-phenotype distribution or entropy across split values.", "label": "dead-broken", "io": "reads uns['tcri_metadata']; joint_distribution", "calls": "joint_distribution (engine); clonotypic_entropy (undefined)", "consolidate_into": "", "notes": "Two defects: (1) `phenotypes` defaults to uns['tcri_metadata']['phenotype_col'], a column-NAME string not a list, so len()/iteration run over characters and angles length mismatches values; (2) the entropy branch (L1420) calls bare `clonotypic_entropy(...)` which is not imported (only imported as `centropy`) \u2192 NameError.", "module": "tcri/plotting/_plotting.py", "namespace": "pl"}, {"name": "SankeyNode", "kind": "class", "signature": "SankeyNode(object)", "purpose": "Drawing primitive representing one rectangular sankey node and the ribbons flowing out of it, used to render the phenotypic-flux sankey.", "io": "-", "calls": "consumed by pl.plot_pheno_sankey in _plotting.py; no model/posterior", "label": "helper", "consolidate_into": "", "notes": "Low-level matplotlib drawing primitive (node rectangle + connection ribbon) supporting the core phenotypic_flux/sankey plot; belongs in a shared plotting-helper module. Not a metric plot itself.", "module": "tcri/plotting/_sankey.py", "namespace": "pl"}, {"name": "SankeyNode.__init__", "kind": "method", "signature": "__init__(self, x, y, val, dx=0.2, color=None, **kwargs)", "purpose": "Compute node bounding box (min/max x,y from center x, base y, width dx, height val) and build the mpatches.Rectangle patch.", "io": "sets self.x/y/dx/dy/x_gap/max_*/min_*/color/patch; builds matplotlib Rectangle (no adata)", "calls": "mpatches.Rectangle; no intra-tcri deps", "label": "helper", "consolidate_into": "", "notes": "Non-trivial __init__ (geometry math + patch construction); part of the sankey drawing-primitive helper.", "module": "tcri/plotting/_sankey.py", "namespace": "pl"}, {"name": "SankeyNode.plot", "kind": "method", "signature": "plot(self, ax)", "purpose": "Render the node's rectangle patch onto the given matplotlib axis.", "io": "side effect: ax.add_patch(self.patch)", "calls": "ax.add_patch; called from plot_pheno_sankey:447", "label": "helper", "consolidate_into": "", "notes": "Trivial one-liner drawing helper; kept because it carries the render side effect for the sankey primitive.", "module": "tcri/plotting/_sankey.py", "namespace": "pl"}, {"name": "SankeyNode.hex_to_rgb", "kind": "method", "signature": "hex_to_rgb(self, hex_color)", "purpose": "Parse a #RRGGBB hex string into a normalized (r,g,b) float tuple in [0,1].", "io": "-", "calls": "-", "label": "dead-broken", "consolidate_into": "", "notes": "0 callers anywhere in package (grep only hits its own def); plot_node_connection uses matplotlib mcolors.to_rgb instead. Repo contract doc (docs/contract/build_tcri_contract.py:373) marks SankeyNode.hex_to_rgb 'deleted, unused method'.", "module": "tcri/plotting/_sankey.py", "namespace": "pl"}, {"name": "SankeyNode.plot_node_connection", "kind": "method", "signature": "plot_node_connection(self, destination_node, ax, **kwargs)", "purpose": "Draw the curved, color-interpolated ribbon (500 fill_between segments, sigmoid-shaped top/bottom edges) from this node to a destination node.", "io": "side effect: many ax.fill_between calls on the axis", "calls": "np.linspace/np.power, mcolors.to_rgb, ax.fill_between; called from plot_pheno_sankey:451; no model/posterior", "label": "helper", "consolidate_into": "", "notes": "Core sankey ribbon-drawing routine for the phenotypic-flux plot; reusable low-level matplotlib drawing helper. Note the **kwargs is accepted but ignored.", "module": "tcri/plotting/_sankey.py", "namespace": "pl"}, {"name": "build_archetypes", "kind": "function", "signature": "build_archetypes(c2p_mat, K=4)", "purpose": "KMeans-cluster clone->phenotype rows into K normalized archetype centroids used as the Dirichlet-mixture prior concentration.", "io": "-", "calls": "sklearn KMeans; consumed by TCRIModel.__init__ to build mixture_concentration", "label": "model-construction", "consolidate_into": "", "notes": "Archetype builder for the model prior; per label def archetype builders are model-construction. Pure numpy/sklearn.", "module": "tcri/model/_model.py", "namespace": "ml"}, {"name": "PhenotypeClassifier", "kind": "class", "signature": "class PhenotypeClassifier(nn.Module)", "purpose": "Temperature-scaled MLP head mapping latent z to phenotype logits.", "io": "-", "calls": "-", "label": "model-construction", "consolidate_into": "", "notes": "nn.Module used inside TCRIModule as self.classifier; model internal.", "module": "tcri/model/_model.py", "namespace": "ml"}, {"name": "PhenotypeClassifier.__init__", "kind": "method", "signature": "__init__(self, n_latent, classifier_hidden, P, num_layers=3, dropout_rate=0.1, temperature=1.0)", "purpose": "Build the stacked Linear/ReLU/Dropout MLP and store the softmax temperature.", "io": "-", "calls": "-", "label": "model-construction", "consolidate_into": "", "notes": "Carries layer-building logic, so not a trivial __init__.", "module": "tcri/model/_model.py", "namespace": "ml"}, {"name": "PhenotypeClassifier.forward", "kind": "method", "signature": "forward(self, x)", "purpose": "Return MLP logits divided by temperature.", "io": "-", "calls": "-", "label": "model-construction", "consolidate_into": "", "notes": "Classifier forward pass; model internal.", "module": "tcri/model/_model.py", "namespace": "ml"}, {"name": "VampPrior", "kind": "class", "signature": "class VampPrior(torch.nn.Module)", "purpose": "VampPrior over latent z: a uniform mixture of encoder posteriors evaluated at learnable pseudo-inputs.", "io": "-", "calls": "shares the module encoder (model forward)", "label": "model-construction", "consolidate_into": "", "notes": "Prior module used as p(z) in TCRIModule.model; model internal.", "module": "tcri/model/_model.py", "namespace": "ml"}, {"name": "VampPrior.__init__", "kind": "method", "signature": "__init__(self, pseudo_inputs, encoder)", "purpose": "Register learnable pseudo-inputs as a Parameter and hold the shared encoder.", "io": "-", "calls": "-", "label": "model-construction", "consolidate_into": "", "notes": "Non-trivial init (creates a Parameter).", "module": "tcri/model/_model.py", "namespace": "ml"}, {"name": "VampPrior.get_mixture", "kind": "method", "signature": "get_mixture(self)", "purpose": "Encode pseudo-inputs and assemble a uniform MixtureSameFamily of Independent Normals as p(z).", "io": "-", "calls": "self.encoder (model forward); pyro Normal/Independent/MixtureSameFamily", "label": "model-construction", "consolidate_into": "", "notes": "Core mechanic of the VampPrior; model internal.", "module": "tcri/model/_model.py", "namespace": "ml"}, {"name": "VampPrior.log_prob", "kind": "method", "signature": "log_prob(self, z)", "purpose": "Log density of z under the VampPrior mixture.", "io": "-", "calls": "self.get_mixture()", "label": "model-construction", "consolidate_into": "", "notes": "Distribution method; model internal.", "module": "tcri/model/_model.py", "namespace": "ml"}, {"name": "VampPrior.sample", "kind": "method", "signature": "sample(self, sample_shape=torch.Size())", "purpose": "Draw samples from the VampPrior mixture.", "io": "-", "calls": "self.get_mixture()", "label": "model-construction", "consolidate_into": "", "notes": "Distribution method; model internal.", "module": "tcri/model/_model.py", "namespace": "ml"}, {"name": "MixtureDirichlet", "kind": "class", "signature": "class MixtureDirichlet(dist.TorchDistribution)", "purpose": "Custom Pyro distribution: a mixture of Dirichlets over the phenotype simplex, used as the clonotype prior p_c.", "io": "-", "calls": "-", "label": "model-construction", "consolidate_into": "", "notes": "File comment tags it 'TODO: Refactor'; custom distribution consumed by TCRIModule.model. Model internal.", "module": "tcri/model/_model.py", "namespace": "ml"}, {"name": "MixtureDirichlet.__init__", "kind": "method", "signature": "__init__(self, mixture_weights, concentration, validate_args=None)", "purpose": "Clamp concentrations, infer batch/B/K shapes, and init the TorchDistribution.", "io": "-", "calls": "-", "label": "model-construction", "consolidate_into": "", "notes": "Shape/constraint setup logic; model internal.", "module": "tcri/model/_model.py", "namespace": "ml"}, {"name": "MixtureDirichlet.sample", "kind": "method", "signature": "sample(self, sample_shape=torch.Size())", "purpose": "Sample a mixture component per batch element then draw from the selected Dirichlet.", "io": "-", "calls": "pyro Categorical/Dirichlet", "label": "model-construction", "consolidate_into": "", "notes": "Distribution sampling; model internal.", "module": "tcri/model/_model.py", "namespace": "ml"}, {"name": "MixtureDirichlet.log_prob", "kind": "method", "signature": "log_prob(self, value)", "purpose": "Log-sum-exp of component Dirichlet log-probs weighted by (log) mixture weights.", "io": "-", "calls": "pyro Dirichlet", "label": "model-construction", "consolidate_into": "", "notes": "Distribution density; model internal.", "module": "tcri/model/_model.py", "namespace": "ml"}, {"name": "MixtureDirichlet.score_parts", "kind": "method", "signature": "score_parts(self, value)", "purpose": "Return (log_prob, zero score-fn, zero entropy) so Pyro treats it as reparam-free.", "io": "-", "calls": "self.log_prob", "label": "model-construction", "consolidate_into": "", "notes": "Pyro ELBO plumbing; model internal.", "module": "tcri/model/_model.py", "namespace": "ml"}, {"name": "MixtureDirichlet.__call__", "kind": "method", "signature": "__call__(self, *args, **kwargs)", "purpose": "Make the distribution callable as an alias for sample().", "io": "-", "calls": "self.sample", "label": "model-construction", "consolidate_into": "", "notes": "Convenience callable; model internal.", "module": "tcri/model/_model.py", "namespace": "ml"}, {"name": "TCRIModule", "kind": "class", "signature": "class TCRIModule(PyroBaseModuleClass)", "purpose": "Pyro CVAE module with hierarchical clonotype->clonotype-covariate Dirichlet priors and a phenotype classifier.", "io": "registers buffers: clone_phen_prior, ct_to_c, c_array, ct_array, ct_to_cov, _target_phenotypes, log_class_weights", "calls": "Encoder, DecoderSCVI, PhenotypeClassifier, VampPrior, MixtureDirichlet; is the model/guide", "label": "model-construction", "consolidate_into": "", "notes": "The pyro model container; model internal per label def.", "module": "tcri/model/_model.py", "namespace": "ml"}, {"name": "TCRIModule.__init__", "kind": "method", "signature": "__init__(self, n_input, n_latent, P, n_batch, global_scale=10.0, local_scale=5.0, ..., mixture_concentration=None, ...)", "purpose": "Construct encoder/decoder/classifier/VampPrior, px_r param, and register empty two-level buffers + class weights.", "io": "creates nn submodules; registers buffers; sets self.kl_weight, reconstruction_loss_scale defaults", "calls": "Encoder, DecoderSCVI, PhenotypeClassifier, VampPrior", "label": "model-construction", "consolidate_into": "", "notes": "Asserts mixture_concentration is not None. Model internal construction.", "module": "tcri/model/_model.py", "namespace": "ml"}, {"name": "TCRIModule.prepare_two_level_params", "kind": "method", "signature": "prepare_two_level_params(self, c_count, ct_count, clone_phen_prior_mat, ct_to_c_array, c_array_for_cells, ct_array_for_cells, target_phenotypes, ct_to_cov_array=None)", "purpose": "Normalize/temperature the clone-phenotype prior and register all two-level index buffers onto the module.", "io": "registers buffers clone_phen_prior, ct_to_c, c_array, ct_array, _target_phenotypes, ct_to_cov; sets c_count/ct_count/n_cells", "calls": "called by TCRIModel.__init__", "label": "model-construction", "consolidate_into": "", "notes": "Model setup/state registration; model internal.", "module": "tcri/model/_model.py", "namespace": "ml"}, {"name": "TCRIModule.use_gate", "kind": "method", "signature": "@property use_gate(self) -> bool", "purpose": "True when gate_prob is not None (selects convex-gate vs additive phenotype combination).", "io": "-", "calls": "-", "label": "model-construction", "consolidate_into": "", "notes": "Module config accessor consumed by training/eval; not a reusable cross-module helper.", "module": "tcri/model/_model.py", "namespace": "ml"}, {"name": "TCRIModule._get_fn_args_from_batch", "kind": "method", "signature": "@staticmethod _get_fn_args_from_batch(tensor_dict)", "purpose": "Extract (x, batch_idx, log_library) tuple from a scvi batch dict for model()/guide().", "io": "reads REGISTRY_KEYS.X_KEY, BATCH_KEY from batch", "calls": "scvi/pyro dataloader plumbing", "label": "model-construction", "consolidate_into": "", "notes": "Required Pyro/scvi arg-mapping hook; model internal.", "module": "tcri/model/_model.py", "namespace": "ml"}, {"name": "TCRIModule.model", "kind": "method", "signature": "@auto_move_data model(self, x, batch_idx, log_library)", "purpose": "Generative model: sample p_c (MixtureDirichlet), p_ct (Dirichlet), latent z (VampPrior), and ZINB gene obs.", "io": "reads clone_phen_prior/ct_to_c/ct_array buffers; kl_weight, reconstruction_loss_scale; pyro sample sites p_c,p_ct,latent,obs", "calls": "MixtureDirichlet, VampPrior.get_mixture, encoder, decoder, classifier; defines posterior model", "label": "model-construction", "consolidate_into": "", "notes": "Pyro model() -> model-construction per label definition (pyro model/guide).", "module": "tcri/model/_model.py", "namespace": "ml"}, {"name": "TCRIModule.guide", "kind": "method", "signature": "@auto_move_data guide(self, x, batch_idx, log_library)", "purpose": "Variational guide: learnable Dirichlet params q(p_c), q(p_ct) and Normal q(z) from the encoder.", "io": "reads/creates pyro params q_p_c_raw, q_p_ct_raw; reads clone_phen_prior; pyro sample sites p_c,p_ct,latent", "calls": "encoder; pyro.param store (q_p_c_raw, q_p_ct_raw)", "label": "model-construction", "consolidate_into": "", "notes": "Pyro guide() -> model-construction per label definition.", "module": "tcri/model/_model.py", "namespace": "ml"}, {"name": "TCRIModule.get_latent", "kind": "method", "signature": "@auto_move_data get_latent(self, tensor_dict)", "purpose": "Encode a batch to the posterior-mean latent z_loc (collapses MC dim if present).", "io": "reads REGISTRY_KEYS.X_KEY, BATCH_KEY; returns z_loc on CPU", "calls": "self.encoder (model forward); used by get_latent_representation and training diagnostics", "label": "core", "consolidate_into": "", "notes": "Implements the core latent-representation output; not model/guide, so core rather than model-construction.", "module": "tcri/model/_model.py", "namespace": "ml"}, {"name": "TCRIModule.get_p_ct", "kind": "method", "signature": "@torch.no_grad get_p_ct(self)", "purpose": "Read learned q_p_ct_raw from the Pyro param store and normalize (with guide-temperature) to the clone-covariate x phenotype posterior.", "io": "reads pyro param store q_p_ct_raw", "calls": "pyro get_param_store; substrate for get_cell_phenotype_probs and TCRIModel.get_p_ct", "label": "core", "consolidate_into": "", "notes": "Core posterior accessor (the p_ct substrate). Requires train() to have populated the param store, else KeyError.", "module": "tcri/model/_model.py", "namespace": "ml"}, {"name": "UnifiedTrainingPlan", "kind": "class", "signature": "class UnifiedTrainingPlan(PyroTrainingPlan)", "purpose": "Training plan adding KL warmup, diagnostics logging, and a validation_step emitting elbo_validation for early stopping.", "io": "logs metrics; mutates module.kl_weight", "calls": "TraceEnum_ELBO/Trace_ELBO; wraps TCRIModule", "label": "model-construction", "consolidate_into": "", "notes": "Training plan -> model-construction per label definition.", "module": "tcri/model/_model.py", "namespace": "ml"}, {"name": "UnifiedTrainingPlan.__init__", "kind": "method", "signature": "__init__(self, module, n_steps_kl_warmup=1000, reconstruction_loss_scale=1e-2, num_particles=5, optimizer_config=None, class_weights=None, **kwargs)", "purpose": "Choose TraceEnum_ELBO vs Trace_ELBO by module.use_enumeration and store optimizer/warmup config.", "io": "-", "calls": "TraceEnum_ELBO/Trace_ELBO; super PyroTrainingPlan", "label": "model-construction", "consolidate_into": "", "notes": "Training-plan setup logic; model internal.", "module": "tcri/model/_model.py", "namespace": "ml"}, {"name": "UnifiedTrainingPlan.loss", "kind": "method", "signature": "@property loss(self)", "purpose": "Expose the configured ELBO loss object.", "io": "-", "calls": "-", "label": "model-construction", "consolidate_into": "", "notes": "Training-plan accessor; model internal.", "module": "tcri/model/_model.py", "namespace": "ml"}, {"name": "UnifiedTrainingPlan.configure_optimizers", "kind": "method", "signature": "configure_optimizers(self)", "purpose": "Build an Adam optimizer over module parameters from optimizer_config.", "io": "-", "calls": "torch.optim.Adam", "label": "model-construction", "consolidate_into": "", "notes": "Lightning/scvi training hook; model internal.", "module": "tcri/model/_model.py", "namespace": "ml"}, {"name": "UnifiedTrainingPlan.training_step", "kind": "method", "signature": "training_step(self, batch, batch_idx)", "purpose": "Apply KL warmup, run the Pyro ELBO step, and log KL/entropy/confidence diagnostics.", "io": "mutates module.kl_weight; self.log kl_divergence_with_prior_train/entropy_train/confidence_train", "calls": "super().training_step (ELBO); module.get_latent, module.classifier, module.get_p_ct", "label": "model-construction", "consolidate_into": "", "notes": "Training loop step; model internal.", "module": "tcri/model/_model.py", "namespace": "ml"}, {"name": "UnifiedTrainingPlan.validation_step", "kind": "method", "signature": "validation_step(self, batch, batch_idx)", "purpose": "Compute validation ELBO and prior-KL, logging elbo_validation for scvi early stopping.", "io": "self.log kl_divergence_with_prior_val, elbo_validation", "calls": "super().training_step (ELBO); module.get_latent, module.classifier, module.get_p_ct", "label": "model-construction", "consolidate_into": "", "notes": "Validation loop step; model internal.", "module": "tcri/model/_model.py", "namespace": "ml"}, {"name": "TCRIModel", "kind": "class", "signature": "class TCRIModel(BaseModelClass)", "purpose": "High-level scvi model API: setup, build, train, and extract latent/phenotype/p_ct outputs.", "io": "reads/writes adata.obs and adata.uns via methods; owns self.module", "calls": "TCRIModule, UnifiedTrainingPlan, TrainRunner", "label": "core", "consolidate_into": "", "notes": "The core MODEL entry point named in the CORE definition.", "module": "tcri/model/_model.py", "namespace": "ml"}, {"name": "TCRIModel.setup_anndata", "kind": "method", "signature": "@classmethod setup_anndata(cls, adata, layer=None, clonotype_key='unique_clone_id', phenotype_key='phenotype_col', covariate_key='timepoint', batch_key='patient', **kwargs)", "purpose": "Register clonotype/phenotype/covariate/batch/count fields with scvi and stash the manager/layer on the AnnData.", "io": "writes adata.obs['indices']; adata.uns['tcri_manager']; adata.uns['tcri_layer'] (or pops); reads the four key columns", "calls": "scvi AnnDataManager, LayerField, CategoricalObsField, register_manager", "label": "core", "consolidate_into": "", "notes": "Explicitly listed as core (MODEL setup).", "module": "tcri/model/_model.py", "namespace": "ml"}, {"name": "TCRIModel.__init__", "kind": "method", "signature": "__init__(self, adata, n_latent=128, n_hidden=128, n_layers=3, ..., K=10, phenotype_weights=None, gate_prob=None, ...)", "purpose": "Build clone->phenotype matrix, archetypes, clonotype-covariate index maps and class weights, then construct+prime TCRIModule.", "io": "reads adata.obs clonotype/phenotype/covariate/batch cols; sets self.c2p_mat, centers, labels, class_weights, module, init_params_, patience", "calls": "build_archetypes; TCRIModule; module.prepare_two_level_params", "label": "core", "consolidate_into": "", "notes": "Core model build step (assembles priors + module).", "module": "tcri/model/_model.py", "namespace": "ml"}, {"name": "TCRIModel.train", "kind": "method", "signature": "train(self, max_epochs=1000, batch_size=1000, lr=1e-3, reconstruction_loss_scale=1e-3, n_steps_kl_warmup=2000, **kwargs)", "purpose": "Split data, build UnifiedTrainingPlan, and run TrainRunner with elbo_validation early stopping.", "io": "sets module.reconstruction_loss_scale; populates self.history_, module params, pyro param store", "calls": "DataSplitter, UnifiedTrainingPlan, TrainRunner", "label": "core", "consolidate_into": "", "notes": "Explicitly listed as core (MODEL train).", "module": "tcri/model/_model.py", "namespace": "ml"}, {"name": "TCRIModel.get_latent_representation", "kind": "method", "signature": "@torch.no_grad get_latent_representation(self, adata=None, indices=None, batch_size=None)", "purpose": "Batched encode of an AnnData to a (n_cells, n_latent) numpy latent matrix.", "io": "reads registry X/BATCH via loader; returns numpy array", "calls": "self._validate_anndata, _make_data_loader, module.get_latent", "label": "core", "consolidate_into": "", "notes": "Explicitly listed as core (MODEL output).", "module": "tcri/model/_model.py", "namespace": "ml"}, {"name": "TCRIModel.use_gate", "kind": "method", "signature": "@property use_gate(self) -> bool", "purpose": "Public passthrough to module.use_gate.", "io": "-", "calls": "self.module.use_gate", "label": "helper", "consolidate_into": "", "notes": "Trivial delegating accessor; internal code uses self.module.use_gate directly, so this public property has no in-file callers.", "module": "tcri/model/_model.py", "namespace": "ml"}, {"name": "TCRIModel.get_p_ct", "kind": "method", "signature": "@torch.no_grad get_p_ct(self)", "purpose": "Return the learned clone-covariate x phenotype posterior p_ct as a numpy array.", "io": "reads pyro param store via module", "calls": "module.get_p_ct", "label": "core", "consolidate_into": "", "notes": "Explicitly listed as core (MODEL output feeding metrics).", "module": "tcri/model/_model.py", "namespace": "ml"}, {"name": "TCRIModel.get_cell_phenotype_probs", "kind": "method", "signature": "@torch.no_grad get_cell_phenotype_probs(self, adata=None, batch_size=256, eps=1e-8) -> np.ndarray", "purpose": "Per-cell phenotype probabilities by combining classifier logits with log p_ct prior (gate or additive), matching training.", "io": "reads registry X/BATCH and module.ct_array; returns (n_cells, P) numpy", "calls": "module.get_p_ct, module.encoder, module.classifier", "label": "core", "consolidate_into": "", "notes": "Explicitly listed as core (MODEL output).", "module": "tcri/model/_model.py", "namespace": "ml"}, {"name": "TCRIModel.boost_phenotype_prior", "kind": "method", "signature": "boost_phenotype_prior(self, phenotype_name, boost_factor=5.0, *, affect_mixture=True)", "purpose": "Multiply one phenotype's column in the clone prior (and optionally mixture centers) by a factor, renormalize, and overwrite module buffers.", "io": "mutates self.c2p_mat, module.clone_phen_prior buffer, self.centers, module.mixture_concentration; prints ANSI status", "calls": "inner _ok; reads adata.obs phenotype col", "label": "model-construction", "consolidate_into": "", "notes": "Optional pre-train prior-adjustment convenience, not part of the core five; functional. Inline ANSI printing (_ok) should move to a shared helper.", "module": "tcri/model/_model.py", "namespace": "ml"}, {"name": "TCRIModel.boost_phenotype_prior._ok", "kind": "inner-function", "signature": "_ok(m)", "purpose": "Print a green check-marked status line.", "io": "stdout (ANSI)", "calls": "-", "label": "helper", "consolidate_into": "", "notes": "Console-printing utility duplicated inline; belongs in a shared console/ANSI helper module.", "module": "tcri/model/_model.py", "namespace": "ml"}, {"name": "TCRIModel.plot_archetypes", "kind": "method", "signature": "plot_archetypes(self)", "purpose": "Heatmap the cluster-ordered clone-phenotype matrix and the archetype centroids.", "io": "reads self.c2p_mat, self.labels, self.centers; matplotlib plt.show side effect", "calls": "matplotlib", "label": "plotting-beyond-core", "consolidate_into": "", "notes": "Model-internal diagnostic plot (archetypes/priors), not a plot of the core metrics or joint distribution.", "module": "tcri/model/_model.py", "namespace": "ml"}, {"name": "TCRIModel.plot_loss", "kind": "method", "signature": "plot_loss(self, log_scale=False)", "purpose": "Plot training/validation ELBO and prior-KL curves from self.history_.", "io": "reads self.history_ keys elbo_train/elbo_validation/kl_divergence_with_prior_train_epoch/kl_divergence_with_prior_val; plt.show", "calls": "matplotlib", "label": "plotting-beyond-core", "consolidate_into": "", "notes": "Training-diagnostic plot, not a core-metric plot. Minor mislabel: axes[1] legend says 'Accuracy' but the series are KL-divergence values.", "module": "tcri/model/_model.py", "namespace": "ml"}, {"name": "stars", "kind": "function", "signature": "stars(p) -> str", "purpose": "Map a p-value to a significance-star string (****/***/**/*/ns).", "io": "-", "calls": "-", "label": "helper", "consolidate_into": "", "notes": "Reusable low-level stats/formatting helper; belongs in a shared _stats module. Currently imported by tcri/plotting/_plotting.py but never actually called (stars( appears only at its own def).", "module": "tcri/utils/_utils.py", "namespace": "ut"}, {"name": "auc_and_label_permutation", "kind": "function", "signature": "auc_and_label_permutation(scores, labels, pos_label=None, n_perm=200000, seed=42, max_exact=200000) -> (obs_auc, p_perm, perm_stats, perm_mode)", "purpose": "Compute observed ROC-AUC plus a two-sided permutation p-value (exact combinations if feasible, else Monte-Carlo).", "io": "-", "calls": "-", "label": "helper", "consolidate_into": "", "notes": "Pure sklearn/numpy stats utility; used only by pl.mi_compare (_plotting.py:142). Should be extracted to a shared _stats helper module.", "module": "tcri/utils/_utils.py", "namespace": "ut"}, {"name": "bootstrap_auc", "kind": "function", "signature": "bootstrap_auc(scores, labels, pos_label=None, n_boot=5000, seed=42) -> np.ndarray[2]", "purpose": "Bootstrap 95% CI (2.5/97.5 quantiles) of ROC-AUC, resampling until both classes present.", "io": "-", "calls": "-", "label": "helper", "consolidate_into": "", "notes": "Pure stats utility; used only by pl.mi_compare (_plotting.py:144). Shared _stats helper.", "module": "tcri/utils/_utils.py", "namespace": "ut"}, {"name": "_ensure_pyro_posterior_params", "kind": "function", "signature": "_ensure_pyro_posterior_params(model, adata) -> None", "purpose": "After load, guarantee the Pyro param 'q_p_ct_raw' exists; if missing, warn and re-init it to a uniform 1/P simplex so posterior metrics can still run.", "io": "reads/writes pyro global param store (q_p_ct_raw); reads model.module (ct_count/classifier), adata.obs phenotype col; emits RuntimeWarning", "calls": "model.get_latent_representation (hits latent/posterior), model.module.classifier forward", "label": "session-io", "consolidate_into": "", "notes": "Part of the session-load restoration path; invoked by load_tcri_session and by tests/test_pyro_params.py. Ensures posterior substrate is present post-load.", "module": "tcri/utils/_utils.py", "namespace": "ut"}, {"name": "_resolve_TCRIModel", "kind": "function", "signature": "_resolve_TCRIModel() -> type", "purpose": "Dynamically locate and import the TCRIModel class from common module paths or sibling files (editable installs).", "io": "module import side effects (may exec sibling _model.py)", "calls": "-", "label": "session-io", "consolidate_into": "", "notes": "Import-resolution plumbing that exists solely to support load_tcri_session.", "module": "tcri/utils/_utils.py", "namespace": "ut"}, {"name": "_disable_scvi_onload_train", "kind": "function", "signature": "@contextmanager _disable_scvi_onload_train()", "purpose": "Context manager that monkey-patches scvi PyroBaseModuleClass.on_load to a no-op during model load, avoiding the one-step warmup train that trips EarlyStopping.", "io": "monkeypatches scvi PyroBaseModuleClass.on_load (restored in finally)", "calls": "defines/uses inner _noop (clears pyro store)", "label": "session-io", "consolidate_into": "", "notes": "Load-time plumbing wrapped around TCRIModel.load inside load_tcri_session.", "module": "tcri/utils/_utils.py", "namespace": "ut"}, {"name": "_disable_scvi_onload_train._noop", "kind": "inner-function", "signature": "_noop(self, *args, **kwargs)", "purpose": "Replacement on_load that swallows all args and just clears the Pyro param store (own params loaded afterward).", "io": "clears pyro global param store", "calls": "pyro.clear_param_store", "label": "session-io", "consolidate_into": "", "notes": "Inner monkeypatch target of _disable_scvi_onload_train; not trivial __init__ \u2014 carries the store-clearing logic.", "module": "tcri/utils/_utils.py", "namespace": "ut"}, {"name": "_pyro_load", "kind": "function", "signature": "_pyro_load(path, map_location=None) -> None", "purpose": "torch.load a Pyro param-store state dict (weights_only=False for constraint objects) and set it into the global param store.", "io": "reads .pt file; writes pyro global param store", "calls": "-", "label": "session-io", "consolidate_into": "", "notes": "Deserialization plumbing for load_tcri_session; weights_only=False is intentional for self-produced artifacts.", "module": "tcri/utils/_utils.py", "namespace": "ut"}, {"name": "load_tcri_session", "kind": "function", "signature": "load_tcri_session(run_dir, *, adata_path=None, map_location=None, layer=None) -> (model, adata)", "purpose": "Reconstruct a trained TCRIModel + AnnData from a saved run dir: read h5ad, restore setup/category order, re-run setup_anndata, load model without warmup, restore pyro store, ensure posterior params.", "io": "reads run_dir/{adata.h5ad,setup.json,pyro_params.pt}; mutates adata.obs category order; writes pyro param store", "calls": "_resolve_TCRIModel, _collect_setup_from_adata_or_model, _restore_category_order, TCRIModel.setup_anndata, TCRIModel.load, _disable_scvi_onload_train, _pyro_load, _ensure_pyro_posterior_params (model load + pyro)", "label": "session-io", "consolidate_into": "", "notes": "Core session-load entry point; covered by tests/test_session_round_trip.py. Uses module constants AD_FILE/SETUP_FILE/PYRO_FILE defined later (resolved at call time, fine).", "module": "tcri/utils/_utils.py", "namespace": "ut"}, {"name": "probabilities", "kind": "function", "signature": "probabilities(adata) -> dict[barcode -> dict[label -> prob]]", "purpose": "Build a per-cell {barcode: {phenotype: prob}} dict from the probability columns and the joint_distribution index.", "io": "reads adata.uns['probability_columns'], adata.obs[those cols], adata.uns['joint_distribution'].index", "calls": "-", "label": "dead-broken", "consolidate_into": "", "notes": "Broken/dead: reads adata.uns['joint_distribution'] which is NEVER written anywhere in the package (only tests/conftest writes uns['mock_joint_distribution']), so it raises KeyError at line 293. Imported by _plotting.py but not invoked in a working path.", "module": "tcri/utils/_utils.py", "namespace": "ut"}, {"name": "build_nested_tcri_pgm", "kind": "function", "signature": "build_nested_tcri_pgm() -> daft.PGM", "purpose": "Construct a daft probabilistic-graphical-model diagram of the nested TCRI generative model (plates for batch/clonotype/ct/data, nodes p_c/p_ct/z/etc.).", "io": "-", "calls": "daft.PGM (external)", "label": "plotting-beyond-core", "consolidate_into": "", "notes": "Visualizes model architecture, not the core metrics/joint distribution/flux; hardcoded layout coordinates. 0 callers except draw_tcri_pgm_nested.", "module": "tcri/utils/_utils.py", "namespace": "ut"}, {"name": "draw_tcri_pgm_nested", "kind": "function", "signature": "draw_tcri_pgm_nested() -> None", "purpose": "Render build_nested_tcri_pgm and save it to a hardcoded PDF, then plt.show().", "io": "writes file 'tcri_model_fully_explicit.pdf' to cwd; shows figure", "calls": "build_nested_tcri_pgm", "label": "plotting-beyond-core", "consolidate_into": "", "notes": "One-off diagram export with hardcoded output filename; 0 callers in package. Belongs in pl as a model-diagram plot, not core-metric plotting.", "module": "tcri/utils/_utils.py", "namespace": "ut"}, {"name": "_ensure_dir", "kind": "function", "signature": "_ensure_dir(path: str) -> None", "purpose": "os.makedirs(path, exist_ok=True) wrapper.", "io": "creates directory on disk", "calls": "-", "label": "session-io", "consolidate_into": "", "notes": "Trivial filesystem plumbing used only by save_tcri_session; part of the session save/load cluster.", "module": "tcri/utils/_utils.py", "namespace": "ut"}, {"name": "_to_jsonable", "kind": "function", "signature": "_to_jsonable(x: Any) -> Any", "purpose": "Recursively coerce arbitrary values (numpy/torch/nested) into JSON-serializable primitives.", "io": "-", "calls": "self (recursive)", "label": "helper", "consolidate_into": "", "notes": "Reusable serialization helper, but currently 0 external callers (only self-recursion); save_tcri_session json.dumps setup/meta directly without it. Extract to shared helpers or wire in / drop.", "module": "tcri/utils/_utils.py", "namespace": "ut"}, {"name": "_pop_nonserializables", "kind": "function", "signature": "_pop_nonserializables(adata) -> Dict[str, Any]", "purpose": "Remove the non-picklable AnnDataManager (uns['tcri_manager']) before writing h5ad, returning a sidecar note.", "io": "mutates adata.uns (pops 'tcri_manager')", "calls": "-", "label": "session-io", "consolidate_into": "", "notes": "Serialization sanitization plumbing used by write_adata_safely.", "module": "tcri/utils/_utils.py", "namespace": "ut"}, {"name": "write_adata_safely", "kind": "function", "signature": "write_adata_safely(adata, path, *, compression='gzip') -> None", "purpose": "Write an AnnData to h5ad after stripping the non-serializable tcri_manager (not restored; rebuilt on load).", "io": "writes .h5ad file; mutates adata.uns via _pop_nonserializables", "calls": "_pop_nonserializables", "label": "session-io", "consolidate_into": "", "notes": "AnnData serialization step of save_tcri_session.", "module": "tcri/utils/_utils.py", "namespace": "ut"}, {"name": "_collect_setup_from_adata_or_model", "kind": "function", "signature": "_collect_setup_from_adata_or_model(adata, model) -> Dict[str, Any]", "purpose": "Assemble the setup dict (phenotype/clone/covariate/batch cols, category lists, layer) from adata.uns['tcri_metadata']/categories and, if available, the model's adata_manager registry.", "io": "reads adata.uns (tcri_metadata, tcri_*_categories, tcri_layer); reads model.adata_manager.registry", "calls": "-", "label": "session-io", "consolidate_into": "", "notes": "Setup-metadata collector used by both save_tcri_session and load_tcri_session (fallback).", "module": "tcri/utils/_utils.py", "namespace": "ut"}, {"name": "_restore_category_order", "kind": "function", "signature": "_restore_category_order(adata, setup) -> None", "purpose": "Re-impose saved categorical ordering on adata.obs phenotype/clone/covariate columns from the setup dict.", "io": "mutates adata.obs[col] into ordered pd.Categorical", "calls": "-", "label": "session-io", "consolidate_into": "", "notes": "Category-order restoration used by load_tcri_session to keep phenotype/covariate encodings aligned.", "module": "tcri/utils/_utils.py", "namespace": "ut"}, {"name": "save_tcri_session", "kind": "function", "signature": "save_tcri_session(model, adata, out_dir, *, save_adata=True, compression='gzip') -> Dict[str, Any]", "purpose": "Persist a trained session: scvi model (weights+registry, no embedded adata), Pyro param store, setup.json, sanitized adata.h5ad, and meta.json with versions; return written paths.", "io": "writes out_dir/{model files, pyro_params.pt, setup.json, adata.h5ad, meta.json}", "calls": "model.save (scvi), pyro param_store.save, _collect_setup_from_adata_or_model, write_adata_safely, _ensure_dir", "label": "session-io", "consolidate_into": "", "notes": "Core session-save entry point; covered by tests/test_session_round_trip.py.", "module": "tcri/utils/_utils.py", "namespace": "ut"}], "mergedCount": 131, "critic": {"missing": [], "mislabeled": [{"name": "clonality (pl)", "current": "core", "corrected": "plotting-beyond-core", "why": "CORE PLOTTING is plots that DIRECTLY visualize the core metrics/joint-distribution/flux. clonality is NOT in the core metric set, and metrics.clonality is itself labeled 'redundant' in this inventory. A core-labeled plot whose only subject is a redundant, non-core metric is inconsistent. tcri/plotting/_plotting.py:661."}, {"name": "tcri_boxplot (pl)", "current": "core", "corrected": "helper", "why": "Generic reusable box/strip engine taking an arbitrary metric `function`; a shared plotting primitive, not a plot of a specific core metric. Its sole caller pl.clonality renders the non-core (redundant) clonality metric. Per the label set this is a reusable helper to extract, not core. tcri/plotting/_plotting.py:598."}, {"name": "TCRIModule.get_latent", "current": "core", "corrected": "model-construction", "why": "PLAUSIBLE consistency flag. The enumerated core output is public TCRIModel.get_latent_representation (correctly core), which delegates to this module-internal encoder method. All sibling TCRIModule methods (model, guide, prepare_two_level_params, use_gate, _get_fn_args_from_batch) are model-construction; labeling this one core is inconsistent. tcri/model/_model.py:502."}, {"name": "TCRIModule.get_p_ct", "current": "core", "corrected": "model-construction", "why": "PLAUSIBLE consistency flag. The core p_ct entrypoints are TCRIModel.get_p_ct + pp.register_model, which wrap/materialize this internal pyro-param-store accessor. For consistency with the other TCRIModule internals it is model-construction. Defensible to keep as core substrate, hence PLAUSIBLE. tcri/model/_model.py:511."}], "duplication": [{"members": "metrics._ok/_info/_warn/_fin (_metrics.py:47-56); preprocessing._ok/_info/_warn/_fin (_preprocessing.py:50-59); plotting._ok/_info/_warn/_fin (_plotting.py:37-46); model.boost_phenotype_prior._ok inline (_model.py:982)", "into": "one shared console/ANSI helper module in utils", "why": "Byte-for-byte triplicated (quadruplicated for _ok) console printers. Several pl notes claim these duplicate 'tcri.utils._utils._ok' but that utils copy does NOT exist (grep confirms no such def in utils/_utils.py). Should be consolidated into utils per the CORE definition. Labeled helper, not yet redundant."}, {"members": "metrics._ascii_hist (_metrics.py:62); preprocessing._ascii_hist (_preprocessing.py:62)", "into": "one shared ASCII-histogram helper", "why": "Two identical copies of the numpy-histogram ASCII bar builder; only the metrics copy is called. Both labeled helper (not redundant)."}, {"members": "preprocessing.group_singletons + inner collapse_singleton (_preprocessing.py:83); preprocessing.group_small_clones (_preprocessing.py:507)", "into": "one parametrized small-clone-collapse util", "why": "Both collapse clones under a size threshold into 'Singleton_{group}' labels writing obs['trb_unique']. group_singletons is helper, group_small_clones is dead-broken; the cluster is not marked redundant, so flagging per instructions."}, {"members": "metrics.dkl module-level (_metrics.py:159); metrics.flux.dkl_func inner (_metrics.py:962)", "into": "one shared KL/distance helper", "why": "Both implement KL divergence for the 'dkl' distance branch. Module-level dkl is dead-broken (0 callers); flux's inner dkl_func reimplements the same kernel and is helper. Cluster not marked redundant."}], "count_by_file": "All defs/classes accounted for; zero missing. metrics/_metrics.py: 21 records (19 top-level + 2 inner). preprocessing/_preprocessing.py: 20 (19 + inner collapse_singleton). plotting/_plotting.py: 24. plotting/_sankey.py: 5. model/_model.py: 42 (incl. inner boost_phenotype_prior._ok). utils/_utils.py: 19 (incl. inner _noop). Each grep'd def/class maps 1:1 to an inventory record.\nCompleteness PASS \u2014 no missing functions. Verified secondary claims: (a) metrics.mutual_information is keyword-only after `covariate` with NO `weighted` param and its posterior=False branch raises NotImplementedError, so metrics.mutual_information(posterior=False) and pl.mutual_information(weighted=...) both error (dead-broken correct). (b) metrics.phenotypic_entropy signature (adata, covariate, *, point_estimate, n_samples, temperature, combine_with_logits) vs pl.phenotypic_entropy calling pentropy(subt,t,ph,...,clones=,normalized=) \u2014 illegal 3rd positional + unknown kwargs \u2192 TypeError (dead-broken correct). (c) pl.flux (plotting-beyond-core) is ALSO broken on its default paint=None path: `pcat` at _plotting.py:1075 is used before assignment (only set inside the paint!=None branch) \u2192 NameError; runs only when paint supplied. Not reclassified since it works with paint, but the record understates that the default call crashes. Borderline-but-not-flagged: metrics.clonality 'redundant' is defensible though it is actively consumed by pl.clonality; TCRIModel.use_gate 'helper' is a trivial passthrough to module.use_gate (minor duplicate, not worth a cluster).\n\n"}, "synthesis": {"target_layout": [{"module": "tcri/__init__.py", "purpose": "Top-level namespace. Keep the existing alias+sys.modules pattern (tl/pp/pl/ut/ml) but add `diag` (new diagnostics) and optionally `get`; consider a _LazyModule proxy for heavy pl/ml so `import tcri` stays light. Declare no `import *`.", "functions": []}, {"module": "tcri/_console.py", "purpose": "NEW shared house-style helper (grafiti plotting/_base + _colors analog for console output). Single canonical copy of the ANSI printers + ASCII histogram, killing the 3-4x triplication across metrics/pp/pl/model.", "functions": ["_ok", "_info", "_warn", "_fin", "_ascii_hist"]}, {"module": "tcri/_stats.py", "purpose": "NEW shared stats helper (grafiti _compute/stats analog). Pure numpy/sklearn significance + AUROC utilities used by pl.mi_compare.", "functions": ["stars", "auc_and_label_permutation", "bootstrap_auc"]}, {"module": "tcri/_distance.py", "purpose": "NEW shared distance-kernel helper. One home for the phenotype-distribution distance metrics used by flux (l1 + KL), replacing the dead module-level dkl and flux's inner dkl_func.", "functions": ["kl_divergence", "l1_distance", "phenotype_distance"]}, {"module": "tcri/_keys.py", "purpose": "NEW shared constants module (grafiti _state/keys analog). Single source of truth for every uns/obs/obsm string (tcri_metadata, tcri_p_ct, tcri_local_scale, tcri_*_categories, X_tcri, X_tcri_logits, ...). Eliminates scattered string literals and legacy-key drift (tcri_clone_key/tcri_phenotype_key vs tcri_metadata).", "functions": []}, {"module": "tcri/model/_model.py", "purpose": "MODEL public API (gf.ml analog). The scvi-style TCRIModel class: the five core methods plus optional pre-train prior boost. plot_loss/plot_archetypes are beyond-core (see triage).", "functions": ["TCRIModel.setup_anndata", "TCRIModel.__init__", "TCRIModel.train", "TCRIModel.get_latent_representation", "TCRIModel.get_p_ct", "TCRIModel.get_cell_phenotype_probs", "TCRIModel.use_gate", "TCRIModel.boost_phenotype_prior"]}, {"module": "tcri/model/_module.py", "purpose": "The Pyro CVAE module (model/guide + latent/p_ct accessors that TCRIModel wraps). Model-construction internals split out of the 1074-line monolith.", "functions": ["TCRIModule.__init__", "TCRIModule.prepare_two_level_params", "TCRIModule.use_gate", "TCRIModule._get_fn_args_from_batch", "TCRIModule.model", "TCRIModule.guide", "TCRIModule.get_latent", "TCRIModule.get_p_ct"]}, {"module": "tcri/model/_priors.py", "purpose": "Prior distributions + archetype builder for the clonotype prior.", "functions": ["build_archetypes", "MixtureDirichlet", "VampPrior"]}, {"module": "tcri/model/_classifier.py", "purpose": "Temperature-scaled MLP phenotype-classifier head.", "functions": ["PhenotypeClassifier"]}, {"module": "tcri/model/_training.py", "purpose": "scvi/Lightning training plan (KL warmup, diagnostics logging, elbo_validation early stopping).", "functions": ["UnifiedTrainingPlan"]}, {"module": "tcri/model/__init__.py", "purpose": "Re-export only the public class(es): TCRIModel (+ TCRIModule if needed for load).", "functions": []}, {"module": "tcri/preprocessing/_register.py", "purpose": "PREPROCESSING (gf.pp analog): materialize model outputs onto the AnnData substrate every metric reads.", "functions": ["register_model", "_compute_logits_and_prior", "register_phenotype_key", "register_clonotype_key"]}, {"module": "tcri/preprocessing/_engine.py", "purpose": "CORE ENGINE: the Bayesian posterior-sampling clone x phenotype substrate. One unified engine (posterior/n_samples switch) that every tl metric reads and pl.sankey draws.", "functions": ["joint_distribution_posterior"]}, {"module": "tcri/preprocessing/_clones.py", "purpose": "Clone-level preprocessing utilities (one parametrized small-clone collapse + per-cell clone size).", "functions": ["group_singletons", "clone_size"]}, {"module": "tcri/preprocessing/__init__.py", "purpose": "Explicit re-exports with aggregate __all__ (register_model, joint_distribution_posterior, group_singletons, clone_size) \u2014 no `import *`.", "functions": []}, {"module": "tcri/metrics/_clonotypic.py", "purpose": "METRICS (gf.tl analog) \u2014 clonotypic-entropy family + its between-covariate change and tidy table. clonotypic_entropy absorbs clonotypic_entropy_base and clonality.", "functions": ["clonotypic_entropy", "delta_clonotypic_entropy", "delta_entropy_table"]}, {"module": "tcri/metrics/_phenotypic.py", "purpose": "METRICS \u2014 phenotypic entropy H[P(phi|c,m)].", "functions": ["phenotypic_entropy"]}, {"module": "tcri/metrics/_mutual_information.py", "purpose": "METRICS \u2014 clone x phenotype mutual information + tidy comparison table; private MI kernel stays module-local.", "functions": ["mutual_information", "mi_compare", "_mi_from_joint"]}, {"module": "tcri/metrics/_flux.py", "purpose": "METRICS \u2014 per-clone phenotypic flux + tidy table (distance kernels come from tcri/_distance.py).", "functions": ["flux", "flux_table"]}, {"module": "tcri/metrics/__init__.py", "purpose": "Explicit re-exports grouped by topic with aggregate __all__ \u2014 drops the `from .._metrics import *` leak.", "functions": []}, {"module": "tcri/plotting/_clonotypic.py", "purpose": "PLOTTING (gf.pl analog) \u2014 pl twins of the clonotypic metrics. clonotypic_entropy_by_phenotype is functional; ridge_delta_entropy needs a 2-line fix (import matplotlib.cm as cm, scipy.stats as st).", "functions": ["clonotypic_entropy_by_phenotype", "ridge_delta_entropy"]}, {"module": "tcri/plotting/_phenotypic.py", "purpose": "PLOTTING \u2014 pl twin of phenotypic_entropy. Needs signature fix: call phenotypic_entropy(adata, covariate, point_estimate=, n_samples=, temperature=, combine_with_logits=) \u2014 drop illegal positional + unknown clones=/normalized= kwargs.", "functions": ["phenotypic_entropy"]}, {"module": "tcri/plotting/_mutual_information.py", "purpose": "PLOTTING \u2014 pl twins for MI. mi_compare is functional; per-covariate mutual_information needs fix: remove the unsupported weighted= kwarg passed to tl.mutual_information.", "functions": ["mutual_information", "mi_compare"]}, {"module": "tcri/plotting/_flux.py", "purpose": "PLOTTING \u2014 the flux/joint-distribution Sankey (core flux visualization). _phenotype_mass_per_clone must be repointed to joint_distribution_posterior after engine consolidation.", "functions": ["phenotypic_flux", "plot_pheno_sankey", "_phenotype_mass_per_clone"]}, {"module": "tcri/plotting/_sankey.py", "purpose": "PLOTTING helper \u2014 low-level Sankey drawing primitive (drop the dead hex_to_rgb method).", "functions": ["SankeyNode"]}, {"module": "tcri/plotting/_base.py", "purpose": "NEW pl house-style layer (grafiti plotting/_base analog): the generic box/strip metric engine + a scanpy-style _finish(fig, ret, show, save) show/save helper.", "functions": ["tcri_boxplot", "_finish"]}, {"module": "tcri/plotting/_colors.py", "purpose": "NEW pl palette helper (grafiti plotting/_colors analog): tcri_colors palette + category color resolution.", "functions": ["set_color_palette"]}, {"module": "tcri/plotting/__init__.py", "purpose": "Explicit re-exports mirroring tl filenames 1:1 with aggregate __all__ \u2014 drops the `from .._plotting import *` leak.", "functions": []}, {"module": "tcri/utils/_session.py", "purpose": "UTILS \u2014 session save/load and its private plumbing (the whole session-io cluster).", "functions": ["save_tcri_session", "load_tcri_session", "write_adata_safely", "_collect_setup_from_adata_or_model", "_restore_category_order", "_ensure_pyro_posterior_params", "_resolve_TCRIModel", "_disable_scvi_onload_train", "_pyro_load", "_ensure_dir", "_pop_nonserializables", "_to_jsonable"]}, {"module": "tcri/utils/__init__.py", "purpose": "Re-export save_tcri_session/load_tcri_session only (stats/console now live in tcri/_stats.py, tcri/_console.py).", "functions": []}, {"module": "tcri/diagnostics/__init__.py", "purpose": "NEW category (gf.diag analog) \u2014 the one tcri lacks. Thin, read-only concordance/sanity checks on finalized model outputs that return small tidy DataFrames and are never plotted; reads model state directly, not routed through the metric/engine loop. Seed it by extracting the posterior/joint-vs-empirical concordance logic from test_joint_distribution_posterior.", "functions": []}], "consolidation_groups": [{"theme": "clonotypic-entropy family", "members": ["metrics.clonotypic_entropy_base", "metrics.clonality"], "into": "metrics.clonotypic_entropy", "rationale": "clonotypic_entropy_base is the single-phenotype/single-draw special case (only extra is posterior=False + weighted) and its sole caller delta_clonotypic_entropy can be rebuilt on clonotypic_entropy(point_estimate=False). clonality is the hard-count complement (1 - H/log2K) of the same H[P(c|phi)] but on observed counts via legacy uns keys (tcri_phenotype_key/tcri_clone_key); it is inconsistent with the tcri_metadata convention and its only consumer (pl.clonality) is beyond-core. Collapse both into clonotypic_entropy as posterior/point modes; drop clonality (or keep as a thin derived helper migrated to tcri_metadata keys \u2014 see open questions)."}, {"theme": "joint-distribution engine", "members": ["preprocessing.joint_distribution"], "into": "preprocessing.joint_distribution_posterior", "rationale": "joint_distribution is the point-estimate/ct-level path producing the same covariate->clone x phenotype table WITHOUT per-cell logits. Unify into one engine with a posterior/n_samples switch. Follow-ups: repoint pl _phenotype_mass_per_clone (sankey) to the unified engine, and wire mutual_information's currently-NotImplementedError posterior=False branch to the engine's point-estimate mode (or drop the posterior kwarg since only the posterior path is real)."}, {"theme": "phenotype assignment algorithm", "members": ["preprocessing.classify_phenotypes"], "into": "preprocessing.register_model", "rationale": "classify_phenotypes writes the same phenotype-probability + hard-label slots register_model produces, via a different archetype-cosine algorithm instead of softmax(logits+prior). It has 0 callers and diverges from the training-consistent assignment. Drop the alternate algorithm; register_model is the single source of per-cell phenotype probabilities."}, {"theme": "small-clone collapse", "members": ["preprocessing.group_small_clones"], "into": "preprocessing.group_singletons", "rationale": "Both collapse sub-threshold clones into 'Singleton_{group}' labels written to obs['trb_unique']. group_small_clones is dead-broken/dataset-specific (0 callers; hardcodes obs['trb'] and requires a precomputed obs['clone_size']). Keep the parametrized group_singletons (clonotype_key/groupby/target_col/min_clone_size) as the single util."}, {"theme": "KL/distance kernel", "members": ["metrics.dkl", "metrics.flux.dkl_func"], "into": "tcri/_distance.py (kl_divergence)", "rationale": "Module-level dkl is dead (0 callers) and flux's inner dkl_func reimplements the same KL kernel for the distance_metric=='dkl' branch. Fold both into one shared kl_divergence in tcri/_distance.py alongside l1_distance, dispatched by phenotype_distance(metric); flux calls the shared helper."}], "deletions": [{"name": "metrics._ent", "reason": "0 callers; would-be entropy helper never wired in \u2014 the entropy metrics inline their own clip/normalise/entropy."}, {"name": "metrics.clone_fraction", "reason": "0 callers; slow reimplementation of obs.groupby(groupby)[clone].value_counts(normalize=True); also uses legacy uns['tcri_clone_key']."}, {"name": "preprocessing.remove_meaningless_genes", "reason": "0 callers AND broken: the include_mtrn/include_hsp branches reassign `genes` from the FULL adata.var.index, silently discarding all prior filters."}, {"name": "preprocessing.get_latent_embedding", "reason": "0 callers; fabricates a spherical Gaussian with a constant posterior_scale (not the model's true latent variance), and default n_samples=0 yields an empty array."}, {"name": "preprocessing.register_probability_columns", "reason": "0 callers; trivial setter whose uns['probability_columns'] key is never read anywhere in the package."}, {"name": "preprocessing.gene_entropy", "reason": "0 callers; per-GENE expression entropy unrelated to core clonotypic/phenotypic entropy; also drops the first count bin on an assumption."}, {"name": "plotting.compare_joint_distribution", "reason": "dead-broken: references an undefined global `model` (model.adata_manager.registry) -> NameError at runtime."}, {"name": "plotting.probability_distribution", "reason": "dead-broken: infinite self-recursion (calls probability_distribution instead of utils.probabilities); also ax indexing fails when splitby=None."}, {"name": "plotting.bayesian_mutual_information", "reason": "dead-broken (passes unsupported weighted= to tl.mutual_information -> TypeError) AND redundant with the functional core pl.mi_compare MI-comparison plot."}, {"name": "plotting.polar_plot", "reason": "dead-broken: phenotypes defaults to a column-NAME string (iterated char-by-char) and the entropy branch calls unimported clonotypic_entropy -> NameError; bespoke radar view, not a core metric plot."}, {"name": "plotting._sankey.SankeyNode.hex_to_rgb", "reason": "0 callers; plot_node_connection uses matplotlib mcolors.to_rgb instead; contract doc already marks it a deleted/unused method."}, {"name": "utils.probabilities", "reason": "dead-broken: reads uns['joint_distribution'] which is never written anywhere in the package -> KeyError."}], "helper_extraction": [{"helper": "_ok, _info, _warn, _fin (ANSI console printers)", "into_module": "tcri/_console.py", "current_copies": "metrics/_metrics.py:47-56; preprocessing/_preprocessing.py:50-59; plotting/_plotting.py:37-46; model/_model.py:982 (TCRIModel.boost_phenotype_prior._ok inline). Byte-for-byte triplicated (quadruplicated for _ok). NOTE: pl notes claim a utils._utils copy exists \u2014 it does not; utils has no such def."}, {"helper": "_ascii_hist (numpy-histogram ASCII bar builder)", "into_module": "tcri/_console.py", "current_copies": "metrics/_metrics.py:62 (the used copy); preprocessing/_preprocessing.py:62 (identical, unused)."}, {"helper": "stars, auc_and_label_permutation, bootstrap_auc (significance + AUROC stats)", "into_module": "tcri/_stats.py", "current_copies": "utils/_utils.py (single copies today; relocate out of the utils monolith into a named shared stats module; stars is currently imported by plotting but never called)."}, {"helper": "kl_divergence + l1_distance + phenotype_distance dispatcher", "into_module": "tcri/_distance.py", "current_copies": "metrics/_metrics.py:159 (dead module-level dkl) and metrics/_metrics.py flux.dkl_func inner kernel \u2014 merge both into one shared KL kernel, add l1, dispatch by distance_metric."}, {"helper": "_mi_from_joint (single-source MI kernel)", "into_module": "tcri/metrics/_mutual_information.py (module-private)", "current_copies": "metrics/_metrics.py:_mi_from_joint (single copy; keep private, called by mutual_information)."}, {"helper": "tcri_boxplot (generic per-phenotype box/strip engine) + _finish (scanpy show/save)", "into_module": "tcri/plotting/_base.py", "current_copies": "plotting/_plotting.py:598 tcri_boxplot; _finish is new (formalize the fig/ax show/save convention repeated inline across pl functions)."}, {"helper": "SankeyNode drawing primitive + _phenotype_mass_per_clone data-prep", "into_module": "tcri/plotting/_sankey.py", "current_copies": "plotting/_sankey.py (SankeyNode already; drop hex_to_rgb); plotting/_plotting.py:_phenotype_mass_per_clone (co-locate with the sankey it feeds)."}, {"helper": "set_color_palette + tcri_colors palette constants", "into_module": "tcri/plotting/_colors.py", "current_copies": "plotting/_plotting.py set_color_palette (also fix: it writes uns on adata.copy() so writes never reach caller)."}, {"helper": "register_phenotype_key, register_clonotype_key, _compute_logits_and_prior", "into_module": "tcri/preprocessing/_register.py", "current_copies": "preprocessing/_preprocessing.py (relocate alongside register_model; migrate the legacy tcri_phenotype_key/tcri_clone_key writes toward the tcri_metadata + _keys.py constants)."}, {"helper": "group_singletons (+ collapse_singleton inner), clone_size", "into_module": "tcri/preprocessing/_clones.py", "current_copies": "preprocessing/_preprocessing.py:83 group_singletons; preprocessing/_preprocessing.py clone_size."}, {"helper": "_to_jsonable (recursive JSON coercion of numpy/torch/nested)", "into_module": "tcri/utils/_session.py", "current_copies": "utils/_utils.py:_to_jsonable (0 external callers today; wire it into save_tcri_session's json.dumps or drop \u2014 currently setup/meta are dumped without it)."}], "plotting_triage": {"core_keep": ["clonotypic_entropy_by_phenotype (functional)", "mi_compare (functional)", "phenotypic_flux (functional; flux Sankey)", "plot_pheno_sankey (functional; flux Sankey engine)", "mutual_information (KEEP+FIX: remove unsupported weighted= kwarg passed to tl.mutual_information)", "phenotypic_entropy (KEEP+FIX: correct call to tl.phenotypic_entropy signature \u2014 drop illegal 3rd positional + unknown clones=/normalized= kwargs)", "ridge_delta_entropy (KEEP+FIX: add `import matplotlib.cm as cm` and `import scipy.stats as st`)"], "beyond_core_drop_or_move": ["compare_phenotypes (move->examples: generic categorical-crosstab heatmap, visualizes no core metric)", "top_clone_umap (move->examples: bespoke top-N-clone UMAP overlay, hardcoded title)", "clone_size_umap (move->examples: bespoke clone-size UMAP overlay; mutates adata.obs as a side effect)", "plot_phenotype_probabilities (move->examples: per-cell phenotype-probability UMAP panels)", "clonality (drop: only plots the merged-away redundant clonality metric)", "flux boxplot (drop: broken on default paint=None path \u2014 pcat used before assignment -> NameError \u2014 plus dataset-specific `paint` overlay; the Sankey is the flux visualization)", "plot_archetypes (drop or keep as optional TCRIModel model-diagnostic method: model prior heatmap, not a core-metric plot)", "plot_loss (drop or keep as optional TCRIModel model-diagnostic method: training ELBO/KL curves)", "build_nested_tcri_pgm (move->docs/examples: daft PGM diagram of the model architecture)", "draw_tcri_pgm_nested (drop: one-off PGM export with a hardcoded output PDF filename)"]}, "counts": "By corrected label (critic overrides applied): core 22, redundant 4, helper 35, dead-broken 17, model-construction 31, session-io 12, plotting-beyond-core 10. TOTAL = 131 (matches file record counts 21+20+24+5+42+19). Dispositions: 22 core survive; 4 redundant merge into their core target (consolidation groups G1-G3); 35 helpers de-duplicate to a single copy each across _console/_stats/_distance/_base/_colors/_sankey and the split pp/pl/utils modules; 31 model-construction records survive but are split across model/_module,_priors,_classifier,_training; 12 session-io survive as utils/_session. Of the 17 dead-broken: 12 deleted, 2 merged away (dkl->_distance G5, group_small_clones->group_singletons G4), and 3 KEPT+FIXED as core pl twins (ridge_delta_entropy, pl.phenotypic_entropy, pl.mutual_information) because the CORE definition requires a plot that directly visualizes each core metric and each is a <=2-line fix. Of the 10 plotting-beyond-core: 6 move to examples/docs, 4 drop (clonality, flux boxplot, draw_tcri_pgm_nested, and optionally plot_loss/plot_archetypes kept as model methods).\n1) Broken core-metric plots: confirm the KEEP+FIX approach (repair pl.mutual_information / pl.phenotypic_entropy / pl.ridge_delta_entropy in place) vs deleting and rebuilding fresh pl twins \u2014 all three are trivial fixes but are currently labeled dead-broken. 2) Engine home: keep joint_distribution_posterior in preprocessing/_engine.py (pp namespace, as today) or move it to a private tcri/_compute/_posterior.py and/or expose it under tl, since it is the substrate every tl metric reads. 3) Grafiti storage layer: adopt the full _state/keys.py + @tl_result cache decorator + get.py accessors now, or defer? tcri metrics currently return live values (no uns caching), so pl functions recompute rather than load_result \u2014 making pl true 'cache renderers' is a larger refactor than the split itself. Recommend: land _keys.py now (cheap, high value), defer @tl_result/get.py. 4) diagnostics content: confirm seeding gf.diag by extracting the posterior/joint-vs-empirical concordance logic from test_joint_distribution_posterior into a read-only DataFrame check. 5) clone_size + group_singletons: their only live consumers are beyond-core UMAP/dataset-prep code \u2014 if those move to examples, do clone_size/group_singletons stay in pp core or move too? 6) plot_loss/plot_archetypes/boost_phenotype_prior: keep as optional TCRIModel methods (model-construction) or relocate to a model-diagnostics/examples surface? 7) clonality metric: fully drop, or keep a thin hard-count clonality derived from clonotypic_entropy (would need migration off legacy uns['tcri_phenotype_key']/['tcri_clone_key'] to tcri_metadata)? 8) After engine unification, decide whether to keep the posterior kwarg on mutual_information/flux at all, or hardwire the posterior path and delete the (currently NotImplementedError) prior-only branch.\n\n"}} \ No newline at end of file diff --git a/example/DO_NOT_READ.md b/example/DO_NOT_READ.md new file mode 100644 index 0000000..8815443 --- /dev/null +++ b/example/DO_NOT_READ.md @@ -0,0 +1,5 @@ +# 🚫 DO NOT READ / DO NOT USE THIS DIRECTORY 🚫 + +Every notebook in `example/` is DISPOSABLE and calls an OLD tcri API. + +**Rules (see docs/contract/REFACTOR_HISTORY.md):** never read these for any decision; never move package functions here; `examples/` is not a destination; notebook usage is NOT evidence to keep a function. diff --git a/example/colitis/comparison.ipynb b/example/colitis/comparison.ipynb index 2970d31..471cdc5 100644 --- a/example/colitis/comparison.ipynb +++ b/example/colitis/comparison.ipynb @@ -1,5 +1,20 @@ { "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# \ud83d\udeab DO NOT READ \u2014 DISPOSABLE NOTEBOOK \ud83d\udeab\n", + "\n", + "**This notebook calls an OLD/DIVERGENT tcri API and is scheduled to be rewritten from scratch against the new API. It is NOT a source of truth for anything.**\n", + "\n", + "- Do NOT read it to decide what functions to keep or drop.\n", + "- Do NOT use it for 'sufficiency' or 'is-it-used' evidence.\n", + "- A function being called here is NOT a reason to keep it.\n", + "\n", + "See `docs/contract/REFACTOR_HISTORY.md` (Hard Rules).\n" + ] + }, { "cell_type": "code", "execution_count": 2, @@ -426,7 +441,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 621/621 [00:00<00:00, 3563.87it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 621/621 [00:00<00:00, 3563.87it/s]\n" ] }, { @@ -441,7 +456,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 192510/192510 [01:03<00:00, 3031.58it/s] \n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 192510/192510 [01:03<00:00, 3031.58it/s] \n" ] }, { @@ -457,7 +472,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 621/621 [00:00<00:00, 2781.35it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 621/621 [00:00<00:00, 2781.35it/s]\n" ] }, { @@ -591,7 +606,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 621/621 [00:00<00:00, 2812810.78it/s]" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 621/621 [00:00<00:00, 2812810.78it/s]" ] }, { @@ -606,7 +621,7 @@ "output_type": "stream", "text": [ "\n", - "100%|██████████| 2687599/2687599 [00:04<00:00, 641730.54it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 2687599/2687599 [00:04<00:00, 641730.54it/s]\n" ] }, { @@ -620,7 +635,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 21160/21160 [00:01<00:00, 18649.30it/s]" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 21160/21160 [00:01<00:00, 18649.30it/s]" ] }, { @@ -660,7 +675,7 @@ { "data": { "text/plain": [ - "AnnData object with n_obs × n_vars = 21160 × 621\n", + "AnnData object with n_obs \u00d7 n_vars = 21160 \u00d7 621\n", " obs: 'va', 'ja', 'cdr3a', 'cdr3a_nucseq', 'vb', 'jb', 'cdr3b', 'cdr3b_nucseq', 'disease', 'disease_int', 'donor', 'donor_int', 'batch', 'tcr_cluster', 'n_genes', 'percent_mito', 'n_counts', 'leiden_gex_for_cd4_vs_cd8', 'cd4_or_cd8'\n", " var: 'gene_ids', 'feature_types', 'n_cells', 'highly_variable', 'means', 'dispersions', 'dispersions_norm', 'mean', 'std'\n", " uns: 'organism', 'batch_keys', 'conga_results', 'conga_stats', 'log1p', 'hvg', 'raw_matrix_is_logged', 'pca', 'neighbors', 'leiden', 'umap'\n", @@ -989,7 +1004,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 21160/21160 [00:00<00:00, 67477.96it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 21160/21160 [00:00<00:00, 67477.96it/s]\n" ] }, { @@ -1004,7 +1019,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 21160/21160 [00:00<00:00, 70818.01it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 21160/21160 [00:00<00:00, 70818.01it/s]\n" ] }, { @@ -1019,7 +1034,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 21160/21160 [00:00<00:00, 71409.30it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 21160/21160 [00:00<00:00, 71409.30it/s]\n" ] }, { @@ -1034,7 +1049,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 21160/21160 [00:00<00:00, 69817.25it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 21160/21160 [00:00<00:00, 69817.25it/s]\n" ] }, { @@ -1049,7 +1064,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 21160/21160 [00:00<00:00, 70503.03it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 21160/21160 [00:00<00:00, 70503.03it/s]\n" ] }, { @@ -1064,7 +1079,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 21160/21160 [00:00<00:00, 72784.57it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 21160/21160 [00:00<00:00, 72784.57it/s]\n" ] } ], @@ -4917,7 +4932,7 @@ "# Calculate the linear regression and stats\n", "slope, intercept, r_value, p_value, std_err = stats.linregress(x, y)\n", "r_squared = r_value**2\n", - "ax[0].text(0.05, 0.95, f'R² = {r_squared:.3f}\\np-value = {p_value:.3e}', transform=ax[0].transAxes, verticalalignment='top')\n", + "ax[0].text(0.05, 0.95, f'R\u00b2 = {r_squared:.3f}\\np-value = {p_value:.3e}', transform=ax[0].transAxes, verticalalignment='top')\n", "\n", "# Second subplot\n", "sns.regplot(x=x, y=w, ax=ax[1], scatter=False)\n", @@ -4927,7 +4942,7 @@ "# Calculate the linear regression and stats for the second plot\n", "slope, intercept, r_value, p_value, std_err = stats.linregress(x, w)\n", "r_squared = r_value**2\n", - "ax[1].text(0.05, 0.95, f'R² = {r_squared:.3f}\\np-value = {p_value:.3e}', transform=ax[1].transAxes, verticalalignment='top')\n", + "ax[1].text(0.05, 0.95, f'R\u00b2 = {r_squared:.3f}\\np-value = {p_value:.3e}', transform=ax[1].transAxes, verticalalignment='top')\n", "\n", "fig.tight_layout()\n", "plt.savefig(\"/Users/ceglian/conga.pdf\")" @@ -5099,7 +5114,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 21160/21160 [00:00<00:00, 68329.44it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 21160/21160 [00:00<00:00, 68329.44it/s]\n" ] }, { @@ -5369,7 +5384,7 @@ " \n", " \n", "\n", - "

441 rows × 5 columns

\n", + "

441 rows \u00d7 5 columns

\n", "" ], "text/plain": [ @@ -5492,4 +5507,4 @@ }, "nbformat": 4, "nbformat_minor": 2 -} +} \ No newline at end of file diff --git a/example/colitis/preprocess_conga.ipynb b/example/colitis/preprocess_conga.ipynb index 70a96c8..33980fb 100644 --- a/example/colitis/preprocess_conga.ipynb +++ b/example/colitis/preprocess_conga.ipynb @@ -1,5 +1,20 @@ { "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# \ud83d\udeab DO NOT READ \u2014 DISPOSABLE NOTEBOOK \ud83d\udeab\n", + "\n", + "**This notebook calls an OLD/DIVERGENT tcri API and is scheduled to be rewritten from scratch against the new API. It is NOT a source of truth for anything.**\n", + "\n", + "- Do NOT read it to decide what functions to keep or drop.\n", + "- Do NOT use it for 'sufficiency' or 'is-it-used' evidence.\n", + "- A function being called here is NOT a reason to keep it.\n", + "\n", + "See `docs/contract/REFACTOR_HISTORY.md` (Hard Rules).\n" + ] + }, { "cell_type": "code", "execution_count": 1, @@ -155,7 +170,7 @@ { "data": { "text/plain": [ - "AnnData object with n_obs × n_vars = 2896 × 1022\n", + "AnnData object with n_obs \u00d7 n_vars = 2896 \u00d7 1022\n", " obs: 'va', 'ja', 'cdr3a', 'cdr3a_nucseq', 'vb', 'jb', 'cdr3b', 'cdr3b_nucseq', 'n_genes', 'percent_mito', 'n_counts', 'clone_sizes', 'gex_variation'\n", " var: 'gene_ids', 'feature_types', 'genome', 'n_cells', 'highly_variable', 'means', 'dispersions', 'dispersions_norm', 'mean', 'std'\n", " uns: 'conga_results', 'conga_stats', 'organism', 'log1p', 'hvg', 'raw_matrix_is_logged', 'pca'\n", @@ -1181,7 +1196,7 @@ { "data": { "text/plain": [ - "AnnData object with n_obs × n_vars = 2896 × 1022\n", + "AnnData object with n_obs \u00d7 n_vars = 2896 \u00d7 1022\n", " obs: 'va', 'ja', 'cdr3a', 'cdr3a_nucseq', 'vb', 'jb', 'cdr3b', 'cdr3b_nucseq', 'n_genes', 'percent_mito', 'n_counts', 'clone_sizes', 'gex_variation', 'louvain_gex', 'clusters_gex', 'louvain_tcr', 'clusters_tcr', 'nndists_gex', 'nndists_tcr', 'is_invariant', 'conga_scores', 'conga_fdr_values'\n", " var: 'gene_ids', 'feature_types', 'genome', 'n_cells', 'highly_variable', 'means', 'dispersions', 'dispersions_norm', 'mean', 'std'\n", " uns: 'conga_results', 'conga_stats', 'organism', 'log1p', 'hvg', 'raw_matrix_is_logged', 'pca', 'neighbors', 'umap', 'louvain', 'clusters_tcr_names', 'rank_genes_good_biclusters'\n", @@ -1639,7 +1654,7 @@ " \n", " \n", "\n", - "

2896 rows × 22 columns

\n", + "

2896 rows \u00d7 22 columns

\n", "" ], "text/plain": [ @@ -1777,13 +1792,13 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 994/994 [00:00<00:00, 5889.40it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 994/994 [00:00<00:00, 5889.40it/s]\n" ] }, { "data": { "text/plain": [ - "View of AnnData object with n_obs × n_vars = 2896 × 994\n", + "View of AnnData object with n_obs \u00d7 n_vars = 2896 \u00d7 994\n", " obs: 'va', 'ja', 'cdr3a', 'cdr3a_nucseq', 'vb', 'jb', 'cdr3b', 'cdr3b_nucseq', 'n_genes', 'percent_mito', 'n_counts', 'clone_sizes', 'gex_variation', 'louvain_gex', 'clusters_gex', 'louvain_tcr', 'clusters_tcr', 'nndists_gex', 'nndists_tcr', 'is_invariant', 'conga_scores', 'conga_fdr_values'\n", " var: 'gene_ids', 'feature_types', 'genome', 'n_cells', 'highly_variable', 'means', 'dispersions', 'dispersions_norm', 'mean', 'std', 'entropy'\n", " uns: 'conga_results', 'conga_stats', 'organism', 'log1p', 'hvg', 'raw_matrix_is_logged', 'pca', 'neighbors', 'umap', 'louvain', 'clusters_tcr_names', 'rank_genes_good_biclusters'\n", @@ -1838,7 +1853,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 2878624/2878624 [00:01<00:00, 1926297.42it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 2878624/2878624 [00:01<00:00, 1926297.42it/s]\n" ] }, { @@ -1889,7 +1904,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 994/994 [00:00<00:00, 6151.26it/s]" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 994/994 [00:00<00:00, 6151.26it/s]" ] }, { @@ -1917,7 +1932,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 493521/493521 [05:59<00:00, 1374.13it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 493521/493521 [05:59<00:00, 1374.13it/s]\n" ] }, { @@ -1933,7 +1948,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 994/994 [00:00<00:00, 1742.86it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 994/994 [00:00<00:00, 1742.86it/s]\n" ] }, { @@ -1977,7 +1992,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 994/994 [00:00<00:00, 3824897.41it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 994/994 [00:00<00:00, 3824897.41it/s]\n" ] }, { @@ -2014,7 +2029,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 994/994 [00:00<00:00, 2948471.13it/s]" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 994/994 [00:00<00:00, 2948471.13it/s]" ] }, { @@ -2029,7 +2044,7 @@ "output_type": "stream", "text": [ "\n", - "100%|██████████| 2878624/2878624 [00:03<00:00, 764303.13it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 2878624/2878624 [00:03<00:00, 764303.13it/s]\n" ] }, { @@ -2043,7 +2058,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 2896/2896 [00:00<00:00, 3314.29it/s]" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 2896/2896 [00:00<00:00, 3314.29it/s]" ] }, { @@ -2083,7 +2098,7 @@ { "data": { "text/plain": [ - "AnnData object with n_obs × n_vars = 2896 × 994\n", + "AnnData object with n_obs \u00d7 n_vars = 2896 \u00d7 994\n", " obs: 'va', 'ja', 'cdr3a', 'cdr3a_nucseq', 'vb', 'jb', 'cdr3b', 'cdr3b_nucseq', 'n_genes', 'percent_mito', 'n_counts', 'clone_sizes', 'gex_variation', 'louvain_gex', 'clusters_gex', 'louvain_tcr', 'clusters_tcr', 'nndists_gex', 'nndists_tcr', 'is_invariant', 'conga_scores', 'conga_fdr_values'\n", " var: 'gene_ids', 'feature_types', 'genome', 'n_cells', 'highly_variable', 'means', 'dispersions', 'dispersions_norm', 'mean', 'std', 'entropy'\n", " uns: 'conga_results', 'conga_stats', 'organism', 'log1p', 'hvg', 'raw_matrix_is_logged', 'pca', 'neighbors', 'umap', 'louvain', 'clusters_tcr_names', 'rank_genes_good_biclusters', 'tcri_clone_key', 'tcri_unique_clonotypes'\n", @@ -2139,7 +2154,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 2896/2896 [00:00<00:00, 63867.16it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 2896/2896 [00:00<00:00, 63867.16it/s]\n" ] }, { @@ -2154,7 +2169,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 2896/2896 [00:00<00:00, 70021.12it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 2896/2896 [00:00<00:00, 70021.12it/s]\n" ] }, { @@ -2169,7 +2184,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 2896/2896 [00:00<00:00, 70206.08it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 2896/2896 [00:00<00:00, 70206.08it/s]\n" ] }, { @@ -2184,7 +2199,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 2896/2896 [00:00<00:00, 70867.17it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 2896/2896 [00:00<00:00, 70867.17it/s]\n" ] } ], @@ -2438,4 +2453,4 @@ }, "nbformat": 4, "nbformat_minor": 2 -} +} \ No newline at end of file diff --git a/example/renal/renal.ipynb b/example/renal/renal.ipynb index 516d0c3..58e4cff 100644 --- a/example/renal/renal.ipynb +++ b/example/renal/renal.ipynb @@ -1,5 +1,20 @@ { "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# \ud83d\udeab DO NOT READ \u2014 DISPOSABLE NOTEBOOK \ud83d\udeab\n", + "\n", + "**This notebook calls an OLD/DIVERGENT tcri API and is scheduled to be rewritten from scratch against the new API. It is NOT a source of truth for anything.**\n", + "\n", + "- Do NOT read it to decide what functions to keep or drop.\n", + "- Do NOT use it for 'sufficiency' or 'is-it-used' evidence.\n", + "- A function being called here is NOT a reason to keep it.\n", + "\n", + "See `docs/contract/REFACTOR_HISTORY.md` (Hard Rules).\n" + ] + }, { "cell_type": "code", "execution_count": 2, @@ -23,7 +38,7 @@ { "data": { "text/plain": [ - "AnnData object with n_obs × n_vars = 12293 × 15729\n", + "AnnData object with n_obs \u00d7 n_vars = 12293 \u00d7 15729\n", " obs: 'orig.ident', 'nCount_RNA', 'nFeature_RNA', 'CTgene', 'CTnt', 'CTaa', 'CTstrict', 'sample', 'barcode', 'Frequency', 'cloneType', 'mito.genes', 'db.class', 'db.score', 'PT.annot', 'PT.score', 'HPCA.labels', 'HPCA.pruned.labels', 'Monaco.labels', 'Monaco.pruned.labels', 'db.weight.score', 'db.ratio', 'HPCA.first.labels', 'functional.cluster', 'functional.cluster.conf', 'ID', 'T.filter', 'Tissue', 'GEO_RNA', 'Cohort', 'Sample', 'Treatments', 'Response', 'T', 'SampleID'\n", " var: 'features', 'entropy-BCT1.3', 'entropy-BCT1.5', 'entropy-BCT10.1', 'entropy-BCT2.1', 'entropy-BCT2.2', 'entropy-BCT2.3', 'entropy-BCT3.1', 'entropy-BCT4.1', 'entropy-BCT4.2', 'entropy-BCT4.3', 'entropy-BCT5.1', 'entropy-BCT5.2', 'entropy-BCT5.3', 'entropy-BCT5.4', 'entropy-BCT5.5', 'entropy-BCT6.1', 'entropy-BCT6.2', 'entropy-BCT6.4', 'entropy-BCT6.5', 'entropy-BCT7.1', 'entropy-BCT7.2', 'entropy-BCT7.3', 'entropy-BCT8.1', 'entropy-BCT8.2', 'entropy-BCT8.3', 'entropy-BCT8.5', 'entropy-BCT9.1', 'entropy-BT10.1', 'entropy-BT10.2', 'entropy-BT11.1', 'entropy-BT12.2', 'entropy-BT13.2', 'entropy-BT14.1', 'entropy-BT14.2', 'entropy-BT16.2', 'entropy-BT17.1', 'entropy-BT17.2', 'entropy-BT18.2', 'entropy-BT19.2', 'entropy-BT20.2', 'entropy-BT21.2', 'entropy-BT22.1', 'entropy-BT22.2', 'entropy-BT23.1', 'entropy-BT23.2', 'entropy-BT25.1', 'entropy-BT26.2', 'entropy-BT27.1', 'entropy-BT28.2', 'entropy-BT29.1', 'entropy-BT30.2', 'entropy-BT31.2', 'entropy-BT32.1', 'entropy-BT32.2', 'entropy-BT33.1', 'entropy-BT33.2', 'entropy-BT35.1', 'entropy-BT35.2', 'entropy-BT36.1', 'entropy-BT36.2', 'entropy-BT37.1', 'entropy-BT38.1', 'entropy-BT38.2', 'entropy-BT39.1', 'entropy-BT4.1', 'entropy-BT4.2', 'entropy-BT41.2', 'entropy-BT42.2', 'entropy-BT44.1', 'entropy-BT45.1', 'entropy-BT5.1', 'entropy-BT5.2', 'entropy-BT6.2', 'entropy-BT8.2', 'entropy-BT9.1', 'entropy-BT9.2', 'entropy-CT1', 'entropy-EN1', 'entropy-EN2', 'entropy-EN3', 'entropy-ESN1', 'entropy-ESN6', 'entropy-ESN7', 'entropy-EST1', 'entropy-EST3', 'entropy-ET1', 'entropy-ET2', 'entropy-HL2.2', 'entropy-HL2.3', 'entropy-HL3.1', 'entropy-HL4.3', 'entropy-HT1.1', 'entropy-HT1.2', 'entropy-HT2.1', 'entropy-HT2.2', 'entropy-HT3.2', 'entropy-HT3.3', 'entropy-HT4.1', 'entropy-HT4.2', 'entropy-LB14.1', 'entropy-LB14.14', 'entropy-LB14.3', 'entropy-LB14.6', 'entropy-LB14.8', 'entropy-LB8', 'entropy-LB9', 'entropy-LJ8', 'entropy-LJ9', 'entropy-LL13.3', 'entropy-LL14.1', 'entropy-LL14.2', 'entropy-LL14.4', 'entropy-LL18.1', 'entropy-LL18.2', 'entropy-LL18.3', 'entropy-LN1', 'entropy-LN10', 'entropy-LN11', 'entropy-LN12', 'entropy-LN14.2', 'entropy-LN14.3', 'entropy-LN14.4', 'entropy-LN14.5', 'entropy-LN16.1', 'entropy-LN16.2', 'entropy-LN18.1', 'entropy-LN18.2', 'entropy-LN18.5', 'entropy-LN18.6', 'entropy-LN18.8', 'entropy-LN19', 'entropy-LN2', 'entropy-LN22', 'entropy-LN3', 'entropy-LN4', 'entropy-LT1', 'entropy-LT11', 'entropy-LT12', 'entropy-LT13.2', 'entropy-LT13.3', 'entropy-LT13.4', 'entropy-LT13.6', 'entropy-LT13.7', 'entropy-LT13.9', 'entropy-LT14.1', 'entropy-LT14.2', 'entropy-LT14.3', 'entropy-LT14.4', 'entropy-LT14.5', 'entropy-LT14.6', 'entropy-LT14.8', 'entropy-LT16.1', 'entropy-LT16.2', 'entropy-LT16.3', 'entropy-LT16.7', 'entropy-LT16.9', 'entropy-LT17', 'entropy-LT18.10', 'entropy-LT18.3', 'entropy-LT18.4', 'entropy-LT18.6', 'entropy-LT18.8', 'entropy-LT18.9', 'entropy-LT19.1', 'entropy-LT19.2', 'entropy-LT19.3', 'entropy-LT2', 'entropy-LT20.2', 'entropy-LT20.3', 'entropy-LT21.1', 'entropy-LT21.2', 'entropy-LT21.3', 'entropy-LT21.4', 'entropy-LT23', 'entropy-LT24.1', 'entropy-LT24.4', 'entropy-LT25.1', 'entropy-LT25.2', 'entropy-LT25.4', 'entropy-LT26.3', 'entropy-LT26.4', 'entropy-LT29', 'entropy-LT31', 'entropy-LT32', 'entropy-LT33', 'entropy-LT35.1', 'entropy-LT35.2', 'entropy-LT36', 'entropy-LT38.1', 'entropy-LT40', 'entropy-LT41', 'entropy-LT42', 'entropy-LT43', 'entropy-LT46', 'entropy-LT49', 'entropy-LT50', 'entropy-LT51', 'entropy-LT52', 'entropy-LT53', 'entropy-LT54.1', 'entropy-LT54.2', 'entropy-LT55.1', 'entropy-LT55.2', 'entropy-LT56.1', 'entropy-LT56.2', 'entropy-LT58.1', 'entropy-LT59', 'entropy-LT60', 'entropy-LT61', 'entropy-LT7', 'entropy-LT8', 'entropy-LT9', 'entropy-MB1', 'entropy-MB2.2', 'entropy-MB3.1', 'entropy-ML4', 'entropy-MT1', 'entropy-MT3', 'entropy-OT3', 'entropy-OT4', 'entropy-OT5', 'entropy-OT6', 'entropy-RB1', 'entropy-RB10', 'entropy-RB11', 'entropy-RB12', 'entropy-RB13.1', 'entropy-RB14.1', 'entropy-RB15.1', 'entropy-RB15.2', 'entropy-RB16', 'entropy-RB2', 'entropy-RB9', 'entropy-RL10', 'entropy-RN1', 'entropy-RN10', 'entropy-RN12', 'entropy-RN13.1', 'entropy-RN14.2', 'entropy-RN15.1', 'entropy-RN3', 'entropy-RN7', 'entropy-RN8', 'entropy-RN9', 'entropy-RT1', 'entropy-RT10.1', 'entropy-RT10.3', 'entropy-RT11.1', 'entropy-RT11.2', 'entropy-RT11.3', 'entropy-RT12.1', 'entropy-RT12.2', 'entropy-RT13.1', 'entropy-RT13.2', 'entropy-RT14.1', 'entropy-RT15.1', 'entropy-RT16.1', 'entropy-RT2', 'entropy-RT5', 'entropy-RT7.1', 'entropy-RT7.2', 'entropy-RT8.3', 'entropy-RT9.1', 'entropy-RT9.2', 'entropy-RT9.3', 'entropy-SCT1.3', 'entropy-SCT1.4', 'entropy-SCT1.6', 'entropy-SCT2.1', 'entropy-SCT3.1', 'entropy-SCT3.2', 'entropy-SCT4.1', 'entropy-SCT4.2'" ] @@ -80,7 +95,7 @@ "name": "stdout", "output_type": "stream", "text": [ - "AnnData object with n_obs × n_vars = 12293 × 325\n", + "AnnData object with n_obs \u00d7 n_vars = 12293 \u00d7 325\n", " obs: 'orig.ident', 'nCount_RNA', 'nFeature_RNA', 'CTgene', 'CTnt', 'CTaa', 'CTstrict', 'sample', 'barcode', 'Frequency', 'cloneType', 'mito.genes', 'db.class', 'db.score', 'PT.annot', 'PT.score', 'HPCA.labels', 'HPCA.pruned.labels', 'Monaco.labels', 'Monaco.pruned.labels', 'db.weight.score', 'db.ratio', 'HPCA.first.labels', 'functional.cluster', 'functional.cluster.conf', 'ID', 'T.filter', 'Tissue', 'GEO_RNA', 'Cohort', 'Sample', 'Treatments', 'Response', 'T', 'SampleID', 'trb'\n", " var: 'features', 'entropy-BCT1.3', 'entropy-BCT1.5', 'entropy-BCT10.1', 'entropy-BCT2.1', 'entropy-BCT2.2', 'entropy-BCT2.3', 'entropy-BCT3.1', 'entropy-BCT4.1', 'entropy-BCT4.2', 'entropy-BCT4.3', 'entropy-BCT5.1', 'entropy-BCT5.2', 'entropy-BCT5.3', 'entropy-BCT5.4', 'entropy-BCT5.5', 'entropy-BCT6.1', 'entropy-BCT6.2', 'entropy-BCT6.4', 'entropy-BCT6.5', 'entropy-BCT7.1', 'entropy-BCT7.2', 'entropy-BCT7.3', 'entropy-BCT8.1', 'entropy-BCT8.2', 'entropy-BCT8.3', 'entropy-BCT8.5', 'entropy-BCT9.1', 'entropy-BT10.1', 'entropy-BT10.2', 'entropy-BT11.1', 'entropy-BT12.2', 'entropy-BT13.2', 'entropy-BT14.1', 'entropy-BT14.2', 'entropy-BT16.2', 'entropy-BT17.1', 'entropy-BT17.2', 'entropy-BT18.2', 'entropy-BT19.2', 'entropy-BT20.2', 'entropy-BT21.2', 'entropy-BT22.1', 'entropy-BT22.2', 'entropy-BT23.1', 'entropy-BT23.2', 'entropy-BT25.1', 'entropy-BT26.2', 'entropy-BT27.1', 'entropy-BT28.2', 'entropy-BT29.1', 'entropy-BT30.2', 'entropy-BT31.2', 'entropy-BT32.1', 'entropy-BT32.2', 'entropy-BT33.1', 'entropy-BT33.2', 'entropy-BT35.1', 'entropy-BT35.2', 'entropy-BT36.1', 'entropy-BT36.2', 'entropy-BT37.1', 'entropy-BT38.1', 'entropy-BT38.2', 'entropy-BT39.1', 'entropy-BT4.1', 'entropy-BT4.2', 'entropy-BT41.2', 'entropy-BT42.2', 'entropy-BT44.1', 'entropy-BT45.1', 'entropy-BT5.1', 'entropy-BT5.2', 'entropy-BT6.2', 'entropy-BT8.2', 'entropy-BT9.1', 'entropy-BT9.2', 'entropy-CT1', 'entropy-EN1', 'entropy-EN2', 'entropy-EN3', 'entropy-ESN1', 'entropy-ESN6', 'entropy-ESN7', 'entropy-EST1', 'entropy-EST3', 'entropy-ET1', 'entropy-ET2', 'entropy-HL2.2', 'entropy-HL2.3', 'entropy-HL3.1', 'entropy-HL4.3', 'entropy-HT1.1', 'entropy-HT1.2', 'entropy-HT2.1', 'entropy-HT2.2', 'entropy-HT3.2', 'entropy-HT3.3', 'entropy-HT4.1', 'entropy-HT4.2', 'entropy-LB14.1', 'entropy-LB14.14', 'entropy-LB14.3', 'entropy-LB14.6', 'entropy-LB14.8', 'entropy-LB8', 'entropy-LB9', 'entropy-LJ8', 'entropy-LJ9', 'entropy-LL13.3', 'entropy-LL14.1', 'entropy-LL14.2', 'entropy-LL14.4', 'entropy-LL18.1', 'entropy-LL18.2', 'entropy-LL18.3', 'entropy-LN1', 'entropy-LN10', 'entropy-LN11', 'entropy-LN12', 'entropy-LN14.2', 'entropy-LN14.3', 'entropy-LN14.4', 'entropy-LN14.5', 'entropy-LN16.1', 'entropy-LN16.2', 'entropy-LN18.1', 'entropy-LN18.2', 'entropy-LN18.5', 'entropy-LN18.6', 'entropy-LN18.8', 'entropy-LN19', 'entropy-LN2', 'entropy-LN22', 'entropy-LN3', 'entropy-LN4', 'entropy-LT1', 'entropy-LT11', 'entropy-LT12', 'entropy-LT13.2', 'entropy-LT13.3', 'entropy-LT13.4', 'entropy-LT13.6', 'entropy-LT13.7', 'entropy-LT13.9', 'entropy-LT14.1', 'entropy-LT14.2', 'entropy-LT14.3', 'entropy-LT14.4', 'entropy-LT14.5', 'entropy-LT14.6', 'entropy-LT14.8', 'entropy-LT16.1', 'entropy-LT16.2', 'entropy-LT16.3', 'entropy-LT16.7', 'entropy-LT16.9', 'entropy-LT17', 'entropy-LT18.10', 'entropy-LT18.3', 'entropy-LT18.4', 'entropy-LT18.6', 'entropy-LT18.8', 'entropy-LT18.9', 'entropy-LT19.1', 'entropy-LT19.2', 'entropy-LT19.3', 'entropy-LT2', 'entropy-LT20.2', 'entropy-LT20.3', 'entropy-LT21.1', 'entropy-LT21.2', 'entropy-LT21.3', 'entropy-LT21.4', 'entropy-LT23', 'entropy-LT24.1', 'entropy-LT24.4', 'entropy-LT25.1', 'entropy-LT25.2', 'entropy-LT25.4', 'entropy-LT26.3', 'entropy-LT26.4', 'entropy-LT29', 'entropy-LT31', 'entropy-LT32', 'entropy-LT33', 'entropy-LT35.1', 'entropy-LT35.2', 'entropy-LT36', 'entropy-LT38.1', 'entropy-LT40', 'entropy-LT41', 'entropy-LT42', 'entropy-LT43', 'entropy-LT46', 'entropy-LT49', 'entropy-LT50', 'entropy-LT51', 'entropy-LT52', 'entropy-LT53', 'entropy-LT54.1', 'entropy-LT54.2', 'entropy-LT55.1', 'entropy-LT55.2', 'entropy-LT56.1', 'entropy-LT56.2', 'entropy-LT58.1', 'entropy-LT59', 'entropy-LT60', 'entropy-LT61', 'entropy-LT7', 'entropy-LT8', 'entropy-LT9', 'entropy-MB1', 'entropy-MB2.2', 'entropy-MB3.1', 'entropy-ML4', 'entropy-MT1', 'entropy-MT3', 'entropy-OT3', 'entropy-OT4', 'entropy-OT5', 'entropy-OT6', 'entropy-RB1', 'entropy-RB10', 'entropy-RB11', 'entropy-RB12', 'entropy-RB13.1', 'entropy-RB14.1', 'entropy-RB15.1', 'entropy-RB15.2', 'entropy-RB16', 'entropy-RB2', 'entropy-RB9', 'entropy-RL10', 'entropy-RN1', 'entropy-RN10', 'entropy-RN12', 'entropy-RN13.1', 'entropy-RN14.2', 'entropy-RN15.1', 'entropy-RN3', 'entropy-RN7', 'entropy-RN8', 'entropy-RN9', 'entropy-RT1', 'entropy-RT10.1', 'entropy-RT10.3', 'entropy-RT11.1', 'entropy-RT11.2', 'entropy-RT11.3', 'entropy-RT12.1', 'entropy-RT12.2', 'entropy-RT13.1', 'entropy-RT13.2', 'entropy-RT14.1', 'entropy-RT15.1', 'entropy-RT16.1', 'entropy-RT2', 'entropy-RT5', 'entropy-RT7.1', 'entropy-RT7.2', 'entropy-RT8.3', 'entropy-RT9.1', 'entropy-RT9.2', 'entropy-RT9.3', 'entropy-SCT1.3', 'entropy-SCT1.4', 'entropy-SCT1.6', 'entropy-SCT2.1', 'entropy-SCT3.1', 'entropy-SCT3.2', 'entropy-SCT4.1', 'entropy-SCT4.2', 'entropy'\n", " uns: 'tcri_clone_key', 'tcri_unique_clonotypes'\n" @@ -89,7 +104,7 @@ { "data": { "text/plain": [ - "View of AnnData object with n_obs × n_vars = 12293 × 325\n", + "View of AnnData object with n_obs \u00d7 n_vars = 12293 \u00d7 325\n", " obs: 'orig.ident', 'nCount_RNA', 'nFeature_RNA', 'CTgene', 'CTnt', 'CTaa', 'CTstrict', 'sample', 'barcode', 'Frequency', 'cloneType', 'mito.genes', 'db.class', 'db.score', 'PT.annot', 'PT.score', 'HPCA.labels', 'HPCA.pruned.labels', 'Monaco.labels', 'Monaco.pruned.labels', 'db.weight.score', 'db.ratio', 'HPCA.first.labels', 'functional.cluster', 'functional.cluster.conf', 'ID', 'T.filter', 'Tissue', 'GEO_RNA', 'Cohort', 'Sample', 'Treatments', 'Response', 'T', 'SampleID', 'trb'\n", " var: 'features', 'entropy-BCT1.3', 'entropy-BCT1.5', 'entropy-BCT10.1', 'entropy-BCT2.1', 'entropy-BCT2.2', 'entropy-BCT2.3', 'entropy-BCT3.1', 'entropy-BCT4.1', 'entropy-BCT4.2', 'entropy-BCT4.3', 'entropy-BCT5.1', 'entropy-BCT5.2', 'entropy-BCT5.3', 'entropy-BCT5.4', 'entropy-BCT5.5', 'entropy-BCT6.1', 'entropy-BCT6.2', 'entropy-BCT6.4', 'entropy-BCT6.5', 'entropy-BCT7.1', 'entropy-BCT7.2', 'entropy-BCT7.3', 'entropy-BCT8.1', 'entropy-BCT8.2', 'entropy-BCT8.3', 'entropy-BCT8.5', 'entropy-BCT9.1', 'entropy-BT10.1', 'entropy-BT10.2', 'entropy-BT11.1', 'entropy-BT12.2', 'entropy-BT13.2', 'entropy-BT14.1', 'entropy-BT14.2', 'entropy-BT16.2', 'entropy-BT17.1', 'entropy-BT17.2', 'entropy-BT18.2', 'entropy-BT19.2', 'entropy-BT20.2', 'entropy-BT21.2', 'entropy-BT22.1', 'entropy-BT22.2', 'entropy-BT23.1', 'entropy-BT23.2', 'entropy-BT25.1', 'entropy-BT26.2', 'entropy-BT27.1', 'entropy-BT28.2', 'entropy-BT29.1', 'entropy-BT30.2', 'entropy-BT31.2', 'entropy-BT32.1', 'entropy-BT32.2', 'entropy-BT33.1', 'entropy-BT33.2', 'entropy-BT35.1', 'entropy-BT35.2', 'entropy-BT36.1', 'entropy-BT36.2', 'entropy-BT37.1', 'entropy-BT38.1', 'entropy-BT38.2', 'entropy-BT39.1', 'entropy-BT4.1', 'entropy-BT4.2', 'entropy-BT41.2', 'entropy-BT42.2', 'entropy-BT44.1', 'entropy-BT45.1', 'entropy-BT5.1', 'entropy-BT5.2', 'entropy-BT6.2', 'entropy-BT8.2', 'entropy-BT9.1', 'entropy-BT9.2', 'entropy-CT1', 'entropy-EN1', 'entropy-EN2', 'entropy-EN3', 'entropy-ESN1', 'entropy-ESN6', 'entropy-ESN7', 'entropy-EST1', 'entropy-EST3', 'entropy-ET1', 'entropy-ET2', 'entropy-HL2.2', 'entropy-HL2.3', 'entropy-HL3.1', 'entropy-HL4.3', 'entropy-HT1.1', 'entropy-HT1.2', 'entropy-HT2.1', 'entropy-HT2.2', 'entropy-HT3.2', 'entropy-HT3.3', 'entropy-HT4.1', 'entropy-HT4.2', 'entropy-LB14.1', 'entropy-LB14.14', 'entropy-LB14.3', 'entropy-LB14.6', 'entropy-LB14.8', 'entropy-LB8', 'entropy-LB9', 'entropy-LJ8', 'entropy-LJ9', 'entropy-LL13.3', 'entropy-LL14.1', 'entropy-LL14.2', 'entropy-LL14.4', 'entropy-LL18.1', 'entropy-LL18.2', 'entropy-LL18.3', 'entropy-LN1', 'entropy-LN10', 'entropy-LN11', 'entropy-LN12', 'entropy-LN14.2', 'entropy-LN14.3', 'entropy-LN14.4', 'entropy-LN14.5', 'entropy-LN16.1', 'entropy-LN16.2', 'entropy-LN18.1', 'entropy-LN18.2', 'entropy-LN18.5', 'entropy-LN18.6', 'entropy-LN18.8', 'entropy-LN19', 'entropy-LN2', 'entropy-LN22', 'entropy-LN3', 'entropy-LN4', 'entropy-LT1', 'entropy-LT11', 'entropy-LT12', 'entropy-LT13.2', 'entropy-LT13.3', 'entropy-LT13.4', 'entropy-LT13.6', 'entropy-LT13.7', 'entropy-LT13.9', 'entropy-LT14.1', 'entropy-LT14.2', 'entropy-LT14.3', 'entropy-LT14.4', 'entropy-LT14.5', 'entropy-LT14.6', 'entropy-LT14.8', 'entropy-LT16.1', 'entropy-LT16.2', 'entropy-LT16.3', 'entropy-LT16.7', 'entropy-LT16.9', 'entropy-LT17', 'entropy-LT18.10', 'entropy-LT18.3', 'entropy-LT18.4', 'entropy-LT18.6', 'entropy-LT18.8', 'entropy-LT18.9', 'entropy-LT19.1', 'entropy-LT19.2', 'entropy-LT19.3', 'entropy-LT2', 'entropy-LT20.2', 'entropy-LT20.3', 'entropy-LT21.1', 'entropy-LT21.2', 'entropy-LT21.3', 'entropy-LT21.4', 'entropy-LT23', 'entropy-LT24.1', 'entropy-LT24.4', 'entropy-LT25.1', 'entropy-LT25.2', 'entropy-LT25.4', 'entropy-LT26.3', 'entropy-LT26.4', 'entropy-LT29', 'entropy-LT31', 'entropy-LT32', 'entropy-LT33', 'entropy-LT35.1', 'entropy-LT35.2', 'entropy-LT36', 'entropy-LT38.1', 'entropy-LT40', 'entropy-LT41', 'entropy-LT42', 'entropy-LT43', 'entropy-LT46', 'entropy-LT49', 'entropy-LT50', 'entropy-LT51', 'entropy-LT52', 'entropy-LT53', 'entropy-LT54.1', 'entropy-LT54.2', 'entropy-LT55.1', 'entropy-LT55.2', 'entropy-LT56.1', 'entropy-LT56.2', 'entropy-LT58.1', 'entropy-LT59', 'entropy-LT60', 'entropy-LT61', 'entropy-LT7', 'entropy-LT8', 'entropy-LT9', 'entropy-MB1', 'entropy-MB2.2', 'entropy-MB3.1', 'entropy-ML4', 'entropy-MT1', 'entropy-MT3', 'entropy-OT3', 'entropy-OT4', 'entropy-OT5', 'entropy-OT6', 'entropy-RB1', 'entropy-RB10', 'entropy-RB11', 'entropy-RB12', 'entropy-RB13.1', 'entropy-RB14.1', 'entropy-RB15.1', 'entropy-RB15.2', 'entropy-RB16', 'entropy-RB2', 'entropy-RB9', 'entropy-RL10', 'entropy-RN1', 'entropy-RN10', 'entropy-RN12', 'entropy-RN13.1', 'entropy-RN14.2', 'entropy-RN15.1', 'entropy-RN3', 'entropy-RN7', 'entropy-RN8', 'entropy-RN9', 'entropy-RT1', 'entropy-RT10.1', 'entropy-RT10.3', 'entropy-RT11.1', 'entropy-RT11.2', 'entropy-RT11.3', 'entropy-RT12.1', 'entropy-RT12.2', 'entropy-RT13.1', 'entropy-RT13.2', 'entropy-RT14.1', 'entropy-RT15.1', 'entropy-RT16.1', 'entropy-RT2', 'entropy-RT5', 'entropy-RT7.1', 'entropy-RT7.2', 'entropy-RT8.3', 'entropy-RT9.1', 'entropy-RT9.2', 'entropy-RT9.3', 'entropy-SCT1.3', 'entropy-SCT1.4', 'entropy-SCT1.6', 'entropy-SCT2.1', 'entropy-SCT3.1', 'entropy-SCT3.2', 'entropy-SCT4.1', 'entropy-SCT4.2', 'entropy'\n", " uns: 'tcri_clone_key', 'tcri_unique_clonotypes'" @@ -148,13 +163,13 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 325/325 [00:00<00:00, 3869.78it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 325/325 [00:00<00:00, 3869.78it/s]\n" ] }, { "data": { "text/plain": [ - "View of AnnData object with n_obs × n_vars = 12293 × 325\n", + "View of AnnData object with n_obs \u00d7 n_vars = 12293 \u00d7 325\n", " obs: 'orig.ident', 'nCount_RNA', 'nFeature_RNA', 'CTgene', 'CTnt', 'CTaa', 'CTstrict', 'sample', 'barcode', 'Frequency', 'cloneType', 'mito.genes', 'db.class', 'db.score', 'PT.annot', 'PT.score', 'HPCA.labels', 'HPCA.pruned.labels', 'Monaco.labels', 'Monaco.pruned.labels', 'db.weight.score', 'db.ratio', 'HPCA.first.labels', 'functional.cluster', 'functional.cluster.conf', 'ID', 'T.filter', 'Tissue', 'GEO_RNA', 'Cohort', 'Sample', 'Treatments', 'Response', 'T', 'SampleID', 'trb'\n", " var: 'features', 'entropy-BCT1.3', 'entropy-BCT1.5', 'entropy-BCT10.1', 'entropy-BCT2.1', 'entropy-BCT2.2', 'entropy-BCT2.3', 'entropy-BCT3.1', 'entropy-BCT4.1', 'entropy-BCT4.2', 'entropy-BCT4.3', 'entropy-BCT5.1', 'entropy-BCT5.2', 'entropy-BCT5.3', 'entropy-BCT5.4', 'entropy-BCT5.5', 'entropy-BCT6.1', 'entropy-BCT6.2', 'entropy-BCT6.4', 'entropy-BCT6.5', 'entropy-BCT7.1', 'entropy-BCT7.2', 'entropy-BCT7.3', 'entropy-BCT8.1', 'entropy-BCT8.2', 'entropy-BCT8.3', 'entropy-BCT8.5', 'entropy-BCT9.1', 'entropy-BT10.1', 'entropy-BT10.2', 'entropy-BT11.1', 'entropy-BT12.2', 'entropy-BT13.2', 'entropy-BT14.1', 'entropy-BT14.2', 'entropy-BT16.2', 'entropy-BT17.1', 'entropy-BT17.2', 'entropy-BT18.2', 'entropy-BT19.2', 'entropy-BT20.2', 'entropy-BT21.2', 'entropy-BT22.1', 'entropy-BT22.2', 'entropy-BT23.1', 'entropy-BT23.2', 'entropy-BT25.1', 'entropy-BT26.2', 'entropy-BT27.1', 'entropy-BT28.2', 'entropy-BT29.1', 'entropy-BT30.2', 'entropy-BT31.2', 'entropy-BT32.1', 'entropy-BT32.2', 'entropy-BT33.1', 'entropy-BT33.2', 'entropy-BT35.1', 'entropy-BT35.2', 'entropy-BT36.1', 'entropy-BT36.2', 'entropy-BT37.1', 'entropy-BT38.1', 'entropy-BT38.2', 'entropy-BT39.1', 'entropy-BT4.1', 'entropy-BT4.2', 'entropy-BT41.2', 'entropy-BT42.2', 'entropy-BT44.1', 'entropy-BT45.1', 'entropy-BT5.1', 'entropy-BT5.2', 'entropy-BT6.2', 'entropy-BT8.2', 'entropy-BT9.1', 'entropy-BT9.2', 'entropy-CT1', 'entropy-EN1', 'entropy-EN2', 'entropy-EN3', 'entropy-ESN1', 'entropy-ESN6', 'entropy-ESN7', 'entropy-EST1', 'entropy-EST3', 'entropy-ET1', 'entropy-ET2', 'entropy-HL2.2', 'entropy-HL2.3', 'entropy-HL3.1', 'entropy-HL4.3', 'entropy-HT1.1', 'entropy-HT1.2', 'entropy-HT2.1', 'entropy-HT2.2', 'entropy-HT3.2', 'entropy-HT3.3', 'entropy-HT4.1', 'entropy-HT4.2', 'entropy-LB14.1', 'entropy-LB14.14', 'entropy-LB14.3', 'entropy-LB14.6', 'entropy-LB14.8', 'entropy-LB8', 'entropy-LB9', 'entropy-LJ8', 'entropy-LJ9', 'entropy-LL13.3', 'entropy-LL14.1', 'entropy-LL14.2', 'entropy-LL14.4', 'entropy-LL18.1', 'entropy-LL18.2', 'entropy-LL18.3', 'entropy-LN1', 'entropy-LN10', 'entropy-LN11', 'entropy-LN12', 'entropy-LN14.2', 'entropy-LN14.3', 'entropy-LN14.4', 'entropy-LN14.5', 'entropy-LN16.1', 'entropy-LN16.2', 'entropy-LN18.1', 'entropy-LN18.2', 'entropy-LN18.5', 'entropy-LN18.6', 'entropy-LN18.8', 'entropy-LN19', 'entropy-LN2', 'entropy-LN22', 'entropy-LN3', 'entropy-LN4', 'entropy-LT1', 'entropy-LT11', 'entropy-LT12', 'entropy-LT13.2', 'entropy-LT13.3', 'entropy-LT13.4', 'entropy-LT13.6', 'entropy-LT13.7', 'entropy-LT13.9', 'entropy-LT14.1', 'entropy-LT14.2', 'entropy-LT14.3', 'entropy-LT14.4', 'entropy-LT14.5', 'entropy-LT14.6', 'entropy-LT14.8', 'entropy-LT16.1', 'entropy-LT16.2', 'entropy-LT16.3', 'entropy-LT16.7', 'entropy-LT16.9', 'entropy-LT17', 'entropy-LT18.10', 'entropy-LT18.3', 'entropy-LT18.4', 'entropy-LT18.6', 'entropy-LT18.8', 'entropy-LT18.9', 'entropy-LT19.1', 'entropy-LT19.2', 'entropy-LT19.3', 'entropy-LT2', 'entropy-LT20.2', 'entropy-LT20.3', 'entropy-LT21.1', 'entropy-LT21.2', 'entropy-LT21.3', 'entropy-LT21.4', 'entropy-LT23', 'entropy-LT24.1', 'entropy-LT24.4', 'entropy-LT25.1', 'entropy-LT25.2', 'entropy-LT25.4', 'entropy-LT26.3', 'entropy-LT26.4', 'entropy-LT29', 'entropy-LT31', 'entropy-LT32', 'entropy-LT33', 'entropy-LT35.1', 'entropy-LT35.2', 'entropy-LT36', 'entropy-LT38.1', 'entropy-LT40', 'entropy-LT41', 'entropy-LT42', 'entropy-LT43', 'entropy-LT46', 'entropy-LT49', 'entropy-LT50', 'entropy-LT51', 'entropy-LT52', 'entropy-LT53', 'entropy-LT54.1', 'entropy-LT54.2', 'entropy-LT55.1', 'entropy-LT55.2', 'entropy-LT56.1', 'entropy-LT56.2', 'entropy-LT58.1', 'entropy-LT59', 'entropy-LT60', 'entropy-LT61', 'entropy-LT7', 'entropy-LT8', 'entropy-LT9', 'entropy-MB1', 'entropy-MB2.2', 'entropy-MB3.1', 'entropy-ML4', 'entropy-MT1', 'entropy-MT3', 'entropy-OT3', 'entropy-OT4', 'entropy-OT5', 'entropy-OT6', 'entropy-RB1', 'entropy-RB10', 'entropy-RB11', 'entropy-RB12', 'entropy-RB13.1', 'entropy-RB14.1', 'entropy-RB15.1', 'entropy-RB15.2', 'entropy-RB16', 'entropy-RB2', 'entropy-RB9', 'entropy-RL10', 'entropy-RN1', 'entropy-RN10', 'entropy-RN12', 'entropy-RN13.1', 'entropy-RN14.2', 'entropy-RN15.1', 'entropy-RN3', 'entropy-RN7', 'entropy-RN8', 'entropy-RN9', 'entropy-RT1', 'entropy-RT10.1', 'entropy-RT10.3', 'entropy-RT11.1', 'entropy-RT11.2', 'entropy-RT11.3', 'entropy-RT12.1', 'entropy-RT12.2', 'entropy-RT13.1', 'entropy-RT13.2', 'entropy-RT14.1', 'entropy-RT15.1', 'entropy-RT16.1', 'entropy-RT2', 'entropy-RT5', 'entropy-RT7.1', 'entropy-RT7.2', 'entropy-RT8.3', 'entropy-RT9.1', 'entropy-RT9.2', 'entropy-RT9.3', 'entropy-SCT1.3', 'entropy-SCT1.4', 'entropy-SCT1.6', 'entropy-SCT2.1', 'entropy-SCT3.1', 'entropy-SCT3.2', 'entropy-SCT4.1', 'entropy-SCT4.2', 'entropy'\n", " uns: 'tcri_clone_key', 'tcri_unique_clonotypes'" @@ -198,7 +213,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 2257452/2257452 [00:01<00:00, 1795570.48it/s]" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 2257452/2257452 [00:01<00:00, 1795570.48it/s]" ] }, { @@ -263,7 +278,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 325/325 [00:00<00:00, 3804.24it/s]" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 325/325 [00:00<00:00, 3804.24it/s]" ] }, { @@ -291,7 +306,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 52650/52650 [01:16<00:00, 690.33it/s] \n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 52650/52650 [01:16<00:00, 690.33it/s] \n" ] }, { @@ -307,7 +322,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 325/325 [00:00<00:00, 4032.71it/s]" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 325/325 [00:00<00:00, 4032.71it/s]" ] }, { @@ -547,7 +562,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 325/325 [00:00<00:00, 1607486.79it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 325/325 [00:00<00:00, 1607486.79it/s]\n" ] }, { @@ -584,8 +599,8 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 325/325 [00:00<00:00, 1456355.56it/s]\n", - "100%|██████████| 2257452/2257452 [00:02<00:00, 896380.63it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 325/325 [00:00<00:00, 1456355.56it/s]\n", + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 2257452/2257452 [00:02<00:00, 896380.63it/s]\n" ] }, { @@ -599,7 +614,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 12293/12293 [00:01<00:00, 9173.44it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 12293/12293 [00:01<00:00, 9173.44it/s]\n" ] }, { @@ -622,7 +637,7 @@ { "data": { "text/plain": [ - "AnnData object with n_obs × n_vars = 12293 × 325\n", + "AnnData object with n_obs \u00d7 n_vars = 12293 \u00d7 325\n", " obs: 'orig.ident', 'nCount_RNA', 'nFeature_RNA', 'CTgene', 'CTnt', 'CTaa', 'CTstrict', 'sample', 'barcode', 'Frequency', 'cloneType', 'mito.genes', 'db.class', 'db.score', 'PT.annot', 'PT.score', 'HPCA.labels', 'HPCA.pruned.labels', 'Monaco.labels', 'Monaco.pruned.labels', 'db.weight.score', 'db.ratio', 'HPCA.first.labels', 'functional.cluster', 'functional.cluster.conf', 'ID', 'T.filter', 'Tissue', 'GEO_RNA', 'Cohort', 'Sample', 'Treatments', 'Response', 'T', 'SampleID', 'trb'\n", " var: 'features', 'entropy-BCT1.3', 'entropy-BCT1.5', 'entropy-BCT10.1', 'entropy-BCT2.1', 'entropy-BCT2.2', 'entropy-BCT2.3', 'entropy-BCT3.1', 'entropy-BCT4.1', 'entropy-BCT4.2', 'entropy-BCT4.3', 'entropy-BCT5.1', 'entropy-BCT5.2', 'entropy-BCT5.3', 'entropy-BCT5.4', 'entropy-BCT5.5', 'entropy-BCT6.1', 'entropy-BCT6.2', 'entropy-BCT6.4', 'entropy-BCT6.5', 'entropy-BCT7.1', 'entropy-BCT7.2', 'entropy-BCT7.3', 'entropy-BCT8.1', 'entropy-BCT8.2', 'entropy-BCT8.3', 'entropy-BCT8.5', 'entropy-BCT9.1', 'entropy-BT10.1', 'entropy-BT10.2', 'entropy-BT11.1', 'entropy-BT12.2', 'entropy-BT13.2', 'entropy-BT14.1', 'entropy-BT14.2', 'entropy-BT16.2', 'entropy-BT17.1', 'entropy-BT17.2', 'entropy-BT18.2', 'entropy-BT19.2', 'entropy-BT20.2', 'entropy-BT21.2', 'entropy-BT22.1', 'entropy-BT22.2', 'entropy-BT23.1', 'entropy-BT23.2', 'entropy-BT25.1', 'entropy-BT26.2', 'entropy-BT27.1', 'entropy-BT28.2', 'entropy-BT29.1', 'entropy-BT30.2', 'entropy-BT31.2', 'entropy-BT32.1', 'entropy-BT32.2', 'entropy-BT33.1', 'entropy-BT33.2', 'entropy-BT35.1', 'entropy-BT35.2', 'entropy-BT36.1', 'entropy-BT36.2', 'entropy-BT37.1', 'entropy-BT38.1', 'entropy-BT38.2', 'entropy-BT39.1', 'entropy-BT4.1', 'entropy-BT4.2', 'entropy-BT41.2', 'entropy-BT42.2', 'entropy-BT44.1', 'entropy-BT45.1', 'entropy-BT5.1', 'entropy-BT5.2', 'entropy-BT6.2', 'entropy-BT8.2', 'entropy-BT9.1', 'entropy-BT9.2', 'entropy-CT1', 'entropy-EN1', 'entropy-EN2', 'entropy-EN3', 'entropy-ESN1', 'entropy-ESN6', 'entropy-ESN7', 'entropy-EST1', 'entropy-EST3', 'entropy-ET1', 'entropy-ET2', 'entropy-HL2.2', 'entropy-HL2.3', 'entropy-HL3.1', 'entropy-HL4.3', 'entropy-HT1.1', 'entropy-HT1.2', 'entropy-HT2.1', 'entropy-HT2.2', 'entropy-HT3.2', 'entropy-HT3.3', 'entropy-HT4.1', 'entropy-HT4.2', 'entropy-LB14.1', 'entropy-LB14.14', 'entropy-LB14.3', 'entropy-LB14.6', 'entropy-LB14.8', 'entropy-LB8', 'entropy-LB9', 'entropy-LJ8', 'entropy-LJ9', 'entropy-LL13.3', 'entropy-LL14.1', 'entropy-LL14.2', 'entropy-LL14.4', 'entropy-LL18.1', 'entropy-LL18.2', 'entropy-LL18.3', 'entropy-LN1', 'entropy-LN10', 'entropy-LN11', 'entropy-LN12', 'entropy-LN14.2', 'entropy-LN14.3', 'entropy-LN14.4', 'entropy-LN14.5', 'entropy-LN16.1', 'entropy-LN16.2', 'entropy-LN18.1', 'entropy-LN18.2', 'entropy-LN18.5', 'entropy-LN18.6', 'entropy-LN18.8', 'entropy-LN19', 'entropy-LN2', 'entropy-LN22', 'entropy-LN3', 'entropy-LN4', 'entropy-LT1', 'entropy-LT11', 'entropy-LT12', 'entropy-LT13.2', 'entropy-LT13.3', 'entropy-LT13.4', 'entropy-LT13.6', 'entropy-LT13.7', 'entropy-LT13.9', 'entropy-LT14.1', 'entropy-LT14.2', 'entropy-LT14.3', 'entropy-LT14.4', 'entropy-LT14.5', 'entropy-LT14.6', 'entropy-LT14.8', 'entropy-LT16.1', 'entropy-LT16.2', 'entropy-LT16.3', 'entropy-LT16.7', 'entropy-LT16.9', 'entropy-LT17', 'entropy-LT18.10', 'entropy-LT18.3', 'entropy-LT18.4', 'entropy-LT18.6', 'entropy-LT18.8', 'entropy-LT18.9', 'entropy-LT19.1', 'entropy-LT19.2', 'entropy-LT19.3', 'entropy-LT2', 'entropy-LT20.2', 'entropy-LT20.3', 'entropy-LT21.1', 'entropy-LT21.2', 'entropy-LT21.3', 'entropy-LT21.4', 'entropy-LT23', 'entropy-LT24.1', 'entropy-LT24.4', 'entropy-LT25.1', 'entropy-LT25.2', 'entropy-LT25.4', 'entropy-LT26.3', 'entropy-LT26.4', 'entropy-LT29', 'entropy-LT31', 'entropy-LT32', 'entropy-LT33', 'entropy-LT35.1', 'entropy-LT35.2', 'entropy-LT36', 'entropy-LT38.1', 'entropy-LT40', 'entropy-LT41', 'entropy-LT42', 'entropy-LT43', 'entropy-LT46', 'entropy-LT49', 'entropy-LT50', 'entropy-LT51', 'entropy-LT52', 'entropy-LT53', 'entropy-LT54.1', 'entropy-LT54.2', 'entropy-LT55.1', 'entropy-LT55.2', 'entropy-LT56.1', 'entropy-LT56.2', 'entropy-LT58.1', 'entropy-LT59', 'entropy-LT60', 'entropy-LT61', 'entropy-LT7', 'entropy-LT8', 'entropy-LT9', 'entropy-MB1', 'entropy-MB2.2', 'entropy-MB3.1', 'entropy-ML4', 'entropy-MT1', 'entropy-MT3', 'entropy-OT3', 'entropy-OT4', 'entropy-OT5', 'entropy-OT6', 'entropy-RB1', 'entropy-RB10', 'entropy-RB11', 'entropy-RB12', 'entropy-RB13.1', 'entropy-RB14.1', 'entropy-RB15.1', 'entropy-RB15.2', 'entropy-RB16', 'entropy-RB2', 'entropy-RB9', 'entropy-RL10', 'entropy-RN1', 'entropy-RN10', 'entropy-RN12', 'entropy-RN13.1', 'entropy-RN14.2', 'entropy-RN15.1', 'entropy-RN3', 'entropy-RN7', 'entropy-RN8', 'entropy-RN9', 'entropy-RT1', 'entropy-RT10.1', 'entropy-RT10.3', 'entropy-RT11.1', 'entropy-RT11.2', 'entropy-RT11.3', 'entropy-RT12.1', 'entropy-RT12.2', 'entropy-RT13.1', 'entropy-RT13.2', 'entropy-RT14.1', 'entropy-RT15.1', 'entropy-RT16.1', 'entropy-RT2', 'entropy-RT5', 'entropy-RT7.1', 'entropy-RT7.2', 'entropy-RT8.3', 'entropy-RT9.1', 'entropy-RT9.2', 'entropy-RT9.3', 'entropy-SCT1.3', 'entropy-SCT1.4', 'entropy-SCT1.6', 'entropy-SCT2.1', 'entropy-SCT3.1', 'entropy-SCT3.2', 'entropy-SCT4.1', 'entropy-SCT4.2', 'entropy'\n", " uns: 'tcri_clone_key', 'tcri_unique_clonotypes', 'neighbors', 'umap'\n", @@ -697,7 +712,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 12293/12293 [00:00<00:00, 42997.86it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 12293/12293 [00:00<00:00, 42997.86it/s]\n" ] }, { @@ -712,7 +727,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 12293/12293 [00:00<00:00, 43713.11it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 12293/12293 [00:00<00:00, 43713.11it/s]\n" ] }, { @@ -727,7 +742,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 12293/12293 [00:00<00:00, 40755.68it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 12293/12293 [00:00<00:00, 40755.68it/s]\n" ] }, { @@ -742,7 +757,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 12293/12293 [00:00<00:00, 34999.65it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 12293/12293 [00:00<00:00, 34999.65it/s]\n" ] } ], @@ -962,7 +977,7 @@ "output_type": "stream", "text": [ "WARNING: Default of the method has been changed to 't-test' from 't-test_overestim_var'\n", - "WARNING: You’re trying to run this on 325 dimensions of `.X`, if you really want this, set `use_rep='X'`.\n", + "WARNING: You\u2019re trying to run this on 325 dimensions of `.X`, if you really want this, set `use_rep='X'`.\n", " Falling back to preprocessing with `sc.pp.pca` and default params.\n", "WARNING: saving figure to file figures/dotplot_gv_degs.pdf\n" ] @@ -1102,7 +1117,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 3/3 [00:00<00:00, 13.00it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 3/3 [00:00<00:00, 13.00it/s]\n" ] }, { @@ -1210,4 +1225,4 @@ }, "nbformat": 4, "nbformat_minor": 2 -} +} \ No newline at end of file diff --git a/example/smith/smith.ipynb b/example/smith/smith.ipynb index 638fc14..369b7bf 100644 --- a/example/smith/smith.ipynb +++ b/example/smith/smith.ipynb @@ -1,5 +1,20 @@ { "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# \ud83d\udeab DO NOT READ \u2014 DISPOSABLE NOTEBOOK \ud83d\udeab\n", + "\n", + "**This notebook calls an OLD/DIVERGENT tcri API and is scheduled to be rewritten from scratch against the new API. It is NOT a source of truth for anything.**\n", + "\n", + "- Do NOT read it to decide what functions to keep or drop.\n", + "- Do NOT use it for 'sufficiency' or 'is-it-used' evidence.\n", + "- A function being called here is NOT a reason to keep it.\n", + "\n", + "See `docs/contract/REFACTOR_HISTORY.md` (Hard Rules).\n" + ] + }, { "cell_type": "code", "execution_count": 4, @@ -1057,7 +1072,7 @@ { "data": { "text/plain": [ - "View of AnnData object with n_obs × n_vars = 183908 × 642\n", + "View of AnnData object with n_obs \u00d7 n_vars = 183908 \u00d7 642\n", " obs: 'trb', 'sample', 'patient', 'batch', 'specificity', 'response', 'response_specific', 'CD8 T_SCORE', 'CD4 T_SCORE', 'CD8 T Pseudo-probability', 'CD4 T Pseudo-probability', 'phenotype'\n", " var: 'gene_ids', 'feature_types', 'entropy'\n", " uns: 'probability_columns'" @@ -1610,7 +1625,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 642/642 [00:02<00:00, 225.29it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 642/642 [00:02<00:00, 225.29it/s]\n" ] }, { @@ -1625,7 +1640,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 205761/205761 [1:53:20<00:00, 30.26it/s] \n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 205761/205761 [1:53:20<00:00, 30.26it/s] \n" ] }, { @@ -1641,7 +1656,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 642/642 [00:00<00:00, 2367.04it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 642/642 [00:00<00:00, 2367.04it/s]\n" ] }, { @@ -1700,7 +1715,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 642/642 [00:00<00:00, 2477224.63it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 642/642 [00:00<00:00, 2477224.63it/s]\n" ] }, { @@ -1746,8 +1761,8 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 642/642 [00:00<00:00, 2892312.75it/s]\n", - "100%|██████████| 75350815/75350815 [02:04<00:00, 604348.54it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 642/642 [00:00<00:00, 2892312.75it/s]\n", + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 75350815/75350815 [02:04<00:00, 604348.54it/s]\n" ] }, { @@ -1761,7 +1776,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 183908/183908 [00:29<00:00, 6272.76it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 183908/183908 [00:29<00:00, 6272.76it/s]\n" ] }, { @@ -1794,7 +1809,7 @@ { "data": { "text/plain": [ - "AnnData object with n_obs × n_vars = 183908 × 642\n", + "AnnData object with n_obs \u00d7 n_vars = 183908 \u00d7 642\n", " obs: 'trb', 'sample', 'patient', 'batch', 'specificity', 'response', 'response_specific', 'CD8 T_SCORE', 'CD4 T_SCORE', 'CD8 T Pseudo-probability', 'CD4 T Pseudo-probability', 'phenotype'\n", " var: 'gene_ids', 'feature_types', 'entropy'\n", " uns: 'probability_columns', 'neighbors', 'umap'\n", @@ -1852,7 +1867,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 183908/183908 [00:02<00:00, 72968.11it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 183908/183908 [00:02<00:00, 72968.11it/s]\n" ] }, { @@ -1867,7 +1882,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 183908/183908 [00:02<00:00, 71177.09it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 183908/183908 [00:02<00:00, 71177.09it/s]\n" ] }, { @@ -1882,7 +1897,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 183908/183908 [00:02<00:00, 72343.20it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 183908/183908 [00:02<00:00, 72343.20it/s]\n" ] }, { @@ -1897,7 +1912,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 183908/183908 [00:02<00:00, 70605.67it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 183908/183908 [00:02<00:00, 70605.67it/s]\n" ] }, { @@ -1912,7 +1927,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 183908/183908 [00:02<00:00, 64835.06it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 183908/183908 [00:02<00:00, 64835.06it/s]\n" ] } ], @@ -3620,7 +3635,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 183908/183908 [00:02<00:00, 69905.04it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 183908/183908 [00:02<00:00, 69905.04it/s]\n" ] }, { @@ -3635,7 +3650,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 183908/183908 [00:02<00:00, 73756.42it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 183908/183908 [00:02<00:00, 73756.42it/s]\n" ] }, { @@ -3650,7 +3665,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 183908/183908 [00:02<00:00, 74134.26it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 183908/183908 [00:02<00:00, 74134.26it/s]\n" ] }, { @@ -3723,7 +3738,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 183908/183908 [00:02<00:00, 70735.80it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 183908/183908 [00:02<00:00, 70735.80it/s]\n" ] }, { @@ -3738,7 +3753,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 183908/183908 [00:02<00:00, 71070.97it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 183908/183908 [00:02<00:00, 71070.97it/s]\n" ] }, { @@ -3753,7 +3768,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 183908/183908 [00:02<00:00, 73386.60it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 183908/183908 [00:02<00:00, 73386.60it/s]\n" ] }, { @@ -4433,7 +4448,7 @@ "name": "stdout", "output_type": "stream", "text": [ - "View of AnnData object with n_obs × n_vars = 1364 × 456\n", + "View of AnnData object with n_obs \u00d7 n_vars = 1364 \u00d7 456\n", " obs: 'orig.ident', 'nCount_RNA', 'nFeature_RNA', 'CTgene', 'CTnt', 'CTaa', 'CTstrict', 'sample', 'barcode', 'Frequency', 'cloneType', 'mito.genes', 'db.class', 'db.score', 'PT.annot', 'PT.score', 'HPCA.labels', 'HPCA.pruned.labels', 'Monaco.labels', 'Monaco.pruned.labels', 'db.weight.score', 'db.ratio', 'HPCA.first.labels', 'functional.cluster', 'functional.cluster.conf', 'ID', 'T.filter', 'Tissue', 'GEO_RNA', 'Cohort', 'Sample', 'Treatments', 'Response', 'T', 'SampleID', 'batch', 'trb', 'specificity', 'response_specific', 'clone_size', 'UMAP1', 'UMAP2', 'log(Clone Size)', 'genevector', 'Naive/CM Pseudo-probability', 'Activated Pseudo-probability', 'Memory Pseudo-probability', 'Dysfunctional Pseudo-probability', 'TRM Pseudo-probability', 'Naive/CM Module Score', 'Activated Module Score', 'Memory Module Score', 'Dysfunctional Module Score', 'TRM Module Score', 'response'\n", " var: 'features', 'entropy-BCT1.3', 'entropy-BCT1.5', 'entropy-BCT10.1', 'entropy-BCT2.1', 'entropy-BCT2.2', 'entropy-BCT2.3', 'entropy-BCT3.1', 'entropy-BCT4.1', 'entropy-BCT4.2', 'entropy-BCT4.3', 'entropy-BCT5.1', 'entropy-BCT5.2', 'entropy-BCT5.3', 'entropy-BCT5.4', 'entropy-BCT5.5', 'entropy-BCT6.1', 'entropy-BCT6.2', 'entropy-BCT6.4', 'entropy-BCT6.5', 'entropy-BCT7.1', 'entropy-BCT7.2', 'entropy-BCT7.3', 'entropy-BCT8.1', 'entropy-BCT8.2', 'entropy-BCT8.3', 'entropy-BCT8.5', 'entropy-BCT9.1', 'entropy-BT10.1', 'entropy-BT10.2', 'entropy-BT11.1', 'entropy-BT12.2', 'entropy-BT13.2', 'entropy-BT14.1', 'entropy-BT14.2', 'entropy-BT16.2', 'entropy-BT17.1', 'entropy-BT17.2', 'entropy-BT18.2', 'entropy-BT19.2', 'entropy-BT20.2', 'entropy-BT21.2', 'entropy-BT22.1', 'entropy-BT22.2', 'entropy-BT23.1', 'entropy-BT23.2', 'entropy-BT25.1', 'entropy-BT26.2', 'entropy-BT27.1', 'entropy-BT28.2', 'entropy-BT29.1', 'entropy-BT30.2', 'entropy-BT31.2', 'entropy-BT32.1', 'entropy-BT32.2', 'entropy-BT33.1', 'entropy-BT33.2', 'entropy-BT35.1', 'entropy-BT35.2', 'entropy-BT36.1', 'entropy-BT36.2', 'entropy-BT37.1', 'entropy-BT38.1', 'entropy-BT38.2', 'entropy-BT39.1', 'entropy-BT4.1', 'entropy-BT4.2', 'entropy-BT41.2', 'entropy-BT42.2', 'entropy-BT44.1', 'entropy-BT45.1', 'entropy-BT5.1', 'entropy-BT5.2', 'entropy-BT6.2', 'entropy-BT8.2', 'entropy-BT9.1', 'entropy-BT9.2', 'entropy-CT1', 'entropy-EN1', 'entropy-EN2', 'entropy-EN3', 'entropy-ESN1', 'entropy-ESN6', 'entropy-ESN7', 'entropy-EST1', 'entropy-EST3', 'entropy-ET1', 'entropy-ET2', 'entropy-HL2.2', 'entropy-HL2.3', 'entropy-HL3.1', 'entropy-HL4.3', 'entropy-HT1.1', 'entropy-HT1.2', 'entropy-HT2.1', 'entropy-HT2.2', 'entropy-HT3.2', 'entropy-HT3.3', 'entropy-HT4.1', 'entropy-HT4.2', 'entropy-LB14.1', 'entropy-LB14.14', 'entropy-LB14.3', 'entropy-LB14.6', 'entropy-LB14.8', 'entropy-LB8', 'entropy-LB9', 'entropy-LJ8', 'entropy-LJ9', 'entropy-LL13.3', 'entropy-LL14.1', 'entropy-LL14.2', 'entropy-LL14.4', 'entropy-LL18.1', 'entropy-LL18.2', 'entropy-LL18.3', 'entropy-LN1', 'entropy-LN10', 'entropy-LN11', 'entropy-LN12', 'entropy-LN14.2', 'entropy-LN14.3', 'entropy-LN14.4', 'entropy-LN14.5', 'entropy-LN16.1', 'entropy-LN16.2', 'entropy-LN18.1', 'entropy-LN18.2', 'entropy-LN18.5', 'entropy-LN18.6', 'entropy-LN18.8', 'entropy-LN19', 'entropy-LN2', 'entropy-LN22', 'entropy-LN3', 'entropy-LN4', 'entropy-LT1', 'entropy-LT11', 'entropy-LT12', 'entropy-LT13.2', 'entropy-LT13.3', 'entropy-LT13.4', 'entropy-LT13.6', 'entropy-LT13.7', 'entropy-LT13.9', 'entropy-LT14.1', 'entropy-LT14.2', 'entropy-LT14.3', 'entropy-LT14.4', 'entropy-LT14.5', 'entropy-LT14.6', 'entropy-LT14.8', 'entropy-LT16.1', 'entropy-LT16.2', 'entropy-LT16.3', 'entropy-LT16.7', 'entropy-LT16.9', 'entropy-LT17', 'entropy-LT18.10', 'entropy-LT18.3', 'entropy-LT18.4', 'entropy-LT18.6', 'entropy-LT18.8', 'entropy-LT18.9', 'entropy-LT19.1', 'entropy-LT19.2', 'entropy-LT19.3', 'entropy-LT2', 'entropy-LT20.2', 'entropy-LT20.3', 'entropy-LT21.1', 'entropy-LT21.2', 'entropy-LT21.3', 'entropy-LT21.4', 'entropy-LT23', 'entropy-LT24.1', 'entropy-LT24.4', 'entropy-LT25.1', 'entropy-LT25.2', 'entropy-LT25.4', 'entropy-LT26.3', 'entropy-LT26.4', 'entropy-LT29', 'entropy-LT31', 'entropy-LT32', 'entropy-LT33', 'entropy-LT35.1', 'entropy-LT35.2', 'entropy-LT36', 'entropy-LT38.1', 'entropy-LT40', 'entropy-LT41', 'entropy-LT42', 'entropy-LT43', 'entropy-LT46', 'entropy-LT49', 'entropy-LT50', 'entropy-LT51', 'entropy-LT52', 'entropy-LT53', 'entropy-LT54.1', 'entropy-LT54.2', 'entropy-LT55.1', 'entropy-LT55.2', 'entropy-LT56.1', 'entropy-LT56.2', 'entropy-LT58.1', 'entropy-LT59', 'entropy-LT60', 'entropy-LT61', 'entropy-LT7', 'entropy-LT8', 'entropy-LT9', 'entropy-MB1', 'entropy-MB2.2', 'entropy-MB3.1', 'entropy-ML4', 'entropy-MT1', 'entropy-MT3', 'entropy-OT3', 'entropy-OT4', 'entropy-OT5', 'entropy-OT6', 'entropy-RB1', 'entropy-RB10', 'entropy-RB11', 'entropy-RB12', 'entropy-RB13.1', 'entropy-RB14.1', 'entropy-RB15.1', 'entropy-RB15.2', 'entropy-RB16', 'entropy-RB2', 'entropy-RB9', 'entropy-RL10', 'entropy-RN1', 'entropy-RN10', 'entropy-RN12', 'entropy-RN13.1', 'entropy-RN14.2', 'entropy-RN15.1', 'entropy-RN3', 'entropy-RN7', 'entropy-RN8', 'entropy-RN9', 'entropy-RT1', 'entropy-RT10.1', 'entropy-RT10.3', 'entropy-RT11.1', 'entropy-RT11.2', 'entropy-RT11.3', 'entropy-RT12.1', 'entropy-RT12.2', 'entropy-RT13.1', 'entropy-RT13.2', 'entropy-RT14.1', 'entropy-RT15.1', 'entropy-RT16.1', 'entropy-RT2', 'entropy-RT5', 'entropy-RT7.1', 'entropy-RT7.2', 'entropy-RT8.3', 'entropy-RT9.1', 'entropy-RT9.2', 'entropy-RT9.3', 'entropy-SCT1.3', 'entropy-SCT1.4', 'entropy-SCT1.6', 'entropy-SCT2.1', 'entropy-SCT3.1', 'entropy-SCT3.2', 'entropy-SCT4.1', 'entropy-SCT4.2', 'entropy'\n", " uns: 'Response_colors', 'SampleID_colors', 'dendrogram_genevector', 'genevector_colors', 'log1p', 'neighbors', 'probability_columns', 'rank_genes_groups', 'response_specific_colors', 'specificity_colors', 'tcri_clone_key', 'tcri_phenotype_key', 'tcri_unique_clonotypes', 'tcri_unique_phenotypes', 'umap', 'joint_distribution'\n", @@ -4569,4 +4584,4 @@ }, "nbformat": 4, "nbformat_minor": 2 -} +} \ No newline at end of file diff --git a/example/synthetic/synthetic1.ipynb b/example/synthetic/synthetic1.ipynb index ef398b7..a6abd29 100644 --- a/example/synthetic/synthetic1.ipynb +++ b/example/synthetic/synthetic1.ipynb @@ -1,5 +1,20 @@ { "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# \ud83d\udeab DO NOT READ \u2014 DISPOSABLE NOTEBOOK \ud83d\udeab\n", + "\n", + "**This notebook calls an OLD/DIVERGENT tcri API and is scheduled to be rewritten from scratch against the new API. It is NOT a source of truth for anything.**\n", + "\n", + "- Do NOT read it to decide what functions to keep or drop.\n", + "- Do NOT use it for 'sufficiency' or 'is-it-used' evidence.\n", + "- A function being called here is NOT a reason to keep it.\n", + "\n", + "See `docs/contract/REFACTOR_HISTORY.md` (Hard Rules).\n" + ] + }, { "cell_type": "code", "execution_count": 18, @@ -67,7 +82,7 @@ { "data": { "text/plain": [ - "AnnData object with n_obs × n_vars = 2000 × 10\n", + "AnnData object with n_obs \u00d7 n_vars = 2000 \u00d7 10\n", " obs: 'clone_id', 'true_phenotype', 'covariate', 'patient'\n", " layers: 'counts'" ] @@ -197,7 +212,7 @@ { "data": { "text/plain": [ - "AnnData object with n_obs × n_vars = 2000 × 10\n", + "AnnData object with n_obs \u00d7 n_vars = 2000 \u00d7 10\n", " obs: 'clone_id', 'true_phenotype', 'covariate', 'patient', 'indices', '_scvi_clonotype_col_in_registry', '_scvi_phenotype_col_in_registry', '_scvi_covariate_col_in_registry', '_scvi_indices', '_scvi_batch'\n", " uns: 'rank_genes_groups', 'dendrogram_true_phenotype', 'log1p', 'pca', 'neighbors', 'umap', 'true_phenotype_colors', 'clone_id_colors', '_scvi_uuid', '_scvi_manager_uuid', 'tcri_manager'\n", " obsm: 'X_pca', 'X_umap'\n", @@ -525,4 +540,4 @@ }, "nbformat": 4, "nbformat_minor": 5 -} +} \ No newline at end of file diff --git a/example/yost/yost.ipynb b/example/yost/yost.ipynb index e5e87ba..cea048a 100644 --- a/example/yost/yost.ipynb +++ b/example/yost/yost.ipynb @@ -1,5 +1,20 @@ { "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# \ud83d\udeab DO NOT READ \u2014 DISPOSABLE NOTEBOOK \ud83d\udeab\n", + "\n", + "**This notebook calls an OLD/DIVERGENT tcri API and is scheduled to be rewritten from scratch against the new API. It is NOT a source of truth for anything.**\n", + "\n", + "- Do NOT read it to decide what functions to keep or drop.\n", + "- Do NOT use it for 'sufficiency' or 'is-it-used' evidence.\n", + "- A function being called here is NOT a reason to keep it.\n", + "\n", + "See `docs/contract/REFACTOR_HISTORY.md` (Hard Rules).\n" + ] + }, { "cell_type": "code", "execution_count": 1, @@ -75,7 +90,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 15828/15828 [00:04<00:00, 3438.56it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 15828/15828 [00:04<00:00, 3438.56it/s]\n" ] }, { @@ -206,7 +221,7 @@ { "data": { "text/plain": [ - "AnnData object with n_obs × n_vars = 13614 × 163\n", + "AnnData object with n_obs \u00d7 n_vars = 13614 \u00d7 163\n", " obs: 'cell.id', 'patient', 'treatment', 'cluster', 'UMAP1', 'UMAP2', 'trb', 'trb2', 'disease', 'response', 'clone_size'\n", " var: 'entropy'\n", " uns: 'cluster_colors', 'tcri_clone_key', 'tcri_unique_clonotypes'\n", @@ -401,7 +416,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 1045/1045 [00:00<00:00, 344796.07it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 1045/1045 [00:00<00:00, 344796.07it/s]\n" ] }, { @@ -436,8 +451,8 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 1045/1045 [00:00<00:00, 3708162.17it/s]\n", - "100%|██████████| 11670957/11670957 [00:18<00:00, 635611.94it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 1045/1045 [00:00<00:00, 3708162.17it/s]\n", + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 11670957/11670957 [00:18<00:00, 635611.94it/s]\n" ] }, { @@ -451,7 +466,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 22452/22452 [00:03<00:00, 5822.29it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 22452/22452 [00:03<00:00, 5822.29it/s]\n" ] }, { @@ -474,7 +489,7 @@ { "data": { "text/plain": [ - "AnnData object with n_obs × n_vars = 22452 × 1045\n", + "AnnData object with n_obs \u00d7 n_vars = 22452 \u00d7 1045\n", " obs: 'cell.id', 'patient', 'treatment', 'cluster', 'UMAP1', 'UMAP2', 'trb', 'trb2', 'disease', 'response'\n", " var: 'entropy'\n", " uns: 'cluster_colors', 'tcri_clone_key', 'tcri_unique_clonotypes', 'neighbors', 'umap'\n", @@ -531,7 +546,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 22452/22452 [00:00<00:00, 70050.41it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 22452/22452 [00:00<00:00, 70050.41it/s]\n" ] }, { @@ -546,7 +561,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 22452/22452 [00:00<00:00, 71271.65it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 22452/22452 [00:00<00:00, 71271.65it/s]\n" ] }, { @@ -561,7 +576,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 22452/22452 [00:00<00:00, 71910.37it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 22452/22452 [00:00<00:00, 71910.37it/s]\n" ] }, { @@ -576,7 +591,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 22452/22452 [00:00<00:00, 72544.94it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 22452/22452 [00:00<00:00, 72544.94it/s]\n" ] } ], @@ -850,7 +865,7 @@ { "data": { "text/plain": [ - "AnnData object with n_obs × n_vars = 22452 × 1045\n", + "AnnData object with n_obs \u00d7 n_vars = 22452 \u00d7 1045\n", " obs: 'cell.id', 'patient', 'treatment', 'cluster', 'UMAP1', 'UMAP2', 'trb', 'trb2', 'disease', 'response', 'clone_size', 'log(Clone Size)', 'genevector', 'response_treatment', 'rdt', 'response_disease', 'Naive Module Score', 'Activated Module Score', 'Memory Module Score', 'Dsyfunctional Module Score', 'Naive Pseudo-probability', 'Activated Pseudo-probability', 'Memory Pseudo-probability', 'Dsyfunctional Pseudo-probability'\n", " var: 'entropy'\n", " uns: 'cluster_colors', 'tcri_clone_key', 'tcri_unique_clonotypes', 'neighbors', 'umap', 'log1p', 'probability_columns', 'genevector_colors', 'treatment_colors', 'response_colors', 'patient_colors', 'response_disease_colors', 'response_treatment_colors', 'tcri_unique_phenotypes', 'tcri_phenotype_key', 'pca'\n", @@ -1610,7 +1625,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 4/4 [02:30<00:00, 37.60s/it]" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 4/4 [02:30<00:00, 37.60s/it]" ] }, { @@ -1653,7 +1668,7 @@ "name": "stderr", "output_type": "stream", "text": [ - " 50%|█████ | 2/4 [01:18<01:18, 39.46s/it]" + " 50%|\u2588\u2588\u2588\u2588\u2588 | 2/4 [01:18<01:18, 39.46s/it]" ] }, { @@ -1995,7 +2010,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 14/14 [00:28<00:00, 2.05s/it]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 14/14 [00:28<00:00, 2.05s/it]\n" ] }, { @@ -2729,4 +2744,4 @@ }, "nbformat": 4, "nbformat_minor": 5 -} +} \ No newline at end of file diff --git a/example/yost/yost_v2.ipynb b/example/yost/yost_v2.ipynb index 5f5d3ac..380cbb4 100644 --- a/example/yost/yost_v2.ipynb +++ b/example/yost/yost_v2.ipynb @@ -1,5 +1,20 @@ { "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# \ud83d\udeab DO NOT READ \u2014 DISPOSABLE NOTEBOOK \ud83d\udeab\n", + "\n", + "**This notebook calls an OLD/DIVERGENT tcri API and is scheduled to be rewritten from scratch against the new API. It is NOT a source of truth for anything.**\n", + "\n", + "- Do NOT read it to decide what functions to keep or drop.\n", + "- Do NOT use it for 'sufficiency' or 'is-it-used' evidence.\n", + "- A function being called here is NOT a reason to keep it.\n", + "\n", + "See `docs/contract/REFACTOR_HISTORY.md` (Hard Rules).\n" + ] + }, { "cell_type": "code", "execution_count": 57, @@ -640,11 +655,11 @@ "\n", " sns.kdeplot(delta, fill=True, ax=ax, alpha=0.4)\n", " ax.axvline(0, color=\"k\", linestyle=\"--\")\n", - " ax.set_title(f\"{level}\\nΔ={delta_mean:.2f}, P(Δ>0)={p_greater:.2f}\")\n", - " ax.set_xlabel(\"Δ Phenotypic Entropy\")\n", + " ax.set_title(f\"{level}\\n\u0394={delta_mean:.2f}, P(\u0394>0)={p_greater:.2f}\")\n", + " ax.set_xlabel(\"\u0394 Phenotypic Entropy\")\n", " ax.set_ylabel(\"Posterior Density\")\n", "\n", - " fig.suptitle(\"Posterior Δ Phenotypic Entropy (Bayesian)\", fontsize=14)\n", + " fig.suptitle(\"Posterior \u0394 Phenotypic Entropy (Bayesian)\", fontsize=14)\n", " fig.tight_layout()\n", "\n", " if save:\n", @@ -689,7 +704,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "Computing mutual information: 100%|██████████| 14/14 [00:04<00:00, 2.85it/s]\n" + "Computing mutual information: 100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 14/14 [00:04<00:00, 2.85it/s]\n" ] }, { @@ -1221,7 +1236,7 @@ " })\n", "\n", " ret_df = df\n", - " # 1) define your x‑ and hue‑levels\n", + " # 1) define your x\u2011 and hue\u2011levels\n", " x_cats = df[splitby].unique().tolist()\n", " hue_order = [\"pre\", \"post\"]\n", " cov_cats = hue_order\n", @@ -1235,7 +1250,7 @@ " n_cov)\n", " cov2off = dict(zip(cov_cats, offsets))\n", " \n", - " # 3) add a numeric x‑position for your scatter\n", + " # 3) add a numeric x\u2011position for your scatter\n", " df[\"_x\"] = df[splitby].map(lambda x: x_cats.index(x)) + df[cov_col].map(cov2off)\n", " \n", " # 4) start the plot\n", @@ -1251,7 +1266,7 @@ " )\n", " ax.legend_.remove()\n", " \n", - " # 5) compute all covariate‑pairs and global y positions\n", + " # 5) compute all covariate\u2011pairs and global y positions\n", " pairs = list(itertools.combinations(cov_cats, 2))\n", " y_max = df[\"Mutual Information\"].max()\n", " y_min = df[\"Mutual Information\"].min()\n", @@ -1259,7 +1274,7 @@ " h_bar = yrange * 0.05\n", " v_step = h_bar * 1.2\n", " \n", - " # 6) for each x‑category, annotate each pair\n", + " # 6) for each x\u2011category, annotate each pair\n", " for i, xcat in enumerate(x_cats):\n", " for j, (g1, g2) in enumerate(pairs):\n", " # slice data for the two levels\n", @@ -1268,14 +1283,14 @@ " if len(d1)==0 or len(d2)==0:\n", " continue\n", " \n", - " # Mann–Whitney U test\n", + " # Mann\u2013Whitney U test\n", " stat, p = mannwhitneyu(d1, d2)\n", " stars = (\"ns\" if p>=0.05 else\n", " \"*\" if p<0.05 else\n", " \"**\" if p<0.01 else\n", " \"***\" if p<0.001 else \"****\")\n", " print(xcat,p,stat)\n", - " # x‑positions of the two boxes\n", + " # x\u2011positions of the two boxes\n", " x1 = i + cov2off[g1]\n", " x2 = i + cov2off[g2]\n", " \n", @@ -1293,7 +1308,7 @@ " fontsize=legend_fontsize,\n", " zorder=4)\n", " \n", - " # 7) now scatter your patient‑colored points on top\n", + " # 7) now scatter your patient\u2011colored points on top\n", " patients = df[batch_col].unique().tolist()\n", " patient_palette = sns.color_palette(\"tab20b\", len(patients))\n", " patient_colors = dict(zip(patients, patient_palette))\n", @@ -1344,17 +1359,17 @@ "output_type": "stream", "text": [ "\n", - "🔍 response = NR — Δ MI (post - pre)\n", - "Δ mean = -0.0244, std = 0.0384\n", + "\ud83d\udd0d response = NR \u2014 \u0394 MI (post - pre)\n", + "\u0394 mean = -0.0244, std = 0.0384\n", "Cohen's d = -0.636\n", "95% HDI: (-0.0889, 0.0523)\n", - "P(Δ > 0) = 0.259, P(Δ < 0) = 0.741\n", + "P(\u0394 > 0) = 0.259, P(\u0394 < 0) = 0.741\n", "\n", - "🔍 response = R — Δ MI (post - pre)\n", - "Δ mean = 0.0696, std = 0.0237\n", + "\ud83d\udd0d response = R \u2014 \u0394 MI (post - pre)\n", + "\u0394 mean = 0.0696, std = 0.0237\n", "Cohen's d = 2.937\n", "95% HDI: (0.0234, 0.1160)\n", - "P(Δ > 0) = 0.997, P(Δ < 0) = 0.003\n" + "P(\u0394 > 0) = 0.997, P(\u0394 < 0) = 0.003\n" ] }, { @@ -1435,11 +1450,11 @@ " p_greater = np.mean(delta > 0)\n", " p_less = np.mean(delta < 0)\n", "\n", - " print(f\"\\n🔍 {splitby} = {group} — Δ MI ({group2} - {group1})\")\n", - " print(f\"Δ mean = {delta_mean:.4f}, std = {delta_std:.4f}\")\n", + " print(f\"\\n\ud83d\udd0d {splitby} = {group} \u2014 \u0394 MI ({group2} - {group1})\")\n", + " print(f\"\u0394 mean = {delta_mean:.4f}, std = {delta_std:.4f}\")\n", " print(f\"Cohen's d = {cohens_d:.3f}\")\n", " print(f\"95% HDI: ({hdi_low:.4f}, {hdi_high:.4f})\")\n", - " print(f\"P(Δ > 0) = {p_greater:.3f}, P(Δ < 0) = {p_less:.3f}\")\n", + " print(f\"P(\u0394 > 0) = {p_greater:.3f}, P(\u0394 < 0) = {p_less:.3f}\")\n", "\n", " results[group] = {\n", " \"delta_samples\": delta,\n", @@ -1453,15 +1468,15 @@ " \"hdi\": (hdi_low, hdi_high)\n", " }\n", "\n", - " # 🎨 Plot panels\n", + " # \ud83c\udfa8 Plot panels\n", " fig, axs = plt.subplots(1, 3, figsize=(15, 4), gridspec_kw={\"width_ratios\": [2, 2, 1]})\n", "\n", - " # Panel A: Posterior ΔMI distributions\n", + " # Panel A: Posterior \u0394MI distributions\n", " for group in groups:\n", " sns.kdeplot(results[group][\"delta_samples\"], fill=True, ax=axs[0], label=group, alpha=0.8)\n", " axs[0].axvline(0, color=\"k\", linestyle=\"--\")\n", - " axs[0].set_title(\"A. Δ Mutual Information (post - pre)\")\n", - " axs[0].set_xlabel(\"Δ Normalized MI\")\n", + " axs[0].set_title(\"A. \u0394 Mutual Information (post - pre)\")\n", + " axs[0].set_xlabel(\"\u0394 Normalized MI\")\n", " axs[0].set_ylabel(\"Posterior density\")\n", " axs[0].legend(title=splitby)\n", "\n", @@ -1474,7 +1489,7 @@ " axs[1].set_ylabel(\"Posterior density\")\n", " axs[1].legend()\n", "\n", - " # Panel C: Δ MI summary bars\n", + " # Panel C: \u0394 MI summary bars\n", " labels = []\n", " means = []\n", " errors = []\n", @@ -1484,10 +1499,10 @@ " errors.append(results[group][\"delta_std\"])\n", " axs[2].bar(labels, means, yerr=errors, capsize=5, color=sns.color_palette(\"Set2\", len(groups)))\n", " axs[2].axhline(0, color=\"k\", linestyle=\"--\")\n", - " axs[2].set_title(\"C. Δ MI (post - pre)\")\n", - " axs[2].set_ylabel(\"Δ Normalized MI\")\n", + " axs[2].set_title(\"C. \u0394 MI (post - pre)\")\n", + " axs[2].set_ylabel(\"\u0394 Normalized MI\")\n", "\n", - " fig.suptitle(\"Bayesian Analysis of Clonotype–Phenotype Coupling\", fontsize=14)\n", + " fig.suptitle(\"Bayesian Analysis of Clonotype\u2013Phenotype Coupling\", fontsize=14)\n", " fig.tight_layout()\n", " fig.savefig(\"bayes_mi_yost.pdf\")\n", " return results\n", @@ -1514,7 +1529,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "Computing clonotypic entropy: 100%|██████████| 7/7 [00:00<00:00, 15.45it/s]\n" + "Computing clonotypic entropy: 100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 7/7 [00:00<00:00, 15.45it/s]\n" ] }, { @@ -1531,7 +1546,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "Computing clonotypic entropy: 100%|██████████| 7/7 [00:00<00:00, 15.84it/s]\n" + "Computing clonotypic entropy: 100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 7/7 [00:00<00:00, 15.84it/s]\n" ] }, { @@ -1559,7 +1574,7 @@ " figsize=(6, 3),\n", " rotation=90,\n", " return_df=False,\n", - " bayesian=True, # ✅ NEW FLAG\n", + " bayesian=True, # \u2705 NEW FLAG\n", " bayes_samples=1000\n", "):\n", " import numpy as np\n", @@ -1640,7 +1655,7 @@ " )\n", " ax.legend_.remove()\n", "\n", - " # 🔁 Bayesian-style posterior annotations\n", + " # \ud83d\udd01 Bayesian-style posterior annotations\n", " y_max = df[\"Clonotypic Entropy\"].max()\n", " y_min = df[\"Clonotypic Entropy\"].min()\n", " yrange = y_max - y_min\n", @@ -1662,7 +1677,7 @@ " hdi_low, hdi_high = np.percentile(delta, [2.5, 97.5])\n", " p_greater = np.mean(delta > 0)\n", "\n", - " label = f\"Δ={mean:.2g}\\nP(Δ>0)={p_greater:.2f}\"\n", + " label = f\"\u0394={mean:.2g}\\nP(\u0394>0)={p_greater:.2f}\"\n", " else:\n", " from scipy.stats import mannwhitneyu\n", " stat, pval = mannwhitneyu(d1, d2, alternative=\"two-sided\")\n", @@ -1719,7 +1734,7 @@ " cdata,\n", " temperature=1,\n", " hue_order=[\"pre\", \"post\"],\n", - " bayesian=True, # ✅ enables posterior comparison\n", + " bayesian=True, # \u2705 enables posterior comparison\n", " return_df=True,\n", " save=\"yost_ce_{}.pdf\".format(x),\n", " )" @@ -1831,7 +1846,7 @@ "# -----------------------------\n", "# 3. Compare Distributions: Plotting Side-by-Side\n", "# -----------------------------\n", - "# We’ll loop over the unique tissues and for each, plot the inferred (model) and empirical distributions.\n", + "# We\u2019ll loop over the unique tissues and for each, plot the inferred (model) and empirical distributions.\n", "unique_tissues = df_model[\"treatment\"].unique()\n", "n_tissues = len(unique_tissues)\n", "\n", @@ -1990,7 +2005,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "Computing phenotypic entropy: 100%|██████████| 2/2 [01:58<00:00, 59.02s/it]\n" + "Computing phenotypic entropy: 100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 2/2 [01:58<00:00, 59.02s/it]\n" ] }, { @@ -2120,7 +2135,7 @@ "def plot_entropy_panels(df, cov_col=\"timepoint\", group_col=\"Group\", n_bootstrap=1000):\n", " fig, axs = plt.subplots(1, 3, figsize=(15, 5))\n", "\n", - " # Panel A: Posterior Δ entropy\n", + " # Panel A: Posterior \u0394 entropy\n", " groups = df[group_col].unique()\n", " for g in groups:\n", " d1 = df[(df[group_col] == g) & (df[cov_col] == \"pre\")][\"Phenotypic Entropy\"].values\n", @@ -2131,8 +2146,8 @@ " sns.kdeplot(delta, fill=True, ax=axs[0], label=g, alpha=0.4)\n", "\n", " axs[0].axvline(0, color=\"k\", linestyle=\"--\")\n", - " axs[0].set_title(\"A. Δ Phenotypic Entropy (post - pre)\")\n", - " axs[0].set_xlabel(\"Δ entropy\")\n", + " axs[0].set_title(\"A. \u0394 Phenotypic Entropy (post - pre)\")\n", + " axs[0].set_xlabel(\"\u0394 entropy\")\n", " axs[0].legend()\n", "\n", " # Panel B: Clone-level violin plot\n", @@ -2226,7 +2241,7 @@ " if phenotype_subset is None:\n", " phenotype_subset = all_phenotypes\n", "\n", - " phenotypes = Phenotypes(phenotype_subset) # ✅ subset phenotype space\n", + " phenotypes = Phenotypes(phenotype_subset) # \u2705 subset phenotype space\n", " cell_probabilities = collections.defaultdict(dict)\n", "\n", " for s in order:\n", @@ -2236,7 +2251,7 @@ " jd = jd.groupby(\"clonotype_id\").mean()\n", " for x in jd.T:\n", " raw_counts = jd.T[x].to_dict()\n", - " filtered_counts = {k: v for k, v in raw_counts.items() if k in phenotype_subset} # ✅ subset\n", + " filtered_counts = {k: v for k, v in raw_counts.items() if k in phenotype_subset} # \u2705 subset\n", " cell_probabilities[s][x] = filtered_counts\n", "\n", " repertoires = {}\n", @@ -2262,7 +2277,7 @@ " ):\n", " if str(seq) != \"nan\" and condition in repertoires and seq in cell_probabilities[condition]:\n", " phenotypes_and_counts = cell_probabilities[condition][seq]\n", - " phenotypes_and_counts = {k: v for k, v in phenotypes_and_counts.items() if k in phenotype_subset} # ✅ filter\n", + " phenotypes_and_counts = {k: v for k, v in phenotypes_and_counts.items() if k in phenotype_subset} # \u2705 filter\n", " if nt:\n", " t = Tcell(\n", " phenotypes=phenotypes,\n", @@ -2353,4 +2368,4 @@ }, "nbformat": 4, "nbformat_minor": 5 -} +} \ No newline at end of file diff --git a/example/zhang/zhang.ipynb b/example/zhang/zhang.ipynb index 6422329..01f19a0 100644 --- a/example/zhang/zhang.ipynb +++ b/example/zhang/zhang.ipynb @@ -1,5 +1,20 @@ { "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# \ud83d\udeab DO NOT READ \u2014 DISPOSABLE NOTEBOOK \ud83d\udeab\n", + "\n", + "**This notebook calls an OLD/DIVERGENT tcri API and is scheduled to be rewritten from scratch against the new API. It is NOT a source of truth for anything.**\n", + "\n", + "- Do NOT read it to decide what functions to keep or drop.\n", + "- Do NOT use it for 'sufficiency' or 'is-it-used' evidence.\n", + "- A function being called here is NOT a reason to keep it.\n", + "\n", + "See `docs/contract/REFACTOR_HISTORY.md` (Hard Rules).\n" + ] + }, { "cell_type": "code", "execution_count": 2, @@ -96,13 +111,13 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 18529/18529 [00:14<00:00, 1317.65it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 18529/18529 [00:14<00:00, 1317.65it/s]\n" ] }, { "data": { "text/plain": [ - "View of AnnData object with n_obs × n_vars = 77030 × 356\n", + "View of AnnData object with n_obs \u00d7 n_vars = 77030 \u00d7 356\n", " obs: 'patient', 'sample', 'celltype', 'cluster', 'trb', 'response', 'timepoint', 'timepoint_simple', 'treatment', 'timepoint_complex', 'patient_timepoint_complex', 'complete_response', 'timepoint_relapse'\n", " var: 'entropy'" ] @@ -151,7 +166,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 15252980/15252980 [00:07<00:00, 1993156.42it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 15252980/15252980 [00:07<00:00, 1993156.42it/s]\n" ] }, { @@ -186,7 +201,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 356/356 [00:00<00:00, 654.47it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 356/356 [00:00<00:00, 654.47it/s]\n" ] }, { @@ -201,7 +216,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 356/356 [00:00<00:00, 5169.30it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 356/356 [00:00<00:00, 5169.30it/s]\n" ] }, { @@ -246,7 +261,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 356/356 [00:00<00:00, 1799002.68it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 356/356 [00:00<00:00, 1799002.68it/s]\n" ] }, { @@ -283,8 +298,8 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 356/356 [00:00<00:00, 3183736.09it/s]\n", - "100%|██████████| 15252980/15252980 [00:24<00:00, 628055.69it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 356/356 [00:00<00:00, 3183736.09it/s]\n", + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 15252980/15252980 [00:24<00:00, 628055.69it/s]\n" ] }, { @@ -298,7 +313,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 77030/77030 [00:05<00:00, 13776.39it/s]" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 77030/77030 [00:05<00:00, 13776.39it/s]" ] }, { @@ -345,7 +360,7 @@ { "data": { "text/plain": [ - "AnnData object with n_obs × n_vars = 77030 × 356\n", + "AnnData object with n_obs \u00d7 n_vars = 77030 \u00d7 356\n", " obs: 'patient', 'sample', 'celltype', 'cluster', 'trb', 'response', 'timepoint', 'timepoint_simple', 'treatment', 'timepoint_complex', 'patient_timepoint_complex', 'complete_response', 'timepoint_relapse'\n", " var: 'entropy'\n", " uns: 'tcri_clone_key', 'tcri_unique_clonotypes', 'neighbors', 'umap'\n", @@ -490,7 +505,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 77030/77030 [00:01<00:00, 73449.57it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 77030/77030 [00:01<00:00, 73449.57it/s]\n" ] }, { @@ -505,7 +520,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 77030/77030 [00:01<00:00, 72344.34it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 77030/77030 [00:01<00:00, 72344.34it/s]\n" ] }, { @@ -520,7 +535,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 77030/77030 [00:01<00:00, 73754.34it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 77030/77030 [00:01<00:00, 73754.34it/s]\n" ] }, { @@ -535,7 +550,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 77030/77030 [00:01<00:00, 71691.81it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 77030/77030 [00:01<00:00, 71691.81it/s]\n" ] } ], @@ -1189,7 +1204,7 @@ " \n", " \n", "\n", - "

4 rows × 41540 columns

\n", + "

4 rows \u00d7 41540 columns

\n", "" ], "text/plain": [ @@ -1701,7 +1716,7 @@ "name": "stdout", "output_type": "stream", "text": [ - "View of AnnData object with n_obs × n_vars = 77030 × 356\n", + "View of AnnData object with n_obs \u00d7 n_vars = 77030 \u00d7 356\n", " obs: 'patient', 'sample', 'celltype', 'cluster', 'trb', 'response', 'timepoint', 'timepoint_simple', 'treatment', 'timepoint_complex', 'patient_timepoint_complex', 'complete_response', 'timepoint_relapse', 'cr', 'genevector', 'instant_response', 'timepoint_instant', 'tc', 'clone_size', 'trbaa', 'pent', 'pentv', 'tcs', 'Naive/CM Pseudo-probability', 'Activated Pseudo-probability', 'Memory Pseudo-probability', 'Dsyfunctional Pseudo-probability'\n", " var: 'entropy'\n", " uns: 'tcri_clone_key', 'tcri_unique_clonotypes', 'neighbors', 'umap', 'probability_columns', 'joint_distribution', 'log1p', 'genevector_colors', 'complete_response_colors', 'response_colors', 'treatment_colors', 'timepoint_relapse_colors', 'timepoint_colors', 'patient_colors', 'instant_response_colors', 'timepoint_instant_colors', 'tcri_unique_phenotypes', 'tcri_phenotype_key', 'tc_colors', 'timepoint_complex_colors'\n", @@ -3044,7 +3059,7 @@ " \n", " \n", "\n", - "

40893 rows × 4 columns

\n", + "

40893 rows \u00d7 4 columns

\n", "" ], "text/plain": [ @@ -3251,7 +3266,7 @@ " \n", " \n", "\n", - "

40893 rows × 3 columns

\n", + "

40893 rows \u00d7 3 columns

\n", "" ], "text/plain": [ @@ -4636,7 +4651,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 4/4 [00:03<00:00, 1.30it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 4/4 [00:03<00:00, 1.30it/s]\n" ] }, { @@ -4874,7 +4889,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 2/2 [00:03<00:00, 1.91s/it]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 2/2 [00:03<00:00, 1.91s/it]\n" ] }, { @@ -4954,7 +4969,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 2/2 [00:51<00:00, 25.98s/it]" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 2/2 [00:51<00:00, 25.98s/it]" ] }, { @@ -5557,7 +5572,7 @@ "name": "stdout", "output_type": "stream", "text": [ - "View of AnnData object with n_obs × n_vars = 12934 × 356\n", + "View of AnnData object with n_obs \u00d7 n_vars = 12934 \u00d7 356\n", " obs: 'patient', 'sample', 'celltype', 'cluster', 'trb', 'response', 'timepoint', 'timepoint_simple', 'treatment', 'timepoint_complex', 'patient_timepoint_complex', 'complete_response', 'timepoint_relapse', 'cr', 'clone_size', 'UMAP1', 'UMAP2', 'log(Clone Size)', 'instant_response', 'timepoint_instant', 'genevector', 'Naive/CM Pseudo-probability', 'Activated Pseudo-probability', 'Memory Pseudo-probability', 'Dsyfunctional Pseudo-probability'\n", " var: 'entropy'\n", " uns: 'tcri_clone_key', 'tcri_unique_clonotypes', 'neighbors', 'umap', 'probability_columns', 'genevector_colors', 'complete_response_colors', 'response_colors', 'treatment_colors', 'timepoint_relapse_colors', 'timepoint_colors', 'patient_colors', 'instant_response_colors', 'timepoint_instant_colors', 'tcri_unique_phenotypes', 'tcri_phenotype_key', 'joint_distribution'\n", @@ -5655,7 +5670,7 @@ "name": "stdout", "output_type": "stream", "text": [ - "View of AnnData object with n_obs × n_vars = 12934 × 356\n", + "View of AnnData object with n_obs \u00d7 n_vars = 12934 \u00d7 356\n", " obs: 'patient', 'sample', 'celltype', 'cluster', 'trb', 'response', 'timepoint', 'timepoint_simple', 'treatment', 'timepoint_complex', 'patient_timepoint_complex', 'complete_response', 'timepoint_relapse', 'cr', 'clone_size', 'UMAP1', 'UMAP2', 'log(Clone Size)', 'instant_response', 'timepoint_instant', 'genevector', 'Naive/CM Pseudo-probability', 'Activated Pseudo-probability', 'Memory Pseudo-probability', 'Dsyfunctional Pseudo-probability'\n", " var: 'entropy'\n", " uns: 'tcri_clone_key', 'tcri_unique_clonotypes', 'neighbors', 'umap', 'probability_columns', 'genevector_colors', 'complete_response_colors', 'response_colors', 'treatment_colors', 'timepoint_relapse_colors', 'timepoint_colors', 'patient_colors', 'instant_response_colors', 'timepoint_instant_colors', 'tcri_unique_phenotypes', 'tcri_phenotype_key', 'joint_distribution'\n", @@ -5927,7 +5942,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 77030/77030 [00:01<00:00, 72445.91it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 77030/77030 [00:01<00:00, 72445.91it/s]\n" ] }, { @@ -5942,7 +5957,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 77030/77030 [00:01<00:00, 72363.85it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 77030/77030 [00:01<00:00, 72363.85it/s]\n" ] }, { @@ -5957,7 +5972,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 77030/77030 [00:01<00:00, 65475.09it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 77030/77030 [00:01<00:00, 65475.09it/s]\n" ] }, { @@ -5972,7 +5987,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 77030/77030 [00:01<00:00, 69571.27it/s]\n" + "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 77030/77030 [00:01<00:00, 69571.27it/s]\n" ] }, { @@ -6230,4 +6245,4 @@ }, "nbformat": 4, "nbformat_minor": 2 -} +} \ No newline at end of file diff --git a/tcri/_console.py b/tcri/_console.py new file mode 100644 index 0000000..38e040d --- /dev/null +++ b/tcri/_console.py @@ -0,0 +1,47 @@ +"""Single home for tcri's colored console helpers. + +Replaces the ``_ok`` / ``_info`` / ``_warn`` / ``_fin`` blocks that were +copy-pasted (with inconsistent ANSI aliases) into ``metrics`` / ``preprocessing`` +/ ``plotting``. Import from here; do not redefine. + +Behavior is preserved verbatim so the dedup is a no-op at runtime. Both spellings +of each color alias used across the old files are exported (``GRN``/``GREEN``, +``CYN``/``CYAN``, ``MAG``/``MAGENT``) so existing direct uses resolve unchanged. + +Note: this stays print-based for the dedup PR. Routing verbosity through scanpy's +logger (silenceable, leveled) is a follow-up refinement (REFACTOR_NOTES). +""" + +# ── ANSI palette (all aliases the legacy files used) ───────────────────────── +RESET = "\x1b[0m" +BOLD = "\x1b[1m" +DIM = "\x1b[2m" +GREEN = GRN = "\x1b[32m" +CYAN = CYN = "\x1b[36m" +MAGENT = MAG = "\x1b[35m" +YLW = "\x1b[33m" +RED = "\x1b[31m" + + +def _ok(msg: str, quiet: bool = False): + """Success mark.""" + if not quiet: + print(f"{GRN}✅ {msg}{RESET}") + + +def _info(key: str, txt, quiet: bool = False): + """Key-value info line.""" + if not quiet: + print(f" {CYN}🎯 {key:<22}{DIM}{txt}{RESET}") + + +def _warn(msg: str, quiet: bool = False): + """Warning line.""" + if not quiet: + print(f"{YLW}⚠️ {msg}{RESET}") + + +def _fin(quiet: bool = False): + """Final flourish.""" + if not quiet: + print(f"{MAG}✨ Done!{RESET}") diff --git a/tcri/_contract.pyi b/tcri/_contract.pyi new file mode 100644 index 0000000..1ac327f --- /dev/null +++ b/tcri/_contract.pyi @@ -0,0 +1,145 @@ +"""TCRI target API contract (frozen) — the machine-checkable public surface. + +The conformance test (``tests/test_contract_conformance.py``) checks each +*implemented* function's live signature against its declaration here; declared- +but-absent functions are the refactor worklist. Prose spec: +``docs/contract/tcri_api_and_responsibilities.md``. + +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``). +- ``tl``/``pp``/``pl``/``diag``/``ut`` are container classes purely for namespacing + (so the ``tl``/``pl`` twins can share a name); ``TCRIModel`` is the real model class. +- Decisions baked in: ``n_samples=250`` default (sampling; ``0`` = opt-in point + estimate); ``weighted=False`` default (kept as a dial); ``use_logits`` (not + ``posterior=``); ``normalize_mode="min"``; American spelling; keyword-only options. +""" +from __future__ import annotations +from typing import Any, Optional +from anndata import AnnData +import pandas as pd + + +# ── model (ml) ─────────────────────────────────────────────────────────────── +class TCRIModel: + @classmethod + def setup_anndata( + cls, adata: AnnData, *, layer: Optional[str] = ..., + clonotype_key: str = ..., phenotype_key: str = ..., + covariate_key: str = ..., batch_key: str = ..., + ) -> None: ... + def train( + self, max_epochs: int = ..., batch_size: int = ..., lr: float = ..., + reconstruction_loss_scale: float = ..., n_steps_kl_warmup: int = ..., + ) -> 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 get_p_ct(self) -> Any: ... + def to_anndata(self, adata: AnnData, *, batch_size: int = ..., compute_umap: bool = ...) -> AnnData: ... + + +# ── preprocessing (pp) ─────────────────────────────────────────────────────── +class pp: + def group_singletons( + adata: AnnData, *, clonotype_key: str = ..., groupby: str = ..., + target_col: str = ..., min_clone_size: int = ..., + ) -> None: ... + def clone_size(adata: AnnData, *, key_added: str = ..., return_counts: bool = ...) -> Any: ... + + +# ── tools / metrics (tl) ───────────────────────────────────────────────────── +class tl: + def joint_distribution( + adata: AnnData, *, covariate: Optional[str] = ..., groupby: Optional[str] = ..., + n_samples: int = ..., use_logits: bool = ..., weighted: bool = ..., + clones: Any = ..., temperature: float = ..., random_state: 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 = ..., + ) -> Any: ... + def phenotypic_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 = ..., + ) -> Any: ... + def mutual_information( + adata_or_jd: Any, *, covariate: Optional[str] = ..., groupby: Optional[str] = ..., + splitby: Optional[str] = ..., n_samples: int = ..., temperature: float = ..., + clones: Any = ..., weighted: bool = ..., normalized: bool = ..., + normalize_mode: str = ..., random_state: Any = ..., + ) -> Any: ... + def phenotypic_flux( + adata: AnnData, *, cov_from: str, cov_to: str, groupby: Optional[str] = ..., + splitby: Optional[str] = ..., n_samples: int = ..., temperature: float = ..., + clones: Any = ..., weighted: bool = ..., distance_metric: str = ..., random_state: Any = ..., + ) -> Any: ... + def compare_groups( + df: pd.DataFrame, *, value: str, splitby: str, reference: Optional[str] = ..., + paired: bool = ..., pair_on: Optional[str] = ..., hdi_prob: float = ..., + alternative: str = ..., + ) -> pd.DataFrame: ... + + +# ── plotting (pl) — twins mirror tl by name ────────────────────────────────── +class pl: + def clonotypic_entropy( + adata: AnnData, *, covariate: Optional[str] = ..., groupby: Optional[str] = ..., + splitby: Optional[str] = ..., n_samples: int = ..., temperature: float = ..., + clones: Any = ..., weighted: bool = ..., normalized: bool = ..., + order: Any = ..., hue_order: Any = ..., palette: Any = ..., ax: Any = ..., + figsize: Any = ..., save: Any = ..., show: Any = ..., return_df: bool = ..., + ) -> Any: ... + def phenotypic_entropy( + adata: AnnData, *, covariate: Optional[str] = ..., groupby: Optional[str] = ..., + splitby: Optional[str] = ..., n_samples: int = ..., temperature: float = ..., + clones: Any = ..., weighted: bool = ..., normalized: bool = ..., + order: Any = ..., hue_order: Any = ..., palette: Any = ..., ax: Any = ..., + figsize: Any = ..., save: Any = ..., show: Any = ..., return_df: bool = ..., + ) -> Any: ... + def mutual_information( + adata: AnnData, *, covariate: Optional[str] = ..., groupby: Optional[str] = ..., + splitby: Optional[str] = ..., n_samples: int = ..., temperature: float = ..., + clones: Any = ..., weighted: bool = ..., normalized: bool = ..., normalize_mode: str = ..., + order: Any = ..., hue_order: Any = ..., palette: Any = ..., ax: Any = ..., + figsize: Any = ..., save: Any = ..., show: Any = ..., return_df: bool = ..., + ) -> Any: ... + def phenotypic_flux( + adata: AnnData, *, order: Any, groupby: Optional[str] = ..., splitby: Optional[str] = ..., + n_samples: int = ..., temperature: float = ..., clones: Any = ..., weighted: bool = ..., + 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: ... + + +# ── diagnostics (diag) — returns DataFrames ────────────────────────────────── +class diag: + def joint_distribution_ppc( + adata: AnnData, *, covariate: Optional[str] = ..., distance_metric: str = ..., + temperature: float = ..., clones: Any = ..., random_state: Any = ..., + ) -> pd.DataFrame: ... + def phenotype_calibration(adata: AnnData, *, n_bins: int = ...) -> pd.DataFrame: ... + def reconstruction_ppc( + model: Any, adata: Optional[AnnData] = ..., *, n_sims: int = ..., random_state: Any = ..., + ) -> pd.DataFrame: ... + def permutation_null( + adata: AnnData, *, metric: str = ..., covariate: Optional[str] = ..., + groupby: Optional[str] = ..., n_perm: int = ..., random_state: Any = ..., + ) -> pd.DataFrame: ... + def loss(model: Any, *, log_scale: bool = ..., ax: Any = ..., save: Any = ...) -> Any: ... + def archetypes(model: Any, *, ax: Any = ..., save: Any = ...) -> Any: ... + + +# ── utils / session (ut) ───────────────────────────────────────────────────── +class ut: + def save_tcri_session( + model: Any, adata: AnnData, out_dir: str, *, save_adata: bool = ..., compression: str = ..., + ) -> dict: ... + def load_tcri_session( + run_dir: str, *, adata_path: Optional[str] = ..., map_location: Any = ..., layer: Optional[str] = ..., + ) -> Any: ... diff --git a/tcri/_distance.py b/tcri/_distance.py new file mode 100644 index 0000000..a62e133 --- /dev/null +++ b/tcri/_distance.py @@ -0,0 +1,60 @@ +"""Distance / divergence kernels for phenotype distributions. + +Single home for the KL kernel (was `metrics.dkl` + `flux.dkl_func`, two copies) +plus L1 and a symmetric Jensen–Shannon option. All operate on 1-D probability +vectors, use log base 2 (bits), and share one eps floor. ``phenotype_distance`` +is the string→callable dispatcher used by ``phenotypic_flux(distance_metric=)``. +""" +from __future__ import annotations + +import numpy as np + +EPS = 1e-12 + + +def _normalize(p): + p = np.clip(np.asarray(p, float), EPS, None) + return p / p.sum() + + +def kl_divergence(p, q, *, base: float = 2.0, eps: float = EPS) -> float: + """KL(p ‖ q) in bits (base 2). Asymmetric.""" + p = np.clip(np.asarray(p, float), eps, None); p = p / p.sum() + q = np.clip(np.asarray(q, float), eps, None); q = q / q.sum() + return float(np.sum(p * (np.log(p / q) / np.log(base)))) + + +def l1_distance(p, q) -> float: + """L1 (Manhattan) distance between two normalized distributions, in [0, 2].""" + return float(np.abs(_normalize(p) - _normalize(q)).sum()) + + +def jensen_shannon(p, q, *, base: float = 2.0, eps: float = EPS) -> float: + """Jensen–Shannon divergence: symmetric, bounded [0, 1] bit — the recommended + symmetric shift measure.""" + p = _normalize(p) + q = _normalize(q) + m = 0.5 * (p + q) + return float(0.5 * kl_divergence(p, m, base=base, eps=eps) + + 0.5 * kl_divergence(q, m, base=base, eps=eps)) + + +_REGISTRY = { + "l1": l1_distance, + "kl": kl_divergence, + "dkl": kl_divergence, + "js": jensen_shannon, + "jsd": jensen_shannon, +} + + +def phenotype_distance(metric): + """Resolve ``distance_metric`` (a name or a callable ``f(p, q) -> float``).""" + if callable(metric): + return metric + key = str(metric).lower() + if key not in _REGISTRY: + raise ValueError( + f"distance_metric must be a callable or one of {sorted(_REGISTRY)}; got {metric!r}" + ) + return _REGISTRY[key] diff --git a/tcri/_keys.py b/tcri/_keys.py new file mode 100644 index 0000000..5bf5940 --- /dev/null +++ b/tcri/_keys.py @@ -0,0 +1,52 @@ +"""Canonical AnnData key registry — the single source of truth for every +``uns`` / ``obsm`` / ``obs`` string tcri reads or writes. + +Import as ``from tcri import _keys as K`` and use ``K.P_CT`` etc. **Never** write +a ``"tcri_*"`` string literal in a function signature or body** — a PR1 +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``. +""" + +# ── uns: metadata + learned priors ─────────────────────────────────────────── +METADATA = "tcri_metadata" # {covariate_col, clone_col, phenotype_col, batch_col} +P_CT = "tcri_p_ct" # learned posterior-mean p_ct, shape (n_ct, P) +LOCAL_SCALE = "tcri_local_scale" # Dirichlet total-concentration scale +GATE_PROB = "tcri_gate_prob" # NEW (Phase 4): classifier/prior gate, scalar or None +CLASSIFIER_TEMPERATURE = "tcri_classifier_temperature" # NEW (Phase 4): classifier temperature + +CT_TO_COV = "tcri_ct_to_cov" # ct -> covariate index +CT_TO_C = "tcri_ct_to_c" # ct -> clonotype index +CT_ARRAY = "tcri_ct_array_for_cells" # per-cell ct index +COV_ARRAY = "tcri_cov_array_for_cells" # per-cell covariate index + +COVARIATE_CATEGORIES = "tcri_covariate_categories" +CLONOTYPE_CATEGORIES = "tcri_clonotype_categories" +PHENOTYPE_CATEGORIES = "tcri_phenotype_categories" + +# ── obsm ───────────────────────────────────────────────────────────────────── +X_TCRI = "X_tcri" # latent posterior mean z +X_LOGITS = "X_tcri_logits" # per-cell classifier logits +X_LOGPOSTERIOR = "X_tcri_logposterior" # logits + log prior +X_PROBABILITIES = "X_tcri_probabilities" # per-cell phenotype probabilities +X_UMAP = "X_umap" + +# ── obs ────────────────────────────────────────────────────────────────────── +PHENOTYPE = "tcri_phenotype" # hard phenotype label +CLONE_SIZE = "clone_size" +INDICES = "indices" # scvi registration glue (kept) + +# ── metadata sub-keys (values inside uns[METADATA]) ────────────────────────── +COVARIATE_COL = "covariate_col" +CLONE_COL = "clone_col" +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 diff --git a/tcri/_stats.py b/tcri/_stats.py new file mode 100644 index 0000000..6468298 --- /dev/null +++ b/tcri/_stats.py @@ -0,0 +1,111 @@ +"""Shared statistics helpers. + +Consolidates the significance / AUROC helpers that lived in ``utils`` (used by +plotting) and adds the posterior-summary primitives the metric layer needs +(true HDI, equal-tailed interval, signed-direction probability). The metric and +plotting layers import from here; ``utils`` will drop its copies during adoption. +""" +from __future__ import annotations + +import itertools +import math + +import numpy as np +from scipy.stats import mannwhitneyu +from sklearn.metrics import roc_auc_score + + +def stars(p: float) -> str: + """Significance stars for a p-value.""" + if p < 1e-4: + return "****" + if p < 1e-3: + return "***" + if p < 1e-2: + return "**" + if p < 0.05: + return "*" + return "ns" + + +def mann_whitney(a, b, *, alternative: str = "two-sided"): + """Mann–Whitney U + two-sided p (thin wrapper for a single import site).""" + return mannwhitneyu(np.asarray(a, float), np.asarray(b, float), alternative=alternative) + + +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).""" + 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_pos = int(y.sum()) + n_exact = math.comb(len(y), 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_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_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 + + +def bootstrap_auc(scores, labels, pos_label=None, n_boot=5000, seed=42): + """Bootstrap 95% CI for AUROC.""" + 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) + rng = np.random.default_rng(seed) + idx = np.arange(len(y)) + aucs = [] + while len(aucs) < n_boot: + samp = rng.choice(idx, size=len(idx), replace=True) + if len(np.unique(y[samp])) < 2: + continue + aucs.append(roc_auc_score(y[samp], scores[samp])) + return np.quantile(aucs, [0.025, 0.975]) + + +# ── posterior-summary primitives (metric layer) ────────────────────────────── +def eti(samples, *, prob: float = 0.94): + """Equal-tailed credible interval (percentile). Stable, transform-invariant.""" + s = np.asarray(samples, float) + lo = (1.0 - prob) / 2.0 + return tuple(np.percentile(s, [100 * lo, 100 * (1 - lo)])) + + +def hdi(samples, *, prob: float = 0.94): + """True highest-density interval: the *narrowest* window holding ``prob`` mass. + + Sounder than ``eti`` for the bounded, skewed entropy/flux posteriors, but + noisier from few draws near a boundary (use ``n_samples ≳ 500`` when tight). + """ + s = np.sort(np.asarray(samples, float)) + n = s.size + if n == 0: + return (np.nan, np.nan) + inc = max(1, int(np.floor(prob * n))) # points spanned by the interval + if inc >= n: + return (float(s[0]), float(s[-1])) + widths = s[inc:] - s[:n - inc] # width of every inc-spanning window + i = int(np.argmin(widths)) + return (float(s[i]), float(s[i + inc])) + + +def prob_direction(delta): + """Signed-contrast probabilities for a difference-draw vector.""" + d = np.asarray(delta, float) + p_gt = float((d > 0).mean()) + return p_gt, 1.0 - p_gt diff --git a/tcri/metrics/_metrics.py b/tcri/metrics/_metrics.py index b38a98b..6ff3d51 100644 --- a/tcri/metrics/_metrics.py +++ b/tcri/metrics/_metrics.py @@ -1,5 +1,6 @@ # Standard library imports import warnings +from .. import _keys as K import numpy as np import pandas as pd import torch @@ -44,17 +45,7 @@ RESET = "\x1b[0m"; BOLD = "\x1b[1m"; DIM = "\x1b[2m" GRN = "\x1b[32m"; CYN = "\x1b[36m"; MAG = "\x1b[35m"; YLW = "\x1b[33m"; RED = "\x1b[31m" -def _ok(msg:str, quiet=False): # success mark - if not quiet: print(f"{GRN}✅ {msg}{RESET}") - -def _info(key:str, txt:str, quiet=False): # key-value info line - if not quiet: print(f" {CYN}🎯 {key:<22}{DIM}{txt}{RESET}") - -def _warn(msg:str, quiet=False): # warning line - if not quiet: print(f"{YLW}⚠️ {msg}{RESET}") - -def _fin(quiet=False): # final flourish - if not quiet: print(f"{MAG}✨ Done!{RESET}") +from .._console import _ok, _info, _warn, _fin # ╰──────────────────────────────────────────────────────────────────────────╯ @@ -93,7 +84,7 @@ def _mi_from_joint(pxy: np.ndarray, normalised: bool, mode: str="average") -> fl 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["tcri_metadata"] + 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"] @@ -234,19 +225,62 @@ def clonotypic_entropy( combine_with_logits: bool = True, _clones: Optional[List[str]] = None, ) -> Union[pd.Series, np.ndarray]: - """ - H[ P(C | P=p, T=covariate) ] for every phenotype p. + 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 ------- - point_estimate=True -> pd.Series indexed by phenotype name, - values are mean H over n_samples posterior draws. - point_estimate=False -> np.ndarray of shape (n_samples, n_phenotypes). + 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["tcri_phenotype_categories"]) + phenotypes = list(adata.uns[K.PHENOTYPE_CATEGORIES]) n_pheno = len(phenotypes) samples = np.empty((n_samples, n_pheno), dtype=float) @@ -373,7 +407,7 @@ def delta_entropy_table( if seed is not None: np.random.seed(seed) - meta = adata.uns["tcri_metadata"] + meta = adata.uns[K.METADATA] clone_col = meta["clone_col"] phen_col = meta["phenotype_col"] @@ -433,19 +467,58 @@ def phenotypic_entropy( temperature: float = 1.0, combine_with_logits: bool = True, ) -> Union[pd.Series, np.ndarray]: - """ - H[ P(P | C=c, T=covariate) ] for every clone c present at the covariate. + 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 ------- - point_estimate=True -> pd.Series indexed by clone name, - values are mean H over n_samples posterior draws. - point_estimate=False -> np.ndarray of shape (n_samples, n_clones). + 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["tcri_metadata"] + meta = adata.uns[K.METADATA] clone_col = meta["clone_col"] covariate_col = meta["covariate_col"] @@ -492,6 +565,34 @@ def phenotypic_entropy( 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() @@ -534,11 +635,64 @@ def mutual_information( verbose: bool = True, graph: bool = False, ) -> Union[float, np.ndarray]: - """ - MI between clonotype and phenotype at one covariate value. + 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. - posterior=True → Dirichlet draw of p_ct (+ optional logits) - posterior=False → prior-only (uses your original joint_distribution). + 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: @@ -625,7 +779,7 @@ def flux_table( if seed is not None: np.random.seed(seed) - meta = adata.uns["tcri_metadata"] + meta = adata.uns[K.METADATA] clone_col = meta["clone_col"] groups = sorted(adata.obs[splitby].dropna().unique().tolist()) @@ -700,20 +854,65 @@ def flux( graph : bool = False, # ASCII histogram seed : Optional[int] = 42 ) -> Union[pd.Series, np.ndarray]: - """ - Flux distance D(p_clone^from , p_clone^to) per clone. + 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 ------- - • `pd.Series` (index = clone_id) if `n_samples == 0` - • `np.ndarray` shape = (n_samples, n_clones) otherwise - (rows correspond to posterior draws) + 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["tcri_metadata"]["clone_col"] + clone_col = adata.uns[K.METADATA]["clone_col"] if clones is None: clones = adata.obs[clone_col].unique().tolist() elif isinstance(clones, str): diff --git a/tcri/plotting/_plotting.py b/tcri/plotting/_plotting.py index adb38a6..1a17053 100644 --- a/tcri/plotting/_plotting.py +++ b/tcri/plotting/_plotting.py @@ -1,4 +1,5 @@ import numpy as np +from .. import _keys as K import pandas as pd import seaborn as sns import matplotlib.pyplot as plt @@ -34,17 +35,7 @@ RESET = "\x1b[0m"; BOLD = "\x1b[1m"; DIM = "\x1b[2m" GRN = "\x1b[32m"; CYN = "\x1b[36m"; MAG = "\x1b[35m"; YLW = "\x1b[33m"; RED = "\x1b[31m" -def _ok(msg:str, quiet=False): # success mark - if not quiet: print(f"{GRN}✅ {msg}{RESET}") - -def _info(key:str, txt:str, quiet=False): # key-value info line - if not quiet: print(f" {CYN}🎯 {key:<22}{DIM}{txt}{RESET}") - -def _warn(msg:str, quiet=False): # warning line - if not quiet: print(f"{YLW}⚠️ {msg}{RESET}") - -def _fin(quiet=False): # final flourish - if not quiet: print(f"{MAG}✨ Done!{RESET}") +from .._console import _ok, _info, _warn, _fin # ╰──────────────────────────────────────────────────────────────────────────╯ red = "#cd442a" @@ -89,7 +80,7 @@ def _fin(quiet=False): # final flourish sns.set_palette(sns.color_palette(tcri_colors)) from ..metrics._metrics import mi_compare as mi_compare_tl -from ..utils._utils import auc_and_label_permutation, bootstrap_auc, stars +from .._stats import auc_and_label_permutation, bootstrap_auc def mi_compare(adata, groupby, groups=None, treatment=None, n_samples=50, @@ -192,7 +183,7 @@ 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["tcri_metadata"]["covariate_col"] + 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) @@ -294,13 +285,13 @@ def _phenotype_mass_per_clone(adata, covariate, clones, normalize): probability vector unweighted; with normalize=False, each row is weighted by the clone's cell count at this covariate. """ - phenotypes = list(adata.uns["tcri_phenotype_categories"]) + 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["tcri_metadata"] + meta = adata.uns[K.METADATA] clone_col = meta["clone_col"] cov_col = meta["covariate_col"] counts_at_cov = ( @@ -347,7 +338,7 @@ def plot_pheno_sankey( construction. Flow geometry preserves the per-clone outer-product semantics of the original implementation. """ - phenotypes = list(adata.uns["tcri_phenotype_categories"]) + phenotypes = list(adata.uns[K.PHENOTYPE_CATEGORIES]) n_phen = len(phenotypes) n_reps = len(covariate_order) if n_reps == 0: @@ -503,7 +494,7 @@ def phenotypic_flux( """ times = list(range(len(order))) if phenotype_colors is None: - phenotype_colors = dict(zip(adata.uns["tcri_phenotype_categories"], tcri_colors)) + phenotype_colors = dict(zip(adata.uns[K.PHENOTYPE_CATEGORIES], tcri_colors)) fig, ax = plot_pheno_sankey( adata, covariate_order=order, @@ -683,7 +674,7 @@ def clonotypic_entropy_by_phenotype( """Box-and-dot plot of clonotypic entropy per phenotype / covariate.""" # ---- meta columns --------------------------------------------- # - meta = adata.uns["tcri_metadata"] + meta = adata.uns[K.METADATA] cov_col = meta["covariate_col"] clone_col = meta["clone_col"] phen_col = meta["phenotype_col"] @@ -811,7 +802,7 @@ def clonotypic_entropy_by_phenotype( 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["tcri_phenotype_categories"] + phenotypes = adata.uns[K.PHENOTYPE_CATEGORIES] prob_labels = [] adata = adata.copy() for y,x in zip(phenotypes, adata.obsm[phenotype_prob_slot].T): @@ -825,7 +816,7 @@ def clone_size_umap(adata, reduction="umap",figsize=(10,8),size=1,alpha=0.7,pale clone_size(adata) df = adata.obs reduction="umap" - sizes = np.log10(adata.obs["clone_size"].to_numpy()) + 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 @@ -975,10 +966,10 @@ def ridge_delta_entropy( 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["tcri_metadata"]["covariate_col"] - clone_col = adata.uns["tcri_metadata"]["clone_col"] - phenotype_col = adata.uns["tcri_metadata"]["phenotype_col"] - batch_col = adata.uns["tcri_metadata"]["batch_col"] + 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() @@ -1098,7 +1089,7 @@ def mutual_information(adata, splitby=None, temperature=1.0, n_samples=0, normal 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"] + 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 @@ -1140,10 +1131,10 @@ def mutual_information(adata, splitby=None, temperature=1.0, n_samples=0, normal palette = tcri_colors # Retrieve metadata from adata - cov_col = adata.uns["tcri_metadata"]["covariate_col"] - clone_col = adata.uns["tcri_metadata"]["clone_col"] - phenotype_col = adata.uns["tcri_metadata"]["phenotype_col"] - batch_col = adata.uns["tcri_metadata"]["batch_col"] + 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() @@ -1247,7 +1238,7 @@ def bayesian_mutual_information( ): np.random.seed(seed) - meta = adata.uns["tcri_metadata"] + meta = adata.uns[K.METADATA] cov_col = meta["covariate_col"] clone_col = meta["clone_col"] @@ -1358,7 +1349,7 @@ def polar_plot(adata, phenotypes=None, statistic="distribution", method="joint_d 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"] + 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" @@ -1366,7 +1357,7 @@ def polar_plot(adata, phenotypes=None, statistic="distribution", method="joint_d 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"] + in adata.uns['tcri_metadata']["covariate_col"] color_dict : dict, optional Dictionary mapping split categories to colors temperature : float, default=1.0 @@ -1388,10 +1379,10 @@ def polar_plot(adata, phenotypes=None, statistic="distribution", method="joint_d >>> ax = tcri.pl.polar_plot(adata, statistic="entropy", color_dict=color_dict) """ if phenotypes is None: - phenotypes = adata.uns["tcri_metadata"]["phenotype_col"] + phenotypes = adata.uns[K.METADATA]["phenotype_col"] if splitby is None: - splitby = adata.uns["tcri_metadata"]["covariate_col"] + splitby = adata.uns[K.METADATA]["covariate_col"] # Get unique splits splits = adata.obs[splitby].unique() @@ -1412,7 +1403,7 @@ def polar_plot(adata, phenotypes=None, statistic="distribution", method="joint_d subset = adata[adata.obs[splitby] == split] values = np.zeros(len(phenotypes)) for j, pheno in enumerate(phenotypes): - mask = subset.obs[adata.uns["tcri_metadata"]["phenotype_col"]] == pheno + mask = subset.obs[adata.uns[K.METADATA]["phenotype_col"]] == pheno values[j] = np.sum(mask) / len(subset) else: # entropy values = np.zeros(len(phenotypes)) diff --git a/tcri/preprocessing/_preprocessing.py b/tcri/preprocessing/_preprocessing.py index e827e81..9cfc0eb 100644 --- a/tcri/preprocessing/_preprocessing.py +++ b/tcri/preprocessing/_preprocessing.py @@ -1,4 +1,5 @@ from scipy.stats import entropy +from .. import _keys as K import numpy as np import tqdm import pandas as pd @@ -47,17 +48,7 @@ RESET = "\x1b[0m"; BOLD = "\x1b[1m"; DIM = "\x1b[2m" GRN = "\x1b[32m"; CYN = "\x1b[36m"; MAG = "\x1b[35m"; YLW = "\x1b[33m"; RED = "\x1b[31m" -def _ok(msg:str, quiet=False): # success mark - if not quiet: print(f"{GRN}✅ {msg}{RESET}") - -def _info(key:str, txt:str, quiet=False): # key-value info line - if not quiet: print(f" {CYN}🎯 {key:<22}{DIM}{txt}{RESET}") - -def _warn(msg:str, quiet=False): # warning line - if not quiet: print(f"{YLW}⚠️ {msg}{RESET}") - -def _fin(quiet=False): # final flourish - if not quiet: print(f"{MAG}✨ Done!{RESET}") +from .._console import _ok, _info, _warn, _fin def _ascii_hist(samples, bins=25, width=40) -> str: hist, edges = np.histogram(samples, bins=bins) @@ -92,14 +83,14 @@ 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="tcri_phenotype"): +def classify_phenotypes(adata, phenotype_prob_slot="X_tcri_phenotypes", phenotype_assignment_obs=K.PHENOTYPE): print("\t...classifying phenotypes...\n") - phenotype_col = adata.uns["tcri_metadata"]["phenotype_col"] - ct_array = adata.uns["tcri_ct_array_for_cells"] + 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["tcri_p_ct"] - phenotypes = adata.uns["tcri_phenotype_categories"] - latent_z = adata.obsm["X_tcri"] + 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([ @@ -164,9 +155,9 @@ def _compute_logits_and_prior(model, adata, batch_size=256, eps=1e-8): @torch.no_grad() def register_model( adata, model, - phenotype_prob_slot="X_tcri_probabilities", - phenotype_assignment_obs="tcri_phenotype", - latent_slot="X_tcri", + 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, @@ -181,12 +172,12 @@ def register_model( print(f"{BOLD}{MAGENT}🔗 Registering TCRi model outputs …{RESET}") # 1) priors & arrays ------------------------------------------------- - adata.uns["tcri_p_ct"] = model.module.get_p_ct().cpu().numpy() - adata.uns["tcri_ct_to_cov"] = model.module.ct_to_cov.cpu().numpy() - adata.uns["tcri_ct_to_c"] = model.module.ct_to_c.cpu().numpy() - adata.uns["tcri_local_scale"] = model.module.local_scale + 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 ("tcri_p_ct","tcri_ct_to_cov","tcri_ct_to_c"): + 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 -------------------------------------------------------- @@ -196,7 +187,7 @@ def register_model( "phenotype_col": model.adata_manager.registry["phenotype_col"], "batch_col": model.adata_manager.registry["batch_col"], } - adata.uns["tcri_metadata"] = meta + adata.uns[K.METADATA] = meta _ok("stored metadata dictionary") # categories @@ -209,9 +200,9 @@ def register_model( # per-cell ct / cov arrays ct_arr = model.module.ct_array.cpu().numpy() - adata.uns["tcri_ct_array_for_cells"] = ct_arr + adata.uns[K.CT_ARRAY] = ct_arr cov_arr = model.module.ct_to_cov.cpu().numpy()[ct_arr] - adata.uns["tcri_cov_array_for_cells"] = cov_arr + adata.uns[K.COV_ARRAY] = cov_arr _ok("stored per-cell ct / cov indices") # 3) latent means ---------------------------------------------------- @@ -223,10 +214,10 @@ def register_model( # 4) logits & log-posterior ----------------------------------------- cls_logits, prior_log = _compute_logits_and_prior(model, adata, batch_size) if store_logits: - adata.obsm["X_tcri_logits"] = cls_logits + adata.obsm[K.X_LOGITS] = cls_logits _info("obsm['X_tcri_logits']", cls_logits.shape) if store_logposterior: - adata.obsm["X_tcri_logposterior"] = cls_logits + prior_log + adata.obsm[K.X_LOGPOSTERIOR] = cls_logits + prior_log _info("obsm['X_tcri_logposterior']", cls_logits.shape) _ok("computed logits & additive log-posterior") @@ -239,7 +230,7 @@ def register_model( adata.obs[phenotype_assignment_obs] = pd.Categorical.from_codes( adata.obsm[phenotype_prob_slot].argmax(1), - categories=adata.uns["tcri_phenotype_categories"], + categories=adata.uns[K.PHENOTYPE_CATEGORIES], ) _ok("stored probabilities and hard labels") @@ -265,12 +256,12 @@ 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["tcri_metadata"]; cov_col = meta["covariate_col"] - clone_col = meta["clone_col"]; ph_cats = adata.uns["tcri_phenotype_categories"] - cov_idx = adata.uns["tcri_covariate_categories"].index(covariate_label) + 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["tcri_ct_array_for_cells"] - cov_per_cell = adata.uns["tcri_cov_array_for_cells"] + 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 @@ -296,8 +287,8 @@ def joint_distribution_posterior( _ok(f"selected {len(idx_cov):,} cells", silent) - p_ct_mean = torch.tensor(adata.uns["tcri_p_ct"]) - local_scale = adata.uns.get("tcri_local_scale", 1.0) + 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] @@ -306,9 +297,9 @@ def joint_distribution_posterior( _ok("sampled one draw from posterior p_ct", silent) if combine_with_logits: - if "X_tcri_logits" not in adata.obsm: + if K.X_LOGITS not in adata.obsm: raise RuntimeError("X_tcri_logits missing in adata.") - logits = adata.obsm["X_tcri_logits"][idx_cov] + 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) @@ -370,15 +361,15 @@ def joint_distribution( weighted: bool = False, ) -> pd.DataFrame: - p_ct = torch.tensor(adata.uns["tcri_p_ct"]) - ct_to_cov = torch.tensor(adata.uns["tcri_ct_to_cov"]) - ct_to_c = torch.tensor(adata.uns["tcri_ct_to_c"]) + 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["tcri_covariate_categories"] - phenotype_categories = adata.uns["tcri_phenotype_categories"] - clonotype_categories = adata.uns["tcri_clonotype_categories"] + covariate_categories = adata.uns[K.COVARIATE_CATEGORIES] + phenotype_categories = adata.uns[K.PHENOTYPE_CATEGORIES] + clonotype_categories = adata.uns[K.CLONOTYPE_CATEGORIES] - metadata = adata.uns["tcri_metadata"] + metadata = adata.uns[K.METADATA] covariate_col = metadata["covariate_col"] # Convert covariate_label to index @@ -400,8 +391,8 @@ def joint_distribution( clone_indices = ct_to_c[chosen_idx].numpy() # Get cell counts for each clonotype-covariate pair (for weighting) - ct_array_for_cells = adata.uns["tcri_ct_array_for_cells"] - cov_array_for_cells = adata.uns["tcri_cov_array_for_cells"] + 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) @@ -442,7 +433,7 @@ def joint_distribution( else: # Sample from Dirichlet distribution - local_scale = adata.uns.get("tcri_local_scale", 1.0) + local_scale = adata.uns.get(K.LOCAL_SCALE, 1.0) conc = local_scale * p_ct_for_cov samples = Dirichlet(conc).sample((n_samples,)) @@ -491,7 +482,7 @@ def joint_distribution( def get_latent_embedding( adata, - latent_slot: str = "X_tcri", + latent_slot: str = K.X_TCRI, n_samples: int = 0, posterior_scale: float = 1.0 ) -> "np.ndarray": @@ -506,7 +497,7 @@ def get_latent_embedding( def group_small_clones(adata, patient_key=""): ct = [] - for x, s, p in zip(adata.obs["trb"], adata.obs["clone_size"], adata.obs[patient_key]): + 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: @@ -545,7 +536,7 @@ def gene_entropy(adata, key_added="entropy", batch_key=None, agg_function=None): aggregated_entropies.append(ent) adata.var[key_added] = aggregated_entropies -def clone_size(adata, key_added="clone_size", return_counts=False): +def clone_size(adata, key_added=K.CLONE_SIZE, return_counts=False): tcr_key = adata.uns["tcri_clone_key"] res = np.unique(adata.obs[tcr_key].tolist(), return_counts=True) clone_sizes = dict(zip(res[0],res[1])) diff --git a/tcri/utils/_utils.py b/tcri/utils/_utils.py index 55e3b3e..dd1c338 100644 --- a/tcri/utils/_utils.py +++ b/tcri/utils/_utils.py @@ -1,4 +1,5 @@ from __future__ import print_function, division +from .. import _keys as K import os import sys import numpy as np @@ -27,61 +28,10 @@ import torch as _torch import pyro as _pyro -import math -import itertools import numpy as np -from sklearn.metrics import roc_auc_score - - -def stars(p): - if p < 1e-4: return "****" - if p < 1e-3: return "***" - if p < 1e-2: return "**" - if p < 0.05: return "*" - return "ns" - - -def auc_and_label_permutation(scores, labels, pos_label=None, - n_perm=200_000, seed=42, max_exact=200_000): - 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_pos = int(y.sum()) - n_exact = math.comb(len(y), 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_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_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 - - -def bootstrap_auc(scores, labels, pos_label=None, n_boot=5000, seed=42): - 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) - rng = np.random.default_rng(seed) - idx = np.arange(len(y)) - aucs = [] - while len(aucs) < n_boot: - samp = rng.choice(idx, size=len(idx), replace=True) - if len(np.unique(y[samp])) < 2: - continue - aucs.append(roc_auc_score(y[samp], scores[samp])) - return np.quantile(aucs, [0.025, 0.975]) + + +# stars / auc_and_label_permutation / bootstrap_auc → moved to tcri/_stats.py (PR1). def _ensure_pyro_posterior_params(model, adata) -> None: @@ -562,7 +512,7 @@ def write_adata_safely(adata: "_ad.AnnData", path: str, *, compression: str = "g def _collect_setup_from_adata_or_model(adata: "_ad.AnnData", model: Any) -> Dict[str, Any]: setup: Dict[str, Any] = {} - meta = adata.uns.get("tcri_metadata", {}) + meta = adata.uns.get(K.METADATA, {}) if meta: setup.update({ "phenotype_col": meta.get("phenotype_col"), @@ -594,9 +544,9 @@ def _collect_setup_from_adata_or_model(adata: "_ad.AnnData", model: Any) -> Dict def _restore_category_order(adata: "_ad.AnnData", setup: Dict[str, Any]) -> None: mapping = [ - ("phenotype_col", "tcri_phenotype_categories"), - ("clone_col", "tcri_clonotype_categories"), - ("covariate_col", "tcri_covariate_categories"), + ("phenotype_col", K.PHENOTYPE_CATEGORIES), + ("clone_col", K.CLONOTYPE_CATEGORIES), + ("covariate_col", K.COVARIATE_CATEGORIES), ] for col_key, cats_key in mapping: col = setup.get(col_key) diff --git a/tests/test_contract_conformance.py b/tests/test_contract_conformance.py new file mode 100644 index 0000000..6cfa858 --- /dev/null +++ b/tests/test_contract_conformance.py @@ -0,0 +1,155 @@ +"""Contract conformance — the interface guardrail for the refactor. + +``tcri/_contract.pyi`` freezes the target public surface. Each *implemented* +function's live signature is checked against its ``.pyi`` declaration (parameter +names, kinds, and which carry defaults). Declared-but-absent functions are the +worklist. Onboard a newly implemented function by adding it to ``IMPLEMENTED`` +as its PR lands. + +AST logic ported from grafiti's ``test_contract_conformance.py``. +""" +import ast +import importlib +import inspect +from pathlib import Path + +import pytest + +import tcri + +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]] = {} + + +def _params_from_ast(a: ast.arguments): + """(name, kind, has_default) per parameter of a .pyi FunctionDef.""" + params = [] + positional = list(a.posonlyargs) + list(a.args) + n_def = len(a.defaults) + for i, arg in enumerate(positional): + kind = "POSITIONAL_ONLY" if arg in a.posonlyargs else "POSITIONAL_OR_KEYWORD" + params.append((arg.arg, kind, i >= len(positional) - n_def)) + if a.vararg: + params.append((a.vararg.arg, "VAR_POSITIONAL", False)) + for arg, default in zip(a.kwonlyargs, a.kw_defaults): + params.append((arg.arg, "KEYWORD_ONLY", default is not None)) + if a.kwarg: + params.append((a.kwarg.arg, "VAR_KEYWORD", False)) + return params + + +def _strip_receiver(params): + """Drop a leading self/cls so contract methods compare to live signatures.""" + if params and params[0][0] in ("self", "cls"): + return params[1:] + return params + + +def _live_params(fn): + return [ + (p.name, p.kind.name, p.default is not inspect.Parameter.empty) + for p in inspect.signature(fn).parameters.values() + ] + + +def _contract_signatures(): + """Parse the .pyi: {"Namespace.func": [(name, kind, has_default), ...]}. + + Namespace classes (tl/pp/pl/diag/ut) and the real TCRIModel class both hold + their functions as method FunctionDefs; keys are ``ClassName.funcname``. + """ + tree = ast.parse(PYI.read_text()) + sigs = {} + for node in tree.body: + if isinstance(node, ast.ClassDef): + for m in node.body: + if isinstance(m, ast.FunctionDef): + sigs[f"{node.name}.{m.name}"] = _strip_receiver(_params_from_ast(m.args)) + elif isinstance(node, ast.FunctionDef): + sigs[node.name] = _params_from_ast(node.args) + return sigs + + +CONTRACT = _contract_signatures() + + +def test_contract_pyi_parses(): + """The frozen contract is present and declares a non-trivial surface.""" + assert PYI.exists(), "tcri/_contract.pyi missing" + assert len(CONTRACT) >= 20, f"contract looks truncated: {len(CONTRACT)} entries" + # every declared key is Namespace-qualified (no accidental bare defs) + assert all("." in k for k in CONTRACT), "unexpected un-namespaced contract entry" + + +@pytest.mark.parametrize("key", sorted(IMPLEMENTED)) +def test_signature_matches_contract(key): + """Each implemented function's live signature matches the frozen contract.""" + module, attr = IMPLEMENTED[key] + obj = importlib.import_module(module) + for part in attr.split("."): + obj = getattr(obj, part) + live = _strip_receiver(_live_params(obj)) + assert live == CONTRACT[key], ( + f"signature drift for {key}:\n contract={CONTRACT[key]}\n live ={live}" + ) + + +def test_report_unimplemented(capsys): + """Informational worklist: declared-but-not-yet-implemented (never fails).""" + todo = sorted(set(CONTRACT) - set(IMPLEMENTED)) + with capsys.disabled(): + print(f"\n[contract] implemented {len(IMPLEMENTED)}/{len(CONTRACT)}; " + f"remaining worklist ({len(todo)}):") + for k in todo: + print(f" ☐ {k}") + assert True + + +def test_import_smoke(): + """`import tcri` is green and the public namespaces resolve.""" + for ns in ("tl", "pp", "pl", "ml", "ut"): + assert hasattr(tcri, ns), f"tcri.{ns} missing" + # diag is added in Phase 8; assert once it lands. + + +# migrated canonical keys (PR1) — must come from _keys.K.*, never a literal. +# legacy keys (tcri_clone_key/…, X_tcri_phenotypes) are exempt until their removal phase. +_MIGRATED_KEYS = [ + "METADATA", "P_CT", "LOCAL_SCALE", "CT_TO_COV", "CT_TO_C", "CT_ARRAY", + "COV_ARRAY", "COVARIATE_CATEGORIES", "CLONOTYPE_CATEGORIES", + "PHENOTYPE_CATEGORIES", "X_LOGITS", "X_LOGPOSTERIOR", "X_PROBABILITIES", + "X_TCRI", "PHENOTYPE", "CLONE_SIZE", +] + + +def test_no_canonical_key_literals(): + """Migrated canonical keys must be accessed via `K.*` **in code** (subscripts / + `.get(...)`), never a string literal. Docstrings and display/log strings may show + the readable key name — this AST check only inspects real key-access code.""" + from tcri import _keys as K + + pkg = Path(tcri.__file__).parent + forbidden = {getattr(K, n) for n in _MIGRATED_KEYS} + offenders = [] + for py in pkg.rglob("*.py"): + if py.name == "_keys.py": + continue + for node in ast.walk(ast.parse(py.read_text())): + key = None + if (isinstance(node, ast.Subscript) and isinstance(node.slice, ast.Constant) + and isinstance(node.slice.value, str)): + key = node.slice.value + elif (isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute) + and node.func.attr == "get" and node.args + and isinstance(node.args[0], ast.Constant) + and isinstance(node.args[0].value, str)): + key = node.args[0].value + if key in forbidden: + offenders.append(f"{py.relative_to(pkg).as_posix()}:{node.lineno}: {key!r}") + assert not offenders, ( + "canonical keys used as code literals (use K.*):\n " + "\n ".join(offenders) + ) diff --git a/tests/test_helpers.py b/tests/test_helpers.py new file mode 100644 index 0000000..137d37d --- /dev/null +++ b/tests/test_helpers.py @@ -0,0 +1,72 @@ +"""Unit tests for the shared helper modules introduced in PR1 +(`_keys`, `_console`, `_stats`, `_distance`). These are pure and fast — no model.""" +import numpy as np +import pytest + +from tcri import _console, _distance as D, _keys as K, _stats as S + + +def test_keys_constants(): + assert K.P_CT == "tcri_p_ct" + 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 + + +def test_console_aliases_and_callables(): + assert _console.MAG == _console.MAGENT + assert _console.GRN == _console.GREEN + assert _console.CYN == _console.CYAN + for fn in (_console._ok, _console._info, _console._warn, _console._fin): + callable(fn) + _console._ok("x", quiet=True) # quiet path prints nothing, must not raise + + +def test_stars_thresholds(): + assert S.stars(1e-5) == "****" + assert S.stars(5e-4) == "***" + assert S.stars(5e-3) == "**" + assert S.stars(0.04) == "*" + assert S.stars(0.2) == "ns" + + +def test_hdi_hugs_skew_and_matches_eti_on_symmetric(): + x = np.array([0, 0, 0, 0, 0.1, 0.2, 5.0]) + hlo, hhi = S.hdi(x, prob=0.8) + elo, ehi = S.eti(x, prob=0.8) + assert (hhi - hlo) <= (ehi - elo) # HDI no wider than ETI on skew + assert hhi < ehi # and it excludes the far tail + g = np.random.default_rng(0).normal(size=20000) + assert np.allclose(S.hdi(g, prob=0.94), S.eti(g, prob=0.94), atol=0.08) + + +def test_prob_direction(): + p_gt, p_lt = S.prob_direction([1, 1, -1, 1.0]) + assert abs(p_gt - 0.75) < 1e-9 and abs(p_lt - 0.25) < 1e-9 + + +def test_auc_helpers_on_perfect_separation(): + scores = np.array([0.1, 0.2, 0.8, 0.9]); labels = np.array([0, 0, 1, 1]) + auc, p_perm, perm, mode = S.auc_and_label_permutation(scores, labels) + assert auc == 1.0 and mode == "exact" and 0.0 <= p_perm <= 1.0 + lo, hi = S.bootstrap_auc(scores, labels) + assert 0.0 <= lo <= hi <= 1.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 + assert D.jensen_shannon([1, 0], [0, 1]) == pytest.approx(1.0, abs=1e-6) # ~1 bit disjoint + assert D.jensen_shannon([0.5, 0.5], [0.5, 0.5]) < 1e-9 # symmetric, self=0 + assert D.l1_distance([1, 0], [0, 1]) == pytest.approx(2.0) + + +def test_distance_dispatch(): + assert D.phenotype_distance("l1")([1, 0], [0, 1]) == pytest.approx(2.0) + assert D.phenotype_distance("dkl") is D.kl_divergence + f = lambda p, q: 0.0 + assert D.phenotype_distance(f) is f + with pytest.raises(ValueError): + D.phenotype_distance("nope")