phylogenie 2.1.25__tar.gz → 2.1.26__tar.gz

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

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

{phylogenie-2.1.25 → 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.25 → 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
@@ -29,6 +29,10 @@ SkylineVector = str | pgt.Scalar | pgt.Many[SkylineParameter] | SkylineVectorMod
 SkylineMatrix = str | pgt.Scalar | pgt.Many[SkylineVector] | SkylineMatrixModel | None
 
 
-class Mutation(StrictBaseModel):
-    rate: Scalar
-    rate_scalers: dict[MutationTargetType, Distribution]
+class Event(StrictBaseModel):
+    state: str | None = None
+    rate: SkylineParameter
+
+
+class Mutation(Event):
+    rate_scalers: dict[EventType, Distribution]

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

@@ -222,11 +222,18 @@ def distribution(x: Distribution, data: dict[str, Any]) -> Distribution:
     return Distribution(type=x.type, **args)
 
 
-def mutation(x: cfg.Mutation, data: dict[str, Any]) -> Mutation:
-    return Mutation(
-        scalar(x.rate, data),
-        {k: distribution(v, data) for k, v in x.rate_scalers.items()},
-    )
+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]:

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

@@ -12,7 +12,7 @@ from phylogenie.generators.dataset import DatasetGenerator, DataType
 from phylogenie.generators.factories import (
     data,
     integer,
-    mutation,
+    mutations,
     scalar,
     skyline_matrix,
     skyline_parameter,
@@ -47,7 +47,7 @@ class ParameterizationType(str, Enum):
 
 class TreeDatasetGenerator(DatasetGenerator):
     data_type: Literal[DataType.TREES] = DataType.TREES
-    mutations: list[cfg.Mutation] | None = None
+    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
@@ -67,14 +67,9 @@ class TreeDatasetGenerator(DatasetGenerator):
             if self.init_state is None
             else self.init_state.format(**data)
         )
-        mutations = (
-            None
-            if self.mutations is None
-            else [mutation(m, data) for m in self.mutations]
-        )
+        states = {e.state for e in self._get_events(data)}
         return simulate_tree(
-            events=self._get_events(data),
-            mutations=mutations,
+            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),

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

@@ -3,7 +3,9 @@ from phylogenie.treesimulator.events import (
     BirthWithContactTracing,
     Death,
     Event,
+    EventType,
     Migration,
+    Mutation,
     Sampling,
     SamplingWithContactTracing,
     get_BD_events,
@@ -13,22 +15,20 @@ from phylogenie.treesimulator.events import (
     get_contact_tracing_events,
     get_epidemiological_events,
     get_FBD_events,
+    get_mutation_id,
 )
 from phylogenie.treesimulator.features import Feature, set_features
 from phylogenie.treesimulator.gillespie import generate_trees, simulate_tree
 from phylogenie.treesimulator.model import get_node_state
-from phylogenie.treesimulator.mutations import Mutation
-from phylogenie.treesimulator.mutations import TargetType as MutationTargetType
-from phylogenie.treesimulator.mutations import get_mutation_id
 
 __all__ = [
     "Birth",
     "BirthWithContactTracing",
     "Death",
     "Event",
+    "EventType",
     "Migration",
     "Mutation",
-    "MutationTargetType",
     "Sampling",
     "SamplingWithContactTracing",
     "get_BD_events",

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

@@ -1,4 +1,4 @@
-from phylogenie.treesimulator.events.base import Event
+from phylogenie.treesimulator.events.base import Event, EventType
 from phylogenie.treesimulator.events.contact_tracing import (
     BirthWithContactTracing,
     SamplingWithContactTracing,
@@ -16,12 +16,14 @@ from phylogenie.treesimulator.events.core import (
     get_epidemiological_events,
     get_FBD_events,
 )
+from phylogenie.treesimulator.events.mutations import Mutation, get_mutation_id
 
 __all__ = [
     "Birth",
     "BirthWithContactTracing",
     "Death",
     "Event",
+    "EventType",
     "Migration",
     "Sampling",
     "SamplingWithContactTracing",
@@ -32,4 +34,6 @@ __all__ = [
     "get_contact_tracing_events",
     "get_epidemiological_events",
     "get_FBD_events",
+    "Mutation",
+    "get_mutation_id",
 ]

{phylogenie-2.1.25 → phylogenie-2.1.26}/phylogenie/treesimulator/events/base.py

@@ -1,4 +1,6 @@
 from abc import ABC, abstractmethod
+from enum import Enum
+from typing import Any
 
 import numpy as np
 from numpy.random import Generator
@@ -7,10 +9,24 @@ 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))
@@ -23,4 +39,6 @@ class Event(ABC):
         return rate * n_individuals
 
     @abstractmethod
-    def apply(self, model: Model, time: float, rng: Generator) -> None: ...
+    def apply(
+        self, model: Model, events: "list[Event]", time: float, rng: Generator
+    ) -> dict[str, Any] | None: ...

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

@@ -1,11 +1,12 @@
 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
+from phylogenie.treesimulator.events.base import Event, EventType
 from phylogenie.treesimulator.events.core import Birth, Death, Migration, Sampling
 from phylogenie.treesimulator.model import Model
 
@@ -13,20 +14,22 @@ 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:
@@ -39,6 +42,8 @@ class BirthWithContactTracing(Event):
 
 
 class SamplingWithContactTracing(Event):
+    type = EventType.SAMPLING_WITH_CT
+
     def __init__(
         self,
         state: str,
@@ -50,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()
@@ -61,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})"
@@ -80,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:
@@ -99,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):
@@ -116,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.25 → phylogenie-2.1.26}/phylogenie/treesimulator/events/core.py

