TCRmeta 0.1.0__py3-none-any.whl

tcrmeta/__init__.py

@@ -0,0 +1,22 @@
+"""TCRmeta: antigen-aware TCR repertoire embeddings and downstream
+repertoire-level analysis (CSS scoring, UMAP projection, energy-distance
+shift), built on a pretrained-ESM2 + contrastive-fine-tuned encoder.
+"""
+from .embedding import embed_repertoire
+from .reference import build_reference, load_reference, save_reference
+from .css import compute_css, load_default_reference
+from .umap_plot import plot_umap
+from .energy import energy_shift
+
+__all__ = [
+    "embed_repertoire",
+    "build_reference",
+    "save_reference",
+    "load_reference",
+    "load_default_reference",
+    "compute_css",
+    "plot_umap",
+    "energy_shift",
+]
+
+__version__ = "0.1.0"

tcrmeta/_dataset.py

@@ -0,0 +1,53 @@
+"""Fixed-length tokenization of (CDR1 + CDR2 + CDR2.5 + CDR3) TCR
+sequences for the ESM2 base encoder.
+"""
+from __future__ import annotations
+
+from typing import Dict, List, Tuple
+
+import pandas as pd
+import torch
+from torch.utils.data import Dataset
+
+DEFAULT_MAX_LEN = 48
+
+
+class TCRDatasetFixedLen(Dataset):
+    def __init__(self, df: pd.DataFrame, token2idx: Dict[str, int], max_len: int = DEFAULT_MAX_LEN):
+        self.df = df.reset_index(drop=True)
+        self.t2i = token2idx
+        self.max_len = max_len
+        self.cls = token2idx["[CLS]"]
+        self.pad = token2idx["[PAD]"]
+        self.eos = token2idx["[EOS]"]
+        self.unk = token2idx["[UNK]"]
+        self.gap = token2idx["X"]  # separator between CDR segments
+
+    def _encode_seq(self, s: str) -> List[int]:
+        return [self.t2i.get(a, self.unk) for a in s]
+
+    def __getitem__(self, i) -> Tuple[Tuple[str, str], List[int]]:
+        r = self.df.iloc[i]
+        v_seq = r.cdr1 + "X" + r.cdr2 + "X" + r.cdr2_5
+        tokens = (
+            [self.cls]
+            + self._encode_seq(v_seq)
+            + [self.gap]
+            + self._encode_seq(r.cdr3aa)
+            + [self.eos]
+        )
+        if len(tokens) < self.max_len:
+            tokens = tokens + [self.pad] * (self.max_len - len(tokens))
+        else:
+            tokens = tokens[: self.max_len]
+        key = (r.cdr3aa, r.v_gene)
+        return key, tokens
+
+    def __len__(self) -> int:
+        return len(self.df)
+
+
+def fixedlen_collate(batch):
+    keys, seqs = zip(*batch)
+    input_ids = torch.tensor(seqs, dtype=torch.long)
+    return keys, input_ids

tcrmeta/_model.py

