phykit 2.1.0__py3-none-any.whl

phykit/services/tree/nearest_neighbor_interchange.py

@@ -0,0 +1,140 @@
+import copy
+from typing import Dict, List, Generator
+import pickle
+
+from Bio import Phylo
+from Bio.Phylo import Newick
+
+from .base import Tree
+
+
+class NearestNeighborInterchange(Tree):
+    def __init__(self, args) -> None:
+        super().__init__(**self.process_args(args))
+
+    def run(self) -> None:
+        tree = self.read_tree_file()
+
+        # Use standard neighbor generation with optimized copying
+        all_nnis = [tree]
+        all_nnis.extend(self.get_neighbors(tree))
+        Phylo.write(all_nnis, self.output_file_path, "newick")
+
+    def process_args(self, args) -> Dict[str, str]:
+        tree_file_path = args.tree
+        output_file_path = \
+            f"{args.output}" if args.output else f"{tree_file_path}.nnis"
+
+        return dict(
+            tree_file_path=tree_file_path,
+            output_file_path=output_file_path
+        )
+
+    def _fast_tree_copy(self, tree: Newick.Tree) -> Newick.Tree:
+        """Fast tree copying using pickle instead of deep copy."""
+        return pickle.loads(pickle.dumps(tree, protocol=pickle.HIGHEST_PROTOCOL))
+
+    def get_neighbors(
+        self,
+        tree: Newick.Tree
+    ) -> List[Newick.Tree]:
+        ### This code is from BioPython (so is this comment)
+        # Get all neighbor trees of the given tree (PRIVATE).
+        # Currently only for binary rooted trees.
+        ###
+        # make child to parent dict
+        parents = {}
+        for clade in tree.find_clades():
+            if clade != tree.root:
+                node_path = tree.get_path(clade)
+                # cannot get the parent if the parent is root. Bug?
+                if len(node_path) == 1:
+                    parents[clade] = tree.root
+                else:
+                    parents[clade] = node_path[-2]
+        neighbors = []
+        root_childs = []
+        for clade in tree.get_nonterminals(order="level"):
+            if clade == tree.root:
+                left = clade.clades[0]
+                right = clade.clades[1]
+                root_childs.append(left)
+                root_childs.append(right)
+                if not left.is_terminal() and not right.is_terminal():
+                    # make changes around the left_left clade
+                    # left_left = left.clades[0]
+                    left_right = left.clades[1]
+                    right_left = right.clades[0]
+                    right_right = right.clades[1]
+                    # neighbor 1 (left_left + right_right)
+                    del left.clades[1]
+                    del right.clades[1]
+                    left.clades.append(right_right)
+                    right.clades.append(left_right)
+                    temp_tree = self._fast_tree_copy(tree)
+                    neighbors.append(temp_tree)
+                    # neighbor 2 (left_left + right_left)
+                    del left.clades[1]
+                    del right.clades[0]
+                    left.clades.append(right_left)
+                    right.clades.append(right_right)
+                    temp_tree = self._fast_tree_copy(tree)
+                    neighbors.append(temp_tree)
+                    # change back (left_left + left_right)
+                    del left.clades[1]
+                    del right.clades[0]
+                    left.clades.append(left_right)
+                    right.clades.insert(0, right_left)
+            elif clade in root_childs:
+                # skip root child
+                continue
+            else:
+                # method for other clades
+                # make changes around the parent clade
+                left = clade.clades[0]
+                right = clade.clades[1]
+                parent = parents[clade]
+                if clade == parent.clades[0]:
+                    sister = parent.clades[1]
+                    # neighbor 1 (parent + right)
+                    del parent.clades[1]
+                    del clade.clades[1]
+                    parent.clades.append(right)
+                    clade.clades.append(sister)
+                    temp_tree = self._fast_tree_copy(tree)
+                    neighbors.append(temp_tree)
+                    # neighbor 2 (parent + left)
+                    del parent.clades[1]
+                    del clade.clades[0]
+                    parent.clades.append(left)
+                    clade.clades.append(right)
+                    temp_tree = self._fast_tree_copy(tree)
+                    neighbors.append(temp_tree)
+                    # change back (parent + sister)
+                    del parent.clades[1]
+                    del clade.clades[0]
+                    parent.clades.append(sister)
+                    clade.clades.insert(0, left)
+                else:
+                    sister = parent.clades[0]
+                    # neighbor 1 (parent + right)
+                    del parent.clades[0]
+                    del clade.clades[1]
+                    parent.clades.insert(0, right)
+                    clade.clades.append(sister)
+                    temp_tree = self._fast_tree_copy(tree)
+                    neighbors.append(temp_tree)
+                    # neighbor 2 (parent + left)
+                    del parent.clades[0]
+                    del clade.clades[0]
+                    parent.clades.insert(0, left)
+                    clade.clades.append(right)
+                    temp_tree = self._fast_tree_copy(tree)
+                    neighbors.append(temp_tree)
+                    # change back (parent + sister)
+                    del parent.clades[0]
+                    del clade.clades[0]
+                    parent.clades.insert(0, sister)
+                    clade.clades.insert(0, left)
+
+        return neighbors

