Add pepsickle proteasomal cleavage predictor

.github/workflows/tests.yml

@@ -31,7 +31,7 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          python -m pip install -e ".[dev]"
+          python -m pip install -e ".[dev,pepsickle]"
 
       - name: Run lint checks
         run: ./lint.sh
@@ -46,6 +46,9 @@ jobs:
             tests/test_binding_prediction_collection.py \
             tests/test_cli_parsing_helpers.py \
             tests/test_mhc_formats.py \
+            tests/test_pred.py \
+            tests/test_processing_predictor.py \
+            tests/test_pepsickle.py \
             tests/test_random.py
 
   integration-netmhc:

mhctools/__init__.py

@@ -11,7 +11,20 @@
 )
 from .mixmhcpred import MixMHCpred
 from .mhcflurry import MHCflurry
+from .processing_predictor import (
+    ProcessingPredictor,
+    SCORING_MODES,
+    resolve_scoring,
+    score_cterm,
+    score_nterm_cterm,
+    score_cterm_anti_max_internal,
+    score_cterm_anti_mean_internal,
+    score_nterm_cterm_anti_max_internal,
+    score_nterm_cterm_anti_mean_internal,
+)
+from .proteasome_predictor import ProteasomePredictor
 from .netchop import NetChop
+from .pepsickle import Pepsickle
 from .netmhc import NetMHC
 from .netmhc3 import NetMHC3
 from .netmhc4 import NetMHC4
@@ -44,7 +57,18 @@
     "IedbNetMHCIIpan",
     "MixMHCpred",
     "MHCflurry",
+    "ProcessingPredictor",
+    "ProteasomePredictor",
+    "SCORING_MODES",
+    "resolve_scoring",
+    "score_cterm",
+    "score_nterm_cterm",
+    "score_cterm_anti_max_internal",
+    "score_cterm_anti_mean_internal",
+    "score_nterm_cterm_anti_max_internal",
+    "score_nterm_cterm_anti_mean_internal",
     "NetChop",
+    "Pepsickle",
     "NetMHC",
     "NetMHC3",
     "NetMHC4",

mhctools/netchop.py

@@ -10,56 +10,93 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-
-import subprocess
 import logging
+import subprocess
 import tempfile
 
+from .proteasome_predictor import ProteasomePredictor
+
+logger = logging.getLogger(__name__)
+
 
-class NetChop(object):
+class NetChop(ProteasomePredictor):
     """
-    Wrapper around netChop tool. Assumes netChop is in your PATH.
+    Wrapper around the netChop command-line tool.
+
+    Assumes ``netChop`` (or a custom *program_name*) is on your PATH.
+
+    Parameters
+    ----------
+    default_peptide_lengths : list of int, optional
+        Peptide lengths used when scanning proteins. Default ``[9]``.
+
+    scoring : callable, optional
+        See :class:`ProcessingPredictor`.  Default:
+        ``score_cterm_anti_max_internal``.
+
+    program_name : str
+        Name or path of the netChop executable (default ``"netChop"``).
     """
 
-    def predict(self, sequences):
-        """
-        Return netChop predictions for each position in each sequence.
+    def __init__(
+            self,
+            default_peptide_lengths=None,
+            scoring=None,
+            program_name="netChop"):
+        ProteasomePredictor.__init__(
+            self,
+            default_peptide_lengths=default_peptide_lengths,
+            scoring=scoring,
+        )
+        self.program_name = program_name
 
-        Parameters
-        -----------
-        sequences : list of string
-            Amino acid sequences to predict cleavage for
+    def __str__(self):
+        return "%s(program_name=%r, scoring=%s)" % (
+            self.__class__.__name__,
+            self.program_name,
+            getattr(self.scoring, "__name__", repr(self.scoring)))
 
-        Returns
-        -----------
-        list of list of float
+    def _predictor_name(self):
+        return "netchop"
 
-        The i'th list corresponds to the i'th sequence. Each list gives
-        the cleavage probability for each position in the sequence.
+    def cleavage_probs(self, sequence):
+        """
+        Run netChop on a single sequence.
+
+        Returns
+        -------
+        list of float
+            Per-position cleavage probabilities.
         """
-        with tempfile.NamedTemporaryFile(suffix=".fsa", mode="w") as input_fd:
-            for (i, sequence) in enumerate(sequences):
-                input_fd.write("> %d\n" % i)
-                input_fd.write(sequence)
-                input_fd.write("\n")
-            input_fd.flush()
+        with tempfile.NamedTemporaryFile(suffix=".fsa", mode="w") as fd:
+            fd.write("> seq\n")
+            fd.write(sequence)
+            fd.write("\n")
+            fd.flush()
             try:
-                output = subprocess.check_output(["netChop", input_fd.name])
+                output = subprocess.check_output([self.program_name, fd.name])
             except subprocess.CalledProcessError as e:
-                logging.error("Error calling netChop: %s:\n%s" % (e, e.output))
+                logger.error(
+                    "Error calling %s: %s:\n%s",
+                    self.program_name, e, e.output)
                 raise
-
         parsed = self.parse_netchop(output)
-        assert len(parsed) == len(sequences), \
-            "Expected %d results but got %d" % (
-                len(sequences), len(parsed))
-        assert [len(x) for x in parsed] == [len(x) for x in sequences]
-        return parsed
+        assert len(parsed) == 1, \
+            "Expected 1 result from netChop, got %d" % len(parsed)
+        assert len(parsed[0]) == len(sequence), \
+            "Expected %d scores, got %d" % (len(sequence), len(parsed[0]))
+        return parsed[0]
 
     @staticmethod
     def parse_netchop(netchop_output):
         """
         Parse netChop stdout.