@@ -0,0 +1,131 @@
+"""Model architectures: the ESM2-backbone masked-LM encoder used during
+pretraining (only its encoder + CLS embedding is used at inference time),
+and the ResMLP Siamese projection head used during contrastive fine-tuning.
+"""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import EsmModel
+
+from ._weights import ESM2_BACKBONE
+
+
+class TCRMaskedModel(nn.Module):
+    """ESM2 encoder + masked-LM head. Only `encoder` (and its CLS-token
+    embedding) is used downstream by TCRmeta; `lm_head` is retained so
+    pretrained checkpoints load without a state_dict mismatch.
+    """
+
+    def __init__(self, backbone_name: str = ESM2_BACKBONE, vocab_size: int = 33):
+        super().__init__()
+        self.encoder = EsmModel.from_pretrained(backbone_name)
+        self.hidden_size = self.encoder.config.hidden_size
+        self.lm_head = nn.Linear(self.hidden_size, vocab_size)
+
+    def forward(self, input_tokens, attention_mask=None, return_embedding=False):
+        outputs = self.encoder(input_ids=input_tokens, attention_mask=attention_mask)
+        sequence_output = outputs.last_hidden_state  # (batch, seq_len, hidden)
+        logits = self.lm_head(sequence_output)
+        return (logits, sequence_output) if return_embedding else logits
+
+
+class ResMLPBlock(nn.Module):
+    """Pre-norm residual MLP block: LayerNorm -> Linear -> GELU -> Dropout
+    -> Linear -> Dropout, added back with a learnable (small-initialized)
+    residual scale so the block starts close to identity.
+    """
+
+    def __init__(self, dim: int, hidden: int = 512, dropout: float = 0.1):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.fc1 = nn.Linear(dim, hidden)
+        self.fc2 = nn.Linear(hidden, dim)
+        self.act = nn.GELU()
+        self.drop = nn.Dropout(dropout)
+        self.res_scale = nn.Parameter(torch.tensor(0.1))
+
+        nn.init.xavier_uniform_(self.fc1.weight)
+        nn.init.zeros_(self.fc2.weight)
+        nn.init.zeros_(self.fc2.bias)
+
+    def forward(self, x):
+        h = self.norm(x)
+        h = self.fc1(h)
+        h = self.act(h)
+        h = self.drop(h)
+        h = self.fc2(h)
+        h = self.drop(h)
+        return x + self.res_scale * h
+
+
+class ResMLPStack(nn.Module):
+    """`depth` stacked ResMLPBlocks."""
+
+    def __init__(self, dim: int, hidden: int = 512, dropout: float = 0.1, depth: int = 2):
+        super().__init__()
+        self.blocks = nn.ModuleList([ResMLPBlock(dim, hidden, dropout) for _ in range(depth)])
+
+    def forward(self, x):
+        for blk in self.blocks:
+            x = blk(x)
+        return x
+
+
+class SiameseOnEmb_ResMLP(nn.Module):
+    """Contrastive fine-tuning projection head: ResMLP re-mixes the
+    480-dim base embedding, then a linear projection + L2 normalization
+    maps it onto the unit sphere in the antigen-aware embedding space.
+    """
+
+    def __init__(
+        self,
+        in_dim: int,
+        proj_dim: int = 64,
+        hidden: int = 768,
+        dropout: float = 0.1,
+        depth: int = 3,
+    ):
+        super().__init__()
+        self.mlp = ResMLPStack(in_dim, hidden=hidden, dropout=dropout, depth=depth)
+        self.proj = nn.Linear(in_dim, proj_dim, bias=False)
+        # Learnable scalars carried over from contrastive training; unused
+        # at inference (encode()/forward() distance is computed downstream),
+        # kept only so pretrained state_dicts load without a mismatch.
+        self.alpha = nn.Parameter(torch.tensor(10.0))
+        self.m = nn.Parameter(torch.tensor(1.0))
+        self.scale = nn.Parameter(torch.tensor(1.0))
+
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.mlp(x)
+        z = self.proj(x)
+        return F.normalize(z, dim=-1)
+
+    def forward(self, x1, x2):
+        z1 = self.encode(x1)
+        z2 = self.encode(x2)
+        d = torch.norm(z1 - z2, dim=-1) * self.scale
+        return z1, z2, d
+
+
+def load_state_dict_flexible(model: nn.Module, ckpt) -> None:
+    """Load a checkpoint saved under any of the common key conventions
+    (raw state_dict, {"model_state_dict": ...}, {"state_dict": ...},
+    an EMA shadow dict, etc.).
+    """
+    if isinstance(ckpt, dict) and all(torch.is_tensor(v) for v in ckpt.values()):
+        model.load_state_dict(ckpt, strict=True)
+        return
+    for key in ("model_state_dict", "ema_state_dict", "state_dict", "model"):
+        if key in ckpt:
+            model.load_state_dict(ckpt[key], strict=True)
+            return
+    if "ema_shadow" in ckpt:
+        sd = model.state_dict()
+        for k, v in ckpt["ema_shadow"].items():
+            if k in sd:
+                sd[k] = v
+        model.load_state_dict(sd, strict=False)
+        return
+    raise RuntimeError(f"Unrecognized checkpoint format, top-level keys: {list(ckpt.keys())}")

tcrmeta/_tokenizer.py

@@ -0,0 +1,23 @@
+"""ESM2 tokenizer -> amino-acid token id lookup, built once and cached."""
+from __future__ import annotations
+
+from functools import lru_cache
+from typing import Dict
+
+from transformers import EsmTokenizer
+
+from ._weights import ESM2_BACKBONE
+
+AMINO_ACIDS = list("ACDEFGHIKLMNPQRSTVWYX")
+
+
+@lru_cache(maxsize=1)
+def build_token2idx(backbone_name: str = ESM2_BACKBONE) -> Dict[str, int]:
+    tokenizer = EsmTokenizer.from_pretrained(backbone_name)
+    token2idx = {aa: tokenizer.convert_tokens_to_ids(aa) for aa in AMINO_ACIDS}
+    token2idx["[PAD]"] = tokenizer.convert_tokens_to_ids("<pad>")
+    token2idx["[MASK]"] = tokenizer.convert_tokens_to_ids("<mask>")
+    token2idx["[UNK]"] = tokenizer.convert_tokens_to_ids("<unk>")
+    token2idx["[CLS]"] = tokenizer.convert_tokens_to_ids("<cls>")
+    token2idx["[EOS]"] = tokenizer.convert_tokens_to_ids("<eos>")
+    return token2idx

