bella-companion 0.0.0__py3-none-any.whl

bella_companion/simulations/figures/utils.py

@@ -0,0 +1,250 @@
+import os
+from itertools import product
+from typing import Any
+
+import matplotlib.pyplot as plt
+import numpy as np
+import polars as pl
+import seaborn as sns
+from joblib import Parallel, delayed
+from lumiere.backend import (
+    ActivationFunction,
+    get_partial_dependence_values,
+    get_shap_features_importance,
+    sigmoid,
+)
+from lumiere.backend.typings import Weights
+from tqdm import tqdm
+
+from src.simulations.features import Feature
+
+
+def _set_xticks(n: int, reverse: bool = False):
+    xticks_labels = range(n)
+    if reverse:
+        xticks_labels = reversed(xticks_labels)
+    plt.xticks(ticks=range(n), labels=list(map(str, xticks_labels)))
+
+
+def step(
+    x: list[float],
+    reverse_xticks: bool = False,
+    **kwargs: dict[str, Any],
+):
+    data = x.copy()
+    data.insert(0, data[0])
+    plt.step(list(range(len(data))), data, **kwargs)
+    _set_xticks(len(data), reverse_xticks)
+    plt.xlabel("Time bin")
+
+
+def _count_time_bins(true_values: dict[str, list[float]]) -> int:
+    assert (
+        len({len(true_value) for true_value in true_values.values()}) == 1
+    ), "All targets must have the same number of change times."
+    return len(next(iter((true_values.values()))))
+
+
+def plot_maes_per_time_bin(
+    logs_summaries: dict[str, pl.DataFrame],
+    true_values: dict[str, list[float]],
+    output_filepath: str,
+    reverse_xticks: bool = False,
+):
+    def _mae(target: str, i: int) -> pl.Expr:
+        return (pl.col(f"{target}i{i}_median") - true_values[target][i]).abs()
+
+    n_time_bins = _count_time_bins(true_values)
+    df = pl.concat(
+        logs_summaries[model]
+        .select(
+            pl.mean_horizontal([_mae(target, i) for target in true_values]).alias("MAE")
+        )
+        .with_columns(pl.lit(i).alias("Time bin"), pl.lit(model).alias("Model"))
+        for i in range(n_time_bins)
+        for model in logs_summaries
+    )
+    sns.violinplot(
+        x="Time bin",
+        y="MAE",
+        hue="Model",
+        data=df,
+        inner=None,
+        cut=0,
+        density_norm="width",
+        legend=False,
+    )
+    _set_xticks(n_time_bins, reverse_xticks)
+    plt.savefig(output_filepath)
+    plt.close()
+
+
+def plot_coverage_per_time_bin(
+    logs_summaries: dict[str, pl.DataFrame],
+    true_values: dict[str, list[float]],
+    output_filepath: str,
+    reverse_xticks: bool = False,
+):
+    def _coverage(model: str, target: str, i: int) -> float:
+        lower_bound = logs_summaries[model][f"{target}i{i}_lower"]
+        upper_bound = logs_summaries[model][f"{target}i{i}_upper"]
+        true_value = true_values[target][i]
+        N = len(logs_summaries[model])
+        return ((lower_bound <= true_value) & (true_value <= upper_bound)).sum() / N
+
+    n_time_bins = _count_time_bins(true_values)
+    for model in logs_summaries:
+        avg_coverage_by_time_bin = [
+            np.mean([_coverage(model, target, i) for target in true_values])
+            for i in range(_count_time_bins(true_values))
+        ]
+        plt.plot(avg_coverage_by_time_bin, marker="o")
+
+    _set_xticks(n_time_bins, reverse_xticks)
+    plt.xlabel("Time bin")
+    plt.ylabel("Coverage")
+    plt.ylim((0, 1.05))
+    plt.savefig(output_filepath)
+    plt.close()
+
+
+def plot_partial_dependencies(
+    weights: list[list[Weights]],  # shape: (n_mcmcs, n_weights_samples, ...)
+    features: dict[str, Feature],
+    output_dir: str,
+    hidden_activation: ActivationFunction = sigmoid,
+    output_activation: ActivationFunction = sigmoid,
+):
+    os.makedirs(output_dir, exist_ok=True)
+    features_grid = [feature.grid for feature in features.values()]
+
+    def _get_median_partial_dependence_values(
+        weights: list[Weights],
+    ) -> list[list[float]]:
+        pdvalues_distribution = [
+            get_partial_dependence_values(
+                weights=w,
+                features_grid=features_grid,
+                hidden_activation=hidden_activation,
+                output_activation=output_activation,
+            )
+            for w in weights
+        ]
+        return [
+            np.median(
+                [pdvalues[feature_idx] for pdvalues in pdvalues_distribution], axis=0
+            ).tolist()
+            for feature_idx in range(len(features))
+        ]
+
+    jobs = Parallel(n_jobs=-1)(
+        delayed(_get_median_partial_dependence_values)(w) for w in weights
+    )
+    pdvalues = [
+        job for job in tqdm(jobs, total=len(weights), desc="Evaluating PDPs")
+    ]  # shape: (n_mcmcs, n_features, n_grid_points)
+    pdvalues = [
+        np.array(mcmc_pds).T for mcmc_pds in zip(*pdvalues)
+    ]  # shape: (n_features, n_grid_points, n_mcmcs)
+
+    for feature_idx, (feature_name, feature) in enumerate(features.items()):
+        color = "red" if feature.is_relevant else "gray"
+        feature_pdvalues = pdvalues[feature_idx]  # shape: (n_grid_points, n_mcmcs)
+        if not feature.is_categorical:
+            median = np.median(feature_pdvalues, axis=1)
+            lower = np.percentile(feature_pdvalues, 2.5, axis=1)
+            high = np.percentile(feature_pdvalues, 100 - 2.5, axis=1)
+            plt.fill_between(feature.grid, lower, high, alpha=0.25, color=color)
+            for mcmc_pds in feature_pdvalues.T:
+                plt.plot(
+                    feature.grid,
+                    mcmc_pds,
+                    color=color,
+                    alpha=0.2,
+                    linewidth=1,
+                )
+            plt.plot(feature.grid, median, color=color, label=feature_name)
+    plt.xlabel("Feature value")
+    plt.ylabel(f"MLP Output")
+    plt.legend()
+    plt.savefig(os.path.join(output_dir, "PDPs-continuous.svg"))
+    plt.close()
+
+    plot_data = []
+    grid_labels = []
+    list_labels = []
+    for feature_idx, (feature_name, feature) in enumerate(features.items()):
+        feature_pdvalues = pdvalues[feature_idx]  # shape: (n_grid_points, n_mcmcs)
+        if feature.is_categorical:
+            for i, grid_point in enumerate(feature.grid):
+                plot_data.extend(feature_pdvalues[i])
+                grid_labels.extend([grid_point] * len(feature_pdvalues[i]))
+                list_labels.extend([feature_name] * len(feature_pdvalues[i]))
+    if not (any(feature.is_categorical for feature in features.values())):
+        return
+    sns.violinplot(
+        x=grid_labels,
+        y=plot_data,
+        hue=list_labels,
+        split=False,
+        cut=0,
+        palette={
+            feature_name: "red" if feature.is_relevant else "gray"
+            for feature_name, feature in features.items()
+            if feature.is_categorical
+        },
+    )
+    plt.xlabel("Feature value")
+    plt.ylabel(f"MLP Output")
+    plt.legend()
+    plt.savefig(os.path.join(output_dir, "PDPs-categorical.svg"))
+    plt.close()
+
+
+def plot_shap_features_importance(
+    weights: list[list[Weights]],  # shape: (n_mcmcs, n_weights_samples, ...)
+    features: dict[str, Feature],
+    output_file: str,
+    hidden_activation: ActivationFunction = sigmoid,
+    output_activation: ActivationFunction = sigmoid,
+):
+    features_grid = [feature.grid for feature in features.values()]
+    inputs = list(product(*features_grid))
+
+    def _get_median_shap_features_importance(
+        weights: list[Weights],
+    ) -> list[list[float]]:
+        features_importance = np.array(
+            [
+                get_shap_features_importance(
+                    weights=w,
+                    inputs=inputs,
+                    hidden_activation=hidden_activation,
+                    output_activation=output_activation,
+                )
+                for w in weights
+            ]
+        )  # shape: (n_weights_samples, n_features)
+        return np.median(features_importance, axis=0).tolist()  # shape: (n_features,)
+
+    jobs = Parallel(n_jobs=-1, return_as="generator_unordered")(
+        delayed(_get_median_shap_features_importance)(w) for w in weights
+    )
+    features_importance_distribution = np.array(
+        [job for job in tqdm(jobs, total=len(weights), desc="Evaluating SHAPs")]
+    )  # shape: (n_mcmcs, n_features)
+    features_importance_distribution /= features_importance_distribution.sum(
+        axis=1, keepdims=True
+    )
+
+    for i, (feature_name, feature) in enumerate(features.items()):
+        sns.violinplot(
+            y=features_importance_distribution[:, i],
+            x=[feature_name] * features_importance_distribution.shape[0],
+            cut=0,
+            color="red" if feature.is_relevant else "gray",
+        )
+    plt.xlabel("Feature")
+    plt.ylabel("Importance")
+    plt.savefig(output_file)
+    plt.close()

