diff --git a/CHANGELOG.md b/CHANGELOG.md index d8d2419..c624cd7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,6 +3,48 @@ All notable changes to `mirpy-lib` (import `mir`). This project follows semantic versioning; the v3 line is a greenfield ML/embedding rewrite (the classical v1.x/v2 toolkit is frozen on branch `legacy-v2`). +## 3.4.0 — 2026-07-18 + +Minor: a command-line interface, a uv-based dev setup, and a documentation overhaul. No public +Python API removed; one optional-dependency group split out. + +### Added + +- **`mir` command-line interface** (`[project.scripts]`, also `python -m mir.cli`) — the two + embedding scales without writing Python: + - `mir embed clonotypes SAMPLE` → a per-clonotype TCREMP embedding table (`e0…`). + - `mir embed repertoires SAMPLE…` → one repertoire vector `Φ(S)` per sample **per chain** on one + shared basis (`phi0…`), with optional `--mmd` pairwise-distance output. + + Inputs are any format `vdjtools.io` reads (AIRR/vdjtools/MiXCR/immunoSEQ/parquet); output is TSV + or Parquet. See `mir embed -h`. Tests in `tests/test_cli.py`. +- **`mir.repertoire.correct_batch`** — Harmony-like cluster-aware batch correction on a stacked + sample×feature Φ matrix. Removes the batch offset *per soft cluster* (batch-diversity-penalised), + so a batch confounded with a biological cluster is corrected without erasing that biology; reduces + exactly to `mir.cohort.residualize` at `n_clusters=1` / `theta=0` (`prop:batch`). +- **`[ann]` optional-dependency group** for the approximate-NN density backend (`pynndescent`). + +### Changed + +- **Development now uses a repo-local `.venv` via [uv](https://docs.astral.sh/uv/)** instead of + conda. `setup.sh` is rewritten (bash/zsh portable; `--dev-parents`, `--docs`, `--tests`); the + conda `environment.yml` is removed. Runtime is unchanged — still a pure-Python `py3-none-any` wheel. +- **`pynndescent` moved from `[bench]` to the new `[ann]` extra.** `[bench]` is now all pure-wheel + (no numba/llvmlite), so `pip install "mirpy-lib[bench]"` resolves cleanly on any Python. Users of + `density.neighbor_enrichment(backend="ann")` should install `"mirpy-lib[ann]"`. +- **`vdjtools>=3.0.0`** (was `>=2.3.0`). +- Documentation overhaul: a use-case-driven user guide, the two CLI commands documented, an + examples/notebooks page, `mir.cohort` and `mir.bench.eval` added to the API reference, a logo, and + the sample-embedding schematic + real depth-robustness figure. Zero-warning Sphinx build. + +### Fixed + +- `DEFAULT_GAP_POSITIONS = (3, 4, -4, -3)` had three independent definitions + (`distances/junction`, `embedding/tcremp`, `ml/bundle`); now defined once in `distances.junction` + and imported, so the coordinate constant cannot drift. +- `cohort.cluster_samples` docstring described itself; `AntigenMetric` and `mir.bench.eval` gained + the docstrings/module-map entries they were missing. + ## 3.3.0 — 2026-07-17 Minor: one new public parameter, nothing removed or changed. diff --git a/CLAUDE.md b/CLAUDE.md index 4b24283..a7d42ee 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -88,6 +88,20 @@ fix were added to `vdjtools` under the owner's direction (this is that owner's e Opt-in `fit_repertoire_space(n_eigs=r)` swaps the second-moment block's full `D₂(D₂+1)/2` upper-triangle for its top-`r` eigenvalues (rotation-invariant spectrum; default `None` = upper-tri, unchanged — but lossy for the *directional* HLA imprint, so the full triangle stays the recommended block; `benchmark_repertoire_spectral.py`). + Also `fit_repertoire_spaces` (one basis **per locus** — the multi-chain fit), `centroid_atypicality` + (per-sample cosine distance to a group centroid — a Φ-geometry op feeding the digital-donor atypicality channel), + and `correct_batch` (Harmony-like cluster-aware batch correction on a stacked Φ matrix; reduces to + `cohort.residualize` at `n_clusters=1`/`theta=0` — `prop:batch`; test in `test_repertoire.py`). +- `cohort.py` — **the digital donor** (T.7): `fit_donor_embeddings`/`DonorCohort` fuse per-chain identity + (kernel-mean, cross-sample PCA-reduced) ‖ diversity ‖ coverage across loci through **one `ChannelBuilder`**, + with an `extra_channels` hook for the analysis's own tissue/clinical blocks; comparability bites twice + (per-locus `prototype_hash` **and** the stored identity PCA) so `save`/`load` verify every hash and + `transform` is the only held-out path. Plus `residualize` (batch cookbook), `cluster_samples` + (MMD→precomputed-metric cluster), `incidence_biomarkers` (subject-incidence Fisher, delegating to + `vdjtools.biomarker.fisher`). Generalizes the analysis-repo `_tcga_embedding.build_embedding` glue. +- `bench/eval.py` — the scorers `channel_report` consumes (kept out of `explain.py` so it stays scorer-free): + `cv_auc` / `held_out_auc` (classification), `cv_cindex` / `km_logrank` (Cox survival, `[bench]`→lifelines), + `kmer_matrix` (baseline). - `explain.py` — **explainable readouts over any feature matrix** (T7). `ChannelSpec`/`ChannelBuilder`/ `stack_embeddings` attach the name→column map `Φ.vector` does not carry (per-chain blocks merge by name); `channel_report(X, spec, scorer, base=, mode="in"|"out"|"both", n_permutations=)` ablates each @@ -99,14 +113,22 @@ fix were added to `vdjtools` under the owner's direction (this is that owner's e has no clonotype pre-image and it raises. Depends one-way on `repertoire.py`; nothing there changed. - `ml/` — Part 2 (torch), neural codecs + `set_encoder.py` (learned repertoire track: Set-Transformer/DeepRC attention pooling + `SetEncoderBundle`). +- `cli.py` — the `mir` console script (argparse, stdlib): `embed clonotypes` (TCREmp table) / + `embed repertoires` (per-locus `fit_repertoire_space`→`sample_embedding` Φ(S), optional `--mmd`). + Reads via `vdjtools.io.read`; writes TSV/parquet. - `resources/` — `prototypes/` (TSVs + manifest), `gene_library/` (region_annotations.txt), `germline_dist/` (baked `.npz`, from `build_germline_dist.py`). ## Build / test / run -- Conda env **`mirpy`** (Python 3.12; do NOT use `.venv` here). `pip install -e .` - (pure-Python hatchling; no C build). Extras: `[bench] [annotate] [build] [ml] [docs] [dev]`. -- Tests: `python -m pytest tests/ -q` (78 pass; `-m "not integration"` for the ~5s fast tier — - the pynndescent ANN parity test carries a one-time JIT cost). All self-contained on bundled resources. +- **Repo-local `.venv` via uv** (Python 3.12; conda retired 2026-07-18). `bash setup.sh` + (bash/zsh; `--dev-parents` editable-installs `../seqtree ../vdjtools ../vdjmatch`, `--docs`, + `--tests`), or `uv pip install -e ".[dev,bench]"`. Pure-Python hatchling; no C build for `mir` + itself. Extras: `[bench] [ann] [annotate] [build] [ml] [docs] [dev] [examples]` (`[ann]` = + pynndescent, split out of `[bench]` so `[bench]` stays numba-free / clean to resolve). +- CLI: `mir embed clonotypes SAMPLE` (→ clonotype embedding table) / `mir embed repertoires + SAMPLE… ` (→ per-sample-per-chain Φ(S)); `mir/cli.py`, `tests/test_cli.py`. +- Tests: `python -m pytest tests/ -q` (`-m "not integration and not benchmark"` for the fast tier; + torch/pynndescent tests skip unless `[ml]`/`[ann]` installed). All self-contained on bundled resources. - Experiments: `python benchmarks/reproduce_supplementary.py` (theory S1–S3), `python benchmarks/benchmark_vdjdb.py` (Table S1). Analyses: `analyze_prototype_counts.py` (geometry saturates by K≈100 — T.1/S4), `analyze_pc_decomposition.py` (V/J η² ≈0.44/0.49, @@ -126,6 +148,15 @@ fix were added to `vdjtools` under the owner's direction (this is that owner's e GPU only in `mir.ml`: `pick_device()` = **CUDA → MPS → CPU** auto, override `device=`/`MIR_DEVICE`. ## Open loops / next steps +- **Embedding-tier roadmap** (`ROADMAP.md`, 2026-07-17) — the "vdjtools at the embedding level" audit + + plan (three verbs: make / measure / generate-decode). **Phase 0** (robustness + optimization quick wins) + and **Phase 1** (cohort tier: `bench/eval.py`, `repertoire.{fit_repertoire_spaces,centroid_atypicality}`, + `cohort.py` digital donor) are **DONE**. Next: **Phase 2** generative loop (`generate.py` `DescriptorDensity` + + `evolve`/`sample`, `CodecBundle.from_unified/from_decoder`), then multimodal encoders + embedding + trajectory. **Analysis-repo follow-up:** refactor `_tcga_embedding.build_embedding` onto + `cohort.fit_donor_embeddings` (+ `extra_channels` for isotype/composition/atypicality) and re-verify the + pan-cancer ΔC numbers. NB Phase 0 flipped the density default to `backend="kdtree"` — re-verify any + recorded balloon-mode baselines (±1 boundary counts). - **v3.0 remaining**: 10X paired benchmark; docs (Sphinx theory section + notebooks); CI; publish `py3-none-any` wheel; regenerate `generate_prototypes.py` via `vdjtools.model.generate`. - **Bench tuning**: raw kneedle eps over-merges; `cluster(eps_factor=0.4)` recovers the paper diff --git a/README.md b/README.md index 45c725f..3d7aa59 100644 --- a/README.md +++ b/README.md @@ -1,3 +1,13 @@ +

+ + + + + mirpy + +

+

mirpy — ML embeddings for immune repertoires

