phylogenie 2.1.0__py3-none-any.whl → 2.1.2__py3-none-any.whl

phylogenie/__init__.py

@@ -10,7 +10,7 @@ from phylogenie.generators import (
     FBDTreeDatasetGenerator,
     TreeDatasetGeneratorConfig,
 )
-from phylogenie.io import load_fasta, load_newick
+from phylogenie.io import dump_newick, load_fasta, load_newick
 from phylogenie.msa import MSA
 from phylogenie.skyline import (
     SkylineMatrix,
@@ -26,12 +26,21 @@ from phylogenie.skyline import (
 )
 from phylogenie.tree import Tree
 from phylogenie.treesimulator import (
+    Birth,
+    BirthWithContactTracing,
+    Death,
     Event,
+    Migration,
+    Mutation,
+    MutationTargetType,
+    Sampling,
+    SamplingWithContactTracing,
     generate_trees,
     get_BD_events,
     get_BDEI_events,
     get_BDSS_events,
     get_canonical_events,
+    get_contact_tracing_events,
     get_epidemiological_events,
     get_FBD_events,
     simulate_tree,
@@ -59,15 +68,25 @@ __all__ = [
     "skyline_vector",
     "Tree",
     "TreeDatasetGeneratorConfig",
+    "Birth",
+    "BirthWithContactTracing",
+    "Death",
     "Event",
+    "Migration",
+    "Mutation",
+    "MutationTargetType",
+    "Sampling",
+    "SamplingWithContactTracing",
     "get_BD_events",
     "get_BDEI_events",
     "get_BDSS_events",
     "get_canonical_events",
+    "get_contact_tracing_events",
     "get_epidemiological_events",
     "get_FBD_events",
     "generate_trees",
     "simulate_tree",
+    "dump_newick",
     "load_fasta",
     "load_newick",
     "MSA",

phylogenie/generators/configs.py

@@ -1,12 +1,6 @@
-from pydantic import BaseModel, ConfigDict
-
 import phylogenie.typings as pgt
-
-
-class Distribution(BaseModel):
-    type: str
-    model_config = ConfigDict(extra="allow")
-
+from phylogenie.treesimulator import MutationTargetType
+from phylogenie.utils import Distribution, StrictBaseModel
 
 Integer = str | int
 Scalar = str | pgt.Scalar
@@ -15,10 +9,6 @@ OneOrManyScalars = Scalar | pgt.Many[Scalar]
 OneOrMany2DScalars = Scalar | pgt.Many2D[Scalar]
 
 
-class StrictBaseModel(BaseModel):
-    model_config = ConfigDict(extra="forbid")
-
-
 class SkylineParameterModel(StrictBaseModel):
     value: ManyScalars
     change_times: ManyScalars
@@ -37,3 +27,12 @@ class SkylineMatrixModel(StrictBaseModel):
 SkylineParameter = Scalar | SkylineParameterModel
 SkylineVector = str | pgt.Scalar | pgt.Many[SkylineParameter] | SkylineVectorModel
 SkylineMatrix = str | pgt.Scalar | pgt.Many[SkylineVector] | SkylineMatrixModel | None
+
+
+class Event(StrictBaseModel):
+    states: str | list[str] | None = None
+    rate: SkylineParameter
+
+
+class Mutation(Event):
+    rate_scalers: dict[MutationTargetType, Distribution]

phylogenie/generators/dataset.py

@@ -10,8 +10,8 @@ import pandas as pd
 from numpy.random import Generator, default_rng
 from tqdm import tqdm
 
-import phylogenie.generators.configs as cfg
-from phylogenie.generators.factories import eval_expression
+from phylogenie.generators.factories import distribution
+from phylogenie.utils import Distribution, StrictBaseModel
 
 
 class DataType(str, Enum):
@@ -23,12 +23,12 @@ DATA_DIRNAME = "data"
 METADATA_FILENAME = "metadata.csv"
 
 
-class DatasetGenerator(ABC, cfg.StrictBaseModel):
+class DatasetGenerator(ABC, StrictBaseModel):
     output_dir: str = "phylogenie-outputs"
     n_samples: int | dict[str, int] = 1
     n_jobs: int = -1
     seed: int | None = None
-    context: dict[str, cfg.Distribution] | None = None
+    context: dict[str, Distribution] | None = None
 
     @abstractmethod
     def _generate_one(
@@ -56,11 +56,8 @@ class DatasetGenerator(ABC, cfg.StrictBaseModel):
         data: list[dict[str, Any]] = [{} for _ in range(n_samples)]
         if self.context is not None:
             for d, (k, v) in product(data, self.context.items()):
-                args = v.model_extra if v.model_extra is not None else {}
-                for arg_name, arg_value in args.items():
-                    if isinstance(arg_value, str):
-                        args[arg_name] = eval_expression(arg_value, d)
-                d[k] = np.array(getattr(rng, v.type)(**args)).tolist()
+                dist = distribution(v, d)
+                d[k] = np.array(getattr(rng, dist.type)(**dist.args)).tolist()
             df = pd.DataFrame([{"file_id": str(i), **d} for i, d in enumerate(data)])
             df.to_csv(os.path.join(output_dir, METADATA_FILENAME), index=False)
 

phylogenie/generators/factories.py

@@ -14,9 +14,10 @@ from phylogenie.skyline import (
     SkylineVector,
     SkylineVectorCoercible,
 )
+from phylogenie.utils import Distribution
 
 
-def eval_expression(expression: str, data: dict[str, Any]) -> Any:
+def _eval_expression(expression: str, data: dict[str, Any]) -> Any:
     return np.array(
         eval(
             expression,
@@ -29,9 +30,17 @@ def eval_expression(expression: str, data: dict[str, Any]) -> Any:
     ).tolist()
 
 
+def distribution(x: Distribution, data: dict[str, Any]) -> Distribution:
+    args = x.args
+    for arg_name, arg_value in args.items():
+        if isinstance(arg_value, str):
+            args[arg_name] = _eval_expression(arg_value, data)
+    return Distribution(type=x.type, **args)
+
+
 def integer(x: cfg.Integer, data: dict[str, Any]) -> int:
     if isinstance(x, str):
-        e = eval_expression(x, data)
+        e = _eval_expression(x, data)
         if isinstance(e, int):
             return e
         raise ValueError(
@@ -42,7 +51,7 @@ def integer(x: cfg.Integer, data: dict[str, Any]) -> int:
 
 def scalar(x: cfg.Scalar, data: dict[str, Any]) -> pgt.Scalar:
     if isinstance(x, str):
-        e = eval_expression(x, data)
+        e = _eval_expression(x, data)
         if isinstance(e, pgt.Scalar):
             return e
         raise ValueError(
@@ -53,7 +62,7 @@ def scalar(x: cfg.Scalar, data: dict[str, Any]) -> pgt.Scalar:
 
 def many_scalars(x: cfg.ManyScalars, data: dict[str, Any]) -> pgt.ManyScalars:
     if isinstance(x, str):
-        e = eval_expression(x, data)
+        e = _eval_expression(x, data)
         if tg.is_many_scalars(e):
             return e
         raise ValueError(
@@ -66,7 +75,7 @@ def one_or_many_scalars(
     x: cfg.OneOrManyScalars, data: dict[str, Any]
 ) -> pgt.OneOrManyScalars:
     if isinstance(x, str):
-        e = eval_expression(x, data)
+        e = _eval_expression(x, data)
         if tg.is_one_or_many_scalars(e):
             return e
         raise ValueError(
@@ -92,7 +101,7 @@ def skyline_vector(
     x: cfg.SkylineVector, data: dict[str, Any]
 ) -> SkylineVectorCoercible:
     if isinstance(x, str):
-        e = eval_expression(x, data)
+        e = _eval_expression(x, data)
         if tg.is_one_or_many_scalars(e):
             return e
         raise ValueError(
@@ -107,7 +116,7 @@ def skyline_vector(
 
     change_times = many_scalars(x.change_times, data)
     if isinstance(x.value, str):
-        e = eval_expression(x.value, data)
+        e = _eval_expression(x.value, data)
         if tg.is_many_one_or_many_scalars(e):
             value = e
         else:
@@ -135,7 +144,7 @@ def one_or_many_2D_scalars(
     x: cfg.OneOrMany2DScalars, data: dict[str, Any]
 ) -> pgt.OneOrMany2DScalars:
     if isinstance(x, str):
-        e = eval_expression(x, data)
+        e = _eval_expression(x, data)
         if tg.is_one_or_many_2D_scalars(e):
             return e
         raise ValueError(
@@ -153,7 +162,7 @@ def skyline_matrix(
         return None
 
     if isinstance(x, str):
-        e = eval_expression(x, data)
+        e = _eval_expression(x, data)
         if tg.is_one_or_many_2D_scalars(e):
             return e
         raise ValueError(
@@ -168,7 +177,7 @@ def skyline_matrix(
 
     change_times = many_scalars(x.change_times, data)
     if isinstance(x.value, str):
-        e = eval_expression(x.value, data)
+        e = _eval_expression(x.value, data)
         if tg.is_many_one_or_many_2D_scalars(e):
             value = e
         else:

phylogenie/generators/trees.py

@@ -9,6 +9,7 @@ from pydantic import Field
 import phylogenie.generators.configs as cfg
 from phylogenie.generators.dataset import DatasetGenerator, DataType
 from phylogenie.generators.factories import (
+    distribution,
     integer,
     scalar,
     skyline_matrix,
@@ -19,6 +20,7 @@ from phylogenie.io import dump_newick
 from phylogenie.tree import Tree
 from phylogenie.treesimulator import (
     Event,
+    Mutation,
     get_BD_events,
     get_BDEI_events,
     get_BDSS_events,
@@ -46,7 +48,6 @@ class TreeDatasetGenerator(DatasetGenerator):
     max_time: cfg.Scalar = np.inf
     init_state: str | None = None
     sampling_probability_at_present: cfg.Scalar = 0.0
-    max_tries: int | None = None
 
     @abstractmethod
     def _get_events(self, data: dict[str, Any]) -> list[Event]: ...
@@ -67,7 +68,6 @@ class TreeDatasetGenerator(DatasetGenerator):
             sampling_probability_at_present=scalar(
                 self.sampling_probability_at_present, data
             ),
-            max_tries=self.max_tries,
             seed=int(rng.integers(2**32)),
         )
 
@@ -127,11 +127,12 @@ class FBDTreeDatasetGenerator(TreeDatasetGenerator):
         )
 
 
-class TreeDatasetGeneratorWithContactTracing(TreeDatasetGenerator):
+class TreeDatasetGeneratorForEpidemiology(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]: ...
@@ -150,10 +151,30 @@ class TreeDatasetGeneratorWithContactTracing(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(TreeDatasetGeneratorWithContactTracing):
+class EpidemiologicalTreeDatasetGenerator(TreeDatasetGeneratorForEpidemiology):
     parameterization: Literal[ParameterizationType.EPIDEMIOLOGICAL] = (
         ParameterizationType.EPIDEMIOLOGICAL
     )
@@ -179,7 +200,7 @@ class EpidemiologicalTreeDatasetGenerator(TreeDatasetGeneratorWithContactTracing
         )
 
 
-class BDTreeDatasetGenerator(TreeDatasetGeneratorWithContactTracing):
+class BDTreeDatasetGenerator(TreeDatasetGeneratorForEpidemiology):
     parameterization: Literal[ParameterizationType.BD] = ParameterizationType.BD
     reproduction_number: cfg.SkylineParameter
     infectious_period: cfg.SkylineParameter
@@ -193,7 +214,7 @@ class BDTreeDatasetGenerator(TreeDatasetGeneratorWithContactTracing):
         )
 
 
-class BDEITreeDatasetGenerator(TreeDatasetGeneratorWithContactTracing):
+class BDEITreeDatasetGenerator(TreeDatasetGeneratorForEpidemiology):
     parameterization: Literal[ParameterizationType.BDEI] = ParameterizationType.BDEI
     reproduction_number: cfg.SkylineParameter
     infectious_period: cfg.SkylineParameter
@@ -209,7 +230,7 @@ class BDEITreeDatasetGenerator(TreeDatasetGeneratorWithContactTracing):
         )
 
 
-class BDSSTreeDatasetGenerator(TreeDatasetGeneratorWithContactTracing):
+class BDSSTreeDatasetGenerator(TreeDatasetGeneratorForEpidemiology):
     parameterization: Literal[ParameterizationType.BDSS] = ParameterizationType.BDSS
     reproduction_number: cfg.SkylineParameter
     infectious_period: cfg.SkylineParameter

phylogenie/io.py

@@ -24,21 +24,28 @@ def _parse_newick(newick: str) -> Tree:
             stack.append(current_nodes)
             current_nodes = []
         else:
-            id = _parse_chars([":", ",", ")", ";"])
+            id = _parse_chars([":", ",", ")", ";", "["])
             branch_length = None
             if newick[i] == ":":
                 i += 1
-                branch_length = _parse_chars([",", ")", ";"])
+                branch_length = float(_parse_chars([",", ")", ";", "["]))
 
-            current_node = Tree(
-                id=id,
-                branch_length=(None if branch_length is None else float(branch_length)),
-            )
+            current_node = Tree(id, branch_length)
             for node in current_children:
                 current_node.add_child(node)
                 current_children = []
             current_nodes.append(current_node)
 
+            if newick[i] == "[":
+                i += 1
+                features = _parse_chars(["]"]).split(":")
+                i += 1
+                if features[0] != "&&NHX":
+                    raise ValueError(f"Expected '&&NHX' for node features.")
+                for feature in features[1:]:
+                    key, value = feature.split("=", 1)
+                    current_node.set(key, eval(value))
+
             if newick[i] == ")":
                 current_children = current_nodes
                 current_nodes = stack.pop()
@@ -47,7 +54,7 @@ def _parse_newick(newick: str) -> Tree:
 
         i += 1
 
-    raise ValueError("Newick string does not end with a semicolon.")
+    raise ValueError("Newick string is invalid.")
 
 
 def load_newick(filepath: str) -> Tree | list[Tree]:
@@ -63,6 +70,10 @@ def _to_newick(tree: Tree) -> str:
         newick = f"({children_newick}){newick}"
     if tree.branch_length is not None:
         newick += f":{tree.branch_length}"
+    if tree.features:
+        reprs = {k: repr(v).replace("'", '"') for k, v in tree.features.items()}
+        features = [f"{k}={repr}" for k, repr in reprs.items()]
+        newick += f"[&&NHX:{':'.join(features)}]"
     return newick
 
 
@@ -83,13 +94,14 @@ def load_fasta(
             if not line.startswith(">"):
                 raise ValueError(f"Invalid FASTA format: expected '>', got '{line[0]}'")
             id = line[1:].strip()
+            time = None
             if extract_time_from_id is not None:
                 time = extract_time_from_id(id)
-            else:
+            elif "|" in id:
                 try:
                     time = float(id.split("|")[-1])
-                except:
-                    time = None
+                except ValueError:
+                    pass
             chars = next(f).strip()
             sequences.append(Sequence(id, chars, time))
     return MSA(sequences)

phylogenie/tree.py

@@ -1,18 +1,41 @@
 from collections.abc import Iterator
+from typing import Any
 
 
 class Tree:
     def __init__(self, id: str = "", branch_length: float | None = None):
         self.id = id
         self.branch_length = branch_length
-        self.parent: Tree | None = None
-        self.children: list[Tree] = []
+        self._parent: Tree | None = None
+        self._children: list[Tree] = []
+        self._features: dict[str, Any] = {}
+
+    @property
+    def children(self) -> tuple["Tree", ...]:
+        return tuple(self._children)
+
+    @property
+    def parent(self) -> "Tree | None":
+        return self._parent
+
+    @property
+    def features(self) -> dict[str, Any]:
+        return self._features
 
     def add_child(self, child: "Tree") -> "Tree":
-        child.parent = self
-        self.children.append(child)
+        child._parent = self
+        self._children.append(child)
         return self
 
+    def remove_child(self, child: "Tree") -> None:
+        self._children.remove(child)
+        child._parent = None
+
+    def set_parent(self, node: "Tree | None"):
+        self._parent = node
+        if node is not None:
+            node._children.append(self)
+
     def preorder_traversal(self) -> Iterator["Tree"]:
         yield self
         for child in self.children:
@@ -29,6 +52,9 @@ class Tree:
                 return node
         raise ValueError(f"Node with id {id} not found.")
 
+    def is_leaf(self) -> bool:
+        return not self.children
+
     def get_leaves(self) -> list["Tree"]:
         return [node for node in self if not node.children]
 
@@ -38,11 +64,19 @@ class Tree:
             raise ValueError(f"Branch length of node {self.id} is not set.")
         return self.branch_length + parent_time
 
-    def is_leaf(self) -> bool:
-        return not self.children
+    def set(self, key: str, value: Any) -> None:
+        self._features[key] = value
+
+    def get(self, key: str) -> Any:
+        return self._features.get(key)
+
+    def delete(self, key: str) -> None:
+        del self._features[key]
 
-    def copy(self) -> "Tree":
+    def copy(self):
         new_tree = Tree(self.id, self.branch_length)
+        for key, value in self._features.items():
+            new_tree.set(key, value)
         for child in self.children:
             new_tree.add_child(child.copy())
         return new_tree
@@ -51,4 +85,4 @@ class Tree:
         return self.preorder_traversal()
 
     def __repr__(self) -> str:
-        return f"TreeNode(id='{self.id}', branch_length={self.branch_length})"
+        return f"TreeNode(id='{self.id}', branch_length={self.branch_length}, features={self.features})"

phylogenie/treesimulator/__init__.py

@@ -1,5 +1,13 @@
 from phylogenie.treesimulator.events import (
+    Birth,
+    BirthWithContactTracing,
+    Death,
     Event,
+    Migration,
+    Mutation,
+    MutationTargetType,
+    Sampling,
+    SamplingWithContactTracing,
     get_BD_events,
     get_BDEI_events,
     get_BDSS_events,
@@ -11,7 +19,15 @@ from phylogenie.treesimulator.events import (
 from phylogenie.treesimulator.gillespie import generate_trees, simulate_tree
 
 __all__ = [
+    "Birth",
+    "BirthWithContactTracing",
+    "Death",
     "Event",
+    "Migration",
+    "Mutation",
+    "MutationTargetType",
+    "Sampling",
+    "SamplingWithContactTracing",
     "get_BD_events",
     "get_BDEI_events",
     "get_BDSS_events",

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