Add NetTCR dataframe input

mhctools/__init__.py

@@ -87,7 +87,7 @@ def __getattr__(name):
     raise AttributeError(
         "module %r has no attribute %r" % (__name__, name))
 
-__version__ = "3.31.2"
+__version__ = "3.31.3"
 
 __all__ = [
     "Prediction",

mhctools/nettcr.py

@@ -351,6 +351,44 @@ def _predict_raw(self, peptides, tcrs):
     # Public API
     # ------------------------------------------------------------------
 
+    @staticmethod
+    def _tcrs_from_dataframe(df, cdr_cols=None):
+        default_cols = {
+            "cdr1a": "cdr1a",
+            "cdr2a": "cdr2a",
+            "cdr3a": "cdr3a",
+            "cdr1b": "cdr1b",
+            "cdr2b": "cdr2b",
+            "cdr3b": "cdr3b",
+        }
+        if cdr_cols is None:
+            cdr_cols = default_cols
+        else:
+            unknown = sorted(set(cdr_cols) - set(default_cols))
+            if unknown:
+                raise ValueError(
+                    "Unknown NetTCR CDR field(s): %s"
+                    % ", ".join(unknown))
+            cdr_cols = {**default_cols, **cdr_cols}
+
+        missing = [
+            column for column in cdr_cols.values()
+            if column not in df.columns
+        ]
+        if missing:
+            raise ValueError(
+                "NetTCR dataframe missing CDR column(s): %s"
+                % ", ".join(sorted(missing)))
+
+        tcrs = []
+        for row in df.itertuples(index=False):
+            row_values = row._asdict()
+            tcrs.append(TCR.from_dict({
+                cdr: row_values[column]
+                for cdr, column in cdr_cols.items()
+            }))
+        return tcrs
+
     def predict_pairs(self, pairs):
         """Score explicit ``(peptide, TCR)`` pairs.
 
@@ -423,10 +461,37 @@ def predict(self, peptides, tcrs):
             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)]
+    def predict_dataframe(
+            self,
+            peptides,
+            tcrs=None,
+            sample_name="",
+            peptide_col="peptide",
+            cdr_cols=None):
+        """Predict flattened to a DataFrame.
+
+        Existing peptide-list calls still use ``predict(peptides, tcrs)``.
+        When *peptides* is a DataFrame, each row is treated as one
+        ``(peptide, TCR)`` pair and TCRs are built from CDR columns.
+        """
+        if isinstance(peptides, pd.DataFrame):
+            if tcrs is not None:
+                raise ValueError(
+                    "tcrs must be omitted when peptides is a DataFrame")
+            if peptide_col not in peptides.columns:
+                raise ValueError(
+                    "NetTCR dataframe missing peptide column %r"
+                    % peptide_col)
+            peptide_list = peptides[peptide_col].tolist()
+            row_tcrs = self._tcrs_from_dataframe(peptides, cdr_cols=cdr_cols)
+            dfs = [pp.to_dataframe(sample_name)
+                   for pp in self.predict_pairs(zip(peptide_list, row_tcrs))]
+        else:
+            if tcrs is None:
+                raise ValueError(
+                    "tcrs is required unless peptides is 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)

mhctools/tcr.py

@@ -29,6 +29,16 @@
 from dataclasses import asdict, dataclass, fields
 
 
+_SHORT_CDR_ALIASES = {
+    "a1": "cdr1a",
+    "a2": "cdr2a",
+    "a3": "cdr3a",
+    "b1": "cdr1b",
+    "b2": "cdr2b",
+    "b3": "cdr3b",
+}
+
+
 @dataclass(frozen=True)
 class TCR:
     """A paired αβ T-cell receptor, described by its six CDR loops.
@@ -117,6 +127,26 @@ def to_dict(self):
 
     @classmethod
     def from_dict(cls, d):
-        """Deserialize from a dict (as produced by :meth:`to_dict`)."""
+        """Deserialize from a dict.
+
+        Accepts canonical ``cdr1a``..``cdr3b`` keys plus NetTCR/IMMREP-style
+        short aliases ``a1``..``b3`` / ``A1``..``B3``. Canonical field names
+        win if both a canonical key and an alias are present.
+        """
         valid = {f.name for f in fields(cls)}
-        return cls(**{k: v for k, v in d.items() if k in valid})
+        values = {}
+        for key, value in d.items():
+            lower = str(key).lower()
+            canonical = lower if lower in valid else _SHORT_CDR_ALIASES.get(lower)
+            if canonical is None:
+                continue
+            if canonical not in values or lower in valid:
+                values[canonical] = value
+        return cls(**values)
+
+    @classmethod
+    def from_series(cls, row):
+        """Deserialize from a pandas Series or other mapping-like row."""
+        if hasattr(row, "to_dict"):
+            row = row.to_dict()
+        return cls.from_dict(row)

tests/test_nettcr.py

@@ -13,6 +13,7 @@
 import os
 
 import numpy as np
+import pandas as pd
 import pytest
 
 from mhctools import TCR
@@ -124,6 +125,138 @@ def test_init_no_models_raises(tmp_path):
     reason="NetTCR-2.2 not installed (set NETTCR_DIR or clone to ~/NetTCR-2.2)")
 
 
