phylogenie 3.1.11__tar.gz → 3.1.18__tar.gz

{phylogenie-3.1.11/src/phylogenie.egg-info → phylogenie-3.1.18}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: phylogenie
-Version: 3.1.11
+Version: 3.1.18
 Summary: Generate phylogenetic datasets with minimal setup effort
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown

{phylogenie-3.1.11 → phylogenie-3.1.18}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "phylogenie"
-version = "3.1.11"
+version = "3.1.18"
 description = "Generate phylogenetic datasets with minimal setup effort"
 readme = "README.md"
 requires-python = ">=3.10"
@@ -29,4 +29,4 @@ phylogenie = "phylogenie.main:main"
 
 [build-system]
 requires = ["setuptools>=42"]
-build-backend = "setuptools.build_meta"
+build-backend = "setuptools.build_meta"

{phylogenie-3.1.11 → phylogenie-3.1.18}/src/phylogenie/__init__.py

@@ -38,6 +38,7 @@ from phylogenie.treesimulator import (
     Tree,
     compute_mean_leaf_pairwise_distance,
     compute_sackin_index,
+    count_hops,
     dump_newick,
     generate_trees,
     get_BD_events,
@@ -50,12 +51,15 @@ from phylogenie.treesimulator import (
     get_FBD_events,
     get_mrca,
     get_mutation_id,
+    get_node_ages,
     get_node_depth_levels,
     get_node_depths,
     get_node_height_levels,
     get_node_heights,
     get_node_leaf_counts,
     get_node_state,
+    get_node_times,
+    get_path,
     load_newick,
     load_nexus,
     set_features,
@@ -100,6 +104,7 @@ __all__ = [
     "Tree",
     "compute_mean_leaf_pairwise_distance",
     "compute_sackin_index",
+    "count_hops",
     "dump_newick",
     "generate_trees",
     "get_BD_events",
@@ -112,12 +117,15 @@ __all__ = [
     "get_FBD_events",
     "get_mrca",
     "get_mutation_id",
+    "get_node_ages",
     "get_node_depth_levels",
     "get_node_depths",
     "get_node_height_levels",
     "get_node_heights",
     "get_node_leaf_counts",
     "get_node_state",
+    "get_node_times",
+    "get_path",
     "load_newick",
     "load_nexus",
     "set_features",

{phylogenie-3.1.11 → phylogenie-3.1.18}/src/phylogenie/draw.py

@@ -1,4 +1,5 @@
 import datetime
+from collections.abc import Mapping
 from dataclasses import dataclass
 from typing import Any, Literal, overload
 
@@ -30,7 +31,7 @@ Color = str | tuple[float, float, float] | tuple[float, float, float, float]
 def draw_tree(
     tree: Tree,
     ax: Axes | None = None,
-    colors: Color | dict[Tree, Color] = "black",
+    colors: Color | Mapping[Tree, Color] = "black",
     backward_time: bool = False,
     branch_kwargs: dict[str, Any] | None = None,
     sampled_ancestor_kwargs: dict[str, Any] | None = None,
@@ -44,8 +45,8 @@ def draw_tree(
         The phylogenetic tree to draw.
     ax : Axes | None, optional
         The matplotlib Axes to draw on. If None, uses the current Axes.
-    colors : Color | dict[Tree, Color], optional
-        A single color for all branches or a dictionary mapping each node to a color.
+    colors : Color | Mapping[Tree, Color], optional
+        A single color for all branches or a mapping from each node to a color.
     backward_time : bool, optional
         If True, the x-axis is inverted to represent time going backward.
     branch_kwargs : dict[str, Any] | None, optional
@@ -67,7 +68,7 @@ def draw_tree(
     if "marker" not in sampled_ancestor_kwargs:
         sampled_ancestor_kwargs["marker"] = "o"
 
-    if not isinstance(colors, dict):
+    if not isinstance(colors, Mapping):
         colors = {node: colors for node in tree}
 
     xs = (
@@ -141,7 +142,7 @@ def draw_dated_tree(
     tree: Tree,
     calibration_nodes: tuple[CalibrationNode, CalibrationNode],
     ax: Axes | None = None,
-    colors: Color | dict[Tree, Color] = "black",
+    colors: Color | Mapping[Tree, Color] = "black",
     branch_kwargs: dict[str, Any] | None = None,
 ) -> Axes:
     """
@@ -155,8 +156,8 @@ def draw_dated_tree(
         Two calibration nodes defining the mapping from depth to date.
     ax : Axes | None, optional
         The matplotlib Axes to draw on. If None, uses the current Axes.
-    colors : Color | dict[Tree, Color], optional
-        A single color for all branches or a dictionary mapping each node to a color.
+    colors : Color | Mapping[Tree, Color], optional
+        A single color for all branches or a mapping from each node to a color.
     branch_kwargs : dict[str, Any] | None, optional
         Additional keyword arguments to pass to the branch drawing functions.
 
@@ -170,7 +171,7 @@ def draw_dated_tree(
     if branch_kwargs is None:
         branch_kwargs = {}
 
-    if not isinstance(colors, dict):
+    if not isinstance(colors, Mapping):
         colors = {node: colors for node in tree}
 
     xs = {
@@ -224,9 +225,9 @@ def _init_colored_tree_categorical(
     color_by: str,
     ax: Axes | None = None,
     default_color: Color = "black",
-    colormap: str | Colormap = "tab20",
+    colormap: str | Mapping[str, Color] | Colormap = "tab20",
     show_legend: bool = True,
-    labels: dict[Any, str] | None = None,
+    labels: Mapping[Any, str] | None = None,
     legend_kwargs: dict[str, Any] | None = None,
 ) -> tuple[Axes, dict[Tree, Color]]:
     """
@@ -242,11 +243,14 @@ def _init_colored_tree_categorical(
         The matplotlib Axes to draw on. If None, uses the current Axes.
     default_color : Color, optional
         The color to use for nodes without the specified metadata.
-    colormap : str | Colormap, optional
-        The colormap to use for coloring categories. Defaults to 'tab20'.
+    colormap : str | Mapping[str, Color] | Colormap, optional
+        The colormap to use for coloring categories.
+        If a string, it is used to get a matplotlib colormap.
+        If a mapping, it maps category values to colors directly.
+        Defaults to 'tab20'.
     show_legend : bool, optional
         Whether to display a legend for the categories.
-    labels : dict[Any, str] | None, optional
+    labels : Mapping[Any, str] | None, optional
         A mapping from category values to labels for the legend.
     legend_kwargs : dict[str, Any] | None, optional
         Additional keyword arguments to pass to the legend.
@@ -259,13 +263,16 @@ def _init_colored_tree_categorical(
     if ax is None:
         ax = plt.gca()
 
+    features = {node: node[color_by] for node in tree if color_by in node.metadata}
     if isinstance(colormap, str):
         colormap = plt.get_cmap(colormap)
+    if isinstance(colormap, Colormap):
+        feature_colors = {
+            f: mcolors.to_hex(colormap(i)) for i, f in enumerate(set(features.values()))
+        }
+    else:
+        feature_colors = colormap
 
-    features = {node: node[color_by] for node in tree if color_by in node.metadata}
-    feature_colors = {
-        f: mcolors.to_hex(colormap(i)) for i, f in enumerate(set(features.values()))
-    }
     colors = {
         node: feature_colors[features[node]] if node in features else default_color
         for node in tree
@@ -294,9 +301,9 @@ def draw_colored_tree_categorical(
     ax: Axes | None = None,
     backward_time: bool = False,
     default_color: Color = "black",
-    colormap: str | Colormap = "tab20",
+    colormap: str | Mapping[str, Color] | Colormap = "tab20",
     show_legend: bool = True,
-    labels: dict[Any, str] | None = None,
+    labels: Mapping[Any, str] | None = None,
     legend_kwargs: dict[str, Any] | None = None,
     branch_kwargs: dict[str, Any] | None = None,
     sampled_ancestor_kwargs: dict[str, Any] | None = None,
@@ -316,11 +323,14 @@ def draw_colored_tree_categorical(
         If True, the x-axis is inverted to represent time going backward.
     default_color : Color, optional
         The color to use for nodes without the specified metadata.
-    colormap : str | Colormap, optional
-        The colormap to use for coloring categories. Defaults to 'tab20'.
+    colormap : str | Mapping[str, Color] | Colormap, optional
+        The colormap to use for coloring categories.
+        If a string, it is used to get a matplotlib colormap.
+        If a mapping, it maps category values to colors directly.
+        Defaults to 'tab20'.
     show_legend : bool, optional
         Whether to display a legend for the categories.
-    labels : dict[Any, str] | None, optional
+    labels : Mapping[Any, str] | None, optional
         A mapping from category values to labels for the legend.
     legend_kwargs : dict[str, Any] | None, optional
         Additional keyword arguments to pass to the legend.
@@ -360,9 +370,9 @@ def draw_colored_dated_tree_categorical(
     color_by: str,
     ax: Axes | None = None,
     default_color: Color = "black",
-    colormap: str | Colormap = "tab20",
+    colormap: str | Mapping[str, Color] | Colormap = "tab20",
     show_legend: bool = True,
-    labels: dict[Any, str] | None = None,
+    labels: Mapping[Any, str] | None = None,
     legend_kwargs: dict[str, Any] | None = None,
     branch_kwargs: dict[str, Any] | None = None,
 ) -> Axes:
@@ -381,11 +391,14 @@ def draw_colored_dated_tree_categorical(
         The matplotlib Axes to draw on. If None, uses the current Axes.
     default_color : Color, optional
         The color to use for nodes without the specified metadata.
-    colormap : str | Colormap, optional
-        The colormap to use for coloring categories. Defaults to 'tab20'.
+    colormap : str | Mapping[str, Color] | Colormap, optional
+        The colormap to use for coloring categories.
+        If a string, it is used to get a matplotlib colormap.
+        If a mapping, it maps category values to colors directly.
+        Defaults to 'tab20'.
     show_legend : bool, optional
         Whether to display a legend for the categories.
-    labels : dict[Any, str] | None, optional
+    labels : Mapping[Any, str] | None, optional
         A mapping from category values to labels for the legend.
     legend_kwargs : dict[str, Any] | None, optional
         Additional keyword arguments to pass to the legend.

{phylogenie-3.1.11 → phylogenie-3.1.18}/src/phylogenie/generators/factories.py

@@ -19,18 +19,22 @@ from phylogenie.skyline import (
 from phylogenie.treesimulator import EventType, Mutation
 
 
-def _eval_expression(expression: str, data: dict[str, Any]) -> Any:
+def eval_expression(
+    expression: str, data: dict[str, Any], extra_globals: dict[str, Any] | None = None
+) -> Any:
+    if extra_globals is None:
+        extra_globals = {}
     return np.array(
         eval(
             expression,
-            {"np": np, **{k: np.array(v) for k, v in data.items()}},
+            {"np": np, **{k: np.array(v) for k, v in data.items()}, **extra_globals},
         )
     ).tolist()
 
 
 def integer(x: cfg.Integer, data: dict[str, Any]) -> int:
     if isinstance(x, str):
-        e = _eval_expression(x, data)
+        e = eval_expression(x, data)
         if isinstance(e, int):
             return e
         raise ValueError(
@@ -41,7 +45,7 @@ def integer(x: cfg.Integer, data: dict[str, Any]) -> int:
 
 def scalar(x: cfg.Scalar, data: dict[str, Any]) -> pgt.Scalar:
     if isinstance(x, str):
-        e = _eval_expression(x, data)
+        e = eval_expression(x, data)
         if isinstance(e, pgt.Scalar):
             return e
         raise ValueError(
@@ -54,13 +58,13 @@ def string(s: Any, data: dict[str, Any]) -> str:
     if not isinstance(s, str):
         return str(s)
     return re.sub(
-        r"\{([^{}]+)\}", lambda match: str(_eval_expression(match.group(1), data)), s
+        r"\{([^{}]+)\}", lambda match: str(eval_expression(match.group(1), data)), s
     )  # Match content inside curly braces
 
 
 def many_scalars(x: cfg.ManyScalars, data: dict[str, Any]) -> pgt.ManyScalars:
     if isinstance(x, str):
-        e = _eval_expression(x, data)
+        e = eval_expression(x, data)
         if tg.is_many_scalars(e):
             return e
         raise ValueError(
@@ -73,7 +77,7 @@ def one_or_many_scalars(
     x: cfg.OneOrManyScalars, data: dict[str, Any]
 ) -> pgt.OneOrManyScalars:
     if isinstance(x, str):
-        e = _eval_expression(x, data)
+        e = eval_expression(x, data)
         if tg.is_one_or_many_scalars(e):
             return e
         raise ValueError(
@@ -99,7 +103,7 @@ def skyline_vector(
     x: cfg.SkylineVector, data: dict[str, Any]
 ) -> SkylineVectorCoercible:
     if isinstance(x, str):
-        e = _eval_expression(x, data)
+        e = eval_expression(x, data)
         if tg.is_one_or_many_scalars(e):
             return e
         raise ValueError(
@@ -114,7 +118,7 @@ def skyline_vector(
 
     change_times = many_scalars(x.change_times, data)
     if isinstance(x.value, str):
-        e = _eval_expression(x.value, data)
+        e = eval_expression(x.value, data)
         if tg.is_many_one_or_many_scalars(e):
             value = e
         else:
@@ -142,7 +146,7 @@ def one_or_many_2D_scalars(
     x: cfg.OneOrMany2DScalars, data: dict[str, Any]
 ) -> pgt.OneOrMany2DScalars:
     if isinstance(x, str):
-        e = _eval_expression(x, data)
+        e = eval_expression(x, data)
         if tg.is_one_or_many_2D_scalars(e):
             return e
         raise ValueError(
@@ -160,7 +164,7 @@ def skyline_matrix(
         return None
 
     if isinstance(x, str):
-        e = _eval_expression(x, data)
+        e = eval_expression(x, data)
         if tg.is_one_or_many_2D_scalars(e):
             return e
         raise ValueError(
@@ -175,7 +179,7 @@ def skyline_matrix(
 
     change_times = many_scalars(x.change_times, data)
     if isinstance(x.value, str):
-        e = _eval_expression(x.value, data)
+        e = eval_expression(x.value, data)
         if tg.is_many_one_or_many_2D_scalars(e):
             value = e
         else:
@@ -213,7 +217,7 @@ def distribution(x: cfg.Distribution, data: dict[str, Any]) -> cfg.Distribution:
     args = x.args
     for arg_name, arg_value in args.items():
         if isinstance(arg_value, str):
-            args[arg_name] = _eval_expression(arg_value, data)
+            args[arg_name] = eval_expression(arg_value, data)
     return cfg.Distribution(type=x.type, **args)
 
 

{phylogenie-3.1.11 → phylogenie-3.1.18}/src/phylogenie/generators/trees.py

@@ -12,6 +12,7 @@ from phylogenie.generators.configs import Distribution
 from phylogenie.generators.dataset import DatasetGenerator, DataType
 from phylogenie.generators.factories import (
     data,
+    eval_expression,
     integer,
     mutations,
     scalar,
@@ -57,6 +58,7 @@ class TreeDatasetGenerator(DatasetGenerator):
     timeout: float = np.inf
     node_features: list[Feature] | None = None
     acceptance_criterion: str | None = None
+    logs: dict[str, str] | None = None
 
     @abstractmethod
     def _get_events(self, data: dict[str, Any]) -> list[Event]: ...
@@ -77,10 +79,21 @@ class TreeDatasetGenerator(DatasetGenerator):
         acceptance_criterion: None | Callable[[Tree], bool] = (
             None
             if self.acceptance_criterion is None
-            else lambda tree: eval(
-                self.acceptance_criterion, {}, {"tree": tree}  # pyright: ignore
+            else lambda tree: eval_expression(
+                self.acceptance_criterion,  # pyright: ignore
+                data,
+                {"tree": tree},
             )
         )
+        logs: None | dict[str, Callable[[Tree], Any]] = (
+            None
+            if self.logs is None
+            else {
+                key: lambda tree, expr=expr: eval_expression(expr, data, {"tree": tree})
+                for key, expr in self.logs.items()
+            }
+        )
+
         return simulate_tree(
             events=events,
             n_tips=None if self.n_tips is None else integer(self.n_tips, data),
@@ -92,6 +105,7 @@ class TreeDatasetGenerator(DatasetGenerator):
             seed=seed,
             timeout=self.timeout,
             acceptance_criterion=acceptance_criterion,
+            logs=logs,
         )
 
     def generate_one(
@@ -111,7 +125,7 @@ class TreeDatasetGenerator(DatasetGenerator):
                 dump_newick(tree, f"{filename}.nwk")
                 break
             except TimeoutError:
-                print("Simulation timed out, retrying with different parameters...")
+                print("Simulation timed out. Retrying with different parameters...")
         return d | metadata
 
 

{phylogenie-3.1.11 → phylogenie-3.1.18}/src/phylogenie/msa.py

@@ -1,4 +1,4 @@
-from collections.abc import Iterator
+from collections.abc import Iterable, Iterator
 from dataclasses import dataclass
 from datetime import date
 
@@ -11,16 +11,23 @@ class Sequence:
     chars: str
     time: float | date | None = None
 
+    def __len__(self) -> int:
+        return len(self.chars)
+
 
 class MSA:
-    def __init__(self, sequences: list[Sequence]):
-        self.sequences = sequences
-        lengths = {len(sequence.chars) for sequence in sequences}
+    def __init__(self, sequences: Iterable[Sequence]):
+        self._sequences = sequences
+        lengths = {len(sequence) for sequence in sequences}
         if len(lengths) > 1:
             raise ValueError(
                 f"All sequences in the alignment must have the same length (got lengths: {lengths})"
             )
 
+    @property
+    def sequences(self) -> tuple[Sequence, ...]:
+        return tuple(self._sequences)
+
     @property
     def ids(self) -> list[str]:
         return [sequence.id for sequence in self.sequences]

{phylogenie-3.1.11 → phylogenie-3.1.18}/src/phylogenie/treesimulator/__init__.py

@@ -25,6 +25,7 @@ from phylogenie.treesimulator.tree import Tree
 from phylogenie.treesimulator.utils import (
     compute_mean_leaf_pairwise_distance,
     compute_sackin_index,
+    count_hops,
     get_distance,
     get_mrca,
     get_node_ages,
@@ -66,6 +67,7 @@ __all__ = [
     "Tree",
     "compute_mean_leaf_pairwise_distance",
     "compute_sackin_index",
+    "count_hops",
     "get_distance",
     "get_mrca",
     "get_node_ages",

{phylogenie-3.1.11 → phylogenie-3.1.18}/src/phylogenie/treesimulator/gillespie.py

@@ -25,6 +25,7 @@ def simulate_tree(
     seed: int | None = None,
     timeout: float = np.inf,
     acceptance_criterion: Callable[[Tree], bool] | None = None,
+    logs: dict[str, Callable[[Tree], Any]] | 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.")
@@ -85,7 +86,13 @@ def simulate_tree(
                 model.sample(individual, current_time, True)
 
         tree = model.get_sampled_tree()
-        if acceptance_criterion is None or acceptance_criterion(tree):
+
+        if acceptance_criterion is not None and not acceptance_criterion(tree):
+            continue
+
+        if logs is not None:
+            for key, func in logs.items():
+                metadata[key] = func(tree)
             return (tree, metadata)
 
 
@@ -102,6 +109,7 @@ def generate_trees(
     n_jobs: int = -1,
     timeout: float = np.inf,
     acceptance_criterion: Callable[[Tree], bool] | None = None,
+    logs: dict[str, Callable[[Tree], Any]] | None = None,
 ) -> pd.DataFrame:
     if isinstance(output_dir, str):
         output_dir = Path(output_dir)
@@ -121,6 +129,7 @@ def generate_trees(
                     seed=seed,
                     timeout=timeout,
                     acceptance_criterion=acceptance_criterion,
+                    logs=logs,
                 )
                 metadata["file_id"] = i
                 if node_features is not None:
@@ -128,7 +137,7 @@ def generate_trees(
                 dump_newick(tree, output_dir / f"{i}.nwk")
                 return metadata
             except TimeoutError:
-                print("Simulation timed out, retrying with a different seed...")
+                print("Simulation timed out. Retrying with a different seed...")
             seed += 1
 
     rng = default_rng(seed)

{phylogenie-3.1.11 → phylogenie-3.1.18/src/phylogenie.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: phylogenie
-Version: 3.1.11
+Version: 3.1.18
 Summary: Generate phylogenetic datasets with minimal setup effort
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown