bella-companion 0.0.45__tar.gz → 0.0.47__tar.gz

{bella_companion-0.0.45 → bella_companion-0.0.47}/PKG-INFO

@@ -1,12 +1,12 @@
 Metadata-Version: 2.4
 Name: bella-companion
-Version: 0.0.45
-Summary: Add your description here
+Version: 0.0.47
+Summary: Bayesian Evolutionary Layered Learning Architectures (BELLA) companion
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 Requires-Dist: arviz>=0.22.0
 Requires-Dist: bella-lumiere>=0.0.14
 Requires-Dist: dotenv>=0.9.9
-Requires-Dist: phylogenie>=3.1.3
+Requires-Dist: phylogenie>=3.1.4
 Requires-Dist: poetry>=2.2.1
 Requires-Dist: seaborn>=0.13.2

{bella_companion-0.0.45 → bella_companion-0.0.47}/bella_companion/cli.py

@@ -2,6 +2,7 @@ import argparse
 import os
 from pathlib import Path
 
+import matplotlib.pyplot as plt
 from dotenv import load_dotenv
 
 from bella_companion.eucovid import run_beast as run_eucovid_beast
@@ -21,12 +22,16 @@ def main():
             os.getenv("JAVA_OPTIONS", ""),
             f"-jar {Path(__file__).parent / 'BELLA.jar'}",
             f"-version_file {Path(__file__).parent / 'version.xml'}",
-            "-seed 42",
             "-overwrite",
             "-statefile /tmp/state",
-            "-DFout /dev/null",
         ]
     )
