diff --git a/mhcflurry/predict_command.py b/mhcflurry/predict_command.py index f468b16807570e1b311926f6a13477d69d950224..cf23a02b9154506893c90b7e0fb7e7a40f1557cd 100644 --- a/mhcflurry/predict_command.py +++ b/mhcflurry/predict_command.py @@ -278,7 +278,7 @@ def run(argv=sys.argv[1:]): del predictions["best_allele"] for col in predictions.columns: - if col not in ("allele", "peptide", "experiment_name"): + if col not in ("allele", "peptide", "experiment_name", "peptide_num"): df[args.prediction_column_prefix + col] = predictions[col] if args.out: diff --git a/test/test_predict_command.py b/test/test_predict_command.py index a4fbdc51d218f132f6aa59c089dc5875ab877ab7..8d73fa7a2391a9e81b74b79938bc6a2026e24584 100644 --- a/test/test_predict_command.py +++ b/test/test_predict_command.py @@ -45,6 +45,7 @@ def test_no_csv(): "--alleles", "HLA-A0201", "H-2-Kb", "--peptides", "SIINFEKL", "DENDREKLLL", "PICKLEEE", "--prediction-column-prefix", "mhcflurry1_", + "--affinity-only", ] deletes = [] @@ -61,7 +62,7 @@ def test_no_csv(): os.unlink(delete) print(result) - assert_equal(result.shape, (6, 6)) + assert_equal(len(result), 6) sub_result1 = result.loc[result.peptide == "SIINFEKL"].set_index("allele") print(sub_result1) assert (