phylogenie 2.0.5__py3-none-any.whl → 2.0.7__py3-none-any.whl

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

@@ -3,16 +3,16 @@ from pydantic import BaseModel, ConfigDict
 import phylogenie.typings as pgt
 
 
-class DistributionConfig(BaseModel):
+class Distribution(BaseModel):
     type: str
     model_config = ConfigDict(extra="allow")
 
 
-IntegerConfig = str | int
-ScalarConfig = str | pgt.Scalar
-ManyScalarsConfig = str | list[ScalarConfig]
-OneOrManyScalarsConfig = ScalarConfig | list[ScalarConfig]
-OneOrMany2DScalarsConfig = ScalarConfig | list[list[ScalarConfig]]
+Integer = str | int
+Scalar = str | pgt.Scalar
+ManyScalars = str | list[Scalar]
+OneOrManyScalars = Scalar | list[Scalar]
+OneOrMany2DScalars = Scalar | list[list[Scalar]]
 
 
 class StrictBaseModel(BaseModel):
@@ -20,24 +20,20 @@ class StrictBaseModel(BaseModel):
 
 
 class SkylineParameterModel(StrictBaseModel):
-    value: ManyScalarsConfig
-    change_times: ManyScalarsConfig
+    value: ManyScalars
+    change_times: ManyScalars
 
 
 class SkylineVectorModel(StrictBaseModel):
-    value: str | list[OneOrManyScalarsConfig]
-    change_times: ManyScalarsConfig
+    value: str | list[OneOrManyScalars]
+    change_times: ManyScalars
 
 
 class SkylineMatrixModel(StrictBaseModel):
-    value: str | list[OneOrMany2DScalarsConfig]
-    change_times: ManyScalarsConfig
+    value: str | list[OneOrMany2DScalars]
+    change_times: ManyScalars
 
 
-SkylineParameterConfig = ScalarConfig | SkylineParameterModel
-SkylineVectorConfig = (
-    str | pgt.Scalar | list[SkylineParameterConfig] | SkylineVectorModel
-)
-SkylineMatrixConfig = (
-    str | pgt.Scalar | list[SkylineVectorConfig] | SkylineMatrixModel | None
-)
+SkylineParameter = Scalar | SkylineParameterModel
+SkylineVector = str | pgt.Scalar | list[SkylineParameter] | SkylineVectorModel
+SkylineMatrix = str | pgt.Scalar | list[SkylineVector] | SkylineMatrixModel | None

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

@@ -1,32 +1,28 @@
-from typing import TypeGuard
+from typing import Any, TypeGuard
 
 import phylogenie.generators.configs as cfg
 import phylogenie.typings as pgt
 
 
-def is_list(x: object) -> TypeGuard[list[object]]:
+def is_list(x: Any) -> TypeGuard[list[Any]]:
     return isinstance(x, list)
 
 
-def is_list_of_scalar_configs(x: object) -> TypeGuard[list[cfg.ScalarConfig]]:
-    return is_list(x) and all(isinstance(v, cfg.ScalarConfig) for v in x)
+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: object,
-) -> TypeGuard[list[cfg.SkylineParameterConfig]]:
-    return is_list(x) and all(isinstance(v, cfg.SkylineParameterConfig) for v in x)
+    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: object,
-) -> TypeGuard[cfg.SkylineVectorConfig]:
+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: object,
-) -> TypeGuard[list[cfg.SkylineVectorConfig]]:
+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/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/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/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/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/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",

phylogenie/treesimulator/events.py

@@ -1,95 +1,84 @@
 from abc import ABC, abstractmethod
+from collections.abc import Sequence
+from dataclasses import dataclass
 
-from numpy.random import Generator
+import numpy as np
 
 from phylogenie.skyline import (
     SkylineMatrixCoercible,
+    SkylineParameter,
     SkylineParameterLike,
     SkylineVectorCoercible,
     skyline_matrix,
     skyline_parameter,
     skyline_vector,
 )
-from phylogenie.treesimulator.model import Model
+from phylogenie.treesimulator.model import Model, get_CT_state
 
 INFECTIOUS_STATE = "I"
 EXPOSED_STATE = "E"
 SUPERSPREADER_STATE = "S"
 
 
+@dataclass
 class Event(ABC):
-    def __init__(
-        self,
-        rate: SkylineParameterLike,
-        state: str | None = None,
-    ):
-        self.rate = skyline_parameter(rate)
-        self.state = state
+    rate: SkylineParameter
+    state: str
 
     def get_propensity(self, model: Model, time: float) -> float:
