phylogenie 2.0.14__py3-none-any.whl → 2.1.1__py3-none-any.whl

phylogenie/__init__.py

@@ -8,12 +8,9 @@ from phylogenie.generators import (
     DatasetGeneratorConfig,
     EpidemiologicalTreeDatasetGenerator,
     FBDTreeDatasetGenerator,
-    SkylineMatrixModel,
-    SkylineParameterModel,
-    SkylineVectorModel,
     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,
@@ -29,11 +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,
@@ -51,27 +58,35 @@ __all__ = [
     "FBDTreeDatasetGenerator",
     "SkylineMatrix",
     "SkylineMatrixCoercible",
-    "SkylineMatrixModel",
     "skyline_matrix",
     "SkylineParameter",
     "SkylineParameterLike",
-    "SkylineParameterModel",
     "skyline_parameter",
     "SkylineVector",
     "SkylineVectorCoercible",
     "SkylineVectorLike",
-    "SkylineVectorModel",
     "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/generators/__init__.py

@@ -3,11 +3,6 @@ from typing import Annotated
 from pydantic import Field
 
 from phylogenie.generators.alisim import AliSimDatasetGenerator
-from phylogenie.generators.configs import (
-    SkylineMatrixModel,
-    SkylineParameterModel,
-    SkylineVectorModel,
-)
 from phylogenie.generators.dataset import DatasetGenerator
 from phylogenie.generators.trees import (
     BDEITreeDatasetGenerator,
@@ -34,7 +29,4 @@ __all__ = [
     "BDTreeDatasetGenerator",
     "BDEITreeDatasetGenerator",
     "BDSSTreeDatasetGenerator",
-    "SkylineMatrixModel",
-    "SkylineParameterModel",
-    "SkylineVectorModel",
 ]

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/generators/dataset.py

@@ -10,7 +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 distribution
+from phylogenie.utils import Distribution, StrictBaseModel
 
 
 class DataType(str, Enum):
@@ -22,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(
@@ -55,19 +56,21 @@ 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 {}
-                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)
 
-        joblib.Parallel(n_jobs=self.n_jobs)(
+        jobs = joblib.Parallel(n_jobs=self.n_jobs, return_as="generator_unordered")(
             joblib.delayed(self.generate_one)(
                 filename=os.path.join(data_dir, str(i)),
                 data=data[i],
                 seed=int(rng.integers(2**32)),
             )
-            for i in tqdm(range(n_samples), desc=f"Generating {data_dir}...")
+            for i in range(n_samples)
         )
+        for _ in tqdm(jobs, total=n_samples, desc=f"Generating {data_dir}..."):
+            pass
 
     def generate(self) -> None:
         rng = default_rng(self.seed)

phylogenie/generators/factories.py

@@ -14,6 +14,7 @@ from phylogenie.skyline import (
     SkylineVector,
     SkylineVectorCoercible,
 )
+from phylogenie.utils import Distribution
 
 
 def _eval_expression(expression: str, data: dict[str, Any]) -> Any:
@@ -29,6 +30,14 @@ 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)

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,10 +20,12 @@ 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,
     get_canonical_events,
+    get_contact_tracing_events,
     get_epidemiological_events,
     get_FBD_events,
     simulate_tree,
@@ -41,16 +44,13 @@ class ParameterizationType(str, Enum):
 class TreeDatasetGenerator(DatasetGenerator):
     data_type: Literal[DataType.TREES] = DataType.TREES
     min_tips: cfg.Integer = 1
-    max_tips: cfg.Integer | None = None
+    max_tips: cfg.Integer = 2**32
     max_time: cfg.Scalar = np.inf
     init_state: str | None = None
     sampling_probability_at_present: cfg.Scalar = 0.0
-    max_tries: int | None = None
-    notification_probability: cfg.Scalar = 0.0
-    max_notified_contacts: cfg.Integer = 1
-    samplable_states_after_notification: list[str] | None = None
-    sampling_rate_after_notification: cfg.SkylineParameter = np.inf
-    contacts_removal_probability: cfg.SkylineParameter = 1
+
+    @abstractmethod
+    def _get_events(self, data: dict[str, Any]) -> list[Event]: ...
 
     def simulate_one(self, rng: Generator, data: dict[str, Any]) -> Tree | None:
         events = self._get_events(data)
@@ -59,34 +59,18 @@ class TreeDatasetGenerator(DatasetGenerator):
             if self.init_state is None
             else self.init_state.format(**data)
         )
-        max_tips = (
-            self.max_tips if self.max_tips is None else integer(self.max_tips, data)
-        )
         return simulate_tree(
             events=events,
             min_tips=integer(self.min_tips, data),
-            max_tips=max_tips,
+            max_tips=integer(self.max_tips, data),
             max_time=scalar(self.max_time, data),
             init_state=init_state,
             sampling_probability_at_present=scalar(
                 self.sampling_probability_at_present, data
             ),
-            notification_probability=scalar(self.notification_probability, data),
-            max_notified_contacts=integer(self.max_notified_contacts, data),
-            samplable_states_after_notification=self.samplable_states_after_notification,
-            sampling_rate_after_notification=skyline_parameter(
-                self.sampling_rate_after_notification, data
-            ),
-            contacts_removal_probability=skyline_parameter(
-                self.contacts_removal_probability, data
-            ),
-            max_tries=self.max_tries,
             seed=int(rng.integers(2**32)),
         )
 
-    @abstractmethod
-    def _get_events(self, data: dict[str, Any]) -> list[Event]: ...
-
     def _generate_one(
         self, filename: str, rng: Generator, data: dict[str, Any]
     ) -> None:
@@ -101,9 +85,9 @@ class CanonicalTreeDatasetGenerator(TreeDatasetGenerator):
     )
     states: list[str]
     sampling_rates: cfg.SkylineVector
+    remove_after_sampling: bool
     birth_rates: cfg.SkylineVector = 0
     death_rates: cfg.SkylineVector = 0
-    removal_probabilities: cfg.SkylineVector = 0
     migration_rates: cfg.SkylineMatrix = None
     birth_rates_among_states: cfg.SkylineMatrix = None
 
@@ -111,9 +95,9 @@ class CanonicalTreeDatasetGenerator(TreeDatasetGenerator):
         return get_canonical_events(
             states=self.states,
             sampling_rates=skyline_vector(self.sampling_rates, data),
+            remove_after_sampling=self.remove_after_sampling,
             birth_rates=skyline_vector(self.birth_rates, data),
             death_rates=skyline_vector(self.death_rates, data),
-            removal_probabilities=skyline_vector(self.removal_probabilities, data),
             migration_rates=skyline_matrix(self.migration_rates, data),
             birth_rates_among_states=skyline_matrix(
                 self.birth_rates_among_states, data
@@ -121,19 +105,87 @@ class CanonicalTreeDatasetGenerator(TreeDatasetGenerator):
         )
 
 
-class EpidemiologicalTreeDatasetGenerator(TreeDatasetGenerator):
+class FBDTreeDatasetGenerator(TreeDatasetGenerator):
+    parameterization: Literal[ParameterizationType.FBD] = ParameterizationType.FBD
+    states: list[str]
+    sampling_proportions: cfg.SkylineVector
+    diversification: cfg.SkylineVector = 0
+    turnover: cfg.SkylineVector = 0
+    migration_rates: cfg.SkylineMatrix = None
+    diversification_between_states: cfg.SkylineMatrix = None
+
+    def _get_events(self, data: dict[str, Any]) -> list[Event]:
+        return get_FBD_events(
+            states=self.states,
+            diversification=skyline_vector(self.diversification, data),
+            turnover=skyline_vector(self.turnover, data),
+            sampling_proportions=skyline_vector(self.sampling_proportions, data),
+            migration_rates=skyline_matrix(self.migration_rates, data),
+            diversification_between_states=skyline_matrix(
+                self.diversification_between_states, data
+            ),
+        )
+
+
+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]: ...
+
+    def _get_events(self, data: dict[str, Any]) -> list[Event]:
+        events = self._get_base_events(data)
+        if self.notification_probability:
+            events = get_contact_tracing_events(
+                events=events,
+                max_notified_contacts=integer(self.max_notified_contacts, data),
+                notification_probability=skyline_parameter(
+                    self.notification_probability, data
+                ),
+                sampling_rate_after_notification=skyline_parameter(
+                    self.sampling_rate_after_notification, data
+                ),
+                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(TreeDatasetGeneratorForEpidemiology):
     parameterization: Literal[ParameterizationType.EPIDEMIOLOGICAL] = (
         ParameterizationType.EPIDEMIOLOGICAL
     )
     states: list[str]
+    sampling_proportions: cfg.SkylineVector
     reproduction_numbers: cfg.SkylineVector = 0
     become_uninfectious_rates: cfg.SkylineVector = 0
-    sampling_proportions: cfg.SkylineVector = 1
-    removal_probabilities: cfg.SkylineVector = 1
     migration_rates: cfg.SkylineMatrix = None
     reproduction_numbers_among_states: cfg.SkylineMatrix = None
 
-    def _get_events(self, data: dict[str, Any]) -> list[Event]:
+    def _get_base_events(self, data: dict[str, Any]) -> list[Event]:
         return get_epidemiological_events(
             states=self.states,
             reproduction_numbers=skyline_vector(self.reproduction_numbers, data),
@@ -141,7 +193,6 @@ class EpidemiologicalTreeDatasetGenerator(TreeDatasetGenerator):
                 self.become_uninfectious_rates, data
             ),
             sampling_proportions=skyline_vector(self.sampling_proportions, data),
-            removal_probabilities=skyline_vector(self.removal_probabilities, data),
             migration_rates=skyline_matrix(self.migration_rates, data),
             reproduction_numbers_among_states=skyline_matrix(
                 self.reproduction_numbers_among_states, data
@@ -149,37 +200,13 @@ class EpidemiologicalTreeDatasetGenerator(TreeDatasetGenerator):
         )
 
 
-class FBDTreeDatasetGenerator(TreeDatasetGenerator):
-    parameterization: Literal[ParameterizationType.FBD] = ParameterizationType.FBD
-    states: list[str]
-    diversification: cfg.SkylineVector = 0
-    turnover: cfg.SkylineVector = 0
-    sampling_proportions: cfg.SkylineVector = 1
-    removal_probabilities: cfg.SkylineVector = 0
-    migration_rates: cfg.SkylineMatrix = None
-    diversification_between_types: cfg.SkylineMatrix = None
-
-    def _get_events(self, data: dict[str, Any]) -> list[Event]:
-        return get_FBD_events(
-            states=self.states,
-            diversification=skyline_vector(self.diversification, data),
-            turnover=skyline_vector(self.turnover, data),
-            sampling_proportions=skyline_vector(self.sampling_proportions, data),
-            removal_probabilities=skyline_vector(self.removal_probabilities, data),
-            migration_rates=skyline_matrix(self.migration_rates, data),
-            diversification_between_types=skyline_matrix(
-                self.diversification_between_types, data
-            ),
-        )
-
-
-class BDTreeDatasetGenerator(TreeDatasetGenerator):
+class BDTreeDatasetGenerator(TreeDatasetGeneratorForEpidemiology):
     parameterization: Literal[ParameterizationType.BD] = ParameterizationType.BD
     reproduction_number: cfg.SkylineParameter
     infectious_period: cfg.SkylineParameter
     sampling_proportion: cfg.SkylineParameter = 1
 
-    def _get_events(self, data: dict[str, Any]) -> list[Event]:
+    def _get_base_events(self, data: dict[str, Any]) -> list[Event]:
         return get_BD_events(
             reproduction_number=skyline_parameter(self.reproduction_number, data),
             infectious_period=skyline_parameter(self.infectious_period, data),
@@ -187,14 +214,14 @@ class BDTreeDatasetGenerator(TreeDatasetGenerator):
         )
 
 
-class BDEITreeDatasetGenerator(TreeDatasetGenerator):
+class BDEITreeDatasetGenerator(TreeDatasetGeneratorForEpidemiology):
     parameterization: Literal[ParameterizationType.BDEI] = ParameterizationType.BDEI
     reproduction_number: cfg.SkylineParameter
     infectious_period: cfg.SkylineParameter
     incubation_period: cfg.SkylineParameter
     sampling_proportion: cfg.SkylineParameter = 1
 
-    def _get_events(self, data: dict[str, Any]) -> list[Event]:
+    def _get_base_events(self, data: dict[str, Any]) -> list[Event]:
         return get_BDEI_events(
             reproduction_number=skyline_parameter(self.reproduction_number, data),
             infectious_period=skyline_parameter(self.infectious_period, data),
@@ -203,7 +230,7 @@ class BDEITreeDatasetGenerator(TreeDatasetGenerator):
         )
 
 
-class BDSSTreeDatasetGenerator(TreeDatasetGenerator):
+class BDSSTreeDatasetGenerator(TreeDatasetGeneratorForEpidemiology):
     parameterization: Literal[ParameterizationType.BDSS] = ParameterizationType.BDSS
     reproduction_number: cfg.SkylineParameter
     infectious_period: cfg.SkylineParameter
@@ -211,7 +238,7 @@ class BDSSTreeDatasetGenerator(TreeDatasetGenerator):
     superspreaders_proportion: cfg.SkylineParameter
     sampling_proportion: cfg.SkylineParameter = 1
 
-    def _get_events(self, data: dict[str, Any]) -> list[Event]:
+    def _get_base_events(self, data: dict[str, Any]) -> list[Event]:
         return get_BDSS_events(
             reproduction_number=skyline_parameter(self.reproduction_number, data),
             infectious_period=skyline_parameter(self.infectious_period, data),

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,12 +70,19 @@ 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
 
 
-def dump_newick(tree: Tree, filepath: str) -> None:
+def dump_newick(trees: Tree | list[Tree], filepath: str) -> None:
+    if isinstance(trees, Tree):
+        trees = [trees]
     with open(filepath, "w") as file:
-        file.write(_to_newick(tree) + ";")
+        for t in trees:
+            file.write(_to_newick(t) + ";\n")
 
 
 def load_fasta(
@@ -80,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/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/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,
@@ -8,10 +16,18 @@ from phylogenie.treesimulator.events import (
     get_epidemiological_events,
     get_FBD_events,
 )
-from phylogenie.treesimulator.gillespie import simulate_tree
+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",
@@ -19,5 +35,6 @@ __all__ = [
     "get_contact_tracing_events",
     "get_epidemiological_events",
     "get_FBD_events",
+    "generate_trees",
     "simulate_tree",
 ]