phylogenie 2.1.29__py3-none-any.whl → 3.0.0__py3-none-any.whl

phylogenie/treesimulator/events/mutations.py

@@ -1,10 +1,9 @@
 import re
 from copy import deepcopy
-from typing import Any
+from typing import Any, Callable
 
 from numpy.random import Generator
 
-from phylogenie.models import Distribution
 from phylogenie.skyline import SkylineParameterLike
 from phylogenie.treesimulator.events.base import Event, EventType
 from phylogenie.treesimulator.events.contact_tracing import (
@@ -42,7 +41,7 @@ class Mutation(Event):
         self,
         state: str,
         rate: SkylineParameterLike,
-        rate_scalers: dict[EventType, Distribution],
+        rate_scalers: dict[EventType, Callable[[], float]],
         rates_to_log: list[EventType] | None = None,
     ):
         super().__init__(state, rate)
@@ -52,16 +51,16 @@ class Mutation(Event):
     def apply(
         self, model: Model, events: list[Event], time: float, rng: Generator
     ) -> dict[str, Any]:
-        if NEXT_MUTATION_ID not in model.context:
-            model.context[NEXT_MUTATION_ID] = 0
-        model.context[NEXT_MUTATION_ID] += 1
-        mutation_id = model.context[NEXT_MUTATION_ID]
+        if NEXT_MUTATION_ID not in model.metadata:
+            model[NEXT_MUTATION_ID] = 0
+        model[NEXT_MUTATION_ID] += 1
+        mutation_id = model[NEXT_MUTATION_ID]
 
         individual = self.draw_individual(model, rng)
         model.migrate(individual, _get_mutated_state(mutation_id, self.state), time)
 
         rate_scalers: dict[EventType, float] = {
-            target_type: getattr(rng, rate_scaler.type)(**rate_scaler.args)
+            target_type: rate_scaler()
             for target_type, rate_scaler in self.rate_scalers.items()
         }
 

phylogenie/treesimulator/features.py

@@ -1,10 +1,10 @@
 from collections.abc import Iterable
 from enum import Enum
 
-from phylogenie.tree import Tree
 from phylogenie.treesimulator.events.mutations import get_mutation_id
 from phylogenie.treesimulator.model import get_node_state