-        return self.rate.get_value_at_time(time) * model.count_leaves(self.state)
+        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
-    def apply(self, rng: Generator, model: Model, time: float) -> None: ...
+    def apply(self, model: Model, time: float) -> None: ...
 
 
+@dataclass
 class BirthEvent(Event):
-    def __init__(
-        self,
-        rate: SkylineParameterLike,
-        state: str | None = None,
-        child_state: str | None = None,
-    ):
-        super().__init__(rate, state)
-        self.child_state = child_state
+    child_state: str
 
-    def apply(self, rng: Generator, model: Model, time: float) -> None:
-        node = model.get_random_leaf(self.state, rng)
-        model.add_child(node, time, True, self.child_state)
+    def apply(self, model: Model, time: float) -> None:
+        individual = model.get_random_individual(self.state)
+        model.birth_from(individual, self.child_state, time)
 
 
 class DeathEvent(Event):
-    def apply(self, rng: Generator, model: Model, time: float) -> None:
-        node = model.get_random_leaf(self.state, rng)
-        model.remove(node)
+    def apply(self, model: Model, time: float) -> None:
+        individual = model.get_random_individual(self.state)
+        model.remove(individual, time)
 
 
+@dataclass
 class MigrationEvent(Event):
-    def __init__(self, rate: SkylineParameterLike, state: str, target_state: str):
-        super().__init__(rate, state)
-        self.target_state = target_state
+    target_state: str
 
-    def apply(self, rng: Generator, model: Model, time: float) -> None:
-        node = model.get_random_leaf(self.state, rng)
-        model.add_child(node, time, False, self.target_state)
+    def apply(self, model: Model, time: float) -> None:
+        individual = model.get_random_individual(self.state)
+        model.migrate(individual, self.target_state, time)
 
 
+@dataclass
 class SamplingEvent(Event):
-    def __init__(
-        self,
-        rate: SkylineParameterLike,
-        state: str | None = None,
-        removal_probability: SkylineParameterLike = 0,
-    ):
-        super().__init__(rate, state)
-        self.removal_probability = skyline_parameter(removal_probability)
+    removal_probability: SkylineParameter
 
-    def apply(self, rng: Generator, model: Model, time: float) -> None:
-        node = model.get_random_leaf(self.state, rng)
-        remove = rng.random() < self.removal_probability.get_value_at_time(time)
-        model.sample(node, time, remove)
+    def apply(self, model: Model, time: float) -> None:
+        individual = model.get_random_individual(self.state)
+        model.sample(individual, time, self.removal_probability.get_value_at_time(time))
 
 
 def get_canonical_events(
+    states: Sequence[str],
     sampling_rates: SkylineVectorCoercible,
     birth_rates: SkylineVectorCoercible = 0,
     death_rates: SkylineVectorCoercible = 0,
     removal_probabilities: SkylineVectorCoercible = 0,
     migration_rates: SkylineMatrixCoercible | None = None,
     birth_rates_among_states: SkylineMatrixCoercible | None = None,
-    states: list[str] | None = None,
 ) -> list[Event]:
-    N = 1 if states is None else len(states)
+    N = len(states)
 
     birth_rates = skyline_vector(birth_rates, N)
     death_rates = skyline_vector(death_rates, N)