bella_companion/simulations/generate_data.py

@@ -0,0 +1,25 @@
+import os
+
+from phylogenie import generate_trees
+
+from bella_companion.simulations.scenarios import SCENARIOS, ScenarioType
+
+N_TREES = 100
+MIN_TIPS = 200
+MAX_TIPS = 500
+
+
+def generate_data():
+    base_output_dir = os.environ["BELLA_SIMULATIONS_DATA_DIR"]
+    for scenario_name, scenario in SCENARIOS.items():
+        generate_trees(
+            output_dir=os.path.join(base_output_dir, scenario_name),
+            n_trees=N_TREES,
+            events=scenario.events,
+            init_state=scenario.init_state,
+            sampling_probability_at_present=int(scenario.type == ScenarioType.FBD),
+            max_time=scenario.max_time,
+            min_tips=MIN_TIPS,
+            max_tips=MAX_TIPS,
+            seed=42,
+        )

bella_companion/simulations/run_beast.py

@@ -0,0 +1,92 @@
+import json
+import os
+from collections import defaultdict
+from glob import glob
+from pathlib import Path
+
+from numpy.random import default_rng
+from phylogenie import Tree, load_newick
+from phylogenie.utils import get_node_depths
+from tqdm import tqdm
+
+import config as cfg
+from bella_companion.simulations.scenarios import SCENARIOS, ScenarioType
+from bella_companion.utils import run_sbatch
+
+
+def main():
+    rng = default_rng(42)
+    job_ids = {}
+    for scenario_name, scenario in SCENARIOS.items():
+        job_ids[scenario_name] = defaultdict(dict)
+        data_dir = cfg.SIMULATED_DATA_DIR / scenario_name
+        inference_configs_dir = (
+            scenario_name.split("_")[0] if "_" in scenario_name else scenario_name
+        )
+        for tree_file in tqdm(
+            glob(str(data_dir / "*.nwk")),
+            desc=f"Submitting BEAST2 jobs for {scenario_name}",
+        ):
+            tree_id = Path(tree_file).stem
+            for model in ["Nonparametric", "GLM"] + [
+                f"MLP-{hidden_nodes}" for hidden_nodes in ["3_2", "16_8", "32_16"]
+            ]:
+                outputs_dir = cfg.BEAST_OUTPUTS_DIR / scenario_name / model
+                os.makedirs(outputs_dir, exist_ok=True)
+                beast_args = [
+                    f"-D treeFile={tree_file},treeID={tree_id}",
+                    f"-prefix {outputs_dir}{os.sep}",
+                ]
+                beast_args.extend(
+                    [
+                        f'-D {key}="{value}"'
+                        for key, value in scenario.beast_args.items()
+                    ]
+                )
+                beast_args.append(
+                    f'-D randomPredictor="{" ".join(map(str, scenario.get_random_predictor(rng)))}"'
+                )
+                if scenario.type == ScenarioType.EPI:
+                    tree = load_newick(tree_file)
+                    assert isinstance(tree, Tree)
+                    beast_args.append(
+                        f"-D lastSampleTime={max(get_node_depths(tree).values())}"
+                    )
+
+                if model in ["Nonparametric", "GLM"]:
+                    command = " ".join(
+                        [
+                            cfg.RUN_BEAST,
+                            *beast_args,
+                            str(
+                                cfg.BEAST_CONFIGS_DIR
+                                / inference_configs_dir
+                                / f"{model}.xml"
+                            ),
+                        ]
+                    )
+                else:
+                    nodes = model.split("-")[1].split("_")
+                    command = " ".join(
+                        [
+                            cfg.RUN_BEAST,
+                            *beast_args,
+                            f'-D nodes="{" ".join(map(str, nodes))}"',
+                            str(
+                                cfg.BEAST_CONFIGS_DIR
+                                / inference_configs_dir
+                                / "MLP.xml"
+                            ),
+                        ]
+                    )
+
+                job_ids[scenario_name][model][tree_id] = run_sbatch(
+                    command, cfg.SBATCH_LOGS_DIR / scenario_name / model / tree_id
+                )
+
+    with open(cfg.BEAST_OUTPUTS_DIR / "simulations_job_ids.json", "w") as f:
+        json.dump(job_ids, f)
+
+
+if __name__ == "__main__":
+    main()

