phylogenie 2.1.4__py3-none-any.whl → 3.1.7__py3-none-any.whl

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/treesimulator/__init__.py

@@ -3,9 +3,9 @@ from phylogenie.treesimulator.events import (
     BirthWithContactTracing,
     Death,
     Event,
+    EventType,
     Migration,
     Mutation,
-    MutationTargetType,
     Sampling,
     SamplingWithContactTracing,
     get_BD_events,
@@ -15,17 +15,33 @@ 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.io import dump_newick, load_newick, load_nexus
+from phylogenie.treesimulator.model import get_node_state
+from phylogenie.treesimulator.tree import Tree
+from phylogenie.treesimulator.utils import (
+    compute_mean_leaf_pairwise_distance,
+    compute_sackin_index,
+    get_distance,
+    get_mrca,
+    get_node_depth_levels,
+    get_node_depths,
+    get_node_height_levels,
+    get_node_heights,
+    get_node_leaf_counts,
+)
 
 __all__ = [
     "Birth",
     "BirthWithContactTracing",
     "Death",
     "Event",
+    "EventType",
     "Migration",
     "Mutation",
-    "MutationTargetType",
     "Sampling",
     "SamplingWithContactTracing",
     "get_BD_events",
@@ -35,6 +51,23 @@ __all__ = [
     "get_contact_tracing_events",
     "get_epidemiological_events",
     "get_FBD_events",
-    "generate_trees",
+    "get_mutation_id",
+    "Feature",
+    "set_features",
     "simulate_tree",
+    "dump_newick",
+    "load_newick",
+    "load_nexus",
+    "generate_trees",
+    "get_node_state",
+    "Tree",
+    "compute_mean_leaf_pairwise_distance",
+    "compute_sackin_index",
+    "get_distance",
+    "get_mrca",
+    "get_node_depth_levels",
+    "get_node_depths",
+    "get_node_height_levels",
+    "get_node_heights",
+    "get_node_leaf_counts",
 ]

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,19 +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 Mutation, get_mutation_id
 
 __all__ = [
     "Birth",
     "BirthWithContactTracing",
     "Death",
     "Event",
+    "EventType",
     "Migration",
-    "Mutation",
     "Sampling",
     "SamplingWithContactTracing",
-    "MutationTargetType",
     "get_BD_events",
     "get_BDEI_events",
     "get_BDSS_events",
@@ -36,4 +34,6 @@ __all__ = [
     "get_contact_tracing_events",
     "get_epidemiological_events",
     "get_FBD_events",
+    "Mutation",
+    "get_mutation_id",
 ]

phylogenie/treesimulator/events/base.py

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

phylogenie/treesimulator/events/contact_tracing.py

@@ -1,43 +1,48 @@
 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
+
     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:
-            model.context[CONTACTS_KEY] = defaultdict(list)
-        model.context[CONTACTS_KEY][individual].append(new_individual)
-        model.context[CONTACTS_KEY][new_individual].append(individual)
+        if CONTACTS_KEY not in model.metadata:
+            model[CONTACTS_KEY] = defaultdict(list)
+        model[CONTACTS_KEY][individual].append(new_individual)
+        model[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):
+    type = EventType.SAMPLING
+
     def __init__(
         self,
         state: str,
@@ -49,19 +54,19 @@ 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()
-        if CONTACTS_KEY not in model.context:
+        if CONTACTS_KEY not in model.metadata:
             return
-        contacts = model.context[CONTACTS_KEY][individual]
+        contacts = model[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)
+                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})"
@@ -70,8 +75,8 @@ class SamplingWithContactTracing(Event):
 def get_contact_tracing_events(
     events: Sequence[Event],
     max_notified_contacts: int = 1,
-    notification_probability: SkylineParameterLike = 1,
-    sampling_rate_after_notification: SkylineParameterLike = np.inf,
+    notification_probability: SkylineParameterLike = 0.0,
+    sampling_rate_after_notification: SkylineParameterLike = 2**32,
     samplable_states_after_notification: list[str] | None = None,
 ) -> list[Event]:
     ct_events: list[Event] = []
@@ -79,17 +84,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 +110,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 +130,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/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)
 
@@ -64,8 +73,8 @@ class Sampling(Event):
 
 def get_canonical_events(
     states: list[str],
-    sampling_rates: SkylineVectorCoercible,
-    remove_after_sampling: bool,
+    sampling_rates: SkylineVectorCoercible = 0,
+    remove_after_sampling: bool = False,
     birth_rates: SkylineVectorCoercible = 0,
     death_rates: SkylineVectorCoercible = 0,
     migration_rates: SkylineMatrixCoercible | None = None,
@@ -100,7 +109,7 @@ def get_canonical_events(
 
 def get_FBD_events(
     states: list[str],
-    sampling_proportions: SkylineVectorCoercible = 1,
+    sampling_proportions: SkylineVectorCoercible = 0,
     diversification: SkylineVectorCoercible = 0,
     turnover: SkylineVectorCoercible = 0,
     migration_rates: SkylineMatrixCoercible | None = None,
@@ -134,7 +143,7 @@ def get_FBD_events(
 
 def get_epidemiological_events(
     states: list[str],
-    sampling_proportions: SkylineVectorCoercible = 1,
+    sampling_proportions: SkylineVectorCoercible,
     reproduction_numbers: SkylineVectorCoercible = 0,
     become_uninfectious_rates: SkylineVectorCoercible = 0,
     migration_rates: SkylineMatrixCoercible | None = None,
@@ -172,7 +181,7 @@ def get_epidemiological_events(
 def get_BD_events(
     reproduction_number: SkylineParameterLike,
     infectious_period: SkylineParameterLike,
-    sampling_proportion: SkylineParameterLike = 1,
+    sampling_proportion: SkylineParameterLike,
 ) -> list[Event]:
     return get_epidemiological_events(
         states=[INFECTIOUS_STATE],
@@ -186,7 +195,7 @@ def get_BDEI_events(
     reproduction_number: SkylineParameterLike,
     infectious_period: SkylineParameterLike,
     incubation_period: SkylineParameterLike,
-    sampling_proportion: SkylineParameterLike = 1,
+    sampling_proportion: SkylineParameterLike,
 ) -> list[Event]:
     return get_epidemiological_events(
         states=[EXPOSED_STATE, INFECTIOUS_STATE],
@@ -202,7 +211,7 @@ def get_BDSS_events(
     infectious_period: SkylineParameterLike,
     superspreading_ratio: SkylineParameterLike,
     superspreaders_proportion: SkylineParameterLike,
-    sampling_proportion: SkylineParameterLike = 1,
+    sampling_proportion: SkylineParameterLike,
 ) -> list[Event]:
     f_SS = superspreaders_proportion
     r_SS = superspreading_ratio

phylogenie/treesimulator/events/mutations.py

@@ -1,26 +1,20 @@
+import re
 from copy import deepcopy
-from enum import Enum
-from typing import Type
+from typing import Any, Callable
 
 from numpy.random import Generator
 
 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
-from phylogenie.utils import Distribution
 
 MUTATION_PREFIX = "MUT-"
-MUTATIONS_KEY = "MUTATIONS"
+NEXT_MUTATION_ID = "NEXT_MUTATION_ID"
 
 
 def _get_mutation(state: str) -> str | None:
@@ -29,77 +23,83 @@ def _get_mutation(state: str) -> str | None:
 
 def _get_mutated_state(mutation_id: int, state: str) -> str:
     if state.startswith(MUTATION_PREFIX):
-        state = state.split(".")[1]
+        _, state = state.split(".")
     return f"{MUTATION_PREFIX}{mutation_id}.{state}"
 
 
-class TargetType(str, Enum):
-    BIRTH = "birth"
-    DEATH = "death"
-    MIGRATION = "migration"
-    SAMPLING = "sampling"
-    MUTATION = "mutation"
+def get_mutation_id(node_name: str) -> int:
+    match = re.search(rf"{MUTATION_PREFIX}(\d+)\.", node_name)
+    if match:
+        return int(match.group(1))
+    return 0
 
 
 class Mutation(Event):
+    type = EventType.MUTATION
+
     def __init__(
         self,
         state: str,
         rate: SkylineParameterLike,
-        rate_scalers: dict[TargetType, Distribution],
+        rate_scalers: dict[EventType, Callable[[], float]],
+        rates_to_log: list[EventType] | None = None,
     ):
         super().__init__(state, rate)
         self.rate_scalers = rate_scalers
+        self.rates_to_log = [] if rates_to_log is None else rates_to_log
 
-    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.metadata:
+            model[NEXT_MUTATION_ID] = 0
+        model[NEXT_MUTATION_ID] += 1
+        mutation_id = model[NEXT_MUTATION_ID]
 
         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)
+        rate_scalers: dict[EventType, float] = {
+            target_type: rate_scaler()
             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)
+                metadata_key = f"birth_rate_from_{event.state}_to_{event.child_state}"
             elif isinstance(event, Migration):
                 event.target_state = _get_mutated_state(mutation_id, event.target_state)
-            elif not isinstance(
+                metadata_key = (
+                    f"migration_rate_from_{event.state}_to_{event.target_state}"
+                )
+            elif isinstance(
                 event, Mutation | Death | Sampling | SamplingWithContactTracing
             ):
+                metadata_key = f"{event.type}_rate_for_{event.state}"
+            else:
                 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
+            event.rate *= rate_scalers.get(event.type, 1)
+            if event.type in self.rates_to_log:
+                metadata[metadata_key] = (
+                    event.rate.value[0]
+                    if len(event.rate.value) == 1
+                    else list(event.rate.value)
+                )
+
+            events.append(event)
 
-            model.add_event(event)
+        return metadata
 
     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,),
-}

phylogenie/treesimulator/features.py

@@ -0,0 +1,49 @@
+from collections.abc import Iterable
+from enum import Enum
+
+from phylogenie.treesimulator.events.mutations import get_mutation_id
+from phylogenie.treesimulator.model import get_node_state
+from phylogenie.treesimulator.tree import Tree
+from phylogenie.treesimulator.utils import (
+    get_node_depth_levels,
+    get_node_depths,
+    get_node_height_levels,
+    get_node_heights,
+    get_node_leaf_counts,
+)
+
+
+def _get_states(tree: Tree) -> dict[Tree, str]:
+    return {node: get_node_state(node.name) for node in tree}
+
+
+def _get_mutations(tree: Tree) -> dict[Tree, int]:
+    return {node: get_mutation_id(node.name) for node in tree}
+
+
+class Feature(str, Enum):
+    DEPTH = "depth"
+    DEPTH_LEVEL = "depth_level"
+    HEIGHT = "height"
+    HEIGHT_LEVEL = "height_level"
+    MUTATION = "mutation"
+    N_LEAVES = "n_leaves"
+    STATE = "state"
+
+
+FEATURES_EXTRACTORS = {
+    Feature.DEPTH: get_node_depths,
+    Feature.DEPTH_LEVEL: get_node_depth_levels,
+    Feature.HEIGHT: get_node_heights,
+    Feature.HEIGHT_LEVEL: get_node_height_levels,
+    Feature.MUTATION: _get_mutations,
+    Feature.N_LEAVES: get_node_leaf_counts,
+    Feature.STATE: _get_states,
+}
+
+
+def set_features(tree: Tree, features: Iterable[Feature]) -> None:
+    for feature in features:
+        feature_maps = FEATURES_EXTRACTORS[feature](tree)
+        for node in tree:
+            node[feature.value] = feature_maps[node]