bella-companion 0.0.4__py3-none-any.whl → 0.0.5__py3-none-any.whl

bella_companion/cli.py

@@ -4,7 +4,12 @@ from pathlib import Path
 
 from dotenv import load_dotenv
 
-from bella_companion.simulations import generate_data, run_beast, summarize_logs
+from bella_companion.simulations import (
+    generate_data,
+    generate_figures,
+    run_beast,
+    summarize_logs,
+)
 
 
 def main():
@@ -18,16 +23,20 @@ def main():
     subparsers = parser.add_subparsers(dest="command", required=True)
 
     subparsers.add_parser(
-        "generate-simulations-data", help="Generate simulated data."
+        "generate-simulation-data", help="Generate simulated data."
     ).set_defaults(func=generate_data)
 
     subparsers.add_parser(
-        "run-beast-simulations", help="Run BEAST2 on simulated data."
+        "beast-run-simulations", help="Run BEAST2 on simulated data."
     ).set_defaults(func=run_beast)
 
     subparsers.add_parser(
-        "summarize-simulation-logs", help="Summarize simulation logs."
+        "summarize-simulation-logs", help="Summarize BEAST2 logs for simulations data."
     ).set_defaults(func=summarize_logs)
 
+    subparsers.add_parser(
+        "generate-simulation-figures", help="Generate figures for simulations data."
+    ).set_defaults(func=generate_figures)
+
     args = parser.parse_args()
     args.func()

bella_companion/simulations/__init__.py

@@ -1,5 +1,6 @@
 from bella_companion.simulations.generate_data import generate_data
+from bella_companion.simulations.generate_figures import generate_figures
 from bella_companion.simulations.run_beast import run_beast
 from bella_companion.simulations.summarize_logs import summarize_logs
 
-__all__ = ["generate_data", "run_beast", "summarize_logs"]
+__all__ = ["generate_data", "generate_figures", "run_beast", "summarize_logs"]

bella_companion/simulations/figures/__init__.py

@@ -0,0 +1,21 @@
+from bella_companion.simulations.figures.epi_multitype_results import (
+    plot_epi_multitype_results,
+)
+from bella_companion.simulations.figures.epi_skyline_results import (
+    plot_epi_skyline_results,
+)
+from bella_companion.simulations.figures.fbd_2traits_results import (
+    plot_fbd_2traits_results,
+)
+from bella_companion.simulations.figures.fbd_no_traits_results import (
+    plot_fbd_no_traits_results,
+)
+from bella_companion.simulations.figures.scenarios import plot_scenarios
+
+__all__ = [
+    "plot_epi_multitype_results",
+    "plot_epi_skyline_results",
+    "plot_fbd_2traits_results",
+    "plot_fbd_no_traits_results",
+    "plot_scenarios",
+]

bella_companion/simulations/figures/epi_multitype_results.py

