phykit 2.1.0__py3-none-any.whl

phykit/services/tree/base.py

@@ -0,0 +1,178 @@
+import sys
+import copy
+from typing import List
+from functools import lru_cache
+import os
+import hashlib
+
+from Bio import Phylo
+
+from ..base import BaseService
+
+
+class Tree(BaseService):
+    def __init__(
+        self,
+        *args,
+        tree_file_path=None,
+        idmap=None,
+        alignment_file_path=None,
+        tree1_file_path=None,
+        outgroup_taxa_file_path=None,
+        output_file_path=None,
+        factor=None,
+        remove=None,
+        verbose=None,
+        reference=None,
+        list_of_taxa=None,
+        trees=None,
+        groups=None,
+        support=None,
+        tip_1=None,
+        tip_2=None,
+        clade=None,
+        keep=None,
+        exclude_gaps=None,
+    ):
+        self.tree_file_path = tree_file_path
+        self.tree1_file_path = tree1_file_path
+        self.alignment_file_path = alignment_file_path
+        self.output_file_path = output_file_path
+        self.outgroup_taxa_file_path = outgroup_taxa_file_path
+        self.tree_format = "newick"
+        self.verbose = verbose
+        self.factor = factor
+        self.remove = remove
+        self.idmap = idmap
+        self.reference = reference
+        self.list_of_taxa = list_of_taxa
+        self.trees = trees
+        self.groups = groups
+        self.support = support
+        self.tip_1 = tip_1
+        self.tip_2 = tip_2
+        self.clade = clade
+        self.keep = keep
+        self.exclude_gaps = exclude_gaps
+
+    @staticmethod
+    @lru_cache(maxsize=32)
+    def _cached_tree_read(file_path: str, tree_format: str, file_hash: str):
+        """Cached tree reading with file hash for cache invalidation."""
+        return Phylo.read(file_path, tree_format)
+
+    @staticmethod
+    def _get_file_hash(file_path: str) -> str:
+        """Get a hash based on file path, size, and modification time."""
+        try:
+            stat = os.stat(file_path)
+            cache_key = f"{file_path}_{stat.st_size}_{stat.st_mtime}"
+            return hashlib.md5(cache_key.encode()).hexdigest()
+        except:
+            return ""
+
+    def read_tree_file(self):
+        try:
+            file_hash = self._get_file_hash(self.tree_file_path)
+            tree = self._cached_tree_read(self.tree_file_path, self.tree_format, file_hash)
+            # Return a deep copy to prevent modifications to the cached tree
+            return copy.deepcopy(tree)
+        except FileNotFoundError:
+            print(f"{self.tree_file_path} corresponds to no such file or directory.")
+            print("Please check filename and pathing")
+            sys.exit(2)
+
+    def read_tree1_file(self):
+        try:
+            file_hash = self._get_file_hash(self.tree1_file_path)
+            tree = self._cached_tree_read(self.tree1_file_path, self.tree_format, file_hash)
+            # Return a deep copy to prevent modifications to the cached tree
+            return copy.deepcopy(tree)
+        except FileNotFoundError:
+            print(f"{self.tree1_file_path} corresponds to no such file or directory.")
+            print("Please check filename and pathing")
+            sys.exit(2)
+
+    def read_reference_tree_file(self):
+        try:
+            file_hash = self._get_file_hash(self.reference)
+            tree = self._cached_tree_read(self.reference, self.tree_format, file_hash)
+            # Return a deep copy to prevent modifications to the cached tree
+            return copy.deepcopy(tree)
+        except FileNotFoundError:
+            print(f"{self.reference} corresponds to no such file or directory.")
+            print("Please check filename and pathing")
+            sys.exit(2)
+
+    def write_tree_file(self, tree, output_file_path):
+        return Phylo.write(tree, output_file_path, self.tree_format)
+
+    def get_tip_names_from_tree(self, tree) -> list:
+        """
+        get tip names from a tree
+        """
+        # Use list comprehension for better performance
+        return [tip.name for tip in tree.get_terminals()]
+
+    def shared_tips(self, a, b):
+        """
+        Determines what tips are shared between two trees
+        -------------------------------------------------
+        argv: a
+            list of tips from one tree
+        argv: b
+            list of tips from a second tree
+        """
+
+        a_set = set(a)
+        b_set = set(b)
+
+        # check length
+        if len(a_set.intersection(b_set)) > 0:
+            return list(a_set.intersection(b_set))
+        else:
+            print("no common tips")
+            sys.exit(2)
+
+    def prune_tree_using_taxa_list(self, tree, taxa_to_prune: list):
+        """
+        prune taxa from tree
+        """
+        for taxon in taxa_to_prune:
+            tree.prune(taxon)
+
+        return tree
+
+    def calculate_treeness(self, tree=None, print_value=False):
+        if not tree:
+            tree = self.read_tree_file()
+
+        inter_len = float(0.0)
+        # determine internal branch lengths
+        for interal in tree.get_nonterminals():
+            # only include if a branch length value is present
+            if interal.branch_length != None:
+                inter_len += interal.branch_length
+        # determine total branch length
+        total_len = tree.total_branch_length()
+
+        try:
+            treeness = float(inter_len / total_len)
+            try:
+                if print_value:
+                    print(f"{treeness}")
+                return treeness
+            except BrokenPipeError:
+                pass
+        except ZeroDivisionError:
+            try:
+                print("Invalid tree. Tree should contain branch lengths")
+                return None
+            except BrokenPipeError:
+                pass
+
+    def get_gap_chars(is_protein: bool) -> List[str]:
+        if is_protein:
+            return ["-", "?", "*", "X", "x"]
+        else:
+            return ["-", "?", "*", "X", "x", "N", "n"]

