phylogenie 2.0.5__tar.gz → 2.0.6__tar.gz

{phylogenie-2.0.5 → phylogenie-2.0.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: phylogenie
-Version: 2.0.5
+Version: 2.0.6
 Summary: Generate phylogenetic datasets with minimal setup effort
 Author: Gabriele Marino
 Author-email: gabmarino.8601@gmail.com
@@ -45,7 +45,7 @@ Phylogenie comes packed with useful features, including:
   Simply specify the number of cores to use, and Phylogenie handles multiprocessing automatically.
 
 - **Pre-implemented parameterizations** 🎯  
-  Include canonical, fossilized birth-death, epidemiological, birth-death with exposed-infectious (BDEI), birth-death with superspreading (BDSS), and more.
+  Include canonical, fossilized birth-death, epidemiological, birth-death with exposed-infectious (BDEI), birth-death with superspreading (BDSS), and contact tracing (CT).
 
 - **Skyline parameter support** 🪜  
   Support for piece-wise constant parameters.

{phylogenie-2.0.5 → phylogenie-2.0.6}/README.md

@@ -27,7 +27,7 @@ Phylogenie comes packed with useful features, including:
   Simply specify the number of cores to use, and Phylogenie handles multiprocessing automatically.
 
 - **Pre-implemented parameterizations** 🎯  
-  Include canonical, fossilized birth-death, epidemiological, birth-death with exposed-infectious (BDEI), birth-death with superspreading (BDSS), and more.
+  Include canonical, fossilized birth-death, epidemiological, birth-death with exposed-infectious (BDEI), birth-death with superspreading (BDSS), and contact tracing (CT).
 
 - **Skyline parameter support** 🪜  
   Support for piece-wise constant parameters.

{phylogenie-2.0.5 → phylogenie-2.0.6}/phylogenie/generators/alisim.py

@@ -46,11 +46,12 @@ class AliSimDatasetGenerator(DatasetGenerator):
         self, filename: str, rng: Generator, data: dict[str, Any]
     ) -> None:
         if self.keep_trees:
-            base_dir = Path(filename).parent
-            file_id = Path(filename).stem
-            tree_filename = os.path.join(base_dir, TREES_DIRNAME, file_id)
+            base_dir, file_id = Path(filename).parent, Path(filename).stem
+            trees_dir = os.path.join(base_dir, TREES_DIRNAME)
             msas_dir = os.path.join(base_dir, MSAS_DIRNAME)
+            os.makedirs(trees_dir, exist_ok=True)
             os.makedirs(msas_dir, exist_ok=True)
+            tree_filename = os.path.join(trees_dir, file_id)
             msa_filename = os.path.join(msas_dir, file_id)
         else:
             tree_filename = f"{filename}.temp-tree"

phylogenie-2.0.6/phylogenie/generators/configs.py