phykit/services/tree/patristic_distances.py

@@ -0,0 +1,113 @@
+from typing import Dict, List, Tuple
+import itertools
+import multiprocessing as mp
+from functools import partial
+import pickle
+import sys
+
+from Bio.Phylo import Newick
+try:
+    from tqdm import tqdm
+except ImportError:
+    # Fallback if tqdm is not installed
+    def tqdm(iterable, *args, **kwargs):
+        return iterable
+
+from .base import Tree
+
+from ...helpers.stats_summary import (
+    calculate_summary_statistics_from_arr,
+    print_summary_statistics,
+)
+
+
+class PatristicDistances(Tree):
+    def __init__(self, args) -> None:
+        super().__init__(**self.process_args(args))
+
+    def run(self):
+        tree = self.read_tree_file()
+        patristic_distances, combos, stats = \
+            self.calculate_patristic_distances(tree)
+
+        if self.verbose:
+            try:
+                for combo, patristic_distance in zip(combos, patristic_distances):
+                    print(f"{combo[0]}\t{combo[1]}\t{round(patristic_distance, 4)}")
+            except BrokenPipeError:
+                pass
+        else:
+            print_summary_statistics(stats)
+
+    def process_args(self, args) -> Dict[str, str]:
+        return dict(tree_file_path=args.tree, verbose=args.verbose)
+
+    def _calculate_distance_batch(self, tree_pickle, combo_batch):
+        """Helper function to calculate distances for a batch of combinations."""
+        tree = pickle.loads(tree_pickle)
+        return [tree.distance(combo[0], combo[1]) for combo in combo_batch]
+
+    def calculate_distance_between_pairs(
+        self,
+        tips: List[str],
+        tree
+    ) -> Tuple[
+        List[Tuple[str, str]],
+        List[float],
+    ]:
+        combos = list(itertools.combinations(tips, 2))
+
+        # For small datasets, use the original single-threaded approach
+        if len(combos) < 100:
+            patristic_distances = [
+                tree.distance(combo[0], combo[1]) for combo in combos
+            ]
+        else:
+            # Use multiprocessing for larger datasets
+            # Serialize the tree once to avoid repeated serialization
+            tree_pickle = pickle.dumps(tree)
+
+            # Determine optimal number of workers
+            num_workers = min(mp.cpu_count(), 8)
+
+            # Split combos into chunks for parallel processing
+            chunk_size = max(1, len(combos) // (num_workers * 4))
+            combo_chunks = [combos[i:i + chunk_size] for i in range(0, len(combos), chunk_size)]
+
+            # Create partial function with the pickled tree
+            calc_func = partial(self._calculate_distance_batch, tree_pickle)
+
+            # Process in parallel with progress bar
+            with mp.Pool(processes=num_workers) as pool:
+                # Only show progress bar if stderr is a tty (not redirected)
+                if sys.stderr.isatty():
+                    results = list(tqdm(
+                        pool.imap(calc_func, combo_chunks),
+                        total=len(combo_chunks),
+                        desc="Calculating patristic distances",
+                        unit="batch"
+                    ))
+                else:
+                    results = pool.map(calc_func, combo_chunks)
+
+            # Flatten the results
+            patristic_distances = [dist for chunk_result in results for dist in chunk_result]
+
+        return combos, patristic_distances
+
+    def calculate_patristic_distances(
+        self,
+        tree: Newick.Tree,
+    ) -> Tuple[
+        List[float],
+        List[Tuple[str, str]],
+        Dict[str, float],
+    ]:
+        tips = self.get_tip_names_from_tree(tree)
+
+        combos, patristic_distances = \
+            self.calculate_distance_between_pairs(tips, tree)
+
+        stats = calculate_summary_statistics_from_arr(patristic_distances)
+
+        return patristic_distances, combos, stats