@@ -0,0 +1,81 @@
+import os
+from functools import partial
+from pathlib import Path
+
+import joblib
+import matplotlib.pyplot as plt
+import numpy as np
+import polars as pl
+from lumiere.backend import relu, sigmoid
+
+from bella_companion.simulations.figures.explain import (
+    plot_partial_dependencies,
+    plot_shap_features_importance,
+)
+from bella_companion.simulations.scenarios.epi_multitype import (
+    MIGRATION_PREDICTOR,
+    MIGRATION_RATE_UPPER,
+    MIGRATION_RATES,
+    SCENARIO,
+)
+
+
+def _plot_predictions(log_summary: pl.DataFrame, output_dir: Path):
+    targets = SCENARIO.targets["migrationRate"]
+    estimates = np.array(
+        [log_summary[f"{target}_median"].median() for target in targets]
+    )
+    lower = np.array([log_summary[f"{target}_lower"].median() for target in targets])
+    upper = np.array([log_summary[f"{target}_upper"].median() for target in targets])
+
+    sort_idx = np.argsort(MIGRATION_PREDICTOR.flatten())
+    predictors = MIGRATION_PREDICTOR.flatten()[sort_idx]
+    rates = MIGRATION_RATES.flatten()[sort_idx]
+    estimates = estimates[sort_idx]
+    lower = lower[sort_idx]
+    upper = upper[sort_idx]
+
+    plt.errorbar(  # pyright: ignore
+        predictors,
+        estimates,
+        yerr=[estimates - lower, upper - estimates],
+        fmt="o",
+        color="C2",
+        elinewidth=2,
+        capsize=5,
+    )
+    plt.plot(  # pyright: ignore
+        predictors, rates, linestyle="--", marker="o", color="k"
+    )
+
+    plt.xlabel("Migration predictor")  # pyright: ignore
+    plt.ylabel("Migration rate")  # pyright: ignore
+    plt.savefig(output_dir / "predictions.svg")  # pyright: ignore
+    plt.close()
+
+
+def plot_epi_multitype_results():
+    output_dir = Path(os.environ["BELLA_FIGURES_DIR"]) / "epi-multitype"
+    os.makedirs(output_dir, exist_ok=True)
+
+    log_dir = Path(os.environ["BELLA_LOG_SUMMARIES_DIR"]) / "epi-multitype"
+    model = "MLP-32_16"
+    log_summary = pl.read_csv(log_dir / f"{model}.csv")
+    weights = joblib.load(log_dir / f"{model}.weights.pkl")
+    weights = [w["migrationRate"] for w in weights]
+
+    _plot_predictions(log_summary, output_dir)
+    plot_partial_dependencies(
+        weights=weights,
+        features=SCENARIO.features["migrationRate"],
+        output_dir=output_dir,
+        hidden_activation=relu,
+        output_activation=partial(sigmoid, upper=MIGRATION_RATE_UPPER),
+    )
+    plot_shap_features_importance(
+        weights=weights,
+        features=SCENARIO.features["migrationRate"],
+        output_file=output_dir / "shap_values.svg",
+        hidden_activation=relu,
+        output_activation=partial(sigmoid, upper=MIGRATION_RATE_UPPER),
+    )

bella_companion/simulations/figures/epi_skyline_results.py

@@ -0,0 +1,46 @@
+import os
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+import polars as pl
+
+from bella_companion.simulations.scenarios.epi_skyline import REPRODUCTION_NUMBERS
+from bella_companion.utils import (
+    plot_coverage_per_time_bin,
+    plot_maes_per_time_bin,
+    step,
+)
+
+
+def plot_epi_skyline_results():
+    output_dir = Path(os.environ["BELLA_FIGURES_DIR"]) / "epi-skyline-results"
+    os.makedirs(output_dir, exist_ok=True)
+
+    for i, reproduction_number in enumerate(REPRODUCTION_NUMBERS, start=1):
+        summaries_dir = Path(os.environ["BELLA_LOG_SUMMARIES_DIR"]) / f"epi-skyline_{i}"
+        logs_summaries = {
+            "Nonparametric": pl.read_csv(summaries_dir / "Nonparametric.csv"),
+            "GLM": pl.read_csv(summaries_dir / "GLM.csv"),
+            "MLP": pl.read_csv(summaries_dir / "MLP-16_8.csv"),
+        }
+        true_values = {"reproductionNumberSP": reproduction_number}
+
+        for log_summary in logs_summaries.values():
+            step(
+                [
+                    float(np.median(log_summary[f"reproductionNumberSPi{i}_median"]))
+                    for i in range(len(reproduction_number))
+                ]
+            )
+        step(reproduction_number, color="k", linestyle="--")
+        plt.ylabel("Reproduction number")  # pyright: ignore
+        plt.savefig(output_dir / f"predictions-{i}.svg")  # pyright: ignore
+        plt.close()
+
+        plot_coverage_per_time_bin(
+            logs_summaries, true_values, output_dir / f"coverage-{i}.svg"
+        )
+        plot_maes_per_time_bin(
+            logs_summaries, true_values, output_dir / f"maes-{i}.svg"
+        )

bella_companion/simulations/figures/explain/__init__.py

@@ -0,0 +1,6 @@
+from bella_companion.simulations.figures.explain.pdp import plot_partial_dependencies
+from bella_companion.simulations.figures.explain.shap import (
+    plot_shap_features_importance,
+)
+
+__all__ = ["plot_partial_dependencies", "plot_shap_features_importance"]

bella_companion/simulations/figures/explain/pdp.py