@@ -0,0 +1,39 @@
+from pydantic import BaseModel, ConfigDict
+
+import phylogenie.typings as pgt
+
+
+class Distribution(BaseModel):
+    type: str
+    model_config = ConfigDict(extra="allow")
+
+
+Integer = str | int
+Scalar = str | pgt.Scalar
+ManyScalars = str | list[Scalar]
+OneOrManyScalars = Scalar | list[Scalar]
+OneOrMany2DScalars = Scalar | list[list[Scalar]]
+
+
+class StrictBaseModel(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+
+class SkylineParameterModel(StrictBaseModel):
+    value: ManyScalars
+    change_times: ManyScalars
+
+
+class SkylineVectorModel(StrictBaseModel):
+    value: str | list[OneOrManyScalars]
+    change_times: ManyScalars
+
+
+class SkylineMatrixModel(StrictBaseModel):
+    value: str | list[OneOrMany2DScalars]
+    change_times: ManyScalars
+
+
+SkylineParameter = Scalar | SkylineParameterModel
+SkylineVector = str | pgt.Scalar | list[SkylineParameter] | SkylineVectorModel
+SkylineMatrix = str | pgt.Scalar | list[SkylineVector] | SkylineMatrixModel | None

{phylogenie-2.0.5 → phylogenie-2.0.6}/phylogenie/generators/dataset.py

@@ -10,7 +10,7 @@ import pandas as pd
 from numpy.random import Generator, default_rng
 from tqdm import tqdm
 
-from phylogenie.generators.configs import DistributionConfig, StrictBaseModel
+import phylogenie.generators.configs as cfg
 
 
 class DataType(str, Enum):
@@ -22,12 +22,12 @@ DATA_DIRNAME = "data"
 METADATA_FILENAME = "metadata.csv"
 
 
-class DatasetGenerator(ABC, StrictBaseModel):
+class DatasetGenerator(ABC, cfg.StrictBaseModel):
     output_dir: str = "phylogenie-outputs"
     n_samples: int | dict[str, int] = 1
     n_jobs: int = -1
     seed: int | None = None
-    context: dict[str, DistributionConfig] | None = None
+    context: dict[str, cfg.Distribution] | None = None
 
     @abstractmethod
     def _generate_one(

{phylogenie-2.0.5 → phylogenie-2.0.6}/phylogenie/generators/factories.py

@@ -29,7 +29,7 @@ def _eval_expression(expression: str, data: dict[str, Any]) -> Any:
     ).tolist()
 
 
-def integer(x: cfg.IntegerConfig, data: dict[str, Any]) -> int:
+def integer(x: cfg.Integer, data: dict[str, Any]) -> int:
     if isinstance(x, str):
         e = _eval_expression(x, data)
         if isinstance(e, int):
@@ -40,7 +40,7 @@ def integer(x: cfg.IntegerConfig, data: dict[str, Any]) -> int:
     return x
 
 
-def scalar(x: cfg.ScalarConfig, data: dict[str, Any]) -> pgt.Scalar:
+def scalar(x: cfg.Scalar, data: dict[str, Any]) -> pgt.Scalar:
     if isinstance(x, str):
         e = _eval_expression(x, data)
         if isinstance(e, pgt.Scalar):
@@ -51,7 +51,7 @@ def scalar(x: cfg.ScalarConfig, data: dict[str, Any]) -> pgt.Scalar:
     return x
 
 
-def many_scalars(x: cfg.ManyScalarsConfig, data: dict[str, Any]) -> pgt.ManyScalars:
+def many_scalars(x: cfg.ManyScalars, data: dict[str, Any]) -> pgt.ManyScalars:
     if isinstance(x, str):
         e = _eval_expression(x, data)
         if tg.is_many_scalars(e):
@@ -63,7 +63,7 @@ def many_scalars(x: cfg.ManyScalarsConfig, data: dict[str, Any]) -> pgt.ManyScal
 
 
 def one_or_many_scalars(
-    x: cfg.OneOrManyScalarsConfig, data: dict[str, Any]
+    x: cfg.OneOrManyScalars, data: dict[str, Any]
 ) -> pgt.OneOrManyScalars:
     if isinstance(x, str):
         e = _eval_expression(x, data)
@@ -78,9 +78,9 @@ def one_or_many_scalars(
 
 
 def skyline_parameter(
-    x: cfg.SkylineParameterConfig, data: dict[str, Any]
+    x: cfg.SkylineParameter, data: dict[str, Any]
 ) -> SkylineParameterLike:
-    if isinstance(x, cfg.ScalarConfig):
+    if isinstance(x, cfg.Scalar):
         return scalar(x, data)
     return SkylineParameter(
         value=many_scalars(x.value, data),
@@ -89,7 +89,7 @@ def skyline_parameter(
 
 
 def skyline_vector(
-    x: cfg.SkylineVectorConfig, data: dict[str, Any]
+    x: cfg.SkylineVector, data: dict[str, Any]
 ) -> SkylineVectorCoercible:
     if isinstance(x, str):
         e = _eval_expression(x, data)
@@ -132,7 +132,7 @@ def skyline_vector(
 
 
 def one_or_many_2D_scalars(
-    x: cfg.OneOrMany2DScalarsConfig, data: dict[str, Any]
+    x: cfg.OneOrMany2DScalars, data: dict[str, Any]
 ) -> pgt.OneOrMany2DScalars:
     if isinstance(x, str):
         e = _eval_expression(x, data)
@@ -147,7 +147,7 @@ def one_or_many_2D_scalars(
 
 
 def skyline_matrix(
-    x: cfg.SkylineMatrixConfig, data: dict[str, Any]
+    x: cfg.SkylineMatrix, data: dict[str, Any]
 ) -> SkylineMatrixCoercible | None:
     if x is None:
         return None

{phylogenie-2.0.5 → phylogenie-2.0.6}/phylogenie/generators/trees.py

@@ -40,15 +40,20 @@ class ParameterizationType(str, Enum):
 
 class TreeDatasetGenerator(DatasetGenerator):
     data_type: Literal[DataType.TREES] = DataType.TREES
-    min_tips: cfg.IntegerConfig = 1
-    max_tips: cfg.IntegerConfig | None = None
-    max_time: cfg.ScalarConfig = np.inf
+    min_tips: cfg.Integer = 1
+    max_tips: cfg.Integer | None = None
+    max_time: cfg.Scalar = np.inf
     init_state: str | None = None
-    sampling_probability_at_present: cfg.ScalarConfig = 0.0
+    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
 
     def simulate_one(self, rng: Generator, data: dict[str, Any]) -> Tree | None:
-        events = self._get_events(rng, data)
+        events = self._get_events(data)
         init_state = (
             self.init_state
             if self.init_state is None
@@ -66,12 +71,21 @@ class TreeDatasetGenerator(DatasetGenerator):
             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, rng: Generator, data: dict[str, Any]) -> list[Event]: ...
+    def _get_events(self, data: dict[str, Any]) -> list[Event]: ...
 
     def _generate_one(
         self, filename: str, rng: Generator, data: dict[str, Any]
@@ -85,15 +99,15 @@ class CanonicalTreeDatasetGenerator(TreeDatasetGenerator):
     parameterization: Literal[ParameterizationType.CANONICAL] = (
         ParameterizationType.CANONICAL
     )
-    sampling_rates: cfg.SkylineVectorConfig
-    birth_rates: cfg.SkylineVectorConfig = 0
-    death_rates: cfg.SkylineVectorConfig = 0
-    removal_probabilities: cfg.SkylineVectorConfig = 0
-    migration_rates: cfg.SkylineMatrixConfig = None
-    birth_rates_among_states: cfg.SkylineMatrixConfig = None
-    states: list[str] | None = None
-
-    def _get_events(self, rng: Generator, data: dict[str, Any]) -> list[Event]:
+    states: list[str]
+    sampling_rates: cfg.SkylineVector
+    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
+
+    def _get_events(self, data: dict[str, Any]) -> list[Event]:
         return get_canonical_events(
             states=self.states,
             sampling_rates=skyline_vector(self.sampling_rates, data),
@@ -111,15 +125,15 @@ class EpidemiologicalTreeDatasetGenerator(TreeDatasetGenerator):
     parameterization: Literal[ParameterizationType.EPIDEMIOLOGICAL] = (
         ParameterizationType.EPIDEMIOLOGICAL
     )
-    states: list[str] | None = None
-    reproduction_numbers: cfg.SkylineVectorConfig = 0
-    become_uninfectious_rates: cfg.SkylineVectorConfig = 0
-    sampling_proportions: cfg.SkylineVectorConfig = 1
-    removal_probabilities: cfg.SkylineVectorConfig = 1
-    migration_rates: cfg.SkylineMatrixConfig = None
-    reproduction_numbers_among_states: cfg.SkylineMatrixConfig = None
-
-    def _get_events(self, rng: Generator, data: dict[str, Any]) -> list[Event]:
+    states: list[str]
+    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]:
         return get_epidemiological_events(
             states=self.states,
             reproduction_numbers=skyline_vector(self.reproduction_numbers, data),
@@ -137,15 +151,15 @@ class EpidemiologicalTreeDatasetGenerator(TreeDatasetGenerator):
 
 class FBDTreeDatasetGenerator(TreeDatasetGenerator):
     parameterization: Literal[ParameterizationType.FBD] = ParameterizationType.FBD
-    states: list[str] | None = None
-    diversification: cfg.SkylineVectorConfig = 0
-    turnover: cfg.SkylineVectorConfig = 0
-    sampling_proportions: cfg.SkylineVectorConfig = 1
-    removal_probabilities: cfg.SkylineVectorConfig = 0
-    migration_rates: cfg.SkylineMatrixConfig = None
-    diversification_between_types: cfg.SkylineMatrixConfig = None
-
-    def _get_events(self, rng: Generator, data: dict[str, Any]) -> list[Event]:
+    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),
@@ -161,11 +175,11 @@ class FBDTreeDatasetGenerator(TreeDatasetGenerator):
 
 class BDTreeDatasetGenerator(TreeDatasetGenerator):
     parameterization: Literal[ParameterizationType.BD] = ParameterizationType.BD
-    reproduction_number: cfg.SkylineParameterConfig
-    infectious_period: cfg.SkylineParameterConfig
-    sampling_proportion: cfg.SkylineParameterConfig = 1
+    reproduction_number: cfg.SkylineParameter
+    infectious_period: cfg.SkylineParameter
+    sampling_proportion: cfg.SkylineParameter = 1
 
-    def _get_events(self, rng: Generator, data: dict[str, Any]) -> list[Event]:
+    def _get_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),
@@ -175,12 +189,12 @@ class BDTreeDatasetGenerator(TreeDatasetGenerator):
 
 class BDEITreeDatasetGenerator(TreeDatasetGenerator):
     parameterization: Literal[ParameterizationType.BDEI] = ParameterizationType.BDEI
-    reproduction_number: cfg.SkylineParameterConfig
-    infectious_period: cfg.SkylineParameterConfig
-    incubation_period: cfg.SkylineParameterConfig
-    sampling_proportion: cfg.SkylineParameterConfig = 1
+    reproduction_number: cfg.SkylineParameter
+    infectious_period: cfg.SkylineParameter
+    incubation_period: cfg.SkylineParameter
+    sampling_proportion: cfg.SkylineParameter = 1
 
-    def _get_events(self, rng: Generator, data: dict[str, Any]) -> list[Event]:
+    def _get_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),
@@ -191,13 +205,13 @@ class BDEITreeDatasetGenerator(TreeDatasetGenerator):
 
 class BDSSTreeDatasetGenerator(TreeDatasetGenerator):
     parameterization: Literal[ParameterizationType.BDSS] = ParameterizationType.BDSS
-    reproduction_number: cfg.SkylineParameterConfig
-    infectious_period: cfg.SkylineParameterConfig
-    superspreading_ratio: cfg.SkylineParameterConfig
-    superspreaders_proportion: cfg.SkylineParameterConfig
-    sampling_proportion: cfg.SkylineParameterConfig = 1
+    reproduction_number: cfg.SkylineParameter
+    infectious_period: cfg.SkylineParameter
+    superspreading_ratio: cfg.SkylineParameter
+    superspreaders_proportion: cfg.SkylineParameter
+    sampling_proportion: cfg.SkylineParameter = 1
 
-    def _get_events(self, rng: Generator, data: dict[str, Any]) -> list[Event]:
+    def _get_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-2.0.6/phylogenie/generators/typeguards.py

@@ -0,0 +1,28 @@
+from typing import Any, TypeGuard
+
+import phylogenie.generators.configs as cfg
+import phylogenie.typings as pgt
+
+
+def is_list(x: Any) -> TypeGuard[list[Any]]:
+    return isinstance(x, list)
+
+
+def is_list_of_scalar_configs(x: Any) -> TypeGuard[list[cfg.Scalar]]:
+    return is_list(x) and all(isinstance(v, cfg.Scalar) for v in x)
+
+
+def is_list_of_skyline_parameter_configs(
+    x: Any,
+) -> TypeGuard[list[cfg.SkylineParameter]]:
+    return is_list(x) and all(isinstance(v, cfg.SkylineParameter) for v in x)
+
+
+def is_skyline_vector_config(x: Any) -> TypeGuard[cfg.SkylineVector]:
+    return isinstance(
+        x, str | pgt.Scalar | cfg.SkylineVectorModel
+    ) or is_list_of_skyline_parameter_configs(x)
+
+
+def is_list_of_skyline_vector_configs(x: Any) -> TypeGuard[list[cfg.SkylineVector]]:
+    return is_list(x) and all(is_skyline_vector_config(v) for v in x)

{phylogenie-2.0.5 → phylogenie-2.0.6}/phylogenie/main.py

@@ -37,3 +37,7 @@ def main() -> None:
     args = parser.parse_args()
 
     run(args.config_path)
+
+
+if __name__ == "__main__":
+    run("/Users/gmarino/phylogenie/examples/5.BD-CT-Trees.yaml")

{phylogenie-2.0.5 → phylogenie-2.0.6}/phylogenie/skyline/matrix.py

@@ -1,5 +1,5 @@
 from collections.abc import Callable, Iterator
-from typing import TypeGuard, Union, overload
+from typing import Any, TypeGuard, Union, overload
 
 import phylogenie.typeguards as tg
 import phylogenie.typings as pgt
@@ -20,7 +20,7 @@ SkylineMatrixOperand = Union[SkylineVectorOperand, "SkylineMatrix"]
 SkylineMatrixCoercible = Union[pgt.OneOrMany[SkylineVectorCoercible], "SkylineMatrix"]
 
 
-def is_skyline_matrix_operand(x: object) -> TypeGuard[SkylineMatrixOperand]:
+def is_skyline_matrix_operand(x: Any) -> TypeGuard[SkylineMatrixOperand]:
     return isinstance(x, SkylineMatrix) or is_skyline_vector_operand(x)
 
 
@@ -142,7 +142,7 @@ class SkylineMatrix:
     def __bool__(self) -> bool:
         return any(self.params)
 
-    def __eq__(self, other: object) -> bool:
+    def __eq__(self, other: Any) -> bool:
         return isinstance(other, SkylineMatrix) and self.params == other.params
 
     def __repr__(self) -> str:

{phylogenie-2.0.5 → phylogenie-2.0.6}/phylogenie/skyline/parameter.py

@@ -1,6 +1,6 @@
 from bisect import bisect_right
 from collections.abc import Callable
-from typing import TypeGuard, Union
+from typing import Any, TypeGuard, Union
 
 import phylogenie.typeguards as tg
 import phylogenie.typings as pgt
@@ -8,12 +8,12 @@ import phylogenie.typings as pgt
 SkylineParameterLike = Union[pgt.Scalar, "SkylineParameter"]
 
 
-def is_skyline_parameter_like(x: object) -> TypeGuard[SkylineParameterLike]:
+def is_skyline_parameter_like(x: Any) -> TypeGuard[SkylineParameterLike]:
     return isinstance(x, pgt.Scalar | SkylineParameter)
 
 
 def is_many_skyline_parameters_like(
-    x: object,
+    x: Any,
 ) -> TypeGuard[pgt.Many[SkylineParameterLike]]:
     return tg.is_many(x) and all(is_skyline_parameter_like(v) for v in x)
 
@@ -107,7 +107,7 @@ class SkylineParameter:
     def __bool__(self) -> bool:
         return any(self.value)
 
-    def __eq__(self, other: object) -> bool:
+    def __eq__(self, other: Any) -> bool:
         return isinstance(other, SkylineParameter) and (
             self.value == other.value and self.change_times == other.change_times
         )

{phylogenie-2.0.5 → phylogenie-2.0.6}/phylogenie/skyline/vector.py

@@ -1,5 +1,5 @@
 from collections.abc import Callable, Iterator
-from typing import TypeGuard, Union, overload
+from typing import Any, TypeGuard, Union, overload
 
 import phylogenie.typeguards as tg
 import phylogenie.typings as pgt
@@ -16,26 +16,24 @@ SkylineVectorLike = Union[pgt.Many[SkylineParameterLike], "SkylineVector"]
 SkylineVectorCoercible = Union[pgt.OneOrMany[SkylineParameterLike], "SkylineVector"]
 
 
-def is_skyline_vector_operand(x: object) -> TypeGuard[SkylineVectorOperand]:
+def is_skyline_vector_operand(x: Any) -> TypeGuard[SkylineVectorOperand]:
     return isinstance(x, SkylineVector) or is_skyline_parameter_like(x)
 
 
-def is_skyline_vector_like(x: object) -> TypeGuard[SkylineVectorLike]:
+def is_skyline_vector_like(x: Any) -> TypeGuard[SkylineVectorLike]:
     return isinstance(x, SkylineVector) or is_many_skyline_parameters_like(x)
 
 
-def is_skyline_vector_coercible(
-    x: object,
-) -> TypeGuard[SkylineVectorCoercible]:
+def is_skyline_vector_coercible(x: Any) -> TypeGuard[SkylineVectorCoercible]:
     return is_skyline_parameter_like(x) or is_skyline_vector_like(x)
 
 
-def is_many_skyline_vectors_like(x: object) -> TypeGuard[pgt.Many[SkylineVectorLike]]:
+def is_many_skyline_vectors_like(x: Any) -> TypeGuard[pgt.Many[SkylineVectorLike]]:
     return tg.is_many(x) and all(is_skyline_vector_like(v) for v in x)
 
 
 def is_many_skyline_vectors_coercible(
-    x: object,
+    x: Any,
 ) -> TypeGuard[pgt.Many[SkylineVectorCoercible]]:
     return tg.is_many(x) and all(is_skyline_vector_coercible(v) for v in x)
 
@@ -131,7 +129,7 @@ class SkylineVector:
     def __bool__(self) -> bool:
         return any(self.params)
 
-    def __eq__(self, other: object) -> bool:
+    def __eq__(self, other: Any) -> bool:
         return isinstance(other, SkylineVector) and self.params == other.params
 
     def __repr__(self) -> str:

{phylogenie-2.0.5 → phylogenie-2.0.6}/phylogenie/tree.py

@@ -1,6 +1,5 @@
-from __future__ import annotations
-
-from typing import Iterator
+from collections.abc import Iterator
+from typing import Any
 
 
 class Tree:
@@ -9,29 +8,30 @@ class Tree:
         self.branch_length = branch_length
         self.parent: Tree | None = None
         self.children: list[Tree] = []
+        self._features: dict[str, Any] = {}
 
-    def add_child(self, child: Tree) -> Tree:
+    def add_child(self, child: "Tree") -> "Tree":
         child.parent = self
         self.children.append(child)
         return self
 
-    def preorder_traversal(self) -> Iterator[Tree]:
+    def preorder_traversal(self) -> Iterator["Tree"]:
         yield self
         for child in self.children:
             yield from child.preorder_traversal()
 
-    def postorder_traversal(self) -> Iterator[Tree]:
+    def postorder_traversal(self) -> Iterator["Tree"]:
         for child in self.children:
             yield from child.postorder_traversal()
         yield self
 
-    def get_node(self, id: str) -> Tree:
+    def get_node(self, id: str) -> "Tree":
         for node in self:
             if node.id == id:
                 return node
         raise ValueError(f"Node with id {id} not found.")
 
-    def get_leaves(self) -> list[Tree]:
+    def get_leaves(self) -> list["Tree"]:
         return [node for node in self if not node.children]
 
     def get_time(self) -> float:
@@ -40,13 +40,22 @@ class Tree:
             raise ValueError(f"Branch length of node {self.id} is not set.")
         return self.branch_length + parent_time
 
-    def copy(self) -> Tree:
+    def is_leaf(self) -> bool:
+        return not self.children
+
+    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 __iter__(self) -> Iterator[Tree]:
+    def get(self, key: str, default: Any = None) -> Any:
+        return self._features.get(key, default)
+
+    def set(self, key: str, value: Any) -> None:
+        self._features[key] = value
+
+    def __iter__(self) -> Iterator["Tree"]:
         return self.preorder_traversal()
 
     def __repr__(self) -> str:

{phylogenie-2.0.5 → phylogenie-2.0.6}/phylogenie/treesimulator/__init__.py

@@ -4,6 +4,7 @@ from phylogenie.treesimulator.events import (
     get_BDEI_events,
     get_BDSS_events,
     get_canonical_events,
+    get_contact_tracing_events,
     get_epidemiological_events,
     get_FBD_events,
 )
@@ -15,6 +16,7 @@ __all__ = [
     "get_BDEI_events",
     "get_BDSS_events",
     "get_canonical_events",
+    "get_contact_tracing_events",
     "get_epidemiological_events",
     "get_FBD_events",
     "simulate_tree",