Consolidate predictor wrappers + BigMHC fixes + docstring cleanup

.github/workflows/tests.yml

@@ -58,6 +58,7 @@ jobs:
             tests/test_deeptap.py \
             tests/test_calis.py \
             tests/test_eramer.py \
+            tests/test_wrapper_base.py \
             tests/test_processing_predictor.py \
             tests/test_pepsickle.py \
             tests/test_bigmhc.py \

mhctools/__init__.py

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

mhctools/annotate.py

@@ -48,7 +48,7 @@
 import pandas as pd
 
 from .allele_normalization import normalize_allele_name_or_raw
-from .pred import Kind, best_direction
+from .pred import Kind, reduce_op
 
 
 # Whitespace, comma, or semicolon separate multiple alleles packed into one
@@ -140,8 +140,7 @@ def build_predictor(self, alleles):
 
     def direction_op(self):
         """``max`` or ``min`` callable for reducing candidate predictions."""
-        return max if best_direction(self.kind, self.prediction_field) == "max" \
-            else min
+        return reduce_op(self.kind, self.prediction_field)
 
 
 def parse_annotation_spec(token):
@@ -338,7 +337,7 @@ def annotate_table(
         predictor = spec.build_predictor(union_alleles or None)
         results = predictor.predict(union_peptides)
         by_pair, by_peptide = _build_lookup(results, spec)
-        reduce_op = spec.direction_op()
+        reducer = spec.direction_op()
         field = spec.prediction_field
 
         values = []
@@ -347,15 +346,13 @@ def annotate_table(
             candidates = [by_pair[(peptide, a)] for a in alleles
                           if (peptide, a) in by_pair]
             if not candidates:
-                # Allele-free predictors (e.g. processing) emit allele-less
-                # predictions indexed by peptide only. Only such predictors
-                # populate ``by_peptide``, so this fallback never masks a
-                # genuine unsupported-allele miss from a binding predictor.
+                # Allele-free predictors (e.g. processing) index by peptide
+                # only; fall back to that lookup.
                 allele_free = by_peptide.get(peptide)
                 if allele_free is not None:
                     candidates = [allele_free]
             if candidates:
-                best = reduce_op(candidates, key=lambda p: getattr(p, field))
+                best = reducer(candidates, key=lambda p: getattr(p, field))
                 values.append(getattr(best, field))
                 best_alleles.append(best.allele or None)
             else:

mhctools/bigmhc.py

@@ -32,11 +32,11 @@
 import torch
 
 from .pred import (
-    COLUMNS,
     Kind,
     PeptideResult,
     Prediction,
 )
+from .wrapper_base import NewModelPredictorMixin
 
 
 def _find_bigmhc_dir(bigmhc_path=None):
@@ -79,7 +79,7 @@ def _import_bigmhc_modules(bigmhc_dir):
             sys.path.remove(src_dir)
 
 
-class BigMHC:
+class BigMHC(NewModelPredictorMixin):
     """Wrapper for BigMHC presentation and immunogenicity predictions.
 
     Models are loaded lazily on the first call to :meth:`predict` and
@@ -130,26 +130,19 @@ def __str__(self):
         return "BigMHC(mode=%s, alleles=%s, %s)" % (
             self.mode, self.alleles, loaded)
 
-    def __repr__(self):
-        return str(self)
-
-    def _pred_kind(self):
+    def _default_pred_kind(self):
         if self.mode == "im":
             return Kind.immunogenicity
         return Kind.pMHC_presentation
 
     def kind_support(self):
         return {
-            self._pred_kind(): {
+            self._default_pred_kind(): {
                 "mhc_dependence": "single_allele",
                 "mhc_class": "I",
             }
         }
 
-    @property
-    def supported_kinds(self):
-        return tuple(self.kind_support())
-
     def _predictor_name(self):
         return "bigmhc_%s" % self.mode
 
@@ -256,8 +249,7 @@ def predict(self, peptides):
         list of PeptideResult
             One entry per peptide; each contains one Prediction per allele.
         """
-        if isinstance(peptides, str):
-            peptides = [peptides]
+        peptides = self._normalize_peptides(peptides)
 
         # Build all (peptide, allele) combinations
         all_peptides = []
@@ -269,7 +261,7 @@ def predict(self, peptides):
 
         scores = self._predict_raw(all_peptides, all_alleles)
 
-        kind = self._pred_kind()
+        kind = self._default_pred_kind()
         name = self._predictor_name()
 
         idx = 0
@@ -288,13 +280,6 @@ def predict(self, peptides):
             results.append(PeptideResult(preds=tuple(preds)))
         return results
 
-    def predict_dataframe(self, peptides, sample_name=""):
-        """``predict()`` flattened to a DataFrame."""
-        dfs = [pp.to_dataframe(sample_name) for pp in self.predict(peptides)]
-        if not dfs:
-            return pd.DataFrame(columns=COLUMNS)
-        return pd.concat(dfs, ignore_index=True)
-
 
 class BigMHC_EL(BigMHC):
     """BigMHC in presentation (eluted-ligand) mode."""

mhctools/calis.py

@@ -19,28 +19,19 @@
 importance, with the anchor positions (P1, P2, and the C-terminus) masked out
 because their identity is driven by MHC binding rather than TCR recognition.
 
-Unlike every other predictor in mhctools, this needs **no external install and
-no downloaded weights** — the ~30 published parameters are reproduced here
-directly (they come from the open-access CC-BY paper, Calis et al.,
-*PLoS Comput. Biol.* 2013, ``10.1371/journal.pcbi.1003266``). It is a fast,
-dependency-free baseline that sits alongside the learned immunogenicity
-predictors (``BigMHC_IM``, ``PRIME``).
-
-The model is **allele-independent** here: it always masks P1/P2/C-terminus, the
-default the IEDB tool uses when no allele is supplied. (The IEDB tool can swap
-in allele-specific anchor masks for a fixed ~50-allele catalogue, but those
-anchor sets almost all reduce to the P1/P2/C-terminus default anyway.)
-
-Note on interpretation: like every current CD8 immunogenicity predictor, Calis
-is useful for ranking but generalizes weakly — independent benchmarks put the
-whole field near AUC 0.5-0.65 on unseen neoepitopes. Treat scores as a
-prioritization aid, not ground truth. A score > 0 leans immunogenic, < 0 leans
-non-immunogenic.
+The ~30 parameters come from the open-access paper (Calis et al.,
+PLoS Comput. Biol. 2013, 10.1371/journal.pcbi.1003266) and are reproduced here,
+so the predictor is self-contained. It is allele-independent: it always masks
+P1/P2/C-terminus, the default the IEDB tool uses when no allele is given. A
+score > 0 leans immunogenic, < 0 non-immunogenic.
+
+Like every current CD8 immunogenicity predictor it ranks better than it
+generalizes (independent benchmarks put the field near AUC 0.5-0.65 on unseen
+neoepitopes) — a prioritization aid, not ground truth.
 """
 
-import pandas as pd
-
-from .pred import COLUMNS, Kind, PeptideResult, Prediction
+from .pred import Kind, PeptideResult, Prediction
+from .wrapper_base import AlleleFreePredictor
 
 # Per-amino-acid log-enrichment score, log(freq(aa | immunogenic) /
 # freq(aa | non-immunogenic)) from the Calis 2013 training set. Higher = more
@@ -67,7 +58,8 @@ def position_weights(peptide_length):
 
     For 9-mers this is :data:`IMMUNOWEIGHT` verbatim. For longer peptides the
     Calis rule pads the interior (after P5) with ``0.30`` weights, matching the
-    IEDB tool. Shorter peptides use the leading slice of :data:`IMMUNOWEIGHT`.
+    IEDB tool. For shorter peptides it returns ``IMMUNOWEIGHT`` unchanged and the
+    caller uses only the leading ``peptide_length`` entries.
     """
     if peptide_length > 9:
         return (
@@ -104,12 +96,12 @@ def immunogenicity_score(peptide):
     return round(score, 5)
 
 
-class Calis:
+class Calis(AlleleFreePredictor):
     """The Calis (IEDB) class-I immunogenicity predictor.
 
-    Self-contained (no external tool, no downloaded weights). Allele-independent:
-    ``predict()`` returns one :class:`~mhctools.pred.Prediction` per peptide with
-    an empty ``allele`` and ``kind == Kind.immunogenicity``.
+    Allele-independent: ``predict()`` returns one
+    :class:`~mhctools.pred.Prediction` per peptide with an empty ``allele`` and
+    ``kind == Kind.immunogenicity``.
 
     Parameters
     ----------
@@ -121,6 +113,8 @@ class Calis:
         but are outside its intended scope).
     """
 
+    mhc_class = "I"
+
     def __init__(self, min_peptide_length=8, max_peptide_length=11):
         if min_peptide_length < 3:
             # With P1/P2/C-term masked, peptides shorter than this have no
@@ -137,27 +131,9 @@ def __str__(self):
         return "Calis(min_peptide_length=%d, max_peptide_length=%d)" % (
             self.min_peptide_length, self.max_peptide_length)
 
-    def __repr__(self):
-        return str(self)
-
-    def _predictor_name(self):
-        return "calis"
-
     def _default_pred_kind(self):
         return Kind.immunogenicity
 
-    def kind_support(self):
-        return {
-            Kind.immunogenicity: {
-                "mhc_dependence": "none",
-                "mhc_class": "I",
-            },
-        }
-
-    @property
-    def supported_kinds(self):
-        return tuple(self.kind_support())
-
     def _check_peptides(self, peptides):
         for peptide in peptides:
             n = len(peptide)
@@ -182,24 +158,13 @@ def predict(self, peptides):
             ``Kind.immunogenicity`` prediction (empty ``allele``; ``score`` is
             the Calis score, higher = more immunogenic).
         """
-        if isinstance(peptides, str):
-            peptides = [peptides]
-        peptide_list = [str(p).strip().upper() for p in peptides]
+        peptide_list = self._normalize_peptides(peptides)
         self._check_peptides(peptide_list)
-
-        results = []
-        for peptide in peptide_list:
-            pred = Prediction(
+        return [
+            PeptideResult(preds=(Prediction(
                 kind=Kind.immunogenicity,
                 score=immunogenicity_score(peptide),
                 peptide=peptide,
-                predictor_name="calis")
-            results.append(PeptideResult(preds=(pred,)))
-        return results
-
-    def predict_dataframe(self, peptides, sample_name=""):
-        """``predict()`` flattened to a DataFrame."""
-        dfs = [pp.to_dataframe(sample_name) for pp in self.predict(peptides)]
-        if not dfs:
-            return pd.DataFrame(columns=COLUMNS)
-        return pd.concat(dfs, ignore_index=True)
+                predictor_name="calis"),))
+            for peptide in peptide_list
+        ]

mhctools/deeptap.py

@@ -20,25 +20,23 @@
 mhctools otherwise has no TAP predictor; DeepTAP fills that gap and emits
 ``Kind.tap_transport``.
 
-Like the cleavage predictors, DeepTAP is **allele-independent**: it scores each
+Like the cleavage predictors, DeepTAP is allele-independent: it scores each
 peptide once, so ``predict()`` returns one prediction per peptide (with an empty
 ``allele``).
 
-DeepTAP ships its pretrained weights in-repo and is Apache-2.0 licensed, but
-pins an old scientific-Python stack (``pytorch-lightning==1.9.2``, torch). To
-keep that stack out of the mhctools environment, this wrapper shells out to
-DeepTAP's own ``deeptap.py`` CLI in a user-provided checkout, run by a
-user-provided interpreter (the Tulip pattern). The checkpoints load fine under
-modern Lightning too, so any interpreter with torch + pytorch-lightning works.
+DeepTAP ships its pretrained weights in-repo (Apache-2.0) but pins an old stack
+(pytorch-lightning==1.9.2, torch). To keep that out of the mhctools environment,
+this wrapper shells out to DeepTAP's own ``deeptap.py`` CLI in a user-provided
+checkout (``DEEPTAP_HOME``), run by a user-provided interpreter
+(``DEEPTAP_PYTHON``, default the current one); the checkpoints also load under
+modern Lightning.
 
 Upstream: https://github.com/zjupgx/DeepTAP
-Cite: Chen et al., *Comput. Biol. Med.* 2023 — "DeepTAP: An RNN-based method of
+Cite: Chen et al., Comput. Biol. Med. 2023 — "DeepTAP: An RNN-based method of
 TAP-binding peptide prediction in the selection of tumor neoantigens".
 
-Note on interpretation: DeepTAP's evaluation is self-reported, and no
-independent TAP benchmark exists for any tool (true of the whole TAP field).
-Treat the score as a useful pathway signal for prioritization, not a validated
-oracle.
+DeepTAP's evaluation is self-reported and no independent TAP benchmark exists —
+a useful pathway signal, not a validated oracle.
 """
 
 import os
@@ -49,8 +47,9 @@
 import pandas as pd
 
 from .cleanup_context import CleanupFiles
-from .pred import COLUMNS, Kind, PeptideResult, Prediction
+from .pred import Kind, PeptideResult, Prediction
 from .process_helpers import run_command
+from .wrapper_base import AlleleFreePredictor
 
 # DeepTAP one-hot encodes the 20 standard amino acids plus "X" (unknown /
 # padding) and pads every peptide to a fixed width of 17; longer peptides
@@ -83,7 +82,7 @@ def _find_deeptap_home(deeptap_home=None):
     return candidate
 
 
-class DeepTAP:
+class DeepTAP(AlleleFreePredictor):
     """Wrapper for the DeepTAP TAP-transport predictor.
 
     Parameters
@@ -131,27 +130,9 @@ def __str__(self):
         return "DeepTAP(task_type=%r, deeptap_home=%r)" % (
             self.task_type, self.deeptap_home)
 
-    def __repr__(self):
-        return str(self)
-
-    def _predictor_name(self):
-        return "deeptap_%s" % self.task_type
-
     def _default_pred_kind(self):
         return Kind.tap_transport
 
-    def kind_support(self):
-        return {
-            Kind.tap_transport: {
-                "mhc_dependence": "none",
-                "mhc_class": "none",
-            },
-        }
-
-    @property
-    def supported_kinds(self):
-        return tuple(self.kind_support())
-
     def _check_peptides(self, peptides):
         for peptide in peptides:
             if not peptide:
@@ -178,9 +159,7 @@ def predict(self, peptides):
             ``"reg"`` mode the prediction's ``value`` is the predicted affinity
             in nM.
         """
-        if isinstance(peptides, str):
-            peptides = [peptides]
-        peptide_list = [str(p).strip().upper() for p in peptides]
+        peptide_list = self._normalize_peptides(peptides)
         self._check_peptides(peptide_list)
         if not peptide_list:
             return []
@@ -220,13 +199,6 @@ def predict(self, peptides):
             for peptide in peptide_list
         ]
 
-    def predict_dataframe(self, peptides, sample_name=""):
-        """``predict()`` flattened to a DataFrame."""
-        dfs = [pp.to_dataframe(sample_name) for pp in self.predict(peptides)]
-        if not dfs:
-            return pd.DataFrame(columns=COLUMNS)
-        return pd.concat(dfs, ignore_index=True)
-
 
 def parse_deeptap_results(filename, task_type="cla"):
     """Parse a DeepTAP prediction CSV into ``{peptide: [Prediction]}``.