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

phylogenie/treesimulator/events/__init__.py

@@ -0,0 +1,39 @@
+from phylogenie.treesimulator.events.contact_tracing import (
+    BirthWithContactTracing,
+    SamplingWithContactTracing,
+    get_contact_tracing_events,
+)
+from phylogenie.treesimulator.events.core import (
+    Birth,
+    Death,
+    Event,
+    Migration,
+    Sampling,
+    get_BD_events,
+    get_BDEI_events,
+    get_BDSS_events,
+    get_canonical_events,
+    get_epidemiological_events,
+    get_FBD_events,
+)
+from phylogenie.treesimulator.events.mutations import Mutation
+from phylogenie.treesimulator.events.mutations import TargetType as MutationTargetType
+
+__all__ = [
+    "Birth",
+    "BirthWithContactTracing",
+    "Death",
+    "Event",
+    "Migration",
+    "Mutation",
+    "Sampling",
+    "SamplingWithContactTracing",
+    "MutationTargetType",
+    "get_BD_events",
+    "get_BDEI_events",
+    "get_BDSS_events",
+    "get_canonical_events",
+    "get_contact_tracing_events",
+    "get_epidemiological_events",
+    "get_FBD_events",
+]

phylogenie/treesimulator/events/contact_tracing.py

@@ -0,0 +1,125 @@
+from collections import defaultdict
+from collections.abc import Sequence
+
+import numpy as np
+from numpy.random import Generator
+
+from phylogenie.skyline import SkylineParameterLike, skyline_parameter
+from phylogenie.treesimulator.events.core import Birth, Death, Migration, Sampling
+from phylogenie.treesimulator.model import Event, Model
+
+CT_POSTFIX = "-CT"
+CONTACTS_KEY = "CONTACTS"
+
+
+def _get_CT_state(state: str) -> str:
+    return f"{state}{CT_POSTFIX}"
+
+
+def _is_CT_state(state: str) -> bool:
+    return state.endswith(CT_POSTFIX)
+
+
+class BirthWithContactTracing(Event):
+    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:
+        individual = self.draw_individual(model, rng)
+        new_individual = model.birth_from(individual, self.child_state, time)
+        if CONTACTS_KEY not in model.context:
+            model.context[CONTACTS_KEY] = defaultdict(list)
+        model.context[CONTACTS_KEY][individual].append(new_individual)
+        model.context[CONTACTS_KEY][new_individual].append(individual)
+
+    def __repr__(self) -> str:
+        return f"BirthWithContactTracing(state={self.state}, rate={self.rate}, child_state={self.child_state})"
+
+
+class SamplingWithContactTracing(Event):
+    def __init__(
+        self,
+        state: str,
+        rate: SkylineParameterLike,
+        max_notified_contacts: int,
+        notification_probability: SkylineParameterLike,
+    ):
+        super().__init__(state, rate)
+        self.max_notified_contacts = max_notified_contacts
+        self.notification_probability = skyline_parameter(notification_probability)
+
+    def apply(self, model: Model, time: float, rng: Generator) -> None:
+        individual = self.draw_individual(model, rng)
+        model.sample(individual, time, True)
+        population = model.get_population()
+        if CONTACTS_KEY not in model.context:
+            return
+        contacts = model.context[CONTACTS_KEY][individual]
+        for contact in contacts[-self.max_notified_contacts :]:
+            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)
+
+    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})"
+
+
+def get_contact_tracing_events(
+    events: Sequence[Event],
+    max_notified_contacts: int = 1,
+    notification_probability: SkylineParameterLike = 1,
+    sampling_rate_after_notification: SkylineParameterLike = np.inf,
+    samplable_states_after_notification: list[str] | None = None,
+) -> list[Event]:
+    ct_events: list[Event] = []
+    notification_probability = skyline_parameter(notification_probability)
+    sampling_rate_after_notification = skyline_parameter(
+        sampling_rate_after_notification
+    )
+    for event in events:
+        state, rate = event.state, event.rate
+        if isinstance(event, Migration):
+            ct_events.append(event)
+            ct_events.append(
+                Migration(_get_CT_state(state), 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)
+            )
+        elif isinstance(event, Sampling):
+            if not event.removal:
+                raise ValueError(
+                    "Contact tracing requires removal to be set for all sampling events."
+                )
+            ct_events.append(
+                SamplingWithContactTracing(
+                    state, rate, max_notified_contacts, notification_probability
+                )
+            )
+        elif isinstance(event, Death):
+            ct_events.append(event)
+        else:
+            raise NotImplementedError(
+                f"Unsupported event type {type(event)} for contact tracing."
+            )
+
+    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(
+            SamplingWithContactTracing(
+                _get_CT_state(state),
+                sampling_rate_after_notification,
+                max_notified_contacts,
+                notification_probability,
+            )
+        )
+
+    return ct_events