@@ -0,0 +1,101 @@
+import os
+from functools import partial
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+import seaborn as sns
+from joblib import Parallel, delayed
+from lumiere.typing import ActivationFunction, Weights
+from tqdm import tqdm
+
+from bella_companion.simulations.features import Feature
+from bella_companion.utils import get_median_partial_dependence_values
+
+
+def plot_partial_dependencies(
+    weights: list[list[Weights]],  # shape: (n_runs, n_weights_samples, ...)
+    features: dict[str, Feature],
+    output_dir: Path,
+    hidden_activation: ActivationFunction,
+    output_activation: ActivationFunction,
+):
+    os.makedirs(output_dir, exist_ok=True)
+
+    continuous_grid: list[float] = np.linspace(0, 1, 10).tolist()
+    features_grid: list[list[float]] = [
+        [0, 1] if feature.is_binary else continuous_grid
+        for feature in features.values()
+    ]
+    jobs = Parallel(n_jobs=-1, return_as="generator_unordered")(
+        delayed(
+            partial(
+                get_median_partial_dependence_values,
+                features_grid=features_grid,
+                hidden_activation=hidden_activation,
+                output_activation=output_activation,
+            )
+        )(w)
+        for w in weights
+    )
+    pdvalues = [
+        job for job in tqdm(jobs, total=len(weights), desc="Evaluating PDPs")
+    ]  # shape: (n_runs, n_features, n_grid_points)
+    pdvalues = [
+        np.array(mcmc_pds).T for mcmc_pds in zip(*pdvalues)
+    ]  # shape: (n_features, n_grid_points, n_runs)
+
+    if any(not feature.is_binary for feature in features.values()):
+        for (feature_name, feature), feature_pdvalues in zip(
+            features.items(), pdvalues
+        ):
+            if not feature.is_binary:
+                color = "red" if feature.is_relevant else "gray"
+                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(  # pyright: ignore
+                    continuous_grid, lower, high, alpha=0.25, color=color
+                )
+                for mcmc_pds in feature_pdvalues.T:
+                    plt.plot(  # pyright: ignore
+                        continuous_grid, mcmc_pds, color=color, alpha=0.2, linewidth=1
+                    )
+                plt.plot(  # pyright: ignore
+                    continuous_grid, median, color=color, label=feature_name
+                )
+        plt.xlabel("Feature value")  # pyright: ignore
+        plt.ylabel("MLP Output")  # pyright: ignore
+        plt.legend()  # pyright: ignore
+        plt.savefig(output_dir / "PDPs-continuous.svg")  # pyright: ignore
+        plt.close()
+
+    if any(feature.is_binary for feature in features.values()):
+        data: list[float] = []
+        grid_labels: list[int] = []
+        feature_labels: list[str] = []
+        for (feature_name, feature), feature_pdvalues in zip(
+            features.items(), pdvalues
+        ):
+            if feature.is_binary:
+                for i in [0, 1]:
+                    data.extend(feature_pdvalues[i])
+                    grid_labels.extend([i] * len(feature_pdvalues[i]))
+                    feature_labels.extend([feature_name] * len(feature_pdvalues[i]))
+        sns.violinplot(
+            x=grid_labels,
+            y=data,
+            hue=feature_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_binary
+            },
+        )
+        plt.xlabel("Feature value")  # pyright: ignore
+        plt.ylabel("MLP Output")  # pyright: ignore
+        plt.legend()  # pyright: ignore
+        plt.savefig(output_dir / "PDPs-categorical.svg")  # pyright: ignore
+        plt.close()

bella_companion/simulations/figures/explain/shap.py

@@ -0,0 +1,56 @@
+from functools import partial
+from itertools import product
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+import seaborn as sns
+from joblib import Parallel, delayed
+from lumiere.typing import ActivationFunction, Weights
+from tqdm import tqdm
+
+from bella_companion.simulations.features import Feature
+from bella_companion.utils import get_median_shap_features_importance
+
+
+def plot_shap_features_importance(
+    weights: list[list[Weights]],  # shape: (n_runs, n_weights_samples, ...)
+    features: dict[str, Feature],
+    output_file: Path,
+    hidden_activation: ActivationFunction,
+    output_activation: ActivationFunction,
+):
+    continuous_grid: list[float] = np.linspace(0, 1, 10).tolist()
+    features_grid: list[list[float]] = [
+        [0, 1] if feature.is_binary else continuous_grid
+        for feature in features.values()
+    ]
+    inputs = list(product(*features_grid))
+
+    jobs = Parallel(n_jobs=-1, return_as="generator_unordered")(
+        delayed(
+            partial(
+                get_median_shap_features_importance,
+                inputs=inputs,
+                hidden_activation=hidden_activation,
+                output_activation=output_activation,
+            )
+        )(w)
+        for w in weights
+    )
+    features_importances = np.array(
+        [job for job in tqdm(jobs, total=len(weights), desc="Evaluating SHAPs")]
+    )  # shape: (n_runs, n_features)
+    features_importances /= features_importances.sum(axis=1, keepdims=True)
+
+    for i, (feature_name, feature) in enumerate(features.items()):
+        sns.violinplot(
+            y=features_importances[:, i],
+            x=[feature_name] * len(features_importances),
+            cut=0,
+            color="red" if feature.is_relevant else "gray",
+        )
+    plt.xlabel("Feature")  # pyright: ignore
+    plt.ylabel("Importance")  # pyright: ignore
+    plt.savefig(output_file)  # pyright: ignore
+    plt.close()

