bella-companion 0.0.9__py3-none-any.whl → 0.0.11__py3-none-any.whl

bella_companion/fbd_empirical/results.py

@@ -0,0 +1,88 @@
+import os
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+import polars as pl
+from numpy.typing import NDArray
+from phylogenie import get_node_depths, load_newick
+
+
+def _plot_predictions(output_dir: Path):
+    log_summaries_dir = Path(os.environ["BELLA_LOG_SUMMARIES_DIR"])
+    log_summary = pl.read_csv(log_summaries_dir / "fbd-empirical" / "MLP.csv")
+
+    data_dir = Path(__file__).parent / "data"
+    change_times = (
+        pl.read_csv(data_dir / "change_times.csv", has_header=False)
+        .to_series()
+        .to_list()
+    )
+    max_time = max(
+        max(get_node_depths(tree).values())
+        for tree in load_newick(data_dir / "trees.nwk")
+    )
+    time_bins = list(reversed([0.0, *change_times, max_time]))
+
+    colors: dict[str, NDArray[np.floating]] = {
+        "birth": plt.cm.Blues(np.linspace(0.4, 0.9, 4)),  # pyright: ignore
+        "death": plt.cm.Oranges(np.linspace(0.4, 0.9, 4)),  # pyright: ignore
+        "diversification": plt.cm.Greens(np.linspace(0.4, 0.9, 4)),  # pyright: ignore
+    }
+    for rate in ["birth", "death", "diversification"]:
+        for state in range(4):
+            if rate == "diversification":
+                estimates = log_summary.select(
+                    [
+                        pl.col(f"birthRateSPi{i}_{state}_median")
+                        - pl.col(f"deathRateSPi{i}_{state}_median")
+                        for i in range(len(change_times) + 1)
+                    ]
+                ).to_numpy()
+            else:
+                estimates = log_summary.select(
+                    [
+                        pl.col(f"{rate}RateSPi{i}_{state}_median")
+                        for i in range(len(change_times) + 1)
+                    ]
+                ).to_numpy()
+
+            median = np.median(estimates, axis=0)
+            lower = np.percentile(estimates, 2.5, axis=0)
+            upper = np.percentile(estimates, 97.5, axis=0)
+
+            color = colors[rate][state]
+
+            plt.fill_between(  # pyright: ignore
+                time_bins,
+                [lower[0], *lower],
+                [upper[0], *upper],
+                step="pre",
+                alpha=0.25,
+                color=color,
+            )
+
+            for estimate in estimates:
+                plt.step(  # pyright: ignore
+                    time_bins, [estimate[0], *estimate], color=color, alpha=0.15
+                )
+
+            plt.step(  # pyright: ignore
+                time_bins, [median[0], *median], color=color, label=state
+            )
+
+        ax = plt.gca()
+        ax.invert_xaxis()
+        plt.legend(title="Body mass")  # pyright: ignore
+        plt.xlabel("Time (mya)")  # pyright: ignore
+        plt.ylabel(  # pyright: ignore
+            r"$\lambda$" if rate == "birth" else r"$\mu$" if rate == "death" else r"$d$"
+        )
+        plt.savefig(output_dir / f"{rate}-predictions.svg")  # pyright: ignore
+        plt.close()
+
+
+def plot_fbd_empirical_results():
+    output_dir = Path(os.environ["BELLA_FIGURES_DIR"]) / "fbd-empirical"
+    os.makedirs(output_dir, exist_ok=True)
+    _plot_predictions(output_dir)

bella_companion/fbd_empirical/run_beast.py

@@ -1,5 +1,6 @@
 import json
 import os
+from itertools import combinations
 from pathlib import Path
 
 import numpy as np
@@ -10,25 +11,49 @@ from tqdm import tqdm
 
 from bella_companion.utils import submit_job
 
-THIS_DIR = Path(__file__).parent
+
+def _get_migration_rates_init(types: list[str], init_rate: float = 1) -> str:
+    mus: list[float] = []
+    for t1, t2 in combinations(types, 2):
+        traits1 = np.array(list(map(int, t1.split("_"))))
+        traits2 = np.array(list(map(int, t2.split("_"))))
+        mus.append(init_rate if np.sum(np.abs(traits1 - traits2)) == 1 else 0)
+    return " ".join(map(str, mus))
 
 
 def run_beast():
