diff --git a/README.md b/README.md index 8c0947d..d03fe40 100644 --- a/README.md +++ b/README.md @@ -163,6 +163,7 @@ The canonical prediction kind strings are defined in `mhctools.pred.Kind`. | `pMHC_affinity` | Peptide-MHC binding affinity | | `pMHC_presentation` | Likelihood of surface presentation (EL/processing) | | `pMHC_stability` | Peptide-MHC complex stability | +| `pMHC_TCR_binding` | TCR recognition of a peptide-MHC (pMHC:TCR binding) | | `immunogenicity` | T-cell immunogenicity | | `antigen_processing` | Combined processing score | | `proteasome_cleavage` | Proteasomal cleavage score | @@ -203,6 +204,7 @@ Examples: | `MHCflurry` haplotype mode | `pMHC_presentation` | `haplotype` | `I` | | `MHCflurry` per-allele panel mode | `pMHC_presentation` | `single_allele` | `I` | | `Pepsickle` | `proteasome_cleavage` | `none` | `none` | +| `NetTCR` | `pMHC_TCR_binding` | `none` | `I` | For MHCflurry presentation, `presentation_allele_mode="haplotype"` treats the requested alleles as one sample genotype and emits one `pMHC_presentation` @@ -268,6 +270,35 @@ Processing predictors use configurable scoring to aggregate per-position cleavage probabilities into peptide-level scores. See `ProcessingPredictor` and `ProteasomePredictor` for details. +### TCR specificity + +| Predictor | Kinds produced | Requires | +|---|---|---| +| `NetTCR` | pMHC:TCR binding | [NetTCR-2.2](https://github.com/mnielLab/NetTCR-2.2) clone (set `NETTCR_DIR`) + a TFLite runtime (`pip install mhctools[nettcr]`) | + +`NetTCR` predicts whether a paired αβ T-cell receptor recognises a +(class-I) peptide. Unlike the MHC-ligand predictors, its input is a peptide +plus a `TCR` (the six CDR loops), not an allele, and it emits the +`pMHC_TCR_binding` kind. NetTCR ships its pretrained weights in its git +repository as small TFLite models; this wrapper runs the pan cross-validation +ensemble in-process and does not need NetTCR's conda environment. + +```python +from mhctools import NetTCR, TCR + +predictor = NetTCR() # resolves NETTCR_DIR / ~/NetTCR-2.2 +tcr = TCR( + cdr1a="NSASQS", cdr2a="VYSSG", cdr3a="VVEGDKVI", + cdr1b="MGHRA", cdr2b="YSYEKL", cdr3b="ASSHSGYEQF", name="clone1") + +# Score explicit (peptide, TCR) pairs... +results = predictor.predict_pairs([("LLWNGPMAV", tcr)]) +results[0].tcr_binding.score # ensemble-mean recognition probability + +# ...or every peptide x TCR combination. +results = predictor.predict(["LLWNGPMAV", "GILGFVFTL"], [tcr]) +``` + ## Commandline examples ### Prediction for user-supplied peptide sequences diff --git a/mhctools/__init__.py b/mhctools/__init__.py index 5956735..af85cf4 100644 --- a/mhctools/__init__.py +++ b/mhctools/__init__.py @@ -14,6 +14,7 @@ preds_from_rows, ) from .sample import MultiSample +from .tcr import TCR from .iedb import ( IedbNetMHCcons, IedbNetMHCpan, @@ -60,6 +61,7 @@ "BigMHC_IM": (".bigmhc", "BigMHC_IM"), "MHCflurry": (".mhcflurry", "MHCflurry"), "MHCflurry_Affinity": (".mhcflurry", "MHCflurry_Affinity"), + "NetTCR": (".nettcr", "NetTCR"), } @@ -75,7 +77,7 @@ def __getattr__(name): raise AttributeError( "module %r has no attribute %r" % (__name__, name)) -__version__ = "3.14.1" +__version__ = "3.15.0" __all__ = [ "Prediction", @@ -90,6 +92,7 @@ def __getattr__(name): "best_direction", "preds_from_rows", "MultiSample", + "TCR", "BindingPrediction", "BindingPredictionCollection", "IedbNetMHCcons", @@ -140,6 +143,7 @@ def __getattr__(name): "BigMHC", "BigMHC_EL", "BigMHC_IM", + "NetTCR", "RandomBindingPredictor", "UnsupportedAllele", ] diff --git a/mhctools/nettcr.py b/mhctools/nettcr.py new file mode 100644 index 0000000..e908d6b --- /dev/null +++ b/mhctools/nettcr.py @@ -0,0 +1,387 @@ +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Wrapper for NetTCR-2.2 (https://github.com/mnielLab/NetTCR-2.2). + +NetTCR-2.2 predicts pMHC:TCR binding — whether a paired αβ T-cell receptor +(given by its six CDR loops) recognises a given (class-I) peptide. This is a +different modality from the MHC-ligand predictors in the rest of mhctools: +the input is ``(peptide, TCR)`` rather than ``(peptide, allele)``, and the +output kind is :attr:`~mhctools.pred.Kind.pMHC_TCR_binding`. + +Unlike the DTU ``netMHC*`` tools, NetTCR ships its pretrained weights +directly in its git repository as small TFLite models (the pan ensemble is +20 models of ~0.4 MB each). This wrapper runs them **in-process** through a +TFLite interpreter — it does *not* need NetTCR's conda environment. We only +require the cloned repository (for the weights) and a TFLite runtime +(``ai-edge-litert``, ``tflite-runtime``, or ``tensorflow``). + +The BLOSUM50 encoding, per-feature padding lengths, and name-based input +tensor assignment here reproduce ``src/predict.py`` from NetTCR-2.2 exactly +(verified to ~1e-7 against its own output). Scores are the mean over the +pan cross-validation ensemble, matching the published usage. +""" + +from __future__ import annotations + +import glob +import os + +import numpy as np +import pandas as pd + +from .pred import COLUMNS, Kind, PeptideResult, Prediction +from .tcr import TCR + + +# BLOSUM50, restricted to the 20 standard amino acids, in NetTCR's column +# order (A R N D C Q E G H I L K M F P S T W Y V). Copied verbatim from +# NetTCR-2.2 `src/keras_utils.py` (`blosum50_20aa`). +_BLOSUM50_20AA = { + 'A': (5, -2, -1, -2, -1, -1, -1, 0, -2, -1, -2, -1, -1, -3, -1, 1, 0, -3, -2, 0), + 'R': (-2, 7, -1, -2, -4, 1, 0, -3, 0, -4, -3, 3, -2, -3, -3, -1, -1, -3, -1, -3), + 'N': (-1, -1, 7, 2, -2, 0, 0, 0, 1, -3, -4, 0, -2, -4, -2, 1, 0, -4, -2, -3), + 'D': (-2, -2, 2, 8, -4, 0, 2, -1, -1, -4, -4, -1, -4, -5, -1, 0, -1, -5, -3, -4), + 'C': (-1, -4, -2, -4, 13, -3, -3, -3, -3, -2, -2, -3, -2, -2, -4, -1, -1, -5, -3, -1), + 'Q': (-1, 1, 0, 0, -3, 7, 2, -2, 1, -3, -2, 2, 0, -4, -1, 0, -1, -1, -1, -3), + 'E': (-1, 0, 0, 2, -3, 2, 6, -3, 0, -4, -3, 1, -2, -3, -1, -1, -1, -3, -2, -3), + 'G': (0, -3, 0, -1, -3, -2, -3, 8, -2, -4, -4, -2, -3, -4, -2, 0, -2, -3, -3, -4), + 'H': (-2, 0, 1, -1, -3, 1, 0, -2, 10, -4, -3, 0, -1, -1, -2, -1, -2, -3, 2, -4), + 'I': (-1, -4, -3, -4, -2, -3, -4, -4, -4, 5, 2, -3, 2, 0, -3, -3, -1, -3, -1, 4), + 'L': (-2, -3, -4, -4, -2, -2, -3, -4, -3, 2, 5, -3, 3, 1, -4, -3, -1, -2, -1, 1), + 'K': (-1, 3, 0, -1, -3, 2, 1, -2, 0, -3, -3, 6, -2, -4, -1, 0, -1, -3, -2, -3), + 'M': (-1, -2, -2, -4, -2, 0, -2, -3, -1, 2, 3, -2, 7, 0, -3, -2, -1, -1, 0, 1), + 'F': (-3, -3, -4, -5, -2, -4, -3, -4, -1, 0, 1, -4, 0, 8, -4, -3, -2, 1, 4, -1), + 'P': (-1, -3, -2, -1, -4, -1, -1, -2, -2, -3, -4, -1, -3, -4, 10, -1, -1, -4, -3, -3), + 'S': (1, -1, 1, 0, -1, 0, -1, 0, -1, -3, -3, 0, -2, -3, -1, 5, 2, -4, -2, -2), + 'T': (0, -1, 0, -1, -1, -1, -1, -2, -2, -1, -1, -1, -1, -2, -1, 2, 5, -3, -2, 0), + 'W': (-3, -3, -4, -5, -5, -1, -3, -3, -3, -3, -2, -3, -1, 1, -4, -4, -3, 15, 2, -3), + 'Y': (-2, -1, -2, -3, -3, -1, -2, -3, 2, -1, -1, -2, 0, 4, -3, -2, -2, 2, 8, -1), + 'V': (0, -3, -3, -4, -1, -3, -3, -4, -4, 4, 1, -3, 1, -1, -3, -2, 0, -3, -1, 5), +} +_BLOSUM50 = {aa: np.array(vec, dtype=np.float32) for aa, vec in _BLOSUM50_20AA.items()} + +# NetTCR encodes residues as BLOSUM50 rows, pads missing positions with a +# constant, then divides the whole array by this factor. Residues become +# BLOSUM50/5 and padded positions become -1 (= -5 / 5). +_PAD_VALUE = -5.0 +_NORM_FACTOR = 5.0 + +# Per-feature maximum sequence lengths (from NetTCR-2.2 `src/predict.py`). +_FEATURE_MAX_LEN = { + "pep": 12, + "a1": 7, + "a2": 8, + "a3": 22, + "b1": 6, + "b2": 7, + "b3": 23, +} + + +def _load_interpreter(model_path): + """Load a TFLite interpreter from whichever runtime is installed. + + Prefers the lightweight LiteRT/tflite runtimes, falling back to the + ``tensorflow.lite`` interpreter. All expose the same interface. + """ + try: + from ai_edge_litert.interpreter import Interpreter + return Interpreter(model_path=model_path) + except ImportError: + pass + try: + from tflite_runtime.interpreter import Interpreter + return Interpreter(model_path=model_path) + except ImportError: + pass + try: + import tensorflow as tf + return tf.lite.Interpreter(model_path=model_path) + except ImportError as e: + raise ImportError( + "NetTCR needs a TFLite runtime. Install one of: " + "`ai-edge-litert` (recommended, lightweight), `tflite-runtime`, " + "or `tensorflow`.") from e + + +def _find_nettcr_dir(nettcr_path=None): + """Resolve the NetTCR-2.2 installation directory. + + Checks, in order: + 1. The *nettcr_path* argument + 2. The ``NETTCR_DIR`` environment variable + 3. ``~/NetTCR-2.2`` and ``~/code/NetTCR-2.2`` + + An explicitly-provided path (argument or ``NETTCR_DIR``) is validated up + front so a typo fails with a clear message rather than later when no + models are found. + """ + clone_hint = "Clone from https://github.com/mnielLab/NetTCR-2.2" + for source, path in ( + ("nettcr_path argument", nettcr_path), + ("NETTCR_DIR", os.environ.get("NETTCR_DIR"))): + if path: + if not os.path.isdir(path): + raise FileNotFoundError( + "NetTCR-2.2 directory from %s does not exist: %s. %s" + % (source, path, clone_hint)) + return path + home = os.path.expanduser("~") + for candidate in ( + os.path.join(home, "NetTCR-2.2"), + os.path.join(home, "code", "NetTCR-2.2")): + if os.path.isdir(candidate): + return candidate + raise FileNotFoundError( + "NetTCR-2.2 not found. Set NETTCR_DIR or pass nettcr_path= to the " + "constructor. %s" % clone_hint) + + +def _encode_feature(sequences, feature): + """BLOSUM50-encode a list of sequences for one NetTCR feature. + + Returns a ``float32`` array of shape ``(n, max_len, 20)`` where residues + are BLOSUM50/5 and padded positions are -1, matching NetTCR's + ``enc_list_bl_max_len(...) / 5``. + """ + max_len = _FEATURE_MAX_LEN[feature] + n = len(sequences) + arr = _PAD_VALUE * np.ones((n, max_len, 20), dtype=np.float32) + for i, seq in enumerate(sequences): + seq = seq.upper() + if len(seq) > max_len: + raise ValueError( + "NetTCR feature %r max length is %d, got %d-mer %r" + % (feature, max_len, len(seq), seq)) + for j, aa in enumerate(seq): + try: + arr[i, j] = _BLOSUM50[aa] + except KeyError: + raise ValueError( + "Unknown amino acid %r in NetTCR feature %r sequence %r" + % (aa, feature, seq)) + arr /= _NORM_FACTOR + return arr + + +class NetTCR(object): + """Wrapper for NetTCR-2.2 pMHC:TCR binding predictions. + + Runs NetTCR-2.2's pan cross-validation ensemble in-process via a TFLite + interpreter; the reported score is the mean over the ensemble. Models are + loaded lazily on the first :meth:`predict` call and kept in memory. + + Parameters + ---------- + nettcr_path : str, optional + Path to the cloned NetTCR-2.2 repository root. If omitted, resolved + from ``NETTCR_DIR`` or ``~/NetTCR-2.2`` / ``~/code/NetTCR-2.2``. + checkpoint_dir : str, optional + Directory of ``*.tflite`` ensemble models. Defaults to the pan model + (``models/nettcr_2_2_pan/checkpoint``). Only the pan model + generalises across arbitrary peptides; the peptide-specific and + pretrained models are out of scope for this wrapper. + + Notes + ----- + NetTCR-2.2 is distributed under an academic software license; this + wrapper only *runs* a user-provided installation and vendors none of it. + + The cached ensemble interpreters are stateful, so a single ``NetTCR`` + instance is not safe to call from multiple threads concurrently; use one + instance per thread. + """ + + def __init__(self, nettcr_path=None, checkpoint_dir=None): + self.nettcr_dir = _find_nettcr_dir(nettcr_path) + if checkpoint_dir is None: + checkpoint_dir = os.path.join( + self.nettcr_dir, "models", "nettcr_2_2_pan", "checkpoint") + self.checkpoint_dir = checkpoint_dir + self._model_paths = sorted(glob.glob( + os.path.join(self.checkpoint_dir, "*.tflite"))) + if not self._model_paths: + raise FileNotFoundError( + "No NetTCR *.tflite models found in %s" % self.checkpoint_dir) + # Lazy-loaded ensemble of interpreters, and the batch size the + # interpreters are currently allocated for (so repeated same-size + # calls skip re-resizing/re-allocating tensors). + self._interpreters = None + self._allocated_n = None + + def __str__(self): + loaded = "loaded" if self._interpreters is not None else "not loaded" + return "NetTCR(models=%d, %s)" % (len(self._model_paths), loaded) + + def __repr__(self): + return str(self) + + def _predictor_name(self): + return "nettcr" + + def kind_support(self): + return { + Kind.pMHC_TCR_binding: { + # NetTCR takes no MHC allele as input; the peptide is + # implicitly class-I-restricted. + "mhc_dependence": "none", + "mhc_class": "I", + } + } + + @property + def supported_kinds(self): + return tuple(self.kind_support()) + + # ------------------------------------------------------------------ + # Model loading (lazy, cached) + # ------------------------------------------------------------------ + + def _ensure_loaded(self): + if self._interpreters is None: + self._interpreters = [ + _load_interpreter(path) for path in self._model_paths] + + # ------------------------------------------------------------------ + # Core prediction (in-process, batched over the ensemble) + # ------------------------------------------------------------------ + + def _predict_raw(self, peptides, tcrs): + """Score parallel lists of peptides and :class:`TCR` objects. + + Returns a 1-D numpy array of ensemble-mean scores, one per pair. + """ + n = len(peptides) + if n == 0: + return np.zeros(0, dtype=np.float32) + self._ensure_loaded() + + encoded = {"pep": _encode_feature(peptides, "pep")} + for key in ("a1", "a2", "a3", "b1", "b2", "b3"): + encoded[key] = _encode_feature( + [t.cdr_dict()[key] for t in tcrs], key) + + # Resize + allocate only when the batch size changed since last call. + reallocate = self._allocated_n != n + + total = np.zeros(n, dtype=np.float64) + for interpreter in self._interpreters: + inputs = interpreter.get_input_details() + output = interpreter.get_output_details()[0] + if reallocate: + for det in inputs: + interpreter.resize_tensor_input( + det["index"], [n, det["shape"][1], det["shape"][2]]) + interpreter.resize_tensor_input( + output["index"], [n, output["shape"][1]]) + interpreter.allocate_tensors() + for det in inputs: + # NetTCR names inputs "serving_default_:0"; match by + # the trailing feature token rather than by tensor order. + key = det["name"].split(":")[0].split("_")[-1] + try: + tensor = encoded[key] + except KeyError: + raise ValueError( + "NetTCR model has an unexpected input tensor %r " + "(parsed feature %r); expected one of %s" + % (det["name"], key, sorted(encoded))) + interpreter.set_tensor(det["index"], tensor) + interpreter.invoke() + total += interpreter.get_tensor(output["index"]).reshape(n) + self._allocated_n = n + return (total / len(self._interpreters)).astype(np.float32) + + # ------------------------------------------------------------------ + # Public API + # ------------------------------------------------------------------ + + def predict_pairs(self, pairs): + """Score explicit ``(peptide, TCR)`` pairs. + + Parameters + ---------- + pairs : iterable of (str, TCR) + + Returns + ------- + list of PeptideResult + One :class:`PeptideResult` per input pair (in order), each + holding a single :class:`Prediction`. + """ + pairs = list(pairs) + peptides = [pep for pep, _ in pairs] + tcrs = [tcr for _, tcr in pairs] + for tcr in tcrs: + if not isinstance(tcr, TCR): + raise TypeError( + "Expected mhctools.TCR instances, got %r" % type(tcr)) + scores = self._predict_raw(peptides, tcrs) + name = self._predictor_name() + results = [] + for pep, tcr, score in zip(peptides, tcrs, scores): + results.append(PeptideResult(preds=(Prediction( + kind=Kind.pMHC_TCR_binding, + score=float(score), + peptide=pep, + tcr=tcr.identifier, + predictor_name=name, + ),))) + return results + + def predict(self, peptides, tcrs): + """Score every ``peptide × TCR`` combination. + + Parameters + ---------- + peptides : str or list of str + tcrs : TCR or list of TCR + + Returns + ------- + list of PeptideResult + One :class:`PeptideResult` per peptide; each holds one + :class:`Prediction` per TCR. + """ + if isinstance(peptides, str): + peptides = [peptides] + if isinstance(tcrs, TCR): + tcrs = [tcrs] + tcrs = list(tcrs) + + flat_peptides = [] + flat_tcrs = [] + for pep in peptides: + for tcr in tcrs: + flat_peptides.append(pep) + flat_tcrs.append(tcr) + + flat_results = self.predict_pairs(zip(flat_peptides, flat_tcrs)) + + results = [] + idx = 0 + for _ in peptides: + preds = [] + for _ in tcrs: + preds.extend(flat_results[idx].preds) + idx += 1 + results.append(PeptideResult(preds=tuple(preds))) + return results + + def predict_dataframe(self, peptides, tcrs, sample_name=""): + """``predict()`` flattened to a DataFrame.""" + dfs = [pp.to_dataframe(sample_name) + for pp in self.predict(peptides, tcrs)] + if not dfs: + return pd.DataFrame(columns=COLUMNS) + return pd.concat(dfs, ignore_index=True) diff --git a/mhctools/pred.py b/mhctools/pred.py index 8c82f75..679afaa 100644 --- a/mhctools/pred.py +++ b/mhctools/pred.py @@ -46,6 +46,7 @@ class Kind: pMHC_affinity = "pMHC_affinity" pMHC_presentation = "pMHC_presentation" pMHC_stability = "pMHC_stability" + pMHC_TCR_binding = "pMHC_TCR_binding" immunogenicity = "immunogenicity" antigen_processing = "antigen_processing" proteasome_cleavage = "proteasome_cleavage" @@ -129,6 +130,7 @@ def best_direction(kind, field) -> str: "predictor_name", "predictor_version", "allele", + "tcr", "kind", "score", "value", @@ -143,6 +145,7 @@ class Prediction: score: float peptide: str = "" allele: str = "" + tcr: str = "" n_flank: str = "" c_flank: str = "" value: Optional[float] = None @@ -156,6 +159,8 @@ def __repr__(self): parts = [self.peptide or "?", self.kind] if self.allele: parts.insert(1, self.allele) + if self.tcr: + parts.insert(1, self.tcr) parts.append("score=%.4g" % self.score) if self.value is not None: parts.append("value=%.4g" % self.value) @@ -179,6 +184,7 @@ def to_row(self, sample_name=""): "predictor_name": self.predictor_name, "predictor_version": self.predictor_version, "allele": self.allele, + "tcr": self.tcr, "kind": self.kind, "score": self.score, "value": self.value, @@ -238,6 +244,11 @@ def alleles(self) -> set: """Set of allele strings present in this result.""" return {p.allele for p in self.preds if p.allele} + @property + def tcrs(self) -> set: + """Set of TCR identifiers present in this result.""" + return {p.tcr for p in self.preds if p.tcr} + # --- kind accessors (best by score, wrapped for safe field access) --- @property @@ -265,6 +276,11 @@ def cleavage(self) -> Optional[Prediction]: """Best proteasomal cleavage prediction, or None.""" return self.best_by_score(Kind.proteasome_cleavage) + @property + def tcr_binding(self) -> Optional[Prediction]: + """Best pMHC:TCR binding prediction, or None.""" + return self.best_by_score(Kind.pMHC_TCR_binding) + # backward compat aliases @property def best_affinity(self) -> Optional[Prediction]: diff --git a/mhctools/tcr.py b/mhctools/tcr.py new file mode 100644 index 0000000..dd97ad3 --- /dev/null +++ b/mhctools/tcr.py @@ -0,0 +1,122 @@ +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""T-cell receptor input type for pMHC:TCR binding predictors. + +mhctools' predictors have historically taken ``(peptide, allele)`` inputs +and scored MHC-ligand binding/presentation/processing. TCR specificity +predictors (e.g. NetTCR) add a new input modality — the receptor itself, +represented by its complementarity-determining regions (CDRs) — and a new +output kind (:attr:`~mhctools.pred.Kind.pMHC_TCR_binding`). + +A :class:`TCR` is a paired αβ receptor described by its six CDR loops. +The field names use the biological convention (``cdr1a`` = CDR1 of the α +chain, ``cdr3b`` = CDR3 of the β chain, ...); the ``a1``/``b3`` short +forms used by NetTCR are exposed as properties. +""" + +from __future__ import annotations + +from dataclasses import asdict, dataclass, fields + + +@dataclass(frozen=True) +class TCR: + """A paired αβ T-cell receptor, described by its six CDR loops. + + Parameters + ---------- + cdr1a, cdr2a, cdr3a : str + CDR1/2/3 amino-acid sequences of the α chain (NetTCR ``A1``/``A2``/``A3``). + cdr1b, cdr2b, cdr3b : str + CDR1/2/3 amino-acid sequences of the β chain (NetTCR ``B1``/``B2``/``B3``). + name : str, optional + Human-readable identifier (e.g. a clonotype name). If omitted, + :attr:`identifier` falls back to the CDR3α/CDR3β pair. + + Notes + ----- + CDR3β (``cdr3b``) carries most of the antigen-contact specificity, but + NetTCR-2.2's pan model expects all six loops; leaving loops empty will + encode as fully-padded (uninformative) features. + """ + + cdr1a: str = "" + cdr2a: str = "" + cdr3a: str = "" + cdr1b: str = "" + cdr2b: str = "" + cdr3b: str = "" + name: str = "" + + # NetTCR-style short aliases ------------------------------------------- + @property + def a1(self) -> str: + return self.cdr1a + + @property + def a2(self) -> str: + return self.cdr2a + + @property + def a3(self) -> str: + return self.cdr3a + + @property + def b1(self) -> str: + return self.cdr1b + + @property + def b2(self) -> str: + return self.cdr2b + + @property + def b3(self) -> str: + return self.cdr3b + + @property + def identifier(self) -> str: + """Stable string identifier for this receptor. + + Uses :attr:`name` if set, otherwise the ``CDR3α/CDR3β`` pair, which + together define the clonotype for most practical purposes. + """ + if self.name: + return self.name + return "%s/%s" % (self.cdr3a, self.cdr3b) + + def cdr_dict(self) -> dict: + """Return the six CDRs keyed by NetTCR feature name (``a1``..``b3``).""" + return { + "a1": self.cdr1a, + "a2": self.cdr2a, + "a3": self.cdr3a, + "b1": self.cdr1b, + "b2": self.cdr2b, + "b3": self.cdr3b, + } + + def __str__(self): + return "TCR(%s)" % self.identifier + + def __repr__(self): + return str(self) + + def to_dict(self): + """Serialize to a JSON-friendly dict.""" + return asdict(self) + + @classmethod + def from_dict(cls, d): + """Deserialize from a dict (as produced by :meth:`to_dict`).""" + valid = {f.name for f in fields(cls)} + return cls(**{k: v for k, v in d.items() if k in valid}) diff --git a/pyproject.toml b/pyproject.toml index 8ad3516..7dfeb16 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,6 +36,12 @@ dependencies = [ pepsickle = [ "pepsickle", ] +nettcr = [ + # A TFLite runtime for NetTCR's bundled models. ai-edge-litert is the + # lightweight option; tensorflow also works. The NetTCR-2.2 repository + # itself (weights) must be cloned separately. + "ai-edge-litert", +] dev = [ "build", "ruff", diff --git a/tests/test_nettcr.py b/tests/test_nettcr.py new file mode 100644 index 0000000..d3eb9a4 --- /dev/null +++ b/tests/test_nettcr.py @@ -0,0 +1,242 @@ +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os + +import numpy as np +import pytest + +from mhctools import TCR +from mhctools.nettcr import NetTCR, _BLOSUM50, _encode_feature +from mhctools.pred import COLUMNS, Kind + + +# --------------------------------------------------------------------------- +# Encoding unit tests — no NetTCR install or TFLite runtime required. +# These pin the exact encoding reproduced from NetTCR-2.2 src/predict.py. +# --------------------------------------------------------------------------- + +def test_encode_shape(): + arr = _encode_feature(["SIINFEKL"], "pep") + assert arr.shape == (1, 12, 20) # pep max length is 12 + assert arr.dtype == np.float32 + + +def test_encode_residues_are_blosum_over_five(): + # First peptide residue 'S' -> BLOSUM50['S'] / 5 + arr = _encode_feature(["SIINFEKL"], "pep") + np.testing.assert_allclose(arr[0, 0], _BLOSUM50["S"] / 5.0, rtol=1e-6) + + +def test_encode_padding_is_minus_one(): + # An 8-mer in a length-12 feature leaves positions 8..11 padded. + arr = _encode_feature(["SIINFEKL"], "pep") + assert np.all(arr[0, 8:] == -1.0) + + +def test_encode_is_case_insensitive(): + lower = _encode_feature(["siinfekl"], "pep") + upper = _encode_feature(["SIINFEKL"], "pep") + np.testing.assert_array_equal(lower, upper) + + +def test_encode_rejects_too_long(): + with pytest.raises(ValueError, match="max length"): + _encode_feature(["A" * 13], "pep") # pep max is 12 + + +def test_encode_rejects_unknown_amino_acid(): + with pytest.raises(ValueError, match="Unknown amino acid"): + _encode_feature(["SIINFEKB"], "pep") + + +def test_encode_batch(): + arr = _encode_feature(["SIINFEKL", "GILGFVFTL"], "pep") + assert arr.shape == (2, 12, 20) + + +# --------------------------------------------------------------------------- +# Constructor error paths — no NetTCR install required. +# --------------------------------------------------------------------------- + +def test_init_missing_path_raises(): + with pytest.raises(FileNotFoundError, match="does not exist"): + NetTCR(nettcr_path="/nonexistent/nettcr") + + +def test_init_no_models_raises(tmp_path): + # Directory exists but contains no *.tflite ensemble. + with pytest.raises(FileNotFoundError, match="No NetTCR"): + NetTCR(nettcr_path=str(tmp_path)) + + +# --------------------------------------------------------------------------- +# Model tests — require a cloned NetTCR-2.2 (weights) and a TFLite runtime. +# --------------------------------------------------------------------------- + +NETTCR_DIR = None +for candidate in [ + os.environ.get("NETTCR_DIR", ""), + os.path.join(os.path.expanduser("~"), "NetTCR-2.2"), + os.path.join(os.path.expanduser("~"), "code", "NetTCR-2.2"), +]: + if candidate and os.path.isdir( + os.path.join(candidate, "models", "nettcr_2_2_pan", "checkpoint")): + NETTCR_DIR = candidate + break + +requires_nettcr = pytest.mark.skipif( + NETTCR_DIR is None, + reason="NetTCR-2.2 not installed (set NETTCR_DIR or clone to ~/NetTCR-2.2)") + + +# Reference ensemble predictions produced by running NetTCR-2.2's OWN +# `src/predict.py` over all 20 pan-model checkpoints and averaging the +# outputs -- the canonical ensemble defined in `src/make_webserver_prediction.py` +# (`avg_prediction / 20`). These come from upstream code, not this wrapper, so +# the test is not circular. Inputs are real TCRs from NetTCR's own +# `data/nettcr_2_2_limited_dataset.csv`. Fields: (peptide, A1, A2, A3, B1, B2, B3). +PUBLISHED_ENSEMBLE = [ + (("SPRWYFYYL", "KALYS", "LLKGGEQ", "GTEIGGGTSYGKLT", "MNHEY", "SMNVEV", "ASGTETQY"), 0.060078), + (("KSKRTPMGF", "DSAIYN", "IQSSQRE", "AVRNYGGATNKLI", "PRHDT", "FYEKMQ", "ASSLTTGGRNEQF"), 0.062686), + (("AVFDRKSDAK", "VGISA", "LSSGK", "AVFNTGNQFY", "SGDLS", "YYNGEE", "ASTPWGRGTDTQY"), 0.566735), + (("RPPIFIRRL", "TTLSN", "LVKSGEV", "AGADAGNNRKLI", "SGHRS", "YFSETQ", "ASSLDQGAYEQY"), 0.011056), + (("KLGGALQAK", "DSAIYN", "IQSSQRE", "AVRPHSGGGADGLT", "SGHDY", "FNNNVP", "ASSPGDYGYT"), 0.410461), + (("GILGFVFTL", "VSGLRG", "LYSAGEE", "AVPTILTGGGNKLT", "LNHNV", "YYDKDF", "ATSRVQETQY"), 0.024195), +] +# AVFDRKSDAK is the one true binder (binder=1) in this sample. +PUBLISHED_BINDER = "AVFDRKSDAK" + + +def _row_to_pair(row): + peptide, a1, a2, a3, b1, b2, b3 = row + return peptide, TCR(cdr1a=a1, cdr2a=a2, cdr3a=a3, + cdr1b=b1, cdr2b=b2, cdr3b=b3) + + +@pytest.fixture(scope="module") +def predictor(): + return NetTCR(nettcr_path=NETTCR_DIR) + + +@requires_nettcr +def test_init_lazy(predictor): + assert predictor._interpreters is None + assert len(predictor._model_paths) > 0 + assert "not loaded" in str(predictor) + + +@requires_nettcr +def test_reproduces_published_ensemble(predictor): + """The wrapper must reproduce NetTCR's own 20-model ensemble output.""" + pairs = [_row_to_pair(row) for row, _ in PUBLISHED_ENSEMBLE] + results = predictor.predict_pairs(pairs) + for (row, expected), pp in zip(PUBLISHED_ENSEMBLE, results): + got = pp.preds[0].score + assert got == pytest.approx(expected, abs=1e-3), ( + "%s: expected %.6f (upstream ensemble), got %.6f" + % (row[0], expected, got)) + + +@requires_nettcr +def test_published_binder_ranks_top(predictor): + """The labeled binder should outscore every non-binder in the sample.""" + pairs = [_row_to_pair(row) for row, _ in PUBLISHED_ENSEMBLE] + scores = {row[0]: pp.preds[0].score + for (row, _), pp in zip(PUBLISHED_ENSEMBLE, + predictor.predict_pairs(pairs))} + binder = scores[PUBLISHED_BINDER] + others = [s for pep, s in scores.items() if pep != PUBLISHED_BINDER] + assert binder > max(others) + + +@requires_nettcr +def test_predict_pairs_kind_and_fields(predictor): + pairs = [_row_to_pair(row) for row, _ in PUBLISHED_ENSEMBLE] + for (row, _), pp in zip(PUBLISHED_ENSEMBLE, predictor.predict_pairs(pairs)): + pep, tcr = _row_to_pair(row) + pred = pp.preds[0] + assert pred.kind == Kind.pMHC_TCR_binding + assert pred.peptide == pep + assert pred.tcr == tcr.identifier + assert pred.allele == "" + assert pred.predictor_name == "nettcr" + assert pred.value is None # no native units for TCR binding + assert 0.0 <= pred.score <= 1.0 + + +@requires_nettcr +def test_predict_cross_product(predictor): + peptides = ["AVFDRKSDAK", "GILGFVFTL"] + tcrs = [_row_to_pair(PUBLISHED_ENSEMBLE[2][0])[1], + _row_to_pair(PUBLISHED_ENSEMBLE[5][0])[1]] + results = predictor.predict(peptides, tcrs) + assert len(results) == len(peptides) + for pp in results: + assert len(pp.preds) == len(tcrs) + assert pp.tcrs == {t.identifier for t in tcrs} + + +@requires_nettcr +def test_predict_single_peptide_single_tcr(predictor): + _, tcr = _row_to_pair(PUBLISHED_ENSEMBLE[2][0]) + results = predictor.predict("AVFDRKSDAK", tcr) + assert len(results) == 1 + assert len(results[0].preds) == 1 + + +@requires_nettcr +def test_batch_matches_single(predictor): + """A batched call and per-pair calls must give identical scores + (guards the allocation-caching path against batch-size bugs).""" + pairs = [_row_to_pair(row) for row, _ in PUBLISHED_ENSEMBLE] + batched = [pp.preds[0].score for pp in predictor.predict_pairs(pairs)] + singly = [predictor.predict_pairs([p])[0].preds[0].score for p in pairs] + np.testing.assert_allclose(batched, singly, atol=1e-6) + + +@requires_nettcr +def test_predict_models_stay_loaded(predictor): + predictor.predict_pairs([_row_to_pair(PUBLISHED_ENSEMBLE[0][0])]) + assert predictor._interpreters is not None + assert "loaded" in str(predictor) + + +@requires_nettcr +def test_predict_repeated_calls_consistent(predictor): + pairs = [_row_to_pair(row) for row, _ in PUBLISHED_ENSEMBLE] + r1 = predictor.predict_pairs(pairs) + r2 = predictor.predict_pairs(pairs) + for a, b in zip(r1, r2): + assert a.preds[0].score == b.preds[0].score + + +@requires_nettcr +def test_predict_empty_tcrs(predictor): + results = predictor.predict(["GILGFVFTL"], []) + assert len(results) == 1 + assert results[0].preds == () + + +@requires_nettcr +def test_predict_dataframe_schema(predictor): + _, tcr = _row_to_pair(PUBLISHED_ENSEMBLE[2][0]) + df = predictor.predict_dataframe(["AVFDRKSDAK"], [tcr]) + assert list(df.columns) == list(COLUMNS) + assert df["tcr"].iloc[0] == tcr.identifier + assert df["kind"].iloc[0] == Kind.pMHC_TCR_binding + + +@requires_nettcr +def test_bad_tcr_type_raises(predictor): + with pytest.raises(TypeError, match="TCR"): + predictor.predict_pairs([("SIINFEKL", "not-a-tcr")]) diff --git a/tests/test_tcr.py b/tests/test_tcr.py new file mode 100644 index 0000000..fd5000f --- /dev/null +++ b/tests/test_tcr.py @@ -0,0 +1,76 @@ +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for the TCR input type (no NetTCR install required).""" + +import json + +from mhctools import TCR + + +def _tcr(): + return TCR( + cdr1a="NSAFQY", cdr2a="TYSSGN", cdr3a="AMSGDGGSQGNLI", + cdr1b="LNHDA", cdr2b="SQIVND", cdr3b="ASSIRAAYEQY", + name="clone1") + + +def test_short_aliases_map_to_cdrs(): + t = _tcr() + assert t.a1 == t.cdr1a + assert t.a2 == t.cdr2a + assert t.a3 == t.cdr3a + assert t.b1 == t.cdr1b + assert t.b2 == t.cdr2b + assert t.b3 == t.cdr3b + + +def test_cdr_dict_keys_and_values(): + t = _tcr() + assert t.cdr_dict() == { + "a1": "NSAFQY", "a2": "TYSSGN", "a3": "AMSGDGGSQGNLI", + "b1": "LNHDA", "b2": "SQIVND", "b3": "ASSIRAAYEQY"} + + +def test_identifier_uses_name_when_present(): + assert _tcr().identifier == "clone1" + + +def test_identifier_falls_back_to_cdr3_pair(): + t = TCR(cdr3a="AAA", cdr3b="BBB") + assert t.identifier == "AAA/BBB" + + +def test_frozen(): + t = _tcr() + try: + t.cdr3b = "X" + assert False, "TCR should be frozen" + except AttributeError: + pass + + +def test_to_dict_round_trip(): + t = _tcr() + t2 = TCR.from_dict(t.to_dict()) + assert t == t2 + + +def test_to_dict_json_serializable(): + t = _tcr() + t2 = TCR.from_dict(json.loads(json.dumps(t.to_dict()))) + assert t == t2 + + +def test_from_dict_ignores_unknown_keys(): + t = TCR.from_dict({"cdr3b": "ASSF", "bogus": 1}) + assert t.cdr3b == "ASSF"