@@ -97,43 +86,38 @@ def get_canonical_events(
     removal_probabilities = skyline_vector(removal_probabilities, N)
 
     events: list[Event] = []
-    for i in range(N):
-        state = None if states is None else states[i]
+    for i, state in enumerate(states):
         events.append(BirthEvent(birth_rates[i], state, state))
         events.append(DeathEvent(death_rates[i], state))
         events.append(SamplingEvent(sampling_rates[i], state, removal_probabilities[i]))
 
-    if states is not None and migration_rates is not None:
+    if migration_rates is not None:
         migration_rates = skyline_matrix(migration_rates, N, N - 1)
         for i, state in enumerate(states):
             for j, other_state in enumerate([s for s in states if s != state]):
                 events.append(MigrationEvent(migration_rates[i, j], state, other_state))
-    elif migration_rates is not None:
-        raise ValueError(f"Migration rates require states to be provided.")
 
-    if states is not None and birth_rates_among_states is not None:
+    if birth_rates_among_states is not None:
         birth_rates_among_states = skyline_matrix(birth_rates_among_states, N, N - 1)
         for i, state in enumerate(states):
             for j, other_state in enumerate([s for s in states if s != state]):
                 events.append(
                     BirthEvent(birth_rates_among_states[i, j], state, other_state)
                 )
-    elif birth_rates_among_states is not None:
-        raise ValueError(f"Birth rates among states require states to be provided.")
 
     return [event for event in events if event.rate]
 
 
 def get_epidemiological_events(
+    states: Sequence[str],
     sampling_proportions: SkylineVectorCoercible = 1,
     reproduction_numbers: SkylineVectorCoercible = 0,
     become_uninfectious_rates: SkylineVectorCoercible = 0,
     removal_probabilities: SkylineVectorCoercible = 1,
     migration_rates: SkylineMatrixCoercible | None = None,
     reproduction_numbers_among_states: SkylineMatrixCoercible | None = None,
-    states: list[str] | None = None,
 ) -> list[Event]:
-    N = 1 if states is None else len(states)
+    N = len(states)
 
     reproduction_numbers = skyline_vector(reproduction_numbers, N)
     become_uninfectious_rates = skyline_vector(become_uninfectious_rates, N)
@@ -143,16 +127,14 @@ def get_epidemiological_events(
     birth_rates = reproduction_numbers * become_uninfectious_rates
     sampling_rates = become_uninfectious_rates * sampling_proportions
     death_rates = become_uninfectious_rates - removal_probabilities * sampling_rates
-    birth_rates_among_states = None
-    if states is None and reproduction_numbers_among_states is not None:
-        raise ValueError(
-            f"Reproduction numbers among states require states to be provided."
-        )
-    elif reproduction_numbers_among_states is not None:
-        birth_rates_among_states = (
+    birth_rates_among_states = (
+        (
             skyline_matrix(reproduction_numbers_among_states, N, N - 1)
             * become_uninfectious_rates
         )
+        if reproduction_numbers_among_states is not None
+        else None
+    )
 
     return get_canonical_events(
         states=states,
@@ -166,15 +148,15 @@ def get_epidemiological_events(
 
 
 def get_FBD_events(
+    states: Sequence[str],
     diversification: SkylineVectorCoercible = 0,
     turnover: SkylineVectorCoercible = 0,
     sampling_proportions: SkylineVectorCoercible = 1,
     removal_probabilities: SkylineVectorCoercible = 0,
     migration_rates: SkylineMatrixCoercible | None = None,
     diversification_between_types: SkylineMatrixCoercible | None = None,
-    states: list[str] | None = None,
 ) -> list[Event]:
-    N = 1 if states is None else len(states)
+    N = len(states)
 
     diversification = skyline_vector(diversification, N)
     turnover = skyline_vector(turnover, N)
@@ -185,15 +167,11 @@ def get_FBD_events(
     death_rates = turnover * birth_rates
     sampling_rates_dividend = 1 - removal_probabilities * sampling_proportions
     sampling_rates = sampling_proportions * death_rates / sampling_rates_dividend
-    birth_rates_among_states = None
-    if states is None and diversification_between_types is not None:
-        raise ValueError(
-            f"Diversification rates among states require states to be provided."
-        )
-    elif diversification_between_types is not None:
-        birth_rates_among_states = (
-            skyline_matrix(diversification_between_types, N, N - 1) + death_rates
-        )
+    birth_rates_among_states = (
+        (skyline_matrix(diversification_between_types, N, N - 1) + death_rates)
+        if diversification_between_types is not None
+        else None
+    )
 
     return get_canonical_events(
         states=states,
@@ -212,6 +190,7 @@ def get_BD_events(
     sampling_proportion: SkylineParameterLike = 1,
 ) -> list[Event]:
     return get_epidemiological_events(
+        states=[INFECTIOUS_STATE],
         reproduction_numbers=reproduction_number,
         become_uninfectious_rates=1 / infectious_period,
         sampling_proportions=sampling_proportion,
@@ -253,3 +232,40 @@ def get_BDSS_events(
         become_uninfectious_rates=1 / infectious_period,
         sampling_proportions=sampling_proportion,
     )
+
+
+def get_contact_tracing_events(
+    events: Sequence[Event],
+    samplable_states_after_notification: Sequence[str] | None = None,
+    sampling_rate_after_notification: SkylineParameterLike = np.inf,
+    contacts_removal_probability: SkylineParameterLike = 1,
+) -> list[Event]:
+    ct_events = list(events)
+    for event in events:
+        if isinstance(event, MigrationEvent):
+            ct_events.append(
+                MigrationEvent(
+                    event.rate,
+                    get_CT_state(event.state),
+                    get_CT_state(event.target_state),
+                )
+            )
+        elif isinstance(event, BirthEvent):
+            ct_events.append(
+                BirthEvent(event.rate, get_CT_state(event.state), event.child_state)
+            )
+
+    for state in (
+        samplable_states_after_notification
+        if samplable_states_after_notification is not None
+        else [e.state for e in events]
+    ):
+        ct_events.append(
+            SamplingEvent(
+                skyline_parameter(sampling_rate_after_notification),
+                get_CT_state(state),
+                skyline_parameter(contacts_removal_probability),
+            )
+        )
+
+    return ct_events

phylogenie/treesimulator/gillespie.py

@@ -3,9 +3,10 @@ from collections.abc import Sequence
 import numpy as np
 from numpy.random import default_rng
 
+from phylogenie.skyline import SkylineParameterLike
 from phylogenie.tree import Tree
-from phylogenie.treesimulator.events import Event
-from phylogenie.treesimulator.model import Model
+from phylogenie.treesimulator.events import Event, get_contact_tracing_events
+from phylogenie.treesimulator.model import Model, is_CT_state
 
 
 def simulate_tree(
@@ -15,6 +16,11 @@ def simulate_tree(
     max_time: float = np.inf,
     init_state: str | None = None,
     sampling_probability_at_present: float = 0.0,
+    notification_probability: float = 0,
+    max_notified_contacts: int = 1,
+    samplable_states_after_notification: Sequence[str] | None = None,
+    sampling_rate_after_notification: SkylineParameterLike = np.inf,
+    contacts_removal_probability: SkylineParameterLike = 1,
     max_tries: int | None = None,
     seed: int | None = None,
 ) -> Tree | None:
@@ -23,15 +29,19 @@ def simulate_tree(
     if max_tips is None and max_time == np.inf:
         raise ValueError("Either max_tips or max_time must be specified.")
 
+    if notification_probability:
+        events = get_contact_tracing_events(
+            events,
+            samplable_states_after_notification,
+            sampling_rate_after_notification,
+            contacts_removal_probability,
+        )
+
     n_tries = 0
-    states = [e.state for e in events if e.state is not None]
-    init_state = (
-        init_state
-        if init_state is not None
-        else str(rng.choice(states)) if states else None
-    )
+    root_states = [e.state for e in events if not is_CT_state(e.state)]
     while max_tries is None or n_tries < max_tries:
-        model = Model(init_state)
+        root_state = init_state if init_state is not None else rng.choice(root_states)
+        model = Model(root_state, max_notified_contacts, notification_probability, rng)
         current_time = 0.0
         change_times = sorted(set(t for e in events for t in e.rate.change_times))
         next_change_time = change_times.pop(0) if change_times else np.inf
@@ -39,6 +49,13 @@ def simulate_tree(
 
         while current_time < max_time and (n_tips is None or model.n_sampled < n_tips):
             rates = [e.get_propensity(model, current_time) for e in events]
+
+            instantaneous_events = [e for e, r in zip(events, rates) if r == np.inf]
+            if instantaneous_events:
+                event = instantaneous_events[rng.integers(len(instantaneous_events))]
+                event.apply(model, current_time)
+                continue
+
             if not any(rates):
                 break
 
@@ -52,11 +69,11 @@ def simulate_tree(
                 continue
 
             event_idx = np.searchsorted(np.cumsum(rates) / sum(rates), rng.random())
-            events[int(event_idx)].apply(rng, model, current_time)
+            events[int(event_idx)].apply(model, current_time)
 
-        for leaf in model.get_leaves():
+        for individual in model.get_population():
             if rng.random() < sampling_probability_at_present:
-                model.sample(leaf, current_time, True)
+                model.sample(individual, current_time, 1)
 
         if model.n_sampled >= min_tips and (
             max_tips is None or model.n_sampled <= max_tips

phylogenie/treesimulator/model.py

@@ -1,71 +1,125 @@
 from collections import defaultdict
+from dataclasses import dataclass, field
+from typing import ClassVar
 
 from numpy.random import Generator, default_rng
 
 from phylogenie.tree import Tree
 
+CT_POSTFIX = "-CT"
 
-class Model:
-    def __init__(self, init_state: str | None = None):
-        self._next_id = 0
-        self._n_sampled = 0
-        self._leaves: dict[str, Tree] = {}
-        self._leaf2state: dict[str, str | None] = {}
-        self._state2leaves: dict[str | None, set[str]] = defaultdict(set)
-        self._tree = self._get_new_node(init_state, None)
 
-    @property
-    def next_id(self) -> int:
-        self._next_id += 1
-        return self._next_id
+def get_CT_state(state: str) -> str:
+    return f"{state}{CT_POSTFIX}"
+
+
+def is_CT_state(state: str) -> bool:
+    return state.endswith(CT_POSTFIX)
+
+
+@dataclass
+class Individual:
+    node: Tree
+    state: str
+    id: int = field(init=False)
+    _id_counter: ClassVar[int] = 0
+
+    def __post_init__(self):
+        Individual._id_counter += 1
+        self.id = Individual._id_counter
+
+
+class Model:
+    def __init__(
+        self,
+        init_state: str,
+        max_notified_contacts: int = 1,
+        notification_probability: float = 0,
+        rng: int | Generator | None = None,
+    ):
+        self._next_node_id = 0
+        self._population: dict[int, Individual] = {}
+        self._states: dict[str, set[int]] = defaultdict(set)
+        self._contacts: dict[int, list[Individual]] = defaultdict(list)
+        self._sampled: set[str] = set()
+        self._tree = self._get_new_individual(init_state).node
+        self._max_notified_contacts = max_notified_contacts
+        self._notification_probability = notification_probability
+        self._rng = rng if isinstance(rng, Generator) else default_rng(rng)
 
     @property
     def n_sampled(self) -> int:
-        return self._n_sampled
-
-    def _get_new_node(self, state: str | None, branch_length: float | None) -> Tree:
-        id = str(self.next_id) if state is None else f"{self.next_id}|{state}"
-        node = Tree(id, branch_length)
-        if branch_length is None:
-            self._leaves[id] = node
-            self._leaf2state[id] = state
-            self._state2leaves[state].add(id)
-        return node
-
-    def remove(self, node_id: str) -> None:
-        self._state2leaves[self._leaf2state[node_id]].remove(node_id)
-        self._leaf2state.pop(node_id, None)
-        self._leaves.pop(node_id)
-
-    def add_child(
-        self,
-        node_id: str,
-        time: float,
-        stem: bool,
-        state: str | None,
-        branch_length: float | None = None,
-    ) -> None:
-        node = self._leaves[node_id]
+        return len(self._sampled)
+
+    def _get_new_node(self, state: str) -> Tree:
+        self._next_node_id += 1
+        return Tree(f"{self._next_node_id}|{state}")
+
+    def _get_new_individual(self, state: str) -> Individual:
+        individual = Individual(self._get_new_node(state), state)
+        self._population[individual.id] = individual
+        self._states[state].add(individual.id)
+        return individual
+
+    def _set_branch_length(self, node: Tree, time: float) -> None:
         if node.branch_length is not None:
-            raise ValueError("Cannot add a child to a node with a set branch length.")
-        node.add_child(self._get_new_node(state, branch_length))
-        if stem:
-            node.add_child(self._get_new_node(self._leaf2state[node.id], None))
+            raise ValueError(f"Branch length of node {node.id} is already set.")
         node.branch_length = (
             time if node.parent is None else time - node.parent.get_time()
         )
-        self.remove(node_id)
 
-    def sample(self, node_id: str, time: float, remove: bool) -> None:
-        self.add_child(node_id, time, not remove, self._leaf2state[node_id], 0.0)
-        self._n_sampled += 1
+    def _stem(self, individual: Individual, time: float) -> None:
+        self._set_branch_length(individual.node, time)
+        stem_node = self._get_new_node(individual.state)
+        individual.node.add_child(stem_node)
+        individual.node = stem_node
+
+    def remove(self, id: int, time: float) -> None:
+        individual = self._population[id]
+        self._set_branch_length(individual.node, time)
+        state = individual.state
+        self._population.pop(id)
+        self._states[state].remove(id)
+
+    def migrate(self, id: int, state: str, time: float) -> None:
+        individual = self._population[id]
+        self._states[individual.state].remove(id)
+        individual.state = state
+        self._states[state].add(id)
+        self._stem(individual, time)
+
+    def birth_from(self, id: int, state: str, time: float) -> None:
+        individual = self._population[id]
+        new_individual = self._get_new_individual(state)
+        individual.node.add_child(new_individual.node)
+        self._stem(individual, time)
+        self._contacts[id].append(new_individual)
+        self._contacts[new_individual.id].append(individual)
+
+    def sample(self, id: int, time: float, removal_probability: float) -> None:
+        individual = self._population[id]
+        if self._rng.random() < removal_probability:
+            self._sampled.add(individual.node.id)
+            self.remove(id, time)
+        else:
+            sample_node = self._get_new_node(individual.state)
+            sample_node.branch_length = 0.0
+            self._sampled.add(sample_node.id)
+            individual.node.add_child(sample_node)
+            self._stem(individual, time)
+
+        for contact in self._contacts[id][-self._max_notified_contacts :]:
+            if (
+                contact.id in self._population
+                and not is_CT_state(contact.state)
+                and self._rng.random() < self._notification_probability
+            ):
+                self.migrate(contact.id, get_CT_state(contact.state), time)
 
     def get_sampled_tree(self) -> Tree:
         tree = self._tree.copy()
         for node in list(tree.postorder_traversal()):
-            if node.branch_length is None or (
-                node.branch_length > 0 and not node.children
-            ):
+            if node.id not in self._sampled and not node.children:
                 if node.parent is None:
                     raise ValueError("No samples in the tree.")
                 else:
@@ -83,18 +137,18 @@ class Model:
                     node.parent.children.remove(node)
         return tree
 
-    def get_random_leaf(
-        self, state: str | None = None, rng: int | Generator | None = None
-    ) -> str:
-        rng = rng if isinstance(rng, Generator) else default_rng(rng)
+    def get_full_tree(self) -> Tree:
+        return self._tree.copy()
+
+    def get_random_individual(self, state: str | None = None) -> int:
         if state is None:
-            return rng.choice(list(self._leaves))
-        return rng.choice(list(self._state2leaves[state]))
+            return self._rng.choice(list(self._population))
+        return self._rng.choice(list(self._states[state]))
 
-    def get_leaves(self) -> list[str]:
-        return list(self._leaves)
+    def get_population(self) -> list[int]:
+        return list(self._population)
 
-    def count_leaves(self, state: str | None = None) -> int:
+    def count_individuals(self, state: str | None = None) -> int:
         if state is None:
-            return len(self._leaves)
-        return len(self._state2leaves[state])
+            return len(self._population)
+        return len(self._states[state])

phylogenie/typeguards.py

@@ -1,42 +1,42 @@
 from collections.abc import Sequence
-from typing import TypeGuard
+from typing import Any, TypeGuard
 
 import phylogenie.typings as pgt
 
 
-def is_many(x: object) -> TypeGuard[pgt.Many[object]]:
+def is_many(x: Any) -> TypeGuard[pgt.Many[Any]]:
     return isinstance(x, Sequence) and not isinstance(x, str)
 
 
-def is_many_scalars(x: object) -> TypeGuard[pgt.Many[pgt.Scalar]]:
+def is_many_scalars(x: Any) -> TypeGuard[pgt.Many[pgt.Scalar]]:
     return is_many(x) and all(isinstance(i, pgt.Scalar) for i in x)
 
 
-def is_many_ints(x: object) -> TypeGuard[pgt.Many[int]]:
+def is_many_ints(x: Any) -> TypeGuard[pgt.Many[int]]:
     return is_many(x) and all(isinstance(i, int) for i in x)
 
 
-def is_one_or_many_scalars(x: object) -> TypeGuard[pgt.OneOrManyScalars]:
+def is_one_or_many_scalars(x: Any) -> TypeGuard[pgt.OneOrManyScalars]:
     return isinstance(x, pgt.Scalar) or is_many_scalars(x)
 
 
-def is_many_one_or_many_scalars(x: object) -> TypeGuard[pgt.Many[pgt.OneOrManyScalars]]:
+def is_many_one_or_many_scalars(x: Any) -> TypeGuard[pgt.Many[pgt.OneOrManyScalars]]:
     return is_many(x) and all(is_one_or_many_scalars(i) for i in x)
 
 
-def is_many_2D_scalars(x: object) -> TypeGuard[pgt.Many2DScalars]:
+def is_many_2D_scalars(x: Any) -> TypeGuard[pgt.Many2DScalars]:
     return is_many(x) and all(is_many_scalars(i) for i in x)
 
 
-def is_one_or_many_2D_scalars(x: object) -> TypeGuard[pgt.OneOrMany2DScalars]:
+def is_one_or_many_2D_scalars(x: Any) -> TypeGuard[pgt.OneOrMany2DScalars]:
     return isinstance(x, pgt.Scalar) or is_many_2D_scalars(x)
 
 
 def is_many_one_or_many_2D_scalars(
-    x: object,
+    x: Any,
 ) -> TypeGuard[pgt.Many[pgt.OneOrMany2DScalars]]:
     return is_many(x) and all(is_one_or_many_2D_scalars(i) for i in x)
 
 
-def is_many_3D_scalars(x: object) -> TypeGuard[pgt.Many3DScalars]:
+def is_many_3D_scalars(x: Any) -> TypeGuard[pgt.Many3DScalars]:
     return is_many(x) and all(is_many_2D_scalars(i) for i in x)

{phylogenie-2.0.5.dist-info → phylogenie-2.0.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: phylogenie
-Version: 2.0.5
+Version: 2.0.7
 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.7.dist-info/RECORD

@@ -0,0 +1,28 @@
+phylogenie/__init__.py,sha256=1w_0H9lg7hI3b-NLjKuzc34GbJJGyjLq9LrlogecTzI,1759
+phylogenie/generators/__init__.py,sha256=zsOxy28-9j9alOQLIgrOAFfmM58NNHO_NEtW-KXQXAY,888
+phylogenie/generators/alisim.py,sha256=dDqlSwLDbRE2u5SZlsq1mArobTBtuk0aeXY3m1N-bWA,2374
+phylogenie/generators/configs.py,sha256=5ZWdKhRUjlNifw7QKXbooKV1fElqfCk_jBGxfcjh8do,969
+phylogenie/generators/dataset.py,sha256=k6RYJpgxOL8a_yMq98WUF-dcJv8TwxaWnde0k13M4J0,2525
+phylogenie/generators/factories.py,sha256=O8wqL-PvZps0Dq6mQa_PTi4vBvky5LkQIy1jjfOUm-4,6944
+phylogenie/generators/trees.py,sha256=jukaVXGcPGzDBEYMGJ1MKqWt4XbAB5EEfuHXDpwKTqM,9173
+phylogenie/generators/typeguards.py,sha256=Qph6ZnQ7wDMUNvB0VWQKlq42f8wkKOnM42cfMqhNov4,862
+phylogenie/io.py,sha256=ZXlofnSh7FX5UJiP0svRHrTraMSNgKa1GiAv0bMz7jU,2854
+phylogenie/main.py,sha256=vtvSpQxBNlYABoFQ25czl-l3fIr4QRo3svWVd-jcArw,1170
+phylogenie/msa.py,sha256=JDGyZUsAq6-m-SQjoCDjAkAZIxfgyl_PDIhdYn5HOow,2064
+phylogenie/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+phylogenie/skyline/__init__.py,sha256=7pF4CUb4ZCLzNYJNhOjpuTOLTRhlK7L6ugfccNqjIGo,620
+phylogenie/skyline/matrix.py,sha256=Gl8OgKjtieG0NwPYiPimKI36gefV8fm_OeorjdXxPTs,9146
+phylogenie/skyline/parameter.py,sha256=EM9qlPt0JhMBy3TbztM0dj24BaGNEy8KWKdTObDKhbI,4644
+phylogenie/skyline/vector.py,sha256=bJP7_FNX_Klt6wXqsyfj0KX3VNj6-dIhzCKSJuQcOV0,7115
+phylogenie/tree.py,sha256=34gcxUoTGfj72EbIlpnrhWGnNFppUVjms3XEn1ZS3-g,1997
+phylogenie/treesimulator/__init__.py,sha256=INPU9LrPdUmt3dYGzWDRoRKrPR9xENcHu44pJVUbyNA,525
+phylogenie/treesimulator/events.py,sha256=X3_0U9qqMpYgh6-7TwQEnlUipANkHz6QTCXlm-qXFQk,9524
+phylogenie/treesimulator/gillespie.py,sha256=Fn-PyVICx3pWtpHko7rf6omf_kqOkkpebSJy56oPKnQ,3216
+phylogenie/treesimulator/model.py,sha256=XpzAicmg2O6K0Trk5YolH-B_HJZxoSauF2wZOMqp-Iw,5559
+phylogenie/typeguards.py,sha256=JtqmbEWJZBRHbWgCvcl6nrWm3VcBfzRbklbTBYHItn0,1325
+phylogenie/typings.py,sha256=O1X6lGKTjJ2YJz3ApQ-rYb_tEJNUIcHdUIeYlSM4s5o,500
+phylogenie-2.0.7.dist-info/LICENSE.txt,sha256=NUrDqElK-eD3I0WqC004CJsy6cs0JgsAoebDv_42-pw,1071
+phylogenie-2.0.7.dist-info/METADATA,sha256=REo_HAAqLn0XQiHeCPyNPbGuSsbT5p1QupjSxg5Zs_U,5472
+phylogenie-2.0.7.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+phylogenie-2.0.7.dist-info/entry_points.txt,sha256=Rt6_usN0FkBX1ZfiqCirjMN9FKOgFLG8rydcQ8kugeE,51
+phylogenie-2.0.7.dist-info/RECORD,,

phylogenie-2.0.5.dist-info/RECORD

@@ -1,28 +0,0 @@
-phylogenie/__init__.py,sha256=1w_0H9lg7hI3b-NLjKuzc34GbJJGyjLq9LrlogecTzI,1759
-phylogenie/generators/__init__.py,sha256=zsOxy28-9j9alOQLIgrOAFfmM58NNHO_NEtW-KXQXAY,888
-phylogenie/generators/alisim.py,sha256=Psh_fAHHrJBCd1kqVtz_BdoVVmNvieHJRMmapQncJhM,2288
-phylogenie/generators/configs.py,sha256=uiqPvdGSwmItQmhRMT-s1xG7DVnfpwaDTZ_wztL1Pwo,1123
-phylogenie/generators/dataset.py,sha256=zJUSC_c1rNfqq6zTkgGQVPmHcSW47HvfQfDdVTamiT8,2557
-phylogenie/generators/factories.py,sha256=lJ517JC0T8fVu_2NBNC1GUbMPQfRNeNpO5rGdxubQOA,6998
-phylogenie/generators/trees.py,sha256=pSqnYslDMPl4OVsIUFgQu_Rtc2htlBVwhIWnYgSAr3s,8735
-phylogenie/generators/typeguards.py,sha256=6WIQIy6huHok4vYqA9ym6g2kvyyXPggBONWwbiHtDiY,925
-phylogenie/io.py,sha256=ZXlofnSh7FX5UJiP0svRHrTraMSNgKa1GiAv0bMz7jU,2854
-phylogenie/main.py,sha256=vtvSpQxBNlYABoFQ25czl-l3fIr4QRo3svWVd-jcArw,1170
-phylogenie/msa.py,sha256=JDGyZUsAq6-m-SQjoCDjAkAZIxfgyl_PDIhdYn5HOow,2064
-phylogenie/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-phylogenie/skyline/__init__.py,sha256=7pF4CUb4ZCLzNYJNhOjpuTOLTRhlK7L6ugfccNqjIGo,620
-phylogenie/skyline/matrix.py,sha256=BnIreEyCymWzo5jiFtLWXE33ck0pbOFNAfcJW7S2m2s,9147
-phylogenie/skyline/parameter.py,sha256=CJ5OEyRQG2Tg1WJWQ1IpfX-6hjJv80Zj8lMoRke5nnQ,4648
-phylogenie/skyline/vector.py,sha256=becZlWLeT12mvpDC5MLGXZPKlgmbkY-DEcuhdQpnsSc,7135
-phylogenie/tree.py,sha256=nNUQrTtUqgc-Z6CHflqROZ8W2ZL1YWGG7xeWmcKwZUU,1684
-phylogenie/treesimulator/__init__.py,sha256=koyVZxKUeMtY1S2WUgqbeIthoGbdUmx9vP4svFi4a6U,459
-phylogenie/treesimulator/events.py,sha256=RfFodTxMAPUtI3ZkDmw5ekXCaRGnQ8Sy3tjR4oJ8CW4,9636
-phylogenie/treesimulator/gillespie.py,sha256=p2bjwc6iyaf4mMI-xregJwOnXkjT2tIhxXbcGncQ5b0,2293
-phylogenie/treesimulator/model.py,sha256=pPl4JYnBIHjWeSgpIypaDOes11tA4bFAGuMRArXbsVo,3577
-phylogenie/typeguards.py,sha256=WBOSJSaOC8VDtrYoA2w_AYEXTpyKdCfmsM29KaKXl3A,1350
-phylogenie/typings.py,sha256=O1X6lGKTjJ2YJz3ApQ-rYb_tEJNUIcHdUIeYlSM4s5o,500
-phylogenie-2.0.5.dist-info/LICENSE.txt,sha256=NUrDqElK-eD3I0WqC004CJsy6cs0JgsAoebDv_42-pw,1071
-phylogenie-2.0.5.dist-info/METADATA,sha256=PHkxj9W4YYScWLa9oNke4sR_dXlT2-W8f6Z4ekuCgnI,5456
-phylogenie-2.0.5.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-phylogenie-2.0.5.dist-info/entry_points.txt,sha256=Rt6_usN0FkBX1ZfiqCirjMN9FKOgFLG8rydcQ8kugeE,51
-phylogenie-2.0.5.dist-info/RECORD,,