bella_companion/simulations/scenarios/__init__.py

@@ -0,0 +1,20 @@
+from bella_companion.simulations.scenarios.epi_multitype import SCENARIO as EPI_MULTITYPE_SCENARIO
+from bella_companion.simulations.scenarios.epi_skyline import SCENARIOS as EPI_SKYLINE_SCENARIOS
+from bella_companion.simulations.scenarios.fbd_2traits import SCENARIO as FBD_2TRAITS_SCENARIO
+from bella_companion.simulations.scenarios.fbd_no_traits import SCENARIOS as FBD_NO_TRAITS_SCENARIOS
+from bella_companion.simulations.scenarios.scenario import Scenario, ScenarioType
+
+SCENARIOS = {
+    **{
+        f"epi-skyline_{i}": scenario
+        for i, scenario in enumerate(EPI_SKYLINE_SCENARIOS, start=1)
+    },
+    "epi-multitype": EPI_MULTITYPE_SCENARIO,
+    **{
+        f"fbd-no-traits_{i}": scenario
+        for i, scenario in enumerate(FBD_NO_TRAITS_SCENARIOS, start=1)
+    },
+    "fbd-2traits": FBD_2TRAITS_SCENARIO,
+}
+
+__all__ = ["SCENARIOS", "Scenario", "ScenarioType"]

