mirpy — ML embeddings for immune repertoires

mirpy v3 turns T-/B-cell receptor sequences into fixed-length numeric vectors you can cluster, visualize, and feed to ML models. It implements TCREMP — embedding each receptor by its alignment distances to a fixed set of prototype sequences — so that Euclidean distance in embedding space approximates pairwise alignment distance (Theory T1).

v3 is a slim, embedding-focused rewrite. The classical repertoire toolkit (parsing, overlap, diversity, TCRnet, GLIPH, …) lives on the legacy-v2 branch (mirpy-lib 2.x) and in the sibling tools vdjtools / vdjmatch.

Install

pip install mirpy-lib            # core: numpy, polars, scipy, scikit-learn, seqtree, vdjtools
pip install "mirpy-lib[bench]"   # + benchmark / theory experiments

Pure-Python wheel; the heavy lifting (alignment, Pgen, sampling) is reused from seqtree and vdjtools.

Quick start

import polars as pl
from mir.embedding.tcremp import TCREmp

model = TCREmp.from_defaults("human", "TRB", n_prototypes=3000)   # mode="vjcdr3" | "cdr123"
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, 9000) float32 — 3 distances × 3000 prototypes

Downstream (cluster antigen-specific TCRs):

from mir.embedding.pca import pca_denoise
from mir.bench.metrics import cluster, cluster_metrics
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).

# 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 omitted. Values are data-driven from the bundled prototypes (prototype geometry saturates by these counts; PC columns are the PCA dims retaining ~95% / ~99% variance):

Use 95% PCs for clustering/visualization (the paper's regime); use 99% PCs when reconstructing sequences with the neural inverse codec (diverse chains like IGH/TRD/TRA lose too much sequence detail at 95%). Programmatically: from mir.embedding import get_preset.

from mir.embedding import get_preset
from mir.embedding.pca import pca_denoise
p = get_preset("human", "IGH")
Xc = pca_denoise(X, n_components=p.n_components)          # clustering
Xr = pca_denoise(X, n_components=p.n_components_recon)    # codec reconstruction

What's inside

Background subtraction & clustering (mir.density)

TCRNET/ALICE find antigen-driven convergent clusters by neighbour enrichment. mir.density does the same test with neighbour-counting in the embedding space instead of on a sequence graph (Theory T6): the enrichment E(z) = f_obs(z)/f_gen(z) is estimated by an adaptive-bandwidth balloon estimator with a per-clonotype Poisson/binomial significance test and BH q-values — no graph, and it scales to whole repertoires.

from mir.density import fit_density_space, neighbor_enrichment, enriched_mask, denoise_and_cluster
from mir.embedding.tcremp import TCREmp

model = TCREmp.from_defaults("human", "TRB", n_prototypes=1000)
# background = a control repertoire (TCRNET) or generate_background(...) (ALICE, P_gen)
space, obs_emb, bg_emb = fit_density_space(model, obs_df, control_df, n_components=20, space="full")
res  = neighbor_enrichment(obs_emb, bg_emb, test="binomial")   # balloon + water-level calibration
hits = obs_df.filter(enriched_mask(res, alpha=0.05))            # background-subtracted clones
labels, mask = denoise_and_cluster(obs_emb, res)               # noise-filter + DBSCAN the hits

Use a biological control as the background when you have one (e.g. pre- vs post-vaccination, 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. The density benchmarks (YFV, ankylosing-spondylitis B27, TCRNET) live in the companion 2026-mirpy-analysis 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; pip install "mirpy-lib[ann]").

Sample-level (repertoire) embedding (mir.repertoire)

One fixed vector Φ(S) per repertoire — an order-invariant multiset of clonotypes with clone sizes — depth-robust into the low-coverage bulk-RNA-seq regime (Theory §T.7). Φ(S) sketches the empirical measure ρ_S = Σ_σ w_σ δ_{φ(σ)} (concave frequency weights, so one hyperexpanded clone can't dominate) in three blocks: an RFF kernel mean (depth-robust, codebook-free — no K, no clustering), a coverage-standardized Hill diversity profile, and a second-moment Fisher vector carrying clonotype co-occurrence (HLA-linked public structure). Repertoire distance is the MMD ‖Φ₁(S) − Φ₁(S')‖.

from mir.repertoire import fit_repertoire_space, sample_embedding, mmd_matrix, class_witness
from mir.embedding.tcremp import TCREmp
import polars as pl

model  = TCREmp.from_defaults("human", "TRB", n_prototypes=1000)
space  = fit_repertoire_space(model, pl.concat(samples))   # ONE basis for the whole 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²)
motifs = class_witness(space, pos_samples, neg_samples, candidates)   # public clones separating two groups

Comparability invariant (as with the codecs / density): every sample in a cohort must be embedded through one prototype set and one PCA+RFF basis, or the measures are incomparable — fit_repertoire_space fits that basis once and RepertoireSpace refuses a prototype-hash mismatch.

Use the unbiased MMD (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. When a nuisance batch is present, compare within-batch contrasts (residualize Φ on the batch indicator): a batch offset is first-order and cancels, while a batch-orthogonal signal (e.g. HLA) survives. The 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). Recorded results and theory (T7) live in the companion 2026-mirpy-analysis repo (benchmarks/{BENCHMARKS,THEORY}.md) alongside the benchmark scripts.

Reproduce the paper

The self-contained theory notebooks run on bundled data:

pip install "mirpy-lib[examples]"
marimo edit examples/theory.py          # supplementary S1–S3 (distance laws, D↔d, prototype robustness)
marimo edit examples/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 — 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.

Performance & parallelism

mirpy is CPU-parallel by default and uses the GPU for the neural codecs. Knobs, by hot path:

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

Repo-local .venv via 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 for the architecture and reuse map.