Add new argument to aggregate report creation to limit predictions to specific alleles

pvactools/lib/aggregate_all_epitopes.py

@@ -20,6 +20,10 @@
 class AggregateAllEpitopes:
     def __init__(self):
         self.hla_types = pd.read_csv(self.input_file, delimiter="\t", usecols=["HLA Allele"])['HLA Allele'].unique()
+        if self.limiting_alleles:
+            extra_alleles = list(set(self.limiting_alleles) - set(self.hla_types))
+            if len(extra_alleles) > 0:
+                raise Exception(f"Alleles specified in the --alleles parameter not found in the input file: {extra_alleles.join(', ')}")
         thresholds = {}
         for hla_type in self.hla_types:
             threshold = PredictionClass.cutoff_for_allele(hla_type)
@@ -387,6 +391,7 @@ def __init__(
             anchor_contribution_threshold=0.8,
             aggregate_inclusion_binding_threshold=5000,
             aggregate_inclusion_count_limit=15,
+            limiting_alleles=None,
         ):
         self.input_file = input_file
         self.output_file = output_file
@@ -414,6 +419,7 @@ def __init__(
             self.mt_top_score_metric = "Best"
             self.wt_top_score_metric = "Corresponding"
         self.top_score_metric2 = top_score_metric2
+        self.limiting_alleles = limiting_alleles
         self.metrics_file = output_file.replace('.tsv', '.metrics.json')
         super().__init__()
         self.anchor_calculator = AnchorResiduePass(binding_threshold, self.use_allele_specific_binding_thresholds, self.allele_specific_binding_thresholds, allele_specific_anchors, anchor_contribution_threshold, self.wt_top_score_metric)
@@ -443,6 +449,8 @@ def calculate_clonal_vaf(self):
     def read_input_file(self, used_columns, dtypes):
         df = pd.read_csv(self.input_file, delimiter='\t', float_precision='high', low_memory=False, na_values="NA", keep_default_na=False, usecols=used_columns, dtype=dtypes)
         df = df.astype({"{} MT IC50 Score".format(self.mt_top_score_metric):'float'})
+        if self.limiting_alleles:
+            df = df[df["HLA Allele"].isin(self.limiting_alleles)]
         return df
 
     def get_sub_df(self, all_epitopes_df, key):
@@ -833,6 +841,7 @@ def __init__(self,
             top_score_metric2=["ic50", "combined_percentile"],
             aggregate_inclusion_binding_threshold=5000,
             aggregate_inclusion_count_limit=15,
+            limiting_alleles=None,
         ):
         self.input_file = input_file
         self.output_file = output_file
@@ -850,6 +859,7 @@ def __init__(self,
         else:
             self.mt_top_score_metric = "Best"
         self.top_score_metric2 = top_score_metric2
+        self.limiting_alleles = limiting_alleles
         self.metrics_file = output_file.replace('.tsv', '.metrics.json')
         super().__init__()
 
@@ -867,6 +877,8 @@ def read_input_file(self, used_columns, dtypes):
         df = pd.read_csv(self.input_file, delimiter='\t', float_precision='high', low_memory=False, na_values="NA", keep_default_na=False, dtype={"Index": str})
         df = df[df["{} IC50 Score".format(self.mt_top_score_metric)] != 'NA']
         df = df.astype({"{} IC50 Score".format(self.mt_top_score_metric):'float'})
+        if self.limiting_alleles:
+            df = df[df["HLA Allele"].isin(self.limiting_alleles)]
         return df
 
     def get_sub_df(self, all_epitopes_df, key):
@@ -951,6 +963,7 @@ def __init__(
         expn_val=0.1,
         aggregate_inclusion_binding_threshold=5000,
         aggregate_inclusion_count_limit=15,
+        limiting_alleles=None,
     ):
         UnmatchedSequenceAggregateAllEpitopes.__init__(
             self,
@@ -966,6 +979,7 @@ def __init__(
             top_score_metric2=top_score_metric2,
             aggregate_inclusion_binding_threshold=aggregate_inclusion_binding_threshold,
             aggregate_inclusion_count_limit=aggregate_inclusion_count_limit,
+            limiting_alleles=limiting_alleles,
         )
         self.read_support = read_support
         self.expn_val = expn_val
@@ -1092,6 +1106,7 @@ def __init__(
         transcript_prioritization_strategy=['canonical', 'mane_select', 'tsl'],
         maximum_transcript_support_level=1,
         allow_incomplete_transcripts=False,
+        limiting_alleles=None,
     ):
         PvacbindAggregateAllEpitopes.__init__(
             self,
@@ -1107,6 +1122,7 @@ def __init__(
             aggregate_inclusion_count_limit=aggregate_inclusion_count_limit,
             top_score_metric=top_score_metric,
             top_score_metric2=top_score_metric2,
+            limiting_alleles=limiting_alleles,
         )
         self.tumor_purity = tumor_purity
         self.trna_vaf = trna_vaf
@@ -1124,8 +1140,11 @@ def get_list_unique_mutation_keys(self, df):
 
     # pvacbind w/ Index instead of Mutation
     def read_input_file(self, used_columns, dtypes):
-        return pd.read_csv(self.input_file, delimiter='\t', float_precision='high', low_memory=False,
+        df = pd.read_csv(self.input_file, delimiter='\t', float_precision='high', low_memory=False,
                            na_values="NA", keep_default_na=False, dtype={"Index": str})
+        if self.limiting_alleles:
+            df = df[df["HLA Allele"].isin(self.limiting_alleles)]
+        return df
 
     def sort_included_df(self, df):
         return PvacspliceBestCandidate(

pvactools/tools/pvacbind/generate_aggregated_report.py

@@ -83,6 +83,11 @@ def define_parser():
              + "Whether the lowest or median is considered for each metric is controlled by the --top-score-metric parameter. ",
         default=['ic50', 'combined_percentile'],
     )
+    parser.add_argument(
+        "--limiting-alleles", type=lambda s:[a for a in s.split(',')],
+        help="Comma-separated list of alleles used for the predictions made in the input file. "
+             + "If specified, only predictions for those alleles will be included in the aggregated report."
+    )
     return parser
 
 def main(args_input = sys.argv[1:]):
@@ -104,6 +109,7 @@ def main(args_input = sys.argv[1:]):
         top_score_metric=args.top_score_metric,
         aggregate_inclusion_binding_threshold=args.aggregate_inclusion_binding_threshold,
         aggregate_inclusion_count_limit=args.aggregate_inclusion_count_limit,
+        limiting_alleles=args.limiting_alleles,
     ).execute()
     print("Completed")
 

pvactools/tools/pvacfuse/generate_aggregated_report.py

@@ -93,6 +93,11 @@ def define_parser():
         help="Expression Cutoff. Expression is meassured as FFPM (fusion fragments per million total reads). When failing this cutoff sites will be binned in the \"LowExpr\" tier.",
         default=0.1
     )
+    parser.add_argument(
+        "--limiting-alleles", type=lambda s:[a for a in s.split(',')],
+        help="Comma-separated list of alleles used for the predictions made in the input file. "
+             + "If specified, only predictions for those alleles will be included in the aggregated report."
+    )
 
     return parser
 
@@ -118,6 +123,7 @@ def main(args_input = sys.argv[1:]):
         expn_val=args.expn_val,
         aggregate_inclusion_binding_threshold=args.aggregate_inclusion_binding_threshold,
         aggregate_inclusion_count_limit=args.aggregate_inclusion_count_limit,
+        limiting_alleles=args.limiting_alleles,
     ).execute()
     print("Completed")
 

pvactools/tools/pvacseq/generate_aggregated_report.py

@@ -137,6 +137,11 @@ def define_parser():
              + " As a result, a higher threshold leads to the inclusion of more positions to be considered anchors.",
         default=0.8
     )
+    parser.add_argument(
+        "--limiting-alleles", type=lambda s:[a for a in s.split(',')],
+        help="Comma-separated list of alleles used for the predictions made in the input file. "
+             + "If specified, only predictions for those alleles will be included in the aggregated report."
+    )
 
     return parser
 
@@ -168,6 +173,7 @@ def main(args_input = sys.argv[1:]):
         anchor_contribution_threshold=args.anchor_contribution_threshold,
         aggregate_inclusion_binding_threshold=args.aggregate_inclusion_binding_threshold,
         aggregate_inclusion_count_limit=args.aggregate_inclusion_count_limit,
+        limiting_alleles=args.limiting_alleles,
     ).execute()
     print("Completed")
 

