phylogenie 1.0.8__py3-none-any.whl → 2.0.0__py3-none-any.whl

phylogenie/generators/__init__.py

@@ -0,0 +1,14 @@
+from typing import Annotated
+
+from pydantic import Field
+
+from phylogenie.generators.alisim import AliSimDatasetGenerator
+from phylogenie.generators.dataset import DatasetGenerator
+from phylogenie.generators.trees import TreeDatasetGeneratorConfig
+
+DatasetGeneratorConfig = Annotated[
+    TreeDatasetGeneratorConfig | AliSimDatasetGenerator,
+    Field(discriminator="data_type"),
+]
+
+__all__ = ["DatasetGeneratorConfig", "DatasetGenerator"]

phylogenie/generators/alisim.py

@@ -0,0 +1,71 @@
+import os
+import subprocess
+from pathlib import Path
+from typing import Any, Literal
+
+from numpy.random import Generator
+
+from phylogenie.generators.dataset import DatasetGenerator, DataType
+from phylogenie.generators.trees import TreeDatasetGeneratorConfig
+from phylogenie.io import dump_newick
+
+MSAS_DIRNAME = "MSAs"
+TREES_DIRNAME = "trees"
+
+
+class AliSimDatasetGenerator(DatasetGenerator):
+    data_type: Literal[DataType.MSAS] = DataType.MSAS
+    trees: TreeDatasetGeneratorConfig
+    keep_trees: bool = False
+    iqtree_path: str = "iqtree2"
+    args: dict[str, str | int | float]
+
+    def _generate_one_from_tree(
+        self, filename: str, tree_file: str, rng: Generator, data: dict[str, Any]
+    ) -> None:
+        command = [
+            self.iqtree_path,
+            "--alisim",
+            filename,
+            "--tree",
+            tree_file,
+            "--seed",
+            str(rng.integers(2**32)),
+        ]
+
+        for key, value in self.args.items():
+            command.extend(
+                [key, value.format(**data) if isinstance(value, str) else str(value)]
+            )
+
+        command.extend(["-af", "fasta"])
+        subprocess.run(command, check=True, stdout=subprocess.DEVNULL)
+        subprocess.run(["rm", f"{tree_file}.log"], check=True)
+
+    def _generate_one(
+        self, filename: str, rng: Generator, data: dict[str, Any]
+    ) -> None:
+        if self.keep_trees:
+            base_dir = Path(filename).parent
+            file_id = Path(filename).stem
+            tree_filename = os.path.join(base_dir, TREES_DIRNAME, file_id)
+            msas_dir = os.path.join(base_dir, MSAS_DIRNAME)
+            os.makedirs(msas_dir, exist_ok=True)
+            msa_filename = os.path.join(msas_dir, file_id)
+        else:
+            tree_filename = f"{filename}.temp-tree"
+            msa_filename = filename
+
+        tree = self.trees.simulate_one(rng, data)
+        if tree is None:
+            return
+
+        for leaf in tree.get_leaves():
+            leaf.id += f"|{leaf.get_time()}"
+        dump_newick(tree, f"{tree_filename}.nwk")
+
+        self._generate_one_from_tree(
+            filename=msa_filename, tree_file=f"{tree_filename}.nwk", rng=rng, data=data
+        )
+        if not self.keep_trees:
+            os.remove(f"{tree_filename}.nwk")

phylogenie/generators/configs.py

