phykit 2.1.0__py3-none-any.whl

phykit/services/tree/rf_distance.py

@@ -0,0 +1,136 @@
+from typing import Dict, List, Set, Tuple
+from functools import lru_cache
+from concurrent.futures import ProcessPoolExecutor
+import pickle
+
+from Bio.Phylo import Newick
+
+from .base import Tree
+
+
+class RobinsonFouldsDistance(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()
+
+        # get shared tree tip names - use sets for efficiency
+        tree_zero_tips = set(self.get_tip_names_from_tree(tree_zero))
+        tree_one_tips = set(self.get_tip_names_from_tree(tree_one))
+        shared_tree_tips = tree_zero_tips & tree_one_tips
+
+        # prune to common set - already have sets
+        tree_zero_tips_to_prune = list(tree_zero_tips - shared_tree_tips)
+        tree_one_tips_to_prune = list(tree_one_tips - shared_tree_tips)
+
+        if tree_zero_tips_to_prune:
+            tree_zero = self.prune_tree_using_taxa_list(tree_zero, tree_zero_tips_to_prune)
+        if tree_one_tips_to_prune:
+            tree_one = self.prune_tree_using_taxa_list(tree_one, tree_one_tips_to_prune)
+
+        # Get first terminal for rooting
+        tip_for_rooting = tree_zero.get_terminals()[0].name
+        tree_zero.root_with_outgroup(tip_for_rooting)
+        tree_one.root_with_outgroup(tip_for_rooting)
+
+        plain_rf, normalized_rf = self.calculate_robinson_foulds_distance(
+            tree_zero, tree_one
+        )
+
+        print(f"{plain_rf}\t{round(normalized_rf, 4)}")
+
+    def process_args(self, args) -> Dict[str, str]:
+        return dict(
+            tree_file_path=args.tree_zero,
+            tree1_file_path=args.tree_one,
+        )
+
+    def calculate_robinson_foulds_distance(self, tree_zero, tree_one):
+        plain_rf = 0
+        plain_rf = self.compare_trees_optimized(plain_rf, tree_zero, tree_one)
+        plain_rf = self.compare_trees_optimized(plain_rf, tree_one, tree_zero)
+
+        tip_count = tree_zero.count_terminals()
+        normalized_rf = plain_rf / (2 * (tip_count - 3))
+
+        return plain_rf, normalized_rf
+
+    def compare_trees_optimized(
+        self,
+        plain_rf: int,
+        tree_zero: Newick.Tree,
+        tree_one: Newick.Tree
+    ) -> int:
+        # Cache tip names for clades to avoid recomputation
+        tip_names_cache = {}
+
+        def get_cached_tips(clade):
+            clade_id = id(clade)
+            if clade_id not in tip_names_cache:
+                tip_names_cache[clade_id] = frozenset(self.get_tip_names_from_tree(clade))
+            return tip_names_cache[clade_id]
+
+        # loop through tree_zero and find similar clade in tree_one
+        for clade_zero in tree_zero.get_nonterminals()[1:]:
+            # Get tip names from tree_zero clade
+            tip_names_zero = get_cached_tips(clade_zero)
+            # get common ancestor of tree_zero tip names in tree_one
+            clade_one = tree_one.common_ancestor(list(tip_names_zero))
+            # Get tip names from tree_one clade
+            tip_names_one = get_cached_tips(clade_one)
+            # compare the list of tip names
+            if tip_names_zero != tip_names_one:
+                plain_rf += 1
+
+        return plain_rf
+
+    @staticmethod
+    def _calculate_rf_batch(tree_pairs_pickle):
+        """Calculate RF distance for a batch of tree pairs in parallel."""
+        tree_pairs = pickle.loads(tree_pairs_pickle)
+        results = []
+
+        for tree_zero, tree_one in tree_pairs:
+            rf_calc = RobinsonFouldsDistance.__new__(RobinsonFouldsDistance)
+            rf_calc.__dict__.update({'tree_format': 'newick'})
+
+            # Calculate bipartitions
+            bipartitions_zero = rf_calc.get_all_bipartitions(tree_zero)
+            bipartitions_one = rf_calc.get_all_bipartitions(tree_one)
+
+            # Calculate RF distance
+            plain_rf = len(bipartitions_zero ^ bipartitions_one)  # Symmetric difference
+            tip_count = tree_zero.count_terminals()
+            normalized_rf = plain_rf / (2 * (tip_count - 3))
+
+            results.append((plain_rf, normalized_rf))
+
+        return results
+
+    def calculate_multiple_rf_distances(self, tree_pairs: List[Tuple]) -> List[Tuple[int, float]]:
+        """Calculate RF distances for multiple tree pairs in parallel."""
+        if len(tree_pairs) < 5:
+            # Sequential for small datasets
+            results = []
+            for tree_zero, tree_one in tree_pairs:
+                plain_rf, normalized_rf = self.calculate_robinson_foulds_distance(tree_zero, tree_one)
+                results.append((plain_rf, normalized_rf))
+            return results
+
+        # Parallel processing for larger datasets
+        batch_size = max(2, len(tree_pairs) // 4)
+        batches = [tree_pairs[i:i + batch_size] for i in range(0, len(tree_pairs), batch_size)]
+
+        with ProcessPoolExecutor(max_workers=min(4, len(batches))) as executor:
+            futures = []
+            for batch in batches:
+                batch_pickle = pickle.dumps(batch)
+                futures.append(executor.submit(self._calculate_rf_batch, batch_pickle))
+
+            all_results = []
+            for future in futures:
+                all_results.extend(future.result())
+
+        return all_results

phykit/services/tree/root_tree.py

@@ -0,0 +1,35 @@
+import copy
+from Bio import Phylo
+
+from .base import Tree
+
+from ...helpers.files import read_single_column_file_to_list
+
+
+class RootTree(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)
+
+        outgroup = \
+            read_single_column_file_to_list(self.outgroup_taxa_file_path)
+
+        Phylo.BaseTree.Tree.root_with_outgroup(tree_copy, outgroup)
+
+        self.write_tree_file(tree_copy, self.output_file_path)
+
+    def process_args(self, args):
+        tree_file_path = args.tree
+
+        output_file_path = \
+            args.output if args.output else f"{tree_file_path}.rooted"
+
+        return dict(
+            tree_file_path=tree_file_path,
+            outgroup_taxa_file_path=args.root,
+            output_file_path=output_file_path,
+        )

phykit/services/tree/saturation.py

@@ -0,0 +1,209 @@
+from enum import Enum
+import itertools
+import sys
+from typing import Dict, List, Tuple
+import multiprocessing as mp
+from functools import partial
+
+from Bio import Align
+from Bio.Phylo import Newick
+import numpy as np
+from sklearn.linear_model import LinearRegression
+
+from .base import Tree
+from ...helpers.files import (
+    get_alignment_and_format as get_alignment_and_format_helper
+)
+
+
+class FileFormat(Enum):
+    fasta = "fasta"
+    clustal = "clustal"
+    maf = "maf"
+    mauve = "mauve"
+    phylip = "phylip"
+    phylip_seq = "phylip-sequential"
+    phylip_rel = "phylip-relaxed"
+    stockholm = "stockholm"
+
+
+class Saturation(Tree):
+    def __init__(self, args) -> None:
+        super().__init__(**self.process_args(args))
+
+    def run(self) -> None:
+        alignment, _, is_protein = get_alignment_and_format_helper(
+            self.alignment_file_path
+        )
+
+        tree = self.read_tree_file()
+
+        tips = self.get_tip_names_from_tree(tree)
+        combos = list(itertools.combinations(tips, 2))
+
+        (
+            patristic_distances,
+            uncorrected_distances,
+        ) = self.loop_through_combos_and_calculate_pds_and_pis(
+            combos, alignment, tree, self.exclude_gaps
+        )
+
+        # calculate slope and fit the y-intercept to zero
+        # Fitting the y-intercept to zero follows Jeffroy et al.
+        # See fig 2 https://www.cell.com/trends/genetics/fulltext/S0168-9525(06)00051-5
+        model = LinearRegression(fit_intercept=False)
+        model.fit(
+            np.array(patristic_distances).reshape(-1, 1),
+            np.array(uncorrected_distances)
+        )
+
+        self.print_res(
+            self.verbose, combos, uncorrected_distances, patristic_distances, model.coef_[0]
+        )
+
+    def process_args(self, args) -> Dict[str, str]:
+        return dict(
+            tree_file_path=args.tree,
+            alignment_file_path=args.alignment,
+            exclude_gaps=args.exclude_gaps,
+            verbose=args.verbose,
+        )
+
+    def _process_combo_batch(self, tree, seq_arrays, gap_mask, exclude_gaps, combo_batch):
+        """Process a batch of combinations in parallel."""
+        results = []
+        for combo in combo_batch:
+            # Calculate patristic distance
+            pd = tree.distance(combo[0], combo[1])
+
+            # Calculate uncorrected distance using numpy operations
+            seq1_arr = seq_arrays[combo[0]]
+            seq2_arr = seq_arrays[combo[1]]
+
+            if exclude_gaps:
+                # Use pre-computed gap masks
+                gap_mask1 = gap_mask[combo[0]]
+                gap_mask2 = gap_mask[combo[1]]
+                valid_positions = ~(gap_mask1 | gap_mask2)
+
+                if np.any(valid_positions):
+                    matches = seq1_arr[valid_positions] == seq2_arr[valid_positions]
+                    identities = np.sum(matches)
+                    adjusted_len = np.sum(valid_positions)
+                    ud = 1 - (identities / adjusted_len)
+                else:
+                    ud = float('nan')
+            else:
+                matches = seq1_arr == seq2_arr
+                identities = np.sum(matches)
+                ud = 1 - (identities / len(seq1_arr))
+
+            results.append((pd, ud))
+        return results
+
+    def loop_through_combos_and_calculate_pds_and_pis(
+        self,
+        combos: List[Tuple[str, str]],
+        alignment: Align.MultipleSeqAlignment,
+        tree: Newick.Tree,
+        exclude_gaps: bool,
+    ) -> Tuple[
+        List[float],
+        List[float]
+    ]:
+        """
+        loop through all taxon combinations and determine
+        their patristic distance and pairwise identity
+        """
+        gap_chars = self.get_gap_chars()
+
+        # Convert sequences to numpy arrays for vectorized operations
+        seq_arrays = {}
+        gap_mask = {}
+        for record in alignment:
+            seq_arr = np.array([c.upper() for c in str(record.seq)], dtype='U1')
+            seq_arrays[record.name] = seq_arr
+            if exclude_gaps:
+                gap_mask[record.name] = np.isin(seq_arr, list(gap_chars))
+
+        # For small datasets, process sequentially
+        if len(combos) < 50:
+            patristic_distances = []
+            uncorrected_distances = []
+            for combo in combos:
+                pd = tree.distance(combo[0], combo[1])
+                patristic_distances.append(pd)
+
+                seq1_arr = seq_arrays[combo[0]]
+                seq2_arr = seq_arrays[combo[1]]
+
+                if exclude_gaps:
+                    gap_mask1 = gap_mask[combo[0]]
+                    gap_mask2 = gap_mask[combo[1]]
+                    valid_positions = ~(gap_mask1 | gap_mask2)
+
+                    if np.any(valid_positions):
+                        matches = seq1_arr[valid_positions] == seq2_arr[valid_positions]
+                        identities = np.sum(matches)
+                        adjusted_len = np.sum(valid_positions)
+                        ud = 1 - (identities / adjusted_len)
+                    else:
+                        ud = float('nan')
+                else:
+                    matches = seq1_arr == seq2_arr
+                    identities = np.sum(matches)
+                    ud = 1 - (identities / len(seq1_arr))
+
+                uncorrected_distances.append(ud)
+        else:
+            # Use multiprocessing for larger datasets
+            num_workers = min(mp.cpu_count(), 8)
+            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
+            process_func = partial(
+                self._process_combo_batch,
+                tree,
+                seq_arrays,
+                gap_mask,
+                exclude_gaps
+            )
+
+            # Process in parallel
+            with mp.Pool(processes=num_workers) as pool:
+                chunk_results = pool.map(process_func, combo_chunks)
+
+            # Flatten results
+            patristic_distances = []
+            uncorrected_distances = []
+            for chunk_result in chunk_results:
+                for pd, ud in chunk_result:
+                    patristic_distances.append(pd)
+                    uncorrected_distances.append(ud)
+
+        return patristic_distances, uncorrected_distances
+
+    def print_res(
+        self,
+        verbose: bool,
+        combos: List[str],
+        uncorrected_distances: List[float],
+        patristic_distances: List[float],
+        slope: float,
+    ) -> None:
+        """
+        print results to stdout
+        """
+        try:
+            if verbose:
+                for cbo, dist, pd in zip(
+                    combos, uncorrected_distances, patristic_distances
+                ):
+                    print(
+                        f"{cbo[0]}\t{cbo[1]}\t{round(dist,4)}\t{round(pd, 4)}"
+                    )
+            else:
+                print(f"{round(slope, 4)}\t{abs(round(1-slope, 4))}")
+        except BrokenPipeError:
+            pass

phykit/services/tree/spurious_sequence.py

@@ -0,0 +1,75 @@
+import statistics as stat
+from typing import Dict, List, Tuple
+
+from Bio.Phylo import Newick
+
+from .base import Tree
+
+
+class SpuriousSequence(Tree):
+    def __init__(self, args) -> None:
+        super().__init__(**self.process_args(args))
+
+    def run(self) -> None:
+        tree = self.read_tree_file()
+        name_and_branch_len, threshold, median = \
+            self.identify_spurious_sequence(
+                tree, self.factor
+            )
+
+        counter = 0
+        for name, length in name_and_branch_len.items():
+            if length >= threshold:
+                try:
+                    print(
+                        f"{name}\t{round(length, 4)}\t{round(threshold, 4)}\t{round(median, 4)}"
+                    )
+                except BrokenPipeError:
+                    pass
+                counter += 1
+
+        if counter == 0:
+            print("None")
+
+    def process_args(self, args) -> Dict[str, str]:
+        return dict(
+            tree_file_path=args.tree,
+            factor=args.factor or 20
+        )
+
+    def identify_spurious_sequence(
+        self,
+        tree: Newick.Tree,
+        factor: float,
+    ) -> Tuple[
+        Dict[str, float],
+        float,
+        float
+    ]:
+        branch_lengths, name_and_branch_len = \
+            self.get_branch_lengths_and_their_names(tree)
+
+        median = stat.median(branch_lengths)
+
+        threshold = median * factor
+
+        return name_and_branch_len, threshold, median
+
+    def get_branch_lengths_and_their_names(
+        self,
+        tree: Newick.Tree,
+    ) -> Tuple[
+        List[float],
+        Dict[str, float],
+    ]:
+        branch_lengths = []
+        name_and_branch_len = {}
+
+        # collect terminal branch lengths only for spurious sequence detection
+        # (internal branches are not considered for spurious sequence detection)
+        for terminal in tree.get_terminals():
+            if terminal.branch_length is not None:
+                branch_lengths.append(terminal.branch_length)
+                name_and_branch_len[terminal.name] = terminal.branch_length
+
+        return branch_lengths, name_and_branch_len

phykit/services/tree/terminal_branch_stats.py

@@ -0,0 +1,87 @@
+import sys
+from typing import Dict, List, Tuple, Union
+
+from Bio.Phylo import Newick
+
+from .base import Tree
+
+from ...helpers.stats_summary import (
+    calculate_summary_statistics_from_arr,
+    print_summary_statistics,
+)
+
+
+class TerminalBranchStats(Tree):
+    def __init__(self, args) -> None:
+        super().__init__(**self.process_args(args))
+
+    def run(self):
+        tree = self.read_tree_file()
+        _, stats, lengths_and_names = \
+            self.calculate_terminal_branch_stats(tree)
+
+        if self.verbose:
+            try:
+                for len_and_name in lengths_and_names:
+                    print(round(len_and_name[0], 4), len_and_name[1])
+            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 get_terminal_branch_lengths(
+        self,
+        tree: Newick.Tree,
+    ) -> List[
+        Union[
+            float,
+            List[List[Union[float, str]]],
+        ]
+    ]:
+        """
+        loop through tree and get all terminal branch lengths
+        """
+        terminal_branch_lengths = []
+        lengths_and_names = []
+        for terminal_branch in tree.get_terminals():
+            if terminal_branch.branch_length is not None:
+                temp = []
+                temp.append(terminal_branch.branch_length)
+                terminal_branch_lengths.append(terminal_branch.branch_length)
+                temp.append(terminal_branch.name)
+                lengths_and_names.append(temp)
+
+        return terminal_branch_lengths, lengths_and_names
+
+    def check_tree_has_branch_lengths(
+        self,
+        terminal_branch_lengths: List[float],
+    ) -> None:
+        """
+        if tree has no branch lengths, exit
+        """
+        if len(terminal_branch_lengths) == 0:
+            print(
+                "Calculating terminal branch statistics requires a phylogeny with branch lengths."
+            )
+            sys.exit(2)
+
+    def calculate_terminal_branch_stats(
+        self,
+        tree: Newick.Tree,
+    ) -> Tuple[
+        List[float],
+        Dict[str, float],
+        List[List[Union[float, str]]],
+    ]:
+        terminal_branch_lengths, lengths_and_names = \
+            self.get_terminal_branch_lengths(tree)
+
+        self.check_tree_has_branch_lengths(terminal_branch_lengths)
+
+        stats = calculate_summary_statistics_from_arr(terminal_branch_lengths)
+
+        return terminal_branch_lengths, stats, lengths_and_names

phykit/services/tree/tip_labels.py

@@ -0,0 +1,18 @@
+from typing import Dict
+
+from .base import Tree
+
+
+class TipLabels(Tree):
+    def __init__(self, args) -> None:
+        super().__init__(**self.process_args(args))
+
+    def run(self) -> None:
+        tree = self.read_tree_file()
+        try:
+            print("\n".join([tip.name for tip in tree.get_terminals()]))
+        except BrokenPipeError:
+            pass
+
+    def process_args(self, args) -> Dict[str, str]:
+        return dict(tree_file_path=args.tree)

phykit/services/tree/tip_to_tip_distance.py

@@ -0,0 +1,41 @@
+import sys
+from typing import Dict
+
+from Bio.Phylo.BaseTree import TreeMixin
+from Bio.Phylo import Newick
+
+from .base import Tree
+
+
+class TipToTipDistance(Tree):
+    def __init__(self, args) -> None:
+        super().__init__(**self.process_args(args))
+
+    def run(self):
+        tree_zero = self.read_tree_file()
+
+        self.check_leaves(tree_zero, self.tip_1, self.tip_2)
+
+        print(round(TreeMixin.distance(tree_zero, self.tip_1, self.tip_2), 4))
+
+    def check_leaves(
+        self,
+        tree_zero: Newick.Tree,
+        tip_1: str,
+        tip_2: str,
+    ) -> None:
+        leaf1 = TreeMixin.find_any(tree_zero, tip_1)
+        if not bool(leaf1):
+            print(tip_1, "not on tree\nExiting...")
+            sys.exit(2)
+        leaf2 = TreeMixin.find_any(tree_zero, tip_2)
+        if not bool(leaf2):
+            print(tip_2, "not on tree\nExiting...")
+            sys.exit(2)
+
+    def process_args(self, args) -> Dict[str, str]:
+        return dict(
+            tree_file_path=args.tree_zero,
+            tip_1=args.tip_1,
+            tip_2=args.tip_2,
+        )

phykit/services/tree/tip_to_tip_node_distance.py

@@ -0,0 +1,41 @@
+import sys
+from typing import Dict
+
+from Bio.Phylo.BaseTree import TreeMixin
+from Bio.Phylo import Newick
+
+from .base import Tree
+
+
+class TipToTipNodeDistance(Tree):
+    def __init__(self, args) -> None:
+        super().__init__(**self.process_args(args))
+
+    def run(self) -> None:
+        tree_zero = self.read_tree_file()
+
+        self.check_leaves(tree_zero, self.tip_1, self.tip_2)
+
+        print(len(TreeMixin.trace(tree_zero, self.tip_1, self.tip_2)))
+
+    def check_leaves(
+        self,
+        tree_zero: Newick.Tree,
+        tip_1: str,
+        tip_2: str,
+    ) -> None:
+        leaf1 = TreeMixin.find_any(tree_zero, tip_1)
+        if not bool(leaf1):
+            print(tip_1, "not on tree\nExiting...")
+            sys.exit(2)
+        leaf2 = TreeMixin.find_any(tree_zero, tip_2)
+        if not bool(leaf2):
+            print(tip_2, "not on tree\nExiting...")
+            sys.exit(2)
+
+    def process_args(self, args) -> Dict[str, str]:
+        return dict(
+            tree_file_path=args.tree_zero,
+            tip_1=args.tip_1,
+            tip_2=args.tip_2,
+        )

phykit/services/tree/total_tree_length.py

@@ -0,0 +1,25 @@
+from typing import Dict, Union
+
+from Bio.Phylo import Newick
+from .base import Tree
+
+
+class TotalTreeLength(Tree):
+    def __init__(self, args) -> None:
+        super().__init__(**self.process_args(args))
+
+    def run(self) -> None:
+        tree = self.read_tree_file()
+        print(round(self.calculate_total_tree_length(tree), 4))
+
+    def process_args(self, args) -> Dict[str, str]:
+        return dict(tree_file_path=args.tree)
+
+    def calculate_total_tree_length(
+        self,
+        tree: Newick.Tree
+    ) -> Union[int, float]:
+        total_len = tree.total_branch_length()
+
+        if isinstance(total_len, (int, float)):
+            return total_len

phykit/services/tree/treeness.py

@@ -0,0 +1,16 @@
+from typing import Dict
+
+from .base import Tree
+
+
+class Treeness(Tree):
+    def __init__(self, args) -> None:
+        super().__init__(**self.process_args(args))
+
+    def run(self) -> None:
+        tree = self.read_tree_file()
+        treeness = self.calculate_treeness(tree)
+        print(round(treeness, 4))
+
+    def process_args(self, args) -> Dict[str, str]:
+        return dict(tree_file_path=args.tree)