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

bella_companion/utils/beast.py

@@ -9,7 +9,7 @@ import numpy as np
 import polars as pl
 from joblib import Parallel, delayed
 from lumiere import read_log_file, read_weights
-from lumiere.backend.typing import Weights
+from lumiere.typing import Weights
 from tqdm import tqdm
 
 from bella_companion.utils.slurm import get_job_metadata

bella_companion/utils/explain.py

@@ -0,0 +1,45 @@
+import numpy as np
+from lumiere import get_partial_dependence_values, get_shap_features_importance
+from lumiere.typing import ActivationFunction, Weights
+from numpy.typing import ArrayLike
+
+
+def get_median_partial_dependence_values(
+    weights: list[Weights],  # shape: (n_weight_samples, ...)
+    features_grid: list[list[float]],
+    hidden_activation: ActivationFunction,
+    output_activation: ActivationFunction,
+) -> list[list[float]]:  # shape: (n_features, n_grid_points)
+    pdvalues = [
+        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([pd[feature_idx] for pd in pdvalues], axis=0).tolist()
+        for feature_idx in range(len(features_grid))
+    ]
+
+
+def get_median_shap_features_importance(
+    weights: list[Weights],
+    inputs: ArrayLike,
+    hidden_activation: ActivationFunction,
+    output_activation: ActivationFunction,
+) -> list[float]:  # length: n_features
+    features_importance = np.array(
+        [
+            get_shap_features_importance(
+                weights=w,
+                inputs=inputs,
+                hidden_activation=hidden_activation,
+                output_activation=output_activation,
+            )
+            for w in weights
+        ]
+    )
+    return np.median(features_importance, axis=0).tolist()

bella_companion/utils/plots.py

