diff --git a/mhctools/__init__.py b/mhctools/__init__.py index f90ef2c..eeddc20 100644 --- a/mhctools/__init__.py +++ b/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", diff --git a/mhctools/nettcr.py b/mhctools/nettcr.py index 1fad974..6e5f96d 100644 --- a/mhctools/nettcr.py +++ b/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) diff --git a/mhctools/tcr.py b/mhctools/tcr.py index dd97ad3..24f598d 100644 --- a/mhctools/tcr.py +++ b/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) diff --git a/tests/test_nettcr.py b/tests/test_nettcr.py index bde05b7..9cf5cf3 100644 --- a/tests/test_nettcr.py +++ b/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` diff --git a/tests/test_tcr.py b/tests/test_tcr.py index fd5000f..b4cd5cb 100644 --- a/tests/test_tcr.py +++ b/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"