bella_companion/simulations/figures/fbd_2traits_results.py

@@ -0,0 +1,83 @@
+import os
+from functools import partial
+from pathlib import Path
+
+import joblib
+import matplotlib.pyplot as plt
+import numpy as np
+import polars as pl
+from lumiere.backend import relu, sigmoid
+
+from bella_companion.simulations.figures.explain import (
+    plot_partial_dependencies,
+    plot_shap_features_importance,
+)
+from bella_companion.simulations.scenarios.fbd_2traits import (
+    FBD_RATE_UPPER,
+    N_TIME_BINS,
+    RATES,
+    SCENARIO,
+    STATES,
+)
+from bella_companion.utils import step
+
+
+def _plot_predictions(log_summary: pl.DataFrame, output_dir: Path):
+    for rate, state_rates in RATES.items():
+        label = r"\lambda" if rate == "birth" else r"\mu"
+        for state in STATES:
+            estimates = [
+                float(np.median(log_summary[f"{rate}RateSPi{i}_{state}_median"]))
+                for i in range(N_TIME_BINS)
+            ]
+            step(
+                estimates,
+                label=rf"${label}_{{{state[0]},{state[1]}}}$",
+                reverse_xticks=True,
+            )
+        step(
+            state_rates["00"],
+            color="k",
+            linestyle="dashed",
+            label=rf"${label}_{{0,0}}$ = ${label}_{{0,1}}$",
+            reverse_xticks=True,
+        )
+        step(
+            state_rates["10"],
+            color="gray",
+            linestyle="dashed",
+            label=rf"${label}_{{1,0}}$ = ${label}_{{1,1}}$",
+            reverse_xticks=True,
+        )
+        plt.legend()  # pyright: ignore
+        plt.ylabel(rf"${label}$")  # pyright: ignore
+        plt.savefig(output_dir / rate / "predictions.svg")  # pyright: ignore
+        plt.close()
+
+
+def plot_fbd_2traits_results():
+    output_dir = Path(os.environ["BELLA_FIGURES_DIR"]) / "fbd-2traits"
+
+    log_dir = Path(os.environ["BELLA_LOG_SUMMARIES_DIR"]) / "fbd-2traits"
+    model = "MLP-32_16"
+    log_summary = pl.read_csv(log_dir / f"{model}.csv")
+    weights = joblib.load(log_dir / f"{model}.weights.pkl")
+
+    for rate in RATES:
+        os.makedirs(output_dir / rate, exist_ok=True)
+        plot_partial_dependencies(
+            weights=[w[f"{rate}Rate"] for w in weights],
+            features=SCENARIO.features[f"{rate}Rate"],
+            output_dir=output_dir / rate,
+            hidden_activation=relu,
+            output_activation=partial(sigmoid, upper=FBD_RATE_UPPER),
+        )
+        plot_shap_features_importance(
+            weights=[w[f"{rate}Rate"] for w in weights],
+            features=SCENARIO.features[f"{rate}Rate"],
+            output_file=output_dir / rate / "shap_values.svg",
+            hidden_activation=relu,
+            output_activation=partial(sigmoid, upper=FBD_RATE_UPPER),
+        )
+
+    _plot_predictions(log_summary, output_dir)

bella_companion/simulations/figures/fbd_no_traits_results.py

