phylogenie 2.1.4__py3-none-any.whl → 3.1.7__py3-none-any.whl

phylogenie/treesimulator/gillespie.py

@@ -1,35 +1,36 @@
-import os
 import time
-from collections.abc import Sequence
+from collections.abc import Iterable, Sequence
+from pathlib import Path
+from typing import Any, Callable
 
 import joblib
 import numpy as np
+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.model import Event, Model
-
-MAX_TIPS = 2**32
+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 = MAX_TIPS,
+    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,
-) -> Tree:
-    if max_time == np.inf and max_tips == MAX_TIPS:
-        raise ValueError("Either max_time or max_tips must be specified.")
-
-    if max_time == np.inf and sampling_probability_at_present:
+    acceptance_criterion: Callable[[Tree], bool] | None = None,
+) -> tuple[Tree, dict[str, Any]]:
+    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}
@@ -45,31 +46,19 @@ def simulate_tree(
     rng = default_rng(seed)
     start_clock = time.perf_counter()
     while True:
-        model = Model(init_state, events)
+        model = Model(init_state)
+        metadata: dict[str, Any] = {}
+        run_events = list(events)
         current_time = 0.0
         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
-        target_n_tips = rng.integers(min_tips, max_tips) if max_time == np.inf else 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.")
 
-            events = model.events
-            rates = [e.get_propensity(model, current_time) for e in events]
-
-            instantaneous_events = [e for e, r in zip(events, rates) if r == np.inf]
-            if instantaneous_events:
-                event = instantaneous_events[rng.integers(len(instantaneous_events))]
-                event.apply(model, current_time, rng)
-                continue
-
-            if (
-                not any(rates)
-                or model.n_sampled > max_tips
-                or target_n_tips is not None
-                and model.n_sampled >= target_n_tips
-            ):
+            rates = [e.get_propensity(model, current_time) for e in run_events]
+            if not any(rates):
                 break
 
             time_step = rng.exponential(1 / sum(rates))
@@ -83,56 +72,71 @@ def simulate_tree(
             current_time += time_step
 
             event_idx = np.searchsorted(np.cumsum(rates) / sum(rates), rng.random())
-            events[int(event_idx)].apply(model, current_time, rng)
+            event = run_events[int(event_idx)]
+            event_metadata = event.apply(model, run_events, current_time, rng)
+            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 <= max_tips:
-            return model.get_sampled_tree()
+        tree = model.get_sampled_tree()
+        if acceptance_criterion is None or acceptance_criterion(tree):
+            return (tree, metadata)
 
 
 def generate_trees(
-    output_dir: str,
+    output_dir: str | Path,
     n_trees: int,
-    events: list[Event],
-    min_tips: int = 1,
-    max_tips: int = 2**32,
+    events: Sequence[Event],
+    n_tips: int | None = None,
     max_time: float = np.inf,
     init_state: str | None = None,
     sampling_probability_at_present: float = 0.0,
+    node_features: Iterable[Feature] | None = None,
     seed: int | None = None,
     n_jobs: int = -1,
     timeout: float = np.inf,
-) -> None:
-    def _simulate_tree(seed: int) -> Tree:
+    acceptance_criterion: Callable[[Tree], bool] | None = None,
+) -> pd.DataFrame:
+    if isinstance(output_dir, str):
+        output_dir = Path(output_dir)
+    if output_dir.exists():
+        raise FileExistsError(f"Output directory {output_dir} already exists")
+    output_dir.mkdir(parents=True)
+
+    def _simulate_tree(i: int, seed: int) -> dict[str, Any]:
         while True:
             try:
-                return simulate_tree(
+                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:
+                    set_features(tree, node_features)
+                dump_newick(tree, output_dir / f"{i}.nwk")
+                return metadata
             except TimeoutError:
                 print("Simulation timed out, retrying with a different seed...")
             seed += 1
 
-    if os.path.exists(output_dir):
-        raise FileExistsError(f"Output directory {output_dir} already exists")
-    os.mkdir(output_dir)
-
     rng = default_rng(seed)
     jobs = joblib.Parallel(n_jobs=n_jobs, return_as="generator_unordered")(
-        joblib.delayed(_simulate_tree)(seed=int(rng.integers(2**32)))
-        for _ in range(n_trees)
+        joblib.delayed(_simulate_tree)(i=i, seed=int(rng.integers(2**32)))
+        for i in range(n_trees)
+    )
+
+    return pd.DataFrame(
+        [md for md in tqdm(jobs, f"Generating trees in {output_dir}...", n_trees)]
     )
-    for i, tree in tqdm(
-        enumerate(jobs), total=n_trees, desc=f"Generating trees in {output_dir}..."
-    ):
-        dump_newick(tree, os.path.join(output_dir, f"{i}.nwk"))

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/treesimulator/io/newick.py

@@ -0,0 +1,104 @@
+import re
+from pathlib import Path
+
+from phylogenie.treesimulator.tree import Tree
+
+
+def parse_newick(newick: str, translations: dict[str, str] | None = None) -> Tree:
+    newick = newick.strip()
+    newick = re.sub(r"^\[\&[^\]]*\]", "", newick).strip()
+
+    stack: list[list[Tree]] = []
+    current_children: list[Tree] = []
+    current_nodes: list[Tree] = []
+    i = 0
+    while True:
+
+        def _read_chars(stoppers: list[str]) -> str:
+            nonlocal i
+            chars = ""
+            while i < len(newick) and newick[i] not in stoppers:
+                chars += newick[i]
+                i += 1
+            if i == len(newick):
+                raise ValueError(f"Expected one of {stoppers}, got end of string")
+            return chars
+
+        if newick[i] == "(":
+            stack.append(current_nodes)
+            current_nodes = []
+            i += 1
+            continue
+
+        name = _read_chars([":", "[", ",", ")", ";"])
+        if translations is not None and name in translations:
+            name = translations[name]
+        current_node = Tree(name)
+
+        if newick[i] == "[":
+            i += 1
+            if newick[i] != "&":
+                raise ValueError("Expected '[&' at the start of node features")
+            i += 1
+            features = re.split(r",(?=[^,]+=)", _read_chars(["]"]))
+            i += 1
+            for feature in features:
+                key, value = feature.split("=")
+                try:
+                    current_node.set(key, eval(value))
+                except Exception:
+                    current_node.set(key, value)
+
+        if newick[i] == ":":
+            i += 1
+            current_node.branch_length = float(_read_chars([",", ")", ";"]))
+
+        for node in current_children:
+            current_node.add_child(node)
+            current_children = []
+        current_nodes.append(current_node)
+
+        if newick[i] == ")":
+            current_children = current_nodes
+            current_nodes = stack.pop()
+        elif newick[i] == ";":
+            return current_node
+
+        i += 1
+
+
+def load_newick(filepath: str | Path) -> Tree | list[Tree]:
+    with open(filepath, "r") as file:
+        trees = [parse_newick(newick) for newick in file]
+    return trees[0] if len(trees) == 1 else trees
+
+
+def to_newick(tree: Tree) -> str:
+    children_newick = ",".join([to_newick(child) for child in tree.children])
+    newick = tree.name
+    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(
+                    f"Invalid feature key `{k}`: keys must not contain ',', '=', or ']'"
+                )
+            if "=" in r or "]" in r:
+                raise ValueError(
+                    f"Invalid value  `{r}` for feature `{k}`: values must not contain '=' or ']'"
+                )
+        features = [f"{k}={repr}" for k, repr in reprs.items()]
+        newick += f"[&{','.join(features)}]"
+    if children_newick:
+        newick = f"({children_newick}){newick}"
+    if tree.branch_length is not None:
+        newick += f":{tree.branch_length}"
+    return newick
+
+
+def dump_newick(trees: Tree | list[Tree], filepath: str | Path) -> None:
+    if isinstance(trees, Tree):
+        trees = [trees]
+    with open(filepath, "w") as file:
+        for t in trees:
+            file.write(to_newick(t) + ";\n")

phylogenie/treesimulator/io/nexus.py

@@ -0,0 +1,50 @@
+import re
+from collections.abc import Iterator
+from pathlib import Path
+
+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:
+        line = line.strip()
+        match = re.match(r"(\d+)\s+['\"]?([^'\",;]+)['\"]?", line)
+        if match is None:
+            if ";" in line:
+                return translations
+            else:
+                raise ValueError("Invalid translate line. Expected '<num> <name>'.")
+        translations[match.group(1)] = match.group(2)
+    raise ValueError("Translate block not terminated with ';'.")
+
+
+def _parse_trees_block(lines: Iterator[str]) -> dict[str, Tree]:
+    trees: dict[str, Tree] = {}
+    translations = {}
+    for line in lines:
+        line = line.strip()
+        if line.upper() == "TRANSLATE":
+            translations = _parse_translate_block(lines)
+        elif line.upper() == "END;":
+            return trees
+        else:
+            match = re.match(r"^TREE\s*\*?\s+(\S+)\s*=\s*(.+)$", line, re.IGNORECASE)
+            if match is None:
+                raise ValueError(
+                    "Invalid tree line. Expected 'TREE <name> = <newick>'."
+                )
+            name = match.group(1)
+            if name in trees:
+                raise ValueError(f"Duplicate tree name found: {name}.")
+            trees[name] = parse_newick(match.group(2), translations)
+    raise ValueError("Unterminated TREES block.")
+
+
+def load_nexus(nexus_file: str | Path) -> dict[str, Tree]:
+    with open(nexus_file, "r") as f:
+        for line in f:
+            if line.strip().upper() == "BEGIN TREES;":
+                return _parse_trees_block(f)
+    raise ValueError("No TREES block found in the NEXUS file.")

phylogenie/treesimulator/model.py

@@ -1,33 +1,8 @@
-from abc import ABC, abstractmethod
 from collections import defaultdict
-from collections.abc import Sequence
 from dataclasses import dataclass
-from typing import Any
 
-import numpy as np
-from numpy.random import Generator
-
-from phylogenie.skyline import SkylineParameterLike, skyline_parameter
-from phylogenie.tree import Tree
-
-
-class Event(ABC):
-    def __init__(self, state: str, rate: SkylineParameterLike):
-        self.state = state
-        self.rate = skyline_parameter(rate)
-
-    def draw_individual(self, model: "Model", rng: Generator) -> int:
-        return rng.choice(model.get_population(self.state))
-
-    def get_propensity(self, model: "Model", time: float) -> float:
-        n_individuals = model.count_individuals(self.state)
-        rate = self.rate.get_value_at_time(time)
-        if rate == np.inf and not n_individuals:
-            return 0
-        return rate * n_individuals
-
-    @abstractmethod
-    def apply(self, model: "Model", time: float, rng: Generator) -> None: ...
+from phylogenie.mixins import MetadataMixin
+from phylogenie.treesimulator.tree import Tree
 
 
 @dataclass
@@ -37,31 +12,36 @@ class Individual:
     state: str
 
 
-class Model:
-    def __init__(self, init_state: str, events: Sequence[Event]):
+def _get_node_name(node_id: int, state: str) -> str:
+    return f"{node_id}|{state}"
+
+
+def get_node_state(node_name: str) -> str:
+    if "|" not in node_name:
+        raise ValueError(
+            f"Invalid node name: {node_name} (expected format 'id|state')."
+        )
+    return node_name.split("|")[-1]
+
+
+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._events = list(events)
-        self.context: dict[str, Any] = {}
 
     @property
     def n_sampled(self) -> int:
         return len(self._sampled)
 
-    @property
-    def events(self) -> tuple[Event, ...]:
-        return tuple(self._events)
-
-    def add_event(self, event: Event) -> None:
-        self._events.append(event)
-
     def _get_new_node(self, state: str) -> Tree:
         self._next_node_id += 1
-        return Tree(f"{self._next_node_id}|{state}")
+        node = Tree(_get_node_name(self._next_node_id, state))
+        return node
 
     def _get_new_individual(self, state: str) -> Individual:
         self._next_individual_id += 1
@@ -74,10 +54,8 @@ class Model:
 
     def _set_branch_length(self, node: Tree, time: float) -> None:
         if node.branch_length is not None:
-            raise ValueError(f"Branch length of node {node.id} is already set.")
-        node.branch_length = (
-            time if node.parent is None else time - node.parent.get_time()
-        )
+            raise ValueError(f"Branch length of node {node.name} is already set.")
+        node.branch_length = time if node.parent is None else time - node.parent.depth
 
     def _stem(self, individual: Individual, time: float) -> None:
         self._set_branch_length(individual.node, time)
@@ -108,12 +86,12 @@ class Model:
     def sample(self, id: int, time: float, removal: bool) -> None:
         individual = self._population[id]
         if removal:
-            self._sampled.add(individual.node.id)
+            self._sampled.add(individual.node.name)
             self.remove(id, time)
         else:
             sample_node = self._get_new_node(individual.state)
             sample_node.branch_length = 0.0
-            self._sampled.add(sample_node.id)
+            self._sampled.add(sample_node.name)
             individual.node.add_child(sample_node)
             self._stem(individual, time)
 
@@ -123,7 +101,7 @@ class Model:
     def get_sampled_tree(self) -> Tree:
         tree = self._tree.copy()
         for node in list(tree.postorder_traversal()):
-            if node.id not in self._sampled and not node.children:
+            if node.name not in self._sampled and not node.children:
                 if node.parent is None:
                     raise ValueError("No samples in the tree.")
                 else:
@@ -131,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/treesimulator/tree.py

@@ -0,0 +1,196 @@
+from collections import deque
+from collections.abc import Callable, Iterator
+from typing import Any
+
+from phylogenie.mixins import MetadataMixin
+
+
+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] = []
+
+    # ----------------
+    # Basic properties
+    # ----------------
+    # Properties related to parent-child relationships.
+
+    @property
+    def children(self) -> tuple["Tree", ...]:
+        return tuple(self._children)
+
+    @property
+    def parent(self) -> "Tree | None":
+        return self._parent
+
+    def add_child(self, child: "Tree") -> "Tree":
+        if child.parent is not None:
+            raise ValueError(f"Node {child.name} already has a parent.")
+        child._parent = self
+        self._children.append(child)
+        return self
+
+    def remove_child(self, child: "Tree") -> None:
+        self._children.remove(child)
+        child._parent = None
+
+    def set_parent(self, parent: "Tree | None"):
+        if self.parent is not None:
+            self.parent.remove_child(self)
+        self._parent = parent
+        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
+            return
+        if len(self.children) != 2:
+            raise ValueError("Inorder traversal is only defined for binary trees.")
+        left, right = self.children
+        yield from left.inorder_traversal()
+        yield self
+        yield from right.inorder_traversal()
+
+    def postorder_traversal(self) -> Iterator["Tree"]:
+        for child in self.children:
+            yield from child.postorder_traversal()
+        yield self
+
+    def breadth_first_traversal(self) -> Iterator["Tree"]:
+        queue: deque["Tree"] = deque([self])
+        while queue:
+            node = queue.popleft()
+            yield node
+            queue.extend(node.children)
+
+    # ---------------
+    # Tree properties
+    # ---------------
+    # Properties and methods related to tree metrics like leaf count, depth, height, etc.
+
+    @property
+    def n_leaves(self) -> int:
+        return len(self.get_leaves())
+
+    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
+
+    @property
+    def depth_level(self) -> int:
+        return 0 if self.parent is None else self.parent.depth_level + 1
+
+    @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()
+
+    @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 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:
+        def _default_key(node: Tree) -> int:
+            return node.n_leaves
+
+        if key is None:
+            key = _default_key
+        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 {name} not found.")
+
+    def copy(self):
+        new_tree = Tree(self.name, self.branch_length)
+        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()
+
+    def __len__(self) -> int:
+        return sum(1 for _ in self)
+
+    def __repr__(self) -> str:
+        return f"TreeNode(name='{self.name}', branch_length={self.branch_length}, metadata={self.metadata})"