bella_companion/simulations/scenarios/common.py

@@ -0,0 +1,29 @@
+from typing import Any
+
+import numpy as np
+from numpy.random import Generator
+
+EPI_MAX_TIME = 250
+EPI_SAMPLING_PROPORTION = 0.15
+BECOME_UNINFECTIOUS_RATE = 0.07
+
+FBD_MAX_TIME = 35
+FBD_SAMPLING_RATE = 0.2
+FBD_RATE_UPPER = 2
+
+
+def get_start_type_prior_probabilities(types: list[str], init_type: str):
+    start_type_prior_probabilities = ["0"] * len(types)
+    start_type_prior_probabilities[types.index(init_type)] = "1"
+    return " ".join(start_type_prior_probabilities)
+
+
+def get_random_time_series_predictor(rng: Generator, n_time_bins: int) -> list[float]:
+    return np.cumsum(rng.normal(size=n_time_bins)).tolist()
+
+
+def get_prior_params(target: str, upper: float, n: int) -> dict[str, Any]:
+    return {
+        f"{target}Upper": upper,
+        f"{target}Init": " ".join([str(upper / 2)] * n),
+    }

bella_companion/simulations/scenarios/epi_multitype.py

@@ -0,0 +1,68 @@
+import numpy as np
+from numpy.random import Generator
+from phylogenie import get_epidemiological_events
+
+from bella_companion.simulations.features import Feature
+from bella_companion.simulations.scenarios.common import (
+    BECOME_UNINFECTIOUS_RATE,
+    EPI_MAX_TIME,
+    EPI_SAMPLING_PROPORTION,
+    get_prior_params,
+    get_start_type_prior_probabilities,
+)
+from bella_companion.simulations.scenarios.scenario import Scenario, ScenarioType
+
+
+def _get_random_predictor(rng: Generator) -> list[float]:
+    return rng.uniform(-1, 1, N_TYPE_PAIRS).tolist()
+
+
+TYPES = ["A", "B", "C", "D", "E"]
+_REPRODUCTION_NUMBERS = [0.8, 1.0, 1.2, 1.4, 1.6]
+_INIT_TYPE = "C"
+N_TYPES = len(TYPES)
+N_TYPE_PAIRS = N_TYPES * (N_TYPES - 1)
+MIGRATION_PREDICTOR = np.random.default_rng(42).uniform(-1, 1, (N_TYPES, N_TYPES - 1))
+_MIGRATION_SIGMOID_AMPLITUDE = 0.04
+_MIGRATION_SIGMOID_SCALE = -8
+MIGRATION_RATES = _MIGRATION_SIGMOID_AMPLITUDE / (
+    1 + np.exp(_MIGRATION_SIGMOID_SCALE * MIGRATION_PREDICTOR)
+)
+MIGRATION_RATE_UPPER = 0.05
+
+SCENARIO = Scenario(
+    type=ScenarioType.EPI,
+    max_time=EPI_MAX_TIME,
+    init_state=_INIT_TYPE,
+    events=get_epidemiological_events(
+        states=TYPES,
+        sampling_proportions=EPI_SAMPLING_PROPORTION,
+        reproduction_numbers=_REPRODUCTION_NUMBERS,
+        become_uninfectious_rates=BECOME_UNINFECTIOUS_RATE,
+        migration_rates=MIGRATION_RATES.tolist(),
+    ),
+    get_random_predictor=_get_random_predictor,
+    beast_args={
+        "types": ",".join(TYPES),
+        "startTypePriorProbs": get_start_type_prior_probabilities(TYPES, _INIT_TYPE),
+        "processLength": EPI_MAX_TIME,
+        **get_prior_params("migrationRate", MIGRATION_RATE_UPPER, N_TYPE_PAIRS),
+        "reproductionNumber": " ".join(map(str, _REPRODUCTION_NUMBERS)),
+        "becomeUninfectiousRate": BECOME_UNINFECTIOUS_RATE,
+        "samplingProportion": EPI_SAMPLING_PROPORTION,
+        "migrationPredictor": " ".join(map(str, MIGRATION_PREDICTOR.flatten())),
+    },
+    targets={
+        "migrationRate": {
+            f"migrationRateSP{t1}_to_{t2}": MIGRATION_RATES[i, j]
+            for i, t1 in enumerate(TYPES)
+            for j, t2 in enumerate([t for t in TYPES if t != t1])
+        }
+    },
+    features={
+        "migrationRate": {
+            "migrationPredictor": Feature(is_binary=False, is_relevant=True),
+            "randomPredictor": Feature(is_binary=False, is_relevant=False),
+        }
+    },
+)

