phylogenie 2.1.30__py3-none-any.whl → 3.1.1__py3-none-any.whl

phylogenie/__init__.py

@@ -11,7 +11,7 @@ from phylogenie.generators import (
     FBDTreeDatasetGenerator,
     TreeDatasetGeneratorConfig,
 )
-from phylogenie.io import dump_newick, load_fasta, load_newick, load_nexus
+from phylogenie.io import load_fasta
 from phylogenie.msa import MSA
 from phylogenie.skyline import (
     SkylineMatrix,
@@ -25,41 +25,47 @@ from phylogenie.skyline import (
     skyline_parameter,
     skyline_vector,
 )
-from phylogenie.tree import Tree
 from phylogenie.treesimulator import (
     Birth,
     BirthWithContactTracing,
     Death,
     Event,
     EventType,
+    Feature,
     Migration,
     Mutation,
     Sampling,
     SamplingWithContactTracing,
+    Tree,
+    compute_mean_leaf_pairwise_distance,
+    compute_sackin_index,
+    dump_newick,
     generate_trees,
     get_BD_events,
     get_BDEI_events,
     get_BDSS_events,
     get_canonical_events,
     get_contact_tracing_events,
+    get_distance,
     get_epidemiological_events,
     get_FBD_events,
-    simulate_tree,
-)
-from phylogenie.utils import (
-    compute_colless_index,
-    compute_mean_leaf_pairwise_distance,
-    compute_sackin_index,
-    get_distance,
     get_mrca,
+    get_mutation_id,
     get_node_depth_levels,
     get_node_depths,
     get_node_height_levels,
     get_node_heights,
     get_node_leaf_counts,
+    get_node_state,
+    load_newick,
+    load_nexus,
+    set_features,
+    simulate_tree,
 )
 
 __all__ = [
+    "Coloring",
+    "draw_tree",
     "AliSimDatasetGenerator",
     "BDEITreeDatasetGenerator",
     "BDSSTreeDatasetGenerator",
@@ -69,51 +75,52 @@ __all__ = [
     "DatasetGeneratorConfig",
     "EpidemiologicalTreeDatasetGenerator",
     "FBDTreeDatasetGenerator",
+    "TreeDatasetGeneratorConfig",
+    "load_fasta",
+    "MSA",
     "SkylineMatrix",
     "SkylineMatrixCoercible",
-    "skyline_matrix",
     "SkylineParameter",
     "SkylineParameterLike",
-    "skyline_parameter",
     "SkylineVector",
     "SkylineVectorCoercible",
     "SkylineVectorLike",
+    "skyline_matrix",
+    "skyline_parameter",
     "skyline_vector",
-    "Tree",
-    "TreeDatasetGeneratorConfig",
     "Birth",
     "BirthWithContactTracing",
     "Death",
     "Event",
     "EventType",
+    "Feature",
     "Migration",
     "Mutation",
     "Sampling",
     "SamplingWithContactTracing",
+    "Tree",
+    "compute_mean_leaf_pairwise_distance",
+    "compute_sackin_index",
+    "dump_newick",
+    "generate_trees",
     "get_BD_events",
     "get_BDEI_events",
     "get_BDSS_events",
     "get_canonical_events",
     "get_contact_tracing_events",
+    "get_distance",
     "get_epidemiological_events",
     "get_FBD_events",
-    "generate_trees",
-    "simulate_tree",
-    "dump_newick",
-    "load_nexus",
-    "load_fasta",
-    "load_newick",
-    "MSA",
-    "Coloring",
-    "draw_tree",
-    "compute_colless_index",
-    "compute_mean_leaf_pairwise_distance",
-    "compute_sackin_index",
-    "get_distance",
     "get_mrca",
-    "get_node_depths",
+    "get_mutation_id",
     "get_node_depth_levels",
-    "get_node_heights",
+    "get_node_depths",
     "get_node_height_levels",
+    "get_node_heights",
     "get_node_leaf_counts",
+    "get_node_state",
+    "load_newick",
+    "load_nexus",
+    "set_features",
+    "simulate_tree",
 ]

phylogenie/draw.py

@@ -1,6 +1,5 @@
 from enum import Enum
-from itertools import islice
-from typing import Any
+from typing import Any, Callable
 
 import matplotlib.colors as mcolors
 import matplotlib.patches as mpatches
@@ -8,8 +7,7 @@ import matplotlib.pyplot as plt
 from matplotlib.axes import Axes
 from mpl_toolkits.axes_grid1.inset_locator import inset_axes  # pyright: ignore
 
-from phylogenie.tree import Tree
-from phylogenie.utils import get_node_depth_levels, get_node_depths
+from phylogenie.treesimulator import Tree, get_node_depth_levels, get_node_depths
 
 
 class Coloring(str, Enum):
@@ -26,7 +24,7 @@ def _draw_colored_tree(tree: Tree, ax: Axes, colors: Color | dict[Tree, Color])
 
     xs = (
         get_node_depth_levels(tree)
-        if any(node.branch_length is None for node in islice(tree, 1, None))
+        if any(node.branch_length is None for node in tree.iter_descendants())
         else get_node_depths(tree)
     )
     ys: dict[Tree, float] = {node: i for i, node in enumerate(tree.get_leaves())}
@@ -50,7 +48,7 @@ def _draw_colored_tree(tree: Tree, ax: Axes, colors: Color | dict[Tree, Color])
 def draw_tree(
     tree: Tree,
     ax: Axes | None = None,
-    color_by: str | None = None,
+    color_by: str | dict[Tree, Any] | None = None,
     coloring: str | Coloring | None = None,
     default_color: Color = "black",
     cmap: str | None = None,
@@ -69,33 +67,35 @@ def draw_tree(
     if color_by is None:
         return _draw_colored_tree(tree, ax, colors=default_color)
 
-    features = [node.get(color_by) for node in tree if color_by in node.features]
+    if isinstance(color_by, dict):
+        features = {node: color_by[node] for node in tree if node in color_by}
+    else:
+        features = {node: node[color_by] for node in tree if color_by in node.metadata}
 
     if coloring is None:
         coloring = (
             Coloring.CONTINUOUS
-            if any(isinstance(f, float) for f in features)
+            if any(isinstance(f, float) for f in features.values())
             else Coloring.DISCRETE
         )
 
+    def _get_colors(feature_map: Callable[[Any], Color]) -> dict[Tree, Color]:
+        return {
+            node: feature_map(features[node]) if node in features else default_color
+            for node in tree
+        }
+
     if coloring == Coloring.DISCRETE:
-        if any(isinstance(f, float) for f in features):
+        if any(isinstance(f, float) for f in features.values()):
             raise ValueError(
                 "Discrete coloring selected but feature values are not all categorical."
             )
 
         colormap = plt.get_cmap("tab20" if cmap is None else cmap)
         feature_colors = {
-            f: mcolors.to_hex(colormap(i)) for i, f in enumerate(set(features))
-        }
-        colors = {
-            node: (
-                feature_colors[node.get(color_by)]
-                if color_by in node.features
-                else default_color
-            )
-            for node in tree
+            f: mcolors.to_hex(colormap(i)) for i, f in enumerate(set(features.values()))
         }
+        colors = _get_colors(lambda f: feature_colors[f])
 
         if show_legend:
             legend_handles = [
@@ -105,7 +105,7 @@ def draw_tree(
                 )
                 for f in feature_colors
             ]
-            if any(color_by not in node.features for node in tree):
+            if any(color_by not in node.metadata for node in tree):
                 legend_handles.append(mpatches.Patch(color=default_color, label="NA"))
             if legend_kwargs is None:
                 legend_kwargs = {}
@@ -114,18 +114,11 @@ def draw_tree(
         return _draw_colored_tree(tree, ax, colors)
 
     if coloring == Coloring.CONTINUOUS:
-        vmin = min(features) if vmin is None else vmin
-        vmax = max(features) if vmax is None else vmax
+        vmin = min(features.values()) if vmin is None else vmin
+        vmax = max(features.values()) if vmax is None else vmax
         norm = mcolors.Normalize(vmin=vmin, vmax=vmax)
         colormap = plt.get_cmap("viridis" if cmap is None else cmap)
-        colors = {
-            node: (
-                colormap(norm(float(node.get(color_by))))
-                if color_by in node.features
-                else default_color
-            )
-            for node in tree
-        }
+        colors = _get_colors(lambda f: colormap(norm(float(f))))
 
         if show_hist:
             default_hist_axes_kwargs = {"width": "25%", "height": "25%"}

phylogenie/generators/alisim.py

@@ -8,8 +8,7 @@ from numpy.random import Generator, default_rng
 from phylogenie.generators.dataset import DatasetGenerator, DataType
 from phylogenie.generators.factories import data, string
 from phylogenie.generators.trees import TreeDatasetGeneratorConfig
-from phylogenie.io import dump_newick
-from phylogenie.utils import get_node_depths
+from phylogenie.treesimulator import dump_newick, get_node_depths
 
 MSAS_DIRNAME = "MSAs"
 TREES_DIRNAME = "trees"
@@ -60,26 +59,26 @@ class AliSimDatasetGenerator(DatasetGenerator):
             tree_filename = f"{filename}.temp-tree"
             msa_filename = filename
 
-        d: dict[str, Any] = {"file_id": Path(msa_filename).stem}
+        md: dict[str, Any] = {"file_id": Path(msa_filename).stem}
         rng = default_rng(seed)
         while True:
-            d.update(data(context, rng))
+            md.update(data(context, rng))
             try:
-                tree, metadata = self.trees.simulate_one(d, seed)
+                tree, metadata = self.trees.simulate_one(md, seed)
                 break
             except TimeoutError:
                 print(
                     "Tree simulation timed out, retrying with different parameters..."
                 )
-        d.update(metadata)
+        md.update(metadata)
 
         times = get_node_depths(tree)
         for leaf in tree.get_leaves():
             leaf.name += f"|{times[leaf]}"
         dump_newick(tree, f"{tree_filename}.nwk")
 
-        self._generate_one_from_tree(msa_filename, f"{tree_filename}.nwk", rng, d)
+        self._generate_one_from_tree(msa_filename, f"{tree_filename}.nwk", rng, md)
         if not self.keep_trees:
             os.remove(f"{tree_filename}.nwk")
 
-        return d
+        return md

phylogenie/generators/configs.py

@@ -1,7 +1,29 @@
+from typing import Any
+
+from numpy.random import Generator
+from pydantic import BaseModel, ConfigDict
+
 import phylogenie.typings as pgt
-from phylogenie.models import Distribution, StrictBaseModel
 from phylogenie.treesimulator import EventType
 
+
+class StrictBaseModel(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+
+class Distribution(BaseModel):
+    type: str
+    model_config = ConfigDict(extra="allow")
+
+    @property
+    def args(self) -> dict[str, Any]:
+        assert self.model_extra is not None
+        return self.model_extra
+
+    def __call__(self, rng: Generator) -> Any:
+        return getattr(rng, self.type)(**self.args)
+
+
 Integer = str | int
 Scalar = str | pgt.Scalar
 ManyScalars = str | pgt.Many[Scalar]

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.models import Distribution, StrictBaseModel
+from phylogenie.generators.configs import Distribution, StrictBaseModel
 
 
 class DataType(str, Enum):

phylogenie/generators/factories.py

@@ -1,5 +1,5 @@
 import re
-from typing import Any
+from typing import Any, Callable
 
 import numpy as np
 from numpy.random import Generator
@@ -8,7 +8,6 @@ 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,
@@ -24,11 +23,7 @@ def _eval_expression(expression: str, data: dict[str, Any]) -> Any:
     return np.array(
         eval(
             expression,
-            {
-                "__builtins__": __builtins__,
-                "np": np,
-                **{k: np.array(v) for k, v in data.items()},
-            },
+            {"np": np, **{k: np.array(v) for k, v in data.items()}},
         )
     ).tolist()
 
@@ -214,12 +209,12 @@ def skyline_matrix(
     return SkylineMatrix(value=value, change_times=change_times)
 
 
-def distribution(x: Distribution, data: dict[str, Any]) -> Distribution:
+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)
-    return Distribution(type=x.type, **args)
+    return cfg.Distribution(type=x.type, **args)
 
 
 def mutations(
@@ -227,11 +222,14 @@ def mutations(
     data: dict[str, Any],
     states: set[str],
     rates_to_log: list[EventType] | None,
+    rng: Generator,
 ) -> list[Mutation]:
     mutations: list[Mutation] = []
     for m in x:
         rate = skyline_parameter(m.rate, data)
-        rate_scalers = {k: distribution(v, data) for k, v in m.rate_scalers.items()}
+        rate_scalers: dict[EventType, Callable[[], float]] = {
+            k: lambda: distribution(v, data)(rng) for k, v in m.rate_scalers.items()
+        }
         if m.state is None:
             mutations.extend(
                 Mutation(s, rate, rate_scalers, rates_to_log) for s in states
@@ -241,11 +239,10 @@ def mutations(
     return mutations
 
 
-def data(context: dict[str, Distribution] | None, rng: Generator) -> dict[str, Any]:
+def data(context: dict[str, cfg.Distribution] | None, rng: Generator) -> dict[str, Any]:
     if context is None:
         return {}
     data: dict[str, Any] = {}
     for k, v in context.items():
-        dist = distribution(v, data)
-        data[k] = np.array(getattr(rng, dist.type)(**dist.args)).tolist()
+        data[k] = np.array(distribution(v, data)(rng)).tolist()
     return data

phylogenie/generators/trees.py

@@ -1,13 +1,14 @@
 from abc import abstractmethod
 from enum import Enum
 from pathlib import Path
-from typing import Annotated, Any, Literal
+from typing import Annotated, Any, Callable, Literal
 
 import numpy as np
 from numpy.random import default_rng
 from pydantic import Field
 
 import phylogenie.generators.configs as cfg
+from phylogenie.generators.configs import Distribution
 from phylogenie.generators.dataset import DatasetGenerator, DataType
 from phylogenie.generators.factories import (
     data,
@@ -18,13 +19,12 @@ from phylogenie.generators.factories import (
     skyline_parameter,
     skyline_vector,
 )
-from phylogenie.io import dump_newick
-from phylogenie.models import Distribution
-from phylogenie.tree import Tree
 from phylogenie.treesimulator import (
     Event,
     EventType,
     Feature,
+    Tree,
+    dump_newick,
     get_BD_events,
     get_BDEI_events,
     get_BDSS_events,
@@ -50,13 +50,13 @@ class TreeDatasetGenerator(DatasetGenerator):
     data_type: Literal[DataType.TREES] = DataType.TREES
     mutations: list[cfg.Mutation] = Field(default_factory=lambda: [])
     rates_to_log: list[EventType] | None = None
-    min_tips: cfg.Integer = 1
-    max_tips: cfg.Integer | None = None
+    n_tips: cfg.Integer | None = None
     max_time: cfg.Scalar = np.inf
     init_state: str | None = None
     sampling_probability_at_present: cfg.Scalar = 0.0
     timeout: float = np.inf
     node_features: list[Feature] | None = None
+    acceptance_criterion: str | None = None
 
     @abstractmethod
     def _get_events(self, data: dict[str, Any]) -> list[Event]: ...
@@ -71,11 +71,19 @@ class TreeDatasetGenerator(DatasetGenerator):
         )
         events = self._get_events(data)
         states = {e.state for e in events}
-        events += mutations(self.mutations, data, states, self.rates_to_log)
+        events += mutations(
+            self.mutations, data, states, self.rates_to_log, default_rng(seed)
+        )
+        acceptance_criterion: None | Callable[[Tree], bool] = (
+            None
+            if self.acceptance_criterion is None
+            else lambda tree: eval(
+                self.acceptance_criterion, {}, {"tree": tree}  # pyright: ignore
+            )
+        )
         return simulate_tree(
             events=events,
-            min_tips=integer(self.min_tips, data),
-            max_tips=None if self.max_tips is None else integer(self.max_tips, data),
+            n_tips=None if self.n_tips is None else integer(self.n_tips, data),
             max_time=scalar(self.max_time, data),
             init_state=init_state,
             sampling_probability_at_present=scalar(
@@ -83,6 +91,7 @@ class TreeDatasetGenerator(DatasetGenerator):
             ),
             seed=seed,
             timeout=self.timeout,
+            acceptance_criterion=acceptance_criterion,
         )
 
     def generate_one(
@@ -157,7 +166,7 @@ class FBDTreeDatasetGenerator(TreeDatasetGenerator):
 class ContactTracingTreeDatasetGenerator(TreeDatasetGenerator):
     max_notified_contacts: cfg.Integer = 1
     notification_probability: cfg.SkylineParameter = 0.0
-    sampling_rate_after_notification: cfg.SkylineParameter = np.inf
+    sampling_rate_after_notification: cfg.SkylineParameter = 2**32
     samplable_states_after_notification: list[str] | None = None
 
     @abstractmethod

phylogenie/io/__init__.py

@@ -1,5 +1,3 @@
 from phylogenie.io.fasta import load_fasta
-from phylogenie.io.newick import dump_newick, load_newick
-from phylogenie.io.nexus import load_nexus
 
-__all__ = ["load_fasta", "load_newick", "dump_newick", "load_nexus"]
+__all__ = ["load_fasta"]

phylogenie/mixins.py

@@ -0,0 +1,41 @@
+from types import MappingProxyType
+from typing import Any, Mapping, Optional
+
+
+class MetadataMixin:
+    """A mixin that provides metadata management with dictionary-like access."""
+
+    def __init__(self) -> None:
+        self._metadata: dict[str, Any] = {}
+
+    @property
+    def metadata(self) -> Mapping[str, Any]:
+        """Return a read-only view of all metadata."""
+        return MappingProxyType(self._metadata)
+
+    def set(self, key: str, value: Any) -> None:
+        """Set or update a metadata value."""
+        self._metadata[key] = value
+
+    def update(self, metadata: Mapping[str, Any]) -> None:
+        """Bulk update metadata values."""
+        self._metadata.update(metadata)
+
+    def get(self, key: str, default: Optional[Any] = None) -> Any:
+        """Get a metadata value, returning `default` if not found."""
+        return self._metadata.get(key, default)
+
+    def delete(self, key: str) -> None:
+        """Delete a metadata if it exists, else do nothing."""
+        self._metadata.pop(key, None)
+
+    def clear(self) -> None:
+        """Remove all metadata."""
+        self._metadata.clear()
+
+    # Dict-like behavior
+    def __getitem__(self, key: str) -> Any:
+        return self._metadata[key]
+
+    def __setitem__(self, key: str, value: Any) -> None:
+        self._metadata[key] = value

phylogenie/treesimulator/__init__.py

@@ -19,7 +19,20 @@ from phylogenie.treesimulator.events import (
 )
 from phylogenie.treesimulator.features import Feature, set_features
 from phylogenie.treesimulator.gillespie import generate_trees, simulate_tree
+from phylogenie.treesimulator.io import dump_newick, load_newick, load_nexus
 from phylogenie.treesimulator.model import get_node_state
+from phylogenie.treesimulator.tree import Tree
+from phylogenie.treesimulator.utils import (
+    compute_mean_leaf_pairwise_distance,
+    compute_sackin_index,
+    get_distance,
+    get_mrca,
+    get_node_depth_levels,
+    get_node_depths,
+    get_node_height_levels,
+    get_node_heights,
+    get_node_leaf_counts,
+)
 
 __all__ = [
     "Birth",
@@ -38,10 +51,23 @@ __all__ = [
     "get_contact_tracing_events",
     "get_epidemiological_events",
     "get_FBD_events",
-    "generate_trees",
-    "simulate_tree",
     "get_mutation_id",
-    "get_node_state",
     "Feature",
     "set_features",
+    "simulate_tree",
+    "dump_newick",
+    "load_newick",
+    "load_nexus",
+    "generate_trees",
+    "get_node_state",
+    "Tree",
+    "compute_mean_leaf_pairwise_distance",
+    "compute_sackin_index",
+    "get_distance",
+    "get_mrca",
+    "get_node_depth_levels",
+    "get_node_depths",
+    "get_node_height_levels",
+    "get_node_heights",
+    "get_node_leaf_counts",
 ]

phylogenie/treesimulator/events/base.py

@@ -2,7 +2,6 @@ from abc import ABC, abstractmethod
 from enum import Enum
 from typing import Any
 
-import numpy as np
 from numpy.random import Generator
 
 from phylogenie.skyline import SkylineParameterLike, skyline_parameter
@@ -32,8 +31,6 @@ class Event(ABC):
     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

phylogenie/treesimulator/events/contact_tracing.py

@@ -2,7 +2,6 @@ from collections import defaultdict
 from collections.abc import Sequence
 from copy import deepcopy
 
-import numpy as np
 from numpy.random import Generator
 
 from phylogenie.skyline import SkylineParameterLike, skyline_parameter
@@ -32,10 +31,10 @@ class BirthWithContactTracing(Event):
     def apply(self, model: Model, events: list[Event], time: float, rng: Generator):
         individual = self.draw_individual(model, rng)
         new_individual = model.birth_from(individual, self.child_state, time)
-        if CONTACTS_KEY not in model.context:
-            model.context[CONTACTS_KEY] = defaultdict(list)
-        model.context[CONTACTS_KEY][individual].append(new_individual)
-        model.context[CONTACTS_KEY][new_individual].append(individual)
+        if CONTACTS_KEY not in model.metadata:
+            model[CONTACTS_KEY] = defaultdict(list)
+        model[CONTACTS_KEY][individual].append(new_individual)
+        model[CONTACTS_KEY][new_individual].append(individual)
 
     def __repr__(self) -> str:
         return f"BirthWithContactTracing(state={self.state}, rate={self.rate}, child_state={self.child_state})"
@@ -59,9 +58,9 @@ class SamplingWithContactTracing(Event):
         individual = self.draw_individual(model, rng)
         model.sample(individual, time, True)
         population = model.get_population()
-        if CONTACTS_KEY not in model.context:
+        if CONTACTS_KEY not in model.metadata:
             return
-        contacts = model.context[CONTACTS_KEY][individual]
+        contacts = model[CONTACTS_KEY][individual]
         for contact in contacts[-self.max_notified_contacts :]:
             if contact in population:
                 state = model.get_state(contact)
@@ -76,8 +75,8 @@ class SamplingWithContactTracing(Event):
 def get_contact_tracing_events(
     events: Sequence[Event],
     max_notified_contacts: int = 1,
-    notification_probability: SkylineParameterLike = 1,
-    sampling_rate_after_notification: SkylineParameterLike = np.inf,
+    notification_probability: SkylineParameterLike = 0.0,
+    sampling_rate_after_notification: SkylineParameterLike = 2**32,
     samplable_states_after_notification: list[str] | None = None,
 ) -> list[Event]:
     ct_events: list[Event] = []