phylogenie 2.1.24__tar.gz → 2.1.26__tar.gz

{phylogenie-2.1.24 → phylogenie-2.1.26}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: phylogenie
-Version: 2.1.24
+Version: 2.1.26
 Summary: Generate phylogenetic datasets with minimal setup effort
 Author: Gabriele Marino
 Author-email: gabmarino.8601@gmail.com

{phylogenie-2.1.24 → phylogenie-2.1.26}/phylogenie/__init__.py

@@ -31,9 +31,9 @@ from phylogenie.treesimulator import (
     BirthWithContactTracing,
     Death,
     Event,
+    EventType,
     Migration,
     Mutation,
-    MutationTargetType,
     Sampling,
     SamplingWithContactTracing,
     generate_trees,
@@ -85,9 +85,9 @@ __all__ = [
     "BirthWithContactTracing",
     "Death",
     "Event",
+    "EventType",
     "Migration",
     "Mutation",
-    "MutationTargetType",
     "Sampling",
     "SamplingWithContactTracing",
     "get_BD_events",

{phylogenie-2.1.24 → phylogenie-2.1.26}/phylogenie/generators/alisim.py

@@ -65,12 +65,13 @@ class AliSimDatasetGenerator(DatasetGenerator):
         while True:
             d.update(data(context, rng))
             try:
-                tree = self.trees.simulate_one(d, seed)
+                tree, metadata = self.trees.simulate_one(d, seed)
                 break
             except TimeoutError:
                 print(
                     "Tree simulation timed out, retrying with different parameters..."
                 )
+        d.update(metadata)
 
         times = get_node_depths(tree)
         for leaf in tree.get_leaves():

{phylogenie-2.1.24 → phylogenie-2.1.26}/phylogenie/generators/configs.py

@@ -1,6 +1,6 @@
 import phylogenie.typings as pgt
 from phylogenie.models import Distribution, StrictBaseModel
-from phylogenie.treesimulator import MutationTargetType
+from phylogenie.treesimulator import EventType
 
 Integer = str | int
 Scalar = str | pgt.Scalar
@@ -30,9 +30,9 @@ SkylineMatrix = str | pgt.Scalar | pgt.Many[SkylineVector] | SkylineMatrixModel
 
 
 class Event(StrictBaseModel):
-    states: str | list[str] | None = None
+    state: str | None = None
     rate: SkylineParameter
 
 
 class Mutation(Event):
-    rate_scalers: dict[MutationTargetType, Distribution]
+    rate_scalers: dict[EventType, Distribution]

{phylogenie-2.1.24 → phylogenie-2.1.26}/phylogenie/generators/factories.py

@@ -17,6 +17,7 @@ from phylogenie.skyline import (
     SkylineVector,
     SkylineVectorCoercible,
 )
+from phylogenie.treesimulator import Mutation
 
 
 def _eval_expression(expression: str, data: dict[str, Any]) -> Any:
@@ -221,6 +222,20 @@ def distribution(x: Distribution, data: dict[str, Any]) -> Distribution:
     return Distribution(type=x.type, **args)
 
 
+def mutations(
+    x: list[cfg.Mutation], data: dict[str, Any], states: set[str]
+) -> list[Mutation]:
+    mutations: list[Mutation] = []
+    for m in x:
+        rate = skyline_parameter(m.rate, data)
+        rate_scalers = {k: distribution(v, data) for k, v in m.rate_scalers.items()}
+        if m.state is None:
+            mutations.extend(Mutation(s, rate, rate_scalers) for s in states)
+        else:
+            mutations.append(Mutation(m.state, rate, rate_scalers))
+    return mutations
+
+
 def data(context: dict[str, Distribution] | None, rng: Generator) -> dict[str, Any]:
     if context is None:
         return {}

{phylogenie-2.1.24 → phylogenie-2.1.26}/phylogenie/generators/trees.py

@@ -11,8 +11,8 @@ import phylogenie.generators.configs as cfg
 from phylogenie.generators.dataset import DatasetGenerator, DataType
 from phylogenie.generators.factories import (
     data,
-    distribution,
     integer,
+    mutations,
     scalar,
     skyline_matrix,
     skyline_parameter,
@@ -24,7 +24,6 @@ from phylogenie.tree import Tree
 from phylogenie.treesimulator import (
     Event,
     Feature,
-    Mutation,
     get_BD_events,
     get_BDEI_events,
     get_BDSS_events,
@@ -48,6 +47,7 @@ class ParameterizationType(str, Enum):
 
 class TreeDatasetGenerator(DatasetGenerator):
     data_type: Literal[DataType.TREES] = DataType.TREES
+    mutations: list[cfg.Mutation] = Field(default_factory=lambda: [])
     min_tips: cfg.Integer = 1
     max_tips: cfg.Integer | None = None
     max_time: cfg.Scalar = np.inf
@@ -59,14 +59,17 @@ class TreeDatasetGenerator(DatasetGenerator):
     @abstractmethod
     def _get_events(self, data: dict[str, Any]) -> list[Event]: ...
 
-    def simulate_one(self, data: dict[str, Any], seed: int | None = None) -> Tree:
+    def simulate_one(
+        self, data: dict[str, Any], seed: int | None = None
+    ) -> tuple[Tree, dict[str, Any]]:
         init_state = (
             self.init_state
             if self.init_state is None
             else self.init_state.format(**data)
         )
+        states = {e.state for e in self._get_events(data)}
         return simulate_tree(
-            events=self._get_events(data),
+            events=self._get_events(data) + mutations(self.mutations, data, states),
             min_tips=integer(self.min_tips, data),
             max_tips=None if self.max_tips is None else integer(self.max_tips, data),
             max_time=scalar(self.max_time, data),
@@ -89,14 +92,14 @@ class TreeDatasetGenerator(DatasetGenerator):
         while True:
             try:
                 d.update(data(context, rng))
-                tree = self.simulate_one(d, seed)
+                tree, metadata = self.simulate_one(d, seed)
                 if self.node_features is not None:
                     set_features(tree, self.node_features)
                 dump_newick(tree, f"{filename}.nwk")
                 break
             except TimeoutError:
                 print("Simulation timed out, retrying with different parameters...")
-        return d
+        return d | metadata
 
 
 class CanonicalTreeDatasetGenerator(TreeDatasetGenerator):
@@ -147,12 +150,11 @@ class FBDTreeDatasetGenerator(TreeDatasetGenerator):
         )
 
 
-class TreeDatasetGeneratorForEpidemiology(TreeDatasetGenerator):
+class ContactTracingTreeDatasetGenerator(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]: ...
@@ -171,30 +173,10 @@ class TreeDatasetGeneratorForEpidemiology(TreeDatasetGenerator):
                 ),
                 samplable_states_after_notification=self.samplable_states_after_notification,
             )
-        all_states = list({e.state for e in events})
-        for mutation in self.mutations:
-            states = mutation.states
-            if isinstance(states, str):
-                states = [states]
-            elif states is None:
-                states = all_states
-            for state in states:
-                if state not in all_states:
-                    raise ValueError(
-                        f"Mutation state '{state}' is not found in states {all_states}."
-                    )
-                rate_scalers = {
-                    t: distribution(r, data) for t, r in mutation.rate_scalers.items()
-                }
-                events.append(
-                    Mutation(
-                        state, skyline_parameter(mutation.rate, data), rate_scalers
-                    )
-                )
         return events
 
 
-class EpidemiologicalTreeDatasetGenerator(TreeDatasetGeneratorForEpidemiology):
+class EpidemiologicalTreeDatasetGenerator(ContactTracingTreeDatasetGenerator):
     parameterization: Literal[ParameterizationType.EPIDEMIOLOGICAL] = (
         ParameterizationType.EPIDEMIOLOGICAL
     )
@@ -220,7 +202,7 @@ class EpidemiologicalTreeDatasetGenerator(TreeDatasetGeneratorForEpidemiology):
         )
 
 
-class BDTreeDatasetGenerator(TreeDatasetGeneratorForEpidemiology):
+class BDTreeDatasetGenerator(ContactTracingTreeDatasetGenerator):
     parameterization: Literal[ParameterizationType.BD] = ParameterizationType.BD
     reproduction_number: cfg.SkylineParameter
     infectious_period: cfg.SkylineParameter
@@ -234,7 +216,7 @@ class BDTreeDatasetGenerator(TreeDatasetGeneratorForEpidemiology):
         )
 
 
