feat: filter to top-N trees by llh in generate_input main loop

Author: johnoooh

src/neoantigen_utils/generate_input.py

@@ -347,7 +347,7 @@ def main(args):
 
     trees = summ_data["trees"]
 
-    for tree in trees:
+    for tree in select_top_trees(trees, args.top_n_trees):
 
         inner_sample_tree_dict = {"topology": [], "score": trees[tree]["llh"]}
         with open("./" + args.tree_directory + "/" + str(tree) + ".json", "r") as f:
@@ -1331,6 +1331,12 @@ def parse_args():
     parser.add_argument(
         "--kD_cutoff", default=500, help="Cutoff value for the kD, default is 500",
     )
+    parser.add_argument(
+        "--top_n_trees",
+        type=int,
+        default=10,
+        help="Number of top trees (by llh, highest first) to include, default is 10",
+    )
 
     return parser.parse_args()
 

tests/test_generate_input.py

@@ -230,3 +230,24 @@ def test_n_zero_returns_empty_list(self):
 
     def test_empty_trees_returns_empty_list(self):
         assert select_top_trees({}, 10) == []
+
+
+class TestTopNTreesArgDefault:
+    def test_default_top_n_trees_is_ten(self):
+        import sys
+        from neoantigen_utils.generate_input import parse_args
+
+        argv = [
+            "prog",
+            "--maf_file", "x", "--summary_file", "x", "--mutation_file", "x",
+            "--tree_directory", "x", "--id", "x", "--patient_id", "x",
+            "--cohort", "x", "--HLA_genes", "x",
+            "--netMHCpan_MUT_input", "x", "--netMHCpan_WT_input", "x",
+        ]
+        old_argv = sys.argv
+        sys.argv = argv
+        try:
+            args = parse_args()
+        finally:
+            sys.argv = old_argv
+        assert args.top_n_trees == 10