+    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
 
     parser = argparse.ArgumentParser(
         prog="bella",

bella_companion-0.0.47/bella_companion/eucovid/run_beast.py

@@ -0,0 +1,118 @@
+import os
+from glob import glob
+from pathlib import Path
+
+from bella_companion.utils import submit_job
+
+
+def run_beast():
+    base_output_dir = Path(os.environ["BELLA_BEAST_OUTPUT_DIR"])
+    output_dir = base_output_dir / "eucovid"
+    os.makedirs(output_dir, exist_ok=True)
+
+    base_log_dir = Path(os.environ["BELLA_SBATCH_LOG_DIR"]) / "eucovid"
+
+    data_dir = Path(__file__).parent / "data"
+    beast_configs_dir = Path(__file__).parent / "beast_configs"
+    msa_file = data_dir / "msa.fasta"
+
+    predictors_dir = data_dir / "predictors"
+    all_predictor_files = ",".join(glob(str(predictors_dir / "all" / "*.tsv")))
+    all_predictors_data = " ".join(
+        [
+            f"-D msa_file={msa_file}",
+            f"-D changeTimesFile={predictors_dir / 'changetimes_all_7e.tsv'}",
+            f"-D predictorFiles={all_predictor_files}",
+        ]
+    )
+    flight_predictor_data = " ".join(
+        [
+            f"-D msa_file={msa_file}",
+            f"-D changeTimesFile={predictors_dir / 'changetimes_flights_4e.tsv'}",
+            f"-D predictorFiles={predictors_dir / 'flight_pop_x_4e_ls.tsv'}",
+        ]
+    )
+
+    for seed in range(1, 3):
+        run_beast_cmd = [os.environ["BELLA_RUN_BEAST_CMD"], f"-seed {42 + seed}"]
+
+        """
+        os.makedirs(output_dir / "Nonparametric", exist_ok=True)
+        submit_job(
+            " ".join(
+                [
+                    *run_beast_cmd,
+                    f"-D aligned_fasta={msa_file}",
+                    f"-prefix {output_dir / 'Nonparametric'}{os.sep}{seed}-",
+                    str(beast_configs_dir / "Nonparametric.xml"),
+                ]
+            ),
+            base_log_dir / "Nonparametric" / str(seed),
+            cpus=128,
+            mem_per_cpu=12000,
+        )
+        """
+
+        os.makedirs(output_dir / "all-predictors-GLM", exist_ok=True)
+        submit_job(
+            " ".join(
+                [
+                    *run_beast_cmd,
+                    all_predictors_data,
+                    f"-prefix {output_dir / 'all-predictors-GLM'}{os.sep}{seed}-",
+                    str(beast_configs_dir / "GLM.xml"),
+                ]
+            ),
+            base_log_dir / "all-predictors-GLM" / str(seed),
+            cpus=128,
+            mem_per_cpu=12000,
+        )
+        """
+        os.makedirs(output_dir / "flights-GLM", exist_ok=True)
+        submit_job(
+            " ".join(
+                [
+                    *run_beast_cmd,
+                    flight_predictor_data,
+                    f"-prefix {output_dir / 'flights-GLM'}{os.sep}{seed}-",
+                    str(beast_configs_dir / "GLM.xml"),
+                ]
+            ),
+            base_log_dir / "flights-GLM" / str(seed),
+            cpus=128,
+            mem_per_cpu=12000,
+        )
+        """
+
+        os.makedirs(output_dir / "all-predictors-MLP", exist_ok=True)
+        submit_job(
+            " ".join(
+                [
+                    *run_beast_cmd,
+                    '-D layersRange="0,1,2",nodes="16 8"',
+                    all_predictors_data,
+                    f"-prefix {output_dir / 'all-predictors-MLP'}{os.sep}{seed}-",
+                    str(beast_configs_dir / "MLP.xml"),
+                ]
+            ),
+            base_log_dir / "all-predictors-MLP" / str(seed),
+            cpus=128,
+            mem_per_cpu=12000,
+        )
+        """
+        os.makedirs(output_dir / "flights-MLP", exist_ok=True)
+        submit_job(
+            " ".join(
+                [
+                    *run_beast_cmd,
+                    '-D layersRange="0,1,2",nodes="16 8"',
+                    flight_predictor_data,
+                    f"-prefix {output_dir / 'flights-MLP'}{os.sep}{seed}-",
+                    str(beast_configs_dir / "MLP.xml"),
+                ]
+            ),
+            base_log_dir / "flights-MLP" / str(seed),
+            cpus=128,
+            mem_per_cpu=12000,
+        )
+        """

bella_companion-0.0.47/bella_companion/platyrrhine/results.py

@@ -0,0 +1,191 @@
+import os
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+import polars as pl
+from matplotlib.colors import LinearSegmentedColormap, ListedColormap
+from numpy.typing import NDArray
+from phylogenie import draw_tree, get_node_depths, load_newick, load_nexus
+
+DATA_DIR = Path(__file__).parent / "data"
+CHANGE_TIMES = (
+    pl.read_csv(DATA_DIR / "change_times.csv", has_header=False).to_series().to_list()
+)
+
+
+def _plot_marginal_rates(output_dir: Path, trees_file: str | Path):
+    trees = load_newick(trees_file)
+    all_node_depths = [get_node_depths(tree) for tree in trees]
+    max_time = max(max(node_depths.values()) for node_depths in all_node_depths)
+    time_bins = np.array(list(reversed([0.0, *CHANGE_TIMES, max_time])))
+    for tree in trees:
+        for node in tree:
+            node["diversificationRateSP"] = node["birthRateSP"] - node["deathRateSP"]
+
+    for i, rate in enumerate(["birthRateSP", "deathRateSP", "diversificationRateSP"]):
+        marginal_rates: list[list[float]] = [[] for _ in range(len(CHANGE_TIMES) + 1)]
+        for tree, node_depths in zip(trees, all_node_depths):
+            root_age = tree.height
+            for node in tree:
+                if node.parent is None:
+                    marginal_rates[0].append(node[rate])
+                    continue
+                parent_age = root_age - node_depths[node.parent]  # pyright: ignore
+                node_age = root_age - node_depths[node]
+                t1 = np.where(time_bins >= parent_age)[0][-1]
+                t2 = np.where(time_bins > node_age)[0][-1]
+                for t in range(t1, t2 + 1):
+                    marginal_rates[t].append(node[rate])
+
+        median = [float(np.median(rates)) for rates in marginal_rates]
+        lower = [float(np.percentile(rates, 2.5)) for rates in marginal_rates]
+        upper = [float(np.percentile(rates, 97.5)) for rates in marginal_rates]
+        plt.step(  # pyright: ignore
+            time_bins, [median[0], *median], color=f"C{i}", label=type
+        )
+        plt.fill_between(  # pyright: ignore
+            time_bins,
+            [lower[0], *lower],
+            [upper[0], *upper],
+            color=f"C{i}",
+            alpha=0.2,
+            step="pre",
+        )
+
+        plt.gca().invert_xaxis()
+        plt.xlabel("Time (mya)")  # pyright: ignore
+        plt.ylabel(  # pyright: ignore
+            r"$\lambda$"
+            if rate == "birthRateSP"
+            else r"$\mu$"
+            if rate == "deathRateSP"
+            else r"$d$"
+        )
+        plt.savefig(output_dir / f"marginal-{rate}.svg")  # pyright: ignore
+        plt.close()
+
+
+def _plot_tree(output_dir: Path, tree_file: str | Path):
+    sample_trees = load_nexus(tree_file)
+    trees = [sample_trees[f"STATE_{i * 100_000}"] for i in range(11, 101)]
+
+    for tree in trees:
+        for node in tree:
+            if node.parent is not None and not node.branch_length:
+                node.parent.remove_child(node)
+
+    avg_tree = trees[0].copy()
+    for node in avg_tree:
+        node.clear()
+    for avg_node, nodes in zip(avg_tree, zip(*trees)):
+        avg_node["birthRateSP"] = np.mean([n["birthRateSP"] for n in nodes])
+        avg_node["deathRateSP"] = np.mean([n["deathRateSP"] for n in nodes])
+        avg_node["diversificationRateSP"] = (
+            avg_node["birthRateSP"] - avg_node["deathRateSP"]
+        )
+        avg_node["type"] = int(np.median([int(n["type"]) for n in nodes]))
+
+    avg_tree.ladderize()
+    cmaps: dict[str, LinearSegmentedColormap] = {
+        "birthRateSP": plt.cm.Blues,  # pyright: ignore
+        "deathRateSP": plt.cm.Oranges,  # pyright: ignore
+        "diversificationRateSP": plt.cm.Greens,  # pyright: ignore
+        "type": plt.cm.Purples,  # pyright: ignore
+    }
+    for color_by, cm in cmaps.items():
+        ax = plt.gca()
+        if color_by == "type":
+            cmap = ListedColormap(cm(np.linspace(0.3, 1, 4)))  # pyright: ignore
+            draw_tree(
+                avg_tree,
+                ax,
+                color_by=color_by,
+                legend_kwargs={"title": "Body mass", "loc": "upper left"},
+                colormap=cmap,
+            )
+        else:
+            cmap = LinearSegmentedColormap.from_list(
+                "cmap",
+                cm(np.linspace(0.2, 1, 256)),  # pyright: ignore
+            )
+            draw_tree(
+                avg_tree,
+                ax,
+                color_by=color_by,
+                hist_axes_kwargs={
+                    "loc": "upper left",
+                    "bbox_to_anchor": (0.06, 0, 1, 1),
+                    "bbox_transform": ax.transAxes,
+                },
+                colormap=cmap,
+            )
+        height = avg_tree.height
+        ticks = np.arange(height, 0, -10)
+        ax.set_xticks(ticks, [str(round(height - t)) for t in ticks])  # pyright: ignore
+        ax.set_xlabel("Time (mya)")  # pyright: ignore
+        plt.savefig(output_dir / f"{color_by}-tree.svg")  # pyright: ignore
+        plt.close()
+
+
+def _plot_predictions(output_dir: Path, log_summary_file: str | Path):
+    log_summary = pl.read_csv(log_summary_file)
+
+    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]))
+
+    gradient = np.linspace(0.4, 0.9, 4)
+    colors: dict[str, NDArray[np.floating]] = {
+        "birthRateSP": plt.cm.Blues(gradient),  # pyright: ignore
+        "deathRateSP": plt.cm.Oranges(gradient),  # pyright: ignore
+        "diversificationRateSP": plt.cm.Greens(gradient),  # pyright: ignore
+    }
+    for rate in colors:
+        for type in [0, 1, 2, 3]:
+            if rate == "diversificationRateSP":
+                estimates = log_summary.select(
+                    [
+                        pl.col(f"birthRateSPi{i}_{type}_median")
+                        - pl.col(f"deathRateSPi{i}_{type}_median")
+                        for i in range(len(CHANGE_TIMES) + 1)
+                    ]
+                ).to_numpy()
+            else:
+                estimates = log_summary.select(
+                    [
+                        pl.col(f"{rate}i{i}_{type}_median")
+                        for i in range(len(CHANGE_TIMES) + 1)
+                    ]
+                ).to_numpy()
+            median = np.median(estimates, axis=0)
+            color = colors[rate][type]
+            plt.step(  # pyright: ignore
+                time_bins, [median[0], *median], color=color, label=type
+            )
+
+        plt.gca().invert_xaxis()
+        plt.legend(title="Body mass")  # pyright: ignore
+        plt.xlabel("Time (mya)")  # pyright: ignore
+        plt.ylabel(  # pyright: ignore
+            r"$\lambda$"
+            if rate == "birthRateSP"
+            else r"$\mu$"
+            if rate == "deathRateSP"
+            else r"$d$"
+        )
+        if rate in ["birthRateSP", "deathRateSP"]:
+            plt.ylim(0, 0.4)  # pyright: ignore
+        plt.savefig(output_dir / f"{rate}-predictions.svg")  # pyright: ignore
+        plt.close()
+
+
+def plot_platyrrhine_results():
+    summaries_dir = Path(os.environ["BELLA_SUMMARIES_DIR"]) / "platyrrhine"
+    output_dir = Path(os.environ["BELLA_FIGURES_DIR"]) / "platyrrhine"
+    os.makedirs(output_dir, exist_ok=True)
+    _plot_predictions(output_dir, summaries_dir / "MLP.csv")
+    _plot_tree(output_dir, summaries_dir / "sample-trees.nexus")
+    _plot_marginal_rates(output_dir, summaries_dir / "trees.nwk")

