bella-companion 0.0.6__tar.gz → 0.0.7__tar.gz

{bella_companion-0.0.6 → bella_companion-0.0.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: bella-companion
-Version: 0.0.6
+Version: 0.0.7
 Summary: 
 Author: gabriele-marino
 Author-email: gabmarino.8601@gmail.com

{bella_companion-0.0.6 → bella_companion-0.0.7}/bella_companion/cli.py

@@ -4,12 +4,10 @@ from pathlib import Path
 
 from dotenv import load_dotenv
 
-from bella_companion.simulations import (
-    generate_data,
-    generate_figures,
-    run_beast,
-    summarize_logs,
-)
+from bella_companion.fbd_empirical import run_beast as run_fbd_empirical
+from bella_companion.simulations import generate_data, generate_figures
+from bella_companion.simulations import run_beast as run_simulations
+from bella_companion.simulations import summarize_logs
 
 
 def main():
@@ -28,7 +26,7 @@ def main():
 
     subparsers.add_parser(
         "sim-run", help="Run BEAST2 analyses on simulation datasets."
-    ).set_defaults(func=run_beast)
+    ).set_defaults(func=run_simulations)
 
     subparsers.add_parser(
         "sim-summary", help="Summarize BEAST2 log outputs for simulations."
@@ -38,5 +36,9 @@ def main():
         "sim-figures", help="Generate plots and figures from simulation results."
     ).set_defaults(func=generate_figures)
 
+    subparsers.add_parser(
+        "fbd-empirical-run", help="Run BEAST2 analyses on empirical FBD datasets."
+    ).set_defaults(func=run_fbd_empirical)
+
     args = parser.parse_args()
     args.func()

bella_companion-0.0.7/bella_companion/fbd_empirical/__init__.py

@@ -0,0 +1,3 @@
+from bella_companion.fbd_empirical.run_beast import run_beast
+
+__all__ = ["run_beast"]

bella_companion-0.0.7/bella_companion/fbd_empirical/run_beast.py

@@ -0,0 +1,68 @@
+import json
+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
+
+from bella_companion.utils import submit_job
+
+THIS_DIR = Path(__file__).parent
+
+
+def run_beast():
+    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"
+    tree_file = data_dir / "trees.nwk"
+    change_times_file = data_dir / "change_times.csv"
+
+    trees = load_newick(str(tree_file))
+    assert isinstance(trees, list)
+    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))
+
+    job_ids = {}
+    for i, tree in enumerate(tqdm(trees)):
+        process_length = max(get_node_depths(tree).values())
+        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 birthRateUpper=2",
+                f"-D deathRateUpper=2",
+                f"-D samplingRateUpper=10",
+                f'-D samplingRateInit="5 5 5 5 5 5 5"',
+                f"-D migrationRateUpper=10",
+                f'-D migrationRateInit="0.5 0 0 0.5 0.5 0 0 0.5 0.5 0 0 0.5"',
+                f'-D nodes="16 8"',
+                f'-D layersRange="0,1,2"',
+                f"-D treeFile={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 processLength={process_length}",
+                f'-D timePredictor="{time_predictor}"',
+                f'-D bodyMassPredictor="{body_mass_predictor}"',
+                f"-prefix {output_dir}{os.sep}",
+                str(Path(os.environ["BELLA_BEAST_CONFIGS_DIR"]) / "fbd-empirical.xml"),
+            ]
+        )
+        job_ids[i] = submit_job(
+            command,
+            Path(os.environ["BELLA_SBATCH_LOG_DIR"]) / "fbd-empirical" / str(i),
+            mem_per_cpu="12000",
+        )
+
+    with open(base_output_dir / "fbd_empirical_job_ids.json", "w") as f:
+        json.dump(job_ids, f)

{bella_companion-0.0.6 → bella_companion-0.0.7}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "bella-companion"
-version = "0.0.6"
+version = "0.0.7"
 description = ""
 authors = ["gabriele-marino <gabmarino.8601@gmail.com>"]
 

bella_companion-0.0.6/bella_companion/fbd_empirical/params.json

@@ -1,11 +0,0 @@
-{
-    "types": "0,1,2,3",
-    "startTypePriorProbs": "0.25 0.25 0.25 0.25",
-    "birthRateUpper": 2,
-    "deathRateUpper": 2,
-    "samplingRateUpper": 10,
-    "samplingRateInit": "5 5 5 5 5 5 5",
-    "migrationRateUpper": 10,
-    "migrationRateInit": "0.5 0 0 0.5 0.5 0 0 0.5 0.5 0 0 0.5",
-    "nodes": "16 8"
-}

bella_companion-0.0.6/bella_companion/fbd_empirical/run_beast.py

@@ -1,54 +0,0 @@
-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
-AADD LAYERS RANGE
-
-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()