tcrmeta/_utils.py

@@ -0,0 +1,185 @@
+"""Shared, dependency-light utilities used across TCRmeta's embedding,
+CSS, UMAP, and energy-shift modules.
+"""
+from __future__ import annotations
+
+from typing import Optional
+
+import numpy as np
+import pandas as pd
+
+
+# --------------------------------------------------------------------------- #
+# Device handling
+# --------------------------------------------------------------------------- #
+def resolve_device(device: Optional[str] = "auto"):
+    """Resolve a user-facing device string into a torch.device.
+
+    Parameters
+    ----------
+    device:
+        One of "auto" (default, uses CUDA if available), "cpu", "cuda",
+        "cuda:0", "mps", or an already-constructed torch.device.
+
+    Note: torch is imported lazily here (rather than at module level) so
+    that the pure numpy/pandas utilities in this module (downsampling,
+    column validation) can be used/tested without requiring torch to be
+    installed.
+    """
+    import torch
+
+    if isinstance(device, torch.device):
+        return device
+    if device is None or device == "auto":
+        if torch.cuda.is_available():
+            return torch.device("cuda")
+        if getattr(torch.backends, "mps", None) is not None and torch.backends.mps.is_available():
+            return torch.device("mps")
+        return torch.device("cpu")
+    return torch.device(device)
+
+
+# --------------------------------------------------------------------------- #
+# Downsampling
+# --------------------------------------------------------------------------- #
+def downsample_multinomial(
+    counts: np.ndarray,
+    target: int,
+    rng: Optional[np.random.Generator] = None,
+    seed: int = 0,
+) -> np.ndarray:
+    """Multinomial downsampling of clone counts to a fixed target depth.
+
+    If the repertoire's total count is already <= target, counts are
+    returned unchanged (no upsampling).
+    """
+    if rng is None:
+        rng = np.random.default_rng(seed)
+    counts = np.asarray(counts, dtype=np.int64)
+    total = int(counts.sum())
+    if total <= target:
+        return counts.copy()
+    p = counts / total
+    return rng.multinomial(target, p)
+
+
+def normalize_weights(counts: np.ndarray) -> Optional[np.ndarray]:
+    """Normalize a count/weight vector to sum to 1. Returns None if the
+    vector sums to <= 0 (nothing usable to weight by)."""
+    counts = np.asarray(counts, dtype=np.float64)
+    s = counts.sum()
+    if s <= 0:
+        return None
+    return counts / s
+
+
+def downsample_then_embed(
+    df: pd.DataFrame,
+    downsample: int,
+    seed: int,
+    device: str,
+    embed_kwargs: dict,
+) -> pd.DataFrame:
+    """Downsample a repertoire's raw counts FIRST, then embed only the
+    clones that survive downsampling — used by compute_css/plot_umap.
+
+    Multinomial downsampling to `downsample` total reads keeps at most
+    `downsample` unique (cdr3aa, v_gene) rows (each surviving row needs
+    >=1 read). Embedding only those rows — instead of embedding the
+    repertoire's full unique-clone set and downsampling afterward — can
+    be a large speedup for repertoires with many more unique clones than
+    `downsample`.
+
+    Caveat: a handful of surviving rows can still get dropped afterward
+    if their V gene has no CDR1/CDR2/CDR2.5 entry in TCRmeta's germline
+    lookup table (attach_embeddings' internal dropna) — for most real
+    repertoires this is rare/negligible, but it means results can differ
+    very slightly from embedding-then-downsampling in the (uncommon)
+    case where such unrecognized-V-gene rows exist and would otherwise
+    have consumed some of the downsampling budget.
+
+    Returns the downsampled + embedded dataframe with an added 'ds'
+    float64 column (post-downsample weight per surviving row).
+
+    Always embeds with embedding_type="final" (the antigen-aware,
+    contrastively fine-tuned embedding) regardless of what's in
+    embed_kwargs — compute_css/plot_umap's reference maps are built on
+    that embedding space, so scoring against them requires it. Only
+    embed_repertoire() itself exposes a choice of embedding_type.
+    """
+    from .embedding import attach_embeddings  # lazy: keeps this module
+
+    # importable/testable without torch/transformers.
+    df = validate_repertoire_columns(df)
+    rng = np.random.default_rng(seed)
+    counts = df["count"].to_numpy(dtype=np.int64)
+    ds = downsample_multinomial(counts, downsample, rng=rng)
+    keep = ds > 0
+    df = df.loc[keep].copy()
+    df["ds"] = ds[keep].astype(np.float64)
+    embed_kwargs = {**embed_kwargs, "embedding_type": "final"}
+    df_e = attach_embeddings(df, device=device, **embed_kwargs)
+    return df_e
+
+
+# --------------------------------------------------------------------------- #
+# Repertoire column validation
+# --------------------------------------------------------------------------- #
+REQUIRED_COLUMNS = ("cdr3aa", "v_gene", "count")
+
+
+def validate_repertoire_columns(df: pd.DataFrame) -> pd.DataFrame:
+    """Validate and lightly normalize a TCR repertoire dataframe.
+
+    TCRmeta requires the exact columns 'cdr3aa', 'v_gene', 'count' —
+    there is no alias-guessing. If your data uses different column
+    names (e.g. Adaptive's 'aminoAcid'/'vGeneName'/'count (templates/reads)',
+    or AIRR's 'junction_aa'/'v_call'/'duplicate_count'), rename them
+    yourself before calling into TCRmeta.
+
+    Normalization performed (does not change column names):
+      - cdr3aa is cast to str and given canonical IMGT flanks (leading
+        C, trailing F) if missing.
+      - v_gene is cast to str, allele suffix (e.g. '*01') stripped.
+      - count is coerced to int; rows with count <= 0 are dropped.
+    """
+    missing = [c for c in REQUIRED_COLUMNS if c not in df.columns]
+    if missing:
+        raise ValueError(
+            f"Repertoire dataframe is missing required column(s): {missing}. "
+            f"TCRmeta requires exactly these columns: {REQUIRED_COLUMNS}. "
+            "Rename your columns before calling this function (no aliases are "
+            "auto-detected)."
+        )
+
+    df = df.copy()
+    df["cdr3aa"] = df["cdr3aa"].astype(str)
+    df["cdr3aa"] = df["cdr3aa"].str.replace(r"^(?!C)", "C", regex=True)
+    df["cdr3aa"] = df["cdr3aa"].str.replace(r"(?<!F)$", "F", regex=True)
+
+    df["v_gene"] = df["v_gene"].astype(str).str.replace(r"\*.*$", "", regex=True).str.strip()
+
+    df["count"] = pd.to_numeric(df["count"], errors="coerce").fillna(0).astype(np.int64)
+    df = df[df["count"] > 0].reset_index(drop=True)
+    return df
+
+
+def aggregate_duplicate_clones(df: pd.DataFrame, sort: bool = True) -> pd.DataFrame:
+    """Collapse duplicate (cdr3aa, v_gene) rows by summing 'count'.
+
+    `sort` controls whether the result is reordered alphabetically by
+    (cdr3aa, v_gene) (pandas groupby's default) or left in first-occurrence
+    order. This matters for reproducibility: multinomial downsampling
+    consumes the RNG stream in row order, so which convention to use
+    depends on what the calling code's original (pre-packaging) script
+    did. energy_shift's original script deduplicated with pandas'
+    default sort=True before downsampling, so energy.py relies on the
+    default here. CSS/UMAP's original scoring scripts never deduplicated
+    rows at all (raw CSV order preserved throughout) — those code paths
+    (attach_embeddings) intentionally skip calling this function.
+    """
+    agg_cols = {"count": "sum"}
+    other_cols = [c for c in df.columns if c not in ("cdr3aa", "v_gene", "count")]
+    for c in other_cols:
+        agg_cols[c] = "first"
+    return df.groupby(["cdr3aa", "v_gene"], as_index=False, sort=sort).agg(agg_cols)