{bella_companion-0.0.45 → bella_companion-0.0.47}/bella_companion/platyrrhine/run_beast.py

@@ -39,22 +39,23 @@ def run_beast():
         command = " ".join(
             [
                 os.environ["BELLA_RUN_BEAST_CMD"],
+                "-seed 42",
                 f'-D types="{",".join(types)}"',
-                f'-D startTypePriorProbs="0.25 0.25 0.25 0.25"',
-                f"-D birthRateUpper=5",
-                f"-D deathRateUpper=5",
-                f'-D samplingChangeTimes="2.58 5.333 23.03"',
-                f"-D samplingRateUpper=5",
-                f'-D samplingRateInit="2.5 2.5 2.5 2.5"',
-                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 nodes="16 8"',
-                f'-D layersRange="0,1,2"',
+                '-D startTypePriorProbs="0.25 0.25 0.25 0.25"',
+                "-D birthRateUpper=5",
+                "-D deathRateUpper=5",
+                '-D samplingChangeTimes="2.58 5.333 23.03"',
+                "-D samplingRateUpper=5",
+                '-D samplingRateInit="2.5 2.5 2.5 2.5"',
+                "-D migrationRateUpper=5",
+                '-D migrationRateInit="2.5 0 0 2.5 2.5 0 0 2.5 2.5 0 0 2.5"',
+                '-D nodes="16 8"',
+                '-D layersRange="0,1,2"',
                 f"-D treeFile={tree_file}",
                 f"-D treeIndex={i}",
                 f"-D changeTimesFile={change_times_file}",
                 f"-D traitsFile={traits_file}",
-                f"-D traitValueCol=3",
+                "-D traitValueCol=3",
                 f"-D processLength={process_length}",
                 f'-D timePredictor="{time_predictor}"',
                 f'-D log10BMPredictor="{log10BM_predictor}"',

{bella_companion-0.0.45 → bella_companion-0.0.47}/bella_companion/simulations/generate_figures.py

@@ -1,5 +1,3 @@
-import matplotlib.pyplot as plt
-
 from bella_companion.simulations.figures import (
     plot_epi_multitype_results,
     plot_epi_skyline_results,
@@ -10,13 +8,6 @@ from bella_companion.simulations.figures import (
 
 
 def generate_figures():
-    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
-
     plot_scenarios()
     plot_epi_skyline_results()
     plot_epi_multitype_results()

{bella_companion-0.0.45 → bella_companion-0.0.47}/bella_companion/simulations/run_beast.py

@@ -56,7 +56,11 @@ def run_beast():
                         f"-D lastSampleTime={max(get_node_depths(tree).values())}"
                     )
 
-                base_command = [os.environ["BELLA_RUN_BEAST_CMD"], *beast_args]
+                base_command = [
+                    os.environ["BELLA_RUN_BEAST_CMD"],
+                    "-seed 42",
+                    *beast_args,
+                ]
                 if model in ["Nonparametric", "GLM"]:
                     command = " ".join(
                         [*base_command, str(inference_configs_dir / f"{model}.xml")]

{bella_companion-0.0.45 → bella_companion-0.0.47}/bella_companion.egg-info/PKG-INFO

@@ -1,12 +1,12 @@
 Metadata-Version: 2.4
 Name: bella-companion
-Version: 0.0.45
-Summary: Add your description here
+Version: 0.0.47
+Summary: Bayesian Evolutionary Layered Learning Architectures (BELLA) companion
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 Requires-Dist: arviz>=0.22.0
 Requires-Dist: bella-lumiere>=0.0.14
 Requires-Dist: dotenv>=0.9.9
-Requires-Dist: phylogenie>=3.1.3
+Requires-Dist: phylogenie>=3.1.4
 Requires-Dist: poetry>=2.2.1
 Requires-Dist: seaborn>=0.13.2

{bella_companion-0.0.45 → bella_companion-0.0.47}/bella_companion.egg-info/requires.txt

@@ -1,6 +1,6 @@
 arviz>=0.22.0
 bella-lumiere>=0.0.14
 dotenv>=0.9.9
-phylogenie>=3.1.3
+phylogenie>=3.1.4
 poetry>=2.2.1
 seaborn>=0.13.2

{bella_companion-0.0.45 → bella_companion-0.0.47}/pyproject.toml

@@ -1,14 +1,14 @@
 [project]
 name = "bella-companion"
-version = "0.0.45"
-description = "Add your description here"
+version = "0.0.47"
+description = "Bayesian Evolutionary Layered Learning Architectures (BELLA) companion"
 readme = "README.md"
 requires-python = ">=3.10"
 dependencies = [
     "arviz>=0.22.0",
     "bella-lumiere>=0.0.14",
     "dotenv>=0.9.9",
-    "phylogenie>=3.1.3",
+    "phylogenie>=3.1.4",
     "poetry>=2.2.1",
     "seaborn>=0.13.2",
 ]
@@ -16,6 +16,7 @@ dependencies = [
 [dependency-groups]
 dev = [
     "joblib-stubs>=1.5.2.0.20250831",
+    "pyright>=1.1.407",
 ]
 
 [tool.pyright]

bella_companion-0.0.45/bella_companion/eucovid/run_beast.py

@@ -1,111 +0,0 @@
-import os
-from glob import glob
-from pathlib import Path
-
-from bella_companion.utils import submit_job
-
-
-def run_beast():
-    base_output_dir = Path(os.environ["BELLA_BEAST_OUTPUT_DIR"])
-    output_dir = base_output_dir / "eucovid"
-    os.makedirs(output_dir, exist_ok=True)
-
-    base_log_dir = Path(os.environ["BELLA_SBATCH_LOG_DIR"]) / "eucovid"
-
-    data_dir = Path(__file__).parent / "data"
-    beast_configs_dir = Path(__file__).parent / "beast_configs"
-    msa_file = data_dir / "msa.fasta"
-
-    predictors_dir = data_dir / "predictors"
-    all_predictor_files = ",".join(glob(str(predictors_dir / "all" / "*.tsv")))
-    all_predictors_data = " ".join(
-        [
-            f"-D msa_file={msa_file}",
-            f"-D changeTimesFile={predictors_dir / 'changetimes_all_7e.tsv'}",
-            f"-D predictorFiles={all_predictor_files}",
-        ]
-    )
-    flight_predictor_data = " ".join(
-        [
-            f"-D msa_file={msa_file}",
-            f"-D changeTimesFile={predictors_dir / 'changetimes_flights_4e.tsv'}",
-            f"-D predictorFiles={predictors_dir / 'flight_pop_x_4e_ls.tsv'}",
-        ]
-    )
-
-    os.makedirs(output_dir / "Nonparametric", exist_ok=True)
-    submit_job(
-        " ".join(
-            [
-                os.environ["BELLA_RUN_BEAST_CMD"],
-                f"-D aligned_fasta={msa_file}",
-                f"-prefix {output_dir / 'Nonparametric'}{os.sep}",
-                str(beast_configs_dir / "Nonparametric.xml"),
-            ]
-        ),
-        base_log_dir / "Nonparametric",
-        cpus=128,
-        mem_per_cpu=12000,
-    )
-
-    os.makedirs(output_dir / "all-predictors-GLM", exist_ok=True)
-    submit_job(
-        " ".join(
-            [
-                os.environ["BELLA_RUN_BEAST_CMD"],
-                all_predictors_data,
-                f"-prefix {output_dir / 'all-predictors-GLM'}{os.sep}",
-                str(beast_configs_dir / "GLM.xml"),
-            ]
-        ),
-        base_log_dir / "all-predictors-GLM",
-        cpus=128,
-        mem_per_cpu=12000,
-    )
-
-    os.makedirs(output_dir / "flights-GLM", exist_ok=True)
-    submit_job(
-        " ".join(
-            [
-                os.environ["BELLA_RUN_BEAST_CMD"],
-                flight_predictor_data,
-                f"-prefix {output_dir / 'flights-GLM'}{os.sep}",
-                str(beast_configs_dir / "GLM.xml"),
-            ]
-        ),
-        base_log_dir / "flights-GLM",
-        cpus=128,
-        mem_per_cpu=12000,
-    )
-
-    os.makedirs(output_dir / "all-predictors-MLP", exist_ok=True)
-    submit_job(
-        " ".join(
-            [
-                os.environ["BELLA_RUN_BEAST_CMD"],
-                f'-D layersRange="0,1,2",nodes="16 8"',
-                all_predictors_data,
-                f"-prefix {output_dir / 'all-predictors-MLP'}{os.sep}",
-                str(beast_configs_dir / "MLP.xml"),
-            ]
-        ),
-        base_log_dir / "all-predictors-MLP",
-        cpus=128,
-        mem_per_cpu=12000,
-    )
-
-    os.makedirs(output_dir / "flights-MLP", exist_ok=True)
-    submit_job(
-        " ".join(
-            [
-                os.environ["BELLA_RUN_BEAST_CMD"],
-                f'-D layersRange="0,1,2",nodes="16 8"',
-                flight_predictor_data,
-                f"-prefix {output_dir / 'flights-MLP'}{os.sep}",
-                str(beast_configs_dir / "MLP.xml"),
-            ]
-        ),
-        base_log_dir / "flights-MLP",
-        cpus=128,
-        mem_per_cpu=12000,
-    )

bella_companion-0.0.45/bella_companion/platyrrhine/results.py

@@ -1,109 +0,0 @@
-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 draw_tree, get_node_depths, load_newick, load_nexus
-
-
-def _plot_tree(output_dir: Path, tree_file: str | Path):
-    sample_trees = load_nexus(tree_file)
-    trees = [sample_trees[f"STATE_{i*100_000}"] for i in range(11, 101)]
-
-    for tree in trees:
-        for node in tree:
-            if node.parent is not None and not node.branch_length:
-                node.parent.remove_child(node)
-
-    avg_tree = trees[0].copy()
-    for node in avg_tree:
-        node.clear()
-    for avg_node, nodes in zip(avg_tree, zip(*trees)):
-        avg_node.set("birthRateSP", np.mean([n["birthRateSP"] for n in nodes]))
-        avg_node.set("deathRateSP", np.mean([n["deathRateSP"] for n in nodes]))
-
-    avg_tree.ladderize()
-    for rate in ["birthRateSP", "deathRateSP"]:
-        ax = plt.gca()
-        draw_tree(
-            avg_tree,
-            ax,
-            color_by=rate,
-            hist_axes_kwargs={
-                "loc": "upper left",
-                "bbox_to_anchor": (0.06, 0, 1, 1),
-                "bbox_transform": ax.transAxes,
-            },
-        )
-        height = avg_tree.height
-        ticks = np.arange(height, 0, -10)
-        ax.set_xticks(ticks, [str(round(height - t)) for t in ticks])  # pyright: ignore
-        plt.xlabel("Time (mya)")  # pyright: ignore
-        plt.savefig(output_dir / f"{rate}-tree.svg")  # pyright: ignore
-        plt.close()
-
-
-def _plot_predictions(output_dir: Path, log_summary_file: str | Path):
-    log_summary = pl.read_csv(log_summary_file)
-
-    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]))
-
-    gradient = np.linspace(0.4, 0.9, 4)
-    colors: dict[str, NDArray[np.floating]] = {
-        "birthRateSP": plt.cm.Blues(gradient),  # pyright: ignore
-        "deathRateSP": plt.cm.Oranges(gradient),  # pyright: ignore
-        "diversificationRateSP": plt.cm.Greens(gradient),  # pyright: ignore
-    }
-    for rate in colors:
-        for type in [0, 1, 2, 3]:
-            if rate == "diversificationRateSP":
-                estimates = log_summary.select(
-                    [
-                        pl.col(f"birthRateSPi{i}_{type}_median")
-                        - pl.col(f"deathRateSPi{i}_{type}_median")
-                        for i in range(len(change_times) + 1)
-                    ]
-                ).to_numpy()
-            else:
-                estimates = log_summary.select(
-                    [
-                        pl.col(f"{rate}i{i}_{type}_median")
-                        for i in range(len(change_times) + 1)
-                    ]
-                ).to_numpy()
-            median = np.median(estimates, axis=0)
-            color = colors[rate][type]
-            plt.step(  # pyright: ignore
-                time_bins, [median[0], *median], color=color, label=type
-            )
-
-        plt.gca().invert_xaxis()
-        plt.legend(title="Body mass")  # pyright: ignore
-        plt.xlabel("Time (mya)")  # pyright: ignore
-        plt.ylabel(  # pyright: ignore
-            r"$\lambda$"
-            if rate == "birthRateSP"
-            else r"$\mu$" if rate == "deathRateSP" else r"$d$"
-        )
-        plt.savefig(output_dir / f"{rate}-predictions.svg")  # pyright: ignore
-        plt.close()
-
-
-def plot_platyrrhine_results():
-    summaries_dir = Path(os.environ["BELLA_SUMMARIES_DIR"]) / "platyrrhine"
-    output_dir = Path(os.environ["BELLA_FIGURES_DIR"]) / "platyrrhine"
-    os.makedirs(output_dir, exist_ok=True)
-    _plot_predictions(output_dir, summaries_dir / "MLP.csv")
-    _plot_tree(output_dir, summaries_dir / "sample-trees.nexus")