@@ -0,0 +1,58 @@
+import os
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+import polars as pl
+
+from bella_companion.simulations.scenarios.fbd_no_traits import RATES
+from bella_companion.utils import (
+    plot_coverage_per_time_bin,
+    plot_maes_per_time_bin,
+    step,
+)
+
+
+def plot_fbd_no_traits_results():
+    output_dir = Path(os.environ["BELLA_FIGURES_DIR"]) / "fbd-no-traits-predictions"
+    os.makedirs(output_dir, exist_ok=True)
+
+    for i, rates in enumerate(RATES, start=1):
+        summaries_dir = (
+            Path(os.environ["BELLA_LOG_SUMMARIES_DIR"]) / f"fbd-no-traits_{i}"
+        )
+        logs_summaries = {
+            "Nonparametric": pl.read_csv(summaries_dir / "Nonparametric.csv"),
+            "GLM": pl.read_csv(summaries_dir / "GLM.csv"),
+            "MLP": pl.read_csv(summaries_dir / "MLP-16_8.csv"),
+        }
+        true_values = {"birthRateSP": rates["birth"], "deathRateSP": rates["death"]}
+
+        for id, rate in true_values.items():
+            for log_summary in logs_summaries.values():
+                step(
+                    [
+                        float(np.median(log_summary[f"{id}i{i}_median"]))
+                        for i in range(len(rate))
+                    ],
+                    reverse_xticks=True,
+                )
+            step(rate, color="k", linestyle="--", reverse_xticks=True)
+            plt.ylabel(  # pyright: ignore
+                r"$\lambda$" if id == "birthRateSP" else r"$\mu$"
+            )
+            plt.savefig(output_dir / f"{id}-predictions-{i}-.svg")  # pyright: ignore
+            plt.close()
+
+        plot_coverage_per_time_bin(
+            logs_summaries,
+            true_values,
+            output_dir / f"coverage-{i}.svg",
+            reverse_xticks=True,
+        )
+        plot_maes_per_time_bin(
+            logs_summaries,
+            true_values,
+            output_dir / f"maes-{i}.svg",
+            reverse_xticks=True,
+        )

bella_companion/simulations/figures/scenarios.py

@@ -1,56 +1,66 @@
 import os
+from pathlib import Path
 
 import matplotlib.pyplot as plt
 import numpy as np
 