tcrmeta/_weights.py

@@ -0,0 +1,119 @@
+"""Model artifact + reference map download/caching.
+
+All non-code artifacts (the fine-tuned base encoder checkpoint, the 7
+projection-head checkpoints, the GPA rotation matrices, the V-gene CDR
+lookup table, and the default shipped reference map) are hosted on the
+Hugging Face Hub rather than bundled into the PyPI package, and are
+downloaded once and cached locally (via huggingface_hub, which handles
+its own on-disk cache).
+
+Set the TCRMETA_HF_REPO environment variable to point at a different
+Hub repo (e.g. a private mirror), and TCRMETA_WEIGHTS_DIR to bypass the
+Hub entirely and load artifacts from a local directory instead (useful
+offline, or before the Hub repo exists).
+
+If the Hub repo is PRIVATE, authentication is required to download from
+it. Either run `huggingface-cli login` once (huggingface_hub then reuses
+that cached token automatically), or set the TCRMETA_HF_TOKEN
+environment variable to an access token explicitly (useful in
+non-interactive environments like CI or a shared cluster).
+"""
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from typing import Optional
+
+# TODO: replace with the actual Hugging Face Hub repo id once created,
+# e.g. "mingyaopan/tcrmeta-weights". Can be overridden without a code
+# change via the TCRMETA_HF_REPO environment variable.
+DEFAULT_HF_REPO = os.environ.get("TCRMETA_HF_REPO", "mingyaopan/tcrmeta-weights")
+
+# Filenames as they should exist in the Hub repo / local weights dir.
+BASE_CKPT_FILE = "base_encoder/best_model.pth"
+ROTATIONS_FILE = "finetune/rotations_gpa.pkl"
+VGENE_CDR_FILE = "reference_data/vgene_cdr.csv"
+PROJ_HEAD_FILES = [f"finetune/proj_head_{i}.pt" for i in range(7)]
+DEFAULT_REFERENCE_FILE = "reference_maps/young_reference.pkl"
+
+ESM2_BACKBONE = "facebook/esm2_t12_35M_UR50D"
+
+
+def _local_weights_dir() -> Optional[Path]:
+    d = os.environ.get("TCRMETA_WEIGHTS_DIR")
+    return Path(d) if d else None
+
+
+def get_cache_dir() -> Path:
+    cache_dir = Path(os.environ.get("TCRMETA_CACHE_DIR", Path.home() / ".cache" / "tcrmeta"))
+    cache_dir.mkdir(parents=True, exist_ok=True)
+    return cache_dir
+
+
+def resolve_artifact(relative_path: str, repo_id: str = DEFAULT_HF_REPO) -> Path:
+    """Return a local filesystem path for a named artifact, downloading
+    (and caching) it from the Hugging Face Hub if it isn't available
+    locally via TCRMETA_WEIGHTS_DIR.
+    """
+    local_dir = _local_weights_dir()
+    if local_dir is not None:
+        local_path = local_dir / relative_path
+        if not local_path.exists():
+            raise FileNotFoundError(
+                f"TCRMETA_WEIGHTS_DIR is set to {local_dir}, but {relative_path} "
+                "was not found under it."
+            )
+        return local_path
+
+    # Common mistake: pointing TCRMETA_HF_REPO at a local filesystem path
+    # (that's what TCRMETA_WEIGHTS_DIR is for) instead of a Hub repo id
+    # ("namespace/repo_name"). Catch it here with a clear message rather
+    # than letting huggingface_hub's HFValidationError surface instead.
+    looks_like_local_path = repo_id.startswith(("/", "~", ".")) or os.path.exists(repo_id)
+    if looks_like_local_path:
+        raise ValueError(
+            f"repo_id resolved to '{repo_id}', which looks like a local filesystem "
+            "path, not a Hugging Face Hub repo id ('namespace/repo_name'). If you "
+            "meant to load weights from a local directory, set the "
+            "TCRMETA_WEIGHTS_DIR environment variable instead of TCRMETA_HF_REPO:\n"
+            f'  os.environ["TCRMETA_WEIGHTS_DIR"] = "{repo_id}"\n'
+            "TCRMETA_HF_REPO is only for overriding which Hub repo to download from."
+        )
+
+    try:
+        from huggingface_hub import hf_hub_download
+    except ImportError as e:  # pragma: no cover
+        raise ImportError(
+            "huggingface_hub is required to auto-download TCRmeta model weights. "
+            "Install it with `pip install huggingface_hub`, or set TCRMETA_WEIGHTS_DIR "
+            "to a local directory containing the weights."
+        ) from e
+
+    path = hf_hub_download(
+        repo_id=repo_id,
+        filename=relative_path,
+        cache_dir=str(get_cache_dir()),
+        token=os.environ.get("TCRMETA_HF_TOKEN"),  # None -> huggingface_hub
+        # falls back to the cached `huggingface-cli login` token, if any.
+    )
+    return Path(path)
+
+
+def resolve_base_checkpoint() -> Path:
+    return resolve_artifact(BASE_CKPT_FILE)
+
+
+def resolve_rotations() -> Path:
+    return resolve_artifact(ROTATIONS_FILE)
+
+
+def resolve_vgene_table() -> Path:
+    return resolve_artifact(VGENE_CDR_FILE)
+
+
+def resolve_projection_heads() -> list[Path]:
+    return [resolve_artifact(f) for f in PROJ_HEAD_FILES]
+
+
+def resolve_default_reference() -> Path:
+    return resolve_artifact(DEFAULT_REFERENCE_FILE)