-from phylogenie.utils import (
+from phylogenie.treesimulator.tree import Tree
+from phylogenie.treesimulator.utils import (
     get_node_depth_levels,
     get_node_depths,
     get_node_height_levels,
@@ -46,4 +46,4 @@ def set_features(tree: Tree, features: Iterable[Feature]) -> None:
     for feature in features:
         feature_maps = FEATURES_EXTRACTORS[feature](tree)
         for node in tree:
-            node.set(feature.value, feature_maps[node])
+            node[feature.value] = feature_maps[node]

phylogenie/treesimulator/gillespie.py

@@ -1,7 +1,7 @@
 import time
 from collections.abc import Iterable, Sequence
 from pathlib import Path
-from typing import Any
+from typing import Any, Callable
 
 import joblib
 import numpy as np
@@ -9,26 +9,28 @@ import pandas as pd
 from numpy.random import default_rng
 from tqdm import tqdm
 
-from phylogenie.io import dump_newick
-from phylogenie.tree import Tree
 from phylogenie.treesimulator.events import Event
 from phylogenie.treesimulator.features import Feature, set_features
+from phylogenie.treesimulator.io import dump_newick
 from phylogenie.treesimulator.model import Model
+from phylogenie.treesimulator.tree import Tree
 
 
 def simulate_tree(
     events: Sequence[Event],
-    min_tips: int = 1,
-    max_tips: int | None = None,
+    n_tips: int | None = None,
     max_time: float = np.inf,
     init_state: str | None = None,
     sampling_probability_at_present: float = 0.0,
     seed: int | None = None,
     timeout: float = np.inf,
+    acceptance_criterion: Callable[[Tree], bool] | None = None,
 ) -> tuple[Tree, dict[str, Any]]:
-    if max_time == np.inf and sampling_probability_at_present:
+    if (max_time != np.inf) == (n_tips is not None):
+        raise ValueError("Exactly one of max_time or n_tips must be specified.")
+    if sampling_probability_at_present and max_time == np.inf:
         raise ValueError(
-            "sampling_probability_at_present cannot be set when max_time is infinite."
+            "sampling_probability_at_present can only be used with max_time."
         )
 
     states = {e.state for e in events if e.state}
@@ -51,32 +53,12 @@ def simulate_tree(
         change_times = sorted(set(t for e in events for t in e.rate.change_times))
         next_change_time = change_times.pop(0) if change_times else np.inf
 
-        if max_time == np.inf:
-            if max_tips is None:
-                raise ValueError("Either max_time or max_tips must be specified.")
-            target_n_tips = rng.integers(min_tips, max_tips + 1)
-        else:
-            target_n_tips = None
-
-        while current_time < max_time:
+        while current_time < max_time and (n_tips is None or model.n_sampled < n_tips):
             if time.perf_counter() - start_clock > timeout:
                 raise TimeoutError("Simulation timed out.")
 
             rates = [e.get_propensity(model, current_time) for e in run_events]
-
-            instantaneous_events = [e for e, r in zip(run_events, rates) if r == np.inf]
-            if instantaneous_events:
-                event = instantaneous_events[rng.integers(len(instantaneous_events))]
-                event.apply(model, run_events, current_time, rng)
-                continue
-
-            if (
-                not any(rates)
-                or max_tips is not None
-                and model.n_sampled >= max_tips
-                or target_n_tips is not None
-                and model.n_sampled >= target_n_tips
-            ):
+            if not any(rates):
                 break
 
             time_step = rng.exponential(1 / sum(rates))
@@ -95,22 +77,23 @@ def simulate_tree(
             if event_metadata is not None:
                 metadata.update(event_metadata)
 
+        if current_time != max_time and model.n_sampled != n_tips:
+            continue
+
         for individual in model.get_population():
             if rng.random() < sampling_probability_at_present:
                 model.sample(individual, current_time, True)
 
-        if min_tips <= model.n_sampled and (
-            max_tips is None or model.n_sampled <= max_tips
-        ):
-            return (model.get_sampled_tree(), metadata)
+        tree = model.get_sampled_tree()
+        if acceptance_criterion is None or acceptance_criterion(tree):
+            return (tree, metadata)
 
 
 def generate_trees(
     output_dir: str | Path,
     n_trees: int,
     events: Sequence[Event],
-    min_tips: int = 1,
-    max_tips: int | None = None,
+    n_tips: int | None = None,
     max_time: float = np.inf,
     init_state: str | None = None,
     sampling_probability_at_present: float = 0.0,
@@ -118,6 +101,7 @@ def generate_trees(
     seed: int | None = None,
     n_jobs: int = -1,
     timeout: float = np.inf,
+    acceptance_criterion: Callable[[Tree], bool] | None = None,
 ) -> pd.DataFrame:
     if isinstance(output_dir, str):
         output_dir = Path(output_dir)
@@ -130,13 +114,13 @@ def generate_trees(
             try:
                 tree, metadata = simulate_tree(
                     events=events,
-                    min_tips=min_tips,
-                    max_tips=max_tips,
+                    n_tips=n_tips,
                     max_time=max_time,
                     init_state=init_state,
                     sampling_probability_at_present=sampling_probability_at_present,
                     seed=seed,
                     timeout=timeout,
+                    acceptance_criterion=acceptance_criterion,
                 )
                 metadata["file_id"] = i
                 if node_features is not None:

phylogenie/treesimulator/io/__init__.py

@@ -0,0 +1,4 @@
+from phylogenie.treesimulator.io.newick import dump_newick, load_newick
+from phylogenie.treesimulator.io.nexus import load_nexus
+
+__all__ = ["dump_newick", "load_newick", "load_nexus"]

phylogenie/{io → treesimulator/io}/newick.py

@@ -1,7 +1,7 @@
 import re
 from pathlib import Path
 
-from phylogenie.tree import Tree
+from phylogenie.treesimulator.tree import Tree
 
 
 def parse_newick(newick: str, translations: dict[str, str] | None = None) -> Tree:
@@ -76,8 +76,8 @@ def load_newick(filepath: str | Path) -> Tree | list[Tree]:
 def to_newick(tree: Tree) -> str:
     children_newick = ",".join([to_newick(child) for child in tree.children])
     newick = tree.name
-    if tree.features:
-        reprs = {k: repr(v).replace("'", '"') for k, v in tree.features.items()}
+    if tree.metadata:
+        reprs = {k: repr(v).replace("'", '"') for k, v in tree.metadata.items()}
         for k, r in reprs.items():
             if "," in k or "=" in k or "]" in k:
                 raise ValueError(

phylogenie/{io → treesimulator/io}/nexus.py

@@ -2,19 +2,20 @@ import re
 from collections.abc import Iterator
 from pathlib import Path
 
-from phylogenie.io.newick import parse_newick
-from phylogenie.tree import Tree
+from phylogenie.treesimulator.io.newick import parse_newick
+from phylogenie.treesimulator.tree import Tree
 
 
 def _parse_translate_block(lines: Iterator[str]) -> dict[str, str]:
     translations: dict[str, str] = {}
     for line in lines:
-        match = re.match(r"\s*(\d+)\s+['\"]?([^'\",;]+)['\"]?", line)
+        line = line.strip()
+        match = re.match(r"(\d+)\s+['\"]?([^'\",;]+)['\"]?", line)
         if match is None:
             if ";" in line:
                 return translations
             else:
-                raise ValueError(f"Invalid translate line: {line.strip()}")
+                raise ValueError(f"Invalid translate line: {line}")
         translations[match.group(1)] = match.group(2)
     raise ValueError("Translate block not terminated with ';'")
 

phylogenie/treesimulator/model.py

@@ -1,8 +1,8 @@
 from collections import defaultdict
 from dataclasses import dataclass
-from typing import Any
 
-from phylogenie.tree import Tree
+from phylogenie.mixins import MetadataMixin
+from phylogenie.treesimulator.tree import Tree
 
 
 @dataclass
@@ -17,23 +17,22 @@ def _get_node_name(node_id: int, state: str) -> str:
 
 
 def get_node_state(node_name: str) -> str:
-    try:
-        return node_name.split("|")[1]
-    except IndexError:
+    if "|" not in node_name:
         raise ValueError(
             f"Invalid node name: {node_name} (expected format 'id|state')."
         )
+    return node_name.split("|")[-1]
 
 
-class Model:
+class Model(MetadataMixin):
     def __init__(self, init_state: str):
+        super().__init__()
         self._next_node_id = 0
         self._next_individual_id = 0
         self._population: dict[int, Individual] = {}
         self._states: dict[str, set[int]] = defaultdict(set)
         self._sampled: set[str] = set()
         self._tree = self._get_new_individual(init_state).node
-        self.context: dict[str, Any] = {}
 
     @property
     def n_sampled(self) -> int:
@@ -110,9 +109,7 @@ class Model:
             elif len(node.children) == 1:
                 (child,) = node.children
                 child.set_parent(node.parent)
-                assert child.branch_length is not None
-                assert node.branch_length is not None
-                child.branch_length += node.branch_length
+                child.branch_length += node.branch_length  # pyright: ignore
                 if node.parent is None:
                     return child
                 else:

phylogenie/{tree.py → treesimulator/tree.py}

@@ -2,14 +2,21 @@ from collections import deque
 from collections.abc import Callable, Iterator
 from typing import Any
 
+from phylogenie.mixins import MetadataMixin
 
-class Tree:
+
+class Tree(MetadataMixin):
     def __init__(self, name: str = "", branch_length: float | None = None):
+        super().__init__()
         self.name = name
         self.branch_length = branch_length
         self._parent: Tree | None = None
         self._children: list[Tree] = []
-        self._features: dict[str, Any] = {}
+
+    # ----------------
+    # Basic properties
+    # ----------------
+    # Properties related to parent-child relationships.
 
     @property
     def children(self) -> tuple["Tree", ...]:
@@ -19,59 +26,6 @@ class Tree:
     def parent(self) -> "Tree | None":
         return self._parent
 
-    @property
-    def features(self) -> dict[str, Any]:
-        return self._features.copy()
-
-    @property
-    def depth(self) -> float:
-        if self.parent is None:
-            return 0 if self.branch_length is None else self.branch_length
-        if self.branch_length is None:
-            raise ValueError(f"Branch length of node {self.name} is not set.")
-        return self.parent.depth + self.branch_length
-
-    @property
-    def depth_level(self) -> int:
-        if self.parent is None:
-            return 0
-        return self.parent.depth_level + 1
-
-    @property
-    def height(self) -> float:
-        if self.is_leaf():
-            return 0.0
-        if any(child.branch_length is None for child in self.children):
-            raise ValueError(
-                f"Branch length of one or more children of node {self.name} is not set."
-            )
-        return max(
-            child.branch_length + child.height  # pyright: ignore
-            for child in self.children
-        )
-
-    @property
-    def height_level(self) -> int:
-        if self.is_leaf():
-            return 0
-        return 1 + max(child.height_level for child in self.children)
-
-    @property
-    def n_leaves(self) -> int:
-        return len(self.get_leaves())
-
-    def set(self, key: str, value: Any) -> None:
-        self._features[key] = value
-
-    def update_features(self, features: dict[str, Any]) -> None:
-        self._features.update(features)
-
-    def get(self, key: str) -> Any:
-        return self._features[key]
-
-    def delete(self, key: str) -> None:
-        del self._features[key]
-
     def add_child(self, child: "Tree") -> "Tree":
         if child.parent is not None:
             raise ValueError(f"Node {child.name} already has a parent.")
@@ -90,6 +44,53 @@ class Tree:
         if parent is not None:
             parent._children.append(self)
 
+    def is_leaf(self) -> bool:
+        return not self.children
+
+    def get_leaves(self) -> tuple["Tree", ...]:
+        return tuple(node for node in self if node.is_leaf())
+
+    def is_internal(self) -> bool:
+        return not self.is_leaf()
+
+    def get_internal_nodes(self) -> tuple["Tree", ...]:
+        return tuple(node for node in self if node.is_internal())
+
+    def is_binary(self) -> bool:
+        return all(len(node.children) in (0, 2) for node in self)
+
+    # --------------
+    # Tree traversal
+    # --------------
+    # Methods for traversing the tree in various orders.
+
+    def iter_ancestors(self, stop: "Tree | None" = None) -> Iterator["Tree"]:
+        node = self
+        while True:
+            if node.parent is None:
+                if stop is None:
+                    return
+                raise ValueError("Reached root without encountering stop node.")
+            node = node.parent
+            if node == stop:
+                return
+            yield node
+
+    def iter_upward(self, stop: "Tree | None" = None) -> Iterator["Tree"]:
+        if self == stop:
+            return
+        yield self
+        yield from self.iter_ancestors(stop=stop)
+
+    def iter_descendants(self) -> Iterator["Tree"]:
+        for child in self.children:
+            yield child
+            yield from child.iter_descendants()
+
+    def preorder_traversal(self) -> Iterator["Tree"]:
+        yield self
+        yield from self.iter_descendants()
+
     def inorder_traversal(self) -> Iterator["Tree"]:
         if self.is_leaf():
             yield self
@@ -101,22 +102,11 @@ class Tree:
         yield self
         yield from right.inorder_traversal()
 
-    def preorder_traversal(self) -> Iterator["Tree"]:
-        yield self
-        for child in self.children:
-            yield from child.preorder_traversal()
-
     def postorder_traversal(self) -> Iterator["Tree"]:
         for child in self.children:
             yield from child.postorder_traversal()
         yield self
 
-    def iter_ancestors(self, stop: "Tree | None" = None) -> Iterator["Tree"]:
-        node = self
-        while node is not None and node is not stop:
-            yield node
-            node = node.parent
-
     def breadth_first_traversal(self) -> Iterator["Tree"]:
         queue: deque["Tree"] = deque([self])
         while queue:
@@ -124,39 +114,75 @@ class Tree:
             yield node
             queue.extend(node.children)
 
-    def get_node(self, name: str) -> "Tree":
-        for node in self:
-            if node.name == name:
-                return node
-        raise ValueError(f"Node with name {name} not found.")
+    # ---------------
+    # Tree properties
+    # ---------------
+    # Properties and methods related to tree metrics like leaf count, depth, height, etc.
 
-    def is_leaf(self) -> bool:
-        return not self.children
+    @property
+    def n_leaves(self) -> int:
+        return len(self.get_leaves())
 
-    def get_leaves(self) -> tuple["Tree", ...]:
-        return tuple(node for node in self if node.is_leaf())
+    def branch_length_or_raise(self) -> float:
+        if self.parent is None:
+            return 0 if self.branch_length is None else self.branch_length
+        if self.branch_length is None:
+            raise ValueError(f"Branch length of node {self.name} is not set.")
+        return self.branch_length
 
-    def is_internal(self) -> bool:
-        return not self.is_leaf()
+    @property
+    def depth_level(self) -> int:
+        return 0 if self.parent is None else self.parent.depth_level + 1
 
-    def get_internal_nodes(self) -> tuple["Tree", ...]:
-        return tuple(node for node in self if node.is_internal())
+    @property
+    def depth(self) -> float:
+        parent_depth = 0 if self.parent is None else self.parent.depth
+        return parent_depth + self.branch_length_or_raise()
 
-    def is_binary(self) -> bool:
-        return all(len(node.children) in (0, 2) for node in self)
+    @property
+    def height_level(self) -> int:
+        if self.is_leaf():
+            return 0
+        return 1 + max(child.height_level for child in self.children)
 
-    def ladderize(self, key: Callable[["Tree"], Any]) -> None:
+    @property
+    def height(self) -> float:
+        if self.is_leaf():
+            return 0.0
+        return max(
+            child.branch_length_or_raise() + child.height for child in self.children
+        )
+
+    # -------------
+    # Miscellaneous
+    # -------------
+    # Other useful miscellaneous methods.
+
+    def ladderize(self, key: Callable[["Tree"], Any] | None = None) -> None:
+        if key is None:
+            key = lambda node: node.n_leaves
         self._children.sort(key=key)
         for child in self.children:
             child.ladderize(key)
 
+    def get_node(self, name: str) -> "Tree":
+        for node in self:
+            if node.name == name:
+                return node
+        raise ValueError(f"Node with name {name} not found.")
+
     def copy(self):
         new_tree = Tree(self.name, self.branch_length)
-        new_tree.update_features(self._features)
+        new_tree.update(self.metadata)
         for child in self.children:
             new_tree.add_child(child.copy())
         return new_tree
 
+    # ----------------
+    # Dunder methods
+    # ----------------
+    # Special methods for standard behaviors like iteration, length, and representation.
+
     def __iter__(self) -> Iterator["Tree"]:
         return self.preorder_traversal()
 
@@ -164,4 +190,4 @@ class Tree:
         return sum(1 for _ in self)
 
     def __repr__(self) -> str:
-        return f"TreeNode(name='{self.name}', branch_length={self.branch_length}, features={self.features})"
+        return f"TreeNode(name='{self.name}', branch_length={self.branch_length}, metadata={self.metadata})"

phylogenie/treesimulator/utils.py

@@ -0,0 +1,108 @@
+import math
+from math import comb
+
+from phylogenie.treesimulator.tree import Tree
+
+
+def get_node_leaf_counts(tree: Tree) -> dict[Tree, int]:
+    n_leaves: dict[Tree, int] = {}
+    for node in tree.postorder_traversal():
+        n_leaves[node] = sum(n_leaves[child] for child in node.children) or 1
+    return n_leaves
+
+
+def get_node_depth_levels(tree: Tree) -> dict[Tree, int]:
+    depth_levels: dict[Tree, int] = {tree: tree.depth_level}
+    for node in tree.iter_descendants():
+        depth_levels[node] = depth_levels[node.parent] + 1  # pyright: ignore
+    return depth_levels
+
+
+def get_node_depths(tree: Tree) -> dict[Tree, float]:
+    depths: dict[Tree, float] = {tree: tree.depth}
+    for node in tree.iter_descendants():
+        parent_depth = depths[node.parent]  # pyright: ignore
+        depths[node] = tree.branch_length_or_raise() + parent_depth
+    return depths
+
+
+def get_node_height_levels(tree: Tree) -> dict[Tree, int]:
+    height_levels: dict[Tree, int] = {}
+    for node in tree.postorder_traversal():
+        height_levels[node] = (
+            0
+            if node.is_leaf()
+            else max(1 + height_levels[child] for child in node.children)
+        )
+    return height_levels
+
+
+def get_node_heights(tree: Tree) -> dict[Tree, float]:
+    heights: dict[Tree, float] = {}
+    for node in tree.postorder_traversal():
+        heights[node] = (
+            0
+            if node.is_leaf()
+            else max(
+                child.branch_length_or_raise() + heights[child]
+                for child in node.children
+            )
+        )
+    return heights
+
+
+def get_mrca(node1: Tree, node2: Tree) -> Tree:
+    node1_ancestors = set(node1.iter_upward())
+    for node2_ancestor in node2.iter_upward():
+        if node2_ancestor in node1_ancestors:
+            return node2_ancestor
+    raise ValueError(f"No common ancestor found between node {node1} and node {node2}.")
+
+
+def get_path(node1: Tree, node2: Tree) -> list[Tree]:
+    mrca = get_mrca(node1, node2)
+    return [
+        *node1.iter_upward(stop=mrca.parent),
+        *reversed(list(node2.iter_upward(stop=mrca))),
+    ]
+
+
+def count_hops(node1: Tree, node2: Tree) -> int:
+    return len(get_path(node1, node2)) - 1
+
+
+def get_distance(node1: Tree, node2: Tree) -> float:
+    mrca = get_mrca(node1, node2)
+    path = get_path(node1, node2)
+    path.remove(mrca)
+    return sum(node.branch_length_or_raise() for node in path)
+
+
+def compute_sackin_index(tree: Tree, normalize: bool = False) -> float:
+    depth_levels = get_node_depth_levels(tree)
+    sackin_index = sum(dl for node, dl in depth_levels.items() if node.is_leaf())
+    if normalize:
+        if not tree.is_binary():
+            raise ValueError(
+                "Normalized Sackin index is only defined for binary trees."
+            )
+        n = tree.n_leaves
+        h = math.floor(math.log2(n))
+        min_sackin_index = n * (h + 2) - 2 ** (h + 1)
+        max_sackin_index = n * (n - 1) / 2
+        return (sackin_index - min_sackin_index) / (max_sackin_index - min_sackin_index)
+    return sackin_index
+
+
+def compute_mean_leaf_pairwise_distance(tree: Tree) -> float:
+    leaves = tree.get_leaves()
+    n_leaves = len(leaves)
+    if n_leaves < 2:
+        return 0.0
+
+    total_distance = sum(
+        get_distance(leaves[i], leaves[j])
+        for i in range(n_leaves)
+        for j in range(i + 1, n_leaves)
+    )
+    return total_distance / comb(n_leaves, 2)

{phylogenie-2.1.29.dist-info → phylogenie-3.0.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: phylogenie
-Version: 2.1.29
+Version: 3.0.0
 Summary: Generate phylogenetic datasets with minimal setup effort
 Author: Gabriele Marino
 Author-email: gabmarino.8601@gmail.com