From 826a51e2e69a32215c99b98cc931b426bf13231d Mon Sep 17 00:00:00 2001
From: Alex Rubinsteyn <alex.rubinsteyn@gmail.com>
Date: Fri, 10 Jul 2026 19:20:26 -0400
Subject: [PATCH] Fix NetMHCstabpan mixed-length peptides

---
 mhctools/__init__.py         |  2 +-
 mhctools/netmhcstabpan.py    |  9 ++------
 tests/test_netmhc_stabpan.py | 42 +++++++++++++++++++++++++++++++++++-
 3 files changed, 44 insertions(+), 9 deletions(-)

diff --git a/mhctools/__init__.py b/mhctools/__init__.py
index 76e843a..0b54d37 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.0"
+__version__ = "3.31.1"
 
 __all__ = [
     "Prediction",
diff --git a/mhctools/netmhcstabpan.py b/mhctools/netmhcstabpan.py
index a557a3c..0892659 100644
--- a/mhctools/netmhcstabpan.py
+++ b/mhctools/netmhcstabpan.py
@@ -36,13 +36,8 @@ def __init__(
             length_flag="-l",
             allele_flag="-a",
             extra_flags=flags,
-            process_limit=process_limit)
-        
-    def predict_peptides(self, peptides):
-        peptide_lengths = set(len(p) for p in peptides)
-        if len(peptide_lengths) > 1:
-            raise ValueError("All peptides must be the same length")
-        return super().predict_peptides(peptides)
+            process_limit=process_limit,
+            group_peptides_by_length=True)
 
     def _default_pred_kind(self):
         return Kind.pMHC_stability
diff --git a/tests/test_netmhc_stabpan.py b/tests/test_netmhc_stabpan.py
index 0aded42..5099bad 100644
--- a/tests/test_netmhc_stabpan.py
+++ b/tests/test_netmhc_stabpan.py
@@ -10,8 +10,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
+
 from numpy.testing import assert_allclose
 from mhctools import NetMHCstabpan
+from mhctools.base_commandline_predictor import BaseCommandlinePredictor
+from mhctools.binding_prediction_collection import BindingPredictionCollection
 
 
 DEFAULT_ALLELE = 'HLA-A*02:01'
@@ -58,4 +62,40 @@ def test_netmhc_stabpan_accuracy():
         # This could be the result of different versions of dependencies or the nature of the ANN itself.
         assert_allclose(expected, actual, atol=0.01, err_msg="Peptide %d: expected %f but got %f" % (i, expected, actual))
 
-    
\ No newline at end of file
+
+def test_netmhc_stabpan_groups_mixed_length_peptides(monkeypatch):
+    def fake_collect(self, commands, input_filenames, temp_dir_list,
+                     sequence_key_mapping=None):
+        seen_groups = []
+        for path in input_filenames:
+            with open(path) as fd:
+                seen_groups.append([line.strip() for line in fd])
+            os.remove(path)
+        for output_file in commands:
+            output_file.close()
+            os.remove(output_file.name)
+        self.seen_groups = seen_groups
+        return BindingPredictionCollection([])
+
+    monkeypatch.setattr(
+        BaseCommandlinePredictor,
+        "_determine_supported_alleles",
+        staticmethod(lambda command, flag: {"HLA-A02:01"}))
+    monkeypatch.setattr(
+        NetMHCstabpan,
+        "_run_commands_and_collect_predictions",
+        fake_collect)
+    monkeypatch.setattr(
+        NetMHCstabpan,
+        "_check_results",
+        lambda self, binding_predictions, peptides, alleles: None)
+
+    predictor = NetMHCstabpan(alleles=[DEFAULT_ALLELE])
+    predictor.predict_peptides(["SIINFEKL", "SIINFEKLL", "SIINFEKLQY"])
+
+    assert predictor.group_peptides_by_length is True
+    assert sorted(predictor.seen_groups) == sorted([
+        ["SIINFEKL"],
+        ["SIINFEKLL"],
+        ["SIINFEKLQY"],
+    ])