bella_companion/simulations/scenarios/epi_skyline.py

@@ -0,0 +1,65 @@
+from functools import partial
+
+import numpy as np
+from phylogenie import SkylineParameter, get_epidemiological_events
+
+from bella_companion.simulations.features import Feature
+from bella_companion.simulations.scenarios.common import (
+    BECOME_UNINFECTIOUS_RATE,
+    EPI_MAX_TIME,
+    EPI_SAMPLING_PROPORTION,
+    get_prior_params,
+    get_random_time_series_predictor,
+)
+from bella_companion.simulations.scenarios.scenario import Scenario, ScenarioType
+
+
+def _get_scenario(reproduction_number: list[float]) -> Scenario:
+    n_time_bins = len(reproduction_number)
+    change_times = np.linspace(0, EPI_MAX_TIME, n_time_bins + 1)[1:-1].tolist()
+    return Scenario(
+        type=ScenarioType.EPI,
+        max_time=EPI_MAX_TIME,
+        events=get_epidemiological_events(
+            states=["X"],
+            sampling_proportions=EPI_SAMPLING_PROPORTION,
+            reproduction_numbers=SkylineParameter(reproduction_number, change_times),
+            become_uninfectious_rates=BECOME_UNINFECTIOUS_RATE,
+        ),
+        get_random_predictor=partial(
+            get_random_time_series_predictor, n_time_bins=n_time_bins
+        ),
+        beast_args={
+            "processLength": EPI_MAX_TIME,
+            "changeTimes": " ".join(map(str, change_times)),
+            **get_prior_params(
+                "reproductionNumber", REPRODUCTION_NUMBER_UPPER, n_time_bins
+            ),
+            "becomeUninfectiousRate": BECOME_UNINFECTIOUS_RATE,
+            "samplingProportion": EPI_SAMPLING_PROPORTION,
+            "timePredictor": " ".join(map(str, np.linspace(0, 1, n_time_bins))),
+        },
+        targets={
+            "reproductionNumber": {
+                f"reproductionNumberSPi{i}": r
+                for i, r in enumerate(reproduction_number)
+            }
+        },
+        features={
+            "reproductionNumber": {
+                "timePredictor": Feature(
+                    is_binary=False, is_relevant=len(set(reproduction_number)) > 1
+                ),
+                "randomPredictor": Feature(is_binary=False, is_relevant=False),
+            }
+        },
+    )
+
+
+REPRODUCTION_NUMBERS: list[list[float]] = [
+    [1.2] * 10,
+    np.linspace(1.5, 1.0, 10).tolist(),
+    np.linspace(1.2, 1.5, 5).tolist() + np.linspace(1.5, 1.0, 5).tolist(),
+]
+REPRODUCTION_NUMBER_UPPER = 5
+SCENARIOS = [_get_scenario(r) for r in REPRODUCTION_NUMBERS]

