From 3f0be84fbb1c041a5929a1b5e42d5aa44f43f6dc Mon Sep 17 00:00:00 2001 From: Nicholas Ceglia Date: Mon, 8 Jun 2026 00:04:12 -0400 Subject: [PATCH] docs: add conceptual data-model overview (D2) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit New docs/concepts/data-model.md — the end-to-end mental model for TCRi: the hierarchical generative model (clonotype prior p_c -> local p_ct per clone x covariate -> per-cell phenotype = classifier logits + log prior), the objects register_model writes to .uns/.obsm/.obs (with shapes), the ct_to_c / ct_to_cov / ct_array indexing, point-estimate vs posterior sampling, how the joint distribution feeds the entropy/MI/flux metrics, and the three temperatures. - index.md: add concepts/data-model to the toctree (after quickstart). - conf.py: enable MyST dollarmath so $...$ math renders. Builds clean (0 warnings); cross-references resolve; math renders via mathjax. Closes Notion D2. Co-Authored-By: Claude Opus 4.8 --- docs/concepts/data-model.md | 152 ++++++++++++++++++++++++++++++++++++ docs/conf.py | 3 + docs/index.md | 1 + 3 files changed, 156 insertions(+) create mode 100644 docs/concepts/data-model.md diff --git a/docs/concepts/data-model.md b/docs/concepts/data-model.md new file mode 100644 index 0000000..8951133 --- /dev/null +++ b/docs/concepts/data-model.md @@ -0,0 +1,152 @@ +# Data model & concepts + +This page is the end-to-end mental model for TCRi: the objects it learns, where +they live on your `AnnData`, and how the hierarchical model connects gene +expression, clonotypes, and phenotypes. Read it once before the API reference and +the rest will make sense. + +## The problem + +TCRi works on **paired single-cell data**: every cell has both a gene-expression +profile and a TCR **clonotype** label, and belongs to a **covariate** group +(e.g. timepoint, response status). The questions it answers are *distributional*: +for a given clonotype, what is its distribution over **phenotypes**, and how does +that distribution shift across covariates? From those distributions TCRi computes +information-theoretic summaries (entropies, mutual information, flux). + +Three index sets recur throughout: + +| Symbol | Meaning | Example | +|--------|---------|---------| +| $c$ | clonotype | a TCR clone | +| $m$ | covariate value | `T1`, `T2` | +| $\phi$ | phenotype | `A`, `B`, `C` | + +A **`ct` pair** is a specific $(c, m)$ combination — clonotype $c$ observed at +covariate $m$. Every cell maps to exactly one `ct` pair. + +## The hierarchical model + +`TCRIModel` is a two-level Bayesian model fit with Pyro on top of an scVI-style +variational autoencoder. The generative story: + +1. **Clonotype prior $p_c$.** Each clonotype draws a distribution over phenotypes + from a mixture-Dirichlet prior: $p_c \sim \text{MixtureDirichlet}(\cdot)$, + shape `(n_clonotypes, P)`. This is the *top* of the hierarchy — what a clone + looks like overall. +2. **Local distribution $p_{ct}$.** For each `ct` pair, a covariate-specific + distribution is drawn, anchored to its clone's prior: + $p_{ct} \sim \text{Dirichlet}(\texttt{local\_scale}\cdot p_{c[ct]})$. This lets + a clone's phenotype mix *shift per covariate* while staying tied to the clone + prior. `local_scale` controls how tightly: large → stay near the prior. +3. **Latent $z$.** Each cell's expression $x$ is encoded to a latent $z$ under a + [VampPrior](https://arxiv.org/abs/1705.07120) mixture, and a ZINB decoder + reconstructs counts (the scVI-style expression model). +4. **Per-cell phenotype.** A classifier maps $z$ to phenotype **logits** + (expression-based evidence). Combined additively with the log of the cell's + local prior $\log p_{ct}$, the **per-cell posterior** is + $\operatorname{softmax}(\text{logits} + \log p_{ct})$. + +So a cell's phenotype call fuses two sources: *what its expression looks like* +(classifier logits) and *what its clone tends to be at this covariate* +($p_{ct}$). The variational **guide** learns approximate posteriors +$q(p_c)$ and $q(p_{ct})$; `TCRIModel`'s `get_p_ct()` returns the posterior-mean +$p_{ct}$. + +## What `register_model` writes + +After training, {func}`register_model ` +materializes the learned quantities onto your `AnnData` so every downstream +function can read them. This is the **data model** the rest of the library +assumes: + +| Location | Key | Meaning | Shape | +|----------|-----|---------|-------| +| `.uns` | `tcri_p_ct` | posterior-mean local phenotype distribution per `ct` pair | `(n_ct, P)` | +| `.uns` | `tcri_ct_to_c` | clonotype index for each `ct` pair | `(n_ct,)` | +| `.uns` | `tcri_ct_to_cov` | covariate index for each `ct` pair | `(n_ct,)` | +| `.uns` | `tcri_ct_array_for_cells` | `ct`-pair index for each **cell** | `(n_cells,)` | +| `.uns` | `tcri_cov_array_for_cells` | covariate index for each **cell** | `(n_cells,)` | +| `.uns` | `tcri_local_scale` | Dirichlet concentration scale for posterior sampling | scalar | +| `.uns` | `tcri_{phenotype,clonotype,covariate}_categories` | category label lists (index ↔ name) | — | +| `.uns` | `tcri_metadata` | column-name mapping (`phenotype_col`, `clone_col`, …) | dict | +| `.obsm` | `X_tcri` | latent means $z$ | `(n_cells, n_latent)` | +| `.obsm` | `X_tcri_logits` | classifier phenotype logits (expression evidence) | `(n_cells, P)` | +| `.obsm` | `X_tcri_logposterior` | `logits + log p_ct` | `(n_cells, P)` | +| `.obsm` | `X_tcri_probabilities` | `softmax(logits + log p_ct)` — per-cell phenotype posterior | `(n_cells, P)` | +| `.obs` | `tcri_phenotype` | hard phenotype label (argmax of the posterior) | `(n_cells,)` | + +```{important} +These per-cell `.uns` arrays (`tcri_ct_array_for_cells`, `tcri_cov_array_for_cells`) +are stored in the **original full-cell space**. Slicing the `AnnData` to a view or +subset shifts `.obs`/`.obsm` but **not** `.uns`, so the indices misalign. Re-run +`register_model` on a subset, or pass the full object and filter with the +`clones=` argument. The metric functions guard against this and raise rather than +return silently-wrong numbers. +``` + +### Indexing, concretely + +The `ct_to_c` / `ct_to_cov` arrays are the join keys. To pull the phenotype +distribution of clonotype $c$ at covariate $m$: + +```text +ct_to_c : ct → c (which clone is this pair?) +ct_to_cov : ct → m (which covariate is this pair?) +ct_array : cell → ct (which pair does this cell belong to?) +tcri_p_ct : ct → p(φ) (the pair's phenotype distribution) +``` + +A cell's local prior is therefore `tcri_p_ct[ct_array[cell]]`, and all cells of a +given covariate are `cov_array_for_cells == m`. + +## Point estimate vs. posterior samples + +Two ways to read the clone→phenotype distributions: + +- **Point estimate** — use the posterior mean `tcri_p_ct` directly. +- **Posterior samples** — draw + $p_{ct} \sim \text{Dirichlet}(\texttt{local\_scale}\cdot \bar p_{ct})$ to + propagate uncertainty into the metrics. + +{func}`joint_distribution ` +and {func}`joint_distribution_posterior ` +expose both; passing `n_samples > 0` returns a stack of samples instead of a single +point estimate. + +## From distributions to metrics + +The **joint distribution** $p(c, \phi)$ for a covariate — clonotypes × phenotypes, +optionally weighted by clone size — is the input to every information-theoretic +metric in `tcri.tl`: + +- **Phenotypic entropy** $H(\phi \mid c)$ — how phenotypically mixed a clone is. +- **Clonotypic entropy** $H(c \mid \phi)$ — how clonally diverse a phenotype is. +- **Mutual information** $I(c; \phi)$ — how much knowing the clonotype tells you + about phenotype (clone–phenotype coupling). +- **Flux** — change in a clone's phenotype distribution between two covariates. + +## Temperatures + +Three distinct "temperatures" sharpen ($T<1$) or flatten ($T>1$) distributions at +different stages — don't conflate them: + +| Parameter | Acts on | +|-----------|---------| +| `prior_temperature` | the fixed clone→phenotype prior, at model setup | +| `guide_temperature` | the learned variational posteriors $q(p_c)$, $q(p_{ct})$ | +| `temperature` (in `joint_distribution*`) | the combined per-cell distribution at query time | + +## Typical flow + +```text +setup_anndata → TCRIModel(...) → model.train(...) → register_model(adata, model) + │ + ┌───────────────────────────────┴───────────────┐ + tcri.tl metrics tcri.pl plots + (entropy, mutual_information, flux) (sankey, probabilities, …) +``` + +With the objects above on your `AnnData`, the API reference for +[preprocessing](../api/preprocessing.md), [metrics](../api/metrics.md), and +[plotting](../api/plotting.md) tells you the exact call signatures. diff --git a/docs/conf.py b/docs/conf.py index aa8d722..564619f 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -47,6 +47,9 @@ def _get_release(): 'myst_parser', ] +# MyST: enable dollar-delimited math ($...$ and $$...$$). +myst_enable_extensions = ["dollarmath"] + # Add any paths that contain templates here, relative to this directory. templates_path = ['_templates'] diff --git a/docs/index.md b/docs/index.md index fc19c93..ed3d7f5 100644 --- a/docs/index.md +++ b/docs/index.md @@ -36,6 +36,7 @@ If you use TCRi in your research, please cite: usage/installation usage/quickstart +concepts/data-model api/model api/preprocessing api/metrics