phylogenie/treesimulator/{events.py → events/core.py}

@@ -1,80 +1,73 @@
-from abc import ABC, abstractmethod
-from collections.abc import Sequence
-from dataclasses import dataclass
-
-import numpy as np
+from numpy.random import Generator
 
 from phylogenie.skyline import (
     SkylineMatrixCoercible,
-    SkylineParameter,
     SkylineParameterLike,
     SkylineVectorCoercible,
     skyline_matrix,
-    skyline_parameter,
     skyline_vector,
 )
-from phylogenie.treesimulator.model import Model, get_CT_state
+from phylogenie.treesimulator.model import Event, Model
 
 INFECTIOUS_STATE = "I"
 EXPOSED_STATE = "E"
 SUPERSPREADER_STATE = "S"
 
 
-@dataclass
-class Event(ABC):
-    rate: SkylineParameter
-    state: str
+class Birth(Event):
+    def __init__(self, state: str, rate: SkylineParameterLike, child_state: str):
+        super().__init__(state, rate)
+        self.child_state = child_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
+    def apply(self, model: Model, time: float, rng: Generator) -> None:
+        individual = self.draw_individual(model, rng)
+        model.birth_from(individual, self.child_state, time)
 
-    @abstractmethod
-    def apply(self, model: Model, time: float) -> None: ...
+    def __repr__(self) -> str:
+        return f"Birth(state={self.state}, rate={self.rate}, child_state={self.child_state})"
 
 
-@dataclass
-class BirthEvent(Event):
-    child_state: str
+class Death(Event):
+    def apply(self, model: Model, time: float, rng: Generator) -> None:
+        individual = self.draw_individual(model, rng)
+        model.remove(individual, time)
 
-    def apply(self, model: Model, time: float) -> None:
-        individual = model.get_random_individual(self.state)
-        model.birth_from(individual, self.child_state, time)
+    def __repr__(self) -> str:
+        return f"Death(state={self.state}, rate={self.rate})"
 
 
-class DeathEvent(Event):
-    def apply(self, model: Model, time: float) -> None:
-        individual = model.get_random_individual(self.state)
-        model.remove(individual, time)
+class Migration(Event):
+    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:
+        individual = self.draw_individual(model, rng)
+        model.migrate(individual, self.target_state, time)
 
-@dataclass
-class MigrationEvent(Event):
-    target_state: str
+    def __repr__(self) -> str:
+        return f"Migration(state={self.state}, rate={self.rate}, target_state={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)
 
+class Sampling(Event):
+    def __init__(self, state: str, rate: SkylineParameterLike, removal: bool):
+        super().__init__(state, rate)
+        self.removal = removal
 
-@dataclass
-class SamplingEvent(Event):
-    removal_probability: SkylineParameter
+    def apply(self, model: Model, time: float, rng: Generator) -> None:
+        individual = self.draw_individual(model, rng)
+        model.sample(individual, time, self.removal)
 
