diff --git a/mhctools/__init__.py b/mhctools/__init__.py index 38ab889..87f131e 100644 --- a/mhctools/__init__.py +++ b/mhctools/__init__.py @@ -78,7 +78,7 @@ def __getattr__(name): raise AttributeError( "module %r has no attribute %r" % (__name__, name)) -__version__ = "3.16.0" +__version__ = "3.17.0" __all__ = [ "Prediction", diff --git a/mhctools/base_commandline_predictor.py b/mhctools/base_commandline_predictor.py index b562d96..459a8ef 100644 --- a/mhctools/base_commandline_predictor.py +++ b/mhctools/base_commandline_predictor.py @@ -14,6 +14,7 @@ from collections import defaultdict import logging +from multiprocessing import cpu_count from subprocess import check_output import tempfile @@ -31,6 +32,17 @@ logger = logging.getLogger(__name__) +# Upper bound on how many alleles the "auto" policy will pack into one +# predictor invocation. Batching amortizes the per-process startup cost +# (~100ms for netMHCpan) across alleles, but the marginal cost of an extra +# allele in a running process is small (~15ms), so the benefit saturates +# quickly. Past this point, larger batches mostly add downside: a single +# failing allele takes out the whole batch, one process holds more output in +# memory, and there are fewer processes to spread across cores. 20 keeps +# ~85% of the amortization while bounding those risks. +AUTO_MAX_ALLELES_PER_COMMAND = 20 + + class BaseCommandlinePredictor(BasePredictor): """ Base class for MHC binding predictors that run a local external @@ -49,6 +61,7 @@ def __init__( tempdir_flag=None, extra_flags=[], max_peptides_per_file=10 ** 4, + max_alleles_per_command=1, process_limit=-1, default_peptide_lengths=[9], group_peptides_by_length=False, @@ -95,6 +108,24 @@ def __init__( max_peptides_per_file : int, optional Maximum number of lines per file when predicting peptides directly. + max_alleles_per_command : int, "auto", or None, optional + How many alleles to pass to a single invocation of the predictor + via a comma-separated allele flag (e.g. ``-a A0201,B3502``). Only + predictors whose allele flag accepts a comma-separated list (e.g. + the netMHCpan family) should batch more than one allele. + + - ``1`` (default): one allele per command, i.e. one process per + (input file, allele). Preserves the historical behavior. + - ``"auto"``: batch alleles to amortize the per-process startup + cost, while keeping enough parallel processes (input files x + allele groups) to use the available cores and capping the + group size at AUTO_MAX_ALLELES_PER_COMMAND. Speeds up + many-allele runs without pessimizing small runs on + otherwise-idle cores or over-batching into fragile mega-calls. + - ``None`` or ``<= 0``: all alleles in a single command + (unbounded batching). + - ``k > 1``: at most ``k`` alleles per command. + process_limit : int, optional Maximum number of parallel processes to start (0 for no limit, -1 for use all available processors) @@ -143,6 +174,12 @@ def __init__( "Maximum number of lines in a peptides input file") self.max_peptides_per_file = max_peptides_per_file + if max_alleles_per_command not in (None, "auto"): + require_integer( + max_alleles_per_command, + "Maximum number of alleles per command") + self.max_alleles_per_command = max_alleles_per_command + require_integer(process_limit, "Maximum number of processes") self.process_limit = process_limit @@ -227,17 +264,64 @@ def prepare_allele_name(self, allele_name): """ return allele_name.replace("*", "") + def _auto_allele_group_size(self, n_alleles, n_input_files): + """ + Group size for max_alleles_per_command="auto": batch alleles to + amortize the per-process startup cost, but (a) keep enough parallel + processes (n_input_files * groups_per_file) to occupy the cores the + command runner will use, and (b) never exceed + AUTO_MAX_ALLELES_PER_COMMAND, past which batching stops paying off. + """ + if self.process_limit and self.process_limit > 0: + target = self.process_limit + else: + target = cpu_count() + n_input_files = max(1, n_input_files) + # allele groups per file needed to reach `target` processes (ceil div) + groups_per_file = max(1, -(-target // n_input_files)) + groups_per_file = min(groups_per_file, n_alleles) + # group size = ceil(n_alleles / groups_per_file), capped + group_size = max(1, -(-n_alleles // groups_per_file)) + return min(group_size, AUTO_MAX_ALLELES_PER_COMMAND) + + def _allele_groups(self, n_input_files=1): + """ + Partition self.alleles into groups, one command (one predictor + process) per group per input file. Group size is controlled by + max_alleles_per_command (see __init__). + + Returns a list of lists of allele names. Empty if there are no + alleles. + """ + alleles = list(self.alleles) + if not alleles: + return [] + n = self.max_alleles_per_command + if n == "auto": + group_size = self._auto_allele_group_size(len(alleles), n_input_files) + elif n is None or n <= 0 or n >= len(alleles): + group_size = len(alleles) + else: + group_size = n + return [alleles[i:i + group_size] + for i in range(0, len(alleles), group_size)] + def _build_command( self, input_filename, - allele, + alleles, length=None, temp_dirname=None, peptide_mode=False): + # accept either a single allele name or a list of them + if isinstance(alleles, str): + alleles = [alleles] args = [self.program_name] if peptide_mode: args.extend(self.peptide_mode_flags) - args.extend([self.allele_flag, self.prepare_allele_name(allele)]) + allele_arg = ",".join( + self.prepare_allele_name(allele) for allele in alleles) + args.extend([self.allele_flag, allele_arg]) if length: args.extend([self.length_flag, str(length)]) if self.tempdir_flag and temp_dirname: @@ -349,7 +433,8 @@ def predict(self, peptides, n_flanks=None, c_flanks=None): dirs = [] for i, input_filename in enumerate(input_filenames): - for j, allele in enumerate(self.alleles): + for j, allele_group in enumerate( + self._allele_groups(n_input_files=len(input_filenames))): if self.tempdir_flag: temp_dirname = tempfile.mkdtemp( prefix="tmp_%d_%d_%s" % (i, j, self.program_name), @@ -364,7 +449,7 @@ def predict(self, peptides, n_flanks=None, c_flanks=None): delete=False) commands[output_file] = self._build_command( input_filename=input_filename, - allele=allele, + alleles=allele_group, peptide_mode=True, temp_dirname=temp_dirname) return self._run_commands_and_collect_preds( @@ -383,7 +468,8 @@ def predict_peptides(self, peptides): dirs = [] for i, input_filename in enumerate(input_filenames): - for j, allele in enumerate(self.alleles): + for j, allele_group in enumerate( + self._allele_groups(n_input_files=len(input_filenames))): if self.tempdir_flag: temp_dirname = tempfile.mkdtemp( prefix="tmp_%d_%d_%s" % ( @@ -392,9 +478,9 @@ def predict_peptides(self, peptides): self.program_name), suffix="XXXXXX") logger.debug( - "Created temporary directory %s for allele %s", + "Created temporary directory %s for alleles %s", temp_dirname, - allele) + allele_group) dirs.append(temp_dirname) else: temp_dirname = None @@ -405,7 +491,7 @@ def predict_peptides(self, peptides): delete=False) commands[output_file] = self._build_command( input_filename=input_filename, - allele=allele, + alleles=allele_group, peptide_mode=True, temp_dirname=temp_dirname) results = self._run_commands_and_collect_predictions( diff --git a/mhctools/netmhc_pan.py b/mhctools/netmhc_pan.py index e82a673..4dcfb14 100644 --- a/mhctools/netmhc_pan.py +++ b/mhctools/netmhc_pan.py @@ -51,7 +51,9 @@ def NetMHCpan( program_name="netMHCpan", process_limit=-1, default_peptide_lengths=[9], - extra_flags=[]): + extra_flags=[], + max_peptides_per_file=10 ** 4, + max_alleles_per_command="auto"): """ Auto-detecting wrapper for any installed version of NetMHCpan. @@ -79,6 +81,8 @@ def NetMHCpan( "program_name": program_name, "process_limit": process_limit, "extra_flags": extra_flags, + "max_peptides_per_file": max_peptides_per_file, + "max_alleles_per_command": max_alleles_per_command, } # Exact match diff --git a/mhctools/netmhc_pan28.py b/mhctools/netmhc_pan28.py index 4c453ca..effe272 100644 --- a/mhctools/netmhc_pan28.py +++ b/mhctools/netmhc_pan28.py @@ -20,7 +20,9 @@ def __init__( default_peptide_lengths=[9], program_name="netMHCpan", process_limit=-1, - extra_flags=[]): + extra_flags=[], + max_peptides_per_file=10 ** 4, + max_alleles_per_command="auto"): BaseCommandlinePredictor.__init__( self, program_name=program_name, @@ -34,4 +36,6 @@ def __init__( length_flag="-l", allele_flag="-a", extra_flags=extra_flags, + max_peptides_per_file=max_peptides_per_file, + max_alleles_per_command=max_alleles_per_command, process_limit=process_limit) diff --git a/mhctools/netmhc_pan3.py b/mhctools/netmhc_pan3.py index f8dc527..1276f45 100644 --- a/mhctools/netmhc_pan3.py +++ b/mhctools/netmhc_pan3.py @@ -21,7 +21,9 @@ def __init__( default_peptide_lengths=[9], program_name="netMHCpan", process_limit=-1, - extra_flags=[]): + extra_flags=[], + max_peptides_per_file=10 ** 4, + max_alleles_per_command="auto"): BaseCommandlinePredictor.__init__( self, program_name=program_name, @@ -34,4 +36,6 @@ def __init__( length_flag="-l", allele_flag="-a", extra_flags=extra_flags, + max_peptides_per_file=max_peptides_per_file, + max_alleles_per_command=max_alleles_per_command, process_limit=process_limit) diff --git a/mhctools/netmhc_pan4.py b/mhctools/netmhc_pan4.py index af2b409..27f2961 100644 --- a/mhctools/netmhc_pan4.py +++ b/mhctools/netmhc_pan4.py @@ -25,7 +25,9 @@ def __init__( program_name="netMHCpan", process_limit=-1, mode="binding_affinity", - extra_flags=[]): + extra_flags=[], + max_peptides_per_file=10 ** 4, + max_alleles_per_command="auto"): """ Wrapper for NetMHCpan4. @@ -54,6 +56,8 @@ def __init__( length_flag="-l", allele_flag="-a", extra_flags=flags + extra_flags, + max_peptides_per_file=max_peptides_per_file, + max_alleles_per_command=max_alleles_per_command, process_limit=process_limit) def kind_support(self): diff --git a/mhctools/netmhc_pan41.py b/mhctools/netmhc_pan41.py index 099ce8f..d232b22 100644 --- a/mhctools/netmhc_pan41.py +++ b/mhctools/netmhc_pan41.py @@ -25,7 +25,9 @@ def __init__( program_name="netMHCpan", process_limit=-1, mode="binding_affinity", - extra_flags=[]): + extra_flags=[], + max_peptides_per_file=10 ** 4, + max_alleles_per_command="auto"): """ Wrapper for NetMHCpan4.1. @@ -54,6 +56,8 @@ def __init__( length_flag="-l", allele_flag="-a", extra_flags=flags + extra_flags, + max_peptides_per_file=max_peptides_per_file, + max_alleles_per_command=max_alleles_per_command, process_limit=process_limit) def kind_support(self): diff --git a/mhctools/netmhc_pan42.py b/mhctools/netmhc_pan42.py index 2852c17..1b6bd2e 100644 --- a/mhctools/netmhc_pan42.py +++ b/mhctools/netmhc_pan42.py @@ -25,7 +25,9 @@ def __init__( program_name="netMHCpan", process_limit=-1, mode="binding_affinity", - extra_flags=[]): + extra_flags=[], + max_peptides_per_file=10 ** 4, + max_alleles_per_command="auto"): """ Wrapper for NetMHCpan 4.2. @@ -55,6 +57,8 @@ def __init__( length_flag="-l", allele_flag="-a", extra_flags=flags + extra_flags, + max_peptides_per_file=max_peptides_per_file, + max_alleles_per_command=max_alleles_per_command, process_limit=process_limit) def kind_support(self): diff --git a/tests/test_allele_batching.py b/tests/test_allele_batching.py new file mode 100644 index 0000000..bc6ce66 --- /dev/null +++ b/tests/test_allele_batching.py @@ -0,0 +1,326 @@ +# 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. + +""" +Unit tests for allele batching in BaseCommandlinePredictor. + +These exercise the command-construction logic (which alleles end up in which +process) WITHOUT running any external predictor binary, so they run on every +platform. End-to-end validation against the real netMHCpan output lives in +tests/test_netmhc_pan.py (which needs the binary) and the multi-allele parser +tests in tests/test_mhc_formats.py. +""" + +import inspect +import os + +import pytest + +from mhctools import ( + NetMHCpan, + NetMHCpan28, + NetMHCpan3, + NetMHCpan4, + NetMHCpan41, + NetMHCpan42, +) +from mhctools.base_commandline_predictor import BaseCommandlinePredictor +from mhctools.binding_prediction_collection import BindingPredictionCollection + + +class _StubPredictor(BaseCommandlinePredictor): + """ + Exercises the pure command-construction logic without a real binary. + + Bypasses BaseCommandlinePredictor.__init__ (which shells out to netMHCpan + to list supported alleles) and sets only the attributes that + _allele_groups(), _build_command() and predict_peptides() read. + """ + def __init__( + self, + alleles, + max_alleles_per_command=1, + prepare=None, + max_peptides_per_file=10 ** 4, + process_limit=-1): + self.program_name = "netMHCpan" + self.peptide_mode_flags = ["-p"] + self.allele_flag = "-a" + self.length_flag = "-l" + self.input_file_flag = "-f" + self.tempdir_flag = None + self.extra_flags = [] + self.alleles = list(alleles) + self.max_alleles_per_command = max_alleles_per_command + self.max_peptides_per_file = max_peptides_per_file + self.process_limit = process_limit + self.group_peptides_by_length = False + # peptide-validation attributes touched by _check_peptide_inputs + self.allow_X_in_peptides = False + self.allow_lowercase_in_peptides = False + self.min_peptide_length = None + self.max_peptide_length = None + self.parse_output_fn = None + self.parse_to_preds_fn = None + self._prepare = prepare + self.captured_commands = None + + def prepare_allele_name(self, allele_name): + if self._prepare is not None: + return self._prepare(allele_name) + return super().prepare_allele_name(allele_name) + + def _run_commands_and_collect_predictions( + self, commands, input_filenames, temp_dir_list, + sequence_key_mapping=None): + # Capture the commands instead of running them, and clean up the temp + # files predict_peptides() created. + self.captured_commands = dict(commands) + for output_file in commands: + output_file.close() + _silent_remove(output_file.name) + for fname in input_filenames: + _silent_remove(fname) + return BindingPredictionCollection([]) + + def _check_results(self, binding_predictions, peptides, alleles): + pass # result completeness is not what these tests check + + +def _silent_remove(path): + try: + os.remove(path) + except OSError: + pass + + +def _allele_arg(command): + """The single value passed after the -a flag in a built command.""" + assert command.count("-a") == 1, \ + "expected exactly one -a flag, got: %s" % (command,) + return command[command.index("-a") + 1] + + +# --------------------------------------------------------------------------- +# _allele_groups(): how self.alleles is partitioned into per-process groups +# --------------------------------------------------------------------------- + +def test_allele_groups_one_per_command_is_default_behavior(): + p = _StubPredictor(["A", "B", "C"], max_alleles_per_command=1) + assert p._allele_groups() == [["A"], ["B"], ["C"]] + + +def test_allele_groups_none_batches_all(): + p = _StubPredictor(["A", "B", "C"], max_alleles_per_command=None) + assert p._allele_groups() == [["A", "B", "C"]] + + +def test_allele_groups_chunk_size(): + p = _StubPredictor(["A", "B", "C", "D", "E"], max_alleles_per_command=2) + assert p._allele_groups() == [["A", "B"], ["C", "D"], ["E"]] + + +def test_allele_groups_non_positive_batches_all(): + p = _StubPredictor(["A", "B", "C"], max_alleles_per_command=0) + assert p._allele_groups() == [["A", "B", "C"]] + + +def test_allele_groups_larger_than_alleles_is_single_group(): + p = _StubPredictor(["A", "B"], max_alleles_per_command=10) + assert p._allele_groups() == [["A", "B"]] + + +def test_allele_groups_empty(): + assert _StubPredictor([], max_alleles_per_command=None)._allele_groups() == [] + assert _StubPredictor([], max_alleles_per_command=1)._allele_groups() == [] + + +# --------------------------------------------------------------------------- +# "auto": batch alleles but keep enough processes to use the cores +# --------------------------------------------------------------------------- + +def test_auto_batches_all_when_files_already_saturate_cores(): + # 8 input files already reach the 4-process target, so all alleles can + # share one command per file (fewest network reloads). + p = _StubPredictor( + ["A", "B", "C", "D", "E"], + max_alleles_per_command="auto", + process_limit=4) + groups = p._allele_groups(n_input_files=8) + assert groups == [["A", "B", "C", "D", "E"]] + + +def test_auto_keeps_parallelism_with_a_single_input_file(): + # 1 input file, 4-process target -> need 4 allele groups to keep cores + # busy, so a single file must not collapse to one process. + p = _StubPredictor( + ["A", "B", "C", "D"], + max_alleles_per_command="auto", + process_limit=4) + groups = p._allele_groups(n_input_files=1) + assert len(groups) == 4 + assert [a for g in groups for a in g] == ["A", "B", "C", "D"] + + +def test_auto_never_more_groups_than_alleles(): + # target far exceeds allele count: at most one allele per group. + p = _StubPredictor( + ["A", "B"], max_alleles_per_command="auto", process_limit=32) + groups = p._allele_groups(n_input_files=1) + assert groups == [["A"], ["B"]] + + +def test_auto_partial_batching_with_few_files(): + # 2 files, target 8 -> 4 allele groups per file; 8 alleles -> size-2 groups. + p = _StubPredictor( + ["A1", "A2", "A3", "A4", "A5", "A6", "A7", "A8"], + max_alleles_per_command="auto", + process_limit=8) + groups = p._allele_groups(n_input_files=2) + assert groups == [ + ["A1", "A2"], ["A3", "A4"], ["A5", "A6"], ["A7", "A8"]] + + +def test_auto_caps_group_size(): + # Many files would let "auto" batch every allele into one call; the cap + # keeps groups bounded so a single call never grows unbounded. + from mhctools.base_commandline_predictor import ( + AUTO_MAX_ALLELES_PER_COMMAND as CAP) + n_alleles = CAP * 3 + 5 + p = _StubPredictor( + ["A%d" % i for i in range(n_alleles)], + max_alleles_per_command="auto", + process_limit=2) + # 100 input files easily saturate the 2-process target, so absent a cap + # auto would put all alleles in one group. + groups = p._allele_groups(n_input_files=100) + assert max(len(g) for g in groups) <= CAP + assert sum(len(g) for g in groups) == n_alleles + + +def test_none_still_batches_all_without_cap(): + # The explicit "unbounded" escape hatch is not subject to the auto cap. + from mhctools.base_commandline_predictor import ( + AUTO_MAX_ALLELES_PER_COMMAND as CAP) + n_alleles = CAP * 3 + p = _StubPredictor( + ["A%d" % i for i in range(n_alleles)], max_alleles_per_command=None) + groups = p._allele_groups(n_input_files=100) + assert len(groups) == 1 + assert len(groups[0]) == n_alleles + + +# --------------------------------------------------------------------------- +# _build_command(): a group becomes one comma-separated -a argument +# --------------------------------------------------------------------------- + +def test_build_command_joins_alleles_into_single_flag(): + p = _StubPredictor(["HLA-A*02:01", "HLA-B*35:02"]) + cmd = p._build_command( + input_filename="peptides.txt", + alleles=["HLA-A*02:01", "HLA-B*35:02"], + peptide_mode=True) + # default prepare_allele_name strips the '*' + assert _allele_arg(cmd) == "HLA-A02:01,HLA-B35:02" + assert cmd[:2] == ["netMHCpan", "-p"] + assert cmd[-2:] == ["-f", "peptides.txt"] + + +def test_build_command_single_allele_string_still_works(): + p = _StubPredictor(["HLA-A*02:01"]) + cmd = p._build_command( + input_filename="peptides.txt", + alleles="HLA-A*02:01", + peptide_mode=True) + assert _allele_arg(cmd) == "HLA-A02:01" + + +def test_build_command_applies_prepare_per_allele_then_joins(): + # netMHC4-style prepare that also strips ':'; must be applied to each + # allele individually before joining, not to the joined string. + def prepare(a): + return a.replace("*", "").replace(":", "") + p = _StubPredictor(["HLA-A*02:01", "HLA-B*35:02"], prepare=prepare) + cmd = p._build_command( + input_filename="peptides.txt", + alleles=["HLA-A*02:01", "HLA-B*35:02"], + peptide_mode=True) + assert _allele_arg(cmd) == "HLA-A0201,HLA-B3502" + + +# --------------------------------------------------------------------------- +# predict_peptides() wiring: number of processes and their -a arguments +# --------------------------------------------------------------------------- + +def test_predict_peptides_batches_all_alleles_into_one_process(): + alleles = ["HLA-A*02:01", "HLA-B*35:02", "HLA-C*07:02"] + p = _StubPredictor(alleles, max_alleles_per_command=None) + p.predict_peptides(["SIINFEKLL", "GILGFVFTL"]) + commands = list(p.captured_commands.values()) + # one input file (2 peptides) x one allele group = one process + assert len(commands) == 1 + got = set(_allele_arg(commands[0]).split(",")) + assert got == {"HLA-A02:01", "HLA-B35:02", "HLA-C07:02"} + + +def test_predict_peptides_one_process_per_allele_when_not_batched(): + alleles = ["HLA-A*02:01", "HLA-B*35:02", "HLA-C*07:02"] + p = _StubPredictor(alleles, max_alleles_per_command=1) + p.predict_peptides(["SIINFEKLL", "GILGFVFTL"]) + commands = list(p.captured_commands.values()) + # one input file x three alleles = three processes, each single-allele + assert len(commands) == 3 + per_command = [_allele_arg(c) for c in commands] + assert all("," not in a for a in per_command) + assert set(per_command) == {"HLA-A02:01", "HLA-B35:02", "HLA-C07:02"} + + +def test_predict_peptides_command_count_is_files_times_groups(): + alleles = ["HLA-A*02:01", "HLA-B*35:02"] + peptides = ["AAAAAAAAA", "CCCCCCCCC", "DDDDDDDDD", "EEEEEEEEE"] + # max_peptides_per_file=2 -> 4 peptides split into 2 input files + batched = _StubPredictor( + alleles, max_alleles_per_command=None, max_peptides_per_file=2) + batched.predict_peptides(peptides) + assert len(batched.captured_commands) == 2 # 2 files x 1 group + + per_allele = _StubPredictor( + alleles, max_alleles_per_command=1, max_peptides_per_file=2) + per_allele.predict_peptides(peptides) + assert len(per_allele.captured_commands) == 4 # 2 files x 2 alleles + + +# --------------------------------------------------------------------------- +# The netMHCpan family batches all alleles by default (perf win out of the box) +# --------------------------------------------------------------------------- + +@pytest.mark.parametrize("factory", [ + NetMHCpan, NetMHCpan28, NetMHCpan3, NetMHCpan4, NetMHCpan41, NetMHCpan42, +]) +def test_pan_wrappers_default_to_auto_batching(factory): + default = inspect.signature(factory).parameters[ + "max_alleles_per_command"].default + assert default == "auto" + + +@pytest.mark.parametrize("factory", [ + NetMHCpan, NetMHCpan28, NetMHCpan3, NetMHCpan4, NetMHCpan41, NetMHCpan42, +]) +def test_pan_wrappers_expose_max_peptides_per_file(factory): + assert "max_peptides_per_file" in inspect.signature(factory).parameters + + +def test_base_predictor_defaults_to_one_allele_per_command(): + # Non-pan command-line predictors keep the historical behavior. + default = inspect.signature( + BaseCommandlinePredictor).parameters["max_alleles_per_command"].default + assert default == 1 diff --git a/tests/test_netmhc_pan.py b/tests/test_netmhc_pan.py index 8cbc60f..92f7a21 100644 --- a/tests/test_netmhc_pan.py +++ b/tests/test_netmhc_pan.py @@ -86,4 +86,38 @@ def test_netmhc_pan_multiple_alleles(): binding_predictions = predictor.predict_subsequences( sequence_dict=sequence_dict) assert len(binding_predictions) == 8, \ - "Expected 4 binding predictions from %s" % (binding_predictions,) + "Expected 8 binding predictions from %s" % (binding_predictions,) + # With allele batching both alleles are produced by a single netMHCpan + # invocation (-a A02:01,B35:02); make sure both are attributed correctly. + observed_alleles = {bp.allele for bp in binding_predictions} + assert observed_alleles == {"HLA-A*02:01", "HLA-B*35:02"}, \ + "Expected both alleles, got %s" % (observed_alleles,) + + +def test_netmhc_pan_batched_matches_per_allele(): + """Batching alleles into one `-a A,B,C` invocation must produce exactly + the same scores as running one process per allele. Uses the real binary's + own output on both paths as non-circular ground truth.""" + alleles = ["HLA-A*02:01", "HLA-B*07:02", "HLA-C*07:02", "HLA-A*01:01"] + peptides = [ + "SIINFEKLL", "GILGFVFTL", "NLVPMVATV", "LLWNGPMAV", "AAAWYLWEV"] + + def scores(max_alleles_per_command): + predictor = NetMHCpan( + alleles=alleles, + max_alleles_per_command=max_alleles_per_command) + return { + (bp.allele, bp.peptide): bp.value + for bp in predictor.predict_peptides(peptides) + } + + batched = scores(None) # all alleles in a single process + per_allele = scores(1) # one process per allele + assert set(batched) == set(per_allele), ( + "Batched and per-allele runs produced different (allele, peptide) " + "pairs") + assert len(batched) == len(alleles) * len(peptides) + for key, value in per_allele.items(): + eq_(value, batched[key], + "Score mismatch for %s: per-allele=%s batched=%s" % ( + key, value, batched[key]))