phykit/services/tree/bipartition_support_stats.py

@@ -0,0 +1,48 @@
+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,
+    print_summary_statistics,
+)
+
+
+class BipartitionSupportStats(Tree):
+    def __init__(self, args) -> None:
+        super().__init__(**self.process_args(args))
+
+    def run(self) -> None:
+        tree = self.read_tree_file()
+        bs_vals, term_names = self.get_bipartition_support_vals(tree)
+
+        if self.verbose:
+            try:
+                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) -> Dict[str, str]:
+        return dict(tree_file_path=args.tree, verbose=args.verbose)
+
+    def get_bipartition_support_vals(
+        self,
+        tree: Newick.Tree,
+    ) -> Tuple[List[float], List[List[str]]]:
+        # Single pass through nonterminals to avoid duplicate tree traversal
+        bs_vals = []
+        term_names = []
+
+        # Cache terminals for each nonterminal in one pass
+        for nonterminal in tree.get_nonterminals():
+            if nonterminal.confidence is not None:
+                bs_vals.append(nonterminal.confidence)
+                # Get terminal names once for this nonterminal
+                term_names.append([term.name for term in nonterminal.get_terminals()])
+
+        return bs_vals, term_names

phykit/services/tree/branch_length_multiplier.py

@@ -0,0 +1,37 @@
+from typing import Dict
+import copy
+
+from Bio.Phylo import Newick
+
+from .base import Tree
+
+
+class BranchLengthMultiplier(Tree):
+    def __init__(self, args) -> None:
+        super().__init__(**self.process_args(args))
+
+    def run(self) -> None:
+        tree = self.read_tree_file()
+        # Make a deep copy to avoid modifying the cached tree
+        tree_copy = copy.deepcopy(tree)
+        self.multiply_branch_lengths_by_factor(tree_copy, self.factor)
+        self.write_tree_file(tree_copy, self.output_file_path)
+
+    def process_args(self, args) -> Dict[str, str]:
+        output_file_path = \
+            args.output or f"{args.tree}.factor_{args.factor}.tre"
+        return dict(
+            tree_file_path=args.tree,
+            factor=args.factor,
+            output_file_path=output_file_path,
+        )
+
+    def multiply_branch_lengths_by_factor(
+        self,
+        tree: Newick.Tree,
+        factor: float,
+    ) -> Newick.Tree:
+        for node in tree.get_nonterminals() + tree.get_terminals():
+            if node.branch_length is not None:
+                node.branch_length *= factor
+        return tree

phykit/services/tree/collapse_branches.py