-    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 __repr__(self) -> str:
+        return f"Sampling(state={self.state}, rate={self.rate}, removal={self.removal})"
 
 
 def get_canonical_events(
-    states: Sequence[str],
+    states: list[str],
     sampling_rates: SkylineVectorCoercible,
+    remove_after_sampling: bool,
     birth_rates: SkylineVectorCoercible = 0,
     death_rates: SkylineVectorCoercible = 0,
-    removal_probabilities: SkylineVectorCoercible = 0,
     migration_rates: SkylineMatrixCoercible | None = None,
     birth_rates_among_states: SkylineMatrixCoercible | None = None,
 ) -> list[Event]:
@@ -83,102 +76,94 @@ def get_canonical_events(
     birth_rates = skyline_vector(birth_rates, N)
     death_rates = skyline_vector(death_rates, N)
     sampling_rates = skyline_vector(sampling_rates, N)
-    removal_probabilities = skyline_vector(removal_probabilities, N)
 
     events: list[Event] = []
     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]))
+        events.append(Birth(state, birth_rates[i], state))
+        events.append(Death(state, death_rates[i]))
+        events.append(Sampling(state, sampling_rates[i], remove_after_sampling))
 
     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))
+                events.append(Migration(state, migration_rates[i, j], other_state))
 
     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)
-                )
+                events.append(Birth(state, birth_rates_among_states[i, j], other_state))
 
     return [event for event in events if event.rate]
 
 
