bella-companion 0.0.0__py3-none-any.whl

bella_companion/fbd_empirical/run_beast.py

@@ -0,0 +1,54 @@
+import os
+from pathlib import Path
+
+import numpy as np
+import polars as pl
+from phylogenie import load_newick
+from phylogenie.utils import get_node_depths
+from tqdm import tqdm
+
+import src.config as cfg
+from src.utils import run_sbatch
+
+THIS_DIR = Path(__file__).parent
+
+
+def main():
+    output_dir = cfg.BEAST_OUTPUTS_DIR / "fbd-empirical"
+    os.makedirs(output_dir, exist_ok=True)
+
+    tree_file = THIS_DIR / "data" / "trees.nwk"
+    change_times_file = THIS_DIR / "data" / "change_times.csv"
+    sampling_change_times_file = THIS_DIR / "data" / "sampling_change_times.csv"
+
+    change_times = (
+        pl.read_csv(change_times_file, has_header=False).to_series().to_numpy()
+    )
+    time_predictor = " ".join(
+        list(map(str, np.repeat(np.insert(change_times, 0, 0), 4)))
+    )
+    body_mass_predictor = " ".join(["0", "1", "2", "3"] * (len(change_times) + 1))
+
+    trees = load_newick(str(tree_file))
+    assert isinstance(trees, list)
+    for i, tree in enumerate(tqdm(trees)):
+        process_length = max(get_node_depths(tree).values())
+        for model in ["hidden-relu", "hidden-tanh"]:
+            command = " ".join(
+                [
+                    cfg.RUN_BEAST,
+                    f'-D treeFile={tree_file},treeIndex={i},typeTraitFile={THIS_DIR / "data" / "body_mass.csv"},changeTimesFile={change_times_file},samplingChangeTimesFile={sampling_change_times_file},processLength={process_length},timePredictor="{time_predictor}",bodyMassPredictor="{body_mass_predictor}"',
+                    f"-DF {THIS_DIR / 'params.json'}",
+                    f"-prefix {output_dir / model}",
+                    str(cfg.BEAST_CONFIGS_DIR / "fbd-empirical" / f"{model}.xml"),
+                ]
+            )
+            run_sbatch(
+                command,
+                cfg.SBATCH_LOGS_DIR / "fbd-empirical" / model / str(i),
+                mem_per_cpu="12000",
+            )
+
+
+if __name__ == "__main__":
+    main()

bella_companion/fbd_empirical/summarize_logs.py

@@ -0,0 +1,50 @@
+import json
+import os
+from pathlib import Path
+
+import joblib
+import polars as pl
+
+from src.config import BEAST_LOGS_SUMMARIES_DIR, BEAST_OUTPUTS_DIR
+from src.utils import summarize_logs
+
+THIS_DIR = Path(__file__).parent
+
+
+def main():
+    summaries_dir = os.path.join(BEAST_LOGS_SUMMARIES_DIR, "fbd-empirical")
+    os.makedirs(summaries_dir, exist_ok=True)
+
+    with open(os.path.join(THIS_DIR, "params", "MLP.json"), "r") as f:
+        params = json.load(f)
+    hidden_nodes = list(map(int, params["nodes"].split()))[:-1]
+    print(hidden_nodes)
+    states = params["types"].split(",")
+    logs_dir = os.path.join(BEAST_OUTPUTS_DIR, "fbd-empirical", "MLP")
+    change_times = (
+        pl.read_csv(
+            os.path.join(THIS_DIR, "data", "change_times.csv"), has_header=False
+        )
+        .to_series()
+        .to_list()
+    )
+    n_time_bins = len(change_times) + 1
+    logs_summary, weights = summarize_logs(
+        logs_dir,
+        target_columns=[
+            f"{rate}Ratei{i}_{s}"
+            for rate in ["birth", "death"]
+            for i in range(n_time_bins)
+            for s in states
+        ],
+        hidden_nodes=hidden_nodes,
+        n_features={f"{rate}Rate": 2 for rate in ["birth", "death"]},
+        layers_range_start=0,
+    )
+
+    logs_summary.write_csv(os.path.join(summaries_dir, f"MLP.csv"))
+    joblib.dump(weights, os.path.join(summaries_dir, "weights.pkl"))
+
+
+if __name__ == "__main__":
+    main()

bella_companion/simulations/features.py

@@ -0,0 +1,7 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class Feature:
+    is_binary: bool
+    is_relevant: bool

bella_companion/simulations/figures/epi_explainations.py

@@ -0,0 +1,101 @@
+import os
+from functools import partial
+
+import joblib
+import matplotlib.pyplot as plt
+import numpy as np
+import polars as pl
+from lumiere.backend import sigmoid
+
+import src.config as cfg
+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 _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

@@ -0,0 +1,58 @@
+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

@@ -0,0 +1,99 @@
+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

@@ -0,0 +1,66 @@
+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()

bella_companion/simulations/figures/scenarios.py

@@ -0,0 +1,87 @@
+import os
+
+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 (
+    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
+
+
+def main():
+    output_dir = os.path.join(cfg.FIGURES_DIR, "scenarios")
+    os.makedirs(output_dir, exist_ok=True)
+
+    set_plt_rcparams()
+
+    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.close()
+
+    sort_idx = np.argsort(MIGRATION_PREDICTOR.flatten())
+    plt.plot(
+        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.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"))
+        plt.close()
+
+    step(
+        BIRTH_RATE_TRAIT1_UNSET,
+        label=r"$\lambda_{0,0} = \lambda_{0,1}$",
+        color="C0",
+        reverse_xticks=True,
+    )
+    step(
+        BIRTH_RATE_TRAIT1_SET,
+        label=r"$\lambda_{1,0} = \lambda_{1,1}$",
+        color="C0",
+        linestyle="dashed",
+        reverse_xticks=True,
+    )
+    step(
+        DEATH_RATE_TRAIT1_UNSET,
+        label=r"$\mu_{0,0} = \mu_{0,1}$",
+        color="C1",
+        reverse_xticks=True,
+    )
+    step(
+        DEATH_RATE_TRAIT1_SET,
+        label=r"$\mu_{1,0} = \mu_{1,1}$",
+        color="C1",
+        linestyle="dashed",
+        reverse_xticks=True,
+    )
+    plt.ylabel("Rate")
+    plt.legend()
+    plt.savefig(os.path.join(output_dir, "fbd-2traits.svg"))
+    plt.close()
+
+
+if __name__ == "__main__":
+    main()