Newer
Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
# Copyright (c) 2016. Mount Sinai School of Medicine
#
# 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.
'''
Run MHCflurry predictor on specified peptide/allele pairs.
Examples:
Write a CSV file containing the contents of INPUT.csv plus an
additional column giving MHCflurry binding affinity predictions:
mhcflurry-predict INPUT.csv --out RESULT.csv
The input CSV file is expected to contain columns 'allele' and 'peptide'.
The predictions are written to a column called 'mhcflurry_prediction'.
These default column names may be changed with the --allele-column,
--peptide-column, and --prediction-column options.
If --out is not specified, results are writtent to standard out.
You can also run on alleles and peptides specified on the commandline, in
which case predictions are written for all combinations of alleles and
peptides:
mhcflurry-predict --alleles HLA-A0201 H-2Kb --peptides SIINFEKL DENDREKLLL
'''
from __future__ import (
print_function,
division,
absolute_import,
)
import sys
import argparse
import itertools
from .class1_affinity_predictor import Class1AffinityPredictor
parser = argparse.ArgumentParser(
description=__doc__,
formatter_class=argparse.RawDescriptionHelpFormatter,
add_help=False)
helper_args = parser.add_argument_group(title="Help")
helper_args.add_argument(
"-h", "--help",
action="help",
help="Show this help message and exit"
)
helper_args.add_argument(
"--list-supported-alleles",
action="store_true",
default=False,
help="Prints the list of supported alleles and exits"
)
helper_args.add_argument(
"--list-supported-peptide-lengths",
action="store_true",
default=False,
help="Prints the list of supported peptide lengths and exits"
)
input_args = parser.add_argument_group(title="Required input arguments")
input_args.add_argument(
"--alleles",
metavar="ALLELE",
nargs="+",
help="Alleles to predict (exclusive with --input)")
"--peptides",
metavar="PEPTIDE",
nargs="+",
help="Peptides to predict (exclusive with --input)")
input_mod_args = parser.add_argument_group(title="Optional input modifiers")
input_mod_args.add_argument(
"--allele-column",
metavar="NAME",
default="allele",
help="Input column name for alleles. Default: '%(default)s'")
"--peptide-column",
metavar="NAME",
default="peptide",
help="Input column name for peptides. Default: '%(default)s'")
input_mod_args.add_argument(
"--no-throw",
action="store_true",
default=False,
help="Return NaNs for unsupported alleles or peptides instead of raising")
output_args = parser.add_argument_group(title="Optional output modifiers")
output_args.add_argument(
"--out",
metavar="OUTPUT.csv",
help="Output CSV")
output_args.add_argument(
default="mhcflurry_",
help="Prefix for output column names. Default: '%(default)s'")
model_args = parser.add_argument_group(title="Optional model settings")
model_args.add_argument(
"Default: %s" % get_path("models_class1", "models", test_exists=False))
"--include-individual-model-predictions",
action="store_true",
default=False,
help="Include predictions from each model in the ensemble"
)
def run(argv=sys.argv[1:]):
args = parser.parse_args(argv)
models_dir = args.models
if models_dir is None:
# The reason we set the default here instead of in the argument parser is that
# we want to test_exists at this point, so the user gets a message instructing
# them to download the models if needed.
models_dir = get_path("models_class1", "models")
predictor = Class1AffinityPredictor.load(models_dir)
# The following two are informative commands that can come
# if a wrapper would like to incorporate input validation
# to not delibaretly make mhcflurry fail
if args.list_supported_alleles:
print("\n".join(predictor.supported_alleles))
return
if args.list_supported_peptide_lengths:
min_len, max_len = predictor.supported_peptide_lengths
print("\n".join([str(l) for l in range(min_len, max_len+1)]))
return
# End of early terminating routines
if args.input:
if args.alleles or args.peptides:
parser.error(
"If an input file is specified, do not specify --alleles "
"or --peptides")
df = pandas.read_csv(args.input)
print("Read input CSV with %d rows, columns are: %s" % (
len(df), ", ".join(df.columns)))
for col in [args.allele_column, args.peptide_column]:
if col not in df.columns:
raise ValueError(
"No such column '%s' in CSV. Columns are: %s" % (
col, ", ".join(["'%s'" % c for c in df.columns])))
else:
if not args.alleles or not args.peptides:
parser.error(
"Specify either an input CSV file or both the "
"--alleles and --peptides arguments")
Alex Rubinsteyn
committed
# split user specified allele and peptide strings in case they
# contain multiple entries separated by commas
alleles = []
for allele_string in args.alleles:
alleles.extend([s.strip() for s in allele_string.split(",")])
peptides = []
for peptide in args.peptides:
peptides.extend(peptide.strip() for p in peptide.split(","))
for peptide in peptides:
if not peptide.isalpha():
raise ValueError(
"Unexpected character(s) in peptide '%s'" % peptide)
pairs = list(itertools.product(alleles, peptides))
df = pandas.DataFrame({
"allele": [p[0] for p in pairs],
"peptide": [p[1] for p in pairs],
})
logging.info(
"Predicting for %d alleles and %d peptides = %d predictions" % (
len(args.alleles), len(args.peptides), len(df)))
predictions = predictor.predict_to_dataframe(
alleles=df[args.allele_column].values,
include_individual_model_predictions=args.include_individual_model_predictions,
throw=not args.no_throw)
for col in predictions.columns:
if col not in ("allele", "peptide"):
df[args.prediction_column_prefix + col] = predictions[col]
if args.out:
df.to_csv(args.out, index=False)
print("Wrote: %s" % args.out)
else:
df.to_csv(sys.stdout, index=False)