phylogenie 2.1.6__tar.gz → 2.1.8__tar.gz

{phylogenie-2.1.6 → phylogenie-2.1.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: phylogenie
-Version: 2.1.6
+Version: 2.1.8
 Summary: Generate phylogenetic datasets with minimal setup effort
 Author: Gabriele Marino
 Author-email: gabmarino.8601@gmail.com
@@ -10,6 +10,7 @@ Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Requires-Dist: joblib (>=1.4.2,<2.0.0)
+Requires-Dist: matplotlib (>=3.10.6,<4.0.0)
 Requires-Dist: pandas (>=2.2.2,<3.0.0)
 Requires-Dist: pydantic (>=2.11.5,<3.0.0)
 Requires-Dist: pyyaml (>=6.0.2,<7.0.0)
@@ -24,6 +25,8 @@ Description-Content-Type: text/markdown
 
 [![AliSim](https://img.shields.io/badge/Powered%20by-AliSim-orange?style=flat-square)](https://iqtree.github.io/doc/AliSim)
 [![PyPI version](https://img.shields.io/pypi/v/phylogenie)](https://pypi.org/project/phylogenie/)
+![Downloads](https://img.shields.io/pypi/dm/phylogenie)
+
 
 Phylogenie is a [Python](https://www.python.org/) package designed to easily simulate phylogenetic datasets—such as trees and multiple sequence alignments (MSAs)—with minimal setup effort. Simply specify the distributions from which your parameters should be sampled, and Phylogenie will handle the rest!
 

{phylogenie-2.1.6 → phylogenie-2.1.8}/README.md

@@ -6,6 +6,8 @@
 
 [![AliSim](https://img.shields.io/badge/Powered%20by-AliSim-orange?style=flat-square)](https://iqtree.github.io/doc/AliSim)
 [![PyPI version](https://img.shields.io/pypi/v/phylogenie)](https://pypi.org/project/phylogenie/)
+![Downloads](https://img.shields.io/pypi/dm/phylogenie)
+
 
 Phylogenie is a [Python](https://www.python.org/) package designed to easily simulate phylogenetic datasets—such as trees and multiple sequence alignments (MSAs)—with minimal setup effort. Simply specify the distributions from which your parameters should be sampled, and Phylogenie will handle the rest!
 

{phylogenie-2.1.6 → phylogenie-2.1.8}/phylogenie/generators/alisim.py

@@ -72,7 +72,7 @@ class AliSimDatasetGenerator(DatasetGenerator):
                 )
 
         for leaf in tree.get_leaves():
-            leaf.id += f"|{leaf.get_time()}"
+            leaf.name += f"|{leaf.get_time()}"
         dump_newick(tree, f"{tree_filename}.nwk")
 
         self._generate_one_from_tree(msa_filename, f"{tree_filename}.nwk", rng, d)

{phylogenie-2.1.6 → phylogenie-2.1.8}/phylogenie/generators/configs.py

@@ -1,6 +1,6 @@
 import phylogenie.typings as pgt
+from phylogenie.models import Distribution, StrictBaseModel
 from phylogenie.treesimulator import MutationTargetType
-from phylogenie.utils import Distribution, StrictBaseModel
 
 Integer = str | int
 Scalar = str | pgt.Scalar

{phylogenie-2.1.6 → phylogenie-2.1.8}/phylogenie/generators/dataset.py

@@ -8,7 +8,7 @@ import pandas as pd
 from numpy.random import Generator, default_rng
 from tqdm import tqdm
 
-from phylogenie.utils import Distribution, StrictBaseModel
+from phylogenie.models import Distribution, StrictBaseModel
 
 
 class DataType(str, Enum):
@@ -31,7 +31,7 @@ class DatasetGenerator(ABC, StrictBaseModel):
     def generate_one(
         self,
         filename: str,
-        context: dict[str, Any] | None = None,
+        context: dict[str, Distribution] | None = None,
         seed: int | None = None,
     ) -> dict[str, Any]: ...
 

{phylogenie-2.1.6 → phylogenie-2.1.8}/phylogenie/generators/factories.py

@@ -8,6 +8,7 @@ import phylogenie.generators.configs as cfg
 import phylogenie.generators.typeguards as ctg
 import phylogenie.typeguards as tg
 import phylogenie.typings as pgt
+from phylogenie.models import Distribution
 from phylogenie.skyline import (
     SkylineMatrix,
     SkylineMatrixCoercible,
@@ -16,7 +17,6 @@ from phylogenie.skyline import (
     SkylineVector,
     SkylineVectorCoercible,
 )
-from phylogenie.utils import Distribution
 
 
 def _eval_expression(expression: str, data: dict[str, Any]) -> Any:

{phylogenie-2.1.6 → phylogenie-2.1.8}/phylogenie/generators/trees.py

@@ -19,9 +19,11 @@ from phylogenie.generators.factories import (
     skyline_vector,
 )
 from phylogenie.io import dump_newick
+from phylogenie.models import Distribution
 from phylogenie.tree import Tree
 from phylogenie.treesimulator import (
     Event,
+    Feature,
     Mutation,
     get_BD_events,
     get_BDEI_events,
@@ -30,6 +32,7 @@ from phylogenie.treesimulator import (
     get_contact_tracing_events,
     get_epidemiological_events,
     get_FBD_events,
+    set_features,
     simulate_tree,
 )
 
@@ -51,6 +54,7 @@ class TreeDatasetGenerator(DatasetGenerator):
     init_state: str | None = None
     sampling_probability_at_present: cfg.Scalar = 0.0
     timeout: float = np.inf
+    node_features: list[Feature] | None = None
 
     @abstractmethod
     def _get_events(self, data: dict[str, Any]) -> list[Event]: ...
@@ -77,7 +81,7 @@ class TreeDatasetGenerator(DatasetGenerator):
     def generate_one(
         self,
         filename: str,
-        context: dict[str, Any] | None = None,
+        context: dict[str, Distribution] | None = None,
         seed: int | None = None,
     ) -> dict[str, Any]:
         d = {"file_id": Path(filename).stem}
@@ -86,6 +90,8 @@ class TreeDatasetGenerator(DatasetGenerator):
             try:
                 d.update(data(context, rng))
                 tree = self.simulate_one(d, seed)
+                if self.node_features is not None:
+                    set_features(tree, self.node_features)
                 dump_newick(tree, f"{filename}.nwk")
                 break
             except TimeoutError:

{phylogenie-2.1.6 → phylogenie-2.1.8}/phylogenie/io.py

@@ -44,7 +44,12 @@ def _parse_newick(newick: str) -> Tree:
                     raise ValueError(f"Expected '&&NHX' for node features.")
                 for feature in features[1:]:
                     key, value = feature.split("=", 1)
-                    current_node.set(key, eval(value))
+                    try:
+                        current_node.set(key, eval(value))
+                    except Exception as e:
+                        raise ValueError(
+                            f"Error setting node feature `{key}` to `{value}`: {e}"
+                        )
 
             if newick[i] == ")":
                 current_children = current_nodes
@@ -65,13 +70,18 @@ def load_newick(filepath: str) -> Tree | list[Tree]:
 
 def _to_newick(tree: Tree) -> str:
     children_newick = ",".join([_to_newick(child) for child in tree.children])
-    newick = tree.id
+    newick = tree.name
     if children_newick:
         newick = f"({children_newick}){newick}"
     if tree.branch_length is not None:
         newick += f":{tree.branch_length}"
     if tree.features:
         reprs = {k: repr(v).replace("'", '"') for k, v in tree.features.items()}
+        for k, r in reprs.items():
+            if ":" in r or "=" in r or "]" in r:
+                raise ValueError(
+                    f"Cannot serialize feature `{k}` with value `{r}`: contains reserved characters."
+                )
         features = [f"{k}={repr}" for k, repr in reprs.items()]
         newick += f"[&&NHX:{':'.join(features)}]"
     return newick

phylogenie-2.1.8/phylogenie/plot.py

@@ -0,0 +1,39 @@
+import matplotlib.colors as mcolors
+import matplotlib.pyplot as plt
+
+from phylogenie import Tree
+from phylogenie.tree import Tree
+from phylogenie.utils import get_times
+
+
+def plot_tree(
+    tree: Tree,
+    ax: plt.Axes | None = None,  # pyright: ignore
+    color_by: str | None = None,
+    default_color: str = "black",
+    cmap: str = "tab20",
+) -> plt.Axes:  # pyright: ignore
+    if ax is None:
+        ax = plt.gca()
+
+    xs = get_times(tree)
+    ys = {node.name: i for i, node in enumerate(tree.inorder_traversal())}
+    if color_by is not None:
+        features = set(node.get(color_by) for node in tree)
+        feature_colors = {
+            f: mcolors.to_hex(plt.get_cmap(cmap, len(features))(i))
+            for i, f in enumerate(features)
+        }
+        colors = {node.name: feature_colors[node.get(color_by)] for node in tree}
+    else:
+        colors = {node.name: default_color for node in tree}
+
+    for node in tree:
+        if node.parent is None:
+            continue
+        x0, y0 = xs[node.parent.name], ys[node.parent.name]
+        x1, y1 = xs[node.name], ys[node.name]
+        ax.plot([x0, x0], [y0, y1], color=colors[node.name])  # pyright: ignore
+        ax.plot([x0, x1], [y1, y1], color=colors[node.name])  # pyright: ignore
+    ax.set_yticks([])  # pyright: ignore
+    return ax

{phylogenie-2.1.6 → phylogenie-2.1.8}/phylogenie/tree.py

@@ -3,8 +3,8 @@ from typing import Any
 
 
 class Tree:
-    def __init__(self, id: str = "", branch_length: float | None = None):
-        self.id = id
+    def __init__(self, name: str = "", branch_length: float | None = None):
+        self.name = name
         self.branch_length = branch_length
         self._parent: Tree | None = None
         self._children: list[Tree] = []
@@ -20,7 +20,7 @@ class Tree:
 
     @property
     def features(self) -> dict[str, Any]:
-        return self._features
+        return self._features.copy()
 
     def add_child(self, child: "Tree") -> "Tree":
         child._parent = self
@@ -36,6 +36,17 @@ class Tree:
         if node is not None:
             node._children.append(self)
 
+    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 preorder_traversal(self) -> Iterator["Tree"]:
         yield self
         for child in self.children:
@@ -46,41 +57,42 @@ class Tree:
             yield from child.postorder_traversal()
         yield self
 
-    def get_node(self, id: str) -> "Tree":
+    def get_node(self, name: str) -> "Tree":
         for node in self:
-            if node.id == id:
+            if node.name == name:
                 return node
-        raise ValueError(f"Node with id {id} not found.")
+        raise ValueError(f"Node with name {name} not found.")
 
     def is_leaf(self) -> bool:
         return not self.children
 
-    def get_leaves(self) -> list["Tree"]:
-        return [node for node in self if not node.children]
+    def get_leaves(self) -> tuple["Tree", ...]:
+        return tuple(node for node in self if not node.children)
+
+    def parse_branch_length(self) -> float:
+        if self.branch_length is None:
+            raise ValueError(f"Branch length of node {self.name} is not set.")
+        return self.branch_length
 
     def get_time(self) -> float:
         parent_time = 0 if self.parent is None else self.parent.get_time()
-        if self.branch_length is None:
-            if self.parent is not None:
-                raise ValueError(
-                    f"Branch length of non-root node {self.id} is not set."
-                )
-            return 0.0
-        return self.branch_length + parent_time
+        return self.parse_branch_length() + parent_time
 
     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.get(key)
+        return self._features[key]
 
     def delete(self, key: str) -> None:
         del self._features[key]
 
     def copy(self):
-        new_tree = Tree(self.id, self.branch_length)
-        for key, value in self._features.items():
-            new_tree.set(key, value)
+        new_tree = Tree(self.name, self.branch_length)
+        new_tree.update_features(self._features)
         for child in self.children:
             new_tree.add_child(child.copy())
         return new_tree
@@ -89,4 +101,4 @@ class Tree:
         return self.preorder_traversal()
 
     def __repr__(self) -> str:
-        return f"TreeNode(id='{self.id}', branch_length={self.branch_length}, features={self.features})"
+        return f"TreeNode(name='{self.name}', branch_length={self.branch_length}, features={self.features})"

{phylogenie-2.1.6 → phylogenie-2.1.8}/phylogenie/treesimulator/__init__.py

@@ -15,8 +15,11 @@ from phylogenie.treesimulator.events import (
     get_contact_tracing_events,
     get_epidemiological_events,
     get_FBD_events,
+    get_mutation_id,
 )
+from phylogenie.treesimulator.features import Feature, set_features
 from phylogenie.treesimulator.gillespie import generate_trees, simulate_tree
+from phylogenie.treesimulator.model import get_node_state
 
 __all__ = [
     "Birth",
@@ -37,4 +40,8 @@ __all__ = [
     "get_FBD_events",
     "generate_trees",
     "simulate_tree",
+    "get_mutation_id",
+    "get_node_state",
+    "Feature",
+    "set_features",
 ]

{phylogenie-2.1.6 → phylogenie-2.1.8}/phylogenie/treesimulator/events/__init__.py

@@ -18,6 +18,7 @@ from phylogenie.treesimulator.events.core import (
 )
 from phylogenie.treesimulator.events.mutations import Mutation
 from phylogenie.treesimulator.events.mutations import TargetType as MutationTargetType
+from phylogenie.treesimulator.events.mutations import get_mutation_id
 
 __all__ = [
     "Birth",
@@ -36,4 +37,5 @@ __all__ = [
     "get_contact_tracing_events",
     "get_epidemiological_events",
     "get_FBD_events",
+    "get_mutation_id",
 ]

{phylogenie-2.1.6 → phylogenie-2.1.8}/phylogenie/treesimulator/events/mutations.py

@@ -1,9 +1,11 @@
+import re
 from copy import deepcopy
 from enum import Enum
 from typing import Type
 
 from numpy.random import Generator
 
+from phylogenie.models import Distribution
 from phylogenie.skyline import SkylineParameterLike
 from phylogenie.treesimulator.events.contact_tracing import (
     BirthWithContactTracing,
@@ -17,7 +19,6 @@ from phylogenie.treesimulator.events.core import (
     Sampling,
 )
 from phylogenie.treesimulator.model import Model
-from phylogenie.utils import Distribution
 
 MUTATION_PREFIX = "MUT-"
 MUTATIONS_KEY = "MUTATIONS"
@@ -29,10 +30,17 @@ def _get_mutation(state: str) -> str | None:
 
 def _get_mutated_state(mutation_id: int, state: str) -> str:
     if state.startswith(MUTATION_PREFIX):
-        state = state.split(".")[1]
+        _, state = state.split(".")
     return f"{MUTATION_PREFIX}{mutation_id}.{state}"
 
 
+def get_mutation_id(node_name: str) -> int:
+    match = re.search(rf"{MUTATION_PREFIX}(\d+)\.", node_name)
+    if match:
+        return int(match.group(1))
+    return 0
+
+
 class TargetType(str, Enum):
     BIRTH = "birth"
     DEATH = "death"
@@ -93,7 +101,7 @@ class Mutation(Event):
             model.add_event(event)
 
     def __repr__(self) -> str:
-        return f"Mutation(state={self.state}, rate={self.rate}, rate_scalers={self.rate_scalers})"
+        return f"Mutation(state={self.state}, rate={self.rate})"
 
 
 TARGETS: dict[TargetType, tuple[Type[Event], ...]] = {

phylogenie-2.1.8/phylogenie/treesimulator/features.py

@@ -0,0 +1,39 @@
+from collections.abc import Iterable
+from enum import Enum
+
+from phylogenie.tree import Tree
+from phylogenie.treesimulator.events import get_mutation_id
+from phylogenie.treesimulator.model import get_node_state
+from phylogenie.utils import get_heights, get_n_tips, get_times
+
+
+def _get_states(tree: Tree) -> dict[str, str]:
+    return {node.name: get_node_state(node.name) for node in tree}
+
+
+def _get_mutations(tree: Tree) -> dict[str, int]:
+    return {node.name: get_mutation_id(node.name) for node in tree}
+
+
+class Feature(str, Enum):
+    STATE = "state"
+    MUTATION = "mutation"
+    N_TIPS = "n_tips"
+    TIME = "time"
+    HEIGHT = "height"
+
+
+FEATURES_EXTRACTORS = {
+    Feature.STATE: _get_states,
+    Feature.MUTATION: _get_mutations,
+    Feature.N_TIPS: get_n_tips,
+    Feature.TIME: get_times,
+    Feature.HEIGHT: get_heights,
+}
+
+
+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.name])

{phylogenie-2.1.6 → phylogenie-2.1.8}/phylogenie/treesimulator/gillespie.py

@@ -1,6 +1,6 @@
 import os
 import time
-from collections.abc import Sequence
+from collections.abc import Iterable, Sequence
 
 import joblib
 import numpy as np
@@ -9,6 +9,7 @@ from tqdm import tqdm
 
 from phylogenie.io import dump_newick
 from phylogenie.tree import Tree
+from phylogenie.treesimulator.features import Feature, set_features
 from phylogenie.treesimulator.model import Event, Model
 
 
@@ -69,7 +70,7 @@ def simulate_tree(
             if (
                 not any(rates)
                 or max_tips is not None
-                and model.n_sampled > max_tips
+                and model.n_sampled >= max_tips
                 or target_n_tips is not None
                 and model.n_sampled >= target_n_tips
             ):
@@ -101,12 +102,13 @@ def simulate_tree(
 def generate_trees(
     output_dir: str,
     n_trees: int,
-    events: list[Event],
+    events: Sequence[Event],
     min_tips: int = 1,
-    max_tips: int = 2**32,
+    max_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,
@@ -114,7 +116,7 @@ def generate_trees(
     def _simulate_tree(seed: int) -> Tree:
         while True:
             try:
-                return simulate_tree(
+                tree = simulate_tree(
                     events=events,
                     min_tips=min_tips,
                     max_tips=max_tips,
@@ -124,6 +126,9 @@ def generate_trees(
                     seed=seed,
                     timeout=timeout,
                 )
+                if node_features is not None:
+                    set_features(tree, node_features)
+                return tree
             except TimeoutError:
                 print("Simulation timed out, retrying with a different seed...")
             seed += 1

{phylogenie-2.1.6 → phylogenie-2.1.8}/phylogenie/treesimulator/model.py

@@ -11,6 +11,13 @@ from phylogenie.skyline import SkylineParameterLike, skyline_parameter
 from phylogenie.tree import Tree
 
 
+@dataclass
+class Individual:
+    id: int
+    node: Tree
+    state: str
+
+
 class Event(ABC):
     def __init__(self, state: str, rate: SkylineParameterLike):
         self.state = state
@@ -30,11 +37,17 @@ class Event(ABC):
     def apply(self, model: "Model", time: float, rng: Generator) -> None: ...
 
 
-@dataclass
-class Individual:
-    id: int
-    node: Tree
-    state: str
+def _get_node_name(node_id: int, state: str) -> str:
+    return f"{node_id}|{state}"
+
+
+def get_node_state(node_name: str) -> str:
+    try:
+        return node_name.split("|")[1]
+    except IndexError:
+        raise ValueError(
+            f"Invalid node name: {node_name} (expected format 'id|state')."
+        )
 
 
 class Model:
@@ -61,7 +74,8 @@ class Model:
 
     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,7 +88,7 @@ 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.")
+            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.get_time()
         )
@@ -108,12 +122,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 +137,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:

phylogenie-2.1.8/phylogenie/utils.py

@@ -0,0 +1,28 @@
+from phylogenie.tree import Tree
+
+
+def get_n_tips(tree: Tree) -> dict[str, int]:
+    n_tips: dict[str, int] = {}
+    for node in tree.postorder_traversal():
+        n_tips[node.name] = (
+            1 if node.is_leaf() else sum(n_tips[child.name] for child in node.children)
+        )
+    return n_tips
+
+
+def get_times(tree: Tree) -> dict[str, float]:
+    times: dict[str, float] = {}
+    for node in tree:
+        parent_time = 0 if node.parent is None else times[node.parent.name]
+        times[node.name] = node.parse_branch_length() + parent_time
+    return times
+
+
+def get_heights(tree: Tree) -> dict[str, int]:
+    heights: dict[str, int] = {}
+    for node in tree.postorder_traversal():
+        if node.is_leaf():
+            heights[node.name] = 0
+        else:
+            heights[node.name] = 1 + max(heights[child.name] for child in node.children)
+    return heights

{phylogenie-2.1.6 → phylogenie-2.1.8}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "phylogenie"
-version = "2.1.6"
+version = "2.1.8"
 description = "Generate phylogenetic datasets with minimal setup effort"
 authors = ["Gabriele Marino <gabmarino.8601@gmail.com>"]
 readme = "README.md"
@@ -9,6 +9,7 @@ readme = "README.md"
 python = "^3.10"
 
 joblib = "^1.4.2"
+matplotlib = "^3.10.6"
 pandas = "^2.2.2"
 pydantic = "^2.11.5"
 pyyaml = "^6.0.2"