-def get_epidemiological_events(
-    states: Sequence[str],
+def get_FBD_events(
+    states: list[str],
     sampling_proportions: SkylineVectorCoercible = 1,
-    reproduction_numbers: SkylineVectorCoercible = 0,
-    become_uninfectious_rates: SkylineVectorCoercible = 0,
-    removal_probabilities: SkylineVectorCoercible = 1,
+    diversification: SkylineVectorCoercible = 0,
+    turnover: SkylineVectorCoercible = 0,
     migration_rates: SkylineMatrixCoercible | None = None,
-    reproduction_numbers_among_states: SkylineMatrixCoercible | None = None,
+    diversification_between_states: SkylineMatrixCoercible | None = None,
 ) -> list[Event]:
     N = len(states)
 
-    reproduction_numbers = skyline_vector(reproduction_numbers, N)
-    become_uninfectious_rates = skyline_vector(become_uninfectious_rates, N)
+    diversification = skyline_vector(diversification, N)
+    turnover = skyline_vector(turnover, N)
     sampling_proportions = skyline_vector(sampling_proportions, N)
-    removal_probabilities = skyline_vector(removal_probabilities, N)
 
-    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 = diversification / (1 - turnover)
+    death_rates = turnover * birth_rates
+    sampling_rates = sampling_proportions * death_rates
     birth_rates_among_states = (
-        (
-            skyline_matrix(reproduction_numbers_among_states, N, N - 1)
-            * become_uninfectious_rates
-        )
-        if reproduction_numbers_among_states is not None
+        (skyline_matrix(diversification_between_states, N, N - 1) + death_rates)
+        if diversification_between_states is not None
         else None
     )
 
     return get_canonical_events(
         states=states,
+        sampling_rates=sampling_rates,
+        remove_after_sampling=False,
         birth_rates=birth_rates,
         death_rates=death_rates,
-        sampling_rates=sampling_rates,
-        removal_probabilities=removal_probabilities,
         migration_rates=migration_rates,
         birth_rates_among_states=birth_rates_among_states,
     )
 
 
-def get_FBD_events(
-    states: Sequence[str],
-    diversification: SkylineVectorCoercible = 0,
-    turnover: SkylineVectorCoercible = 0,
+def get_epidemiological_events(
+    states: list[str],
     sampling_proportions: SkylineVectorCoercible = 1,
-    removal_probabilities: SkylineVectorCoercible = 0,
+    reproduction_numbers: SkylineVectorCoercible = 0,
+    become_uninfectious_rates: SkylineVectorCoercible = 0,
     migration_rates: SkylineMatrixCoercible | None = None,
-    diversification_between_types: SkylineMatrixCoercible | None = None,
+    reproduction_numbers_among_states: SkylineMatrixCoercible | None = None,
 ) -> list[Event]:
     N = len(states)
 
-    diversification = skyline_vector(diversification, N)
-    turnover = skyline_vector(turnover, N)
+    reproduction_numbers = skyline_vector(reproduction_numbers, N)
+    become_uninfectious_rates = skyline_vector(become_uninfectious_rates, N)
     sampling_proportions = skyline_vector(sampling_proportions, N)
-    removal_probabilities = skyline_vector(removal_probabilities, N)
 
-    birth_rates = diversification / (1 - turnover)
-    death_rates = turnover * birth_rates
-    sampling_rates_dividend = 1 - removal_probabilities * sampling_proportions
-    sampling_rates = sampling_proportions * death_rates / sampling_rates_dividend
+    birth_rates = reproduction_numbers * become_uninfectious_rates
+    sampling_rates = become_uninfectious_rates * sampling_proportions
+    death_rates = become_uninfectious_rates - sampling_rates
     birth_rates_among_states = (
-        (skyline_matrix(diversification_between_types, N, N - 1) + death_rates)
-        if diversification_between_types is not None
+        (
+            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,
+        sampling_rates=sampling_rates,
+        remove_after_sampling=True,
         birth_rates=birth_rates,
         death_rates=death_rates,
-        sampling_rates=sampling_rates,
-        removal_probabilities=removal_probabilities,
         migration_rates=migration_rates,
         birth_rates_among_states=birth_rates_among_states,
     )
@@ -232,40 +217,3 @@ 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/events/mutations.py

@@ -0,0 +1,105 @@
+from copy import deepcopy
+from enum import Enum
+from typing import Type
+
+from numpy.random import Generator
+
+from phylogenie.skyline import SkylineParameterLike
+from phylogenie.treesimulator.events.contact_tracing import (
+    BirthWithContactTracing,
+    SamplingWithContactTracing,
+)
+from phylogenie.treesimulator.events.core import (
+    Birth,
+    Death,
+    Event,
+    Migration,
+    Sampling,
+)
+from phylogenie.treesimulator.model import Model
+from phylogenie.utils import Distribution
+
+MUTATION_PREFIX = "MUT-"
+MUTATIONS_KEY = "MUTATIONS"
+
+
+def _get_mutation(state: str) -> str | None:
+    return state.split(".")[0] if state.startswith(MUTATION_PREFIX) else None
+
+
+def _get_mutated_state(mutation_id: int, state: str) -> str:
+    if state.startswith(MUTATION_PREFIX):
+        state = state.split(".")[1]
+    return f"{MUTATION_PREFIX}{mutation_id}.{state}"
+
+
+class TargetType(str, Enum):
+    BIRTH = "birth"
+    DEATH = "death"
+    MIGRATION = "migration"
+    SAMPLING = "sampling"
+    MUTATION = "mutation"
+
+
+class Mutation(Event):
+    def __init__(
+        self,
+        state: str,
+        rate: SkylineParameterLike,
+        rate_scalers: dict[TargetType, 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]
+
+        individual = self.draw_individual(model, rng)
+        model.migrate(individual, _get_mutated_state(mutation_id, self.state), time)
+
+        rate_scalers = {
+            target_type: getattr(rng, rate_scaler.type)(**rate_scaler.args)
+            for target_type, rate_scaler in self.rate_scalers.items()
+        }
+
+        for event in [
+            deepcopy(e)
+            for e in model.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):
+                event.target_state = _get_mutated_state(mutation_id, event.target_state)
+            elif not isinstance(
+                event, Mutation | Death | Sampling | SamplingWithContactTracing
+            ):
+                raise ValueError(
+                    f"Mutation not defined 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)
+
+    def __repr__(self) -> str:
+        return f"Mutation(state={self.state}, rate={self.rate}, rate_scalers={self.rate_scalers})"
+
+
+TARGETS: dict[TargetType, tuple[Type[Event], ...]] = {
+    TargetType.BIRTH: (Birth, BirthWithContactTracing),
+    TargetType.DEATH: (Death,),
+    TargetType.MIGRATION: (Migration,),
+    TargetType.SAMPLING: (Sampling, SamplingWithContactTracing),
+    TargetType.MUTATION: (Mutation,),
+}