phyloom 0.0.0__py3-none-any.whl

phyloom/configs/__init__.py

@@ -0,0 +1,15 @@
+from phyloom.configs.events import EventCfg, EventRuleCfgBase, ScheduledEventCfg
+from phyloom.configs.generator import GENERATOR_REGISTRY, GeneratorBase, GeneratorCfg
+from phyloom.configs.hazards import HAZARD_REGISTRY, HazardCfg, HazardCfgBase
+
+__all__ = [
+    "GENERATOR_REGISTRY",
+    "HAZARD_REGISTRY",
+    "EventCfg",
+    "EventRuleCfgBase",
+    "GeneratorBase",
+    "GeneratorCfg",
+    "HazardCfg",
+    "HazardCfgBase",
+    "ScheduledEventCfg",
+]

phyloom/configs/events.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any, Generic, TypeVar
+
+from pydantic import BaseModel, ConfigDict
+
+from phyloom.configs.hazards import HazardCfg
+from phyloom.configs.roots import FloatListCfg, ModelActionCfg
+from phyloom.simulator.core import Event, ModelT_contra, ScheduledEvent
+
+if TYPE_CHECKING:
+    from phyloom.simulator.core import EventRule, ModelT
+    from phyloom.types import Context
+
+RuleCfgT_co = TypeVar("RuleCfgT_co", bound="EventRuleCfgBase[Any]", covariant=True)
+
+
+class EventRuleCfgBase(ABC, BaseModel, Generic[ModelT_contra]):
+    model_config = ConfigDict(extra="forbid")
+
+    @abstractmethod
+    def contextualize(self, context: Context) -> EventRule[ModelT_contra]: ...
+
+
+class EventCfg(BaseModel, Generic[RuleCfgT_co]):
+    model_config = ConfigDict(extra="forbid", frozen=True)
+
+    rule: RuleCfgT_co
+    hazard: HazardCfg
+
+    def contextualize(
+        self: EventCfg[EventRuleCfgBase[ModelT]], context: Context
+    ) -> Event[ModelT]:
+        return Event(
+            self.hazard.contextualize(context), self.rule.contextualize(context)
+        )
+
+
+class ScheduledEventCfg(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    times: FloatListCfg
+    action: ModelActionCfg
+
+    def contextualize(self, context: Context) -> ScheduledEvent[Any]:
+        return ScheduledEvent(
+            self.times.contextualize(context), self.action.contextualize(context)
+        )

phyloom/configs/generator.py

@@ -0,0 +1,143 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from functools import partial
+from pathlib import Path
+from random import Random
+from typing import TYPE_CHECKING, Annotated
+
+import pandas as pd
+from numpy.random import default_rng
+from pydantic import BaseModel, ConfigDict
+
+from phyloom.configs.roots import (
+    ContextCfg,
+    FloatCfg,
+    ModelExpressionCfg,
+    ModelPredicateCfg,
+    OutputHookCfg,
+    contextualize_optional,
+)
+from phyloom.simulator import (
+    MaxRejectionsExceededError,
+    SimulationTimeoutError,
+    run_simulations,
+)
+from phyloom.simulator.core import Model
+from phyloom.utils import Registry
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+    from typing import Any, TypeAlias
+
+    from phyloom.types import Context, Metadata, SampleID, Seed
+
+
+DATA_DIRNAME = "data"
+METADATA_FILENAME = "metadata.csv"
+
+
+class GeneratorBase(ABC, BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    output_dir: str = "phyloom-outputs"
+    n_samples: int | dict[str, int] = 1
+    n_jobs: int = -1
+    seed: int | None = None
+    context: ContextCfg = ContextCfg()
+
+    max_time: FloatCfg | None = None
+    stop_criterion: ModelPredicateCfg | None = None
+    acceptance_criterion: ModelPredicateCfg | None = None
+    max_rejections: int | None = None
+    timeout: float | None = None
+    logs: dict[str, ModelExpressionCfg] = {}
+
+    output_hooks: tuple[OutputHookCfg, ...]
+
+    @abstractmethod
+    def contextualize(self, context: Context) -> Model: ...
+
+    def _contextualize_logs(
+        self, context: Context
+    ) -> Callable[[Any], dict[str, object]]:
+        return lambda model: {
+            key: value.contextualize(context)(model) for key, value in self.logs.items()
+        }
+
+    def _generate_sample(
+        self, sample_id: SampleID, seed: Seed, *, data_dir: Path
+    ) -> Metadata:
+        timed_out_contexts: list[Context] = []
+        max_rejections_exceeded_contexts: list[Context] = []
+        while True:
+            context = self.context.materialize(default_rng(seed))
+            model = self.contextualize(context)
+
+            try:
+                metadata = model.run(
+                    max_time=contextualize_optional(self.max_time, context),
+                    stop_criterion=contextualize_optional(self.stop_criterion, context),
+                    acceptance_criterion=contextualize_optional(
+                        self.acceptance_criterion, context
+                    ),
+                    logs=self._contextualize_logs(context),
+                    timeout=self.timeout,
+                    max_rejections=self.max_rejections,
+                    seed=seed,
+                )
+                output_dir = data_dir / str(sample_id)
+                output_dir.mkdir(parents=True)
+                for hook in self.output_hooks:
+                    hook.contextualize(context)(model, output_dir)
+            except SimulationTimeoutError:
+                timed_out_contexts.append(context)
+                seed += 1
+                continue
+            except MaxRejectionsExceededError:
+                max_rejections_exceeded_contexts.append(context)
+                seed += 1
+                continue
+
+            return {
+                **context,
+                **metadata,
+                "timed_out_contexts": timed_out_contexts,
+                "max_rejections_exceeded_contexts": max_rejections_exceeded_contexts,
+            }
+
+    def _generate_split(self, n_samples: int, output_dir: Path, seed: Seed):
+        if output_dir.exists():
+            print(f"Output directory {output_dir} already exists. Skipping.")
+            return
+
+        data_dir = output_dir / DATA_DIRNAME
+        df = run_simulations(
+            n_samples,
+            partial(self._generate_sample, data_dir=data_dir),
+            seed=seed,
+            n_jobs=self.n_jobs,
+        )
+
+        for name in ["timed_out_contexts", "max_rejections_exceeded_contexts"]:
+            meta_df = pd.DataFrame([d for row in df[name] for d in row])
+            if not meta_df.empty:
+                meta_df.to_csv(output_dir / f"{name}.csv", index=False)
+            df = df.drop(columns=[name])
+
+        df.to_csv(output_dir / METADATA_FILENAME, index=False)
+
+    def generate(self):
+        rng = Random(self.seed)
+        output_dir = Path(self.output_dir)
+        if isinstance(self.n_samples, dict):
+            for split, n_samples in self.n_samples.items():
+                self._generate_split(n_samples, output_dir / split, rng.getrandbits(32))
+        else:
+            self._generate_split(self.n_samples, output_dir, rng.getrandbits(32))
+
+
+GENERATOR_REGISTRY: Registry[GeneratorBase] = Registry(
+    "Generator", GeneratorBase, discriminator="model"
+)
+GeneratorCfg: TypeAlias = Annotated[GeneratorBase, GENERATOR_REGISTRY.validator]

phyloom/configs/hazards.py

@@ -0,0 +1,88 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Annotated
+
+from pydantic import BaseModel, ConfigDict
+
+from phyloom.configs.roots import (
+    FloatCfg,
+    FloatListCfg,
+    IntCfg,
+    TimeFnCfg,
+)
+from phyloom.simulator.hazards import (
+    ConstantHazard,
+    DeltaHazard,
+    ExponentialHazard,
+    NumericHazard,
+    SkylineHazard,
+)
+from phyloom.utils import Registry
+
+if TYPE_CHECKING:
+    from typing import TypeAlias
+
+    from phyloom.simulator.hazards import Hazard
+    from phyloom.types import Context
+
+
+class HazardCfgBase(ABC, BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    @abstractmethod
+    def contextualize(self, context: Context) -> Hazard: ...
+
+
+HAZARD_REGISTRY: Registry[HazardCfgBase] = Registry("Hazard", HazardCfgBase)
+HazardCfg: TypeAlias = Annotated[HazardCfgBase, HAZARD_REGISTRY.validator]
+
+
+@HAZARD_REGISTRY.register("constant")
+class ConstantHazardCfg(HazardCfgBase):
+    rate: FloatCfg
+
+    def contextualize(self, context: Context) -> ConstantHazard:
+        return ConstantHazard(self.rate.contextualize(context))
+
+
+@HAZARD_REGISTRY.register("delta")
+class DeltaHazardCfg(HazardCfgBase):
+    def contextualize(self, context: Context) -> DeltaHazard:
+        return DeltaHazard()
+
+
+@HAZARD_REGISTRY.register("exponential")
+class ExponentialHazardCfg(HazardCfgBase):
+    rate: FloatCfg
+    decay: FloatCfg
+
+    def contextualize(self, context: Context) -> ExponentialHazard:
+        return ExponentialHazard(
+            self.rate.contextualize(context), self.decay.contextualize(context)
+        )
+
+
+@HAZARD_REGISTRY.register("skyline")
+class SkylineHazardCfg(HazardCfgBase):
+    rates: FloatListCfg
+    break_times: FloatListCfg
+
+    def contextualize(self, context: Context) -> SkylineHazard:
+        return SkylineHazard(
+            self.rates.contextualize(context), self.break_times.contextualize(context)
+        )
+
+
+@HAZARD_REGISTRY.register("numeric")
+class NumericHazardCfg(HazardCfgBase):
+    rate_fn: TimeFnCfg
+    max_bracketing_steps: IntCfg = IntCfg.model_validate(16)
+    initial_dt: FloatCfg = FloatCfg.model_validate(1.0)
+
+    def contextualize(self, context: Context) -> NumericHazard:
+        return NumericHazard(
+            self.rate_fn.contextualize(context),
+            max_bracketing_steps=self.max_bracketing_steps.contextualize(context),
+            initial_dt=self.initial_dt.contextualize(context),
+        )

phyloom/configs/models/bdsh.py

@@ -0,0 +1,141 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Annotated
+
+from phyloom.configs.events import EventCfg, EventRuleCfgBase, ScheduledEventCfg
+from phyloom.configs.generator import GENERATOR_REGISTRY, GeneratorBase
+from phyloom.configs.roots import StrCfg, StrPairCfg, contextualize_optional
+from phyloom.simulator.bdsh import (
+    LDH,
+    LGH,
+    LNH,
+    BDSHModel,
+    Birth,
+    Death,
+    Migration,
+    Sampling,
+)
+from phyloom.utils import Registry
+
+if TYPE_CHECKING:
+    from typing import TypeAlias
+
+    from phyloom.types import Context
+
+BDSH_EVENT_RULE_REGISTRY: Registry[EventRuleCfgBase[BDSHModel]] = Registry(
+    "BDSHEventRule", EventRuleCfgBase
+)
+BDSHRuleCfg: TypeAlias = Annotated[
+    EventRuleCfgBase[BDSHModel], BDSH_EVENT_RULE_REGISTRY.validator
+]
+
+
+@BDSH_EVENT_RULE_REGISTRY.register("birth")
+class BirthCfg(EventRuleCfgBase[BDSHModel]):
+    parent_state: StrCfg | None = None
+    child_states: tuple[StrCfg | None, StrCfg | None] = (None, None)
+
+    def contextualize(self, context: Context) -> Birth:
+        child_state1, child_state2 = self.child_states
+        return Birth(
+            parent_state=contextualize_optional(self.parent_state, context),
+            child_states=(
+                contextualize_optional(child_state1, context),
+                contextualize_optional(child_state2, context),
+            ),
+        )
+
+
+@BDSH_EVENT_RULE_REGISTRY.register("death")
+class DeathCfg(EventRuleCfgBase[BDSHModel]):
+    state: StrCfg | None = None
+
+    def contextualize(self, context: Context) -> Death:
+        return Death(state=contextualize_optional(self.state, context))
+
+
+@BDSH_EVENT_RULE_REGISTRY.register("migration")
+class MigrationCfg(EventRuleCfgBase[BDSHModel]):
+    target_state: StrCfg
+    source_state: StrCfg | None = None
+
+    def contextualize(self, context: Context) -> Migration:
+        return Migration(
+            target_state=self.target_state.contextualize(context),
+            source_state=contextualize_optional(self.source_state, context),
+        )
+
+
+@BDSH_EVENT_RULE_REGISTRY.register("sampling")
+class SamplingCfg(EventRuleCfgBase[BDSHModel]):
+    removal: bool
+    state: StrCfg | None = None
+
+    def contextualize(self, context: Context) -> Sampling:
+        return Sampling(
+            removal=self.removal, state=contextualize_optional(self.state, context)
+        )
+
+
+@BDSH_EVENT_RULE_REGISTRY.register("ldh")
+class LDHCfg(EventRuleCfgBase[BDSHModel]):
+    parent_states: StrPairCfg | None = None
+    hybrid_state: StrCfg | None = None
+
+    def contextualize(self, context: Context) -> LDH:
+        return LDH(
+            parent_states=contextualize_optional(self.parent_states, context),
+            hybrid_state=contextualize_optional(self.hybrid_state, context),
+        )
+
+
+@BDSH_EVENT_RULE_REGISTRY.register("lnh")
+class LNHCfg(EventRuleCfgBase[BDSHModel]):
+    continuing_parent_state: StrCfg | None = None
+    terminating_parent_state: StrCfg | None = None
+    new_parent_state: StrCfg | None = None
+    hybrid_state: StrCfg | None = None
+
+    def contextualize(self, context: Context) -> LNH:
+        return LNH(
+            continuing_parent_state=contextualize_optional(
+                self.continuing_parent_state, context
+            ),
+            terminating_parent_state=contextualize_optional(
+                self.terminating_parent_state, context
+            ),
+            new_parent_state=contextualize_optional(self.new_parent_state, context),
+            hybrid_state=contextualize_optional(self.hybrid_state, context),
+        )
+
+
+@BDSH_EVENT_RULE_REGISTRY.register("lgh")
+class LGHCfg(EventRuleCfgBase[BDSHModel]):
+    parent_states: StrPairCfg | None = None
+    hybrid_state: StrCfg | None = None
+    new_parent_states: tuple[StrCfg | None, StrCfg | None] = (None, None)
+
+    def contextualize(self, context: Context) -> LGH:
+        new_parent_state1, new_parent_state2 = self.new_parent_states
+        return LGH(
+            parent_states=contextualize_optional(self.parent_states, context),
+            hybrid_state=contextualize_optional(self.hybrid_state, context),
+            new_parent_states=(
+                contextualize_optional(new_parent_state1, context),
+                contextualize_optional(new_parent_state2, context),
+            ),
+        )
+
+
+@GENERATOR_REGISTRY.register("bdsh")
+class BDSHGenerator(GeneratorBase):
+    init_state: StrCfg
+    events: tuple[EventCfg[BDSHRuleCfg], ...] = ()
+    schedule: tuple[ScheduledEventCfg, ...] = ()
+
+    def contextualize(self, context: Context) -> BDSHModel:
+        return BDSHModel(
+            init_state=self.init_state.contextualize(context),
+            events=[e.contextualize(context) for e in self.events],
+            schedule=[e.contextualize(context) for e in self.schedule],
+        )

phyloom/configs/models/coalescent.py

@@ -0,0 +1,62 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Annotated
+
+from phyloom.configs.events import EventCfg, EventRuleCfgBase, ScheduledEventCfg
+from phyloom.configs.generator import GENERATOR_REGISTRY, GeneratorBase
+from phyloom.configs.roots import IntCfg, StrCfg, StrPairCfg, contextualize_optional
+from phyloom.simulator.coalescent import Coalesce, CoalescentModel, Retrace
+from phyloom.utils import Registry
+
+if TYPE_CHECKING:
+    from typing import TypeAlias
+
+    from phyloom.types import Context
+
+COALESCENT_EVENT_RULE_REGISTRY: Registry[EventRuleCfgBase[CoalescentModel]] = Registry(
+    "CoalescentEventRule", EventRuleCfgBase
+)
+CoalescentRuleCfg: TypeAlias = Annotated[
+    EventRuleCfgBase[CoalescentModel], COALESCENT_EVENT_RULE_REGISTRY.validator
+]
+
+
+@COALESCENT_EVENT_RULE_REGISTRY.register("retrace")
+class RetraceCfg(EventRuleCfgBase[CoalescentModel]):
+    ancestral_state: StrCfg
+    current_state: StrCfg | None = None
+
+    def contextualize(self, context: Context) -> Retrace:
+        return Retrace(
+            ancestral_state=self.ancestral_state.contextualize(context),
+            current_state=contextualize_optional(self.current_state, context),
+        )
+
+
+@COALESCENT_EVENT_RULE_REGISTRY.register("coalesce")
+class CoalesceCfg(EventRuleCfgBase[CoalescentModel]):
+    current_states: StrPairCfg | None = None
+    ancestral_state: StrCfg | None = None
+
+    def contextualize(self, context: Context) -> Coalesce:
+        return Coalesce(
+            current_states=contextualize_optional(self.current_states, context),
+            ancestral_state=contextualize_optional(self.ancestral_state, context),
+        )
+
+
+@GENERATOR_REGISTRY.register("coalescent")
+class CoalescentGenerator(GeneratorBase):
+    init_population: dict[str, IntCfg]
+    events: tuple[EventCfg[CoalescentRuleCfg], ...] = ()
+    schedule: tuple[ScheduledEventCfg, ...] = ()
+
+    def contextualize(self, context: Context) -> CoalescentModel:
+        return CoalescentModel(
+            init_population={
+                state: size.contextualize(context)
+                for state, size in self.init_population.items()
+            },
+            events=[e.contextualize(context) for e in self.events],
+            schedule=[e.contextualize(context) for e in self.schedule],
+        )

phyloom/configs/models/epidemic.py

@@ -0,0 +1,102 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Annotated
+
+from phyloom.configs.events import EventCfg, EventRuleCfgBase, ScheduledEventCfg
+from phyloom.configs.generator import GENERATOR_REGISTRY, GeneratorBase
+from phyloom.configs.roots import IntCfg, StrCfg, contextualize_optional
+from phyloom.simulator.epidemic import (
+    EpidemicModel,
+    Flow,
+    Infection,
+    Migration,
+    Recovery,
+)
+from phyloom.utils import Registry
+
+if TYPE_CHECKING:
+    from typing import TypeAlias
+
+    from phyloom.types import Context
+
+EPIDEMIC_EVENT_RULE_REGISTRY: Registry[EventRuleCfgBase[EpidemicModel]] = Registry(
+    "EpidemicEventRule", EventRuleCfgBase
+)
+EpidemicRuleCfg: TypeAlias = Annotated[
+    EventRuleCfgBase[EpidemicModel], EPIDEMIC_EVENT_RULE_REGISTRY.validator
+]
+
+
+@EPIDEMIC_EVENT_RULE_REGISTRY.register("infection")
+class InfectionCfg(EventRuleCfgBase[EpidemicModel]):
+    recipient_pool: StrCfg
+    source_state: StrCfg | None = None
+    child_states: tuple[StrCfg | None, StrCfg | None] = (None, None)
+
+    def contextualize(self, context: Context) -> Infection:
+        child_state1, child_state2 = self.child_states
+        return Infection(
+            recipient_pool=self.recipient_pool.contextualize(context),
+            source_state=contextualize_optional(self.source_state, context),
+            child_states=(
+                contextualize_optional(child_state1, context),
+                contextualize_optional(child_state2, context),
+            ),
+        )
+
+
+@EPIDEMIC_EVENT_RULE_REGISTRY.register("recovery")
+class RecoveryCfg(EventRuleCfgBase[EpidemicModel]):
+    destination_pool: StrCfg
+    source_state: StrCfg | None = None
+    sample: bool = False
+
+    def contextualize(self, context: Context) -> Recovery:
+        return Recovery(
+            destination_pool=self.destination_pool.contextualize(context),
+            source_state=contextualize_optional(self.source_state, context),
+            sample=self.sample,
+        )
+
+
+@EPIDEMIC_EVENT_RULE_REGISTRY.register("migration")
+class MigrationCfg(EventRuleCfgBase[EpidemicModel]):
+    target_state: StrCfg
+    source_state: StrCfg | None = None
+
+    def contextualize(self, context: Context) -> Migration:
+        return Migration(
+            target_state=self.target_state.contextualize(context),
+            source_state=contextualize_optional(self.source_state, context),
+        )
+
+
+@EPIDEMIC_EVENT_RULE_REGISTRY.register("flow")
+class FlowCfg(EventRuleCfgBase[EpidemicModel]):
+    source_pool: StrCfg
+    destination_pool: StrCfg
+
+    def contextualize(self, context: Context) -> Flow:
+        return Flow(
+            source_pool=self.source_pool.contextualize(context),
+            destination_pool=self.destination_pool.contextualize(context),
+        )
+
+
+@GENERATOR_REGISTRY.register("epidemic")
+class EpidemicGenerator(GeneratorBase):
+    init_state: StrCfg
+    init_pools: dict[str, IntCfg]
+    events: tuple[EventCfg[EpidemicRuleCfg], ...] = ()
+    schedule: tuple[ScheduledEventCfg, ...] = ()
+
+    def contextualize(self, context: Context) -> EpidemicModel:
+        return EpidemicModel(
+            init_state=self.init_state.contextualize(context),
+            init_pools={
+                pool: count.contextualize(context)
+                for pool, count in self.init_pools.items()
+            },
+            events=[e.contextualize(context) for e in self.events],
+            schedule=[e.contextualize(context) for e in self.schedule],
+        )

phyloom/configs/roots/__init__.py

@@ -0,0 +1,32 @@
+from phyloom.configs.roots.context import ContextCfg, DistributionCfg
+from phyloom.configs.roots.float import FloatCfg, FloatListCfg
+from phyloom.configs.roots.int import IntCfg
+from phyloom.configs.roots.source import (
+    ModelActionCfg,
+    ModelExpressionCfg,
+    ModelPredicateCfg,
+    NumPyExpressionCfg,
+    OutputHookCfg,
+    SourceCfg,
+    TimeFnCfg,
+)
+from phyloom.configs.roots.str import StrCfg, StrPairCfg
+from phyloom.configs.roots.utils import contextualize_optional
+
+__all__ = [
+    "ContextCfg",
+    "DistributionCfg",
+    "FloatCfg",
+    "FloatListCfg",
+    "IntCfg",
+    "ModelActionCfg",
+    "ModelExpressionCfg",
+    "ModelPredicateCfg",
+    "NumPyExpressionCfg",
+    "OutputHookCfg",
+    "SourceCfg",
+    "StrCfg",
+    "StrPairCfg",
+    "TimeFnCfg",
+    "contextualize_optional",
+]

phyloom/configs/roots/context.py

@@ -0,0 +1,39 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from pydantic import BaseModel, ConfigDict, Field, RootModel
+
+from phyloom.configs.roots.source import NumPyExpressionCfg
+from phyloom.configs.roots.str import StrCfg
+
+if TYPE_CHECKING:
+    from numpy.random import Generator
+
+    from phyloom.types import Context
+
+
+class DistributionCfg(BaseModel):
+    model_config = ConfigDict(extra="allow")
+    type: StrCfg
+    __pydantic_extra__: dict[str, NumPyExpressionCfg] = Field(init=False)  # pyright: ignore[reportIncompatibleVariableOverride]
+
+    def draw(self, context: Context, rng: Generator) -> object:
+        args = {
+            name: value.contextualize(context)
+            for name, value in self.__pydantic_extra__.items()
+        }
+        return getattr(rng, self.type.contextualize(context))(**args)
+
+
+class ContextCfg(RootModel[dict[str, DistributionCfg | NumPyExpressionCfg]]):
+    root: dict[str, DistributionCfg | NumPyExpressionCfg] = {}
+
+    def materialize(self, rng: Generator) -> Context:
+        values: Context = {}
+        for key, value in self.root.items():
+            if isinstance(value, NumPyExpressionCfg):
+                values[key] = value.contextualize(values)
+            else:
+                values[key] = value.draw(values, rng)
+        return values