@@ -7,7 +7,7 @@ from phylogenie.skyline import (
     skyline_matrix,
     skyline_vector,
 )
-from phylogenie.treesimulator.events.base import Event
+from phylogenie.treesimulator.events.base import Event, EventType
 from phylogenie.treesimulator.model import Model
 
 INFECTIOUS_STATE = "I"
@@ -16,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)
 
@@ -29,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)
 
@@ -38,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)
 
@@ -51,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.25/phylogenie/treesimulator → phylogenie-2.1.26/phylogenie/treesimulator/events}/mutations.py

@@ -1,22 +1,17 @@
 import re
 from copy import deepcopy
-from enum import Enum
-from typing import Any, 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-"
@@ -40,26 +35,16 @@ def get_mutation_id(node_name: str) -> int:
     return 0
 
 
-class TargetType(str, Enum):
-    BIRTH = "birth"
-    DEATH = "death"
-    MIGRATION = "migration"
-    SAMPLING = "sampling"
-
-
-EVENT_TARGET_TYPES: dict[Type[Event], TargetType] = {
-    Birth: TargetType.BIRTH,
-    BirthWithContactTracing: TargetType.BIRTH,
-    Death: TargetType.DEATH,
-    Migration: TargetType.MIGRATION,
-    Sampling: TargetType.SAMPLING,
-    SamplingWithContactTracing: TargetType.SAMPLING,
-}
+class Mutation(Event):
+    type = EventType.MUTATION
 
-
-class Mutation:
-    def __init__(self, rate: float, rate_scalers: dict[TargetType, Distribution]):
-        self.rate = rate
+    def __init__(
+        self,
+        state: str,
+        rate: SkylineParameterLike,
+        rate_scalers: dict[EventType, Distribution],
+    ):
+        super().__init__(state, rate)
         self.rate_scalers = rate_scalers
 
     def apply(
@@ -70,11 +55,10 @@ class Mutation:
         model.context[NEXT_MUTATION_ID] += 1
         mutation_id = model.context[NEXT_MUTATION_ID]
 
-        individual = rng.choice(model.get_population())
-        state = model.get_state(individual)
-        model.migrate(individual, _get_mutated_state(mutation_id, state), time)
+        individual = self.draw_individual(model, rng)
+        model.migrate(individual, _get_mutated_state(mutation_id, self.state), time)
 
-        rate_scalers: dict[TargetType, float] = {
+        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()
         }
@@ -83,22 +67,24 @@ class Mutation:
         for event in [
             deepcopy(e)
             for e in events
-            if _get_mutation(state) == _get_mutation(e.state)
+            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):
                 event.target_state = _get_mutated_state(mutation_id, event.target_state)
-            elif not isinstance(event, Death | Sampling | SamplingWithContactTracing):
+            elif not isinstance(
+                event, Mutation | Death | Sampling | SamplingWithContactTracing
+            ):
                 raise ValueError(
                     f"Mutation not implemented for event of type {type(event)}."
                 )
 
-            target_type = EVENT_TARGET_TYPES[type(event)]
-            if target_type in rate_scalers:
-                event.rate *= rate_scalers[target_type]
-                metadata[f"{MUTATION_PREFIX}{mutation_id}.{target_type}.rate.value"] = (
+            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)
@@ -109,4 +95,4 @@ class Mutation:
         return metadata
 
     def __repr__(self) -> str:
-        return f"Mutation(rate={self.rate}, rate_scalers={self.rate_scalers})"
+        return f"Mutation(state={self.state}, rate={self.rate}, rate_scalers={self.rate_scalers})"

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

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

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

@@ -14,12 +14,10 @@ from phylogenie.tree import Tree
 from phylogenie.treesimulator.events import Event
 from phylogenie.treesimulator.features import Feature, set_features
 from phylogenie.treesimulator.model import Model
-from phylogenie.treesimulator.mutations import Mutation
 
 
 def simulate_tree(
     events: Sequence[Event],
-    mutations: Sequence[Mutation] | None = None,
     min_tips: int = 1,
     max_tips: int | None = None,
     max_time: float = np.inf,
@@ -43,9 +41,6 @@ def simulate_tree(
     elif init_state not in states:
         raise ValueError(f"Init state {init_state} not found in event states: {states}")
 
-    if mutations is None:
-        mutations = []
-
     rng = default_rng(seed)
     start_clock = time.perf_counter()
     while True:
@@ -72,7 +67,7 @@ def simulate_tree(
             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 (
@@ -84,7 +79,6 @@ def simulate_tree(
             ):
                 break
 
-            rates.extend(m.rate for m in mutations)
             time_step = rng.exponential(1 / sum(rates))
             if current_time + time_step >= next_change_time:
                 current_time = next_change_time
@@ -95,13 +89,11 @@ def simulate_tree(
                 break
             current_time += time_step
 
-            targets = run_events + list(mutations)
-            target_idx = np.searchsorted(np.cumsum(rates) / sum(rates), rng.random())
-            target = targets[int(target_idx)]
-            if isinstance(target, Event):
-                target.apply(model, current_time, rng)
-            else:
-                metadata.update(target.apply(model, run_events, current_time, rng))
+            event_idx = np.searchsorted(np.cumsum(rates) / sum(rates), rng.random())
+            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:

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

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