bella_companion/simulations/scenarios/fbd_2traits.py

@@ -0,0 +1,101 @@
+import numpy as np
+from numpy.random import Generator
+from phylogenie import SkylineVector, get_canonical_events
+
+from bella_companion.simulations.features import Feature
+from bella_companion.simulations.scenarios.common import (
+    FBD_MAX_TIME,
+    FBD_RATE_UPPER,
+    FBD_SAMPLING_RATE,
+    get_prior_params,
+    get_random_time_series_predictor,
+    get_start_type_prior_probabilities,
+)
+from bella_companion.simulations.scenarios.scenario import Scenario, ScenarioType
+
+
+def _get_random_predictor(rng: Generator) -> list[float]:
+    return np.repeat(
+        get_random_time_series_predictor(rng, N_TIME_BINS), N_STATES
+    ).tolist()
+
+
+STATES = ["00", "01", "10", "11"]
+_INIT_STATE = "00"
+N_STATES = len(STATES)
+N_TIME_BINS = 10
+_CHANGE_TIMES = np.linspace(0, FBD_MAX_TIME, N_TIME_BINS + 1)[1:-1].tolist()
+BIRTH_RATE_TRAIT1_UNSET = np.linspace(0.6, 0.1, N_TIME_BINS).tolist()
+BIRTH_RATE_TRAIT1_SET = np.linspace(0.3, 0.05, N_TIME_BINS).tolist()
+DEATH_RATE_TRAIT1_UNSET = np.linspace(0.1, 0.4, N_TIME_BINS).tolist()
+DEATH_RATE_TRAIT1_SET = np.linspace(0.1, 0.2, N_TIME_BINS).tolist()
+BIRTH_RATES = {
+    "00": BIRTH_RATE_TRAIT1_UNSET,
+    "01": BIRTH_RATE_TRAIT1_UNSET,
+    "10": BIRTH_RATE_TRAIT1_SET,
+    "11": BIRTH_RATE_TRAIT1_SET,
+}
+DEATH_RATES = {
+    "00": DEATH_RATE_TRAIT1_UNSET,
+    "01": DEATH_RATE_TRAIT1_UNSET,
+    "10": BIRTH_RATE_TRAIT1_SET,
+    "11": BIRTH_RATE_TRAIT1_SET,
+}
+RATES = {
+    "birth": BIRTH_RATES,
+    "death": DEATH_RATES,
+}
+_MIGRATION_RATES = (
+    np.array([[1, 1, 0], [1, 0, 1], [1, 0, 1], [0, 1, 1]]) * 0.1
+).tolist()
+
+SCENARIO = Scenario(
+    type=ScenarioType.FBD,
+    max_time=FBD_MAX_TIME,
+    init_state=_INIT_STATE,
+    events=get_canonical_events(
+        states=STATES,
+        sampling_rates=FBD_SAMPLING_RATE,
+        remove_after_sampling=False,
+        birth_rates=SkylineVector(
+            value=list(zip(*BIRTH_RATES.values())), change_times=_CHANGE_TIMES
+        ),
+        death_rates=SkylineVector(
+            value=list(zip(*DEATH_RATES.values())), change_times=_CHANGE_TIMES
+        ),
+        migration_rates=_MIGRATION_RATES,
+    ),
+    get_random_predictor=_get_random_predictor,
+    beast_args={
+        "types": ",".join(STATES),
+        "startTypePriorProbs": get_start_type_prior_probabilities(STATES, _INIT_STATE),
+        "processLength": FBD_MAX_TIME,
+        "changeTimes": " ".join(map(str, _CHANGE_TIMES)),
+        **get_prior_params("birthRate", FBD_RATE_UPPER, N_TIME_BINS * N_STATES),
+        **get_prior_params("deathRate", FBD_RATE_UPPER, N_TIME_BINS * N_STATES),
+        "samplingRate": FBD_SAMPLING_RATE,
+        "migrationRate": " ".join(map(str, np.array(_MIGRATION_RATES).flatten())),
+        "timePredictor": " ".join(
+            list(map(str, np.repeat(np.linspace(0, 1, N_TIME_BINS), N_STATES)))
+        ),
+        "trait1Predictor": " ".join(map(str, [1, 1, 0, 0] * N_TIME_BINS)),
+        "trait2Predictor": " ".join(map(str, [0, 1, 0, 1] * N_TIME_BINS)),
+    },
+    targets={
+        f"{rate}Rate": {
+            f"{rate}RateSPi{i}_{s}": values[s][i]
+            for i in range(N_TIME_BINS)
+            for s in STATES
+        }
+        for rate, values in RATES.items()
+    },
+    features={
+        f"{rate}Rate": {
+            "timePredictor": Feature(is_binary=False, is_relevant=True),
+            "trait1Predictor": Feature(is_binary=True, is_relevant=True),
+            "trait2Predictor": Feature(is_binary=True, is_relevant=False),
+            "randomPredictor": Feature(is_binary=False, is_relevant=False),
+        }
+        for rate in RATES
+    },
+)