+    from dotenv import load_dotenv
+
+    load_dotenv()
     base_output_dir = Path(os.environ["BELLA_BEAST_OUTPUT_DIR"])
     output_dir = base_output_dir / "fbd-empirical"
     os.makedirs(output_dir, exist_ok=True)
 
-    data_dir = THIS_DIR / "data"
+    data_dir = Path(__file__).parent / "data"
     tree_file = data_dir / "trees.nwk"
     change_times_file = data_dir / "change_times.csv"
+    traits_file = data_dir / "traits.csv"
 
     trees = load_newick(str(tree_file))
     assert isinstance(trees, list)
+
+    traits = pl.read_csv(traits_file, separator="\t", null_values=["NA"])
+
     change_times = (
         pl.read_csv(change_times_file, has_header=False).to_series().to_numpy()
     )
-    time_predictor = " ".join(list(map(str, np.repeat([0, *change_times], 4))))
-    body_mass_predictor = " ".join(["0", "1", "2", "3"] * (len(change_times) + 1))
+
+    types: list[str] = sorted(traits["type"].unique())
+    types.remove("?")
+    N = len(types)
+
+    time_predictor = " ".join(list(map(str, np.repeat([0, *change_times], N))))
+    log10BM_predictor = " ".join(
+        [t.split("_")[0] for t in types] * (len(change_times) + 1)
+    )
+    midlat_predictor = " ".join(
+        [t.split("_")[1] for t in types] * (len(change_times) + 1)
+    )
 
     job_ids = {}
     for i, tree in enumerate(tqdm(trees)):
