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

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
@@ -11,7 +12,19 @@ from tqdm import tqdm
 from bella_companion.utils import submit_job
 
 
+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)
@@ -19,14 +32,28 @@ def run_beast():
     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)):
@@ -34,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-0.0.10.dist-info → bella_companion-0.0.11.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: bella-companion
-Version: 0.0.10
+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.10.dist-info → bella_companion-0.0.11.dist-info}/RECORD

@@ -1,12 +1,11 @@
 bella_companion/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 bella_companion/cli.py,sha256=D-kQuj0I-g5Vm5-mXdOZoio68mhOBpJ3jVG8B2Kwvn4,2190
 bella_companion/fbd_empirical/__init__.py,sha256=FgzjaEpwZImvBsJs17ghF0ltZvhFmgmkTGGF3HgdsGo,284
-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/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=ZBTtBhSz3wDmPZQJZaY9_o56x_FkL9lKiopWVV_ItX8,2523
+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
@@ -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.10.dist-info/METADATA,sha256=gLQu59HiP0oxF5V7o3cehkcRwhcaDcIumiYIC7D0eXc,577
-bella_companion-0.0.10.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-bella_companion-0.0.10.dist-info/entry_points.txt,sha256=rSeKoAhmjnQqAYFcXBv0gAM2ViJfJe0D8_dD-fWrXeg,50
-bella_companion-0.0.10.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,,