From e5086a91f2080516f3b1d3535fd69a6f95a9e4ae Mon Sep 17 00:00:00 2001 From: Alex Rubinsteyn Date: Thu, 9 Jul 2026 11:42:15 -0400 Subject: [PATCH] Add TULIP-TCR pMHC:TCR binding predictor MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds a `Tulip` predictor wrapping TULIP-TCR (https://github.com/barthelemymp/TULIP-TCR), producing pMHC_TCR_binding predictions like NetTCR but also taking the presenting MHC allele as input. License / isolation: TULIP-TCR is GPLv3 and pinned to transformers==4.32.1; mhctools is Apache-2.0 and depends on neither torch nor transformers. So the wrapper vendors NONE of TULIP — no source, weights, or tokenizers. It runs a user-provided checkout out-of-process, in a separate interpreter, via TULIP's own predict.py (the same "shell out to a user-provided install" pattern used for the DTU netMHC tools and NetTCR). Nothing here imports TULIP's GPL code. Two things are supplied via constructor args or env vars: * TULIP_HOME — a TULIP-TCR checkout (predict.py, src/, tokenizers, weights) * TULIP_PYTHON — an isolated Python 3.11 interpreter with torch and transformers==4.32.1. 3.11 matters: transformers 4.32.1 resolves tokenizers 0.13.x, which has no cp312 wheel and would otherwise build from source (Rust); 3.11 has a prebuilt wheel, so the install needs no compiler. scripts/setup_tulip_env.sh builds that env (uv or venv+pip) and clones TULIP. The wrapper writes an input CSV, invokes predict.py in the checkout, and maps its per-peptide output scores back to (peptide, MHC, CDR3a, CDR3b) by position (predict.py doesn't echo the MHC column). Scores are TULIP's log-likelihood (higher = more likely to bind). Tests (tests/test_tulip.py): binary-free unit tests mock the subprocess to cover input-CSV construction, per-peptide output parsing, position/score mapping, MHC pass-through, dedup, and error propagation (added to the public CI subset); end-to-end tests run only when TULIP_HOME + TULIP_PYTHON are set. A new integration-tulip CI job builds the isolated sidecar (Python 3.11), clones the public TULIP repo, and runs the wrapper end-to-end. Full suite: 526 passed, 38 skipped, 2 xfailed (netMHCpan 4.2 + TULIP sidecar). Bump version to 3.20.0. Claude-Session: https://claude.ai/code/session_01LZahFhBSCiehXTESCYQ7wG --- .github/workflows/tests.yml | 40 ++++ README.md | 30 +++ mhctools/__init__.py | 4 +- mhctools/tulip.py | 398 ++++++++++++++++++++++++++++++++++++ scripts/setup_tulip_env.sh | 79 +++++++ tests/test_tulip.py | 277 +++++++++++++++++++++++++ 6 files changed, 827 insertions(+), 1 deletion(-) create mode 100644 mhctools/tulip.py create mode 100755 scripts/setup_tulip_env.sh create mode 100644 tests/test_tulip.py diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index f9ad91e..46ee141 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -56,6 +56,7 @@ jobs: tests/test_pepsickle.py \ tests/test_bigmhc.py \ tests/test_unparseable_alleles.py \ + tests/test_tulip.py \ tests/test_random.py integration-netmhc: @@ -109,3 +110,42 @@ jobs: mkdir -p "$NETMHC_BUNDLE_TMPDIR" export PATH="$PATH:$NETMHC_BUNDLE_HOME/bin" ./test.sh + + integration-tulip: + # Exercises the TULIP wrapper end-to-end: mhctools runs in one interpreter + # and shells out to an ISOLATED sidecar env (torch + transformers==4.32.1) + # holding a public TULIP-TCR checkout. TULIP is GPLv3; we clone (not vendor) + # it here, same as running any user-provided external tool. + runs-on: ubuntu-22.04 + steps: + - name: Checkout repository + uses: actions/checkout@v4 + + # Python 3.11 so transformers==4.32.1's tokenizers install from a + # prebuilt wheel (no cp312 wheel exists; a Rust build would be needed). + - name: Set up Python 3.11 + uses: actions/setup-python@v5 + with: + python-version: "3.11" + cache: pip + + - name: Install mhctools + run: | + python -m pip install --upgrade pip + python -m pip install -e ".[dev]" + + - name: Build isolated TULIP sidecar env + clone TULIP-TCR + env: + # This job's `python` is already 3.11; use it for the sidecar env. + TULIP_SETUP_PYTHON: python + run: | + python -m pip install uv + scripts/setup_tulip_env.sh "${{ github.workspace }}/tulip-env" \ + "${{ github.workspace }}/TULIP-TCR" + + - name: Run TULIP wrapper tests (incl. end-to-end) + env: + TULIP_HOME: ${{ github.workspace }}/TULIP-TCR + TULIP_PYTHON: ${{ github.workspace }}/tulip-env/bin/python + run: | + python -m pytest tests/test_tulip.py -ra diff --git a/README.md b/README.md index 73f2873..cd799c2 100644 --- a/README.md +++ b/README.md @@ -208,6 +208,36 @@ Examples: | `NetCleave_I` | `proteasome_cleavage` | `none` | `I` | | `NetCleave_II` | `endolysosomal_cleavage` | `none` | `II` | | `NetTCR` | `pMHC_TCR_binding` | `none` | `I` | +| `Tulip` | `pMHC_TCR_binding` | `single_allele` | `I` | + +### TCR predictors (`NetTCR`, `Tulip`) + +`NetTCR` and `Tulip` predict pMHC:TCR binding — whether a paired αβ T-cell +receptor (an `mhctools.TCR`, described by its CDR loops) recognises a peptide. +Both take `(peptide, TCR)` inputs; `Tulip` additionally takes the presenting +MHC allele. + +```python +from mhctools import Tulip, TCR + +tcr = TCR(cdr3a="CAGASGNTGKLIF", cdr3b="CASSIRASYEQYF", name="clone1") +predictor = Tulip() # needs TULIP_HOME + TULIP_PYTHON +results = predictor.predict(["GILGFVFTL"], [tcr], mhc="HLA-A*02:01") +results[0].preds[0].score # higher = more likely binding +``` + +[TULIP-TCR](https://github.com/barthelemymp/TULIP-TCR) is **GPLv3** and pinned to +`transformers==4.32.1`; mhctools is Apache-2.0 and depends on neither torch nor +transformers. The `Tulip` wrapper therefore vendors none of TULIP — it runs a +user-provided checkout out-of-process, in an isolated interpreter, via TULIP's +own `predict.py`. Set two things up first (see `scripts/setup_tulip_env.sh`, +which does both): + +- `TULIP_HOME` — a clone of TULIP-TCR (provides `predict.py`, `src/`, tokenizers, + and the released `model_weights/`); +- `TULIP_PYTHON` — an isolated **Python 3.11** interpreter with `torch` and + `transformers==4.32.1` (3.11 so `tokenizers` installs from a prebuilt wheel and + needs no Rust toolchain). For MHCflurry presentation, `presentation_allele_mode="haplotype"` treats the requested alleles as one sample genotype and emits one `pMHC_presentation` diff --git a/mhctools/__init__.py b/mhctools/__init__.py index 91aaf39..492a562 100644 --- a/mhctools/__init__.py +++ b/mhctools/__init__.py @@ -63,6 +63,7 @@ "MHCflurry": (".mhcflurry", "MHCflurry"), "MHCflurry_Affinity": (".mhcflurry", "MHCflurry_Affinity"), "NetTCR": (".nettcr", "NetTCR"), + "Tulip": (".tulip", "Tulip"), } @@ -78,7 +79,7 @@ def __getattr__(name): raise AttributeError( "module %r has no attribute %r" % (__name__, name)) -__version__ = "3.19.0" +__version__ = "3.20.0" __all__ = [ "Prediction", @@ -148,6 +149,7 @@ def __getattr__(name): "BigMHC_EL", "BigMHC_IM", "NetTCR", + "Tulip", "RandomBindingPredictor", "UnsupportedAllele", ] diff --git a/mhctools/tulip.py b/mhctools/tulip.py new file mode 100644 index 0000000..5c6dc96 --- /dev/null +++ b/mhctools/tulip.py @@ -0,0 +1,398 @@ +# 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 TULIP-TCR (https://github.com/barthelemymp/TULIP-TCR). + +TULIP predicts pMHC:TCR binding — how well a paired αβ T-cell receptor +(given by its CDR3α/CDR3β loops) recognises a peptide presented by an MHC +allele. Like :class:`~mhctools.nettcr.NetTCR` it produces the +:attr:`~mhctools.pred.Kind.pMHC_TCR_binding` kind, but unlike NetTCR it also +takes the presenting **MHC allele** as an input. + +Licensing / isolation +---------------------- +TULIP-TCR is distributed under the **GPLv3** and is coupled to a specific, +old transformers release (``transformers==4.32.1``); mhctools is Apache-2.0 +and depends on neither torch nor transformers. To keep the two apart in both +senses, this wrapper **vendors none of TULIP** — no source, no weights, no +tokenizers. It runs the user's own TULIP checkout as a black box, in a +separate interpreter, by invoking TULIP's own ``predict.py`` via subprocess +(the same "shell out to a user-provided install" pattern used for the DTU +``netMHC*`` tools and NetTCR). Nothing here imports TULIP's GPL code. + +You therefore need two things, supplied via constructor arguments or +environment variables: + +* ``TULIP_HOME`` — a clone of https://github.com/barthelemymp/TULIP-TCR + (provides ``predict.py``, ``src/``, ``aatok/``, ``mhctok/``, ``configs/`` + and the released ``model_weights/``). +* ``TULIP_PYTHON`` — a Python interpreter that can import TULIP, i.e. an + isolated environment with ``torch`` and ``transformers==4.32.1``. Because + ``tokenizers 0.13.x`` (which that transformers pins) has no cp312 wheel, + build this env on **Python 3.11** so the install uses a prebuilt wheel and + needs no Rust toolchain. See ``scripts/setup_tulip_env.sh``. + +Scores are TULIP's per-pair log-likelihood score (higher = more likely to +bind), read back from the CSVs ``predict.py`` writes. +""" + +from __future__ import annotations + +import csv +import os +import subprocess +import sys +import tempfile + +import pandas as pd + +from .pred import COLUMNS, Kind, PeptideResult, Prediction +from .tcr import TCR + + +# TULIP's missing-value token; used for an absent CDR3 chain or MHC allele. +MISSING_TOKEN = "" + +_CLONE_HINT = "Clone from https://github.com/barthelemymp/TULIP-TCR" +_ENV_HINT = ( + "Build an isolated env (Python 3.11, torch + transformers==4.32.1) and " + "point TULIP_PYTHON at its interpreter; see scripts/setup_tulip_env.sh") + + +def _find_tulip_home(tulip_home=None): + """Resolve the TULIP-TCR checkout directory. + + Checks, in order: the *tulip_home* argument, ``TULIP_HOME``, then + ``~/TULIP-TCR`` and ``~/code/TULIP-TCR``. An explicitly-provided path is + validated up front (including that it actually contains ``predict.py``) + so a typo fails clearly rather than deep inside a subprocess. + """ + for source, path in ( + ("tulip_home argument", tulip_home), + ("TULIP_HOME", os.environ.get("TULIP_HOME"))): + if path: + if not os.path.isdir(path): + raise FileNotFoundError( + "TULIP-TCR directory from %s does not exist: %s. %s" + % (source, path, _CLONE_HINT)) + if not os.path.isfile(os.path.join(path, "predict.py")): + raise FileNotFoundError( + "TULIP-TCR directory from %s has no predict.py: %s. %s" + % (source, path, _CLONE_HINT)) + return path + home = os.path.expanduser("~") + for candidate in ( + os.path.join(home, "TULIP-TCR"), + os.path.join(home, "code", "TULIP-TCR")): + if os.path.isfile(os.path.join(candidate, "predict.py")): + return candidate + raise FileNotFoundError( + "TULIP-TCR not found. Set TULIP_HOME or pass tulip_home= to the " + "constructor. %s" % _CLONE_HINT) + + +def _find_tulip_python(tulip_python=None): + """Resolve the interpreter used to run TULIP. + + Checks the *tulip_python* argument then ``TULIP_PYTHON``. Falls back to + the interpreter running mhctools only as a last resort — that works only + if torch and transformers==4.32.1 happen to be importable there, which is + exactly the coupling this wrapper exists to avoid, so it is not the + recommended path. + """ + for source, path in ( + ("tulip_python argument", tulip_python), + ("TULIP_PYTHON", os.environ.get("TULIP_PYTHON"))): + if path: + if not (os.path.isfile(path) and os.access(path, os.X_OK)): + raise FileNotFoundError( + "TULIP_PYTHON from %s is not an executable file: %s. %s" + % (source, path, _ENV_HINT)) + return path + return sys.executable + + +class Tulip(object): + """Wrapper for TULIP-TCR pMHC:TCR binding predictions. + + Parameters + ---------- + tulip_home : str, optional + Path to the cloned TULIP-TCR repository root. If omitted, resolved + from ``TULIP_HOME`` or ``~/TULIP-TCR`` / ``~/code/TULIP-TCR``. + tulip_python : str, optional + Interpreter able to import TULIP (torch + transformers==4.32.1). If + omitted, resolved from ``TULIP_PYTHON``, else the current interpreter. + model_config : str, optional + Path to the model-architecture JSON. Defaults to + ``/configs/shallow.config.json``, which matches the + released weights. + weights : str, optional + Path to the pretrained ``state_dict``. Defaults to + ``/model_weights/pytorch_model.bin``. + batch_size : int + Passed through to ``predict.py``. + + Notes + ----- + TULIP-TCR is GPLv3; this wrapper vendors none of it and only runs a + user-provided installation out-of-process. + """ + + def __init__( + self, + tulip_home=None, + tulip_python=None, + model_config=None, + weights=None, + batch_size=512): + self.tulip_home = _find_tulip_home(tulip_home) + self.tulip_python = _find_tulip_python(tulip_python) + if model_config is None: + model_config = os.path.join( + self.tulip_home, "configs", "shallow.config.json") + if weights is None: + weights = os.path.join( + self.tulip_home, "model_weights", "pytorch_model.bin") + for label, path in (("model_config", model_config), ("weights", weights)): + if not os.path.isfile(path): + raise FileNotFoundError( + "TULIP %s file not found: %s" % (label, path)) + self.model_config = model_config + self.weights = weights + self.batch_size = batch_size + + def __str__(self): + return "Tulip(home=%s, python=%s)" % (self.tulip_home, self.tulip_python) + + def __repr__(self): + return str(self) + + def _predictor_name(self): + return "tulip" + + def kind_support(self): + return { + Kind.pMHC_TCR_binding: { + # TULIP takes the presenting MHC allele as an input. + "mhc_dependence": "single_allele", + "mhc_class": "I", + } + } + + @property + def supported_kinds(self): + return tuple(self.kind_support()) + + # ------------------------------------------------------------------ + # Subprocess bridge to TULIP's own predict.py + # ------------------------------------------------------------------ + + @staticmethod + def _clean(value): + """Normalize a sequence/allele cell: strip, uppercase-safe, and map + empty values to TULIP's missing token.""" + if value is None: + return MISSING_TOKEN + text = str(value).strip() + return text if text else MISSING_TOKEN + + def _run_predict(self, rows): + """Run TULIP ``predict.py`` on *rows* and return a score dict. + + *rows* is a list of ``(peptide, mhc, cdr3a, cdr3b)`` tuples (already + cleaned). Returns ``{(peptide, mhc, cdr3a, cdr3b): score}``. Because + predict.py writes one CSV per unique peptide, in the (unshuffled) + input order of the rows carrying that peptide, we map scores back by + position within each peptide group — robust even to duplicate CDR3s. + """ + if not rows: + return {} + with tempfile.TemporaryDirectory(prefix="mhctools_tulip_") as tmp: + input_csv = os.path.join(tmp, "input.csv") + # predict.py builds output paths as ``.csv``, so + # the output prefix must end in a path separator. + output_prefix = os.path.join(tmp, "out") + os.sep + os.makedirs(output_prefix, exist_ok=True) + + with open(input_csv, "w", newline="") as fh: + writer = csv.writer(fh) + writer.writerow(["CDR3a", "CDR3b", "peptide", "MHC", "binder"]) + for peptide, mhc, cdr3a, cdr3b in rows: + writer.writerow([cdr3a, cdr3b, peptide, mhc, 1]) + + cmd = [ + self.tulip_python, + os.path.join(self.tulip_home, "predict.py"), + "--test_dir", input_csv, + "--modelconfig", self.model_config, + "--load", self.weights, + "--output", output_prefix, + "--batch_size", str(self.batch_size), + ] + proc = subprocess.run( + cmd, + cwd=self.tulip_home, # aatok/, mhctok/, src/ resolve here + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + universal_newlines=True) + if proc.returncode != 0: + raise RuntimeError( + "TULIP predict.py failed (exit %d).\nCommand: %s\n" + "stderr:\n%s" % ( + proc.returncode, " ".join(cmd), proc.stderr[-4000:])) + + # Group input rows by peptide, preserving order, to align with + # predict.py's per-peptide output ordering. + order = {} + for row in rows: + order.setdefault(row[0], []).append(row) + + scores = {} + for peptide, group in order.items(): + out_csv = os.path.join(output_prefix, peptide + ".csv") + if not os.path.isfile(out_csv): + raise RuntimeError( + "TULIP produced no output for peptide %r (expected %s)." + " predict.py stdout tail:\n%s" + % (peptide, out_csv, proc.stdout[-2000:])) + out = pd.read_csv(out_csv) + if len(out) != len(group): + raise RuntimeError( + "TULIP output row count %d != input %d for peptide %r" + % (len(out), len(group), peptide)) + for (pep, mhc, cdr3a, cdr3b), score in zip( + group, out["score"].tolist()): + scores[(pep, mhc, cdr3a, cdr3b)] = float(score) + return scores + + # ------------------------------------------------------------------ + # Public API + # ------------------------------------------------------------------ + + def predict_pairs(self, pairs): + """Score explicit ``(peptide, TCR)`` or ``(peptide, TCR, mhc)`` items. + + Parameters + ---------- + pairs : iterable + Each element is ``(peptide, TCR)`` or ``(peptide, TCR, mhc)``. + When the MHC allele is omitted it is treated as missing + (TULIP's ````), i.e. MHC-agnostic scoring. + + Returns + ------- + list of PeptideResult + One :class:`PeptideResult` per input item, in order, each holding + a single :class:`Prediction`. + """ + items = [] + for pair in pairs: + if len(pair) == 3: + peptide, tcr, mhc = pair + elif len(pair) == 2: + (peptide, tcr), mhc = pair, None + else: + raise ValueError( + "Expected (peptide, TCR) or (peptide, TCR, mhc), got %r" + % (pair,)) + if not isinstance(tcr, TCR): + raise TypeError( + "Expected mhctools.TCR instances, got %r" % type(tcr)) + items.append((peptide, tcr, mhc)) + + # Build cleaned rows and remember each item's lookup key. + rows = [] + keys = [] + for peptide, tcr, mhc in items: + key = ( + self._clean(peptide), + self._clean(mhc), + self._clean(tcr.cdr3a), + self._clean(tcr.cdr3b)) + keys.append(key) + rows.append(key) + + scores = self._run_predict(list(dict.fromkeys(rows))) + + name = self._predictor_name() + results = [] + for (peptide, tcr, mhc), key in zip(items, keys): + allele = "" if key[1] == MISSING_TOKEN else key[1] + results.append(PeptideResult(preds=(Prediction( + kind=Kind.pMHC_TCR_binding, + score=scores[key], + peptide=peptide, + allele=allele, + tcr=tcr.identifier, + predictor_name=name, + ),))) + return results + + def predict(self, peptides, tcrs, mhc=None): + """Score every ``peptide × TCR`` combination at a given MHC allele. + + Parameters + ---------- + peptides : str or list of str + tcrs : TCR or list of TCR + mhc : str or list of str, optional + Presenting MHC allele(s). A single string applies to all + peptides; a list must be parallel to *peptides*. Omit for + MHC-agnostic scoring. + + Returns + ------- + list of PeptideResult + One :class:`PeptideResult` per peptide; each holds one + :class:`Prediction` per TCR. + """ + if isinstance(peptides, str): + peptides = [peptides] + peptides = list(peptides) + if isinstance(tcrs, TCR): + tcrs = [tcrs] + tcrs = list(tcrs) + + if mhc is None or isinstance(mhc, str): + mhc_list = [mhc] * len(peptides) + else: + mhc_list = list(mhc) + if len(mhc_list) != len(peptides): + raise ValueError( + "mhc list length %d != peptides length %d" + % (len(mhc_list), len(peptides))) + + flat = [] + for peptide, allele in zip(peptides, mhc_list): + for tcr in tcrs: + flat.append((peptide, tcr, allele)) + + flat_results = self.predict_pairs(flat) + + 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, mhc=None, sample_name=""): + """``predict()`` flattened to a DataFrame.""" + dfs = [pp.to_dataframe(sample_name) + for pp in self.predict(peptides, tcrs, mhc=mhc)] + if not dfs: + return pd.DataFrame(columns=COLUMNS) + return pd.concat(dfs, ignore_index=True) diff --git a/scripts/setup_tulip_env.sh b/scripts/setup_tulip_env.sh new file mode 100755 index 0000000..1df5495 --- /dev/null +++ b/scripts/setup_tulip_env.sh @@ -0,0 +1,79 @@ +#!/usr/bin/env bash +# Build an isolated Python environment for running TULIP-TCR out-of-process, +# for use with mhctools' `Tulip` predictor (mhctools/tulip.py). +# +# Why isolated: TULIP-TCR is GPLv3 and pinned to transformers==4.32.1 (which in +# turn needs an old tokenizers). mhctools is Apache-2.0 and depends on neither +# torch nor transformers, so we never install those into your mhctools +# environment -- we build a *separate* interpreter and shell out to it. +# +# Why Python 3.11: transformers==4.32.1 resolves tokenizers 0.13.x, which has +# NO cp312 wheel (it predates Python 3.12) and would otherwise be built from +# source (requiring a Rust toolchain). Python 3.11 has a prebuilt wheel, so the +# install is wheels-only and needs no compiler. +# +# Usage: +# scripts/setup_tulip_env.sh [ENV_DIR] [TULIP_HOME] +# +# ENV_DIR where to create the venv (default: ./tulip-env) +# TULIP_HOME where to clone/find TULIP-TCR (default: ./TULIP-TCR) +# +# After it runs, export the two variables it prints, e.g.: +# export TULIP_HOME=/path/to/TULIP-TCR +# export TULIP_PYTHON=/path/to/tulip-env/bin/python + +set -euo pipefail + +ENV_DIR="${1:-./tulip-env}" +TULIP_HOME="${2:-./TULIP-TCR}" +PY311="${TULIP_SETUP_PYTHON:-python3.11}" +TULIP_REPO="https://github.com/barthelemymp/TULIP-TCR.git" + +log() { printf '[setup_tulip_env] %s\n' "$*" >&2; } + +if ! command -v "$PY311" >/dev/null 2>&1; then + log "ERROR: '$PY311' not found. Install Python 3.11 (brew install python@3.11)" + log "or set TULIP_SETUP_PYTHON to a 3.11 interpreter." + exit 1 +fi + +# Clone TULIP-TCR (GPLv3) if it isn't already present. We never redistribute +# it; you obtain it yourself under its own license. +if [[ ! -f "$TULIP_HOME/predict.py" ]]; then + log "Cloning TULIP-TCR (GPLv3) into $TULIP_HOME" + git clone --depth 1 "$TULIP_REPO" "$TULIP_HOME" +else + log "Found existing TULIP-TCR at $TULIP_HOME" +fi + +# Create the isolated env. Prefer uv (fast) but fall back to venv+pip. +DEPS=(torch "transformers==4.32.1" scikit-learn pandas numpy) +if command -v uv >/dev/null 2>&1; then + log "Creating venv with uv at $ENV_DIR (Python 3.11)" + uv venv --python "$PY311" "$ENV_DIR" + log "Installing: ${DEPS[*]}" + VIRTUAL_ENV="$ENV_DIR" uv pip install "${DEPS[@]}" +else + log "uv not found; using venv + pip at $ENV_DIR" + "$PY311" -m venv "$ENV_DIR" + "$ENV_DIR/bin/python" -m pip install --upgrade pip + log "Installing: ${DEPS[*]}" + "$ENV_DIR/bin/python" -m pip install "${DEPS[@]}" +fi + +ENV_PY="$(cd "$ENV_DIR" && pwd)/bin/python" +TULIP_ABS="$(cd "$TULIP_HOME" && pwd)" + +# Smoke test: TULIP's model code must import in the isolated interpreter. +log "Verifying TULIP imports in the isolated interpreter ..." +( cd "$TULIP_ABS" && "$ENV_PY" -c "import src.multiTrans; import torch, transformers; \ +print('ok: transformers', transformers.__version__, '| torch', torch.__version__)" ) + +cat <.csv`` in the SAME row order, assigning each row + a score equal to its 0-based position in the input CSV. That lets tests + assert the wrapper maps scores back to the right (peptide, mhc, cdr3) key + purely by position -- exactly how it aligns with real predict.py output. + """ + def fake_run(cmd, **kwargs): + test_dir = cmd[cmd.index("--test_dir") + 1] + output = cmd[cmd.index("--output") + 1] + if fail: + return subprocess.CompletedProcess(cmd, 1, "", "boom: TULIP failed") + df = pd.read_csv(test_dir).reset_index(drop=True) + df["__score"] = df.index.astype(float) + for peptide, grp in df.groupby("peptide", sort=False): + out = pd.DataFrame({ + "CDR3a": grp["CDR3a"].values, + "CDR3b": grp["CDR3b"].values, + "peptide": peptide, + "score": grp["__score"].values, + "rank": range(len(grp)), + }) + out.to_csv(output + str(peptide) + ".csv") + return subprocess.CompletedProcess(cmd, 0, "ok", "") + + monkeypatch.setattr(subprocess, "run", fake_run) + + +TCR1 = TCR(cdr3a="CAAA", cdr3b="CBBB", name="clone1") +TCR2 = TCR(cdr3a="CAAAAA", cdr3b="CBBBBB", name="clone2") + + +def test_predict_pairs_maps_scores_by_position(tmp_path, monkeypatch): + _install_fake_predict(monkeypatch) + t = _stub_tulip(tmp_path) + pairs = [ + ("GILGFVFTL", TCR1, "HLA-A*02:01"), + ("GILGFVFTL", TCR2, "HLA-A*02:01"), + ("NLVPMVATV", TCR1, "HLA-A*02:01"), + ] + results = t.predict_pairs(pairs) + # Score equals the row's position in the (deduped) input CSV: 0, 1, 2. + assert [r.preds[0].score for r in results] == [0.0, 1.0, 2.0] + for r, (pep, tcr, mhc) in zip(results, pairs): + p = r.preds[0] + assert p.kind == Kind.pMHC_TCR_binding + assert p.peptide == pep + assert p.tcr == tcr.identifier + assert p.allele == mhc + assert p.predictor_name == "tulip" + + +def test_predict_pairs_same_peptide_different_mhc_stay_distinct(tmp_path, monkeypatch): + # Same peptide + same TCR, two MHCs => two distinct keys, mapped by the + # position they occupy in the per-peptide output (MHC isn't echoed back). + _install_fake_predict(monkeypatch) + t = _stub_tulip(tmp_path) + pairs = [ + ("GILGFVFTL", TCR1, "HLA-A*02:01"), + ("GILGFVFTL", TCR1, "HLA-B*07:02"), + ] + results = t.predict_pairs(pairs) + assert [r.preds[0].score for r in results] == [0.0, 1.0] + assert [r.preds[0].allele for r in results] == ["HLA-A*02:01", "HLA-B*07:02"] + + +def test_predict_pairs_missing_mhc_gives_empty_allele(tmp_path, monkeypatch): + _install_fake_predict(monkeypatch) + t = _stub_tulip(tmp_path) + results = t.predict_pairs([("GILGFVFTL", TCR1)]) # no MHC + p = results[0].preds[0] + assert p.allele == "" # missing token surfaced as empty allele + # And the input the wrapper wrote used the missing token, not a blank. + # (verified indirectly: it produced a score, i.e. a well-formed CSV row) + assert p.score == 0.0 + + +def test_predict_pairs_dedupes_identical_items(tmp_path, monkeypatch): + # Two identical (peptide, TCR, MHC) items => one subprocess row, and both + # results get that row's score. + _install_fake_predict(monkeypatch) + t = _stub_tulip(tmp_path) + pair = ("GILGFVFTL", TCR1, "HLA-A*02:01") + results = t.predict_pairs([pair, pair]) + assert [r.preds[0].score for r in results] == [0.0, 0.0] + + +def test_predict_grid(tmp_path, monkeypatch): + _install_fake_predict(monkeypatch) + t = _stub_tulip(tmp_path) + results = t.predict(["GILGFVFTL", "NLVPMVATV"], [TCR1, TCR2], mhc="HLA-A*02:01") + assert len(results) == 2 # one PeptideResult per peptide + for r in results: + assert len(r.preds) == 2 # one Prediction per TCR + assert {p.tcr for p in r.preds} == {"clone1", "clone2"} + assert all(p.allele == "HLA-A*02:01" for p in r.preds) + + +def test_predict_grid_parallel_mhc_list(tmp_path, monkeypatch): + _install_fake_predict(monkeypatch) + t = _stub_tulip(tmp_path) + results = t.predict( + ["GILGFVFTL", "NLVPMVATV"], [TCR1], mhc=["HLA-A*02:01", "HLA-B*07:02"]) + assert results[0].preds[0].allele == "HLA-A*02:01" + assert results[1].preds[0].allele == "HLA-B*07:02" + + +def test_predict_mhc_list_length_mismatch_raises(tmp_path, monkeypatch): + _install_fake_predict(monkeypatch) + t = _stub_tulip(tmp_path) + with pytest.raises(ValueError, match="mhc list length"): + t.predict(["GILGFVFTL", "NLVPMVATV"], [TCR1], mhc=["HLA-A*02:01"]) + + +def test_predict_failure_raises_with_stderr(tmp_path, monkeypatch): + _install_fake_predict(monkeypatch, fail=True) + t = _stub_tulip(tmp_path) + with pytest.raises(RuntimeError, match="boom: TULIP failed"): + t.predict_pairs([("GILGFVFTL", TCR1, "HLA-A*02:01")]) + + +def test_bad_tcr_type_raises(tmp_path, monkeypatch): + _install_fake_predict(monkeypatch) + t = _stub_tulip(tmp_path) + with pytest.raises(TypeError, match="mhctools.TCR"): + t.predict_pairs([("GILGFVFTL", "not-a-tcr", "HLA-A*02:01")]) + + +def test_predict_dataframe_schema(tmp_path, monkeypatch): + _install_fake_predict(monkeypatch) + t = _stub_tulip(tmp_path) + df = t.predict_dataframe(["GILGFVFTL"], [TCR1, TCR2], mhc="HLA-A*02:01") + assert list(df.columns) == list(COLUMNS) + assert len(df) == 2 + assert set(df["kind"]) == {Kind.pMHC_TCR_binding} + + +# --------------------------------------------------------------------------- +# End-to-end -- requires a real TULIP checkout + isolated interpreter. +# --------------------------------------------------------------------------- + +TULIP_HOME = os.environ.get("TULIP_HOME") +TULIP_PYTHON = os.environ.get("TULIP_PYTHON") + +requires_tulip = pytest.mark.skipif( + not (TULIP_HOME and os.path.isfile(os.path.join(TULIP_HOME or "", "predict.py")) + and TULIP_PYTHON and os.access(TULIP_PYTHON or "", os.X_OK)), + reason="TULIP not installed (set TULIP_HOME + TULIP_PYTHON; " + "see scripts/setup_tulip_env.sh)") + + +# Two real TCRs from TULIP's own data/VDJ_test_2.csv; clone1 is a known +# GILGFVFTL (influenza M1) binder. +_E2E_TCR1 = TCR(cdr3a="CAGASGNTGKLIF", cdr3b="CASSIRASYEQYF", name="clone1") +_E2E_TCR2 = TCR(cdr3a="CALSGETSGSRLTF", cdr3b="CASGLVPGGLVYEQYF", name="clone2") + + +@pytest.fixture(scope="module") +def predictor(): + return Tulip() + + +@requires_tulip +def test_e2e_predict_grid(predictor): + results = predictor.predict( + ["GILGFVFTL", "NLVPMVATV"], [_E2E_TCR1, _E2E_TCR2], mhc="HLA-A*02:01") + assert len(results) == 2 + for r in results: + assert len(r.preds) == 2 + for p in r.preds: + assert p.kind == Kind.pMHC_TCR_binding + assert p.allele == "HLA-A*02:01" + assert isinstance(p.score, float) + + +@requires_tulip +def test_e2e_known_binder_scores_higher(predictor): + # For the influenza epitope GILGFVFTL, the known binder (clone1) should + # score higher than an unrelated receptor (clone2). + results = predictor.predict("GILGFVFTL", [_E2E_TCR1, _E2E_TCR2], mhc="HLA-A*02:01") + by_tcr = {p.tcr: p.score for p in results[0].preds} + assert by_tcr["clone1"] > by_tcr["clone2"] + + +@requires_tulip +def test_e2e_mhc_changes_score(predictor): + # Supplying the MHC allele should change the score vs. MHC-agnostic. + with_mhc = predictor.predict_pairs( + [("GILGFVFTL", _E2E_TCR1, "HLA-A*02:01")])[0].preds[0].score + without_mhc = predictor.predict_pairs( + [("GILGFVFTL", _E2E_TCR1)])[0].preds[0].score + assert with_mhc != without_mhc