"""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")
[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,
):
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
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 _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]``.
"""
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)")
# '_' (the legacy vdjtools out-of-frame marker) isn't in seqtree's amino-acid alphabet and
# crashes gapblock with an opaque error; catch it here with a message naming the fix.
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.preprocess.filter_functional"
)
n = clonotypes.height
out = np.empty((n, 3 * self.n_prototypes), dtype=np.float32)
for slot, (comp, gene_col, proto_attr) in enumerate(_MODE_SPEC[self.mode]):
out[:, slot::3] = self._germline.matrix(
comp, clonotypes[gene_col].to_list(), getattr(self, proto_attr)
)
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")