@@ -0,0 +1,98 @@
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+import polars as pl
+import seaborn as sns
+
+
+def _set_time_bin_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)))  # pyright: ignore
+    plt.xlabel("Time bin")  # pyright: ignore
+
+
+def step(
+    x: list[float],
+    reverse_xticks: bool = False,
+    label: str | None = None,
+    color: str | None = None,
+    linestyle: str | None = None,
+):
+    x = [x[0], *x]
+    n = len(x)
+    plt.step(  # pyright: ignore
+        list(range(n)), x, label=label, color=color, linestyle=linestyle
+    )
+    _set_time_bin_xticks(n, reverse_xticks)
+
+
+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: Path,
+    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_time_bin_xticks(n_time_bins, reverse_xticks)
+    plt.savefig(output_filepath)  # pyright: ignore
+    plt.close()
+
+
+def plot_coverage_per_time_bin(
+    logs_summaries: dict[str, pl.DataFrame],
+    true_values: dict[str, list[float]],
+    output_filepath: Path,
+    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")  # pyright: ignore
+
+    _set_time_bin_xticks(n_time_bins, reverse_xticks)
+    plt.ylabel("Coverage")  # pyright: ignore
+    plt.ylim((0, 1.05))  # pyright: ignore
+    plt.savefig(output_filepath)  # pyright: ignore
+    plt.close()

{bella_companion-0.0.3.dist-info → bella_companion-0.0.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: bella-companion
-Version: 0.0.3
+Version: 0.0.5
 Summary: 
 Author: gabriele-marino
 Author-email: gabmarino.8601@gmail.com
@@ -10,6 +10,7 @@ Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Requires-Dist: arviz (>=0.22.0,<0.23.0)
-Requires-Dist: bella-lumiere (>=0.0.10,<0.0.11)
+Requires-Dist: bella-lumiere (>=0.0.13,<0.0.14)
 Requires-Dist: dotenv (>=0.9.9,<0.10.0)
-Requires-Dist: phylogenie (>=2.1.27,<3.0.0)
+Requires-Dist: phylogenie (>=2.1.28,<3.0.0)
+Requires-Dist: seaborn (>=0.13.2,<0.14.0)

{bella_companion-0.0.3.dist-info → bella_companion-0.0.5.dist-info}/RECORD

@@ -1,5 +1,5 @@
 bella_companion/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-bella_companion/cli.py,sha256=0sPnzGyUGo2OBZ0rj17ZGzMdwNH0o-BXKsYtCJjzGvQ,968
+bella_companion/cli.py,sha256=fjCIIsguBUnm_u1dyuo42INGTDa5lh8laiW1yZN65So,1178
 bella_companion/fbd_empirical/data/body_mass.csv,sha256=-UkKNtm9m3g4PjY3BcfdP6z5nL_I6p9cq6cgZ-bWKI8,30360
 bella_companion/fbd_empirical/data/change_times.csv,sha256=zmc9_z91-XMwKyIoP9v9dVlLcf4MeIHkQiHLjoMriOo,120
 bella_companion/fbd_empirical/data/sampling_change_times.csv,sha256=Gwi9RcMFy89RyvfxKVZ_MoKVRHOZLuwB_3LEaq8asMQ,32
@@ -9,29 +9,34 @@ bella_companion/fbd_empirical/notbooks.ipynb,sha256=O45kmz0lZENRDFbKXEWPsIKATfF5
 bella_companion/fbd_empirical/params.json,sha256=hU23LniClZL_GSBAxIEJUJgMa93AM8zdtFOq6mt3vkI,311
 bella_companion/fbd_empirical/run_beast.py,sha256=2sV2UmxOfWmbueiU6D0p3lueMYiZyIkSKYoblTMrYuA,1935
 bella_companion/fbd_empirical/summarize_logs.py,sha256=O6rhE606Wa98a8b1KKlLPjUOro1pfyqVTLdQksQMG0g,1439
-bella_companion/simulations/__init__.py,sha256=i6Fe7l5sUJY9hPxdg6L_FVhwbSPhNxQNMb-m33JlfxI,258
+bella_companion/simulations/__init__.py,sha256=EBZAcI8skNPKjrA7CjrqH9ea7DTntmydAD0RqsxNUMM,352
 bella_companion/simulations/features.py,sha256=DZOBpJGlQ0UinqUZYbEtoemZ2eQGVLV_i-DfpW31qJI,104
-bella_companion/simulations/figures/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-bella_companion/simulations/figures/epi_explainations.py,sha256=omiJgyIY-I6zcJAcyOF7GJ2pba6pMZySLkWy7OFrjFY,3093
-bella_companion/simulations/figures/epi_predictions.py,sha256=4yXwOBKxUv4kgZdI9zAMEhZ0QCNKZdkAafRQ1RTeaWg,1835
-bella_companion/simulations/figures/fbd_explainations.py,sha256=9Uj7yttpn_TH5HqycW8R-Nlky9A9aFXDXRpXQuT1L4s,3037
-bella_companion/simulations/figures/fbd_predictions.py,sha256=jdXYCLledZEWoPCIuTLhHEPMdeG6YXvf5xZnEOslv-U,2119
-bella_companion/simulations/figures/scenarios.py,sha256=vyybn3Qhfq96N8tvW0wSzpFoHHP8EIc8dkOz63o_Atw,2492
-bella_companion/simulations/figures/utils.py,sha256=sY8wFBg02fv5ugpJ80EqQishD_HEdLwhqsw2LfM7wEo,8539
+bella_companion/simulations/figures/__init__.py,sha256=aBYbJntH4egFmkSSWiVMYDEApXPYxJD7eA3TCPNNegM,658
+bella_companion/simulations/figures/epi_multitype_results.py,sha256=j85WgvN5AyAtX-CalMegr2lwlAZBmzyJxkikBPXRjCc,2629
+bella_companion/simulations/figures/epi_skyline_results.py,sha256=Ej1iGHLnLkUohx7oC9gRTLe8T--5WS-JtyViLbxraLA,1647
+bella_companion/simulations/figures/explain/__init__.py,sha256=DnmVIWO65nTT5VsWnbS7NyYgKEY_eo4oMCtCY_ML2Vk,260
+bella_companion/simulations/figures/explain/pdp.py,sha256=bSrH-TZcT2mVZIBXTUeoRvShsX2bIUyhUvgmbzXLIbg,3925
+bella_companion/simulations/figures/explain/shap.py,sha256=tjYX5H3C-IiygG9scXhLgtlCRxEBlsQ2OW5gVWL15kY,1936
+bella_companion/simulations/figures/fbd_2traits_results.py,sha256=YP_ksmh5Sj2WnoxqRuSLt-gd-W7k2lUHs4idAbvbdhw,2708
+bella_companion/simulations/figures/fbd_no_traits_results.py,sha256=O6hx_OZVSHmW0xq9T4q4oz6eskC_xcFZGMEdX8ZilrU,1942
+bella_companion/simulations/figures/scenarios.py,sha256=OKh9_-ZvzNgWsO3-Vd0Aw3ndjVf76i_OuCvsKI-5r2s,2795
 bella_companion/simulations/generate_data.py,sha256=H8OV4ZlTGZB-jXaROTPmOsK3UxRiU-GrX40l-shliw8,728
+bella_companion/simulations/generate_figures.py,sha256=layMgoj3Bfl78Ceb1oE7YirAQ8zhjDyD9IrxDRXf6go,657
 bella_companion/simulations/run_beast.py,sha256=xOuwE0w4IbOqqCSym6kHsAEhfGT2mWdA-jmUZuviMbc,3121
 bella_companion/simulations/scenarios/__init__.py,sha256=3Kl1lKcFpfb3vLX64DmSW4XCF5kXU1ZoHtstFH-ZIzU,876
 bella_companion/simulations/scenarios/common.py,sha256=_ddaSuTvEVdttGkXB4HPc2B7IB1F_GBOCW3cVOPZ-ZM,807
 bella_companion/simulations/scenarios/epi_multitype.py,sha256=GWGIiqvYwX_FrT_3RXkZKYGDht9nZ7ceHRBKUvXDPnA,2432
 bella_companion/simulations/scenarios/epi_skyline.py,sha256=JqnOVATECxBUqEbkR5lBlMI2O8k4hO6ipR8k9cHUsm0,2365
-bella_companion/simulations/scenarios/fbd_2traits.py,sha256=sCtdWyV6GQQOIhnL9Dd8NIbAR-StTwUTD9-b_BalmFQ,3552
+bella_companion/simulations/scenarios/fbd_2traits.py,sha256=G24WCAHrWPwvQeElsy4UMl1I9ALFnVQp6wXuc25Ie-g,3552
 bella_companion/simulations/scenarios/fbd_no_traits.py,sha256=R6CH0fVeQg-Iesl39pq2uY8ICVEO4VZbvUVUCGwauJU,2520
 bella_companion/simulations/scenarios/scenario.py,sha256=_FRWAyOFbw94lAzd3zCD-1ek4TrssoiXfXRQPShLiIA,620
 bella_companion/simulations/summarize_logs.py,sha256=5IdzR9IwjeF2LgZmzpuK0rPfYMct2OUgEp0QyUbUS7g,1263
-bella_companion/utils/__init__.py,sha256=_5tLPH_3GHtimNcH0Yd9Z6yIM3WkWkNApNGLzFnF6nY,222
-bella_companion/utils/beast.py,sha256=RG-iSEFuL92K6yxUV2nxdmcVqfrEiPhaYTmReW4ZoWk,2189
+bella_companion/utils/__init__.py,sha256=bb0-pCjQNwGaUgG2v9htgOLlOMiB1-kdZnSQ_QY2QAo,647
+bella_companion/utils/beast.py,sha256=5Vsv98VTE9HrY56WzUSMECjD_rIPxHTMRMD1ZzmA6wY,2181
+bella_companion/utils/explain.py,sha256=uP7HPyn2YiykAI69BQV3RooDpC6qKoCLXfp3Uibp4zk,1475
+bella_companion/utils/plots.py,sha256=dB_GiJ1HGrZ93cqODz6kB-HeDRPwlm2MkMe9rJZGnfs,3117
 bella_companion/utils/slurm.py,sha256=v5DaG7YHVyK8KRFptgGDC6I8jxEhyJuMVK9N08pZSAI,1812
-bella_companion-0.0.3.dist-info/METADATA,sha256=ARm9evUES-6JEG46mregbCHmiNOMW59VGT_msd2V6Bg,534
-bella_companion-0.0.3.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-bella_companion-0.0.3.dist-info/entry_points.txt,sha256=rSeKoAhmjnQqAYFcXBv0gAM2ViJfJe0D8_dD-fWrXeg,50
-bella_companion-0.0.3.dist-info/RECORD,,
+bella_companion-0.0.5.dist-info/METADATA,sha256=IIRIUzr8WcpcLOApct-WoFLrEBkzAxhjXPIY6HucbOk,576
+bella_companion-0.0.5.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+bella_companion-0.0.5.dist-info/entry_points.txt,sha256=rSeKoAhmjnQqAYFcXBv0gAM2ViJfJe0D8_dD-fWrXeg,50
+bella_companion-0.0.5.dist-info/RECORD,,

bella_companion/simulations/figures/epi_explainations.py

@@ -1,109 +0,0 @@
-import os
-from functools import partial
-
-import joblib
-import matplotlib.pyplot as plt
-import numpy as np
-import polars as pl
-import src.config as cfg
-from lumiere.backend import sigmoid
-from src.simulations.figures.utils import (
-    plot_partial_dependencies,
-    plot_shap_features_importance,
-)
-from src.simulations.scenarios.epi_multitype import (
-    MIGRATION_PREDICTOR,
-    MIGRATION_RATE_UPPER,
-    MIGRATION_RATES,
-    SCENARIO,
-)
-from src.utils import set_plt_rcparams
-
-
-def set_plt_rcparams():
-    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
-
-
-def _plot_predictions(log_summary: pl.DataFrame, output_dir: str):
-    sort_idx = np.argsort(MIGRATION_PREDICTOR.flatten())
-
-    estimates = np.array(
-        [
-            log_summary[f"{target}_median"].median()
-            for target in SCENARIO.targets["migrationRate"]
-        ]
-    )
-    lower = np.array(
-        [
-            log_summary[f"{target}_lower"].median()
-            for target in SCENARIO.targets["migrationRate"]
-        ]
-    )
-    upper = np.array(
-        [
-            log_summary[f"{target}_upper"].median()
-            for target in SCENARIO.targets["migrationRate"]
-        ]
-    )
-    plt.errorbar(
-        MIGRATION_PREDICTOR.flatten()[sort_idx],
-        estimates[sort_idx],
-        yerr=[
-            estimates[sort_idx] - lower[sort_idx],
-            upper[sort_idx] - estimates[sort_idx],
-        ],
-        marker="o",
-        color="C2",
-    )
-    plt.plot(
-        MIGRATION_PREDICTOR.flatten()[sort_idx],
-        estimates[sort_idx],
-        marker="o",
-        color="C2",
-    )
-    plt.plot(
-        MIGRATION_PREDICTOR.flatten()[sort_idx],
-        MIGRATION_RATES.flatten()[sort_idx],
-        linestyle="dashed",
-        marker="o",
-        color="k",
-    )
-    plt.xlabel("Migration predictor")
-    plt.ylabel("Migration rate")
-    plt.savefig(os.path.join(output_dir, "predictions.svg"))
-    plt.close()
-
-
-def main():
-    output_dir = os.path.join(cfg.FIGURES_DIR, "epi-explainations")
-    os.makedirs(output_dir, exist_ok=True)
-
-    log_dir = os.path.join(cfg.BEAST_LOGS_SUMMARIES_DIR, "epi-multitype")
-    model = "MLP-32_16"
-    log_summary = pl.read_csv(os.path.join(log_dir, f"{model}.csv"))
-    weights = joblib.load(os.path.join(log_dir, f"{model}_weights.pkl"))
-
-    set_plt_rcparams()
-
-    _plot_predictions(log_summary, output_dir)
-    plot_partial_dependencies(
-        weights=weights["migrationRate"],
-        features=SCENARIO.features["migrationRate"],
-        output_dir=output_dir,
-        output_activation=partial(sigmoid, upper=MIGRATION_RATE_UPPER),
-    )
-    plot_shap_features_importance(
-        weights=weights["migrationRate"],
-        features=SCENARIO.features["migrationRate"],
-        output_file=os.path.join(output_dir, "shap_values.svg"),
-        output_activation=partial(sigmoid, upper=MIGRATION_RATE_UPPER),
-    )
-
-
-if __name__ == "__main__":
-    main()

bella_companion/simulations/figures/epi_predictions.py

@@ -1,58 +0,0 @@
-import os
-
-import matplotlib.pyplot as plt
-import polars as pl
-
-import src.config as cfg
-from src.simulations.figures.utils import (
-    plot_coverage_per_time_bin,
-    plot_maes_per_time_bin,
-    step,
-)
-from src.simulations.scenarios.epi_skyline import REPRODUCTION_NUMBERS
-from src.utils import set_plt_rcparams
-
-
-def main():
-    output_dir = os.path.join(cfg.FIGURES_DIR, "epi-predictions")
-    os.makedirs(output_dir, exist_ok=True)
-
-    set_plt_rcparams()
-
-    for i, reproduction_number in enumerate(REPRODUCTION_NUMBERS, start=1):
-        summaries_dir = os.path.join(cfg.BEAST_LOGS_SUMMARIES_DIR, f"epi-skyline_{i}")
-        logs_summaries = {
-            "Nonparametric": pl.read_csv(
-                os.path.join(summaries_dir, "Nonparametric.csv")
-            ),
-            "GLM": pl.read_csv(os.path.join(summaries_dir, "GLM.csv")),
-            "MLP": pl.read_csv(os.path.join(summaries_dir, "MLP-16_8.csv")),
-        }
-        true_values = {"reproductionNumber": reproduction_number}
-
-        for log_summary in logs_summaries.values():
-            step(
-                [
-                    log_summary[f"reproductionNumberi{i}_median"].median()
-                    for i in range(len(reproduction_number))
-                ]
-            )
-        step(reproduction_number, color="k", linestyle="--")
-        plt.ylabel("Reproduction number")
-        plt.savefig(os.path.join(output_dir, f"epi-skyline_{i}-predictions.svg"))
-        plt.close()
-
-        plot_coverage_per_time_bin(
-            logs_summaries,
-            true_values,
-            os.path.join(output_dir, f"epi-skyline_{i}-coverage.svg"),
-        )
-        plot_maes_per_time_bin(
-            logs_summaries,
-            true_values,
-            os.path.join(output_dir, f"epi-skyline_{i}-maes.svg"),
-        )
-
-
-if __name__ == "__main__":
-    main()

bella_companion/simulations/figures/fbd_explainations.py

@@ -1,99 +0,0 @@
-import ast
-import os
-from functools import partial
-
-import joblib
-import matplotlib.pyplot as plt
-import polars as pl
-from joblib import Parallel, delayed
-from lumiere.backend import sigmoid
-
-import src.config as cfg
-from src.figures.utils import (
-    plot_partial_dependencies,
-    plot_shap_features_importance,
-    step,
-)
-from src.simulations.scenarios.fbd_2traits import (
-    BIRTH_RATE_TRAIT1_SET,
-    BIRTH_RATE_TRAIT1_UNSET,
-    DEATH_RATE_TRAIT1_SET,
-    DEATH_RATE_TRAIT1_UNSET,
-    FBD_RATE_UPPER,
-    N_TIME_BINS,
-    SCENARIO,
-    STATES,
-)
-from src.utils import set_plt_rcparams
-
-
-def _plot_predictions(log_summary: pl.DataFrame, output_dir: str):
-    for rate in ["birth", "death"]:
-        label = r"\lambda" if rate == "birth" else r"\mu"
-        rate_trait_1_set = (
-            BIRTH_RATE_TRAIT1_UNSET if rate == "birth" else DEATH_RATE_TRAIT1_UNSET
-        )
-        rate_trait_1_unset = (
-            BIRTH_RATE_TRAIT1_SET if rate == "birth" else DEATH_RATE_TRAIT1_SET
-        )
-        for state in STATES:
-            estimates = [
-                log_summary[f"{rate}Ratei{i}_{state}_median"].median()
-                for i in range(N_TIME_BINS)
-            ]
-            step(
-                estimates,
-                label=rf"${label}_{{{state[0]},{state[1]}}}$",
-                reverse_xticks=True,
-            )
-        step(
-            rate_trait_1_unset,
-            color="k",
-            linestyle="dashed",
-            label=rf"${label}_{{0,0}}$ = ${label}_{{0,1}}$",
-            reverse_xticks=True,
-        )
-        step(
-            rate_trait_1_set,
-            color="gray",
-            linestyle="dashed",
-            label=rf"${label}_{{1,0}}$ = ${label}_{{1,1}}$",
-            reverse_xticks=True,
-        )
-        plt.legend()
-        plt.ylabel(rf"${label}$")
-        plt.savefig(os.path.join(output_dir, rate, "predictions.svg"))
-        plt.close()
-
-
-def main():
-    output_dir = os.path.join(cfg.FIGURES_DIR, "fbd-explainations")
-    for rate in ["birth", "death"]:
-        os.makedirs(os.path.join(output_dir, rate), exist_ok=True)
-
-    log_dir = os.path.join(cfg.BEAST_LOGS_SUMMARIES_DIR, "fbd-2traits")
-    model = "MLP-32_16"
-    log_summary = pl.read_csv(os.path.join(log_dir, f"{model}.csv"))
-    weights = joblib.load(os.path.join(log_dir, f"{model}_weights.pkl"))
-
-    set_plt_rcparams()
-
-    _plot_predictions(log_summary, output_dir)
-
-    for rate in ["birth", "death"]:
-        plot_partial_dependencies(
-            weights=weights[f"{rate}Rate"],
-            features=SCENARIO.features[f"{rate}Rate"],
-            output_dir=os.path.join(output_dir, rate),
-            output_activation=partial(sigmoid, upper=FBD_RATE_UPPER),
-        )
-        plot_shap_features_importance(
-            weights=weights[f"{rate}Rate"],
-            features=SCENARIO.features[f"{rate}Rate"],
-            output_file=os.path.join(output_dir, rate, "shap_values.svg"),
-            output_activation=partial(sigmoid, upper=FBD_RATE_UPPER),
-        )
-
-
-if __name__ == "__main__":
-    main()

bella_companion/simulations/figures/fbd_predictions.py

@@ -1,66 +0,0 @@
-import os
-
-import matplotlib.pyplot as plt
-import polars as pl
-
-import src.config as cfg
-from src.simulations.figures.utils import (
-    plot_coverage_per_time_bin,
-    plot_maes_per_time_bin,
-    step,
-)
-from src.simulations.scenarios.fbd_no_traits import BIRTH_RATES, DEATH_RATES
-from src.utils import set_plt_rcparams
-
-
-def main():
-    output_dir = os.path.join(cfg.FIGURES_DIR, "fbd-predictions")
-    os.makedirs(output_dir, exist_ok=True)
-
-    set_plt_rcparams()
-
-    for i, (birth_rate, death_rate) in enumerate(
-        zip(BIRTH_RATES, DEATH_RATES), start=1
-    ):
-        summaries_dir = os.path.join(cfg.BEAST_LOGS_SUMMARIES_DIR, f"fbd-no-traits_{i}")
-        logs_summaries = {
-            "Nonparametric": pl.read_csv(
-                os.path.join(summaries_dir, "Nonparametric.csv")
-            ),
-            "GLM": pl.read_csv(os.path.join(summaries_dir, "GLM.csv")),
-            "MLP": pl.read_csv(os.path.join(summaries_dir, "MLP-16_8.csv")),
-        }
-        true_values = {"birthRate": birth_rate, "deathRate": death_rate}
-
-        for id, rate in true_values.items():
-            for log_summary in logs_summaries.values():
-                step(
-                    [
-                        log_summary[f"{id}i{i}_median"].median()
-                        for i in range(len(rate))
-                    ],
-                    reverse_xticks=True,
-                )
-            step(rate, color="k", linestyle="--", reverse_xticks=True)
-            plt.ylabel(r"$\lambda$" if id == "birthRate" else r"$\mu$")
-            plt.savefig(
-                os.path.join(output_dir, f"fbd-no-traits_{i}-predictions-{id}.svg")
-            )
-            plt.close()
-
-        plot_coverage_per_time_bin(
-            logs_summaries,
-            true_values,
-            os.path.join(output_dir, f"fbd-no-traits_{i}-coverage.svg"),
-            reverse_xticks=True,
-        )
-        plot_maes_per_time_bin(
-            logs_summaries,
-            true_values,
-            os.path.join(output_dir, f"fbd-no-traits_{i}-maes.svg"),
-            reverse_xticks=True,
-        )
-
-
-if __name__ == "__main__":
-    main()