+class _FakeNetTCR(NetTCR):
+    def __init__(self):
+        self.calls = []
+
+    def _predict_raw(self, peptides, tcrs):
+        self.calls.append((
+            list(peptides),
+            [tcr.cdr_dict() for tcr in tcrs],
+        ))
+        return np.array(
+            [0.25 + 0.1 * i for i in range(len(peptides))],
+            dtype=np.float32)
+
+
+def _nettcr_dataframe(column_names=None):
+    values = {
+        "peptide": ["SPRWYFYYL", "AVFDRKSDAK"],
+        "cdr1a": ["KALYS", "VGISA"],
+        "cdr2a": ["LLKGGEQ", "LSSGK"],
+        "cdr3a": ["GTEIGGGTSYGKLT", "AVFNTGNQFY"],
+        "cdr1b": ["MNHEY", "SGDLS"],
+        "cdr2b": ["SMNVEV", "YYNGEE"],
+        "cdr3b": ["ASGTETQY", "ASTPWGRGTDTQY"],
+    }
+    if column_names:
+        values = {column_names.get(key, key): value
+                  for key, value in values.items()}
+    return pd.DataFrame(values)
+
+
+def test_predict_dataframe_builds_tcrs_from_canonical_columns():
+    predictor = _FakeNetTCR()
+    df = _nettcr_dataframe()
+
+    out = predictor.predict_dataframe(df, sample_name="sample1")
+
+    assert predictor.calls == [(
+        ["SPRWYFYYL", "AVFDRKSDAK"],
+        [
+            {
+                "a1": "KALYS",
+                "a2": "LLKGGEQ",
+                "a3": "GTEIGGGTSYGKLT",
+                "b1": "MNHEY",
+                "b2": "SMNVEV",
+                "b3": "ASGTETQY",
+            },
+            {
+                "a1": "VGISA",
+                "a2": "LSSGK",
+                "a3": "AVFNTGNQFY",
+                "b1": "SGDLS",
+                "b2": "YYNGEE",
+                "b3": "ASTPWGRGTDTQY",
+            },
+        ],
+    )]
+    assert list(out.columns) == list(COLUMNS)
+    assert out["sample_name"].tolist() == ["sample1", "sample1"]
+    assert out["peptide"].tolist() == ["SPRWYFYYL", "AVFDRKSDAK"]
+    assert out["tcr"].tolist() == [
+        "GTEIGGGTSYGKLT/ASGTETQY",
+        "AVFNTGNQFY/ASTPWGRGTDTQY",
+    ]
+    assert out["score"].tolist() == pytest.approx([0.25, 0.35])
+
+
+def test_predict_dataframe_accepts_custom_column_mapping():
+    predictor = _FakeNetTCR()
+    df = _nettcr_dataframe({
+        "peptide": "epitope",
+        "cdr1a": "A1",
+        "cdr2a": "A2",
+        "cdr3a": "A3",
+        "cdr1b": "B1",
+        "cdr2b": "B2",
+        "cdr3b": "B3",
+    })
+
+    out = predictor.predict_dataframe(
+        df,
+        peptide_col="epitope",
+        cdr_cols={
+            "cdr1a": "A1",
+            "cdr2a": "A2",
+            "cdr3a": "A3",
+            "cdr1b": "B1",
+            "cdr2b": "B2",
+            "cdr3b": "B3",
+        })
+
+    assert predictor.calls[0][0] == ["SPRWYFYYL", "AVFDRKSDAK"]
+    assert predictor.calls[0][1][0]["a3"] == "GTEIGGGTSYGKLT"
+    assert out["peptide"].tolist() == ["SPRWYFYYL", "AVFDRKSDAK"]
+
+
+def test_predict_dataframe_preserves_existing_cross_product_api():
+    predictor = _FakeNetTCR()
+    tcr = TCR(cdr3a="CAVR", cdr3b="CASS")
+
+    out = predictor.predict_dataframe(["SPRWYFYYL"], [tcr])
+
+    assert predictor.calls == [(
+        ["SPRWYFYYL"],
+        [{
+            "a1": "",
+            "a2": "",
+            "a3": "CAVR",
+            "b1": "",
+            "b2": "",
+            "b3": "CASS",
+        }],
+    )]
+    assert out["tcr"].tolist() == ["CAVR/CASS"]
+
+
+def test_predict_dataframe_missing_column_raises():
+    predictor = _FakeNetTCR()
+    df = _nettcr_dataframe().drop(columns=["cdr2b"])
+
+    with pytest.raises(ValueError, match="cdr2b"):
+        predictor.predict_dataframe(df)
+
+
+def test_predict_dataframe_unknown_cdr_mapping_key_raises():
+    predictor = _FakeNetTCR()
+    df = _nettcr_dataframe()
+
+    with pytest.raises(ValueError, match="cdr4a"):
+        predictor.predict_dataframe(df, cdr_cols={"cdr4a": "cdr4a"})
+
+
 # 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`

