phylogenie 2.1.0__tar.gz → 2.1.1__tar.gz

{phylogenie-2.1.0 → phylogenie-2.1.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: phylogenie
-Version: 2.1.0
+Version: 2.1.1
 Summary: Generate phylogenetic datasets with minimal setup effort
 Author: Gabriele Marino
 Author-email: gabmarino.8601@gmail.com
@@ -24,7 +24,6 @@ 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/)
-[![PyPI downloads](https://shields.io/pypi/dm/phylogenie)](https://pypi.org/project/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.0 → phylogenie-2.1.1}/README.md

@@ -6,7 +6,6 @@
 
 [![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/)
-[![PyPI downloads](https://shields.io/pypi/dm/phylogenie)](https://pypi.org/project/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.0 → phylogenie-2.1.1}/phylogenie/__init__.py

@@ -10,7 +10,7 @@ from phylogenie.generators import (
     FBDTreeDatasetGenerator,
     TreeDatasetGeneratorConfig,
 )
-from phylogenie.io import load_fasta, load_newick
+from phylogenie.io import dump_newick, load_fasta, load_newick
 from phylogenie.msa import MSA
 from phylogenie.skyline import (
     SkylineMatrix,
@@ -26,12 +26,21 @@ from phylogenie.skyline import (
 )
 from phylogenie.tree import Tree
 from phylogenie.treesimulator import (
+    Birth,
+    BirthWithContactTracing,
+    Death,
     Event,
+    Migration,
+    Mutation,
+    MutationTargetType,
+    Sampling,
+    SamplingWithContactTracing,
     generate_trees,
     get_BD_events,
     get_BDEI_events,
     get_BDSS_events,
     get_canonical_events,
+    get_contact_tracing_events,
     get_epidemiological_events,
     get_FBD_events,
     simulate_tree,
@@ -59,15 +68,25 @@ __all__ = [
     "skyline_vector",
     "Tree",
     "TreeDatasetGeneratorConfig",
+    "Birth",
+    "BirthWithContactTracing",
+    "Death",
     "Event",
+    "Migration",
+    "Mutation",
+    "MutationTargetType",
+    "Sampling",
+    "SamplingWithContactTracing",
     "get_BD_events",
     "get_BDEI_events",
     "get_BDSS_events",
     "get_canonical_events",
+    "get_contact_tracing_events",
     "get_epidemiological_events",
     "get_FBD_events",
     "generate_trees",
     "simulate_tree",
+    "dump_newick",
     "load_fasta",
     "load_newick",
     "MSA",

{phylogenie-2.1.0 → phylogenie-2.1.1}/phylogenie/generators/configs.py

@@ -1,12 +1,6 @@
-from pydantic import BaseModel, ConfigDict
-
 import phylogenie.typings as pgt
-
-
-class Distribution(BaseModel):
-    type: str
-    model_config = ConfigDict(extra="allow")
-
+from phylogenie.treesimulator import MutationTargetType
+from phylogenie.utils import Distribution, StrictBaseModel
 
 Integer = str | int
 Scalar = str | pgt.Scalar
@@ -15,10 +9,6 @@ OneOrManyScalars = Scalar | pgt.Many[Scalar]
 OneOrMany2DScalars = Scalar | pgt.Many2D[Scalar]
 
 
-class StrictBaseModel(BaseModel):
-    model_config = ConfigDict(extra="forbid")
-
-
 class SkylineParameterModel(StrictBaseModel):
     value: ManyScalars
     change_times: ManyScalars
@@ -37,3 +27,12 @@ class SkylineMatrixModel(StrictBaseModel):
 SkylineParameter = Scalar | SkylineParameterModel
 SkylineVector = str | pgt.Scalar | pgt.Many[SkylineParameter] | SkylineVectorModel
 SkylineMatrix = str | pgt.Scalar | pgt.Many[SkylineVector] | SkylineMatrixModel | None
+
+
+class Event(StrictBaseModel):
+    states: str | list[str] | None = None
+    rate: SkylineParameter
+
+
+class Mutation(Event):
+    rate_scalers: dict[MutationTargetType, Distribution]

{phylogenie-2.1.0 → phylogenie-2.1.1}/phylogenie/generators/dataset.py

@@ -10,8 +10,8 @@ import pandas as pd
 from numpy.random import Generator, default_rng
 from tqdm import tqdm
 
-import phylogenie.generators.configs as cfg
-from phylogenie.generators.factories import eval_expression
+from phylogenie.generators.factories import distribution
+from phylogenie.utils import Distribution, StrictBaseModel
 
 
 class DataType(str, Enum):
@@ -23,12 +23,12 @@ DATA_DIRNAME = "data"
 METADATA_FILENAME = "metadata.csv"
 
 
-class DatasetGenerator(ABC, cfg.StrictBaseModel):
+class DatasetGenerator(ABC, StrictBaseModel):
     output_dir: str = "phylogenie-outputs"
     n_samples: int | dict[str, int] = 1
     n_jobs: int = -1
     seed: int | None = None
-    context: dict[str, cfg.Distribution] | None = None
+    context: dict[str, Distribution] | None = None
 
     @abstractmethod
     def _generate_one(
@@ -56,11 +56,8 @@ class DatasetGenerator(ABC, cfg.StrictBaseModel):
         data: list[dict[str, Any]] = [{} for _ in range(n_samples)]
         if self.context is not None:
             for d, (k, v) in product(data, self.context.items()):
-                args = v.model_extra if v.model_extra is not None else {}
-                for arg_name, arg_value in args.items():
-                    if isinstance(arg_value, str):
-                        args[arg_name] = eval_expression(arg_value, d)
-                d[k] = np.array(getattr(rng, v.type)(**args)).tolist()
+                dist = distribution(v, d)
+                d[k] = np.array(getattr(rng, dist.type)(**dist.args)).tolist()
             df = pd.DataFrame([{"file_id": str(i), **d} for i, d in enumerate(data)])
             df.to_csv(os.path.join(output_dir, METADATA_FILENAME), index=False)
 

{phylogenie-2.1.0 → phylogenie-2.1.1}/phylogenie/generators/factories.py

@@ -14,9 +14,10 @@ from phylogenie.skyline import (
     SkylineVector,
     SkylineVectorCoercible,
 )
+from phylogenie.utils import Distribution
 
 
-def eval_expression(expression: str, data: dict[str, Any]) -> Any:
+def _eval_expression(expression: str, data: dict[str, Any]) -> Any:
     return np.array(
         eval(
             expression,
@@ -29,9 +30,17 @@ def eval_expression(expression: str, data: dict[str, Any]) -> Any:
     ).tolist()
 
 
+def distribution(x: Distribution, data: dict[str, Any]) -> 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)
+
+
 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(
@@ -42,7 +51,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(
@@ -53,7 +62,7 @@ def scalar(x: cfg.Scalar, data: dict[str, Any]) -> pgt.Scalar:
 
 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(
@@ -66,7 +75,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(
@@ -92,7 +101,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(
@@ -107,7 +116,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:
@@ -135,7 +144,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(
@@ -153,7 +162,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(
@@ -168,7 +177,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:

{phylogenie-2.1.0 → phylogenie-2.1.1}/phylogenie/generators/trees.py

@@ -9,6 +9,7 @@ from pydantic import Field
 import phylogenie.generators.configs as cfg
 from phylogenie.generators.dataset import DatasetGenerator, DataType
 from phylogenie.generators.factories import (
+    distribution,
     integer,
     scalar,
     skyline_matrix,
@@ -19,6 +20,7 @@ from phylogenie.io import dump_newick
 from phylogenie.tree import Tree
 from phylogenie.treesimulator import (
     Event,
+    Mutation,
     get_BD_events,
     get_BDEI_events,
     get_BDSS_events,
@@ -46,7 +48,6 @@ class TreeDatasetGenerator(DatasetGenerator):
     max_time: cfg.Scalar = np.inf
     init_state: str | None = None
     sampling_probability_at_present: cfg.Scalar = 0.0
-    max_tries: int | None = None
 
     @abstractmethod
     def _get_events(self, data: dict[str, Any]) -> list[Event]: ...
@@ -67,7 +68,6 @@ class TreeDatasetGenerator(DatasetGenerator):
             sampling_probability_at_present=scalar(
                 self.sampling_probability_at_present, data
             ),
-            max_tries=self.max_tries,
             seed=int(rng.integers(2**32)),
         )
 
@@ -127,11 +127,12 @@ class FBDTreeDatasetGenerator(TreeDatasetGenerator):
         )
 
 
-class TreeDatasetGeneratorWithContactTracing(TreeDatasetGenerator):
+class TreeDatasetGeneratorForEpidemiology(TreeDatasetGenerator):
     max_notified_contacts: cfg.Integer = 1
     notification_probability: cfg.SkylineParameter = 0.0
     sampling_rate_after_notification: cfg.SkylineParameter = np.inf
     samplable_states_after_notification: list[str] | None = None
+    mutations: tuple[cfg.Mutation, ...] = Field(default_factory=tuple)
 
     @abstractmethod
     def _get_base_events(self, data: dict[str, Any]) -> list[Event]: ...
@@ -150,10 +151,30 @@ class TreeDatasetGeneratorWithContactTracing(TreeDatasetGenerator):
                 ),
                 samplable_states_after_notification=self.samplable_states_after_notification,
             )
+        all_states = list({e.state for e in events})
+        for mutation in self.mutations:
+            states = mutation.states
+            if isinstance(states, str):
+                states = [states]
+            elif states is None:
+                states = all_states
+            for state in states:
+                if state not in all_states:
+                    raise ValueError(
+                        f"Mutation state '{state}' is not found in states {all_states}."
+                    )
+                rate_scalers = {
+                    t: distribution(r, data) for t, r in mutation.rate_scalers.items()
+                }
+                events.append(
+                    Mutation(
+                        state, skyline_parameter(mutation.rate, data), rate_scalers
+                    )
+                )
         return events
 
 
-class EpidemiologicalTreeDatasetGenerator(TreeDatasetGeneratorWithContactTracing):
+class EpidemiologicalTreeDatasetGenerator(TreeDatasetGeneratorForEpidemiology):
     parameterization: Literal[ParameterizationType.EPIDEMIOLOGICAL] = (
         ParameterizationType.EPIDEMIOLOGICAL
     )
@@ -179,7 +200,7 @@ class EpidemiologicalTreeDatasetGenerator(TreeDatasetGeneratorWithContactTracing
         )
 
 
-class BDTreeDatasetGenerator(TreeDatasetGeneratorWithContactTracing):
+class BDTreeDatasetGenerator(TreeDatasetGeneratorForEpidemiology):
     parameterization: Literal[ParameterizationType.BD] = ParameterizationType.BD
     reproduction_number: cfg.SkylineParameter
     infectious_period: cfg.SkylineParameter
@@ -193,7 +214,7 @@ class BDTreeDatasetGenerator(TreeDatasetGeneratorWithContactTracing):
         )
 
 
-class BDEITreeDatasetGenerator(TreeDatasetGeneratorWithContactTracing):
+class BDEITreeDatasetGenerator(TreeDatasetGeneratorForEpidemiology):
     parameterization: Literal[ParameterizationType.BDEI] = ParameterizationType.BDEI
     reproduction_number: cfg.SkylineParameter
     infectious_period: cfg.SkylineParameter
@@ -209,7 +230,7 @@ class BDEITreeDatasetGenerator(TreeDatasetGeneratorWithContactTracing):
         )
 
 
-class BDSSTreeDatasetGenerator(TreeDatasetGeneratorWithContactTracing):
+class BDSSTreeDatasetGenerator(TreeDatasetGeneratorForEpidemiology):
     parameterization: Literal[ParameterizationType.BDSS] = ParameterizationType.BDSS
     reproduction_number: cfg.SkylineParameter
     infectious_period: cfg.SkylineParameter

{phylogenie-2.1.0 → phylogenie-2.1.1}/phylogenie/io.py

@@ -24,21 +24,28 @@ def _parse_newick(newick: str) -> Tree:
             stack.append(current_nodes)
             current_nodes = []
         else:
-            id = _parse_chars([":", ",", ")", ";"])
+            id = _parse_chars([":", ",", ")", ";", "["])
             branch_length = None
             if newick[i] == ":":
                 i += 1
-                branch_length = _parse_chars([",", ")", ";"])
+                branch_length = float(_parse_chars([",", ")", ";", "["]))
 
-            current_node = Tree(
-                id=id,
-                branch_length=(None if branch_length is None else float(branch_length)),
-            )
+            current_node = Tree(id, branch_length)
             for node in current_children:
                 current_node.add_child(node)
                 current_children = []
             current_nodes.append(current_node)
 
+            if newick[i] == "[":
+                i += 1
+                features = _parse_chars(["]"]).split(":")
+                i += 1
+                if features[0] != "&&NHX":
+                    raise ValueError(f"Expected '&&NHX' for node features.")
+                for feature in features[1:]:
+                    key, value = feature.split("=", 1)
+                    current_node.set(key, eval(value))
+
             if newick[i] == ")":
                 current_children = current_nodes
                 current_nodes = stack.pop()
@@ -47,7 +54,7 @@ def _parse_newick(newick: str) -> Tree:
 
         i += 1
 
-    raise ValueError("Newick string does not end with a semicolon.")
+    raise ValueError("Newick string is invalid.")
 
 
 def load_newick(filepath: str) -> Tree | list[Tree]:
@@ -63,6 +70,10 @@ def _to_newick(tree: Tree) -> str:
         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()}
+        features = [f"{k}={repr}" for k, repr in reprs.items()]
+        newick += f"[&&NHX:{':'.join(features)}]"
     return newick
 
 
@@ -83,13 +94,14 @@ def load_fasta(
             if not line.startswith(">"):
                 raise ValueError(f"Invalid FASTA format: expected '>', got '{line[0]}'")
             id = line[1:].strip()
+            time = None
             if extract_time_from_id is not None:
                 time = extract_time_from_id(id)
-            else:
+            elif "|" in id:
                 try:
                     time = float(id.split("|")[-1])
-                except:
-                    time = None
+                except ValueError:
+                    pass
             chars = next(f).strip()
             sequences.append(Sequence(id, chars, time))
     return MSA(sequences)

{phylogenie-2.1.0 → phylogenie-2.1.1}/phylogenie/tree.py

@@ -1,18 +1,41 @@
 from collections.abc import Iterator
+from typing import Any
 
 
 class Tree:
     def __init__(self, id: str = "", branch_length: float | None = None):
         self.id = id
         self.branch_length = branch_length
-        self.parent: Tree | None = None
-        self.children: list[Tree] = []
+        self._parent: Tree | None = None
+        self._children: list[Tree] = []
+        self._features: dict[str, Any] = {}
+
+    @property
+    def children(self) -> tuple["Tree", ...]:
+        return tuple(self._children)
+
+    @property
+    def parent(self) -> "Tree | None":
+        return self._parent
+
+    @property
+    def features(self) -> dict[str, Any]:
+        return self._features
 
     def add_child(self, child: "Tree") -> "Tree":
-        child.parent = self
-        self.children.append(child)
+        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, node: "Tree | None"):
+        self._parent = node
+        if node is not None:
+            node._children.append(self)
+
     def preorder_traversal(self) -> Iterator["Tree"]:
         yield self
         for child in self.children:
@@ -29,6 +52,9 @@ class Tree:
                 return node
         raise ValueError(f"Node with id {id} 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]
 
@@ -38,17 +64,17 @@ class Tree:
             raise ValueError(f"Branch length of node {self.id} is not set.")
         return self.branch_length + parent_time
 
-    def is_leaf(self) -> bool:
-        return not self.children
+    def set(self, key: str, value: Any) -> None:
+        self._features[key] = value
 
-    def copy(self) -> "Tree":
-        new_tree = Tree(self.id, self.branch_length)
-        for child in self.children:
-            new_tree.add_child(child.copy())
-        return new_tree
+    def get(self, key: str) -> Any:
+        return self._features.get(key)
+
+    def delete(self, key: str) -> None:
+        del self._features[key]
 
     def __iter__(self) -> Iterator["Tree"]:
         return self.preorder_traversal()
 
     def __repr__(self) -> str:
-        return f"TreeNode(id='{self.id}', branch_length={self.branch_length})"
+        return f"TreeNode(id='{self.id}', branch_length={self.branch_length}, features={self.features})"

{phylogenie-2.1.0 → phylogenie-2.1.1}/phylogenie/treesimulator/__init__.py

@@ -1,5 +1,13 @@
 from phylogenie.treesimulator.events import (
+    Birth,
+    BirthWithContactTracing,
+    Death,
     Event,
+    Migration,
+    Mutation,
+    MutationTargetType,
+    Sampling,
+    SamplingWithContactTracing,
     get_BD_events,
     get_BDEI_events,
     get_BDSS_events,
@@ -11,7 +19,15 @@ from phylogenie.treesimulator.events import (
 from phylogenie.treesimulator.gillespie import generate_trees, simulate_tree
 
 __all__ = [
+    "Birth",
+    "BirthWithContactTracing",
+    "Death",
     "Event",
+    "Migration",
+    "Mutation",
+    "MutationTargetType",
+    "Sampling",
+    "SamplingWithContactTracing",
     "get_BD_events",
     "get_BDEI_events",
     "get_BDSS_events",

phylogenie-2.1.1/phylogenie/treesimulator/events/__init__.py

@@ -0,0 +1,39 @@
+from phylogenie.treesimulator.events.contact_tracing import (
+    BirthWithContactTracing,
+    SamplingWithContactTracing,
+    get_contact_tracing_events,
+)
+from phylogenie.treesimulator.events.core import (
+    Birth,
+    Death,
+    Event,
+    Migration,
+    Sampling,
+    get_BD_events,
+    get_BDEI_events,
+    get_BDSS_events,
+    get_canonical_events,
+    get_epidemiological_events,
+    get_FBD_events,
+)
+from phylogenie.treesimulator.events.mutations import Mutation
+from phylogenie.treesimulator.events.mutations import TargetType as MutationTargetType
+
+__all__ = [
+    "Birth",
+    "BirthWithContactTracing",
+    "Death",
+    "Event",
+    "Migration",
+    "Mutation",
+    "Sampling",
+    "SamplingWithContactTracing",
+    "MutationTargetType",
+    "get_BD_events",
+    "get_BDEI_events",
+    "get_BDSS_events",
+    "get_canonical_events",
+    "get_contact_tracing_events",
+    "get_epidemiological_events",
+    "get_FBD_events",
+]