@@ -36,24 +61,24 @@ def run_beast():
         command = " ".join(
             [
                 os.environ["BELLA_RUN_BEAST_CMD"],
-                f"-D types=0,1,2,3",
-                f'-D startTypePriorProbs="0.25 0.25 0.25 0.25"',
+                f'-D types="{",".join(types)}"',
+                f'-D startTypePriorProbs="{" ".join([str(1/N)] * N)}"',
                 f"-D birthRateUpper=5",
                 f"-D deathRateUpper=5",
                 f"-D samplingRateUpper=5",
-                f'-D samplingRateInit="2.5 2.5 2.5 2.5 2.5 2.5 2.5"',
+                f'-D samplingRateInit="{" ".join(["2.5"] * N)}"',
                 f"-D migrationRateUpper=5",
-                f'-D migrationRateInit="2.5 0 0 2.5 2.5 0 0 2.5 2.5 0 0 2.5"',
+                f'-D migrationRateInit="{_get_migration_rates_init(types, 2.5)}"',
                 f'-D nodes="16 8"',
                 f'-D layersRange="0,1,2"',
-                f"-D treeFile={tree_file}",
+                f"-D tree_file={tree_file}",
                 f"-D treeIndex={i}",
                 f"-D changeTimesFile={change_times_file}",
-                f"-D samplingChangeTimesFile={data_dir / 'sampling_change_times.csv'}",
-                f"-D typeTraitFile={data_dir / 'body_mass.csv'}",
+                f"-D traitsFile={traits_file}",
                 f"-D processLength={process_length}",
                 f'-D timePredictor="{time_predictor}"',
-                f'-D bodyMassPredictor="{body_mass_predictor}"',
+                f'-D log10BM_predictor="{log10BM_predictor}"',
+                f'-D midlat_predictor="{midlat_predictor}"',
                 f"-prefix {output_dir}{os.sep}",
                 str(Path(os.environ["BELLA_BEAST_CONFIGS_DIR"]) / "fbd-empirical.xml"),
             ]

bella_companion/fbd_empirical/summarize_logs.py

@@ -6,13 +6,10 @@ import polars as pl
 
 from bella_companion.utils import read_weights_dir, summarize_logs_dir
 
-THIS_DIR = Path(__file__).parent
-
 
 def summarize_logs():
-    change_times = pl.read_csv(
-        Path(THIS_DIR) / "data" / "change_times.csv", has_header=False
-    )
+    data_dir = Path(__file__).parent / "data"
+    change_times = pl.read_csv(data_dir / "change_times.csv", has_header=False)
     n_time_bins = len(change_times) + 1
 
     logs_dir = Path(os.environ["BELLA_BEAST_OUTPUT_DIR"]) / "fbd-empirical"

bella_companion/simulations/figures/utils.py

@@ -39,7 +39,7 @@ def _count_time_bins(true_values: dict[str, list[float]]) -> int:
 def plot_maes_per_time_bin(
     logs_summaries: dict[str, pl.DataFrame],
     true_values: dict[str, list[float]],
-    output_filepath: Path,
+    output_filepath: str | Path,
     reverse_xticks: bool = False,
 ):
     def _mae(target: str, i: int) -> pl.Expr:
@@ -73,7 +73,7 @@ def plot_maes_per_time_bin(
 def plot_coverage_per_time_bin(
     logs_summaries: dict[str, pl.DataFrame],
     true_values: dict[str, list[float]],
-    output_filepath: Path,
+    output_filepath: str | Path,
     reverse_xticks: bool = False,
 ):
     def _coverage(model: str, target: str, i: int) -> float:

{bella_companion-0.0.9.dist-info → bella_companion-0.0.11.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: bella-companion
-Version: 0.0.9
+Version: 0.0.11
 Summary: 
 Author: gabriele-marino
 Author-email: gabmarino.8601@gmail.com
@@ -12,5 +12,5 @@ Classifier: Programming Language :: Python :: 3.12
 Requires-Dist: arviz (>=0.22.0,<0.23.0)
 Requires-Dist: bella-lumiere (>=0.0.13,<0.0.14)
 Requires-Dist: dotenv (>=0.9.9,<0.10.0)
-Requires-Dist: phylogenie (>=2.1.29,<3.0.0)
+Requires-Dist: phylogenie (>=2.1.30,<3.0.0)
 Requires-Dist: seaborn (>=0.13.2,<0.14.0)

{bella_companion-0.0.9.dist-info → bella_companion-0.0.11.dist-info}/RECORD

@@ -1,13 +1,12 @@
 bella_companion/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-bella_companion/cli.py,sha256=ykq62Qjt9Kk1g0HZZMMA_nAqoG5E7vnWEWGfqDJx3hU,1937
-bella_companion/fbd_empirical/__init__.py,sha256=nQdZfS56ZsU26davkyK4isTqvKjMVw9OniHqjWRt3q0,177
-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
-bella_companion/fbd_empirical/data/trees.nwk,sha256=zhvLvPLZelhMThVmvOENkmi3p2aPAARb8KMdHTm6mss,4645318
-bella_companion/fbd_empirical/notbooks.ipynb,sha256=0ptmZC-_QfXtfCBrrag7NLtupiVn4gdG-oqVCworT_U,81588
-bella_companion/fbd_empirical/run_beast.py,sha256=108c_vcEA1iHxR2wLx-OwyZu0XdYEIHr3Z809FWmuVE,2544
-bella_companion/fbd_empirical/summarize_logs.py,sha256=TBLdW5LWS5ensEQFFQ8hrzJFJ3t0l-Mzu_SNa-Q8D_w,972
+bella_companion/cli.py,sha256=D-kQuj0I-g5Vm5-mXdOZoio68mhOBpJ3jVG8B2Kwvn4,2190
+bella_companion/fbd_empirical/__init__.py,sha256=FgzjaEpwZImvBsJs17ghF0ltZvhFmgmkTGGF3HgdsGo,284
+bella_companion/fbd_empirical/data/change_times.csv,sha256=tedobHbaY7-6M3V6F6HSqGAU9dpeeajR5mu1XGJfp6w,68
+bella_companion/fbd_empirical/data/traits.csv,sha256=cxNKYmgtBCKHsm9qnOpS29Rkx9cQKe0oxpgaXSIukJg,5422
+bella_companion/fbd_empirical/data/trees.nwk,sha256=JAsdfGhTIzGmJz6g6-0s5SVfqNuxawOMc9gbt0srue4,539385
+bella_companion/fbd_empirical/results.py,sha256=0FENY5EpgNQ6rwPsEiOT2ZEO5oBbIAK6TzPvk11SDnk,3116
+bella_companion/fbd_empirical/run_beast.py,sha256=Vvh0xognC92tAI7GhUyzU901N1x6t8UjUDccvB2JcVA,3309
+bella_companion/fbd_empirical/summarize_logs.py,sha256=OoT43DMjim9o4qsjYCgGKnBKp3dHXzVP6sVtzqR8lMQ,955
 bella_companion/simulations/__init__.py,sha256=ShYRdp1iSjnS_SzcsH-8jbqXz6P1nyRQZqAtPZJCMVE,454
 bella_companion/simulations/features.py,sha256=DZOBpJGlQ0UinqUZYbEtoemZ2eQGVLV_i-DfpW31qJI,104
 bella_companion/simulations/figures/__init__.py,sha256=aBYbJntH4egFmkSSWiVMYDEApXPYxJD7eA3TCPNNegM,658
@@ -19,7 +18,7 @@ bella_companion/simulations/figures/explain/shap.py,sha256=qRfOMNwkU-hsBy5MHMTfV
 bella_companion/simulations/figures/fbd_2traits_results.py,sha256=JiXrbYkH1HwdJQhTHXj6KhMEXYgQmQ6LmDboAIO1CPA,2728
 bella_companion/simulations/figures/fbd_no_traits_results.py,sha256=fLsgpV3IbLLtxQEFNOL9K4jEEJrG324ziUM0rxIv7_k,1962
 bella_companion/simulations/figures/scenarios.py,sha256=gbMz1TUxxT2RSIq2kQlFioNdgSHk-gQY2OQuf6-7Fww,2817
-bella_companion/simulations/figures/utils.py,sha256=dB_GiJ1HGrZ93cqODz6kB-HeDRPwlm2MkMe9rJZGnfs,3117
+bella_companion/simulations/figures/utils.py,sha256=0M5OrxaEuqcj9rR2uAc_O7utQvhEceZGH0sKrGRWaWs,3129
 bella_companion/simulations/generate_data.py,sha256=ZFQnusmGWDjRixKOid0l49Aj4qfrI5zxyf1n21wf0p0,768
 bella_companion/simulations/generate_figures.py,sha256=layMgoj3Bfl78Ceb1oE7YirAQ8zhjDyD9IrxDRXf6go,657
 bella_companion/simulations/metrics.py,sha256=TwLgK4Eui4DlG10V9m6zEbmuGfUIli2dE0Ph-gX_6bU,1925
@@ -36,7 +35,7 @@ bella_companion/utils/__init__.py,sha256=UtMwPK9dWf9NAl0ic8CSsgdW7aSm-5J49OqgvD7
 bella_companion/utils/beast.py,sha256=TBa0cLklX1_tXqoQE4LRYvds7mLg_9fS2-6U6OHENHo,2184
 bella_companion/utils/explain.py,sha256=uP7HPyn2YiykAI69BQV3RooDpC6qKoCLXfp3Uibp4zk,1475
 bella_companion/utils/slurm.py,sha256=v5DaG7YHVyK8KRFptgGDC6I8jxEhyJuMVK9N08pZSAI,1812
-bella_companion-0.0.9.dist-info/METADATA,sha256=SGm-iWd_hwWEA74pd2SQIj2HhVacYElxBKa_NfX4GPg,576
-bella_companion-0.0.9.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-bella_companion-0.0.9.dist-info/entry_points.txt,sha256=rSeKoAhmjnQqAYFcXBv0gAM2ViJfJe0D8_dD-fWrXeg,50
-bella_companion-0.0.9.dist-info/RECORD,,
+bella_companion-0.0.11.dist-info/METADATA,sha256=NQPoh1JxQYhkfvNRr_CrdJx1_t46j5o8P6lHlnEy83k,577
+bella_companion-0.0.11.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+bella_companion-0.0.11.dist-info/entry_points.txt,sha256=rSeKoAhmjnQqAYFcXBv0gAM2ViJfJe0D8_dD-fWrXeg,50
+bella_companion-0.0.11.dist-info/RECORD,,