[![PyPI](https://img.shields.io/pypi/v/mirpy-lib.svg)](https://pypi.org/project/mirpy-lib/) @@ -50,6 +60,24 @@ labels = cluster(pca_denoise(X, n_components=50)) Paired chains concatenate per-chain embeddings via `PairedTCREmp`. Input/output are AIRR polars frames keyed by `vdjtools.io.schema` column names. +## Command line + +`pip install mirpy-lib` also installs a `mir` command for the two embedding scales — no Python +needed. Inputs are any format `vdjtools.io` reads (AIRR TSV, vdjtools, MiXCR, immunoSEQ, parquet). + +```bash +# one repertoire -> per-clonotype embedding table (e0…), the input to clustering / ML +mir embed clonotypes sample.tsv --pca 50 -o clonotypes.parquet + +# a dataset of repertoires -> one fingerprint Φ(S) per sample, per chain (phi0…), on one +# shared basis so the rows are mutually comparable; --mmd also writes the pairwise MMD matrix +mir embed repertoires cohort/*.tsv.gz -o phi.tsv --mmd mmd.tsv +``` + +`mir embed clonotypes -h` / `mir embed repertoires -h` list every flag (species, locus, +prototype count, weight, Φ blocks, …). Sample id defaults to the filename stem; the locus is +inferred per file (or restrict with `--locus`). + ## Recommended presets `TCREmp.from_defaults(species, locus)` uses the per-chain preset when `n_prototypes` is @@ -87,9 +115,11 @@ Xr = pca_denoise(X, n_components=p.n_components_recon) # codec reconstruction | `mir.embedding.tcremp` | `TCREmp` / `PairedTCREmp` — the prototype embedding | | `mir.embedding.pca` | PCA denoising of embeddings | | `mir.distances` | junction distance (`seqtree.gapblock`; `metric`/`matrix`/`alignment` options) + baked germline distances | -| `mir.bench` | VDJdb loader, clustering (`cluster(method=…)`: DBSCAN/HDBSCAN/OPTICS) + F1/retention, theory experiments (incl. `codec_losslessness`) | +| `mir.bench` | VDJdb loader, clustering (`cluster(method=…)`: DBSCAN/HDBSCAN/OPTICS) + F1/retention, theory experiments (incl. `codec_losslessness`), cohort scorers (`bench.eval`: `cv_auc`/`cv_cindex`/`km_logrank`) | | `mir.density` | continuous-density TCRNET/ALICE — enrichment (+ clonal-abundance channel, `backend=` exact/kdtree/ann) + noise-filtering (Theory T6) | -| `mir.repertoire` | sample-level (repertoire) embedding — RFF kernel mean ‖ Hill diversity ‖ second moment; MMD / HLA-stratified distance; motif witness (Theory §T.7) | +| `mir.repertoire` | sample-level (repertoire) embedding — RFF kernel mean ‖ Hill diversity ‖ second moment; MMD / HLA-stratified distance; motif witness; `centroid_atypicality`, multi-locus `fit_repertoire_spaces` (Theory §T.7) | +| `mir.explain` | named-channel fusion (`ChannelBuilder`) + scorer-agnostic ablation (`channel_report`/`channel_drivers`) — which part of Φ carries the signal (§T.7) | +| `mir.cohort` | the **digital donor** — multi-chain `fit_donor_embeddings`/`DonorCohort` (+ `transform`/`save`/`load`) + `residualize` / `cluster_samples` / `incidence_biomarkers` (§T.7) | | `mir.ml` | neural codecs (forward/inverse/Pgen/unified) + learned repertoire `set_encoder` (Set-Transformer/DeepRC) — Part 2, experimental; `[ml]` extra | ## Background subtraction & clustering (`mir.density`) @@ -116,11 +146,12 @@ Use a **biological control** as the background when you have one (e.g. pre- vs p patient vs healthy) — differential enrichment cancels generic public convergence and isolates the antigen-specific response. With no control, `generate_background(locus, n)` samples the vdjtools P_gen model (the ALICE regime); the "water level" of a naive repertoire is handled by the -empirical-null calibration. See `experiments/benchmark_density_{yfv,ankspond,tcrnet}.py`. +empirical-null calibration. The density benchmarks (YFV, ankylosing-spondylitis B27, TCRNET) +live in the companion [`2026-mirpy-analysis`](https://github.com/antigenomics) repo. At whole-repertoire scale, pass `neighbor_enrichment(..., backend="kdtree")` (exact scipy cKDTree, 5–9× faster than the default BallTree) or `backend="ann"` (approximate pynndescent, ~30× faster -past ~10⁵ clones, trading a small conservative undercount); see `experiments/benchmark_ann.py`. +past ~10⁵ clones, trading a small conservative undercount; `pip install "mirpy-lib[ann]"`). ## Sample-level (repertoire) embedding (`mir.repertoire`) @@ -155,16 +186,23 @@ a batch offset is first-order and cancels, while a batch-orthogonal signal (e.g. empirical rule of thumb — **diversity for how-even, the embedding for which-clones**: clone-size phenotypes (age, CMV) are a diversity summary's turf, while clonotype identity (HLA — strongest in TRA and class II) lives in the second moment / witness. A learned co-equal set encoder -(Set-Transformer / DeepRC) is in `mir.ml.set_encoder` (`[ml]` extra). See -`experiments/benchmark_repertoire_*.py`, `BENCHMARKS.md`, and `THEORY.md` T7. +(Set-Transformer / DeepRC) is in `mir.ml.set_encoder` (`[ml]` extra). See [`BENCHMARKS.md`](BENCHMARKS.md) +and [`THEORY.md`](THEORY.md) T7 (the benchmark scripts themselves live in the analysis repo). ## Reproduce the paper +The self-contained theory notebooks run on bundled data: + ```bash -python experiments/reproduce_supplementary.py # supplementary S1–S3 -python experiments/benchmark_vdjdb.py # Table S1 (needs a VDJdb dump) +pip install "mirpy-lib[examples]" +marimo edit notebooks/theory.py # supplementary S1–S3 (distance laws, D↔d, prototype robustness) +marimo edit notebooks/quickstart.py # embed + cluster VDJdb antigens ``` +The full benchmark suite (VDJdb Table S1, density, repertoire/TCGA) and result docs live in the +companion analysis repo [`2026-mirpy-analysis`](https://github.com/antigenomics) — this repo is the +library + CI tests only. + Method: Kremlyakova *et al.*, *TCREMP: a bioinformatic pipeline for efficient embedding of T-cell receptor sequences*, **J Mol Biol** 437 (2025) 169205. @@ -178,12 +216,14 @@ mirpy is CPU-parallel by default and uses the GPU for the neural codecs. Knobs, | Density kNN / balloon | `neighbor_enrichment(..., backend=…)` | `"exact"` (BallTree, **1 core**) | `backend="kdtree"` = exact scipy cKDTree, **all cores** (`workers=-1`), 5–9× faster; `backend="ann"` = pynndescent, auto all-core, ~30× at ≥1e5. Prefer `kdtree` for multicore exact. | | Clustering | `cluster(..., n_jobs=-1)` | sklearn default (1) | forwarded to DBSCAN/OPTICS/HDBSCAN via `**kwargs`; parallelizes the neighbour search. | | BLAS (PCA, RFF, matmul) | `OMP_NUM_THREADS` / `OPENBLAS_NUM_THREADS` env | all cores | numpy/sklearn use the platform BLAS; cap via env if oversubscribed. | -| Neural codecs (`mir.ml`) | `pick_device()` / `device=` / `MIR_DEVICE` env | **CUDA → MPS → CPU**, auto | every `train_*` / codec / bundle takes `device=`; e.g. `MIR_DEVICE=cuda:1 python experiments/train_forward_encoder.py`. Torch-free paths (`density`, `repertoire`) never touch the GPU. | +| Neural codecs (`mir.ml`) | `pick_device()` / `device=` / `MIR_DEVICE` env | **CUDA → MPS → CPU**, auto | every `train_*` / codec / bundle takes `device=`; e.g. `MIR_DEVICE=cuda:1` pins the second GPU. Torch-free paths (`density`, `repertoire`) never touch the GPU. | Rule of thumb: leave `threads=0` (all cores) for embedding; switch density to `backend="kdtree"` for exact multicore or `"ann"` at whole-repertoire scale; the GPU is used only by `mir.ml`. ## Development -Conda env `mirpy` (Python 3.12): `bash setup.sh`. Tests: `python -m pytest tests/ -q`. -See [`CLAUDE.md`](CLAUDE.md) for the architecture and reuse map. +Repo-local `.venv` via [uv](https://docs.astral.sh/uv/) (bash/zsh): `bash setup.sh` — add +`--dev-parents` to editable-install the sibling `seqtree` / `vdjtools` / `vdjmatch` checkouts and +`--tests` to run the fast suite. Tests: `python -m pytest tests/ -q`. See [`CLAUDE.md`](CLAUDE.md) +for the architecture and reuse map. diff --git a/ROADMAP.md b/ROADMAP.md new file mode 100644 index 0000000..fffa630 --- /dev/null +++ b/ROADMAP.md @@ -0,0 +1,91 @@ +# mirpy ROADMAP — "vdjtools, at the embedding level" + +The goal: mirpy is the antigenomics group's **embedding / deep-learning** repertoire library — the +things vdjtools does with counts and graphs (diversity, overlap, biomarkers, generation), re-expressed +**in the language of embeddings**. This file is the durable audit + plan; per-session working notes live +in `~/.claude/plans/`. + +## The organizing idea — three verbs + +Everything the library does is one of three verbs, and the API is shaped to make that explicit: + +- **make** an embedding — `TCREmp.embed` (clonotype) · `density.fit_density_space` (cloud) · + `repertoire.sample_embedding` / `sample_descriptor` (repertoire) · `cohort.fit_donor_embeddings` + (digital donor) · `ml.SetEncoder` (learned). +- **measure across / from** embeddings — `density.neighbor_enrichment` · `repertoire.mmd_distance` / + `mmd_matrix` · `repertoire.class_witness` · `cohort.incidence_biomarkers` · `explain.channel_report`. +- **generate + decode back** — `ml.decoder.SequenceDecoder` · `repertoire.decode_metrics` · + `density.generate_background` · (planned) `generate.*`. + +**Division of labour (the house rule):** the library owns **geometry, fusion, serialization under the +prototype-hash comparability contract**; the analysis owns **which clinical/biological channels and which +statistical model**. `explain.py` ships no scorers and never sees `y`; `bench/eval.py` holds the scorers +the analysis hands it. + +## Tiers (current) + +| Tier | Module(s) | What it makes / measures | +|---|---|---| +| Clonotype | `embedding/`, `distances/`, `ml/` | prototype embedding, junction/germline distance, seq↔emb codec, Pgen regressor | +| Density (cloud) | `density.py` | TCRNET/ALICE neighbour-enrichment in embedding space, spike-in, motif denoise | +| Repertoire | `repertoire.py` | Φ(S) = kernel-mean ‖ Hill ‖ 2nd-moment; MMD; `class_witness`; descriptor | +| Explain | `explain.py` | named-channel fusion + scorer-agnostic ablation | +| **Cohort / digital donor** | `cohort.py` | multi-chain `DonorCohort`, `residualize`, `cluster_samples`, `incidence_biomarkers` | +| Scorers | `bench/eval.py` | `cv_auc` / `cv_cindex` / `km_logrank` / `kmer_matrix` (the `channel_report` closures) | +| Generative | *(planned `generate.py`)* | density over descriptors → evolve/sample → decode to clonotypes | + +## vdjtools → embedding parity + +- **DONE:** IO/schema (reuse vdjtools) · overlap→**MMD** · diversity→**Hill channel** · + TCRNET/ALICE→**density enrichment** · co-occurrence→**2nd-moment block** · Pgen→learned regressor · + biomarker-Fisher→**`cohort.incidence_biomarkers`** (delegate) · sample-clustering→**`cohort.cluster_samples`**. +- **PARTIAL:** vdjmatch annotation→nearest-epitope (benchmark only, no `annotate()`) · k-mer/physchem→named + channels (bus ready) · paired α/β (clonotype concat yes; repertoire-Φ via per-locus `DonorCohort`). +- **MISSING (planned):** clonotype-tracking→**embedding trajectory** · **preprocess** in embedding space + (downsample/batch-correct beyond `residualize`) · the **repertoire generative loop**. + +## Downstream analysis modes + +- **digital immunome / digital donor** — `cohort.fit_donor_embeddings` → `DonorCohort` (multi-chain, + hash-serialized). ✅ +- **disease / exposure classification** — `DonorCohort` + `explain.channel_report(scorer=cv_auc)`. ✅ +- **cancer prognosis & survival** — `channel_report(scorer=cv_cindex)` + `km_logrank`; the TME-aware + channels (isotype/composition/atypicality) injected by the analysis via `extra_channels`. ✅ (library); + full pan-cancer number reproduction is an analysis-repo step (needs `~/hf/airr_tcga`). +- **motif / biomarker detection** — `cohort.incidence_biomarkers` (Fisher, wins at low donor n) beside + `repertoire.class_witness` (geometry). ✅ +- **generative decode to real repertoires/clonotypes** — clonotype codec DONE; repertoire generative loop + **planned** (below). + +## Roadmap (by leverage / risk) + +- **Phase 0 — robustness + optimization quick wins. ✅ DONE** (commit `fix(phase0)`): degenerate-input + guards (unbiased-MMD singleton, empty/zero-count sample, negative-`n` prototypes, zero radius, empty + density trees); density default `backend="kdtree"`; auto-chunk + frame-sample fits; vectorized + `hla_stratified_mmd`; precomputed-witness fast-path. +- **Phase 1 — cohort tier / digital donor. ✅ DONE**: `bench/eval.py` scorers; + `repertoire.{fit_repertoire_spaces, centroid_atypicality}`; `cohort.py` + (`DonorCohort`/`fit_donor_embeddings`/`transform`/`save`/`load`, `residualize`, `cluster_samples`, + `incidence_biomarkers`). **Follow-up (analysis repo):** refactor `_tcga_embedding.build_embedding` onto + `fit_donor_embeddings` (+ `extra_channels` for isotype/composition/atypicality) and re-verify the + pan-cancer ΔC numbers (SKCM +0.039 / BLCA +0.025 / HNSC +0.022 / LGG +0.016). +- **Phase 2 — generative loop, mechanical half** (priority frontier): `CodecBundle.from_unified/from_decoder`; + `generate.py` `DescriptorDensity` (conditional Gaussian) + `evolve` + `sample`, promoting the shipped + in-silico-evolution manifold. Accept: reproduce the coupled in-silico slopes through the library. +- **Phase 3 — generative loop, research half**: RepertoireSpace-PCA ↔ codec-PCA basis bridge → + `invert_kernel_mean` (herd from real candidates + count model + germline-sampled V/J) + Pgen filter. + Accept: re-embedding the generated multiset lands near the target μ. *Fund, iterate.* +- **Phase 4 — multimodal encoders**: `modalities.isotype_fractions` / `hla_indicator` / + `fit_epitope_annotator` (nearest-epitope `annotate`) + paired-Φ recipe. (GEX / pMHC-groove deferred.) +- **Phase 5 — embedding trajectory** (research): `track.repertoire_trajectory` (Φ velocity) + + `clonotype_flux` (differential enrichment). + +## Non-goals / risks + +- **Stays in vdjtools** (delegate, never reimplement): AIRR schema/IO, `downsample`, P_gen generation, + `kmer_profile`, the `fisher_association` engine. +- **Stays analysis-local:** tissue/TME feature engineering (isotype/composition), Cox penalizer / CV scheme. +- **Comparability bites twice** in `DonorCohort`: per-locus `prototype_hash` **and** the stored identity + PCA — `load` verifies all hashes; a `residualize`d `X` is a different coordinate system. +- **Generation** can go off-manifold — herd from real candidates + Pgen filter bound (don't eliminate) the + failure rate; junction-only decoding means V/J are germline-sampled, not jointly learned. diff --git a/docs/_static/custom.css b/docs/_static/custom.css new file mode 100644 index 0000000..46c2579 --- /dev/null +++ b/docs/_static/custom.css @@ -0,0 +1,9 @@ +/* Figures with dark ink on a transparent background (schematics, plots exported from LaTeX/gnuplot) + stay legible in dark mode by rendering on a white card. Scoped to `.mir-fig` so the theme-aware + logo (only-light / only-dark) is untouched. */ +img.mir-fig { + background: #ffffff; + padding: 10px 14px; + border-radius: 8px; + max-width: 100%; +} diff --git a/docs/_static/depth_robustness.png b/docs/_static/depth_robustness.png new file mode 100644 index 0000000..360421c Binary files /dev/null and b/docs/_static/depth_robustness.png differ diff --git a/docs/_static/mirpy_dark.png b/docs/_static/mirpy_dark.png new file mode 100644 index 0000000..097aca2 Binary files /dev/null and b/docs/_static/mirpy_dark.png differ diff --git a/docs/_static/mirpy_dark.svg b/docs/_static/mirpy_dark.svg new file mode 100644 index 0000000..14549df --- /dev/null +++ b/docs/_static/mirpy_dark.svg @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/docs/_static/mirpy_light.png b/docs/_static/mirpy_light.png new file mode 100644 index 0000000..3286ce7 Binary files /dev/null and b/docs/_static/mirpy_light.png differ diff --git a/docs/_static/mirpy_light.svg b/docs/_static/mirpy_light.svg new file mode 100644 index 0000000..a08e415 --- /dev/null +++ b/docs/_static/mirpy_light.svg @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/docs/_static/sample_embedding_diagram.png b/docs/_static/sample_embedding_diagram.png new file mode 100644 index 0000000..8b9eef9 Binary files /dev/null and b/docs/_static/sample_embedding_diagram.png differ diff --git a/docs/api.rst b/docs/api.rst index 33bbf81..d8851da 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -112,6 +112,17 @@ Which named channel of ``Φ`` carries the signal, and which clonotypes drive it :undoc-members: :show-inheritance: +Cohort / digital donor (``mir.cohort``) +--------------------------------------- + +Fuse per-chain repertoire embeddings into one hash-verified, serialisable donor matrix; batch +residualisation, sample clustering, and incidence biomarkers (Theory T7). + +.. automodule:: mir.cohort + :members: + :undoc-members: + :show-inheritance: + Benchmark harness (``mir.bench``) --------------------------------- @@ -145,6 +156,16 @@ Reproduced supplementary theory (S1–S3, T5–T6, codec losslessness). :undoc-members: :show-inheritance: +``mir.bench.eval`` +~~~~~~~~~~~~~~~~~~ + +Scorers for the explainable readout: cross-validated AUC, Cox C-index, log-rank. + +.. automodule:: mir.bench.eval + :members: + :undoc-members: + :show-inheritance: + Neural codecs and learned encoders (``mir.ml``) ----------------------------------------------- @@ -241,3 +262,26 @@ Allele normalisation with default-allele cascade. :undoc-members: :show-inheritance: +Command-line interface (``mir``) +-------------------------------- + +``pip install mirpy-lib`` installs a ``mir`` console script (also ``python -m mir.cli``) with two +commands, one per embedding scale: + +``mir embed clonotypes INPUT`` + One repertoire's clonotype table → a per-clonotype TCREMP embedding table (columns ``e0…``). + Flags: ``--species``, ``--locus`` (inferred when the file has one locus), ``--n-prototypes``, + ``--mode {vjcdr3,cdr123}``, ``--pca K`` (compact the table), ``--threads``, ``-o`` (``.tsv`` / + ``.parquet``; default stdout TSV). + +``mir embed repertoires INPUT...`` + A dataset of clonotype tables → one repertoire vector ``Φ(S)`` per sample **per chain** on one + shared basis (columns ``phi0…``; sample id = filename stem). Flags: ``--locus`` (restrict), + ``--weight {log1p,anscombe,distinct}``, ``--blocks mean,diversity[,second]``, ``--n-rff``, + ``--n-components``, ``--mmd OUT`` (also write the per-chain pairwise unbiased-MMD matrix), + ``--threads``, ``--seed``, ``-o``. + +.. automodule:: mir.cli + :members: main, build_parser + :show-inheritance: + diff --git a/docs/conf.py b/docs/conf.py index 8281f53..84fec2f 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -13,7 +13,7 @@ try: version = release = _pkg_version("mirpy-lib") except PackageNotFoundError: - version = release = "3.1.0" + version = release = "3.4.0" extensions = [ "sphinx.ext.autodoc", @@ -49,6 +49,8 @@ intersphinx_mapping = {"python": ("https://docs.python.org/3", None)} templates_path = ["_templates"] +html_static_path = ["_static"] +html_css_files = ["custom.css"] exclude_patterns = ["_build", "**.ipynb_checkpoints"] html_theme = "pydata_sphinx_theme" diff --git a/docs/examples.rst b/docs/examples.rst new file mode 100644 index 0000000..6d95430 --- /dev/null +++ b/docs/examples.rst @@ -0,0 +1,31 @@ +Examples +======== + +mirpy ships three runnable `marimo `_ notebooks under ``notebooks/``. They are +self-contained (they run on the bundled prototypes / test assets — no downloads) and double as +living documentation for the three tiers. Install the extra and open one: + +.. code-block:: bash + + pip install "mirpy-lib[examples]" # marimo, matplotlib, umap-learn + marimo edit notebooks/quickstart.py # interactive; or `marimo run …` for read-only + +.. list-table:: + :header-rows: 1 + :widths: 22 78 + + * - Notebook + - What it shows + * - ``notebooks/quickstart.py`` + - Clonotype embedding end-to-end: ``TCREmp.embed`` a VDJdb-style set, PCA-denoise, cluster + antigen-specific TCRs, and a UMAP coloured by epitope. + * - ``notebooks/density.py`` + - Background subtraction (Theory T6): fit a density space, run balloon ``neighbor_enrichment`` + against a P_gen / control background, and pull out the enriched convergent family. + * - ``notebooks/theory.py`` + - Reproduces the supplementary results S1–S3 on bundled data — the distance laws + (Gamma / extreme-value), the D↔d correlation, and prototype-source robustness. + +The full benchmark suite (VDJdb Table S1, density, repertoire / TCGA cohorts) and its result docs +live in the companion `2026-mirpy-analysis `_ repository; this repo +keeps the library, its tests, and these bundled-data examples. diff --git a/docs/index.rst b/docs/index.rst index f0fa647..32068bd 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -1,3 +1,15 @@ +.. image:: _static/mirpy_light.png + :class: only-light + :align: center + :width: 340 + :alt: mirpy + +.. image:: _static/mirpy_dark.png + :class: only-dark + :align: center + :width: 340 + :alt: mirpy + mirpy ===== @@ -9,31 +21,56 @@ Pure-Python and built on the antigenomics ecosystem `vdjtools `_, `arda `_). -.. note:: - - **v3.1.0** — the prototype (TCREMP) clonotype embedding plus the Part-2 tier: neural codecs - (``mir.ml``), continuous-density background subtraction (``mir.density``, T6), and a sample-level - repertoire embedding (``mir.repertoire``, T7). Coordinates are arda-native and versioned. The - classical v1.x/v2 toolkit is frozen on the ``legacy-v2`` branch (``mirpy-lib`` 2.x). +mirpy turns receptor sequences into fixed-length vectors at **three scales** — the *clonotype* +(one receptor → a point), the *repertoire* (a whole sample → one fingerprint), and the *cohort* +(many donors → a comparable table) — and gives you the operations to measure across them +(distance, enrichment, MMD), read out what drives a signal, and decode embeddings back to +sequences. The classical v1.x/v2 toolkit is frozen on the ``legacy-v2`` branch (``mirpy-lib`` 2.x). -New here? The :doc:`user guide ` has runnable examples for each module; the mathematical -theory (T1–T7) lives in ``THEORY.md`` and recorded benchmark numbers in ``BENCHMARKS.md``. +New here? Start with the :doc:`user guide ` (install → a real result, one task per section), +skim the runnable :doc:`examples `, and reach for the :doc:`API reference ` for +signatures. The mathematical theory (T1–T7) lives in ``THEORY.md`` and recorded benchmark numbers +in ``BENCHMARKS.md``. -Quickstart — clonotype embedding --------------------------------- +30-second quickstart +-------------------- .. code-block:: python + import polars as pl from mir.embedding.tcremp import TCREmp model = TCREmp.from_defaults("human", "TRB", n_prototypes=1000) - X = model.embed(df) # polars frame (v_call / j_call / junction_aa) -> (N, 3K) float32 + df = pl.DataFrame({ + "v_call": ["TRBV10-3*01", "TRBV20-1*01"], + "j_call": ["TRBJ2-7*01", "TRBJ1-2*01"], + "junction_aa": ["CASSIRSSYEQYF", "CSARVSGYYGYTF"], + }) + X = model.embed(df) # (2, 3K) float32 — distances to 3K prototypes, per V/J/junction Distance in the prototype-embedding space approximates the pairwise alignment distance (Theory T1); ``TCREmp`` computes the junction part via ``seqtree.gapblock`` and adds baked germline V/J distances. -Sample-level (repertoire) embedding ------------------------------------ +Or from the shell — ``pip install mirpy-lib`` ships a ``mir`` command: + +.. code-block:: bash + + mir embed clonotypes sample.tsv -o clonotypes.parquet # per-clonotype embedding table + mir embed repertoires cohort/*.tsv.gz -o phi.tsv --mmd mmd.tsv # one Φ(S) per sample, per chain + +From a bag of receptors to one comparable fingerprint +----------------------------------------------------- + +.. image:: _static/sample_embedding_diagram.png + :align: center + :width: 760 + :class: mir-fig + :alt: A repertoire is a bag of receptors; each receptor is a point; the weighted cloud is + summarised into one fixed-length fingerprint Φ(S); two people are compared by the + distance between their fingerprints. + +The sample-level embedding ``Φ(S)`` (``mir.repertoire``) sketches a repertoire as a weighted cloud +of clonotype points — depth-robust down into the RNA-seq regime (Theory T7): .. code-block:: python @@ -41,7 +78,7 @@ Sample-level (repertoire) embedding space = fit_repertoire_space(model, pooled_clonotypes) # ONE basis for the cohort embs = [sample_embedding(space, s) for s in samples] # Φ(S): mean ‖ diversity ‖ second moment - D = mmd_matrix(embs, unbiased=True) # pairwise repertoire distance (unbiased MMD²) + D = mmd_matrix(embs, unbiased=True) # pairwise repertoire distance (MMD²) Capabilities (see the :doc:`API reference `) ------------------------------------------------- @@ -53,9 +90,12 @@ Capabilities (see the :doc:`API reference `) - **Repertoire embedding** — one fixed vector per sample: RFF kernel mean, coverage-standardised Hill diversity, and a second-moment Fisher block; MMD / HLA-stratified distance and a motif witness (``mir.repertoire``, T7). +- **Digital donor / cohort** — fuse per-chain repertoire embeddings into one multi-chain donor + matrix, hash-verified and serialisable, with batch residualisation and sample clustering + (``mir.cohort``, T7). - **Explainable readouts** — name ``Φ``'s channels and ablate them under *your* scorer to see which - one carries the signal (and whether it is merely redundant), then hop from the winning kernel-mean - channel to the clonotypes driving it (``mir.explain``, T7). + one carries the signal, then hop from the winning kernel-mean channel to the clonotypes driving it + (``mir.explain`` + scorers in ``mir.bench.eval``, T7). - **Neural codecs** — forward / inverse / Pgen / unified codecs and a learned repertoire set encoder; CUDA → MPS → CPU device selection (``mir.ml``, ``[ml]`` extra). - **Benchmark harness** — VDJdb clustering + F1/retention and reproduced theory (``mir.bench``). @@ -65,4 +105,5 @@ Capabilities (see the :doc:`API reference `) self usage + examples api diff --git a/docs/usage.rst b/docs/usage.rst index 499ca36..1e0ae81 100644 --- a/docs/usage.rst +++ b/docs/usage.rst @@ -1,20 +1,38 @@ User guide ========== -Runnable examples for each part of ``mir``. All embeddings operate on ``polars`` frames keyed by the -AIRR column names ``v_call`` / ``j_call`` / ``junction_aa`` (and ``duplicate_count`` for clone sizes). +One task per section, cheapest-to-run first. Every embedding operates on ``polars`` frames keyed by +the AIRR column names ``v_call`` / ``j_call`` / ``junction_aa`` (and ``duplicate_count`` for clone +sizes). Sections build on each other but each is self-contained. -Install -------- +Install & the data model +------------------------- .. code-block:: bash - pip install mirpy-lib # core: import mir + pip install mirpy-lib # core: import mir, and the `mir` CLI + pip install "mirpy-lib[bench]" # + benchmark harness (matplotlib, huggingface_hub, lifelines) + pip install "mirpy-lib[ann]" # + approximate-NN density backend (pynndescent) pip install "mirpy-lib[ml]" # + neural codecs (torch) - pip install "mirpy-lib[bench]" # + benchmark harness (huggingface_hub, pynndescent, ...) + pip install "mirpy-lib[examples]"# + marimo notebooks (marimo, matplotlib, umap-learn) -Requires ``vdjtools>=2.3.0`` and ``seqtree>=0.3.0``; ``mir`` itself is a pure-Python -``py3-none-any`` wheel. +Requires ``vdjtools>=3.0.0`` and ``seqtree>=0.3.0``; ``mir`` itself is a pure-Python +``py3-none-any`` wheel. A clonotype table is a polars frame; clone sizes live in +``duplicate_count`` (repertoire-level embeddings weight by them): + +.. code-block:: python + + import polars as pl + + sample = pl.DataFrame({ + "v_call": ["TRBV10-3*01", "TRBV20-1*01", "TRBV28*01"], + "j_call": ["TRBJ2-7*01", "TRBJ1-2*01", "TRBJ2-1*01"], + "junction_aa": ["CASSIRSSYEQYF", "CSARVSGYYGYTF", "CASSLGQAYEQFF"], + "duplicate_count": [120, 40, 12], + }) + +To read a file in any format instead, use ``vdjtools.io.read`` (AIRR TSV, vdjtools, MiXCR, +immunoSEQ, parquet — the format is sniffed): ``sample = vdjtools.io.read("sample.tsv")``. Clonotype embedding ------------------- @@ -28,7 +46,7 @@ pairwise alignment distance (Theory T1). from mir.embedding.tcremp import TCREmp model = TCREmp.from_defaults("human", "TRB", n_prototypes=1000) - X = model.embed(df) # (N, 3K) float32, interleaved [v, j, junction] + X = model.embed(sample) # (3, 3K) float32, interleaved [v, j, junction] # paired chains: a dict of per-locus frames -> concatenated embedding from mir.embedding.tcremp import PairedTCREmp @@ -45,6 +63,32 @@ Pick prototype counts / PCA dims from the per-chain presets, and denoise with PC preset = get_preset("human", "TRB") # n_prototypes, n_components (95% var), recon dims Xd = pca_denoise(X, n_components=preset.n_components) +**From the shell** — the ``mir embed clonotypes`` command does exactly this on a file: + +.. code-block:: bash + + mir embed clonotypes sample.tsv --pca 50 -o clonotypes.parquet + # input: any format vdjtools.io reads; locus inferred (or pass --locus) + # output: one row per clonotype (id columns + e0…), TSV or (recommended for the wide raw + # embedding) Parquet. --n-prototypes / --mode / --pca / --threads for the knobs. + +Clustering antigen-specific TCRs +-------------------------------- + +Antigen-specific receptors form tight clusters in the embedding. ``mir.bench`` clusters them and +scores per-antigen F1 / retention against known labels (e.g. a VDJdb dump): + +.. code-block:: python + + from mir.bench.metrics import cluster, cluster_metrics + + labels = cluster(pca_denoise(X, n_components=50), method="dbscan") # or "hdbscan" / "optics" + # with epitope labels aligned to the rows: + metrics = cluster_metrics(labels, epitopes) # {epitope: AntigenMetric} + +DBSCAN (default) is tightest/purest; HDBSCAN trades precision for ~3× coverage on variable-density +data. See ``notebooks/quickstart.py`` for the end-to-end VDJdb example with a UMAP. + Density and background subtraction ---------------------------------- @@ -63,14 +107,16 @@ adaptive-bandwidth balloon estimator with a Poisson/binomial test and BH q-value labels, mask = denoise_and_cluster(obs_emb, res) # noise-filter + cluster the hits Prefer a biological control (e.g. pre/post-vaccination) over the P_gen background — differential -enrichment cancels generic public convergence and isolates the antigen-specific response. +enrichment cancels generic public convergence and isolates the antigen-specific response. At +whole-repertoire scale, pass ``backend="kdtree"`` (exact, multicore) or ``backend="ann"`` +(approximate, ``[ann]`` extra). See ``notebooks/density.py``. -Sample-level (repertoire) embedding ------------------------------------ +Repertoire embedding Φ(S) + MMD +------------------------------- ``mir.repertoire`` embeds a whole repertoire — an order-invariant multiset of clonotypes with clone sizes — into one fixed vector ``Φ(S)`` (kernel mean ‖ Hill diversity ‖ second moment), depth-robust -into the RNA-seq regime (Theory T7). Every sample in a cohort must share one basis. +into the RNA-seq regime (Theory T7). Every sample in a cohort must share **one** basis, fit once: .. code-block:: python @@ -79,31 +125,80 @@ into the RNA-seq regime (Theory T7). Every sample in a cohort must share one bas import polars as pl space = fit_repertoire_space(model, pl.concat(samples)) # ONE basis for the cohort - embs = [sample_embedding(space, s) for s in samples] # Φ(S) + embs = [sample_embedding(space, s) for s in samples] # Φ(S); each .vector is the tensor D = mmd_matrix(embs, unbiased=True) # pairwise MMD (unbiased when depth varies) # supervised motif finder: public clones separating two groups motifs = class_witness(space, pos_samples, neg_samples, candidates) -Use ``unbiased=True`` whenever samples differ in depth/diversity. For a batch-confounded contrast, -compare *within-batch* (residualise ``Φ`` on the batch indicator): a batch offset is first-order and -cancels, while a batch-orthogonal signal (e.g. HLA) survives. +Read only a few hundred receptors and Φ is already close to the full-depth fingerprint — the +gap shrinks as ``n_eff^{-1/2}`` (Theory ``prop:kme``): + +.. image:: _static/depth_robustness.png + :align: center + :width: 460 + :class: mir-fig + :alt: The gap between a size-n subsample fingerprint and the full-depth fingerprint falls from + ~0.18 at n=50 to ~0.03 at n=700 (real mir measurement, human TRB). + +Use ``unbiased=True`` whenever samples differ in depth/diversity — the biased V-statistic's +``1/n_eff`` self-term otherwise inflates low-diversity samples and fakes a signal. For a +batch-confounded contrast, compare *within-batch* (residualise ``Φ`` on the batch indicator, see the +cohort section): a batch offset is first-order and cancels, while a batch-orthogonal signal (e.g. +HLA) survives. + +**From the shell** — ``mir embed repertoires`` fits one shared basis per chain and writes the Φ table: -Explainable readouts --------------------- +.. code-block:: bash + + mir embed repertoires cohort/*.tsv.gz -o phi.tsv --mmd mmd.tsv + # one input file per repertoire (sample id = filename stem); one output row per sample per + # chain (sample_id, locus, n_clonotypes, phi0…). --blocks mean,diversity[,second] and + # --weight / --n-rff pick the Φ layout; --mmd also writes the pairwise unbiased-MMD matrix. + +Multi-chain digital donor (``mir.cohort``) +------------------------------------------ + +``mir.cohort`` fuses per-chain repertoire embeddings into one **digital donor** matrix — an +identity (kernel-mean) block plus diversity and coverage channels, per locus, merged by name. The +result is hash-verified (the prototypes *and* the fitted PCA must match) and serialisable, so a +held-out donor is projected through the same basis: + +.. code-block:: python + + from mir.repertoire import fit_repertoire_spaces, correct_batch + from mir.cohort import fit_donor_embeddings, residualize, cluster_samples + + spaces = fit_repertoire_spaces(models, cohort_frames) # one RepertoireSpace per locus + cohort = fit_donor_embeddings(spaces, donor_frames) # DonorCohort{X, spec, ...} + cohort.save("cohort.npz") # carries every locus' prototype hash + + Xc = residualize(cohort.X, batch) # first-order: subtract each batch mean + Xc = correct_batch(cohort.X, batch) # Harmony-like cluster-aware, when the + # batch is confounded with biology (T7 prop:batch) + lab = cluster_samples(embs, unbiased=True) # MMD-cluster repertoires into states + +``residualize`` removes a *global* batch offset; ``correct_batch`` removes it **per soft cluster** +(reducing to ``residualize`` at ``n_clusters=1``), so a batch confounded with a biological cluster is +corrected without erasing that biology. + +Explainable channel readout (``mir.explain``) +--------------------------------------------- ``Φ.vector`` is an anonymous concatenation, so "the classifier found something" has no noun. ``mir.explain`` attaches the names and asks which of them carries the signal. The scorer is yours — -the library never sees the labels, so the same call serves a cross-validated AUC or a Cox C-index: +the library never sees the labels, so the same call serves a cross-validated AUC or a Cox C-index +(scorers live in :mod:`mir.bench.eval`): .. code-block:: python - from mir.explain import channel_report, stack_embeddings + from mir.explain import stack_embeddings, channel_report + from mir.bench.eval import cv_auc - X, spec = stack_embeddings(embs) # X[i] is exactly embs[i].vector; names attached + X, spec = stack_embeddings(embs) # X[i] is exactly embs[i].vector; names attached rep = channel_report(X, spec, lambda B: cv_auc(B, y)[0], base=0.5) - rep.best # -> 'second' (the HLA imprint lives here) - rep.frame() # channel | n_columns | score | delta | rank + rep.best # -> 'second' (the HLA imprint lives here) + rep.frame() # channel | n_columns | score | delta | rank Leave-one-**in** (the default) asks whether a channel carries the signal *alone*; it is marginal, so correlated channels both look important. ``mode="both"`` adds the conditional half — a channel with a @@ -114,6 +209,7 @@ and mark the kernel-mean blocks attributable so the readout can go one hop furth .. code-block:: python from mir.explain import ChannelBuilder, channel_drivers + from mir.bench.eval import cv_cindex b = ChannelBuilder() for c in chains: @@ -126,8 +222,24 @@ and mark the kernel-mean blocks attributable so the readout can go one hop furth Only a kernel-mean channel has a clonotype pre-image. Asking which clones drive a Hill number is a category error, and ``channel_drivers`` raises rather than answer it. -Neural codecs -------------- +Survival / classification scorers (``mir.bench.eval``) +------------------------------------------------------ + +The scorers the readout consumes are plain functions of a feature block and your labels — usable on +their own: + +.. code-block:: python + + from mir.bench.eval import cv_auc, cv_cindex, km_logrank + + mean, std = cv_auc(X, y) # repeated stratified-CV AUC (mean, std) + c = cv_cindex(rows, block=X, base=base) # Cox C-index gain of X over base covariates + p = km_logrank(durations, events, groups) # multivariate log-rank p-value + +Needs the ``[bench]`` extra (scikit-learn is core; ``cv_cindex`` / ``km_logrank`` use ``lifelines``). + +Neural codecs (``mir.ml``) +-------------------------- The optional ``mir.ml`` tier (``[ml]`` extra, torch) trains fast neural approximations of the embedding: a forward encoder (sequence → code), an inverse decoder (code → sequence), a Pgen @@ -137,18 +249,21 @@ automatic (CUDA → MPS → CPU; override with ``device=`` or ``MIR_DEVICE``). .. code-block:: python from mir.ml.bundle import CodecBundle + bundle = CodecBundle.load("path/to/codec") # refuses a prototype-hash mismatch - codes = bundle.encode(X) + encoder = bundle.forward_encoder() # a ForwardEncoder + codes = encoder.encode(sample["junction_aa"].to_list()) # CDR3 strings -> compact codes -Benchmark harness ------------------ +Training scripts and shipped bundles live in the companion analysis repo; this tier is experimental. -``mir.bench`` provides the VDJdb clustering benchmark (F1 / retention / purity) with selectable -clustering methods, and the reproduced theory experiments. +Benchmark harness & reproducing the paper +------------------------------------------ + +``mir.bench`` provides the VDJdb clustering benchmark (F1 / retention / purity) and the reproduced +theory experiments (S1–S3, T5–T6, codec losslessness). The self-contained theory runs on bundled +data via ``notebooks/theory.py``; the full benchmark suite lives in the ``2026-mirpy-analysis`` repo. .. code-block:: python from mir.bench.vdjdb import load_vdjdb from mir.bench.metrics import cluster, cluster_metrics - - labels = cluster(X, method="dbscan") # or "hdbscan" / "optics" diff --git a/environment.yml b/environment.yml deleted file mode 100644 index 8032983..0000000 --- a/environment.yml +++ /dev/null @@ -1,15 +0,0 @@ -# Conda environment for mirpy v3 development. -# -# Minimal: Python + pip only (v3 is a pure-Python package — no C build). All -# dependencies come from the editable pip install; `setup.sh` also editable-installs -# the co-developed siblings (../seqtree, ../vdjtools) when present. -# -# conda env create -f environment.yml -# conda activate mirpy -# bash setup.sh # or: pip install -e ".[dev,bench]" -name: mirpy -channels: - - conda-forge -dependencies: - - python=3.12 - - pip diff --git a/mir/__init__.py b/mir/__init__.py index 66a9dac..f6e07b2 100644 --- a/mir/__init__.py +++ b/mir/__init__.py @@ -15,7 +15,7 @@ import os -__version__ = "3.3.0" +__version__ = "3.4.0" __all__ = ["__version__", "get_resource_path", "TCREmp", "PairedTCREmp"] diff --git a/mir/bench/__init__.py b/mir/bench/__init__.py index 9122cc7..e77d4b1 100644 --- a/mir/bench/__init__.py +++ b/mir/bench/__init__.py @@ -5,8 +5,10 @@ (reproduces the paper's Table 1 / S1). * :mod:`mir.bench.theory` — reproduce supplementary S1–S3 (distribution laws, dissimilarity↔distance correlation, real-vs-model prototype robustness). +* :mod:`mir.bench.eval` — scorers for the explainable readout (cross-validated AUC, + Cox C-index, log-rank) that :func:`mir.explain.channel_report` consumes. -Requires the ``[bench]`` extra (kneed, matplotlib, BioPython for the SW baseline). +Requires the ``[bench]`` extra (kneed, matplotlib, lifelines; BioPython for the SW baseline). """ from mir.bench.metrics import AntigenMetric, cluster, cluster_metrics, estimate_dbscan_eps diff --git a/mir/bench/eval.py b/mir/bench/eval.py new file mode 100644 index 0000000..8b4fedf --- /dev/null +++ b/mir/bench/eval.py @@ -0,0 +1,219 @@ +"""Scorers for the explainable cohort readout (:mod:`mir.explain`) — the closures you hand to +:func:`mir.explain.channel_report`. + +``mir.explain`` is deliberately scorer-agnostic: it slices channels and never sees ``y``, so the +*model choice* lives here, in the caller's closure. This module collects the recurring choices — +cross-validated classification AUC and Cox survival C-index — as small, reusable functions so every +cohort benchmark stops re-gluing them. A scorer takes a feature block and returns a float where +**higher is better**:: + + from mir.bench.eval import cv_auc, cv_cindex + from mir.explain import stack_embeddings, channel_report + + X, spec = stack_embeddings(embs) + rep = channel_report(X, spec, lambda B: cv_auc(B, y)[0], base=0.5) # classification + rep = channel_report(X, spec, lambda B: cv_cindex(dur, evt, base=C, block=B, n_pc=8), + base=cv_cindex(dur, evt, base=C, block=None), mode="both") # survival + +Needs the ``[bench]`` extra (scikit-learn always; lifelines for the survival scorers; vdjtools for +the k-mer baseline). Everything is lazily imported so importing this module stays cheap. +""" + +from __future__ import annotations + +import numpy as np + + +# --------------------------------------------------------------- classification + + +def held_out_auc(Xtr, ytr, Xte, yte, *, pca_cols: int = 0) -> float: + """Held-out ROC-AUC of a logistic head; the first ``pca_cols`` columns get in-fold PCA(0.9 var). + + Wide blocks (kernel mean, second moment, k-mer) overfit unless reduced *inside* the fold; the few + diversity/coverage features pass through raw so their signal isn't buried. ``pca_cols=0`` = no PCA. + """ + from sklearn.compose import ColumnTransformer + from sklearn.decomposition import PCA + from sklearn.linear_model import LogisticRegression + from sklearn.metrics import roc_auc_score + from sklearn.pipeline import make_pipeline + from sklearn.preprocessing import StandardScaler + + Xtr = Xtr.reshape(-1, 1) if Xtr.ndim == 1 else Xtr + Xte = Xte.reshape(-1, 1) if Xte.ndim == 1 else Xte + pca = make_pipeline(StandardScaler(), PCA(n_components=0.9, svd_solver="full", random_state=0)) + n = Xtr.shape[1] + if pca_cols and pca_cols < n: + pre = ColumnTransformer([("m", pca, list(range(pca_cols))), + ("r", StandardScaler(), list(range(pca_cols, n)))]) + elif pca_cols: + pre = pca + else: + pre = StandardScaler() + clf = make_pipeline(pre, LogisticRegression(max_iter=2000, C=1.0)).fit(Xtr, ytr) + return float(roc_auc_score(yte, clf.predict_proba(Xte)[:, 1])) + + +def cv_auc(X, y, *, pca_cols: int = 0, n_splits: int = 5, n_repeats: int = 10, seed: int = 0): + """Repeated stratified k-fold AUC as ``(mean, std)`` — a CI, not a single-split point estimate. + + A single 70/30 split at small ``n`` has AUC SD ≈ 0.5/√n_test (~0.1 for n_test≈30), so point + estimates are near-meaningless; repeated CV exposes whether two methods' intervals separate. + Pass ``cv_auc(B, y)[0]`` as the ``channel_report`` scorer. + """ + from sklearn.model_selection import RepeatedStratifiedKFold + + X = X.reshape(-1, 1) if X.ndim == 1 else X + rskf = RepeatedStratifiedKFold(n_splits=n_splits, n_repeats=n_repeats, random_state=seed) + aucs = [held_out_auc(X[tr], y[tr], X[te], y[te], pca_cols=pca_cols) for tr, te in rskf.split(X, y)] + return float(np.mean(aucs)), float(np.std(aucs)) + + +# --------------------------------------------------------------------- survival + + +def _design(base, block, rows): + """Assemble a pandas design frame from a base covariate matrix and a feature block.""" + import pandas as pd + + cols = {} + if base is not None: + base = np.asarray(base, dtype=np.float64) + base = base.reshape(-1, 1) if base.ndim == 1 else base + for j in range(base.shape[1]): + cols[f"c{j}"] = base[:, j] + if block is not None: + block = np.asarray(block, dtype=np.float64) + block = block.reshape(-1, 1) if block.ndim == 1 else block + for j in range(block.shape[1]): + cols[f"z{j}"] = block[:, j] + return pd.DataFrame(cols, index=np.arange(rows)) + + +def _nondegen(df): + """Drop zero-variance columns (excluding the survival targets) so Cox converges.""" + keep = [c for c in df.columns if c in ("_T", "_E") or df[c].std() > 1e-9] + return df[keep] + + +def cv_cindex(durations, events, *, base=None, block=None, n_pc: int = 0, + n_splits: int = 5, seed: int = 0, penalizer: float = 0.1) -> float: + """5-fold CV Cox C-index of ``base`` covariates + an optional feature ``block``. + + The survival analog of :func:`cv_auc`, and the scorer for survival channel reports. ``base`` is + the clinical design (e.g. age+sex+stage+log-reads — whatever the study built; this function is + schema-agnostic) and ``block`` is the feature block being scored; when ``block`` is wider than + ``n_pc`` it is PCA-reduced to ``n_pc`` components **inside each fold** (train-only fit) so a wide + kernel-mean block doesn't overfit. Returns the mean concordance across folds (nan-robust). + + Use as ``lambda B: cv_cindex(dur, evt, base=C, block=B, n_pc=8)`` with base score + ``cv_cindex(dur, evt, base=C, block=None)``. + """ + from lifelines import CoxPHFitter + from lifelines.utils import concordance_index + from sklearn.decomposition import PCA + from sklearn.model_selection import KFold + from sklearn.preprocessing import StandardScaler + + dur = np.asarray(durations, dtype=np.float64) + evt = np.asarray(events, dtype=np.float64) + n = dur.shape[0] + if base is None and block is None: + return float("nan") # a Cox with no covariates is undefined + + block = None if block is None else ( + np.asarray(block).reshape(-1, 1) if np.asarray(block).ndim == 1 else np.asarray(block)) + sc = [] + for tr, te in KFold(n_splits, shuffle=True, random_state=seed).split(np.arange(n)): + Btr = Bte = None + if block is not None: + if n_pc and block.shape[1] > n_pc: + st = StandardScaler().fit(block[tr]) + pca = PCA(min(n_pc, len(tr) - 1), random_state=seed).fit(st.transform(block[tr])) + Btr, Bte = pca.transform(st.transform(block[tr])), pca.transform(st.transform(block[te])) + else: + Btr, Bte = block[tr], block[te] + dtr = _design(None if base is None else np.asarray(base)[tr], Btr, len(tr)) + dte = _design(None if base is None else np.asarray(base)[te], Bte, len(te)) + dtr["_T"], dtr["_E"] = dur[tr], evt[tr] + dtr = _nondegen(dtr) + dte = dte[[c for c in dtr.columns if c not in ("_T", "_E")]] + try: + cph = CoxPHFitter(penalizer=penalizer).fit(dtr, "_T", "_E") + risk = cph.predict_partial_hazard(dte) + sc.append(concordance_index(dur[te], -risk, evt[te])) + except Exception: + sc.append(np.nan) + return float(np.nanmean(sc)) if np.isfinite(sc).any() else float("nan") + + +def km_logrank(durations, events, groups) -> float: + """Multivariate log-rank p-value across ``groups`` — do the KM survival curves differ? + + The test behind a Kaplan–Meier stratification (e.g. TME states from ``cluster_samples``). + Returns the p-value. + """ + from lifelines.statistics import multivariate_logrank_test + + return float(multivariate_logrank_test(durations, groups, events).p_value) + + +# --------------------------------------------------------------------- baseline + + +def kmer_matrix(frames, k: int = 3) -> np.ndarray: + """Sample × k-mer frequency matrix (vdjtools ``kmer_profile`` long-form → wide) — the classic baseline. + + Args: + frames: One clonotype :class:`polars.DataFrame` per sample (with ``junction_aa`` + counts). + k: k-mer length. + + Returns: + ``(len(frames), n_kmers)`` frequency matrix over the pooled vocabulary. + """ + from vdjtools.features import kmer_profile + + dicts = [] + for df in frames: + kp = kmer_profile(df, k=k, weight="freq", by_locus=False) + dicts.append(dict(zip(kp["kmer"].to_list(), kp["weight"].to_list()))) + vocab = sorted({x for d in dicts for x in d}) + idx = {x: i for i, x in enumerate(vocab)} + M = np.zeros((len(dicts), len(vocab))) + for i, d in enumerate(dicts): + for x, v in d.items(): + M[i, idx[x]] = v + return M + + +def _demo() -> None: + """Self-check on synthetic data: the scorers recover a planted class / survival signal.""" + rng = np.random.default_rng(0) + n = 240 + + # classification: a signal feature carries y, a noise feature does not + y = rng.integers(0, 2, n).astype(float) + signal = y + rng.normal(0, 0.6, n) + noise = rng.normal(0, 1, n) + auc_sig = cv_auc(signal, y, n_repeats=3)[0] + auc_noise = cv_auc(noise, y, n_repeats=3)[0] + assert auc_sig > 0.7 > auc_noise, (auc_sig, auc_noise) + + # survival: risk drives the hazard; C-index of base+risk beats base-only, and the log-rank splits + risk = rng.normal(0, 1, n) + base = rng.normal(0, 1, (n, 2)) # uninformative clinical covariates + dur = rng.exponential(np.exp(-0.9 * risk)) # higher risk -> shorter survival + evt = (rng.random(n) < 0.7).astype(float) + c_base = cv_cindex(dur, evt, base=base, block=None) + c_full = cv_cindex(dur, evt, base=base, block=risk) + assert c_full > c_base and c_full > 0.6, (c_base, c_full) + p = km_logrank(dur, evt, (risk > np.median(risk)).astype(int)) + assert p < 0.05, p + + print(f"mir.bench.eval OK; AUC signal {auc_sig:.2f} > noise {auc_noise:.2f}; " + f"C-index base {c_base:.2f} -> +risk {c_full:.2f}; log-rank p {p:.1e}") + + +if __name__ == "__main__": + _demo() diff --git a/mir/bench/metrics.py b/mir/bench/metrics.py index aceedc7..873dcde 100644 --- a/mir/bench/metrics.py +++ b/mir/bench/metrics.py @@ -70,6 +70,17 @@ def cluster(X: np.ndarray, eps: float | None = None, min_samples: int = 3, @dataclass class AntigenMetric: + """Per-antigen clustering quality: F1 of the recovered cluster plus retention. + + Attributes: + epitope: The antigen epitope this metric scores. + n: Number of clonotypes labelled with this epitope. + f1: F1 of the best-matching predicted cluster against the epitope's clonotypes. + precision: Precision of that best-matching cluster. + recall: Recall of that best-matching cluster. + retention: Fraction of the epitope's clonotypes that were clustered (not noise). + """ + epitope: str n: int f1: float diff --git a/mir/cli.py b/mir/cli.py new file mode 100644 index 0000000..45b38ea --- /dev/null +++ b/mir/cli.py @@ -0,0 +1,207 @@ +"""``mir`` command-line interface — turn receptor tables into embeddings. + +Two commands cover the two scales mirpy embeds at: + +* ``mir embed clonotypes SAMPLE`` — one repertoire's clonotype table → a per-clonotype + TCREMP embedding table (``e0…``), the input to clustering / ML. +* ``mir embed repertoires SAMPLE…`` — a *dataset* of clonotype tables → one repertoire + vector ``Φ(S)`` (``phi0…``) per sample, per chain, on one shared basis (so the rows are + mutually comparable / MMD-able). + +Inputs are any format ``vdjtools.io.read`` sniffs (AIRR TSV, vdjtools, MiXCR, immunoSEQ, +parquet, …). Output is TSV (default / ``.tsv``) or Parquet (``.parquet`` — recommended for +the wide raw embedding); ``-o -`` (or no ``-o``) writes TSV to stdout. + +Run ``mir embed clonotypes -h`` / ``mir embed repertoires -h`` for the full flag list. +""" +from __future__ import annotations + +import argparse +import sys + +import polars as pl + +import mir + + +# --- IO helpers ------------------------------------------------------------ +def _read(path: str) -> pl.DataFrame: + """Read a clonotype file into a normalized AIRR frame (any format vdjtools sniffs).""" + from vdjtools import io + + return io.read(path) + + +def _with_locus(df: pl.DataFrame) -> pl.DataFrame: + """Ensure a ``locus`` column (derive from ``v_call`` when absent).""" + if "locus" in df.columns and df["locus"].null_count() < df.height: + return df + from vdjtools import io + + try: + return io.add_locus(df) + except Exception: + # Fallback: the IMGT locus is the v_call's first 3 characters (TRB/TRA/IGH/…). + return df.with_columns(pl.col("v_call").str.slice(0, 3).alias("locus")) + + +def _emb_frame(X, prefix: str) -> pl.DataFrame: + """(N, d) float matrix → a polars frame with columns ``{prefix}0…{prefix}{d-1}``.""" + return pl.from_numpy(X, schema=[f"{prefix}{i}" for i in range(X.shape[1])]) + + +def _write(df: pl.DataFrame, path: str | None) -> None: + if path is None or path == "-": + sys.stdout.write(df.write_csv(separator="\t")) + elif path.endswith(".parquet"): + df.write_parquet(path) + else: + df.write_csv(path, separator="\t") + + +def _sample_id(path: str) -> str: + """Sample id = filename up to the first dot (``P1.TRB.tsv.gz`` → ``P1``).""" + import os + + return os.path.basename(path).split(".")[0] + + +def _pick_locus(df: pl.DataFrame, requested: str | None) -> str: + loci = [x for x in df["locus"].unique().to_list() if x] + if requested: + return requested + if len(loci) == 1: + return loci[0] + raise SystemExit( + f"multiple loci present ({', '.join(sorted(loci))}); pass --locus to pick one" + ) + + +# --- commands -------------------------------------------------------------- +def cmd_clonotypes(a: argparse.Namespace) -> None: + from mir.embedding.pca import pca_denoise + from mir.embedding.tcremp import TCREmp + + df = _with_locus(_read(a.input)) + locus = _pick_locus(df, a.locus) + sub = df.filter(pl.col("locus") == locus) + if sub.is_empty(): + raise SystemExit(f"no clonotypes for locus {locus!r} in {a.input}") + + model = TCREmp.from_defaults(a.species, locus, n_prototypes=a.n_prototypes, + mode=a.mode, threads=a.threads) + X = model.embed(sub) + if a.pca: + X = pca_denoise(X, n_components=a.pca) + + if (a.output is None or a.output == "-" or a.output.endswith(".tsv")) and X.shape[1] > 500: + print(f"[mir] {X.shape[1]} embedding columns — consider --pca K or a .parquet output.", + file=sys.stderr) + + id_cols = [c for c in ("junction_aa", "v_call", "j_call", "duplicate_count") if c in sub.columns] + out = sub.select(id_cols).hstack(_emb_frame(X, "e")) + _write(out, a.output) + print(f"[mir] embedded {X.shape[0]} {locus} clonotypes → {X.shape[1]}-d " + f"({'PCA ' if a.pca else ''}table)", file=sys.stderr) + + +def cmd_repertoires(a: argparse.Namespace) -> None: + from collections import defaultdict + + from mir.embedding.tcremp import TCREmp + from mir.repertoire import fit_repertoire_space, mmd_matrix, sample_embedding + + blocks = tuple(b.strip() for b in a.blocks.split(",") if b.strip()) + n_rff_second = a.n_rff_second if "second" in blocks else 0 + + # Load every sample, split its clonotypes by locus. + by_locus: dict[str, list] = defaultdict(list) + for path in a.input: + df = _with_locus(_read(path)) + sid = _sample_id(path) + for locus in [x for x in df["locus"].unique().to_list() if x]: + if a.locus and locus != a.locus: + continue + by_locus[locus].append((sid, df.filter(pl.col("locus") == locus))) + + if not by_locus: + raise SystemExit("no samples/loci to embed (check inputs / --locus)") + + rows: list[dict] = [] + vectors: list = [] + for locus in sorted(by_locus): + items = by_locus[locus] + model = TCREmp.from_defaults(a.species, locus, n_prototypes=a.n_prototypes, threads=a.threads) + pooled = pl.concat([sub for _, sub in items]) + space = fit_repertoire_space(model, pooled, n_rff=a.n_rff, n_rff_second=n_rff_second, + n_components=a.n_components, seed=a.seed) + embs = [sample_embedding(space, sub, weight=a.weight, blocks=blocks) for _, sub in items] + for (sid, sub), se in zip(items, embs): + rows.append({"sample_id": sid, "locus": locus, "n_clonotypes": sub.height}) + vectors.append(se.vector) + if a.mmd: + D = mmd_matrix(embs, unbiased=True) + ids = [sid for sid, _ in items] + mmd_df = pl.DataFrame({"sample_id": ids}).with_columns( + [pl.Series(ids[j], D[:, j]) for j in range(len(ids))]) + out = a.mmd if len(by_locus) == 1 else a.mmd.replace(".", f".{locus}.", 1) + _write(mmd_df, out) + print(f"[mir] {locus}: {len(items)} samples → Φ dim {len(embs[0].vector)}", file=sys.stderr) + + import numpy as np + + meta = pl.DataFrame(rows) + out = meta.hstack(_emb_frame(np.vstack(vectors), "phi")) + _write(out, a.output) + + +# --- parser ---------------------------------------------------------------- +def build_parser() -> argparse.ArgumentParser: + p = argparse.ArgumentParser(prog="mir", description=__doc__.splitlines()[0]) + p.add_argument("--version", action="version", version=f"mir {mir.__version__}") + sub = p.add_subparsers(dest="cmd", required=True) + + embed = sub.add_parser("embed", help="compute embeddings").add_subparsers(dest="what", required=True) + + c = embed.add_parser("clonotypes", help="repertoire → per-clonotype embedding table") + c.add_argument("input", help="a clonotype table (AIRR/vdjtools/MiXCR/parquet/…)") + c.add_argument("-o", "--output", help="output .tsv/.parquet (default: stdout TSV)") + c.add_argument("--species", default="human") + c.add_argument("--locus", help="chain to embed (inferred if the file has one locus)") + c.add_argument("--n-prototypes", type=int, default=None, + help="prototype count (default: per-chain preset)") + c.add_argument("--mode", default="vjcdr3", choices=("vjcdr3", "cdr123")) + c.add_argument("--pca", type=int, default=None, metavar="K", + help="PCA-denoise the embedding to K dims (compact table)") + c.add_argument("--threads", type=int, default=0, help="0 = all cores") + c.set_defaults(func=cmd_clonotypes) + + r = embed.add_parser("repertoires", help="dataset of clonotype tables → per-sample Φ(S), by chain") + r.add_argument("input", nargs="+", help="one clonotype file per repertoire (sample id = filename stem)") + r.add_argument("-o", "--output", help="output .tsv/.parquet (default: stdout TSV)") + r.add_argument("--species", default="human") + r.add_argument("--locus", help="restrict to one chain (default: all loci present, one basis each)") + r.add_argument("--n-prototypes", type=int, default=None) + r.add_argument("--weight", default="log1p", choices=("log1p", "anscombe", "distinct"), + help="clone-size weight g (frequencies w = g(a)/Σg)") + r.add_argument("--blocks", default="mean,diversity", + help="Φ blocks: mean,diversity[,second] (second = heavy HLA-interaction block)") + r.add_argument("--n-rff", type=int, default=1024, help="mean-block RFF dimension") + r.add_argument("--n-rff-second", type=int, default=128, help="second-moment RFF dimension (if used)") + r.add_argument("--n-components", type=int, default=None, + help="clonotype-PCA dims for the shared basis (default: preset)") + r.add_argument("--mmd", metavar="OUT", help="also write the per-chain pairwise unbiased-MMD matrix") + r.add_argument("--threads", type=int, default=0, help="0 = all cores") + r.add_argument("--seed", type=int, default=0) + r.set_defaults(func=cmd_repertoires) + + return p + + +def main(argv: list[str] | None = None) -> None: + args = build_parser().parse_args(argv) + args.func(args) + + +if __name__ == "__main__": + main() diff --git a/mir/cohort.py b/mir/cohort.py new file mode 100644 index 0000000..f3e5f65 --- /dev/null +++ b/mir/cohort.py @@ -0,0 +1,431 @@ +"""The digital donor: one fixed, explainable descriptor per subject, fused across chains (§T.7). + +Where :mod:`mir.repertoire` embeds one repertoire (one chain of one sample) and :mod:`mir.explain` +names the pieces of a feature matrix, this module assembles the **cohort** object the clinical +pipelines actually consume — a *digital donor* whose vector concatenates, per locus, the identity +(kernel mean), diversity (Hill) and coverage (log receptor load) channels into one row, ready for +survival / classification via :func:`mir.explain.channel_report`. + +The library owns the **geometry, fusion and serialization**; the study owns **which extra channels** +(isotype, composition, HLA, …) and **which scorer**. So :func:`fit_donor_embeddings` builds the +per-chain geometry and fuses it through :class:`mir.explain.ChannelBuilder`, and takes an +``extra_channels`` hook for the analysis to inject its own tissue/clinical blocks into the *same* +matrix — the channels then flow through ``channel_report`` indistinguishably. + +Comparability bites twice here (two bases, not one): each locus carries its own prototype hash **and** +a cross-sample identity-PCA rotation. :class:`DonorCohort` stores both; :meth:`DonorCohort.load` +verifies *every* locus's prototype hash, and :meth:`DonorCohort.transform` is the only way to project +held-out donors into the fitted basis. A batch-``residualize``d ``X`` is likewise incomparable to the +raw one — record the correction in provenance. + +Torch-free (numpy / sklearn / polars); vdjtools is lazy (only :func:`incidence_biomarkers`). +""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +import numpy as np +import polars as pl + +from mir.explain import ChannelBuilder, ChannelSpec +from mir.repertoire import RepertoireSpace, mmd_matrix, sample_embedding + +_COUNT = "duplicate_count" + + +# ------------------------------------------------------------- per-donor block assembly + + +def _locus_blocks(spaces, donor_frames, *, min_clones): + """Per (donor, locus): kernel mean, Hill diversity (4), log receptor load. ``nan`` where absent. + + Returns ``(means, div, cov)`` where ``means[locus] = [(row_idx, mean_vec), …]`` (only present + donors), and ``div``/``cov`` are ``(n, 4)`` / ``(n, 1)`` with ``nan`` rows for donors lacking the + chain (later imputed by :class:`~mir.explain.ChannelBuilder`). + """ + n = len(donor_frames) + loci = list(spaces) + means = {c: [] for c in loci} + div = {c: np.full((n, 4), np.nan) for c in loci} + cov = {c: np.full((n, 1), np.nan) for c in loci} + for i, frames in enumerate(donor_frames): + for c in loci: + f = frames.get(c) + if f is not None and f.height >= min_clones: + e = sample_embedding(spaces[c], f, blocks=("mean", "diversity")) + means[c].append((i, e.mean)) + div[c][i] = e.diversity + cov[c][i, 0] = np.log1p(float(f[_COUNT].sum())) + return means, div, cov + + +def _identity_matrix(means_c, n, id_pca, *, pca=None, seed=0): + """``(n, id_pca)`` identity block: PCA-reduce the per-donor kernel means (``nan`` where absent). + + Fits a ``StandardScaler`` → ``PCA`` when ``pca is None`` (and enough donors carry the chain); + otherwise reuses the supplied fitted reducer (held-out :meth:`DonorCohort.transform`). + Returns ``(matrix, fitted_pca_or_None)``. + """ + from sklearn.decomposition import PCA + from sklearn.pipeline import make_pipeline + from sklearn.preprocessing import StandardScaler + + out = np.full((n, id_pca), np.nan) + if not means_c: + return out, pca + idxs, M = zip(*means_c) + M = np.stack(M) + if pca is None: + if M.shape[0] <= id_pca: + return out, None # too few to reduce; leave holes (imputed) + pca = make_pipeline(StandardScaler(), PCA(id_pca, random_state=seed)).fit(M) + Mr = pca.transform(M) + for j, i in enumerate(idxs): + out[i] = Mr[j] + return out, pca + + +# --------------------------------------------------------------------- DonorCohort + + +@dataclass +class DonorCohort: + """A fitted cohort of digital donors: the fused matrix plus everything needed to reproduce it. + + Attributes: + X: ``(n_donors, width)`` fused, imputed, z-scored feature matrix. + spec: :class:`~mir.explain.ChannelSpec` — the channel → column map for ``X``. + spaces: ``{locus: RepertoireSpace}`` — the per-locus clonotype bases (hash-verified on load). + identity_pca: ``{locus: fitted StandardScaler→PCA}`` — the cross-sample identity reducer, + stored (not thrown away) so held-out donors project comparably. + rows: Per-donor metadata, row-aligned to ``X`` (``None`` if not supplied). + meta: Provenance — per-locus prototype hashes, ``id_pca``, block order, standardization stats. + """ + + X: np.ndarray + spec: ChannelSpec + spaces: dict + identity_pca: dict + rows: list | None + meta: dict = field(default_factory=dict) + + def transform(self, donor_frames: list[dict], *, extra: dict | None = None) -> np.ndarray: + """Project held-out donors into this fitted basis (the only comparable path for new donors). + + Rebuilds the per-locus identity/diversity/coverage blocks through the *stored* spaces and + identity PCAs, re-adds any ``extra`` channels the caller supplies (same names/widths as at + fit — the study owns those), then applies the *fit-cohort* impute medians and z-scores. + + Args: + donor_frames: Per held-out donor, ``{locus: chain clonotype frame}``. + extra: ``{name: (m, k) array}`` for the non-core channels used at fit (isotype/…). Required + iff the cohort was fit with ``extra_channels``; the widths must match. + + Returns: + ``(len(donor_frames), width)`` matrix in the fitted basis. + """ + n = len(donor_frames) + means, div, cov = _locus_blocks(self.spaces, donor_frames, min_clones=self.meta["min_clones"]) + blocks: dict[str, np.ndarray] = {} + for c in self.spaces: + ident, _ = _identity_matrix(means[c], n, self.meta["id_pca"], pca=self.identity_pca.get(c)) + blocks[f"identity:{c}"] = ident + blocks[f"diversity:{c}"] = div[c] + if self.meta["coverage"]: + blocks[f"coverage:{c}"] = cov[c] + extra = extra or {} + need = set(self.meta["extra_names"]) + if set(extra) != need: + raise ValueError(f"transform needs extra channels {sorted(need)}, got {sorted(extra)}") + cols = [] + for name, width in self.meta["order"]: + mat = blocks.get(name, extra.get(name)) + mat = np.asarray(mat, dtype=np.float64) + mat = mat.reshape(-1, 1) if mat.ndim == 1 else mat + if mat.shape != (n, width): + raise ValueError(f"block {name!r} has shape {mat.shape}, expected {(n, width)}") + cols.append(mat) + raw = np.hstack(cols) + med, mu, sd = self.meta["median"], self.meta["mu"], self.meta["sd"] + for j in range(raw.shape[1]): + bad = ~np.isfinite(raw[:, j]) + if bad.any(): + raw[bad, j] = med[j] + return (raw - mu) / sd + + def save(self, path) -> None: + """Pickle the cohort: ``X``/``spec``/``rows``/``meta``/identity PCAs + each locus's basis.""" + import pickle + + spaces_ser = {c: {"meta": sp.meta, "space": sp.clono.space, "scaler": sp.clono.scaler, + "pca": sp.clono.pca, "rff": sp.rff, "rff2": sp.rff2} + for c, sp in self.spaces.items()} + with open(path, "wb") as fh: + pickle.dump({"X": self.X, "spec": self.spec, "rows": self.rows, "meta": self.meta, + "identity_pca": self.identity_pca, "spaces": spaces_ser}, fh) + + @classmethod + def load(cls, path, *, verify: bool = True) -> "DonorCohort": + """Load a cohort, rebuilding every locus basis and verifying **all** prototype hashes.""" + import pickle + + from mir.density import DensitySpace + from mir.embedding.tcremp import TCREmp + from mir.ml.bundle import prototype_hash + + with open(path, "rb") as fh: + d = pickle.load(fh) + spaces = {} + for c, s in d["spaces"].items(): + m = s["meta"] + if verify: + cur = prototype_hash(m["species"], m["locus"], m["n_prototypes"]) + if cur != m["prototype_hash"]: + raise ValueError( + f"prototype hash mismatch for locus {c} ({m['species']}_{m['locus']}): this " + "cohort was fit on a different prototype set — its embeddings are NOT " + "comparable to the current prototypes. Pass verify=False to override." + ) + model = TCREmp.from_defaults(m["species"], m["locus"], m["n_prototypes"], mode=m["mode"], + metric=m["metric"], gap_positions=tuple(m["gap_positions"])) + clono = DensitySpace(model=model, space=s["space"], scaler=s["scaler"], pca=s["pca"]) + spaces[c] = RepertoireSpace(clono, s["rff"], s["rff2"], m) + return cls(d["X"], d["spec"], spaces, d["identity_pca"], d["rows"], d["meta"]) + + +def fit_donor_embeddings( + spaces: dict, + donor_frames: list[dict], + *, + rows: list | None = None, + id_pca: int = 8, + min_clones: int = 5, + coverage: bool = True, + extra_channels=None, + standardize: bool = True, + impute: bool = True, + seed: int = 0, +) -> DonorCohort: + """Fuse per-chain repertoire embeddings into one digital-donor matrix (§T.7). + + For each donor and each fitted locus, computes the kernel-mean identity (cross-sample PCA-reduced), + Hill diversity and log receptor-load coverage; per-chain blocks of the same kind merge under one + channel name. Missing chains leave holes that :class:`~mir.explain.ChannelBuilder` imputes. + + Args: + spaces: ``{locus: RepertoireSpace}`` from :func:`mir.repertoire.fit_repertoire_spaces`. + donor_frames: Per donor, ``{locus: chain clonotype frame}`` (already canonical-filtered / + grouped by the study; the library only embeds). Row-aligned to ``rows``. + rows: Optional per-donor metadata, carried onto the cohort. + id_pca: Cross-sample identity-PCA dimensionality per locus. + min_clones: Skip a chain for a donor with fewer clonotypes (its blocks become holes). + coverage: Include the per-chain log-receptor-load channel. + extra_channels: Optional ``(rows, identity_concat) -> {name: (n, k) array}`` — the study's own + blocks (isotype / composition / atypicality …), fused into the *same* matrix. All are + treated as non-attributable. ``identity_concat`` is the ``(n, Σ id_pca)`` per-chain + identity, so e.g. :func:`mir.repertoire.centroid_atypicality` can be computed against it. + standardize / impute: Passed through to the final assembly (z-score / median-impute). + seed: RNG seed for the identity PCA. + + Returns: + A :class:`DonorCohort`; ``.X`` / ``.spec`` feed :func:`mir.explain.channel_report`. + """ + if not spaces: + raise ValueError("no fitted spaces; fit_repertoire_spaces returned nothing to embed") + n = len(donor_frames) + means, div, cov = _locus_blocks(spaces, donor_frames, min_clones=min_clones) + + identity_pca: dict = {} + ident: dict = {} + for c in spaces: + ident[c], pca = _identity_matrix(means[c], n, id_pca, pca=None, seed=seed) + if pca is not None: + identity_pca[c] = pca + identity_concat = np.hstack([ident[c] for c in spaces]) if spaces else np.zeros((n, 0)) + + # ordered blocks: per-locus identity(attributable) / diversity / coverage, then study extras + ordered = [] + for c in spaces: + ordered.append((f"identity:{c}", ident[c], True)) + ordered.append((f"diversity:{c}", div[c], False)) + if coverage: + ordered.append((f"coverage:{c}", cov[c], False)) + extra_names = [] + if extra_channels is not None: + for name, mat in extra_channels(rows, identity_concat).items(): + mat = np.asarray(mat, dtype=np.float64) + mat = mat.reshape(-1, 1) if mat.ndim == 1 else mat + ordered.append((name, mat, False)) + extra_names.append(name) + + # merge per-chain blocks under their bare channel name (identity:TRB -> "identity"), as build_embedding does + b = ChannelBuilder() + for name, mat, attr in ordered: + b.add(name.split(":", 1)[0], mat, attributable=attr) + raw, spec = b.build(standardize=False, impute=False) + + median = np.zeros(raw.shape[1]) + for j in range(raw.shape[1]): + good = raw[np.isfinite(raw[:, j]), j] + median[j] = float(np.median(good)) if good.size else 0.0 + if impute: + bad = ~np.isfinite(raw[:, j]) + raw[bad, j] = median[j] + mu = raw.mean(axis=0) if standardize else np.zeros(raw.shape[1]) + sd = raw.std(axis=0) if standardize else np.ones(raw.shape[1]) + sd[sd == 0] = 1.0 + X = (raw - mu) / sd + + meta = { + "prototype_hashes": {c: sp.meta["prototype_hash"] for c, sp in spaces.items()}, + "id_pca": id_pca, "min_clones": min_clones, "coverage": coverage, + "order": [(name, mat.shape[1]) for name, mat, _ in ordered], + "extra_names": extra_names, + "median": median, "mu": mu, "sd": sd, "standardize": standardize, "impute": impute, + } + return DonorCohort(X=X, spec=spec, spaces=spaces, identity_pca=identity_pca, rows=rows, meta=meta) + + +# --------------------------------------------------------------- cohort operations + + +def residualize(X: np.ndarray, group: np.ndarray) -> np.ndarray: + """Subtract each group's mean vector — remove the first-order batch offset (Prop. ``prop:batch``). + + The detect→correct→verify batch cookbook's correction step, applied to a stacked donor matrix. + NB the corrected matrix is a *different* coordinate system: never compare a residualized ``X`` to + an uncorrected one. + """ + X = np.asarray(X, dtype=np.float64) + out = X.copy() + for g in np.unique(group): + m = group == g + out[m] -= X[m].mean(axis=0) + return out + + +def cluster_samples(embs, *, unbiased: bool = True, method: str = "dbscan", + eps: float | None = None, min_samples: int = 3, k: int = 4, **kwargs): + """Cluster repertoires by MMD — the repertoire-level analog of :func:`mir.bench.metrics.cluster` (TME states etc.). + + Builds the MMD distance matrix (unbiased by default, since cohorts differ in depth) and clusters + it with a precomputed-metric density estimator (:func:`mir.bench.metrics.cluster`). + + Args: + embs: Per-sample :class:`~mir.repertoire.SampleEmbedding` (mean block present). + unbiased: Use the diagonal-removed MMD² (recommended for unequal-depth cohorts). + method: ``"dbscan"`` / ``"hdbscan"`` / ``"optics"``. + eps: DBSCAN radius; ``None`` → median of each sample's ``k``-th nearest MMD (self excluded). + min_samples / k: density / eps-estimation neighbours. + **kwargs: forwarded to the estimator. + + Returns: + Cluster labels (``-1`` = noise), length ``len(embs)``. + """ + from mir.bench.metrics import cluster + + D = mmd_matrix(embs, unbiased=unbiased) + if method == "dbscan" and eps is None: + kth = np.sort(D, axis=1)[:, min(k, D.shape[0] - 1)] # k-th nearest (row 0 is self) + eps = float(np.median(kth)) or 1.0 + return cluster(D, eps=eps, min_samples=min_samples, method=method, metric="precomputed", **kwargs) + + +def incidence_biomarkers(cohort, phenotype, *, pheno_col: str, match: str = "1mm", + min_incidence: int = 3, **kwargs): + """Cohort biomarker detection: per-clonotype **subject-incidence Fisher** test (Emerson 2017). + + A thin delegate to :func:`vdjtools.biomarker.fisher.fisher_association` — the presence/absence + biomarker that complements the geometry witness (:func:`mir.repertoire.class_witness`) and, at + realistic donor counts, recovers public motifs the witness misses. ``match="1mm"`` groups + single-mismatch metaclonotypes (the paper's method; exact-match typically returns ~0 hits). + + Args: + cohort: Pooled clonotype frame with a sample-id column (vdjtools cohort form). + phenotype: Per-sample frame carrying the binary ``pheno_col`` label. + pheno_col: Phenotype column name in ``phenotype``. + match: ``"1mm"`` (metaclonotypes) or ``"exact"``. + min_incidence: Minimum subject incidence to test a feature. + **kwargs: forwarded to ``fisher_association`` (``key``, ``scope``, ``alternative``, …). + + Returns: + One row per feature: incidence, odds-ratio, ``p_value``, BH ``q_value``, direction. + """ + from vdjtools.biomarker.fisher import fisher_association + + return fisher_association(cohort, phenotype, pheno_col=pheno_col, match=match, + min_incidence=min_incidence, **kwargs) + + +# --------------------------------------------------------------------- self-check + + +def _demo() -> None: + """Self-check on bundled prototypes: a multi-chain cohort fuses, serializes, and a planted + prognostic channel wins the report.""" + from mir.embedding.prototypes import list_available_prototypes + from mir.embedding.tcremp import TCREmp + from mir.repertoire import centroid_atypicality, fit_repertoire_spaces + + avail = [loc for sp, loc in list_available_prototypes() if sp == "human"] + loci = [c for c in ("TRB", "TRA", "IGH") if c in avail][:2] or [avail[0]] + models = {c: TCREmp.from_defaults("human", c, n_prototypes=300) for c in loci} + protos = {c: pl.DataFrame({"v_call": m._proto_v, "j_call": m._proto_j, + "junction_aa": m._proto_junction}).unique() for c, m in models.items()} + spaces = fit_repertoire_spaces(models, protos, n_rff=512, n_rff_second=0, n_components=15, seed=0) + + # 40 donors in two "risk" groups: group 1 gets a public expansion in the identity space + donor_frames, rows = [], [] + for i in range(40): + grp = i % 2 + frames = {} + for c in spaces: + base = protos[c].sample(120, seed=i, with_replacement=True).with_columns( + pl.lit(1.0).alias(_COUNT)) + if grp == 1: + spike = protos[c].slice(0, 5).with_columns(pl.lit(400.0).alias(_COUNT)) + base = pl.concat([base, spike]) + frames[c] = base + donor_frames.append(frames) + rows.append({"group": grp}) + + def extra(rows, identity): + aty = centroid_atypicality(identity, np.array([r["group"] for r in rows])) + return {"atypicality": aty} + + coh = fit_donor_embeddings(spaces, donor_frames, rows=rows, id_pca=6, extra_channels=extra, seed=0) + assert coh.X.shape[0] == 40 + assert "identity" in coh.spec and "atypicality" in coh.spec + assert "identity" in coh.spec.attributable + + # the identity channel separates the two groups; a channel_report with a CV-AUC scorer ranks it top + from mir.bench.eval import cv_auc + from mir.explain import channel_report + + y = np.array([r["group"] for r in coh.rows], dtype=float) + rep = channel_report(coh.X, coh.spec, lambda B: cv_auc(B, y, n_repeats=3, pca_cols=B.shape[1])[0], + base=0.5, mode="in") + assert rep.best == "identity", rep.frame() + + # transform held-out donors + serialize with hash verification + import tempfile, os + held = donor_frames[:4] + Xt = coh.transform(held, extra={"atypicality": np.zeros(4)}) + assert Xt.shape == (4, coh.X.shape[1]) + with tempfile.TemporaryDirectory() as td: + p = os.path.join(td, "cohort.pkl") + coh.save(p) + back = DonorCohort.load(p) + assert np.allclose(back.X, coh.X) + assert back.spec.names == coh.spec.names + + # cluster_samples on the two-group embeddings + embs = [sample_embedding(spaces[loci[0]], f[loci[0]], blocks=("mean",)) for f in donor_frames] + labels = cluster_samples(embs) + print(f"mir.cohort OK; donors {coh.X.shape[0]} x {coh.X.shape[1]}, channels {coh.spec.names}; " + f"best={rep.best}; clusters {len(set(labels)) - (1 if -1 in labels else 0)}") + + +if __name__ == "__main__": + _demo() diff --git a/mir/density.py b/mir/density.py index adca50f..918ce30 100644 --- a/mir/density.py +++ b/mir/density.py @@ -69,6 +69,9 @@ _REQUIRED_COLS = ("v_call", "j_call", "junction_aa") _AA = "ACDEFGHIKLMNPQRSTVWY" +_AUTO_CHUNK_ROWS = 500_000 +"""Above this pooled row count, :func:`fit_density_space` auto-routes to the chunked path: the +single-shot path would materialize a >10 GB raw matrix (and ~2× again on the float64 upcast).""" def _slice(emb: np.ndarray, space: str) -> np.ndarray: @@ -210,6 +213,14 @@ def fit_density_space( ``(density_space, obs_emb, bg_emb)`` — the fitted :class:`DensitySpace` and the two reduced ``float`` arrays, row-aligned to ``obs_df`` / ``bg_df``. """ + if chunk_size is None and obs_df.height + bg_df.height > _AUTO_CHUNK_ROWS: + import warnings + + chunk_size = 200_000 # the single-shot path would materialize a >10 GB raw matrix + warnings.warn( + f"pooled rows ({obs_df.height + bg_df.height}) exceed {_AUTO_CHUNK_ROWS}: auto-enabling " + f"the chunked fit (chunk_size={chunk_size}) to bound peak memory; the PCA is fit on a " + "200k-row sample. Pass chunk_size explicitly to control this.", stacklevel=2) if chunk_size is not None: return _fit_density_space_chunked( model, obs_df, bg_df, n_components=n_components, space=space, seed=seed, @@ -321,7 +332,15 @@ def calibrate_radius( pl.Series("junction_aa", [_mutate1(s, rng) for s in junc]) ) drift = np.linalg.norm(space.transform(sample_df) - space.transform(mutated_df), axis=1) - return float(np.quantile(drift, quantile)) + r = float(np.quantile(drift, quantile)) + if not (r > 0): + raise ValueError( + "calibrate_radius produced a non-positive radius: the calibration junctions did not " + "move under a single substitution (too short, or all identical), so every enrichment " + "ball would be empty. Supply a sample_df with longer/real junctions, or pass an " + "explicit radius to neighbor_enrichment." + ) + return r def _ann_neighbors(obs, bg, radius, lambda0, n_ref, *, k_max: int = 96, seed: int = 0): @@ -333,7 +352,7 @@ def _ann_neighbors(obs, bg, radius, lambda0, n_ref, *, k_max: int = 96, seed: in kNN graph (k=``k_max``) thresholded by radius — *approximate*: recall < 1 undercounts, biasing enrichment **down** (conservative). A saturated ball (all ``k_max`` neighbours inside ``rad``) is undercounted and warned. For large N where exact trees are slow; use ``backend='exact'`` for - small or reproducibility-critical runs. Needs ``pynndescent`` (the ``[bench]`` extra). + small or reproducibility-critical runs. Needs ``pynndescent`` (the ``[ann]`` extra). """ from pynndescent import NNDescent @@ -383,7 +402,7 @@ def neighbor_enrichment( abundance: np.ndarray | None = None, weight: str = "log1p", orphan: bool = True, - backend: str = "exact", + backend: str = "kdtree", ) -> EnrichmentResult: """Continuous neighbour-enrichment test in one embedding coordinate system. @@ -431,14 +450,14 @@ def neighbor_enrichment( weight: Concave size transform ``g`` — ``"log1p"`` (default), ``"anscombe"``, or ``"distinct"`` (``g≡1``, ignore sizes even if ``abundance`` is given). orphan: When abundance-aware, combine the size p-value with the breadth p-value by Fisher. - backend: neighbour engine. ``"exact"`` (default, BallTree — exact, reproducible), - ``"kdtree"`` (scipy cKDTree — **also exact**, multithreaded, 5–9× faster; the - recommended speedup), or ``"ann"`` (approximate pynndescent kNN for whole-repertoire - scale — ~30× faster than BallTree at ≥40k but recall < 1 undercounts, biasing enrichment - conservatively; see :func:`_ann_neighbors`). ``"kdtree"`` is bit-identical to - ``"exact"`` at a *fixed* ``radius``; in balloon mode counts differ by at most ±1 at the - ball boundary (a float-epsilon difference in the computed k-th-neighbour distance), - which is negligible — prefer ``"kdtree"`` unless reproducing the BallTree baseline exactly. + backend: neighbour engine. ``"kdtree"`` (**default**, scipy cKDTree — exact, multithreaded, + 5–9× faster than BallTree), ``"exact"`` (BallTree — exact, single-threaded; the historical + default, kept for bit-reproducing the BallTree baseline), or ``"ann"`` (approximate + pynndescent kNN for whole-repertoire scale — ~30× faster at ≥40k but recall < 1 undercounts, + biasing enrichment conservatively; see :func:`_ann_neighbors`). ``"kdtree"`` is bit-identical + to ``"exact"`` at a *fixed* ``radius``; in balloon mode counts differ by at most ±1 at the + ball boundary (float-epsilon in the computed k-th-neighbour distance), which is negligible — + pass ``backend="exact"`` only to reproduce the BallTree baseline exactly. Returns: An :class:`EnrichmentResult`; ``fold`` is the (calibrated) density ratio ``E(z)``. In @@ -449,6 +468,16 @@ def neighbor_enrichment( bg = np.ascontiguousarray(bg_emb, dtype=np.float64) n_obs_total, n_bg_total = obs.shape[0], bg.shape[0] n_ref = max(n_obs_total - 1, 1) + if n_obs_total == 0 or n_bg_total == 0: + raise ValueError( + f"neighbor_enrichment requires non-empty obs and bg embeddings " + f"(got N_obs={n_obs_total}, N_bg={n_bg_total})." + ) + if pseudocount <= 0: + raise ValueError( + f"pseudocount must be > 0 (got {pseudocount}); it stabilizes p_bg so a clonotype in a " + "zero-background ball cannot produce an infinite fold." + ) if backend == "exact": obs_tree, bg_tree = BallTree(obs), BallTree(bg) @@ -493,7 +522,7 @@ def _lists_obs(): rad, radius_out, n_bg, _count_obs, _lists_obs = _ann_neighbors( obs, bg, radius, lambda0, n_ref) else: - raise ValueError(f"backend must be 'exact' or 'ann', got {backend!r}") + raise ValueError(f"backend must be 'exact', 'kdtree', or 'ann', got {backend!r}") p_bg = (n_bg + pseudocount) / (n_bg_total + pseudocount) expected = n_ref * p_bg diff --git a/mir/embedding/prototypes.py b/mir/embedding/prototypes.py index 647f827..816a2cc 100644 --- a/mir/embedding/prototypes.py +++ b/mir/embedding/prototypes.py @@ -78,7 +78,7 @@ def load_prototypes(species: str, locus: str, n: int | None = None) -> pl.DataFr generation order. Raises: - ValueError: If ``n > N_PROTOTYPES``. + ValueError: If ``n`` is not in ``[1, N_PROTOTYPES]``. FileNotFoundError: If no prototype file exists for the species/locus. Example: @@ -88,10 +88,11 @@ def load_prototypes(species: str, locus: str, n: int | None = None) -> pl.DataFr >>> len(df) 100 """ - if n is not None and n > N_PROTOTYPES: + if n is not None and (n <= 0 or n > N_PROTOTYPES): raise ValueError( - f"n={n} exceeds the maximum number of prototypes ({N_PROTOTYPES}); " - f"use n <= {N_PROTOTYPES} or None to load all." + f"n={n} must be in [1, {N_PROTOTYPES}] or None to load all. A non-positive n would " + f"silently change the prototype set — and thus its hash — via df.head(n), breaking " + f"embedding comparability." ) species_c = normalize_species_alias(species) diff --git a/mir/embedding/tcremp.py b/mir/embedding/tcremp.py index e50b9af..f01536b 100644 --- a/mir/embedding/tcremp.py +++ b/mir/embedding/tcremp.py @@ -49,11 +49,10 @@ from mir.aliases import normalize_locus_alias, normalize_species_alias from mir.distances.germline import GermlineDistances, load_germline_distances -from mir.distances.junction import junction_distance_matrix +from mir.distances.junction import DEFAULT_GAP_POSITIONS, junction_distance_matrix from mir.embedding.prototypes import load_prototypes MODES = ("vjcdr3", "cdr123") -DEFAULT_GAP_POSITIONS: tuple[int, ...] = (3, 4, -4, -3) # mode -> ((component, query_gene_column, prototype_gene_attr), ...) for the two # germline slots (0, 1). The junction is always slot 2. diff --git a/mir/ml/bundle.py b/mir/ml/bundle.py index 59e17b9..ab27ce5 100644 --- a/mir/ml/bundle.py +++ b/mir/ml/bundle.py @@ -20,6 +20,7 @@ from pathlib import Path import mir +from mir.distances.junction import DEFAULT_GAP_POSITIONS from mir.embedding.prototypes import load_prototypes @@ -48,7 +49,7 @@ def from_forward( species: str, locus: str, n_prototypes: int, - gap_positions: tuple[int, ...] = (3, 4, -4, -3), + gap_positions: tuple[int, ...] = DEFAULT_GAP_POSITIONS, kind: str = "forward", ) -> "CodecBundle": meta = { diff --git a/mir/repertoire.py b/mir/repertoire.py index 0ebee7a..225026d 100644 --- a/mir/repertoire.py +++ b/mir/repertoire.py @@ -51,6 +51,22 @@ _COUNT = "duplicate_count" +def _sample_weights(sample_df: pl.DataFrame, weight: str): + """Validated ``(counts a, concave g, normalized weights w=g/Σg)`` for one sample. + + Raises on the degenerate inputs that would otherwise propagate a silent ``NaN``/``inf`` through + every block: an empty repertoire (no clonotypes) or all-zero ``duplicate_count`` (``g.sum()==0`` + ⇒ ``w`` all-``NaN``, ``n_eff=inf``). + """ + if sample_df.height == 0: + raise ValueError("empty repertoire: sample_df has no clonotypes") + a = sample_df[_COUNT].to_numpy().astype(np.float64) + if not (a.sum() > 0): + raise ValueError("degenerate repertoire: duplicate_count must sum to a positive value") + g = _WEIGHTS[weight](a) + return a, g, g / g.sum() + + # --------------------------------------------------------------------------- RFF @@ -112,9 +128,8 @@ def sample_cloud(self, df: pl.DataFrame, *, weight: str = "log1p"): The raw material for both the fixed kernel mean (:func:`sample_embedding`) and the learned set network (:mod:`mir.ml.set_encoder`) — same basis, so their embeddings are comparable. """ - a = df[_COUNT].to_numpy().astype(np.float64) - g = _WEIGHTS[weight](a) - return self.transform_clonotypes(df), g / g.sum() + _, _, w = _sample_weights(df, weight) + return self.transform_clonotypes(df), w def save(self, path) -> None: """Pickle the basis (scaler + PCA + RFF + meta); the model is reconstructed on load.""" @@ -198,14 +213,17 @@ def fit_repertoire_space( if n_components is None: n_components = get_preset(model.species, model.locus).n_components - X = _embed(model, cohort_df, space) - if pca_fit_cap is not None and X.shape[0] > pca_fit_cap: + # Fit the scaler+PCA on at most pca_fit_cap pooled rows. Sample the FRAME rows and embed only + # those — bit-identical to embedding all then subsetting (embedding is row-wise, same seed picks + # the same indices) but never materializes the full raw matrix for a huge cohort. + if pca_fit_cap is not None and cohort_df.height > pca_fit_cap: rng = np.random.default_rng(seed) - X_fit = X[rng.choice(X.shape[0], pca_fit_cap, replace=False)] + rows = rng.choice(cohort_df.height, pca_fit_cap, replace=False) + X_fit = _embed(model, cohort_df[rows], space) else: - X_fit = X + X_fit = _embed(model, cohort_df, space) scaler = StandardScaler().fit(X_fit) - k = min(n_components, X_fit.shape[0], X.shape[1]) + k = min(n_components, X_fit.shape[0], X_fit.shape[1]) pca = PCA(n_components=k, random_state=seed).fit(scaler.transform(X_fit)) clono = DensitySpace(model=model, space=space, scaler=scaler, pca=pca) @@ -224,6 +242,38 @@ def fit_repertoire_space( return RepertoireSpace(clono, rff, rff2, meta) +def fit_repertoire_spaces( + models: dict, + cohort_frames: dict, + *, + min_clonotypes: int = 50, + **kwargs, +) -> dict: + """Fit one :class:`RepertoireSpace` **per locus** — the multi-chain (digital-donor) basis. + + Each locus gets an independent basis (its own prototype set + PCA + RFF), so per-chain kernel + means are only ever compared within their own chain, honouring the comparability contract. + Loci whose pooled cloud is too small to fit a stable PCA are skipped (returned dict omits them). + + Args: + models: ``{locus: TCREmp}`` — one fitted single-chain model per locus. + cohort_frames: ``{locus: pooled clonotype frame}`` (``v_call``/``j_call``/``junction_aa``), + the pooled cloud for that locus across the whole cohort. + min_clonotypes: Skip a locus whose pooled frame has fewer rows (PCA would be degenerate). + **kwargs: Forwarded to :func:`fit_repertoire_space` (``n_rff``, ``n_components``, ``seed`` …). + + Returns: + ``{locus: RepertoireSpace}`` for the loci that could be fit. + """ + spaces: dict = {} + for loc, model in models.items(): + pool = cohort_frames.get(loc) + if pool is None or pool.height < min_clonotypes: + continue + spaces[loc] = fit_repertoire_space(model, pool, **kwargs) + return spaces + + # --------------------------------------------------------------------- embedding @@ -295,9 +345,7 @@ def sample_embedding( Returns: A :class:`SampleEmbedding`; ``.vector`` is the concatenated fixed-width tensor. """ - a = sample_df[_COUNT].to_numpy().astype(np.float64) - g = _WEIGHTS[weight](a) - w = g / g.sum() + a, g, w = _sample_weights(sample_df, weight) n_eff = float(1.0 / np.sum(w * w)) mean = div = sec = None @@ -373,9 +421,7 @@ def sample_descriptor(space: RepertoireSpace, sample_df: pl.DataFrame, *, diversity and clonality are all smooth coordinates of the *same* object — the representation the in-silico-evolution / embedding-simulation workflow perturbs. """ - a = sample_df[_COUNT].to_numpy().astype(np.float64) - g = _WEIGHTS[weight](a) - w = g / g.sum() + a, g, w = _sample_weights(sample_df, weight) mean = w @ space.rff.transform(space.transform_clonotypes(sample_df)) sw2 = float(np.sum(w * w)) return RepertoireDescriptor(log_mass=float(np.log1p(a.sum())), @@ -405,6 +451,11 @@ def mmd_distance(a: SampleEmbedding, b: SampleEmbedding, *, unbiased: bool = Fal """ if not unbiased: return float(np.linalg.norm(a.mean - b.mean)) + if a.n_eff <= 1.0 or b.n_eff <= 1.0: + raise ValueError( + "unbiased MMD is undefined for a single-clonotype sample (n_eff ≤ 1): the diagonal " + "cannot be removed from a point mass. Use unbiased=False, or drop degenerate samples." + ) sa, sb = 1.0 / a.n_eff, 1.0 / b.n_eff # Σwᵢ² ; RFF self-similarity k(z,z)≈1 haa = (float(a.mean @ a.mean) - sa) / (1.0 - sa) # diagonal-removed ‖μ‖² hbb = (float(b.mean @ b.mean) - sb) / (1.0 - sb) @@ -422,6 +473,12 @@ def mmd_matrix(embs: list[SampleEmbedding], *, unbiased: bool = False) -> np.nda sq = np.diag(G).copy() if unbiased: s = np.array([1.0 / e.n_eff for e in embs]) # per-sample Σwᵢ² + if np.any(s >= 1.0): + raise ValueError( + "unbiased MMD is undefined for single-clonotype samples (n_eff ≤ 1); " + f"{int(np.sum(s >= 1.0))} of {len(embs)} samples are degenerate. " + "Use unbiased=False, or drop them before building the matrix." + ) h = (sq - s) / (1.0 - s) # diagonal-removed self-inner-products d2 = h[:, None] + h[None, :] - 2.0 * G else: @@ -439,6 +496,7 @@ def class_witness( *, weight: str = "log1p", top: int = 30, + witness: np.ndarray | None = None, ) -> pl.DataFrame: """Rank clonotypes by how much they drive the ``pos`` vs ``neg`` group difference (Prop. ``prop:witness``). @@ -453,17 +511,20 @@ def class_witness( candidates: Clonotype frame to score (e.g. all clonotypes seen in ``pos``). weight: Clone-size weight for the per-sample kernel means. top: Number of top motifs to return. + witness: Optional precomputed witness direction ``μ_pos − μ_neg`` (``(D,)``). When given, + ``pos``/``neg`` are not re-embedded — reuse it to score several candidate sets cheaply. Returns: ``candidates`` with a ``witness_score`` column, sorted descending, truncated to ``top``. """ - def group_mean(frames): - return np.mean([w @ space.rff.transform(Z) - for Z, w in (space.sample_cloud(f, weight=weight) for f in frames)], axis=0) + if witness is None: + def group_mean(frames): + return np.mean([w @ space.rff.transform(Z) + for Z, w in (space.sample_cloud(f, weight=weight) for f in frames)], axis=0) - w = group_mean(pos) - group_mean(neg) + witness = group_mean(pos) - group_mean(neg) psi = space.rff.transform(space.transform_clonotypes(candidates)) # (n, D) - scores = psi @ w + scores = psi @ witness return (candidates.with_columns(pl.Series("witness_score", scores)) .sort("witness_score", descending=True).head(top)) @@ -476,14 +537,134 @@ def hla_stratified_mmd(embs: list[SampleEmbedding], hla: list[set]) -> np.ndarra HLA allele are masked to ``nan`` rather than compared. """ D = mmd_matrix(embs) - out = np.full_like(D, np.nan) - for i in range(len(embs)): - for j in range(len(embs)): - if hla[i] & hla[j]: - out[i, j] = D[i, j] + alleles = sorted({al for s in hla for al in s}) + col = {al: k for k, al in enumerate(alleles)} + A = np.zeros((len(hla), len(alleles)), dtype=np.int8) # sample × allele indicator + for i, s in enumerate(hla): + for al in s: + A[i, col[al]] = 1 + shared = A @ A.T # #shared alleles per pair + return np.where(shared > 0, D, np.nan) + + +def centroid_atypicality(X: np.ndarray, groups: np.ndarray) -> np.ndarray: + """Per-sample cosine distance of its identity vector to its group's centroid — a Φ-geometry op. + + A sample far from its group's mean identity has an atypical clonotype composition (selection / + divergence). Grouping is caller-supplied (tumour type, cohort, …); the geometry — centroid then + ``1 − cos`` — is the library concern. Feeds an ``atypicality`` channel of a digital-donor embedding. + + Args: + X: ``(n, d)`` identity block (a per-sample kernel-mean or its PCA reduction). + groups: ``(n,)`` group label per row; centroids are computed within each group. + + Returns: + ``(n,)`` atypicality in ``[0, 2]`` (0 = on the group centroid direction). + """ + X = np.asarray(X, dtype=np.float64) + groups = np.asarray(groups) + out = np.zeros(X.shape[0]) + for g in np.unique(groups): + m = groups == g + cen = X[m].mean(axis=0) + cn = np.linalg.norm(cen) + 1e-9 + xn = np.linalg.norm(X[m], axis=1) + 1e-9 + out[m] = 1.0 - (X[m] @ cen) / (xn * cn) return out +def correct_batch( + X: np.ndarray, + batch: np.ndarray, + *, + covariates: np.ndarray | None = None, + n_clusters: int = 8, + theta: float = 1.0, + sigma: float = 0.1, + max_iter: int = 10, + ridge: float = 1.0, + seed: int = 0, +) -> np.ndarray: + """Harmony-like cluster-aware batch correction on a stacked sample×feature ``Φ`` matrix. + + Removes a batch offset **per soft cluster** rather than globally, so a batch confounded with a + biological cluster is corrected without erasing that biology — the failure mode of plain + per-group mean subtraction (:func:`mir.cohort.residualize`). Follows Harmony (Korsunsky et al. + 2019, *Nat. Methods*): soft-cluster the samples with a batch-diversity penalty ``theta`` + (clusters pushed toward batch-balanced membership), then subtract each cluster's + membership-weighted batch offset (covariates retained), and iterate. **Reduces exactly to** + ``residualize`` at ``n_clusters=1`` or ``theta=0``. + + Args: + X: Stacked ``(n_samples, n_features)`` matrix (e.g. :func:`mir.explain.stack_embeddings`). + batch: Length-``n_samples`` batch label per row. + covariates: Optional ``(n_samples, k)`` biological covariates to *retain* (never removed). + n_clusters: Number of soft clusters (Harmony's ``K``); ``<=1`` disables clustering. + theta: Batch-diversity penalty strength; ``0`` disables clustering (→ ``residualize``). + sigma: Soft-assignment temperature. + max_iter: Correction iterations (converges in a few — the batch fit vanishes once removed). + ridge: L2 ridge on the per-cluster batch regression (stabilises small clusters). + seed: KMeans seed for the soft-cluster initialisation. + + Returns: + The corrected ``(n_samples, n_features)`` matrix — a new coordinate system (as with + ``residualize``, never compare a corrected ``X`` to an uncorrected one). + """ + # ponytail: Harmony-lite — fixed K, a single KMeans-seeded soft-cluster init, fixed theta. The + # upgrade path is Harmony's full automatic-K + block-coordinate objective; this covers the + # confounded-batch case and reduces exactly to residualize when K<=1 / theta==0. + X = np.asarray(X, dtype=np.float64) + n = X.shape[0] + batch = np.asarray(batch) + _, binv = np.unique(batch, return_inverse=True) + Phi = np.eye(int(binv.max()) + 1)[binv] # (n, nb) batch one-hot + pr_b = Phi.mean(0) # global batch proportions + + if n_clusters <= 1 or theta <= 0: # == residualize (exact) + out = X.copy() + for b in range(Phi.shape[1]): + m = binv == b + out[m] -= X[m].mean(0) + return out + + from sklearn.cluster import KMeans + + # Cosine-geometry soft clustering (Harmony normalises to the unit sphere). + Z = X - X.mean(0) + Z = Z / (np.linalg.norm(Z, axis=1, keepdims=True) + 1e-9) + K = min(n_clusters, n) + centers = KMeans(n_clusters=K, n_init=10, random_state=seed).fit(Z).cluster_centers_ + centers = centers / (np.linalg.norm(centers, axis=1, keepdims=True) + 1e-9) + dist = 2.0 * (1.0 - Z @ centers.T) # (n, K) cosine distance + + # Batch-diversity-penalised soft assignment: up-weight clusters under-represented for a row's + # batch, so clusters become batch-balanced (Harmony's key term) rather than batch-defined. + R = np.exp(-dist / sigma) + R = R / (R.sum(1, keepdims=True) + 1e-9) + for _ in range(3): + O = R.T @ Phi # (K, nb) batch mass per cluster + E = R.sum(0)[:, None] * pr_b[None, :] + 1e-9 # expected under global freq + pen = (E / (O + 1e-9)) ** theta # (K, nb) diversity re-weight + R = np.exp(-dist / sigma) * (Phi @ pen.T) # (n, K) + R = R / (R.sum(1, keepdims=True) + 1e-9) + + # Retain intercept + covariates, remove the batch design; drop one batch column for identifiability. + Cov = np.zeros((n, 0)) if covariates is None else np.asarray(covariates, np.float64).reshape(n, -1) + design = np.hstack([np.ones((n, 1)), Cov, Phi[:, 1:]]) + n_keep = 1 + Cov.shape[1] # columns kept; the rest (batch) removed + Xc = X.copy() + for _ in range(max_iter): + acc = np.zeros_like(Xc) + for k in range(K): + w = R[:, k] + Wd = design * w[:, None] + A = Wd.T @ design + ridge * np.eye(design.shape[1]) + beta = np.linalg.solve(A, Wd.T @ Xc) # (p, d) weighted ridge fit + acc += w[:, None] * (Xc - design[:, n_keep:] @ beta[n_keep:]) + Xc = acc / (R.sum(1, keepdims=True) + 1e-9) + return Xc + + # --------------------------------------------------------------------- self-check diff --git a/pyproject.toml b/pyproject.toml index 2a5b51b..c532642 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "mirpy-lib" -version = "3.3.0" +version = "3.4.0" authors = [ { name="ISALGO lab", email="mikhail.shugay@gmail.com"} ] @@ -30,7 +30,7 @@ dependencies = [ "seqtree>=0.3.0", # AIRR data model (polars schema + IO), germline reference, VDJ-rearrangement # Pgen + synthetic sampling. mirpy has no data-model layer of its own. - "vdjtools>=2.3.0", + "vdjtools>=3.0.0", "tqdm", ] license = { text = "GPL-3.0-or-later" } @@ -41,12 +41,17 @@ annotate = ["vdjmatch"] # Build-time only: regenerate region_annotations.txt (arda) and the baked # germline distance matrices (BioPython). Runtime needs neither. arda comes # via vdjtools[model] (single arda version across the ecosystem), not pinned directly. -build = ["vdjtools[model]>=2.3.0", "biopython"] +build = ["vdjtools[model]>=3.0.0", "biopython"] # Part 2: neural codecs and density methods. ml = ["torch"] # Benchmark harness: DBSCAN eps selection + plotting + HuggingFace dataset fetch + -# survival analysis (CoxPH for the TCGA prognostic benchmark). -bench = ["kneed", "matplotlib", "seaborn", "huggingface_hub", "pynndescent", "lifelines"] +# survival analysis (CoxPH for the TCGA prognostic benchmark). All pure-wheel — no +# compiled/numba deps, so it installs cleanly on any Python via pip or uv. +bench = ["kneed", "matplotlib", "seaborn", "huggingface_hub", "lifelines"] +# Approximate-NN backend for mir.density (backend="ann"). Split out of [bench] because +# pynndescent pulls numba/llvmlite, whose version resolution is Python-version-sensitive; +# the exact "kdtree"/"exact" backends need nothing here. +ann = ["pynndescent"] # Interactive marimo notebooks (docs/examples). examples = ["marimo", "matplotlib", "umap-learn"] dev = ["pytest>=8", "pytest-cov", "psutil>=5"] @@ -58,6 +63,9 @@ docs = [ "ipython>=8,<10", ] +[project.scripts] +mir = "mir.cli:main" + [project.urls] Homepage = "https://github.com/antigenomics/mirpy" Repository = "https://github.com/antigenomics/mirpy" diff --git a/requirements.txt b/requirements.txt index 09b85f5..1536bc0 100644 --- a/requirements.txt +++ b/requirements.txt @@ -5,7 +5,7 @@ numpy>=1.26 polars>=1.0 scikit-learn>=1.3.1 seqtree>=0.3.0 -vdjtools +vdjtools>=3.0.0 tqdm # Optional groups live in pyproject extras: [bench] [annotate] [build] [ml] [docs] [dev]. diff --git a/setup.sh b/setup.sh index 3771428..6c4ba5c 100755 --- a/setup.sh +++ b/setup.sh @@ -1,90 +1,88 @@ #!/usr/bin/env bash -# mirpy v3 bootstrap — reproducible install (conda-based, pure Python, no C build). +# mirpy v3 bootstrap — reproducible install into a repo-local .venv with uv. +# +# Portable: runs under bash OR zsh (bash setup.sh / zsh setup.sh / ./setup.sh). Not fish. +# +# mirpy itself is a pure-Python `py3-none-any` package (no C build). The heavy machinery +# (alignment, Pgen, sampling) is reused from the compiled `seqtree` and `vdjtools` wheels, +# which pip resolves from PyPI — unless you pass --dev-parents to editable-install the +# co-developed sibling checkouts from ../ instead. # # Steps: -# 1. Create/update the `mirpy` conda environment from environment.yml. -# 2. pip install -e ../seqtree ../vdjtools (if present) then -e . -# 3. Optionally install docs deps / run tests. +# 1. Create/activate a repo-local .venv (uv if present, else python -m venv). +# 2. (optional) editable-install co-developed sibling parents from ../ if present. +# 3. pip install -e ".[dev,bench]". # # Flags: -# --no-conda Use the already-active environment instead of creating `mirpy`. -# --docs Install docs deps. -# --test Install [dev,bench] and run the test suite. -# --test-all Same as --test (no separate benchmark tier in v3 yet). +# --dev-parents Editable-install ../seqtree ../vdjtools ../vdjmatch if they exist locally +# (they are co-developed; otherwise the PyPI releases are used). Building the +# siblings compiles their C++ _core extensions (needs a C++ toolchain). +# --docs Also install the [docs] extra. +# --tests Run the fast test suite after install. # -# Usage: -# bash setup.sh [--no-conda] [--docs] [--test] [--test-all] +# Requirements: a C++ toolchain (Xcode Command Line Tools on macOS, build-essential on Linux) +# is needed ONLY when --dev-parents rebuilds seqtree/vdjtools from source. The `[build]` extra +# (arda, BioPython) is for regenerating bundled resources and is not installed here. +# +# Usage: bash setup.sh [--dev-parents] [--docs] [--tests] set -euo pipefail -ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -cd "$ROOT" -ENV_NAME="mirpy" -USE_CONDA=1 +ROOT="$(cd "$(dirname "$0")" && pwd)" # $0, not ${BASH_SOURCE}: works in bash AND zsh +DEV_PARENTS=0 INSTALL_DOCS=0 -RUN_TESTS=0 -RUN_HEAVY=0 +DO_TESTS=0 for arg in "$@"; do case "$arg" in - --no-conda) USE_CONDA=0 ;; - --docs) INSTALL_DOCS=1 ;; - --test) RUN_TESTS=1 ;; - --test-all) RUN_TESTS=1; RUN_HEAVY=1 ;; - --help|-h) sed -n '2,16p' "$0"; exit 0 ;; + --dev-parents) DEV_PARENTS=1 ;; + --docs) INSTALL_DOCS=1 ;; + --tests) DO_TESTS=1 ;; + --no-conda) ;; # accepted for backward-compat; conda is no longer used (no-op) + --help|-h) sed -n '2,30p' "$0"; exit 0 ;; *) echo "Unknown flag: $arg" >&2; exit 2 ;; esac done log() { printf '\033[1;34m[mirpy]\033[0m %s\n' "$*"; } -# --- 1. conda environment -------------------------------------------------- -if [[ "$USE_CONDA" -eq 1 ]]; then - if ! command -v conda >/dev/null 2>&1; then - echo "conda not found on PATH; install miniconda/anaconda or pass --no-conda." >&2 - exit 1 - fi - if conda env list | awk '{print $1}' | grep -qx "$ENV_NAME"; then - log "conda env '$ENV_NAME' exists — updating from environment.yml" - conda env update -n "$ENV_NAME" -f "$ROOT/environment.yml" --prune - else - log "creating conda env '$ENV_NAME' from environment.yml" - conda env create -f "$ROOT/environment.yml" - fi - RUN="conda run -n $ENV_NAME" +# --- 1. repo-local .venv (uv preferred) ------------------------------------ +VENV="$ROOT/.venv" +if command -v uv >/dev/null 2>&1; then + PIP="uv pip" + [ -d "$VENV" ] || { log "creating .venv with uv (Python 3.12)"; uv venv --python 3.12 "$VENV"; } else - RUN="" + log "uv not found — using python -m venv + pip (install uv for faster installs: https://docs.astral.sh/uv/)" + PIP="python -m pip" + [ -d "$VENV" ] || { log "creating .venv"; python -m venv "$VENV"; } fi +# shellcheck disable=SC1091 +. "$VENV/bin/activate" # activate script is bash/zsh compatible -# --- 2. editable install (pure Python — no C build) ------------------------ -# Co-developed siblings first if present, else PyPI resolves them. -for parent in seqtree vdjtools; do - if [[ -d "$ROOT/../$parent" ]]; then - log "pip install -e ../$parent" - $RUN python -m pip install -e "$ROOT/../$parent" - fi -done -log "pip install -e ." -$RUN python -m pip install -e "$ROOT" - -# --- 3. optional docs ------------------------------------------------------ -if [[ "$INSTALL_DOCS" -eq 1 ]]; then - log "installing docs deps" - $RUN python -m pip install -e ".[docs]" +# --- 2. co-developed sibling parents (optional) ---------------------------- +if [ "$DEV_PARENTS" -eq 1 ]; then + for parent in seqtree vdjtools vdjmatch; do + if [ -f "$ROOT/../$parent/pyproject.toml" ]; then + log "editable-install ../$parent (compiles its _core extension)" + $PIP install -e "$ROOT/../$parent" + fi + done fi +# --- 3. editable install (pure Python — no C build for mir itself) --------- +EXTRAS="dev,bench" +[ "$INSTALL_DOCS" -eq 1 ] && EXTRAS="$EXTRAS,docs" +log "$PIP install -e .[$EXTRAS]" +$PIP install -e "$ROOT[$EXTRAS]" + # --- 4. verification ------------------------------------------------------- log "verifying install" -$RUN python -c "import mir; from mir.embedding.tcremp import TCREmp; print('mir import OK')" +python -c "import mir; from mir.embedding.tcremp import TCREmp; print('mir', mir.__version__, 'import OK')" # --- 5. optional tests ----------------------------------------------------- -if [[ "$RUN_TESTS" -eq 1 ]]; then - log "installing test + bench tooling" - $RUN python -m pip install -e ".[dev,bench]" - log "running test suite" - $RUN python -m pytest "$ROOT/tests" -q +if [ "$DO_TESTS" -eq 1 ]; then + log "running fast tests" + python -m pytest "$ROOT/tests" -q -m "not integration and not benchmark" fi log "done." -if [[ "$USE_CONDA" -eq 1 ]]; then - echo " conda activate $ENV_NAME" -fi +echo " source $VENV/bin/activate" diff --git a/tests/test_cli.py b/tests/test_cli.py new file mode 100644 index 0000000..51884b5 --- /dev/null +++ b/tests/test_cli.py @@ -0,0 +1,57 @@ +"""CLI smoke tests — the two `mir embed` commands on tiny bundled-geometry frames.""" + +import polars as pl +import pytest + +from mir.cli import main + + +def _write(path, rows, loci=None): + """Write a tiny AIRR TSV. `rows` = list of (v, j, junction, count).""" + df = pl.DataFrame( + {"v_call": [r[0] for r in rows], "j_call": [r[1] for r in rows], + "junction_aa": [r[2] for r in rows], "duplicate_count": [r[3] for r in rows]} + ) + df.write_csv(path, separator="\t") + + +TRB = [ + ("TRBV10-3*01", "TRBJ2-7*01", "CASSIRSSYEQYF", 120), + ("TRBV20-1*01", "TRBJ1-2*01", "CSARVSGYYGYTF", 40), + ("TRBV28*01", "TRBJ2-1*01", "CASSLGQAYEQFF", 12), + ("TRBV19*01", "TRBJ2-3*01", "CASSISGGADTQYF", 7), +] + + +def test_embed_clonotypes_writes_embedding_table(tmp_path): + src = tmp_path / "S.tsv" + out = tmp_path / "emb.tsv" + _write(src, TRB) + main(["embed", "clonotypes", str(src), "--n-prototypes", "300", "--pca", "3", "-o", str(out)]) + + got = pl.read_csv(out, separator="\t") + assert got.height == 4 # one row per clonotype + assert {"junction_aa", "v_call", "j_call", "e0", "e1", "e2"} <= set(got.columns) + assert got.select(pl.col("e0")).dtypes[0].is_numeric() + + +def test_embed_repertoires_one_row_per_sample(tmp_path): + s1, s2 = tmp_path / "P1.tsv", tmp_path / "P2.tsv" + out = tmp_path / "phi.tsv" + _write(s1, TRB) + _write(s2, TRB[:3]) + main(["embed", "repertoires", str(s1), str(s2), "--n-prototypes", "300", + "--n-rff", "32", "-o", str(out)]) + + got = pl.read_csv(out, separator="\t") + assert got.height == 2 # one Φ(S) per sample + assert got["sample_id"].to_list() == ["P1", "P2"] # id = filename stem + assert got["locus"].unique().to_list() == ["TRB"] + assert any(c.startswith("phi") for c in got.columns) + + +def test_multiple_loci_without_flag_errors(tmp_path): + src = tmp_path / "mixed.tsv" + _write(src, [TRB[0], ("TRAV1-2*01", "TRAJ33*01", "CAVMDSNYQLIW", 5)]) + with pytest.raises(SystemExit): + main(["embed", "clonotypes", str(src), "--n-prototypes", "300"]) diff --git a/tests/test_cohort.py b/tests/test_cohort.py new file mode 100644 index 0000000..4fca3e7 --- /dev/null +++ b/tests/test_cohort.py @@ -0,0 +1,131 @@ +"""Tests for mir.cohort (the digital-donor multi-chain fusion, §T.7). + +Self-contained on bundled TRB+TRA prototypes (no network, no torch). +""" + +import numpy as np +import polars as pl +import pytest + +from mir.cohort import DonorCohort, cluster_samples, fit_donor_embeddings, residualize +from mir.embedding.prototypes import list_available_prototypes +from mir.embedding.tcremp import TCREmp +from mir.repertoire import centroid_atypicality, fit_repertoire_spaces, sample_embedding + +_LOCI = [c for c in ("TRB", "TRA") if ("human", c) in list_available_prototypes()] + +pytestmark = pytest.mark.skipif(len(_LOCI) < 2, reason="need >=2 bundled human loci") + + +@pytest.fixture(scope="module") +def spaces(): + models = {c: TCREmp.from_defaults("human", c, n_prototypes=300) for c in _LOCI} + protos = {c: pl.DataFrame({"v_call": m._proto_v, "j_call": m._proto_j, + "junction_aa": m._proto_junction}).unique() for c, m in models.items()} + return fit_repertoire_spaces(models, protos, n_rff=512, n_rff_second=0, n_components=15, seed=0), protos + + +def _cohort_frames(protos, n=24, drop_locus_for=None): + """n donors in two groups; group 1 gets a public expansion. Optionally drop a locus for one donor.""" + donor_frames, rows = [], [] + for i in range(n): + grp = i % 2 + frames = {} + for c in protos: + if drop_locus_for == (i, c): + continue # this donor lacks this chain + base = protos[c].sample(120, seed=i).with_columns(pl.lit(1.0).alias("duplicate_count")) + if grp == 1: + spike = protos[c].slice(0, 5).with_columns(pl.lit(400.0).alias("duplicate_count")) + base = pl.concat([base, spike]) + frames[c] = base + donor_frames.append(frames) + rows.append({"group": grp}) + return donor_frames, rows + + +def test_fit_core_channels_and_shape(spaces): + sp, protos = spaces + frames, rows = _cohort_frames(protos) + coh = fit_donor_embeddings(sp, frames, rows=rows, id_pca=6, seed=0) + assert coh.X.shape[0] == 24 + for ch in ("identity", "diversity", "coverage"): + assert ch in coh.spec + assert "identity" in coh.spec.attributable # kernel-mean block is attributable + assert np.isfinite(coh.X).all() # imputed + z-scored, no holes + # identity merges both chains: 2 loci x id_pca=6 columns + assert len(coh.spec["identity"]) == len(_LOCI) * 6 + + +def test_transform_reproduces_fit_rows(spaces): + # a core-only cohort: transforming the SAME donors through the stored basis reproduces X exactly. + sp, protos = spaces + frames, rows = _cohort_frames(protos) + coh = fit_donor_embeddings(sp, frames, rows=rows, id_pca=6, seed=0) + Xt = coh.transform(frames) + assert np.allclose(Xt, coh.X, atol=1e-6) + + +def test_missing_chain_is_imputed(spaces): + sp, protos = spaces + frames, rows = _cohort_frames(protos, drop_locus_for=(3, _LOCI[1])) # donor 3 lacks the 2nd locus + coh = fit_donor_embeddings(sp, frames, rows=rows, id_pca=6, seed=0) + assert np.isfinite(coh.X).all() # the hole is imputed, not NaN + + +def test_extra_channels_fuse_and_stay_non_attributable(spaces): + sp, protos = spaces + frames, rows = _cohort_frames(protos) + + def extra(rows, identity): + return {"atypicality": centroid_atypicality(identity, np.array([r["group"] for r in rows]))} + + coh = fit_donor_embeddings(sp, frames, rows=rows, id_pca=6, extra_channels=extra, seed=0) + assert "atypicality" in coh.spec + assert "atypicality" not in coh.spec.attributable # a geometry summary has no clonotype pre-image + # transform now requires the extra block re-supplied + with pytest.raises(ValueError, match="extra channels"): + coh.transform(frames[:3]) + Xt = coh.transform(frames[:3], extra={"atypicality": np.zeros(3)}) + assert Xt.shape == (3, coh.X.shape[1]) + + +def test_save_load_roundtrip_and_cross_basis_refusal(spaces, tmp_path): + sp, protos = spaces + frames, rows = _cohort_frames(protos) + coh = fit_donor_embeddings(sp, frames, rows=rows, id_pca=6, seed=0) + p = tmp_path / "cohort.pkl" + coh.save(p) + back = DonorCohort.load(p) + assert np.allclose(back.X, coh.X) + assert back.spec.names == coh.spec.names and set(back.spaces) == set(sp) + + import pickle + with open(p, "rb") as fh: + d = pickle.load(fh) + any_locus = next(iter(d["spaces"])) + d["spaces"][any_locus]["meta"]["prototype_hash"] = "deadbeefdeadbeef" + with open(p, "wb") as fh: + pickle.dump(d, fh) + with pytest.raises(ValueError, match="prototype hash mismatch"): + DonorCohort.load(p) + DonorCohort.load(p, verify=False) # explicit override allowed + + +def test_residualize_removes_group_offset(): + rng = np.random.default_rng(0) + X = rng.normal(0, 1, (30, 5)) + batch = np.array([0, 1, 2] * 10) + X = X + np.array([0, 5, -5])[batch][:, None] # a large per-batch offset + Xr = residualize(X, batch) + for g in (0, 1, 2): + assert np.allclose(Xr[batch == g].mean(0), 0.0, atol=1e-9) # offset removed + + +def test_cluster_samples_runs(spaces): + sp, protos = spaces + frames, _ = _cohort_frames(protos, n=12) + c0 = _LOCI[0] + embs = [sample_embedding(sp[c0], f[c0], blocks=("mean",)) for f in frames] + labels = cluster_samples(embs) + assert len(labels) == 12 and labels.dtype.kind in "iu" diff --git a/tests/test_density.py b/tests/test_density.py index 5af2fab..a40b693 100644 --- a/tests/test_density.py +++ b/tests/test_density.py @@ -104,6 +104,17 @@ def test_unknown_test_raises(): neighbor_enrichment(obs, bg, radius=1.0, test="chi2") +def test_empty_inputs_and_bad_pseudocount_raise(): + rng = np.random.default_rng(0) + obs, bg = rng.standard_normal((10, 4)), rng.standard_normal((10, 4)) + with pytest.raises(ValueError, match="non-empty"): + neighbor_enrichment(obs[:0], bg, radius=0.5) # empty obs -> friendly error, not a tree crash + with pytest.raises(ValueError, match="non-empty"): + neighbor_enrichment(obs, bg[:0], radius=0.5) # empty bg + with pytest.raises(ValueError, match="pseudocount"): + neighbor_enrichment(obs, bg, radius=0.5, pseudocount=0) # would make fold = n_obs/0 + + def test_unknown_backend_raises(): obs, bg = np.zeros((3, 2)), np.ones((3, 2)) with pytest.raises(ValueError, match="backend must be"): @@ -196,6 +207,16 @@ def test_calibrate_radius_positive(model): assert np.isfinite(r) and r > 0 +def test_calibrate_radius_rejects_nonmutable_junctions(model): + # junctions of length <= 2 never mutate -> drift 0 -> a 0 radius would silently null the test. + obs, bg = _load_olga(150), _load_olga(150, offset=400) + space, _, _ = fit_density_space(model, obs, bg, n_components=20, space="junction") + short = pl.DataFrame({"v_call": [obs["v_call"][0]] * 4, "j_call": [obs["j_call"][0]] * 4, + "junction_aa": ["CF", "CW", "CA", "CY"]}) + with pytest.raises(ValueError, match="non-positive radius"): + calibrate_radius(space, sample_df=short) + + def test_denoise_and_cluster_groups_injected_family(model): # observed = OLGA background + a convergent 1-substitution family; background # subtraction keeps the family, and clustering recovers it as one (non-noise) group. diff --git a/tests/test_eval.py b/tests/test_eval.py new file mode 100644 index 0000000..f24dee6 --- /dev/null +++ b/tests/test_eval.py @@ -0,0 +1,38 @@ +"""Tests for mir.bench.eval (channel_report scorers). Needs [bench]: sklearn (+ lifelines for survival).""" + +import numpy as np +import pytest + +pytest.importorskip("sklearn") + +from mir.bench.eval import cv_auc, km_logrank + + +def test_cv_auc_separates_signal_from_noise(): + rng = np.random.default_rng(0) + n = 200 + y = rng.integers(0, 2, n).astype(float) + signal = y + rng.normal(0, 0.6, n) + noise = rng.normal(0, 1, n) + m_sig, s_sig = cv_auc(signal, y, n_repeats=3) + m_noise, _ = cv_auc(noise, y, n_repeats=3) + assert m_sig > 0.7 > m_noise + assert 0.0 <= s_sig < 0.2 # returns (mean, std) — a CI, not a point + + +@pytest.mark.integration +def test_survival_scorers_recover_risk(): + pytest.importorskip("lifelines") + from mir.bench.eval import cv_cindex + + rng = np.random.default_rng(0) + n = 240 + risk = rng.normal(0, 1, n) + base = rng.normal(0, 1, (n, 2)) + dur = rng.exponential(np.exp(-0.9 * risk)) + evt = (rng.random(n) < 0.7).astype(float) + c_base = cv_cindex(dur, evt, base=base, block=None) + c_full = cv_cindex(dur, evt, base=base, block=risk) + assert c_full > c_base and c_full > 0.6 # the risk block adds concordance + p = km_logrank(dur, evt, (risk > np.median(risk)).astype(int)) + assert p < 0.05 diff --git a/tests/test_prototypes.py b/tests/test_prototypes.py index bfc7b49..0fa3033 100644 --- a/tests/test_prototypes.py +++ b/tests/test_prototypes.py @@ -30,6 +30,13 @@ def test_n_cap_raises(): load_prototypes("human", "TRB", n=N_PROTOTYPES + 1) +def test_n_nonpositive_raises(): + # a non-positive n silently changes the prototype set (df.head(-k)) and its hash -> reject it. + for bad in (0, -5): + with pytest.raises(ValueError): + load_prototypes("human", "TRB", n=bad) + + def test_unknown_locus_file(): # IGK has no prototype file? it does; use a locus with no file: mouse IGH with pytest.raises(FileNotFoundError): diff --git a/tests/test_repertoire.py b/tests/test_repertoire.py index 74febc0..0683ece 100644 --- a/tests/test_repertoire.py +++ b/tests/test_repertoire.py @@ -12,6 +12,7 @@ from mir.repertoire import ( RepertoireSpace, _make_rff, + centroid_atypicality, class_witness, decode_metrics, fit_repertoire_space, @@ -189,6 +190,69 @@ def test_hla_stratified_masks_mismatched_pairs(space): assert np.isnan(S[0, 2]) and np.isnan(S[1, 2]) # mismatched pairs masked +def test_hla_stratified_multiallele_and_empty_set(space): + # vectorized indicator: partial-overlap pairs match; a donor with no HLA matches nobody. + embs = [sample_embedding(space, _sample(_clonotypes(50, offset=o)), blocks=("mean",)) + for o in (0, 100, 200)] + hla = [{"A*02:01", "B*07:02"}, {"B*07:02"}, set()] # 0&1 share B*07; 2 has none + S = hla_stratified_mmd(embs, hla) + assert np.isfinite(S[0, 1]) and np.isfinite(S[1, 0]) # partial overlap still matched + assert np.isnan(S[0, 2]) and np.isnan(S[2, 2]) # empty-HLA donor masked everywhere + assert np.isfinite(S[0, 0]) # self compared when it has any allele + + +def test_unbiased_mmd_rejects_singleton(space): + """Unbiased MMD is undefined at n_eff ≤ 1 (a point mass): raise instead of silently dividing by 0.""" + single = sample_embedding(space, _sample(_clonotypes(1, offset=0)), blocks=("mean",)) + ok = sample_embedding(space, _sample(_clonotypes(60, offset=100)), blocks=("mean",)) + assert single.n_eff == 1.0 + assert np.isfinite(mmd_distance(single, ok)) # biased path still works + with pytest.raises(ValueError, match="single-clonotype"): + mmd_distance(single, ok, unbiased=True) + with pytest.raises(ValueError, match="single-clonotype"): + mmd_matrix([single, ok], unbiased=True) + + +def test_empty_and_zero_count_samples_raise(space): + empty = _sample(_clonotypes(0)) + with pytest.raises(ValueError, match="empty repertoire"): + sample_embedding(space, empty) + zeros = _sample(_clonotypes(10, offset=0), np.zeros(10)) + with pytest.raises(ValueError, match="degenerate"): + sample_embedding(space, zeros) + with pytest.raises(ValueError, match="degenerate"): + sample_descriptor(space, zeros) + + +def test_centroid_atypicality_flags_outliers(): + # a point on its group centroid -> ~0; an outlier -> large; grouping is respected. + X = np.array([[1.0, 0.0], [1.0, 0.02], [1.0, -0.02], [-1.0, 0.0], # group 0 (last is the outlier) + [0.0, 1.0], [0.02, 1.0]]) # group 1 + g = np.array([0, 0, 0, 0, 1, 1]) + a = centroid_atypicality(X, g) + assert a[:3].max() < 0.1 and a[3] > 1.5 # in-group tight vs the flipped outlier + assert a[4] < 0.1 and a[5] < 0.1 # group 1 is computed against its own centroid + + +def test_class_witness_precomputed_matches(space): + """Passing witness= must give identical scores to computing it internally (the sweep fast-path).""" + base = _clonotypes(400, offset=0) + motif = _clonotypes(1, offset=300) + pos = [pl.concat([_sample(base.sample(120, seed=s)), + _sample(motif, lambda n: np.full(n, 400.0))]) for s in range(4)] + neg = [_sample(base.sample(120, seed=s + 50)) for s in range(4)] + cand = pl.concat([base.sample(120, seed=0), motif]).unique() + + def group_mean(frames): + return np.mean([w @ space.rff.transform(Z) + for Z, w in (space.sample_cloud(f) for f in frames)], axis=0) + witness = group_mean(pos) - group_mean(neg) + + a = class_witness(space, pos, neg, cand, top=20) + b = class_witness(space, [], [], cand, top=20, witness=witness) # pos/neg ignored + assert np.allclose(a["witness_score"].to_numpy(), b["witness_score"].to_numpy()) + + def test_class_witness_ranks_injected_motif(space): # a public motif seeded into every 'pos' sample must surface at the top of the witness motif = _clonotypes(1, offset=300) # one specific clonotype @@ -251,3 +315,33 @@ def test_save_load_roundtrip_and_cross_basis_refusal(space, tmp_path): with pytest.raises(ValueError, match="prototype hash mismatch"): RepertoireSpace.load(p) RepertoireSpace.load(p, verify=False) # explicit override is allowed + + +def test_correct_batch_reduces_to_residualize_and_beats_it_under_confound(): + """Harmony-lite correct_batch: == residualize at K=1; preserves biology a global + mean-subtraction destroys when batch is confounded with a biological cluster.""" + import numpy as np + from mir.repertoire import correct_batch + from mir.cohort import residualize + + rng = np.random.default_rng(0) + n, d = 200, 10 + bio = rng.integers(0, 2, n) # biological cluster (dominant signal) + batch = np.where(rng.random(n) < 0.8, bio, 1 - bio) # 80% confounded with bio + X = rng.normal(0, 0.1, (n, d)) + X[:, 0] += np.where(bio == 1, 4.0, -4.0) # axis 0 = biology + X[:, 1] += np.where(batch == 1, 1.5, -1.5) # axis 1 = batch offset + + # K=1 reduces exactly to residualize + assert np.allclose(correct_batch(X, batch, n_clusters=1), residualize(X, batch)) + + Xc = correct_batch(X, batch, n_clusters=2, seed=0) + Xr = residualize(X, batch) + + def gap(M, lab, ax): + return abs(M[lab == 1, ax].mean() - M[lab == 0, ax].mean()) + + # batch offset (axis 1) is removed by the cluster-aware correction + assert gap(Xc, batch, 1) < 0.5 * gap(X, batch, 1) + # biology (axis 0) survives the confound better than plain per-group mean subtraction + assert gap(Xc, bio, 0) > gap(Xr, bio, 0)