pvactools/tools/pvacsplice/generate_aggregated_report.py

@@ -118,6 +118,11 @@ def define_parser():
         default=1,
         choices=[1,2,3,4,5]
     )
+    parser.add_argument(
+        "--limiting-alleles", type=lambda s:[a for a in s.split(',')],
+        help="Comma-separated list of alleles used for the predictions made in the input file. "
+             + "If specified, only predictions for those alleles will be included in the aggregated report."
+    )
 
     return parser
 
@@ -148,6 +153,7 @@ def main(args_input = sys.argv[1:]):
         top_score_metric2=args.top_score_metric2,
         aggregate_inclusion_binding_threshold=args.aggregate_inclusion_binding_threshold,
         aggregate_inclusion_count_limit=args.aggregate_inclusion_count_limit,
+        limiting_alleles=args.limiting_alleles,
     ).execute()
     print("Completed")
 

tests/test_aggregate_all_epitopes.py

@@ -98,6 +98,37 @@ def test_aggregate_all_epitopes_HCC1395_pvacseq_normalized_percentiles_runs_and_
             self.assertTrue(os.path.isfile(pvacview_file))
             os.remove(pvacview_file)
 
+    def test_aggregate_all_epitopes_HCC1395_pvacseq_limiting_alleles_runs_and_produces_expected_output(self):
+        self.assertTrue(py_compile.compile(self.executable))
+        output_file = tempfile.NamedTemporaryFile(suffix='.tsv')
+        self.assertFalse(PvacseqAggregateAllEpitopes(
+            os.path.join(self.test_data_dir, 'HCC1395_TUMOR_DNA.all_epitopes.short.tsv'),
+            output_file.name,
+            limiting_alleles=["HLA-C*06:02"]
+        ).execute())
+        self.assertTrue(cmp(
+            output_file.name,
+            os.path.join(self.test_data_dir, "HCC1395.limiting_alleles.output.tsv"),
+        ))
+
+        metrics_file = output_file.name.replace('.tsv', '.metrics.json')
+        self.assertTrue(cmp(
+            metrics_file,
+            os.path.join(self.test_data_dir, "HCC1395.limiting_alleles.output.metrics.json"),
+        ))
+        os.remove(metrics_file)
+
+        for i in self.pvacview_r_files:
+            pvacview_file = os.path.join(os.path.dirname(output_file.name), i)
+            self.assertTrue(os.path.isfile(pvacview_file))
+            os.remove(pvacview_file)
+
+        for i in ["anchor.jpg", "pVACview_logo.png", "pVACview_logo_mini.png"]:
+            pvacview_file = os.path.join(os.path.dirname(output_file.name), "www", i)
+            self.assertTrue(os.path.isfile(pvacview_file))
+            os.remove(pvacview_file)
+
+
     def test_aggregate_all_epitopes_all_class_i_pvacseq_runs_and_produces_expected_output(self):
         self.assertTrue(py_compile.compile(self.executable))
         output_file = tempfile.NamedTemporaryFile(suffix='.tsv')