updated bss function to be more efficient and faster

phykit/services/tree/bipartition_support_stats.py

@@ -1,3 +1,7 @@
+from typing import Dict, List, Tuple
+
+from Bio.Phylo import Newick
+
 from .base import Tree
 from ...helpers.stats_summary import (
     calculate_summary_statistics_from_arr,
@@ -9,35 +13,34 @@ class BipartitionSupportStats(Tree):
     def __init__(self, args) -> None:
         super().__init__(**self.process_args(args))
 
-    def run(self):
+    def run(self) -> None:
         tree = self.read_tree_file()
         bs_vals, term_names = self.get_bipartition_support_vals(tree)
 
         if self.verbose:
             try:
-                for bs_val, names in zip(bs_vals, term_names):
-                    print(bs_val, ";".join(names))
+                for i in range(len(bs_vals)):
+                    print(bs_vals[i], ";".join(term_names[i]))
             except BrokenPipeError:
                 pass
         else:
             stats = calculate_summary_statistics_from_arr(bs_vals)
             print_summary_statistics(stats)
 
-    def process_args(self, args):
+    def process_args(self, args) -> Dict[str, str]:
         return dict(tree_file_path=args.tree, verbose=args.verbose)
 
-    def get_bipartition_support_vals(self, tree):
-        # initialize list to hold bootstrap values
-        bs_vals = []
-        term_names = []
-
-        # populate bs_vals with bootstrap values
-        for nonterminal in tree.get_nonterminals():
-            # only include if a bootstrap value is present
-            if nonterminal.confidence != None:
-                bs_vals.append(nonterminal.confidence)
-                temp = []
-                for term in nonterminal.get_terminals():
-                    temp.append(term.name)
-                term_names.append(temp)
+    def get_bipartition_support_vals(
+        self,
+        tree: Newick.Tree,
+    ) -> Tuple[List[float], List[List[str]]]:
+        bs_vals = [
+            nonterminal.confidence for nonterminal in tree.get_nonterminals()
+            if nonterminal.confidence is not None
+        ]
+        term_names = [
+            [term.name for term in nonterminal.get_terminals()]
+            for nonterminal in tree.get_nonterminals()
+            if nonterminal.confidence is not None
+        ]
         return bs_vals, term_names

tests/integration/tree/test_bipartition_support_stats_integration.py

@@ -1,9 +1,7 @@
 import pytest
 import sys
-from math import isclose
 from mock import patch, call
 from pathlib import Path
-from textwrap import dedent
 
 from phykit.phykit import Phykit
 
@@ -61,7 +59,7 @@ def test_bipartition_support_incorrect_input_file(self, mocked_print):
             "bipartition_support_stats",
             f"{here.parent.parent.parent}/sample_files/small_Aspergillus_tree.tr"
         ]
-        
+
         with pytest.raises(SystemExit) as pytest_wrapped_e:
             Phykit()