tests/test_tcr.py

@@ -14,6 +14,8 @@
 
 import json
 
+import pandas as pd
+
 from mhctools import TCR
 
 
@@ -74,3 +76,33 @@ def test_to_dict_json_serializable():
 def test_from_dict_ignores_unknown_keys():
     t = TCR.from_dict({"cdr3b": "ASSF", "bogus": 1})
     assert t.cdr3b == "ASSF"
+
+
+def test_from_dict_accepts_short_aliases_case_insensitive():
+    t = TCR.from_dict({
+        "A1": "NSAFQY",
+        "a2": "TYSSGN",
+        "A3": "AMSGDGGSQGNLI",
+        "b1": "LNHDA",
+        "B2": "SQIVND",
+        "b3": "ASSIRAAYEQY",
+    })
+    assert t.cdr1a == "NSAFQY"
+    assert t.cdr2a == "TYSSGN"
+    assert t.cdr3a == "AMSGDGGSQGNLI"
+    assert t.cdr1b == "LNHDA"
+    assert t.cdr2b == "SQIVND"
+    assert t.cdr3b == "ASSIRAAYEQY"
+
+
+def test_from_dict_canonical_keys_win_over_aliases():
+    t = TCR.from_dict({"A1": "ALIAS", "cdr1a": "CANONICAL"})
+    assert t.cdr1a == "CANONICAL"
+
+
+def test_from_series_accepts_aliases():
+    row = pd.Series({"A3": "CAVR", "B3": "CASS", "name": "clone1"})
+    t = TCR.from_series(row)
+    assert t.cdr3a == "CAVR"
+    assert t.cdr3b == "CASS"
+    assert t.identifier == "clone1"