+
+        Returns
+        -------
+        list of list of float
+            One inner list per input sequence, with per-position
+            cleavage scores.
         """
         line_iterator = iter(netchop_output.decode().split("\n"))
         scores = []

mhctools/pepsickle.py

@@ -0,0 +1,111 @@
+# 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.
+
+from .proteasome_predictor import ProteasomePredictor
+
+
+class Pepsickle(ProteasomePredictor):
+    """
+    Wrapper around the pepsickle proteasomal cleavage predictor.
+
+    Parameters
+    ----------
+    default_peptide_lengths : list of int, optional
+        Peptide lengths used when scanning proteins. Default ``[9]``.
+
+    scoring : callable, optional
+        See :class:`ProcessingPredictor`.  Default:
+        ``score_cterm_anti_max_internal``.
+
+    model_type : str
+        Pepsickle model to use. One of ``"epitope"`` (default),
+        ``"in-vitro"`` (gradient-boosted), or ``"in-vitro-2"`` (neural net).
+
+    proteasome_type : str
+        ``"C"`` for constitutive (default) or ``"I"`` for immunoproteasome.
+        Only used by in-vitro models; ignored by the epitope model.
+
+    threshold : float
+        Cleavage probability threshold used by pepsickle internally
+        (default 0.5).
+
+    human_only : bool
+        If True, use human-only trained models instead of all-mammal.
+    """
+
+    VALID_MODEL_TYPES = ("epitope", "in-vitro", "in-vitro-2")
+    VALID_PROTEASOME_TYPES = ("C", "I")
+
+    def __init__(
+            self,
+            default_peptide_lengths=None,
+            scoring=None,
+            model_type="epitope",
+            proteasome_type="C",
+            threshold=0.5,
+            human_only=False):
+        if model_type not in self.VALID_MODEL_TYPES:
+            raise ValueError(
+                "model_type must be one of %s, got %r" % (
+                    self.VALID_MODEL_TYPES, model_type))
+        if proteasome_type not in self.VALID_PROTEASOME_TYPES:
+            raise ValueError(
+                "proteasome_type must be 'C' or 'I', got %r" % proteasome_type)
+        ProteasomePredictor.__init__(
+            self,
+            default_peptide_lengths=default_peptide_lengths,
+            scoring=scoring,
+        )
+        self.model_type = model_type
+        self.proteasome_type = proteasome_type
+        self.threshold = threshold
+        self.human_only = human_only
+        self._model = None
+
+    def __str__(self):
+        return "%s(model_type=%r, proteasome_type=%r, scoring=%s)" % (
+            self.__class__.__name__,
+            self.model_type,
+            self.proteasome_type,
+            getattr(self.scoring, "__name__", repr(self.scoring)))
+
+    def _predictor_name(self):
+        return "pepsickle"
+
+    def _load_model(self):
+        if self._model is None:
+            from pepsickle.model_functions import (
+                initialize_epitope_model,
+                initialize_digestion_model,
+                initialize_digestion_gb_model,
+            )
+            if self.model_type == "epitope":
+                self._model = initialize_epitope_model(
+                    human_only=self.human_only)
+            elif self.model_type == "in-vitro":
+                self._model = initialize_digestion_gb_model()
+            elif self.model_type == "in-vitro-2":
+                self._model = initialize_digestion_model(
+                    human_only=self.human_only)
+        return self._model
+
+    def cleavage_probs(self, sequence):
+        from pepsickle.model_functions import predict_protein_cleavage_locations
+        model = self._load_model()
+        preds_raw = predict_protein_cleavage_locations(
+            sequence,
+            model,
+            mod_type=self.model_type,
+            proteasome_type=self.proteasome_type,
+            threshold=self.threshold,
+        )
+        return [entry[2] for entry in preds_raw]