@@ -0,0 +1,41 @@
+from pydantic import BaseModel, ConfigDict
+
+import phylogenie.typings as pgt
+
+
+class DistributionConfig(BaseModel):
+    type: str
+    model_config = ConfigDict(extra="allow")
+
+
+IntegerConfig = str | int
+ScalarConfig = str | pgt.Scalar
+ManyScalarsConfig = str | list[ScalarConfig]
+OneOrManyScalarsConfig = ScalarConfig | list[ScalarConfig]
+OneOrMany2DScalarsConfig = ScalarConfig | list[list[ScalarConfig]]
+
+
+class StrictBaseModel(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+
+class SkylineParameterModel(StrictBaseModel):
+    value: ManyScalarsConfig
+    change_times: ManyScalarsConfig
+
+
+class SkylineVectorModel(StrictBaseModel):
+    value: str | list[OneOrManyScalarsConfig]
+    change_times: ManyScalarsConfig
+
+
+class SkylineMatrixModel(StrictBaseModel):
+    value: str | list[OneOrMany2DScalarsConfig]
+    change_times: ManyScalarsConfig
+
+
+SkylineParameterConfig = ScalarConfig | SkylineParameterModel
+SkylineVectorConfig = (
+    str | pgt.Scalar | list[SkylineParameterConfig] | SkylineVectorModel
+)
+SkylineMatrixConfig = str | pgt.Scalar | list[SkylineVectorConfig] | SkylineMatrixModel

phylogenie/{core → generators}/dataset.py

@@ -1,15 +1,16 @@
 import os
 from abc import ABC, abstractmethod
 from enum import Enum
+from itertools import product
+from typing import Any
 
 import joblib
+import numpy as np
 import pandas as pd
 from numpy.random import Generator, default_rng
 from tqdm import tqdm
 
-import phylogenie.typings as pgt
-from phylogenie.configs import StrictBaseModel
-from phylogenie.core.context import ContextConfig, context_factory
+from phylogenie.generators.configs import DistributionConfig, StrictBaseModel
 
 
 class DataType(str, Enum):
@@ -17,51 +18,52 @@ class DataType(str, Enum):
     MSAS = "msas"
 
 
+DATA_DIRNAME = "data"
+METADATA_FILENAME = "metadata.csv"
+
+
 class DatasetGenerator(ABC, StrictBaseModel):
-    output_dir: str = "phylogenie-out"
-    data_dir: str = "data"
-    metadata_filename: str = "metadata.csv"
+    output_dir: str = "phylogenie-outputs"
     n_samples: int | dict[str, int] = 1
     n_jobs: int = -1
     seed: int | None = None
-    context: ContextConfig | None = None
+    context: dict[str, DistributionConfig] | None = None
 
     @abstractmethod
-    def _generate_one(self, filename: str, rng: Generator, data: pgt.Data) -> None: ...
+    def _generate_one(
+        self, filename: str, rng: Generator, data: dict[str, Any]
+    ) -> None: ...
 
     def generate_one(
-        self, filename: str, data: pgt.Data | None = None, seed: int | None = None
+        self, filename: str, data: dict[str, Any] | None = None, seed: int | None = None
     ) -> None:
         data = {} if data is None else data
         self._generate_one(filename=filename, rng=default_rng(seed), data=data)
 
     def _generate(self, rng: Generator, n_samples: int, output_dir: str) -> None:
-        data_dir = os.path.join(output_dir, self.data_dir)
-        metadata_file = os.path.join(output_dir, self.metadata_filename)
+        data_dir = os.path.join(output_dir, DATA_DIRNAME)
         if os.path.exists(data_dir):
             print(f"Output directory {data_dir} already exists. Skipping.")
             return
         os.makedirs(data_dir)
 
-        data = [
-            {} if self.context is None else context_factory(self.context, rng)
-            for _ in range(n_samples)
-        ]
+        data: list[dict[str, Any]] = [{}] * 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 {}
+                d[k] = np.array(getattr(rng, v.type)(**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)
 
         joblib.Parallel(n_jobs=self.n_jobs)(
             joblib.delayed(self.generate_one)(
                 filename=os.path.join(data_dir, str(i)),
-                data=d,
-                seed=int(rng.integers(0, 2**32)),
-            )
-            for i, d in tqdm(
-                enumerate(data), total=n_samples, desc=f"Generating {data_dir}..."
+                data=data[i],
+                seed=int(rng.integers(2**32)),
             )
+            for i in tqdm(range(n_samples), desc=f"Generating {data_dir}...")
         )
 
-        df = pd.DataFrame([{"file_id": str(i), **d} for i, d in enumerate(data)])
-        df.to_csv(metadata_file, index=False)
-
     def generate(self) -> None:
         rng = default_rng(self.seed)
         if isinstance(self.n_samples, dict):

phylogenie/{core → generators}/factories.py

@@ -2,8 +2,8 @@ from typing import Any
 
 import numpy as np
 
-import phylogenie.core.configs as cfg
-import phylogenie.core.typeguards as ctg
+import phylogenie.generators.configs as cfg
+import phylogenie.generators.typeguards as ctg
 import phylogenie.typeguards as tg
 import phylogenie.typings as pgt
 from phylogenie.skyline import (
@@ -16,7 +16,7 @@ from phylogenie.skyline import (
 )
 
 
-def _eval_expression(expression: str, data: pgt.Data) -> Any:
+def _eval_expression(expression: str, data: dict[str, Any]) -> Any:
     return np.array(
         eval(
             expression,
@@ -29,7 +29,7 @@ def _eval_expression(expression: str, data: pgt.Data) -> Any:
     ).tolist()
 
 
-def int_factory(x: cfg.IntConfig, data: pgt.Data) -> int:
+def integer(x: cfg.IntegerConfig, data: dict[str, Any]) -> int:
     if isinstance(x, str):
         e = _eval_expression(x, data)
         if isinstance(e, int):
@@ -40,7 +40,7 @@ def int_factory(x: cfg.IntConfig, data: pgt.Data) -> int:
     return x
 
 
-def scalar_factory(x: cfg.ScalarConfig, data: pgt.Data) -> pgt.Scalar:
+def scalar(x: cfg.ScalarConfig, data: dict[str, Any]) -> pgt.Scalar:
     if isinstance(x, str):
         e = _eval_expression(x, data)
         if isinstance(e, pgt.Scalar):
@@ -51,18 +51,7 @@ def scalar_factory(x: cfg.ScalarConfig, data: pgt.Data) -> pgt.Scalar:
     return x
 
 
-def many_ints_factory(x: cfg.ManyIntsConfig, data: pgt.Data) -> pgt.Many[int]:
-    if isinstance(x, str):
-        e = _eval_expression(x, data)
-        if tg.is_many_ints(e):
-            return e
-        raise ValueError(
-            f"Expression '{x}' evaluated to {e} of type {type(e)}, expected a sequence of integers."
-        )
-    return [int_factory(v, data) for v in x]
-
-
-def many_scalars_factory(x: cfg.ManyScalarsConfig, data: pgt.Data) -> pgt.ManyScalars:
+def many_scalars(x: cfg.ManyScalarsConfig, data: dict[str, Any]) -> pgt.ManyScalars:
     if isinstance(x, str):
         e = _eval_expression(x, data)
         if tg.is_many_scalars(e):
@@ -70,11 +59,11 @@ def many_scalars_factory(x: cfg.ManyScalarsConfig, data: pgt.Data) -> pgt.ManySc
         raise ValueError(
             f"Expression '{x}' evaluated to {e} of type {type(e)}, expected a sequence of scalars."
         )
-    return [scalar_factory(v, data) for v in x]
+    return [scalar(v, data) for v in x]
 
 
-def one_or_many_scalars_factory(
-    x: cfg.OneOrManyScalarsConfig, data: pgt.Data
+def one_or_many_scalars(
+    x: cfg.OneOrManyScalarsConfig, data: dict[str, Any]
 ) -> pgt.OneOrManyScalars:
     if isinstance(x, str):
         e = _eval_expression(x, data)
@@ -85,22 +74,22 @@ def one_or_many_scalars_factory(
         )
     if isinstance(x, pgt.Scalar):
         return x
-    return many_scalars_factory(x, data)
+    return many_scalars(x, data)
 
 
-def skyline_parameter_like_factory(
-    x: cfg.SkylineParameterLikeConfig, data: pgt.Data
+def skyline_parameter(
+    x: cfg.SkylineParameterConfig, data: dict[str, Any]
 ) -> SkylineParameterLike:
     if isinstance(x, cfg.ScalarConfig):
-        return scalar_factory(x, data)
+        return scalar(x, data)
     return SkylineParameter(
-        value=many_scalars_factory(x.value, data),
-        change_times=many_scalars_factory(x.change_times, data),
+        value=many_scalars(x.value, data),
+        change_times=many_scalars(x.change_times, data),
     )
 
 
-def skyline_vector_coercible_factory(
-    x: cfg.SkylineVectorCoercibleConfig, data: pgt.Data
+def skyline_vector(
+    x: cfg.SkylineVectorConfig, data: dict[str, Any]
 ) -> SkylineVectorCoercible:
     if isinstance(x, str):
         e = _eval_expression(x, data)
@@ -111,12 +100,12 @@ def skyline_vector_coercible_factory(
         )
     if isinstance(x, pgt.Scalar):
         return x
-    if ctg.is_list_of_skyline_parameter_like_configs(x):
-        return [skyline_parameter_like_factory(p, data) for p in x]
+    if ctg.is_list_of_skyline_parameter_configs(x):
+        return [skyline_parameter(p, data) for p in x]
 
-    assert isinstance(x, cfg.SkylineVectorValueModel)
+    assert isinstance(x, cfg.SkylineVectorModel)
 
-    change_times = many_scalars_factory(x.change_times, data)
+    change_times = many_scalars(x.change_times, data)
     if isinstance(x.value, str):
         e = _eval_expression(x.value, data)
         if tg.is_many_one_or_many_scalars(e):
@@ -126,7 +115,7 @@ def skyline_vector_coercible_factory(
                 f"Expression '{x.value}' evaluated to {e} of type {type(e)}, which cannot be coerced to a valid value for a SkylineVector (expected a sequence composed of scalars and/or sequences of scalars)."
             )
     else:
-        value = [one_or_many_scalars_factory(v, data) for v in x.value]
+        value = [one_or_many_scalars(v, data) for v in x.value]
 
     if tg.is_many_scalars(value):
         return SkylineParameter(value=value, change_times=change_times)
@@ -142,8 +131,8 @@ def skyline_vector_coercible_factory(
     return SkylineVector(value=value, change_times=change_times)
 
 
-def one_or_many_2D_scalars_factory(
-    x: cfg.OneOrMany2DScalarsConfig, data: pgt.Data
+def one_or_many_2D_scalars(
+    x: cfg.OneOrMany2DScalarsConfig, data: dict[str, Any]
 ) -> pgt.OneOrMany2DScalars:
     if isinstance(x, str):
         e = _eval_expression(x, data)
@@ -154,11 +143,11 @@ def one_or_many_2D_scalars_factory(
         )
     if isinstance(x, pgt.Scalar):
         return x
-    return [many_scalars_factory(v, data) for v in x]
+    return [many_scalars(v, data) for v in x]
 
 
-def skyline_matrix_coercible_factory(
-    x: cfg.SkylineMatrixCoercibleConfig, data: pgt.Data
+def skyline_matrix(
+    x: cfg.SkylineMatrixConfig, data: dict[str, Any]
 ) -> SkylineMatrixCoercible:
     if isinstance(x, str):
         e = _eval_expression(x, data)
@@ -169,12 +158,12 @@ def skyline_matrix_coercible_factory(
         )
     if isinstance(x, pgt.Scalar):
         return x
-    if ctg.is_list_of_skyline_vector_coercible_configs(x):
-        return [skyline_vector_coercible_factory(v, data) for v in x]
+    if ctg.is_list_of_skyline_vector_configs(x):
+        return [skyline_vector(v, data) for v in x]
 
-    assert isinstance(x, cfg.SkylineMatrixValueModel)
+    assert isinstance(x, cfg.SkylineMatrixModel)
 
-    change_times = many_scalars_factory(x.change_times, data)
+    change_times = many_scalars(x.change_times, data)
     if isinstance(x.value, str):
         e = _eval_expression(x.value, data)
         if tg.is_many_one_or_many_2D_scalars(e):
@@ -184,26 +173,27 @@ def skyline_matrix_coercible_factory(
                 f"Expression '{x.value}' evaluated to {e} of type {type(e)}, which cannot be coerced to a valid value for a SkylineMatrix (expected a sequence composed of scalars and/or nested (2D) sequences of scalars)."
             )
     else:
-        value = [one_or_many_2D_scalars_factory(v, data) for v in x.value]
+        value = [one_or_many_2D_scalars(v, data) for v in x.value]
 
     if tg.is_many_scalars(value):
         return SkylineParameter(value=value, change_times=change_times)
 
-    Ns: set[int] = set()
+    shapes: set[tuple[int, int]] = set()
     for elem in value:
         if tg.is_many_2D_scalars(elem):
-            n_rows = len(elem)
-            if any(len(row) != n_rows for row in elem):
+            Ms = len(elem)
+            Ns = {len(row) for row in elem}
+            if len(Ns) > 1:
                 raise ValueError(
-                    f"All elements in the value of a SkylineMatrix config must be scalars or square matrices (config {x.value} yeilded a non-square matrix: {elem})."
+                    f"The values of a SkylineMatrix config must be scalars or nested (2D) lists of them with a consistent row length (config {x.value} yielded element {elem} with row lengths {Ns})."
                 )
-            Ns.add(n_rows)
+            shapes.add((Ms, Ns.pop()))
 
-    if len(Ns) > 1:
+    if len(shapes) > 1:
         raise ValueError(
-            f"All elements in the value of a SkylineMatrix config must be scalars or have the same square shape (config {x.value} yielded value={value} with inconsistent lengths {Ns})."
+            f"All elements in the value of a SkylineMatrix config must be scalars or nested (2D) lists of them with the same shape (config {x.value} yielded value={value} with inconsistent shapes {shapes})."
         )
-    (N,) = Ns
-    value = [[[p] * N] * N if isinstance(p, pgt.Scalar) else p for p in value]
+    ((M, N),) = shapes
+    value = [[[e] * N] * M if isinstance(e, pgt.Scalar) else e for e in value]
 
     return SkylineMatrix(value=value, change_times=change_times)

phylogenie/generators/trees.py

@@ -0,0 +1,220 @@
+from abc import abstractmethod
+from enum import Enum
+from typing import Annotated, Any, Literal
+
+import numpy as np
+from numpy.random import Generator
+from pydantic import Field
+
+import phylogenie.generators.configs as cfg
+from phylogenie.generators.dataset import DatasetGenerator, DataType
+from phylogenie.generators.factories import (
+    integer,
+    scalar,
+    skyline_matrix,
+    skyline_parameter,
+    skyline_vector,
+)
+from phylogenie.io import dump_newick
+from phylogenie.tree import Tree
+from phylogenie.treesimulator import (
+    Event,
+    get_BD_events,
+    get_BDEI_events,
+    get_BDSS_events,
+    get_canonical_events,
+    get_epidemiological_events,
+    get_FBD_events,
+    simulate_tree,
+)
+
+
+class ParameterizationType(str, Enum):
+    CANONICAL = "canonical"
+    EPIDEMIOLOGICAL = "epidemiological"
+    FBD = "FBD"
+    BD = "BD"
+    BDEI = "BDEI"
+    BDSS = "BDSS"
+
+
+class TreeDatasetGenerator(DatasetGenerator):
+    data_type: Literal[DataType.TREES] = DataType.TREES
+    min_tips: cfg.IntegerConfig = 1
+    max_tips: cfg.IntegerConfig | None = None
+    max_time: cfg.ScalarConfig = np.inf
+    init_state: str | None = None
+    sampling_probability_at_present: cfg.ScalarConfig = 0.0
+    max_tries: int | None = None
+
+    def simulate_one(self, rng: Generator, data: dict[str, Any]) -> Tree | None:
+        events = self._get_events(rng, data)
+        init_state = (
+            self.init_state
+            if self.init_state is None
+            else self.init_state.format(**data)
+        )
+        max_tips = (
+            self.max_tips if self.max_tips is None else integer(self.max_tips, data)
+        )
+        return simulate_tree(
+            events=events,
+            min_tips=integer(self.min_tips, data),
+            max_tips=max_tips,
+            max_time=scalar(self.max_time, data),
+            init_state=init_state,
+            sampling_probability_at_present=scalar(
+                self.sampling_probability_at_present, data
+            ),
+            max_tries=self.max_tries,
+            seed=int(rng.integers(2**32)),
+        )
+
+    @abstractmethod
+    def _get_events(self, rng: Generator, data: dict[str, Any]) -> list[Event]: ...
+
+    def _generate_one(
+        self, filename: str, rng: Generator, data: dict[str, Any]
+    ) -> None:
+        tree = self.simulate_one(rng, data)
+        if tree is not None:
+            dump_newick(tree, f"{filename}.nwk")
+
+
+class CanonicalTreeDatasetGenerator(TreeDatasetGenerator):
+    parameterization: Literal[ParameterizationType.CANONICAL] = (
+        ParameterizationType.CANONICAL
+    )
+    sampling_rates: cfg.SkylineVectorConfig
+    birth_rates: cfg.SkylineVectorConfig = 0
+    death_rates: cfg.SkylineVectorConfig = 0
+    removal_probabilities: cfg.SkylineVectorConfig = 0
+    migration_rates: cfg.SkylineMatrixConfig = 0
+    birth_rates_among_states: cfg.SkylineMatrixConfig = 0
+    states: list[str] | None = None
+
+    def _get_events(self, rng: Generator, data: dict[str, Any]) -> list[Event]:
+        return get_canonical_events(
+            states=self.states,
+            sampling_rates=skyline_vector(self.sampling_rates, data),
+            birth_rates=skyline_vector(self.birth_rates, data),
+            death_rates=skyline_vector(self.death_rates, data),
+            removal_probabilities=skyline_vector(self.removal_probabilities, data),
+            migration_rates=skyline_matrix(self.migration_rates, data),
+            birth_rates_among_states=skyline_matrix(
+                self.birth_rates_among_states, data
+            ),
+        )
+
+
+class EpidemiologicalTreeDatasetGenerator(TreeDatasetGenerator):
+    parameterization: Literal[ParameterizationType.EPIDEMIOLOGICAL] = (
+        ParameterizationType.EPIDEMIOLOGICAL
+    )
+    states: list[str] | None = None
+    reproduction_numbers: cfg.SkylineVectorConfig = 0
+    become_uninfectious_rates: cfg.SkylineVectorConfig = 0
+    sampling_proportions: cfg.SkylineVectorConfig = 1
+    removal_probabilities: cfg.SkylineVectorConfig = 1
+    migration_rates: cfg.SkylineMatrixConfig = 0
+    reproduction_numbers_among_states: cfg.SkylineMatrixConfig = 0
+
+    def _get_events(self, rng: Generator, data: dict[str, Any]) -> list[Event]:
+        return get_epidemiological_events(
+            states=self.states,
+            reproduction_numbers=skyline_vector(self.reproduction_numbers, data),
+            become_uninfectious_rates=skyline_vector(
+                self.become_uninfectious_rates, data
+            ),
+            sampling_proportions=skyline_vector(self.sampling_proportions, data),
+            removal_probabilities=skyline_vector(self.removal_probabilities, data),
+            migration_rates=skyline_matrix(self.migration_rates, data),
+            reproduction_numbers_among_states=skyline_matrix(
+                self.reproduction_numbers_among_states, data
+            ),
+        )
+
+
+class FBDTreeDatasetGenerator(TreeDatasetGenerator):
+    parameterization: Literal[ParameterizationType.FBD] = ParameterizationType.FBD
+    states: list[str] | None = None
+    diversification: cfg.SkylineVectorConfig = 0
+    turnover: cfg.SkylineVectorConfig = 0
+    sampling_proportions: cfg.SkylineVectorConfig = 1
+    removal_probabilities: cfg.SkylineVectorConfig = 0
+    migration_rates: cfg.SkylineMatrixConfig = 0
+    diversification_between_types: cfg.SkylineMatrixConfig = 0
+
+    def _get_events(self, rng: Generator, data: dict[str, Any]) -> list[Event]:
+        return get_FBD_events(
+            states=self.states,
+            diversification=skyline_vector(self.diversification, data),
+            turnover=skyline_vector(self.turnover, data),
+            sampling_proportions=skyline_vector(self.sampling_proportions, data),
+            removal_probabilities=skyline_vector(self.removal_probabilities, data),
+            migration_rates=skyline_matrix(self.migration_rates, data),
+            diversification_between_types=skyline_matrix(
+                self.diversification_between_types, data
+            ),
+        )
+
+
+class BDTreeDatasetGenerator(TreeDatasetGenerator):
+    parameterization: Literal[ParameterizationType.BD] = ParameterizationType.BD
+    reproduction_number: cfg.SkylineParameterConfig
+    infectious_period: cfg.SkylineParameterConfig
+    sampling_proportion: cfg.SkylineParameterConfig = 1
+
+    def _get_events(self, rng: Generator, data: dict[str, Any]) -> list[Event]:
+        return get_BD_events(
+            reproduction_number=skyline_parameter(self.reproduction_number, data),
+            infectious_period=skyline_parameter(self.infectious_period, data),
+            sampling_proportion=skyline_parameter(self.sampling_proportion, data),
+        )
+
+
+class BDEITreeDatasetGenerator(TreeDatasetGenerator):
+    parameterization: Literal[ParameterizationType.BDEI] = ParameterizationType.BDEI
+    reproduction_number: cfg.SkylineParameterConfig
+    infectious_period: cfg.SkylineParameterConfig
+    incubation_period: cfg.SkylineParameterConfig
+    sampling_proportion: cfg.SkylineParameterConfig = 1
+
+    def _get_events(self, rng: Generator, data: dict[str, Any]) -> list[Event]:
+        return get_BDEI_events(
+            reproduction_number=skyline_parameter(self.reproduction_number, data),
+            infectious_period=skyline_parameter(self.infectious_period, data),
+            incubation_period=skyline_parameter(self.incubation_period, data),
+            sampling_proportion=skyline_parameter(self.sampling_proportion, data),
+        )
+
+
+class BDSSTreeDatasetGenerator(TreeDatasetGenerator):
+    parameterization: Literal[ParameterizationType.BDSS] = ParameterizationType.BDSS
+    reproduction_number: cfg.SkylineParameterConfig
+    infectious_period: cfg.SkylineParameterConfig
+    superspreading_ratio: cfg.SkylineParameterConfig
+    superspreaders_proportion: cfg.SkylineParameterConfig
+    sampling_proportion: cfg.SkylineParameterConfig = 1
+
+    def _get_events(self, rng: Generator, data: dict[str, Any]) -> list[Event]:
+        return get_BDSS_events(
+            reproduction_number=skyline_parameter(self.reproduction_number, data),
+            infectious_period=skyline_parameter(self.infectious_period, data),
+            superspreading_ratio=skyline_parameter(self.superspreading_ratio, data),
+            superspreaders_proportion=skyline_parameter(
+                self.superspreaders_proportion, data
+            ),
+            sampling_proportion=skyline_parameter(self.sampling_proportion, data),
+        )
+
+
+TreeDatasetGeneratorConfig = Annotated[
+    CanonicalTreeDatasetGenerator
+    | EpidemiologicalTreeDatasetGenerator
+    | FBDTreeDatasetGenerator
+    | BDTreeDatasetGenerator
+    | BDEITreeDatasetGenerator
+    | BDSSTreeDatasetGenerator,
+    Field(discriminator="parameterization"),
+]

phylogenie/generators/typeguards.py

@@ -0,0 +1,32 @@
+from typing import TypeGuard
+
+import phylogenie.generators.configs as cfg
+import phylogenie.typings as pgt
+
+
+def is_list(x: object) -> TypeGuard[list[object]]:
+    return isinstance(x, list)
+
+
+def is_list_of_scalar_configs(x: object) -> TypeGuard[list[cfg.ScalarConfig]]:
+    return is_list(x) and all(isinstance(v, cfg.ScalarConfig) for v in x)
+
+
+def is_list_of_skyline_parameter_configs(
+    x: object,
+) -> TypeGuard[list[cfg.SkylineParameterConfig]]:
+    return is_list(x) and all(isinstance(v, cfg.SkylineParameterConfig) for v in x)
+
+
+def is_skyline_vector_config(
+    x: object,
+) -> TypeGuard[cfg.SkylineVectorConfig]:
+    return isinstance(
+        x, str | pgt.Scalar | cfg.SkylineVectorModel
+    ) or is_list_of_skyline_parameter_configs(x)
+
+
+def is_list_of_skyline_vector_configs(
+    x: object,
+) -> TypeGuard[list[cfg.SkylineVectorConfig]]:
+    return is_list(x) and all(is_skyline_vector_config(v) for v in x)