"""TCREMP: distance-vector embeddings for TCR/BCR clonotypes.
TCREMP (T-Cell Receptor EMbedding with Prototypes) embeds each clonotype as a flat
vector of distances to a fixed set of *prototype* clonotypes. For clonotype *i* and
prototype *k*, three distances are emitted; the first two depend on the ``mode``:
* ``mode="vjcdr3"`` (default): **V-germline**, **J-germline**, **junction**.
* ``mode="cdr123"``: **CDR1**, **CDR2** (both germline V-gene-determined), **junction**.
The embedding vector (vjcdr3 shown) is::
[v_i1, j_i1, junc_i1, v_i2, j_i2, junc_i2, ..., v_iK, j_iK, junc_iK]
for K prototypes; the output matrix is ``(n_clonotypes, 3*K)`` ``float32``.
Junction (CDR3) distances are computed with :func:`seqtree.gapblock.score_matrix`
(best of contiguous gap-block placements at ``(3, 4, -4, -3)``, BLOSUM62 Gram
penalty). The result ``d`` is a *negative-type semimetric* — a squared Hilbert
distance ``d = ‖ψ(a)−ψ(b)‖²`` (symmetric, ``d(a,a)=0``, non-negative) — not itself a
metric; the induced genuine metric is ``ρ = √d`` (see ``mir/distances/junction.py``).
V/J and CDR1/CDR2 distances are looked up
from baked matrices (:class:`mir.distances.germline.GermlineDistances`), same Gram
construction. Input is a polars DataFrame with the ``vdjtools.io.schema`` columns
``v_call``, ``j_call``, ``junction_aa``.
Coordinate-system options (all default to the published v3 space): ``metric="sqrt"`` maps
every block to the metric ``ρ=√d`` (benchmarked a wash, so ``"squared"`` is default);
``matrix=`` a custom ``seqtree.SubstitutionMatrix`` and
``alignment="sw"`` (paper-exact Smith-Waterman, slow) reshape the junction block — see
:func:`mir.distances.junction.junction_distance_matrix`.
Typical usage::
import polars as pl
from mir.embedding.tcremp import TCREmp
model = TCREmp.from_defaults("human", "TRB", n_prototypes=1000)
df = pl.DataFrame({"v_call": ["TRBV10-3*01"], "j_call": ["TRBJ2-7*01"],
"junction_aa": ["CASSIRSSYEQYF"]})
X = model.embed(df) # shape: (1, 3000)
"""
from __future__ import annotations
from pathlib import Path
import numpy as np
import polars as pl
from mir.aliases import normalize_locus_alias, normalize_species_alias
from mir.distances.germline import GermlineDistances, load_germline_distances
from mir.distances.junction import DEFAULT_GAP_POSITIONS, junction_distance_matrix
from mir.embedding.prototypes import load_prototypes
MODES = ("vjcdr3", "cdr123")
# mode -> ((component, query_gene_column, prototype_gene_attr), ...) for the two
# germline slots (0, 1). The junction is always slot 2.
_MODE_SPEC: dict[str, tuple[tuple[str, str, str], tuple[str, str, str]]] = {
"vjcdr3": (("V", "v_call", "_proto_v"), ("J", "j_call", "_proto_j")),
"cdr123": (("CDR1", "v_call", "_proto_v"), ("CDR2", "v_call", "_proto_v")),
}
_REQUIRED_COLS = ("v_call", "j_call", "junction_aa")
#: The 20 standard amino acids, anchored -- the same predicate ``vdjtools.signature.blocks``
#: sanitises on. seqtree's alphabet is WIDER than this, and that is the problem: a stop codon
#: ``*`` and the ambiguity codes ``X``/``B``/``Z`` do not crash the distance code, they return a
#: finite, meaningless distance, which is strictly worse than crashing.
_STANDARD_AA = r"^[ACDEFGHIKLMNPQRSTVWY]+$"
#: The 20 amino acids plus the two characters a *well-formed* AIRR table can legitimately carry:
#: ``*`` for a stop codon and ``_`` for the legacy out-of-frame marker. A junction outside this is
#: not a biological category anyone chose to keep -- it is a broken file. Measured across the
#: 6,047,716 rows of this project's clinical AIRR store: **zero** violations, and zero ``_``; every
#: non-plain junction was a stop codon. So the character class is not a guess about what data looks
#: like, and a violation is worth an exception rather than a filter.
_PARSEABLE_AA = r"^[ACDEFGHIKLMNPQRSTVWY*_]+$"
[docs]
class TCREmp:
"""Prototype distance-vector embedding for one locus."""
def __init__(
self,
species: str,
locus: str,
prototypes: pl.DataFrame,
germline: GermlineDistances,
mode: str = "vjcdr3",
gap_positions: tuple[int, ...] = DEFAULT_GAP_POSITIONS,
threads: int = 0,
metric: str = "squared",
matrix=None,
alignment: str = "gapblock",
replicate: int = 0,
shm_scale: float = 1.0,
):
if mode not in MODES:
raise ValueError(f"mode must be one of {MODES}, got {mode!r}")
if metric not in ("squared", "sqrt"):
raise ValueError(f"metric must be 'squared' or 'sqrt', got {metric!r}")
missing = set(_REQUIRED_COLS) - set(prototypes.columns)
if missing:
raise ValueError(f"prototypes missing columns: {sorted(missing)}")
self.species = species
self.locus = locus
self.mode = mode
self.threads = threads
self.metric = metric
# provenance only — the prototypes are passed in; this records *which* bundled draw they
# are so the prototype hash (and every save/load comparability check) can see it.
self.replicate = replicate
self._matrix = matrix
self._alignment = alignment
self._germline = germline
# How hard somatic mutation pushes a clonotype away from its germline V. 1.0 puts the SHM
# penalty on the same footing as the germline-to-germline distance, both being BLOSUM62
# squared distances over the same alphabet; 0.0 reproduces the T-cell behaviour exactly,
# which is what makes this safe to leave on.
self.shm_scale = shm_scale
self._proto_v = prototypes["v_call"].to_list()
self._proto_j = prototypes["j_call"].to_list()
self._proto_junction = prototypes["junction_aa"].to_list()
self.n_prototypes = len(self._proto_junction)
self._gap_positions = gap_positions
# fail fast if the locus lacks a component this mode needs
for comp, _, _ in _MODE_SPEC[mode]:
if not germline.has(comp):
raise ValueError(
f"germline distances for {locus} lack component {comp!r} "
f"required by mode {mode!r}"
)
[docs]
@classmethod
def from_defaults(
cls,
species: str,
locus: str,
n_prototypes: int | None = None,
mode: str = "vjcdr3",
replicate: int = 0,
**kwargs,
) -> "TCREmp":
"""Build from the bundled prototypes and baked germline distances.
Args:
species: Species identifier (aliases resolved).
locus: Receptor locus (aliases resolved).
n_prototypes: Prototype count; ``None`` uses the per-chain recommended preset
(:func:`mir.embedding.presets.get_preset`).
mode: ``"vjcdr3"`` (default) or ``"cdr123"``.
replicate: Which bundled prototype draw to use. ``0`` (default) is *the* prototype
set — every preset, bundled codec and published number uses it. ``r > 0`` is an
independent disjoint draw of the same size, for measuring how much a result
depends on the prototype sample (see :func:`mir.embedding.prototypes.n_replicates`).
Embeddings from different replicates are **not** mutually comparable.
"""
from mir.embedding.presets import get_preset
species_c = normalize_species_alias(species)
locus_c = normalize_locus_alias(locus)
if n_prototypes is None:
n_prototypes = get_preset(species_c, locus_c).n_prototypes
prototypes = load_prototypes(species_c, locus_c, n=n_prototypes, replicate=replicate)
germline = load_germline_distances(species_c, locus_c)
return cls(species_c, locus_c, prototypes, germline, mode=mode, replicate=replicate,
**kwargs)
[docs]
@classmethod
def from_file(
cls,
prototypes_path: str | Path,
species: str,
locus: str,
mode: str = "vjcdr3",
**kwargs,
) -> "TCREmp":
"""Build from a custom prototype TSV (``v_call``/``j_call``/``junction_aa``).
Embeddings from a custom prototype set are **not** comparable to those from
the bundled prototypes — the coordinate system differs.
"""
species_c = normalize_species_alias(species)
locus_c = normalize_locus_alias(locus)
prototypes = pl.read_csv(prototypes_path, separator="\t", columns=list(_REQUIRED_COLS))
germline = load_germline_distances(species_c, locus_c)
return cls(species_c, locus_c, prototypes, germline, mode=mode, **kwargs)
@property
def n_features(self) -> int:
return 3 * self.n_prototypes
def _shm_penalties(self, clonotypes: pl.DataFrame):
"""Per-clonotype SHM penalty, or ``None`` when the frame carries no mutation evidence.
Reads ``v_mutations`` (a mutation spec like ``"A23V,S31N"``) in preference to
``v_identity`` (a scalar), because it is the better evidence. Returns ``None`` rather
than a zero vector when neither is present: a zero vector asserts every cell is germline,
which for a BCR sample is a claim, not a default.
Wrapping happens here rather than in the constructor so that the same ``TCREmp`` handles
a frame with mutation data and one without -- a cohort mixes them.
"""
from mir.distances.shm import MutatedGermlineDistances, shm_penalty_batch
muts = clonotypes["v_mutations"].to_list() if "v_mutations" in clonotypes.columns else None
iden = clonotypes["v_identity"].to_list() if "v_identity" in clonotypes.columns else None
if muts is None and iden is None:
return None
if not isinstance(self._germline, MutatedGermlineDistances):
self._germline = MutatedGermlineDistances(self._germline, scale=self.shm_scale)
return shm_penalty_batch(muts, iden)
def _junction_distances(self, junctions: list[str]) -> np.ndarray:
# squared d always; embed() applies the sqrt metric uniformly across all three blocks
return junction_distance_matrix(
junctions, self._proto_junction,
gap_positions=self._gap_positions, threads=self.threads,
matrix=self._matrix, alignment=self._alignment,
)
[docs]
def embed(self, clonotypes: pl.DataFrame) -> np.ndarray:
"""Embed a clonotype frame into ``(n_clonotypes, 3 * n_prototypes)`` float32.
Args:
clonotypes: polars DataFrame with columns ``v_call``, ``j_call``,
``junction_aa`` (``vdjtools.io.schema`` names).
Returns:
``float32`` array of shape ``(len(clonotypes), 3 * n_prototypes)``,
interleaved per prototype as ``[slot0, slot1, junction]``.
Raises:
ValueError: If a required column is missing, or ``junction_aa`` carries a null, a
stop codon, an out-of-frame marker, or any character outside the 20 amino acids.
**There is no opt-out.** mirpy filters non-productive rearrangements on every
read; reaching this guard means something bypassed that, and embedding anyway
would return a finite, meaningless number rather than fail.
"""
missing = set(_REQUIRED_COLS) - set(clonotypes.columns)
if missing:
raise ValueError(f"clonotypes missing columns: {sorted(missing)}")
# a null junction reaches seqtree as None and dies there without naming the column/row
n_null = clonotypes["junction_aa"].is_null().sum()
if n_null:
raise ValueError(f"junction_aa has {n_null} null value(s); drop or impute them first "
"(v_call/j_call nulls are fine — they take the allele fallback)")
# CORRUPTION, not a category of receptor. An ambiguity code (X/B/Z), a lowercase residue,
# an empty string or a stray character is not something an AIRR pipeline decided to emit --
# it means the table is damaged or was written by something that did not agree on the
# alphabet. There is no legitimate reason to embed it, and silently filtering it would
# hide a broken input -- so this is an error even though it is not a receptor question.
corrupt = ~clonotypes["junction_aa"].str.contains(_PARSEABLE_AA).fill_null(False)
n_corrupt = int(corrupt.sum())
if n_corrupt:
ex = clonotypes.filter(corrupt)["junction_aa"].unique().head(5).to_list()
raise ValueError(
f"junction_aa has {n_corrupt} of {clonotypes.height} UNPARSEABLE value(s): outside "
f"the 20 amino acids and outside '*' and '_', e.g. {ex}. A well-formed AIRR table "
"carries nothing else, so this is a data problem, not a filtering decision — "
"ambiguity codes (X/B/Z), lowercase residues, empty strings and stray characters "
"all land here. Fix the source rather than filtering it away."
)
# '_' (the legacy vdjtools out-of-frame marker) isn't in seqtree's amino-acid alphabet and
# crashes gapblock with an opaque error.
n_oof = clonotypes["junction_aa"].str.contains("_", literal=True).fill_null(False).sum()
if n_oof:
raise ValueError(
f"junction_aa has {n_oof} value(s) containing '_' (legacy out-of-frame marker); "
"seqtree's alphabet doesn't include it and will crash — drop non-coding "
"clonotypes first with vdjtools.signature.blocks.sanitise"
)
# A stop codon does not crash -- '*' is in seqtree's alphabet, so it returns a finite,
# meaningless distance. That is strictly worse than crashing, and it is why this guard has
# NO opt-out in mirpy: the library's contract is that it cannot embed a non-productive
# rearrangement meaningfully. vdjtools, which only computes statistics, leaves the same
# filter optional.
bad = ~clonotypes["junction_aa"].str.contains(_STANDARD_AA).fill_null(False)
n_bad = int(bad.sum())
if n_bad:
ex = clonotypes.filter(bad)["junction_aa"].head(3).to_list()
raise ValueError(
f"junction_aa has {n_bad} of {clonotypes.height} non-productive value(s), e.g. "
f"{ex}. '*' is in seqtree's alphabet, so this does NOT raise downstream -- it "
"embeds to a finite, meaningless distance and contaminates the geometry silently. "
"mirpy therefore refuses it with no opt-out. Filter first with "
"vdjtools.preprocess.filter_productive(); if you want the non-productive fraction "
"itself, that is a vdjtools question, not a mirpy one."
)
n = clonotypes.height
out = np.empty((n, 3 * self.n_prototypes), dtype=np.float32)
# A B-cell frame may carry its somatic mutation load, and where it does the V slot must
# use it: otherwise a hypermutated IGHV3-23 and a germline IGHV3-23 resolve to the same
# allele row and embed to identical V coordinates, discarding the whole germinal-centre
# signal. Opt-in by the column being present, so T-cell embeddings are unchanged.
shm = self._shm_penalties(clonotypes)
for slot, (comp, gene_col, proto_attr) in enumerate(_MODE_SPEC[self.mode]):
kw = {"shm": shm} if (shm is not None and comp == "V") else {}
out[:, slot::3] = self._germline.matrix(
comp, clonotypes[gene_col].to_list(), getattr(self, proto_attr), **kw
)
out[:, 2::3] = self._junction_distances(clonotypes["junction_aa"].to_list())
if self.metric == "sqrt":
# all three blocks are the same Gram dissimilarity d; sqrt uniformly -> the metric
# ρ = √d, homogeneous across V/J/junction.
np.sqrt(np.clip(out, 0.0, None, out=out), out=out)
return out
[docs]
class PairedTCREmp:
"""Paired-chain embedding: concatenate two per-chain :class:`TCREmp` embeddings."""
def __init__(self, chains: dict[str, TCREmp]):
# dict preserves insertion order; embedding = concat in that order
self._chains = chains
[docs]
@classmethod
def from_defaults(
cls,
species: str,
loci: tuple[str, str] = ("TRA", "TRB"),
n_prototypes: int | None = None,
mode: str = "vjcdr3",
replicate: int = 0,
**kwargs,
) -> "PairedTCREmp":
# n_prototypes=None -> each locus uses its own recommended preset
chains = {
locus: TCREmp.from_defaults(species, locus, n_prototypes, mode, replicate, **kwargs)
for locus in loci
}
return cls(chains)
@property
def n_features(self) -> int:
return sum(m.n_features for m in self._chains.values())
[docs]
def embed(self, frames: dict[str, pl.DataFrame]) -> np.ndarray:
"""Embed per-chain frames (keyed by locus, row-aligned) and concatenate.
Args:
frames: Mapping ``locus -> DataFrame``; must cover every locus in the
model, all with the same row count (row *i* is the same clonotype).
"""
missing = set(self._chains) - set(frames)
if missing:
raise ValueError(f"missing frames for loci: {sorted(missing)}")
parts = [self._chains[locus].embed(frames[locus]) for locus in self._chains]
heights = {p.shape[0] for p in parts}
if len(heights) != 1:
raise ValueError(f"per-chain frames must have equal row counts, got {heights}")
return np.hstack(parts)
if __name__ == "__main__":
df = pl.DataFrame(
{
"v_call": ["TRBV10-3*01", "TRBV20-1*01"],
"j_call": ["TRBJ2-7*01", "TRBJ1-2*01"],
"junction_aa": ["CASSIRSSYEQYF", "CSARVSGYYGYTF"],
}
)
for mode in MODES:
m = TCREmp.from_defaults("human", "TRB", n_prototypes=50, mode=mode)
X = m.embed(df)
assert X.shape == (2, 150), X.shape
assert X.dtype == np.float32
assert np.isfinite(X).all()
print(f" mode={mode}: X {X.shape} range [{X.min():.1f}, {X.max():.1f}]")
# metric="sqrt": ρ = √d elementwise over all blocks
Xd = TCREmp.from_defaults("human", "TRB", n_prototypes=50).embed(df)
Xs = TCREmp.from_defaults("human", "TRB", n_prototypes=50, metric="sqrt").embed(df)
assert np.allclose(Xs, np.sqrt(np.clip(Xd, 0.0, None)), atol=1e-4)
# custom matrix (pam250) and Smith-Waterman alignment both run and give finite embeddings
import seqtree
Xp = TCREmp.from_defaults("human", "TRB", n_prototypes=50,
matrix=seqtree.SubstitutionMatrix.pam250()).embed(df)
assert np.isfinite(Xp).all() and Xp.shape == (2, 150)
try: # alignment="sw" needs BioPython ([build] extra) — optional, so skip when absent
Xw = TCREmp.from_defaults("human", "TRB", n_prototypes=50, alignment="sw").embed(df)
assert np.isfinite(Xw).all() and Xw.shape == (2, 150)
sw = f"SW junc range [{Xw[:, 2::3].min():.1f}, {Xw[:, 2::3].max():.1f}]"
except ImportError:
sw = "SW skipped (no BioPython)"
print(f" metric=sqrt ρ==√d OK; pam250 OK; {sw}")
print("mir.embedding.tcremp self-check OK")