@@ -0,0 +1,27 @@
+from typing import Dict
+import copy
+
+from .base import Tree
+
+
+class CollapseBranches(Tree):
+    def __init__(self, args) -> None:
+        super().__init__(**self.process_args(args))
+
+    def run(self):
+        tree = self.read_tree_file()
+        # Make a deep copy to avoid modifying the cached tree
+        tree_copy = copy.deepcopy(tree)
+        tree_copy.collapse_all(
+            lambda c: c.confidence and c.confidence < self.support
+        )
+        self.write_tree_file(tree_copy, self.output_file_path)
+
+    def process_args(self, args) -> Dict[str, str]:
+        output_file_path = \
+            args.output or f"{args.tree}.collapsed_{args.support}.tre"
+        return dict(
+            tree_file_path=args.tree,
+            support=args.support,
+            output_file_path=output_file_path,
+        )

phykit/services/tree/covarying_evolutionary_rates.py

@@ -0,0 +1,272 @@
+import copy
+import numpy as np
+from concurrent.futures import ProcessPoolExecutor, as_completed
+from functools import lru_cache
+import pickle
+
+from scipy.stats import pearsonr, zscore
+
+from .base import Tree
+
+
+class CovaryingEvolutionaryRates(Tree):
+    def __init__(self, args) -> None:
+        super().__init__(**self.process_args(args))
+
+    def run(self):
+        tree_zero = self.read_tree_file()
+        tree_one = self.read_tree1_file()
+        tree_ref = self.read_reference_tree_file()
+
+        # - Calculate correlation between two gene trees
+        # and save results to an array, corrArr.
+        # - Branch lengths will also be part of output
+
+        # get tree tip names
+        tree_zero_tips = self.get_tip_names_from_tree(tree_zero)
+        tree_one_tips = self.get_tip_names_from_tree(tree_one)
+        tree_ref_tips = self.get_tip_names_from_tree(tree_ref)
+
+        # get shared tips between the two trees
+        shared_tree_tips = self.shared_tips(tree_zero_tips, tree_one_tips)
+
+        # find differences between tree tips and shared tips
+        # to determine what tips to prune
+        tree_zero_tips_to_prune = list(set(tree_zero_tips) - set(shared_tree_tips))
+        tree_one_tips_to_prune = list(set(tree_one_tips) - set(shared_tree_tips))
+        tree_ref_tips_to_prune = list(set(tree_ref_tips) - set(shared_tree_tips))
+
+        # get a set of pruned trees
+        tree_zero = self.prune_tips(tree_zero, tree_zero_tips_to_prune)
+        tree_one = self.prune_tips(tree_one, tree_one_tips_to_prune)
+        tree_ref = self.prune_tips(tree_ref, tree_ref_tips_to_prune)
+
+        # obtain corrected branch lengths where branch lengths
+        # are corrected by the species tree branch length
+        (
+            tree_zero_corr_branch_lengths,
+            tree_one_corr_branch_lengths,
+            tip_names,
+        ) = self.correct_branch_lengths(tree_zero, tree_one, tree_ref)
+
+        # remove corrected BLs greater than 5
+        outlier_indices = []
+        outlier_indices = self.get_indices_of_outlier_branch_lengths(
+            tree_zero_corr_branch_lengths, outlier_indices
+        )
+        outlier_indices = self.get_indices_of_outlier_branch_lengths(
+            tree_one_corr_branch_lengths, outlier_indices
+        )
+
+        tree_zero_corr_branch_lengths = self.remove_outliers_based_on_indices(
+            tree_zero_corr_branch_lengths, outlier_indices
+        )
+        tree_one_corr_branch_lengths = self.remove_outliers_based_on_indices(
+            tree_one_corr_branch_lengths, outlier_indices
+        )
+        tip_names = self.remove_outliers_based_on_indices(tip_names, outlier_indices)
+
+        # standardize values for final correction
+        tree_zero_corr_branch_lengths = zscore(tree_zero_corr_branch_lengths)
+        tree_one_corr_branch_lengths = zscore(tree_one_corr_branch_lengths)
+
+        # Calculate correlation and append to results array
+        # also keep a list of p values
+        corr = list(
+            pearsonr(tree_zero_corr_branch_lengths, tree_one_corr_branch_lengths)
+        )
+
+        try:
+            if self.verbose:
+                for val_zero, val_one, tip_name in zip(
+                    tree_zero_corr_branch_lengths,
+                    tree_one_corr_branch_lengths,
+                    tip_names,
+                ):
+                    print(
+                        f"{round(val_zero, 4)}\t{round(val_one, 4)}\t{';'.join(tip_name)}"
+                    )
+            else:
+                print(f"{round(corr[0], 4)}\t{round(corr[1], 6)}")
+        except BrokenPipeError:
+            pass
+
+    def process_args(self, args):
+        return dict(
+            tree_file_path=args.tree_zero,
+            tree1_file_path=args.tree_one,
+            reference=args.reference,
+            verbose=args.verbose,
+        )
+
+    def get_indices_of_outlier_branch_lengths(
+        self, corr_branch_lengths, outlier_indices
+    ):
+        """
+        create index for branch lengths that
+        have an absolute value greater than 5
+        """
+        # Convert to numpy array for vectorized operations
+        arr = np.array(corr_branch_lengths, dtype=float)
+
+        # Find outliers using vectorized operations
+        new_outliers = np.where((np.abs(arr) > 5) | np.isnan(arr))[0]
+
+        # Combine with existing outliers
+        all_outliers = set(outlier_indices)
+        all_outliers.update(new_outliers.tolist())
+
+        return list(all_outliers)
+
+    def remove_outliers_based_on_indices(self, corr_branch_lengths, outlier_indices):
+        """
+        remove value if the value is an outlier according
+        to the outlier indices list
+        """
+        if not outlier_indices:
+            return corr_branch_lengths
+
+        # Use numpy for efficient filtering
+        mask = np.ones(len(corr_branch_lengths), dtype=bool)
+        mask[list(outlier_indices)] = False
+
+        if isinstance(corr_branch_lengths[0], (list, tuple)):
+            # Handle list of lists (tip_names)
+            return [item for i, item in enumerate(corr_branch_lengths) if mask[i]]
+        else:
+            # Handle numeric lists
+            return [item for i, item in enumerate(corr_branch_lengths) if mask[i]]
+
+    def prune_tips(self, tree, tips):
+        """
+        prune tips from trees
+        """
+
+        for tip in tips:
+            tree.prune(tip)
+
+        return tree
+
+    @staticmethod
+    def _process_terminal_batch(tree0_pickle, tree1_pickle, terminals_data):
+        """Process a batch of terminals in parallel."""
+        t0 = pickle.loads(tree0_pickle)
+        t1 = pickle.loads(tree1_pickle)
+
+        results = []
+        for terminal_name, terminal_bl, sp_tips in terminals_data:
+            try:
+                newtree = t0.common_ancestor(terminal_name)
+                newtree1 = t1.common_ancestor(terminal_name)
+
+                bl0 = round(newtree.branch_length / terminal_bl, 6) if newtree.branch_length else None
+                bl1 = round(newtree1.branch_length / terminal_bl, 6) if newtree1.branch_length else None
+
+                if bl0 is not None and bl1 is not None:
+                    results.append((bl0, bl1, sp_tips))
+            except:
+                continue
+        return results
+
+    @staticmethod
+    def _process_nonterminal_batch(tree0_pickle, tree1_pickle, nonterminals_data):
+        """Process a batch of nonterminals in parallel."""
+        t0 = pickle.loads(tree0_pickle)
+        t1 = pickle.loads(tree1_pickle)
+
+        results = []
+        for sp_tips, nonterminal_bl in nonterminals_data:
+            try:
+                newtree = t0.common_ancestor(sp_tips)
+                newtree1 = t1.common_ancestor(sp_tips)
+
+                if newtree.branch_length and newtree1.branch_length and nonterminal_bl:
+                    bl0 = round(newtree.branch_length / nonterminal_bl, 6)
+                    bl1 = round(newtree1.branch_length / nonterminal_bl, 6)
+                    results.append((bl0, bl1, sp_tips))
+            except:
+                continue
+        return results
+
+    def correct_branch_lengths(self, t0, t1, sp):
+        """
+        obtain a list of corrected branch lengths with parallel processing
+        """
+        l0 = []
+        l1 = []
+        tip_names = []
+
+        # Collect terminal data
+        terminals = sp.get_terminals()
+        nonterminals = sp.get_nonterminals()
+
+        # Process sequentially if small dataset or use parallel processing
+        if len(terminals) + len(nonterminals) < 50:
+            # Original sequential processing for small datasets
+            for i in terminals:
+                sp_tips = self.get_tip_names_from_tree(i)
+                tip_names.append(sp_tips)
+                try:
+                    newtree = t0.common_ancestor(i.name)
+                    newtree1 = t1.common_ancestor(i.name)
+                    if newtree.branch_length and i.branch_length:
+                        l0.append(round(newtree.branch_length / i.branch_length, 6))
+                        l1.append(round(newtree1.branch_length / i.branch_length, 6))
+                except:
+                    continue
+
+            for i in nonterminals:
+                sp_tips = self.get_tip_names_from_tree(i)
+                try:
+                    newtree = t0.common_ancestor(sp_tips)
+                    newtree1 = t1.common_ancestor(sp_tips)
+                    if newtree.branch_length and newtree1.branch_length and i.branch_length:
+                        l0.append(round(newtree.branch_length / i.branch_length, 6))
+                        l1.append(round(newtree1.branch_length / i.branch_length, 6))
+                        tip_names.append(sp_tips)
+                except:
+                    continue
+        else:
+            # Parallel processing for large datasets
+            tree0_pickle = pickle.dumps(t0)
+            tree1_pickle = pickle.dumps(t1)
+
+            # Prepare terminal data
+            terminals_data = []
+            for i in terminals:
+                sp_tips = self.get_tip_names_from_tree(i)
+                if i.branch_length:
+                    terminals_data.append((i.name, i.branch_length, sp_tips))
+
+            # Prepare nonterminal data
+            nonterminals_data = []
+            for i in nonterminals:
+                if i.branch_length:
+                    sp_tips = self.get_tip_names_from_tree(i)
+                    nonterminals_data.append((sp_tips, i.branch_length))
+
+            # Process in batches
+            batch_size = max(10, (len(terminals_data) + len(nonterminals_data)) // 4)
+
+            with ProcessPoolExecutor(max_workers=min(4, len(terminals_data) + len(nonterminals_data) // 10)) as executor:
+                futures = []
+
+                # Submit terminal batches
+                for i in range(0, len(terminals_data), batch_size):
+                    batch = terminals_data[i:i+batch_size]
+                    futures.append(executor.submit(self._process_terminal_batch, tree0_pickle, tree1_pickle, batch))
+
+                # Submit nonterminal batches
+                for i in range(0, len(nonterminals_data), batch_size):
+                    batch = nonterminals_data[i:i+batch_size]
+                    futures.append(executor.submit(self._process_nonterminal_batch, tree0_pickle, tree1_pickle, batch))
+
+                # Collect results
+                for future in as_completed(futures):
+                    batch_results = future.result()
+                    for bl0, bl1, sp_tips in batch_results:
+                        l0.append(bl0)
+                        l1.append(bl1)
+                        tip_names.append(sp_tips)
+
+        return (l0, l1, tip_names)

phykit/services/tree/dvmc.py

@@ -0,0 +1,37 @@
+import math
+from typing import Dict
+import numpy as np
+
+from Bio.Phylo import Newick
+
+from .base import Tree
+
+
+class DVMC(Tree):
+    def __init__(self, args) -> None:
+        super().__init__(**self.process_args(args))
+
+    def run(self):
+        tree = self.read_tree_file()
+        dvmc = self.determine_dvmc(tree)
+        print(round(dvmc, 4))
+
+    def process_args(self, args) -> Dict[str, str]:
+        return dict(tree_file_path=args.tree)
+
+    def determine_dvmc(self, tree: Newick.Tree) -> float:
+        num_spp = tree.count_terminals()
+
+        # Collect all distances at once for vectorized operations
+        distances = np.array([tree.distance(term) for term in tree.get_terminals()])
+
+        # Calculate statistics using numpy
+        sum_dist = np.sum(distances)
+        sumi2N = np.sum(distances ** 2)
+        avg_dist = np.mean(distances)
+
+        # Calculate variance more efficiently
+        squared_diff_sum = sumi2N - num_spp * (avg_dist ** 2)
+
+        # Return standard deviation
+        return np.sqrt(squared_diff_sum / (num_spp - 1))

phykit/services/tree/evolutionary_rate.py

@@ -0,0 +1,17 @@
+from typing import Dict
+
+from .base import Tree
+
+
+class EvolutionaryRate(Tree):
+    def __init__(self, args) -> None:
+        super().__init__(**self.process_args(args))
+
+    def run(self) -> None:
+        tree = self.read_tree_file()
+        total_tree_length = tree.total_branch_length()
+        num_terminals = tree.count_terminals()
+        print(round(total_tree_length / num_terminals, 4))
+
+    def process_args(self, args) -> Dict[str, str]:
+        return dict(tree_file_path=args.tree)