"""Theory-validation experiments (reproduce TCREMP supplementary S1–S3).
These numerically check the properties the embedding is claimed to have, using the
actual v3 pipeline distances (``seqtree.gapblock`` junction dissimilarity):
* **S2 / Theory T1** — Euclidean distance ``D_ij`` in embedding space tracks the
pairwise junction dissimilarity ``d_ij`` (:func:`s2_dissimilarity_distance_correlation`;
paper Pearson R ≈ 0.56). Using the sequences as their own prototypes, the embedding
of sequence *i* is row *i* of the dissimilarity matrix and ``D_ij = ‖d_i − d_j‖₂``.
* **S1 / Theory T4** — ``d_ij`` follows a Gamma law (Gaussian fails); ``D_ij`` follows
a Fréchet (GEV, shape ξ>0) law (:func:`fit_distributions`).
* **S3** — distances from *real* vs *model* prototypes agree (:func:`prototype_source_correlation`;
paper Pearson R ≈ 0.96), i.e. the prototype source barely matters.
"""
from __future__ import annotations
from dataclasses import dataclass
import numpy as np
from scipy import stats
from scipy.spatial.distance import pdist
from mir.distances.junction import junction_distance_matrix
[docs]
def junction_dissimilarity(cdr3s, threads: int = 0) -> np.ndarray:
"""Symmetric ``(n, n)`` junction dissimilarity ``d_ij`` (v3 gapblock; a
negative-type semimetric / squared Hilbert distance, not a metric — ρ=√d is)."""
return junction_distance_matrix(cdr3s, cdr3s, threads=threads)
[docs]
def junction_dissimilarity_sw(cdr3s) -> np.ndarray:
"""Paper-exact ``d_ij`` via Smith-Waterman BLOSUM62 (supplementary S1/S2).
``d_ij = s_ii + s_jj − 2·s_ij`` from a local BLOSUM62 alignment (linear gap).
Requires BioPython (the ``[build]`` extra). O(n²) alignments — use for
the theory validation on a few thousand sequences, not at scale.
"""
from Bio.Align import PairwiseAligner, substitution_matrices
aligner = PairwiseAligner()
aligner.mode = "local"
aligner.substitution_matrix = substitution_matrices.load("BLOSUM62")
aligner.open_gap_score = -1.0
aligner.extend_gap_score = -1.0 # linear gap penalty
seqs = list(cdr3s)
n = len(seqs)
s = np.empty((n, n), dtype=np.float64)
for i in range(n):
s[i, i] = aligner.score(seqs[i], seqs[i])
for i in range(n):
for j in range(i + 1, n):
s[i, j] = s[j, i] = aligner.score(seqs[i], seqs[j])
return s.diagonal()[:, None] + s.diagonal()[None, :] - 2.0 * s
[docs]
@dataclass
class CorrelationResult:
n: int
n_pairs: int
pearson: float
d: np.ndarray # flat i>j dissimilarities
D: np.ndarray # flat i>j embedding-space Euclidean distances
[docs]
def s2_dissimilarity_distance_correlation(
cdr3s, threads: int = 0, dissimilarity: str = "gapblock"
) -> CorrelationResult:
"""Correlate junction dissimilarity ``d_ij`` with embedding distance ``D_ij``.
Uses the *self-prototype* construction from supplementary Fig. S2: the embedding
of each sequence is its row of the dissimilarity matrix, and ``D_ij`` is the
Euclidean distance between rows. Returns the flat i>j vectors and Pearson R.
Args:
dissimilarity: ``"gapblock"`` (the v3 pipeline metric) or ``"sw"``
(paper-exact Smith-Waterman BLOSUM62).
"""
if dissimilarity == "sw":
d = junction_dissimilarity_sw(cdr3s).astype(np.float64)
else:
d = junction_dissimilarity(cdr3s, threads=threads).astype(np.float64)
n = d.shape[0]
iu = np.triu_indices(n, k=1)
d_flat = d[iu] # dissimilarities, i>j order
D_flat = pdist(d, metric="euclidean") # same i>j order as triu_indices(n,1)
r = float(np.corrcoef(d_flat, D_flat)[0, 1])
return CorrelationResult(n=n, n_pairs=d_flat.size, pearson=r, d=d_flat, D=D_flat)
def _aic(logpdf_sum: float, k: int) -> float:
return 2 * k - 2 * logpdf_sum
def _fit_one(data: np.ndarray, dist, floc: bool, init: tuple = ()) -> dict:
if init:
params = dist.fit(data, *init, loc=float(np.median(data)), scale=float(data.std()))
elif floc:
params = dist.fit(data, floc=0)
else:
params = dist.fit(data)
logL = float(np.sum(dist.logpdf(data, *params)))
ks = float(stats.kstest(data, dist.cdf, args=params).statistic)
return {"params": params, "aic": _aic(logL, len(params)), "ks": ks}
[docs]
def fit_distributions(
d_flat: np.ndarray, D_flat: np.ndarray, sample: int = 200_000, seed: int = 0
) -> dict:
"""Fit S1 distributions: Gamma vs Normal for ``d``, GEV/Fréchet vs Normal for ``D``.
Returns a nested dict of ``{ks, aic, params}`` per fit plus the GEV shape ``xi``
(``xi > 0`` ⇒ Fréchet). Large inputs are subsampled to *sample* points.
"""
rng = np.random.default_rng(seed)
def _sub(x):
return x if x.size <= sample else rng.choice(x, sample, replace=False)
d, D = _sub(np.asarray(d_flat, float)), _sub(np.asarray(D_flat, float))
d = d[d > 0]
# genextreme MLE is unstable without a good shape init; c=-xi, seed near Frechet
gev = _fit_one(D, stats.genextreme, floc=False, init=(-0.1,))
return {
# free location: pairwise d is concentrated far from 0, so forcing loc=0 misfits
"d_gamma": _fit_one(d, stats.gamma, floc=False),
"d_normal": _fit_one(d, stats.norm, floc=False),
"D_gev": gev,
"D_normal": _fit_one(D, stats.norm, floc=False),
"D_gev_xi": float(-gev["params"][0]), # scipy genextreme c = -xi
}
_AA = "ACDEFGHIKLMNPQRSTVWY"
def _mutate(seq: str, k: int, rng) -> str:
"""Apply *k* random interior amino-acid substitutions (an SHM proxy)."""
s = list(seq)
interior = list(range(1, len(s) - 1)) # keep the conserved C…[FW] ends
if not interior:
return seq
for p in rng.choice(interior, min(k, len(interior)), replace=False):
choices = _AA.replace(s[p], "") # a *different* residue, as in density._mutate1
s[p] = choices[int(rng.integers(len(choices)))]
return "".join(s)
[docs]
def shm_embedding_drift(
cdr3s, prototypes, max_mut: int = 8, n_rep: int = 3, seed: int = 0, threads: int = 0
) -> dict[int, tuple[float, float]]:
"""Embedding drift vs mutation load (Theory T5, the IGH/SHM case).
Applies ``k`` random interior substitutions (a somatic-hypermutation proxy) to each
CDR3 and measures the junction-embedding Euclidean distance to the original,
``D_k = ‖φ(mutated) − φ(original)‖``. Returns ``k -> (mean D_k, std)``. The claim: the
drift is *bounded by the mutation load* — ``D_k`` grows ~linearly in ``k`` — so heavily
hypermutated IGH sequences sit controllably far from germline in embedding space.
"""
rng = np.random.default_rng(seed)
seqs = list(cdr3s)
base = junction_distance_matrix(seqs, prototypes, threads=threads).astype(np.float64)
out: dict[int, tuple[float, float]] = {0: (0.0, 0.0)}
for k in range(1, max_mut + 1):
drift: list[float] = []
for _ in range(n_rep):
emb = junction_distance_matrix([_mutate(s, k, rng) for s in seqs],
prototypes, threads=threads).astype(np.float64)
drift.extend(np.linalg.norm(emb - base, axis=1).tolist())
out[k] = (float(np.mean(drift)), float(np.std(drift)))
return out
[docs]
def prototype_source_correlation(
query_cdr3s, real_prototypes, model_prototypes, threads: int = 0
) -> dict:
"""Correlate embedding distances built from *real* vs *model* prototypes (S3).
Embeds ``query_cdr3s`` against each prototype set (junction only), takes pairwise
Euclidean distances under each embedding, and returns their Pearson R (paper ≈ 0.96).
"""
Xr = junction_distance_matrix(query_cdr3s, real_prototypes, threads=threads).astype(np.float64)
Xm = junction_distance_matrix(query_cdr3s, model_prototypes, threads=threads).astype(np.float64)
Dr, Dm = pdist(Xr), pdist(Xm)
return {"pearson": float(np.corrcoef(Dr, Dm)[0, 1]), "n_pairs": Dr.size}
def _hamming1_counts(seqs) -> np.ndarray:
"""Discrete Hamming-1 neighbour count per sequence (equal-length, exactly one mismatch)."""
seqs = list(seqs)
out = np.zeros(len(seqs), dtype=np.int64)
for i, a in enumerate(seqs):
la, c = len(a), 0
for j, b in enumerate(seqs):
if i == j or len(b) != la:
continue
m = 0
for x, y in zip(a, b):
if x != y:
m += 1
if m > 1:
break
if m == 1:
c += 1
out[i] = c
return out
[docs]
def tcrnet_convergence(
obs_cdr3s, bg_cdr3s, prototypes, *, n_components: int = 25,
scales=(0.5, 1.0, 1.5, 2.0, 3.0), seed: int = 0, threads: int = 0
) -> dict:
"""T6: continuous embedding enrichment converges to discrete Hamming-1 enrichment.
Embeds observed and background CDR3s against the prototype set (junction only), reduces
with one shared ``StandardScaler → PCA``, and at radius ``scale × r₁`` (``r₁`` = the
median one-substitution embedding drift) counts each observed clonotype's neighbours in
embedding space. Returns the Spearman correlation between those continuous counts and the
discrete Hamming-1 neighbour counts. The correlation is high at small radii and fades as
the radius grows past one substitution — numerically confirming that graph
neighbour-enrichment (TCRNET/ALICE) is the ``r→0`` limit of the density ratio ``E(z)``.
Args:
obs_cdr3s: Observed junction sequences.
bg_cdr3s: Background junction sequences (P_gen sample / control).
prototypes: Prototype junctions defining the embedding.
n_components: PCA dimensionality of the shared coordinate system.
scales: Radius multiples of the one-substitution scale ``r₁`` to evaluate.
seed: RNG seed for the calibration mutations and PCA solver.
Returns:
``{radius_1sub, hamming1_mean, spearman_by_scale, spearman_at_1sub}``.
"""
from sklearn.decomposition import PCA
from sklearn.neighbors import BallTree
from sklearn.preprocessing import StandardScaler
obs = junction_distance_matrix(obs_cdr3s, prototypes, threads=threads).astype(np.float64)
bg = junction_distance_matrix(bg_cdr3s, prototypes, threads=threads).astype(np.float64)
scaler = StandardScaler().fit(np.vstack([obs, bg]))
k = min(n_components, obs.shape[0] + bg.shape[0], obs.shape[1])
pca = PCA(n_components=k, random_state=seed).fit(scaler.transform(np.vstack([obs, bg])))
obs_emb = pca.transform(scaler.transform(obs))
rng = np.random.default_rng(seed)
mut = [_mutate(s, 1, rng) for s in obs_cdr3s]
mut_emb = pca.transform(scaler.transform(
junction_distance_matrix(mut, prototypes, threads=threads).astype(np.float64)))
r1 = float(np.median(np.linalg.norm(obs_emb - mut_emb, axis=1)))
discrete = _hamming1_counts(obs_cdr3s)
tree = BallTree(obs_emb)
corr = {}
for f in scales:
n_obs = tree.query_radius(obs_emb, r1 * f, count_only=True) - 1
corr[round(float(f), 2)] = float(stats.spearmanr(n_obs, discrete).statistic)
return {
"radius_1sub": r1,
"hamming1_mean": float(discrete.mean()),
"spearman_by_scale": corr,
"spearman_at_1sub": corr[1.0],
}
# ---------------------------------------------------------------------------
# Codec losslessness (appendix draft codec_losslessness.tex): the three-level hierarchy
# ---------------------------------------------------------------------------
def _levenshtein(a: str, b: str) -> int:
"""Levenshtein edit distance (iterative two-row DP)."""
if a == b:
return 0
if not a or not b:
return len(a) or len(b)
prev = list(range(len(b) + 1))
for i, ca in enumerate(a, 1):
cur = [i]
for j, cb in enumerate(b, 1):
cur.append(min(prev[j] + 1, cur[j - 1] + 1, prev[j - 1] + (ca != cb)))
prev = cur
return prev[-1]
[docs]
def codec_losslessness(codes, seqs, recon=None, *, eps: float = 1e-6) -> dict:
"""Measurable losslessness of a codec, split into decoder-independent vs -dependent parts.
**Informational ceiling** (decoder-independent, from ``codes`` alone). The compact code is
injective iff distinct sequences map to distinct codes; two distinct sequences whose codes
coincide can never *both* be recovered, so ``exact_ceiling = 1 − collision_rate`` upper-bounds
any decoder's exact-match — the information the code retains, separate from how well a decoder
reads it. Encoding is deterministic (identical sequences ⇒ identical codes), so duplicate
sequences are collapsed first and a *collision* is a near-coincident code pair between
**distinct** sequences (nearest-code distance < ``eps``). ``nn_dist_{median,min}`` report the
code-resolution scale so ``eps`` can be read in context.
**Reconstructive loss** (needs ``recon``, the decoded strings row-aligned to ``seqs``):
``exact_match``, mean **edit distance** (a *graded* distortion, unlike binary exact-match), and
**per-position** token accuracy over the length-40 frame — localizing *where* loss falls
(conserved C…[FW] anchors vs the specificity-bearing variable middle). Per-position accuracy
re-encodes ``recon`` into the canonical frame, so it measures positional token agreement.
Args:
codes: ``(n, m)`` compact codes, one per sequence (e.g. whitened-PCA junction code).
seqs: the ``n`` true junction strings.
recon: optional ``n`` decoded strings (e.g. ``InverseDecoder.decode(codes)``). When
``None``, only the decoder-independent ceiling is returned.
eps: code-distance below which two distinct sequences count as colliding.
Returns:
Always ``{n, n_unique, collision_rate, exact_ceiling, nn_dist_median, nn_dist_min}``; and
when ``recon`` is given, additionally ``{exact_match, mean_edit, token_acc, anchor_acc,
middle_acc, pos_acc}`` (``pos_acc`` a length-40 array).
"""
from sklearn.neighbors import BallTree
seqs = list(seqs)
codes = np.asarray(codes, dtype=np.float64)
n = len(seqs)
# informational ceiling: dedup sequences, then nearest *distinct*-code distance
uniq = list({s: i for i, s in enumerate(seqs)}.values())
C = codes[uniq]
n_unique = len(uniq)
if n_unique >= 2:
nn = BallTree(C).query(C, k=2)[0][:, 1] # distance to nearest distinct code
collision_rate = float((nn < eps).mean())
nn_median, nn_min = float(np.median(nn)), float(nn.min())
else:
collision_rate, nn_median, nn_min = 0.0, float("nan"), float("nan")
out = {
"n": n, "n_unique": n_unique, "collision_rate": collision_rate,
"exact_ceiling": 1.0 - collision_rate,
"nn_dist_median": nn_median, "nn_dist_min": nn_min,
}
if recon is None:
return out
from mir.ml.tokenize import FIXED_LEN, _LEFT_ANCHOR, _RIGHT_ANCHOR, encode_indices
recon = list(recon)
exact = float(np.mean([r == s for r, s in zip(recon, seqs)]))
mean_edit = float(np.mean([_levenshtein(r, s) for r, s in zip(recon, seqs)]))
pos_acc = (encode_indices(seqs) == encode_indices(recon)).mean(axis=0) # (40,)
anchor = np.zeros(FIXED_LEN, dtype=bool)
anchor[:_LEFT_ANCHOR] = anchor[FIXED_LEN - _RIGHT_ANCHOR:] = True
out.update({
"exact_match": exact, "mean_edit": mean_edit, "token_acc": float(pos_acc.mean()),
"anchor_acc": float(pos_acc[anchor].mean()), "middle_acc": float(pos_acc[~anchor].mean()),
"pos_acc": pos_acc,
})
return out