-class BDEITreeDatasetGenerator(TreeDatasetGeneratorForEpidemiology):
+class BDEITreeDatasetGenerator(ContactTracingTreeDatasetGenerator):
     parameterization: Literal[ParameterizationType.BDEI] = ParameterizationType.BDEI
     reproduction_number: cfg.SkylineParameter
     infectious_period: cfg.SkylineParameter
@@ -250,7 +232,7 @@ class BDEITreeDatasetGenerator(TreeDatasetGeneratorForEpidemiology):
         )
 
 
-class BDSSTreeDatasetGenerator(TreeDatasetGeneratorForEpidemiology):
+class BDSSTreeDatasetGenerator(ContactTracingTreeDatasetGenerator):
     parameterization: Literal[ParameterizationType.BDSS] = ParameterizationType.BDSS
     reproduction_number: cfg.SkylineParameter
     infectious_period: cfg.SkylineParameter

{phylogenie-2.1.24 → phylogenie-2.1.26}/phylogenie/skyline/matrix.py

@@ -33,7 +33,7 @@ class SkylineMatrix:
     ):
         if params is not None and value is None and change_times is None:
             if is_many_skyline_vectors_like(params):
-                self.params = [
+                self._params = [
                     p if isinstance(p, SkylineVector) else SkylineVector(p)
                     for p in params
                 ]
@@ -41,7 +41,7 @@ class SkylineMatrix:
                 raise TypeError(
                     f"It is impossible to create a SkylineMatrix from `params` {params} of type {type(params)}. Please provide a sequence composed of SkylineVectorLike objects (a SkylineVectorLike object can either be a SkylineVector or a sequence of scalars and/or SkylineParameters)."
                 )
-            lengths = {len(p) for p in self.params}
+            lengths = {len(p) for p in self._params}
             if len(lengths) > 1:
                 raise ValueError(
                     f"All `params` must have the same length to create a SkylineMatrix (got params={params} with lengths {lengths})."
@@ -57,7 +57,7 @@ class SkylineMatrix:
                 raise TypeError(
                     f"It is impossible to create a SkylineMatrix from `value` {value} of type {type(value)}. Please provide a nested (3D) sequence of scalar values."
                 )
-            self.params = [
+            self._params = [
                 SkylineVector(
                     value=[matrix[i] for matrix in value], change_times=change_times
                 )
@@ -68,6 +68,10 @@ class SkylineMatrix:
                 "Either `params` or both `value` and `change_times` must be provided to create a SkylineMatrix."
             )
 
+    @property
+    def params(self) -> tuple[SkylineVector, ...]:
+        return tuple(self._params)
+
     @property
     def n_rows(self) -> int:
         return len(self.params)
@@ -82,14 +86,14 @@ class SkylineMatrix:
 
     @property
     def change_times(self) -> pgt.Vector1D:
-        return sorted(set([t for row in self.params for t in row.change_times]))
+        return tuple(sorted(set([t for row in self.params for t in row.change_times])))
 
     @property
     def value(self) -> pgt.Vector3D:
-        return [self.get_value_at_time(t) for t in (0, *self.change_times)]
+        return tuple(self.get_value_at_time(t) for t in (0, *self.change_times))
 
     def get_value_at_time(self, time: pgt.Scalar) -> pgt.Vector2D:
-        return [param.get_value_at_time(time) for param in self.params]
+        return tuple(param.get_value_at_time(time) for param in self.params)
 
     def _operate(
         self,
@@ -185,7 +189,7 @@ class SkylineMatrix:
             raise TypeError(
                 f"It is impossible to set item of SkylineMatrix to value {value} of type {type(value)}. Please provide a SkylineVectorLike object (i.e., a SkylineVector or a sequence of scalars and/or SkylineParameters)."
             )
-        self.params[item] = skyline_vector(value, self.n_cols)
+        self._params[item] = skyline_vector(value, self.n_cols)
 
 
 def skyline_matrix(

{phylogenie-2.1.24 → phylogenie-2.1.26}/phylogenie/skyline/parameter.py

@@ -52,12 +52,20 @@ class SkylineParameter:
                 f"`change_times` must be non-negative (got change_times={change_times})."
             )
 
-        self.value = [value[0]]
-        self.change_times: list[pgt.Scalar] = []
+        self._value = [value[0]]
+        self._change_times: list[pgt.Scalar] = []
         for i in range(1, len(value)):
             if value[i] != value[i - 1]:
-                self.value.append(value[i])
-                self.change_times.append(change_times[i - 1])
+                self._value.append(value[i])
+                self._change_times.append(change_times[i - 1])
+
+    @property
+    def value(self) -> pgt.Vector1D:
+        return tuple(self._value)
+
+    @property
+    def change_times(self) -> pgt.Vector1D:
+        return tuple(self._change_times)
 
     def get_value_at_time(self, t: pgt.Scalar) -> pgt.Scalar:
         if t < 0:

{phylogenie-2.1.24 → phylogenie-2.1.26}/phylogenie/skyline/vector.py

@@ -47,7 +47,7 @@ class SkylineVector:
     ):
         if params is not None and value is None and change_times is None:
             if is_many_skyline_parameters_like(params):
-                self.params = [skyline_parameter(param) for param in params]
+                self._params = [skyline_parameter(param) for param in params]
             else:
                 raise TypeError(
                     f"It is impossible to create a SkylineVector from `params` {params} of type {type(params)}. Please provide a sequence of SkylineParameterLike objects (a SkylineParameterLike object can either be a SkylineParameter or a scalar)."
@@ -63,7 +63,7 @@ class SkylineVector:
                 raise TypeError(
                     f"It is impossible to create a SkylineVector from `value` {value} of type {type(value)}. Please provide a nested (2D) sequence of scalar values."
                 )
-            self.params = [
+            self._params = [
                 SkylineParameter([vector[i] for vector in value], change_times)
                 for i in range(len(value[0]))
             ]
@@ -72,20 +72,26 @@ class SkylineVector:
                 "Either `params` or both `value` and `change_times` must be provided to create a SkylineVector."
             )
 
+    @property
+    def params(self) -> tuple[SkylineParameter, ...]:
+        return tuple(self._params)
+
     @property
     def change_times(self) -> pgt.Vector1D:
-        return sorted(set(t for param in self.params for t in param.change_times))
+        return tuple(
+            sorted(set(t for param in self.params for t in param.change_times))
+        )
 
     @property
     def value(self) -> pgt.Vector2D:
-        return [self.get_value_at_time(t) for t in (0, *self.change_times)]
+        return tuple(self.get_value_at_time(t) for t in (0, *self.change_times))
 
     @property
     def N(self) -> int:
         return len(self.params)
 
     def get_value_at_time(self, t: pgt.Scalar) -> pgt.Vector1D:
-        return [param.get_value_at_time(t) for param in self.params]
+        return tuple(param.get_value_at_time(t) for param in self.params)
 
     def _operate(
         self,
@@ -154,7 +160,7 @@ class SkylineVector:
             raise TypeError(
                 f"It is impossible to set item {item} of SkylineVector with value {value} of type {type(value)}. Please provide a SkylineParameterLike object (i.e., a scalar or a SkylineParameter)."
             )
-        self.params[item] = skyline_parameter(value)
+        self._params[item] = skyline_parameter(value)
 
 
 def skyline_vector(x: SkylineVectorCoercible, N: int) -> SkylineVector:

{phylogenie-2.1.24 → phylogenie-2.1.26}/phylogenie/treesimulator/__init__.py

@@ -3,9 +3,9 @@ from phylogenie.treesimulator.events import (
     BirthWithContactTracing,
     Death,
     Event,
+    EventType,
     Migration,
     Mutation,
-    MutationTargetType,
     Sampling,
     SamplingWithContactTracing,
     get_BD_events,
@@ -26,9 +26,9 @@ __all__ = [
     "BirthWithContactTracing",
     "Death",
     "Event",
+    "EventType",
     "Migration",
     "Mutation",
-    "MutationTargetType",
     "Sampling",
     "SamplingWithContactTracing",
     "get_BD_events",

{phylogenie-2.1.24 → phylogenie-2.1.26}/phylogenie/treesimulator/events/__init__.py

@@ -1,3 +1,4 @@
+from phylogenie.treesimulator.events.base import Event, EventType
 from phylogenie.treesimulator.events.contact_tracing import (
     BirthWithContactTracing,
     SamplingWithContactTracing,
@@ -6,7 +7,6 @@ from phylogenie.treesimulator.events.contact_tracing import (
 from phylogenie.treesimulator.events.core import (
     Birth,
     Death,
-    Event,
     Migration,
     Sampling,
     get_BD_events,
@@ -16,20 +16,17 @@ from phylogenie.treesimulator.events.core import (
     get_epidemiological_events,
     get_FBD_events,
 )
-from phylogenie.treesimulator.events.mutations import Mutation
-from phylogenie.treesimulator.events.mutations import TargetType as MutationTargetType
-from phylogenie.treesimulator.events.mutations import get_mutation_id
+from phylogenie.treesimulator.events.mutations import Mutation, get_mutation_id
 
 __all__ = [
     "Birth",
     "BirthWithContactTracing",
     "Death",
     "Event",
+    "EventType",
     "Migration",
-    "Mutation",
     "Sampling",
     "SamplingWithContactTracing",
-    "MutationTargetType",
     "get_BD_events",
     "get_BDEI_events",
     "get_BDSS_events",
@@ -37,5 +34,6 @@ __all__ = [
     "get_contact_tracing_events",
     "get_epidemiological_events",
     "get_FBD_events",
+    "Mutation",
     "get_mutation_id",
 ]

phylogenie-2.1.26/phylogenie/treesimulator/events/base.py

@@ -0,0 +1,44 @@
+from abc import ABC, abstractmethod
+from enum import Enum
+from typing import Any
+
+import numpy as np
+from numpy.random import Generator
+
+from phylogenie.skyline import SkylineParameterLike, skyline_parameter
+from phylogenie.treesimulator.model import Model
+
+
+class EventType(str, Enum):
+    BIRTH = "birth"
+    DEATH = "death"
+    MIGRATION = "migration"
+    SAMPLING = "sampling"
+    MUTATION = "mutation"
+    BIRTH_WITH_CT = "birth_with_contact_tracing"
+    SAMPLING_WITH_CT = "sampling_with_contact_tracing"
+
+
+class Event(ABC):
+    type: EventType
+
+    def __init__(self, state: str, rate: SkylineParameterLike):
+        self.state = state
+        self.rate = skyline_parameter(rate)
+        if any(v < 0 for v in self.rate.value):
+            raise ValueError("Event rates must be non-negative.")
+
+    def draw_individual(self, model: Model, rng: Generator) -> int:
+        return rng.choice(model.get_population(self.state))
+
+    def get_propensity(self, model: Model, time: float) -> float:
+        n_individuals = model.count_individuals(self.state)
+        rate = self.rate.get_value_at_time(time)
+        if rate == np.inf and not n_individuals:
+            return 0
+        return rate * n_individuals
+
+    @abstractmethod
+    def apply(
+        self, model: Model, events: "list[Event]", time: float, rng: Generator
+    ) -> dict[str, Any] | None: ...

{phylogenie-2.1.24 → phylogenie-2.1.26}/phylogenie/treesimulator/events/contact_tracing.py

@@ -1,31 +1,35 @@
 from collections import defaultdict
 from collections.abc import Sequence
+from copy import deepcopy
 
 import numpy as np
 from numpy.random import Generator
 
 from phylogenie.skyline import SkylineParameterLike, skyline_parameter
+from phylogenie.treesimulator.events.base import Event, EventType
 from phylogenie.treesimulator.events.core import Birth, Death, Migration, Sampling
-from phylogenie.treesimulator.model import Event, Model
+from phylogenie.treesimulator.model import Model
 
 CT_POSTFIX = "-CT"
 CONTACTS_KEY = "CONTACTS"
 
 
-def _get_CT_state(state: str) -> str:
+def get_CT_state(state: str) -> str:
     return f"{state}{CT_POSTFIX}"
 
 
-def _is_CT_state(state: str) -> bool:
+def is_CT_state(state: str) -> bool:
     return state.endswith(CT_POSTFIX)
 
 
 class BirthWithContactTracing(Event):
+    type = EventType.BIRTH_WITH_CT
+
     def __init__(self, state: str, rate: SkylineParameterLike, child_state: str):
         super().__init__(state, rate)
         self.child_state = child_state
 
-    def apply(self, model: Model, time: float, rng: Generator) -> None:
+    def apply(self, model: Model, events: list[Event], time: float, rng: Generator):
         individual = self.draw_individual(model, rng)
         new_individual = model.birth_from(individual, self.child_state, time)
         if CONTACTS_KEY not in model.context:
@@ -38,6 +42,8 @@ class BirthWithContactTracing(Event):
 
 
 class SamplingWithContactTracing(Event):
+    type = EventType.SAMPLING_WITH_CT
+
     def __init__(
         self,
         state: str,
@@ -49,7 +55,7 @@ class SamplingWithContactTracing(Event):
         self.max_notified_contacts = max_notified_contacts
         self.notification_probability = skyline_parameter(notification_probability)
 
-    def apply(self, model: Model, time: float, rng: Generator) -> None:
+    def apply(self, model: Model, events: list[Event], time: float, rng: Generator):
         individual = self.draw_individual(model, rng)
         model.sample(individual, time, True)
         population = model.get_population()
@@ -60,8 +66,8 @@ class SamplingWithContactTracing(Event):
             if contact in population:
                 state = model.get_state(contact)
                 p = self.notification_probability.get_value_at_time(time)
-                if not _is_CT_state(state) and rng.random() < p:
-                    model.migrate(contact, _get_CT_state(state), time)
+                if not is_CT_state(state) and rng.random() < p:
+                    model.migrate(contact, get_CT_state(state), time)
 
     def __repr__(self) -> str:
         return f"SamplingWithContactTracing(state={self.state}, rate={self.rate}, max_notified_contacts={self.max_notified_contacts}, notification_probability={self.notification_probability})"
@@ -79,17 +85,24 @@ def get_contact_tracing_events(
     sampling_rate_after_notification = skyline_parameter(
         sampling_rate_after_notification
     )
-    for event in events:
-        state, rate = event.state, event.rate
+    for event in [deepcopy(e) for e in events]:
         if isinstance(event, Migration):
             ct_events.append(event)
             ct_events.append(
-                Migration(_get_CT_state(state), rate, _get_CT_state(event.target_state))
+                Migration(
+                    get_CT_state(event.state),
+                    event.rate,
+                    get_CT_state(event.target_state),
+                )
             )
         elif isinstance(event, Birth):
-            ct_events.append(BirthWithContactTracing(state, rate, event.child_state))
             ct_events.append(
-                BirthWithContactTracing(_get_CT_state(state), rate, event.child_state)
+                BirthWithContactTracing(event.state, event.rate, event.child_state)
+            )
+            ct_events.append(
+                BirthWithContactTracing(
+                    get_CT_state(event.state), event.rate, event.child_state
+                )
             )
         elif isinstance(event, Sampling):
             if not event.removal:
@@ -98,7 +111,10 @@ def get_contact_tracing_events(
                 )
             ct_events.append(
                 SamplingWithContactTracing(
-                    state, rate, max_notified_contacts, notification_probability
+                    event.state,
+                    event.rate,
+                    max_notified_contacts,
+                    notification_probability,
                 )
             )
         elif isinstance(event, Death):
@@ -115,7 +131,7 @@ def get_contact_tracing_events(
     ):
         ct_events.append(
             SamplingWithContactTracing(
-                _get_CT_state(state),
+                get_CT_state(state),
                 sampling_rate_after_notification,
                 max_notified_contacts,
                 notification_probability,

{phylogenie-2.1.24 → phylogenie-2.1.26}/phylogenie/treesimulator/events/core.py

@@ -7,7 +7,8 @@ from phylogenie.skyline import (
     skyline_matrix,
     skyline_vector,
 )
-from phylogenie.treesimulator.model import Event, Model
+from phylogenie.treesimulator.events.base import Event, EventType
+from phylogenie.treesimulator.model import Model
 
 INFECTIOUS_STATE = "I"
 EXPOSED_STATE = "E"
@@ -15,11 +16,13 @@ SUPERSPREADER_STATE = "S"
 
 
 class Birth(Event):
+    type = EventType.BIRTH
+
     def __init__(self, state: str, rate: SkylineParameterLike, child_state: str):
         super().__init__(state, rate)
         self.child_state = child_state
 
-    def apply(self, model: Model, time: float, rng: Generator) -> None:
+    def apply(self, model: Model, events: list[Event], time: float, rng: Generator):
         individual = self.draw_individual(model, rng)
         model.birth_from(individual, self.child_state, time)
 
@@ -28,7 +31,9 @@ class Birth(Event):
 
 
 class Death(Event):
-    def apply(self, model: Model, time: float, rng: Generator) -> None:
+    type = EventType.DEATH
+
+    def apply(self, model: Model, events: list[Event], time: float, rng: Generator):
         individual = self.draw_individual(model, rng)
         model.remove(individual, time)
 
@@ -37,11 +42,13 @@ class Death(Event):
 
 
 class Migration(Event):
+    type = EventType.MIGRATION
+
     def __init__(self, state: str, rate: SkylineParameterLike, target_state: str):
         super().__init__(state, rate)
         self.target_state = target_state
 
-    def apply(self, model: Model, time: float, rng: Generator) -> None:
+    def apply(self, model: Model, events: list[Event], time: float, rng: Generator):
         individual = self.draw_individual(model, rng)
         model.migrate(individual, self.target_state, time)
 
@@ -50,11 +57,13 @@ class Migration(Event):
 
 
 class Sampling(Event):
+    type = EventType.SAMPLING
+
     def __init__(self, state: str, rate: SkylineParameterLike, removal: bool):
         super().__init__(state, rate)
         self.removal = removal
 
-    def apply(self, model: Model, time: float, rng: Generator) -> None:
+    def apply(self, model: Model, events: list[Event], time: float, rng: Generator):
         individual = self.draw_individual(model, rng)
         model.sample(individual, time, self.removal)
 

{phylogenie-2.1.24 → phylogenie-2.1.26}/phylogenie/treesimulator/events/mutations.py

@@ -1,27 +1,21 @@
 import re
 from copy import deepcopy
-from enum import Enum
-from typing import Type
+from typing import Any
 
 from numpy.random import Generator
 
 from phylogenie.models import Distribution
 from phylogenie.skyline import SkylineParameterLike
+from phylogenie.treesimulator.events.base import Event, EventType
 from phylogenie.treesimulator.events.contact_tracing import (
     BirthWithContactTracing,
     SamplingWithContactTracing,
 )
-from phylogenie.treesimulator.events.core import (
-    Birth,
-    Death,
-    Event,
-    Migration,
-    Sampling,
-)
+from phylogenie.treesimulator.events.core import Birth, Death, Migration, Sampling
 from phylogenie.treesimulator.model import Model
 
 MUTATION_PREFIX = "MUT-"
-MUTATIONS_KEY = "MUTATIONS"
+NEXT_MUTATION_ID = "NEXT_MUTATION_ID"
 
 
 def _get_mutation(state: str) -> str | None:
@@ -41,44 +35,42 @@ def get_mutation_id(node_name: str) -> int:
     return 0
 
 
-class TargetType(str, Enum):
-    BIRTH = "birth"
-    DEATH = "death"
-    MIGRATION = "migration"
-    SAMPLING = "sampling"
-    MUTATION = "mutation"
-
-
 class Mutation(Event):
+    type = EventType.MUTATION
+
     def __init__(
         self,
         state: str,
         rate: SkylineParameterLike,
-        rate_scalers: dict[TargetType, Distribution],
+        rate_scalers: dict[EventType, Distribution],
     ):
         super().__init__(state, rate)
         self.rate_scalers = rate_scalers
 
-    def apply(self, model: Model, time: float, rng: Generator) -> None:
-        if MUTATIONS_KEY not in model.context:
-            model.context[MUTATIONS_KEY] = 0
-        model.context[MUTATIONS_KEY] += 1
-        mutation_id = model.context[MUTATIONS_KEY]
+    def apply(
+        self, model: Model, events: list[Event], time: float, rng: Generator
+    ) -> dict[str, Any]:
+        if NEXT_MUTATION_ID not in model.context:
+            model.context[NEXT_MUTATION_ID] = 0
+        model.context[NEXT_MUTATION_ID] += 1
+        mutation_id = model.context[NEXT_MUTATION_ID]
 
         individual = self.draw_individual(model, rng)
         model.migrate(individual, _get_mutated_state(mutation_id, self.state), time)
 
-        rate_scalers = {
+        rate_scalers: dict[EventType, float] = {
             target_type: getattr(rng, rate_scaler.type)(**rate_scaler.args)
             for target_type, rate_scaler in self.rate_scalers.items()
         }
 
+        metadata: dict[str, Any] = {}
         for event in [
             deepcopy(e)
-            for e in model.events
+            for e in events
             if _get_mutation(self.state) == _get_mutation(e.state)
         ]:
             event.state = _get_mutated_state(mutation_id, event.state)
+
             if isinstance(event, Birth | BirthWithContactTracing):
                 event.child_state = _get_mutated_state(mutation_id, event.child_state)
             elif isinstance(event, Migration):
@@ -87,27 +79,20 @@ class Mutation(Event):
                 event, Mutation | Death | Sampling | SamplingWithContactTracing
             ):
                 raise ValueError(
-                    f"Mutation not defined for event of type {type(event)}."
+                    f"Mutation not implemented for event of type {type(event)}."
                 )
 
-            for target_type, rate_scaler in rate_scalers.items():
-                if target_type not in TARGETS:
-                    raise ValueError(
-                        f"Unsupported target type {target_type} for mutation."
-                    )
-                if isinstance(event, TARGETS[target_type]):
-                    event.rate *= rate_scaler
-
-            model.add_event(event)
+            if event.type in rate_scalers:
+                event.rate *= rate_scalers[event.type]
+                metadata[f"{MUTATION_PREFIX}{mutation_id}.{event.type}.rate.value"] = (
+                    event.rate.value[0]
+                    if len(event.rate.value) == 1
+                    else list(event.rate.value)
+                )
 
-    def __repr__(self) -> str:
-        return f"Mutation(state={self.state}, rate={self.rate})"
+            events.append(event)
 
+        return metadata
 
-TARGETS: dict[TargetType, tuple[Type[Event], ...]] = {
-    TargetType.BIRTH: (Birth, BirthWithContactTracing),
-    TargetType.DEATH: (Death,),
-    TargetType.MIGRATION: (Migration,),
-    TargetType.SAMPLING: (Sampling, SamplingWithContactTracing),
-    TargetType.MUTATION: (Mutation,),
-}
+    def __repr__(self) -> str:
+        return f"Mutation(state={self.state}, rate={self.rate}, rate_scalers={self.rate_scalers})"

{phylogenie-2.1.24 → phylogenie-2.1.26}/phylogenie/treesimulator/features.py

@@ -2,7 +2,7 @@ from collections.abc import Iterable
 from enum import Enum
 
 from phylogenie.tree import Tree
-from phylogenie.treesimulator.events import get_mutation_id
+from phylogenie.treesimulator.events.mutations import get_mutation_id
 from phylogenie.treesimulator.model import get_node_state
 from phylogenie.utils import (
     get_node_depth_levels,

{phylogenie-2.1.24 → phylogenie-2.1.26}/phylogenie/treesimulator/gillespie.py

@@ -1,16 +1,19 @@
 import os
 import time
 from collections.abc import Iterable, Sequence
+from typing import Any
 
 import joblib
 import numpy as np
+import pandas as pd
 from numpy.random import default_rng
 from tqdm import tqdm
 
 from phylogenie.io import dump_newick
 from phylogenie.tree import Tree
+from phylogenie.treesimulator.events import Event
 from phylogenie.treesimulator.features import Feature, set_features
-from phylogenie.treesimulator.model import Event, Model
+from phylogenie.treesimulator.model import Model
 
 
 def simulate_tree(
@@ -22,10 +25,7 @@ def simulate_tree(
     sampling_probability_at_present: float = 0.0,
     seed: int | None = None,
     timeout: float = np.inf,
-) -> Tree:
-    if max_time == np.inf and max_tips is None:
-        raise ValueError("Either max_time or max_tips must be specified.")
-
+) -> tuple[Tree, dict[str, Any]]:
     if max_time == np.inf and sampling_probability_at_present:
         raise ValueError(
             "sampling_probability_at_present cannot be set when max_time is infinite."
@@ -44,12 +44,16 @@ def simulate_tree(
     rng = default_rng(seed)
     start_clock = time.perf_counter()
     while True:
-        model = Model(init_state, events)
+        model = Model(init_state)
+        metadata: dict[str, Any] = {}
+        run_events = list(events)
         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
+
         if max_time == np.inf:
-            assert max_tips is not None
+            if max_tips is None:
+                raise ValueError("Either max_time or max_tips must be specified.")
             target_n_tips = rng.integers(min_tips, max_tips + 1)
         else:
             target_n_tips = None
@@ -58,13 +62,12 @@ def simulate_tree(
             if time.perf_counter() - start_clock > timeout:
                 raise TimeoutError("Simulation timed out.")
 
-            events = model.events
-            rates = [e.get_propensity(model, current_time) for e in events]
+            rates = [e.get_propensity(model, current_time) for e in run_events]
 
-            instantaneous_events = [e for e, r in zip(events, rates) if r == np.inf]
+            instantaneous_events = [e for e, r in zip(run_events, rates) if r == np.inf]
             if instantaneous_events:
                 event = instantaneous_events[rng.integers(len(instantaneous_events))]
-                event.apply(model, current_time, rng)
+                event.apply(model, run_events, current_time, rng)
                 continue
 
             if (
@@ -87,7 +90,10 @@ def simulate_tree(
             current_time += time_step
 
             event_idx = np.searchsorted(np.cumsum(rates) / sum(rates), rng.random())
-            events[int(event_idx)].apply(model, current_time, rng)
+            event = run_events[int(event_idx)]
+            event_metadata = event.apply(model, run_events, current_time, rng)
+            if event_metadata is not None:
+                metadata.update(event_metadata)
 
         for individual in model.get_population():
             if rng.random() < sampling_probability_at_present:
@@ -96,7 +102,7 @@ def simulate_tree(
         if min_tips <= model.n_sampled and (
             max_tips is None or model.n_sampled <= max_tips
         ):
-            return model.get_sampled_tree()
+            return (model.get_sampled_tree(), metadata)
 
 
 def generate_trees(
@@ -112,11 +118,11 @@ def generate_trees(
     seed: int | None = None,
     n_jobs: int = -1,
     timeout: float = np.inf,
-) -> None:
-    def _simulate_tree(seed: int) -> Tree:
+) -> pd.DataFrame:
+    def _simulate_tree(seed: int) -> tuple[Tree, dict[str, Any]]:
         while True:
             try:
-                tree = simulate_tree(
+                tree, metadata = simulate_tree(
                     events=events,
                     min_tips=min_tips,
                     max_tips=max_tips,
@@ -128,7 +134,7 @@ def generate_trees(
                 )
                 if node_features is not None:
                     set_features(tree, node_features)
-                return tree
+                return (tree, metadata)
             except TimeoutError:
                 print("Simulation timed out, retrying with a different seed...")
             seed += 1
@@ -142,7 +148,11 @@ def generate_trees(
         joblib.delayed(_simulate_tree)(seed=int(rng.integers(2**32)))
         for _ in range(n_trees)
     )
-    for i, tree in tqdm(
+
+    df: list[dict[str, Any]] = []
+    for i, (tree, metadata) in tqdm(
         enumerate(jobs), total=n_trees, desc=f"Generating trees in {output_dir}..."
     ):
+        df.append({"file_id": i} | metadata)
         dump_newick(tree, os.path.join(output_dir, f"{i}.nwk"))
+    return pd.DataFrame(df)

{phylogenie-2.1.24 → phylogenie-2.1.26}/phylogenie/treesimulator/model.py

@@ -1,13 +1,7 @@
-from abc import ABC, abstractmethod
 from collections import defaultdict
-from collections.abc import Sequence
 from dataclasses import dataclass
 from typing import Any
 
-import numpy as np
-from numpy.random import Generator
-
-from phylogenie.skyline import SkylineParameterLike, skyline_parameter
 from phylogenie.tree import Tree
 
 
@@ -18,25 +12,6 @@ class Individual:
     state: str
 
 
-class Event(ABC):
-    def __init__(self, state: str, rate: SkylineParameterLike):
-        self.state = state
-        self.rate = skyline_parameter(rate)
-
-    def draw_individual(self, model: "Model", rng: Generator) -> int:
-        return rng.choice(model.get_population(self.state))
-
-    def get_propensity(self, model: "Model", time: float) -> float:
-        n_individuals = model.count_individuals(self.state)
-        rate = self.rate.get_value_at_time(time)
-        if rate == np.inf and not n_individuals:
-            return 0
-        return rate * n_individuals
-
-    @abstractmethod
-    def apply(self, model: "Model", time: float, rng: Generator) -> None: ...
-
-
 def _get_node_name(node_id: int, state: str) -> str:
     return f"{node_id}|{state}"
 
@@ -51,27 +26,19 @@ def get_node_state(node_name: str) -> str:
 
 
 class Model:
-    def __init__(self, init_state: str, events: Sequence[Event]):
+    def __init__(self, init_state: str):
         self._next_node_id = 0
         self._next_individual_id = 0
         self._population: dict[int, Individual] = {}
         self._states: dict[str, set[int]] = defaultdict(set)
         self._sampled: set[str] = set()
         self._tree = self._get_new_individual(init_state).node
-        self._events = list(events)
         self.context: dict[str, Any] = {}
 
     @property
     def n_sampled(self) -> int:
         return len(self._sampled)
 
-    @property
-    def events(self) -> tuple[Event, ...]:
-        return tuple(self._events)
-
-    def add_event(self, event: Event) -> None:
-        self._events.append(event)
-
     def _get_new_node(self, state: str) -> Tree:
         self._next_node_id += 1
         node = Tree(_get_node_name(self._next_node_id, state))

{phylogenie-2.1.24 → phylogenie-2.1.26}/phylogenie/typings.py

@@ -15,6 +15,6 @@ OneOrMany2DScalars = OneOrMany2D[Scalar]
 Many2DScalars = Many2D[Scalar]
 Many3DScalars = Many3D[Scalar]
 
-Vector1D = list[Scalar]
-Vector2D = list[Vector1D]
-Vector3D = list[Vector2D]
+Vector1D = tuple[Scalar, ...]
+Vector2D = tuple[Vector1D, ...]
+Vector3D = tuple[Vector2D, ...]

{phylogenie-2.1.24 → phylogenie-2.1.26}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "phylogenie"
-version = "2.1.24"
+version = "2.1.26"
 description = "Generate phylogenetic datasets with minimal setup effort"
 authors = ["Gabriele Marino <gabmarino.8601@gmail.com>"]
 readme = "README.md"