-import src.config as cfg
-from src.simulations.figures.utils import step
-from src.simulations.scenarios.epi_multitype import MIGRATION_PREDICTOR, MIGRATION_RATES
-from src.simulations.scenarios.epi_skyline import REPRODUCTION_NUMBERS
-from src.simulations.scenarios.fbd_2traits import (
+from bella_companion.simulations.scenarios.epi_multitype import (
+    MIGRATION_PREDICTOR,
+    MIGRATION_RATES,
+)
+from bella_companion.simulations.scenarios.epi_skyline import REPRODUCTION_NUMBERS
+from bella_companion.simulations.scenarios.fbd_2traits import (
     BIRTH_RATE_TRAIT1_SET,
     BIRTH_RATE_TRAIT1_UNSET,
     DEATH_RATE_TRAIT1_SET,
     DEATH_RATE_TRAIT1_UNSET,
 )
-from src.simulations.scenarios.fbd_no_traits import BIRTH_RATES, DEATH_RATES
-from src.utils import set_plt_rcparams
+from bella_companion.simulations.scenarios.fbd_no_traits import RATES
+from bella_companion.utils import step
 
 
-def main():
-    output_dir = os.path.join(cfg.FIGURES_DIR, "scenarios")
+def plot_scenarios():
+    output_dir = Path(os.environ["BELLA_FIGURES_DIR"]) / "targets"
     os.makedirs(output_dir, exist_ok=True)
 
-    set_plt_rcparams()
-
+    # -----------
+    # epi-skyline
+    # -----------
     for i, reproduction_number in enumerate(REPRODUCTION_NUMBERS, start=1):
         step(reproduction_number, color="k")
-        plt.ylabel("Reproduction number")
-        plt.savefig(os.path.join(output_dir, f"epi-skyline_{i}.svg"))
+        plt.ylabel("Reproduction number")  # pyright: ignore
+        plt.savefig(output_dir / f"epi-skyline_{i}.svg")  # pyright: ignore
         plt.close()
 
+    # -------------
+    # epi-multitype
+    # -------------
     sort_idx = np.argsort(MIGRATION_PREDICTOR.flatten())
-    plt.plot(
+    plt.plot(  # pyright: ignore
         MIGRATION_PREDICTOR.flatten()[sort_idx],
         MIGRATION_RATES.flatten()[sort_idx],
         marker="o",
         color="k",
     )
-    plt.xlabel("Migration predictor")
-    plt.ylabel("Migration rate")
-    plt.savefig(os.path.join(output_dir, "epi-multitype.svg"))
+    plt.xlabel("Migration predictor")  # pyright: ignore
+    plt.ylabel("Migration rate")  # pyright: ignore
+    plt.savefig(output_dir / "epi-multitype.svg")  # pyright: ignore
     plt.close()
 
-    for i, (birth_rate, death_rate) in enumerate(
-        zip(BIRTH_RATES, DEATH_RATES), start=1
-    ):
-        step(birth_rate, label=r"$\lambda$", reverse_xticks=True)
-        step(death_rate, label=r"$\mu$", reverse_xticks=True)
-        plt.ylabel("Rate")
-        plt.legend()
-        plt.savefig(os.path.join(output_dir, f"fbd-no-traits_{i}.svg"))
+    # -------------
+    # fbd-no-traits
+    # -------------
+    for i, rates in enumerate(RATES, start=1):
+        step(rates["birth"], label=r"$\lambda$", reverse_xticks=True)
+        step(rates["death"], label=r"$\mu$", reverse_xticks=True)
+        plt.ylabel("Rate")  # pyright: ignore
+        plt.legend()  # pyright: ignore
+        plt.savefig(output_dir / f"fbd-no-traits_{i}.svg")  # pyright: ignore
         plt.close()
 
+    # -----------
+    # fbd-2traits
+    # -----------
     step(
         BIRTH_RATE_TRAIT1_UNSET,
         label=r"$\lambda_{0,0} = \lambda_{0,1}$",
@@ -77,11 +87,7 @@ def main():
         linestyle="dashed",
         reverse_xticks=True,
     )
-    plt.ylabel("Rate")
-    plt.legend()
-    plt.savefig(os.path.join(output_dir, "fbd-2traits.svg"))
+    plt.ylabel("Rate")  # pyright: ignore
+    plt.legend()  # pyright: ignore
+    plt.savefig(output_dir / "fbd-2traits.svg")  # pyright: ignore
     plt.close()
-
-
-if __name__ == "__main__":
-    main()

bella_companion/simulations/generate_figures.py

@@ -0,0 +1,24 @@
+import matplotlib.pyplot as plt
+
+from bella_companion.simulations.figures import (
+    plot_epi_multitype_results,
+    plot_epi_skyline_results,
+    plot_fbd_2traits_results,
+    plot_fbd_no_traits_results,
+    plot_scenarios,
+)
+
+
+def generate_figures():
+    plt.rcParams["pdf.fonttype"] = 42
+    plt.rcParams["xtick.labelsize"] = 14
+    plt.rcParams["ytick.labelsize"] = 14
+    plt.rcParams["font.size"] = 14
+    plt.rcParams["figure.constrained_layout.use"] = True
+    plt.rcParams["lines.linewidth"] = 3
+
+    plot_scenarios()
+    plot_epi_skyline_results()
+    plot_epi_multitype_results()
+    plot_fbd_no_traits_results()
+    plot_fbd_2traits_results()

bella_companion/simulations/scenarios/fbd_2traits.py

@@ -38,8 +38,8 @@ BIRTH_RATES = {
 DEATH_RATES = {
     "00": DEATH_RATE_TRAIT1_UNSET,
     "01": DEATH_RATE_TRAIT1_UNSET,
-    "10": BIRTH_RATE_TRAIT1_SET,
-    "11": BIRTH_RATE_TRAIT1_SET,
+    "10": DEATH_RATE_TRAIT1_SET,
+    "11": DEATH_RATE_TRAIT1_SET,
 }
 RATES = {
     "birth": BIRTH_RATES,

bella_companion/utils/__init__.py

@@ -1,4 +1,23 @@
 from bella_companion.utils.beast import summarize_log, summarize_logs, summarize_weights
+from bella_companion.utils.explain import (
+    get_median_partial_dependence_values,
+    get_median_shap_features_importance,
+)
+from bella_companion.utils.plots import (
+    plot_coverage_per_time_bin,
+    plot_maes_per_time_bin,
+    step,
+)
 from bella_companion.utils.slurm import submit_job
 
-__all__ = ["submit_job", "summarize_log", "summarize_logs", "summarize_weights"]
+__all__ = [
+    "summarize_log",
+    "summarize_logs",
+    "summarize_weights",
+    "get_median_partial_dependence_values",
+    "get_median_shap_features_importance",
+    "plot_coverage_per_time_bin",
+    "plot_maes_per_time_bin",
+    "step",
+    "submit_job",
+]