phylogenie 2.1.2__py3-none-any.whl → 2.1.3__py3-none-any.whl

phylogenie/generators/alisim.py

@@ -3,9 +3,10 @@ import subprocess
 from pathlib import Path
 from typing import Any, Literal
 
-from numpy.random import Generator
+from numpy.random import Generator, default_rng
 
 from phylogenie.generators.dataset import DatasetGenerator, DataType
+from phylogenie.generators.factories import data
 from phylogenie.generators.trees import TreeDatasetGeneratorConfig
 from phylogenie.io import dump_newick
 
@@ -42,9 +43,12 @@ class AliSimDatasetGenerator(DatasetGenerator):
         subprocess.run(command, check=True, stdout=subprocess.DEVNULL)
         subprocess.run(["rm", f"{tree_file}.log"], check=True)
 
-    def _generate_one(
-        self, filename: str, rng: Generator, data: dict[str, Any]
-    ) -> None:
+    def generate_one(
+        self,
+        filename: str,
+        context: dict[str, Any] | None = None,
+        seed: int | None = None,
+    ) -> dict[str, Any]:
         if self.keep_trees:
             base_dir, file_id = Path(filename).parent, Path(filename).stem
             trees_dir = os.path.join(base_dir, TREES_DIRNAME)
@@ -57,16 +61,24 @@ class AliSimDatasetGenerator(DatasetGenerator):
             tree_filename = f"{filename}.temp-tree"
             msa_filename = filename
 
-        tree = self.trees.simulate_one(rng, data)
-        if tree is None:
-            return
+        d: dict[str, Any] = {"file_id": Path(msa_filename).stem}
+        rng = default_rng(seed)
+        while True:
+            d.update(data(context, rng))
+            try:
+                tree = self.trees.simulate_one(d, seed)
+                break
+            except TimeoutError:
+                print(
+                    "Tree simulation timed out, retrying with different parameters..."
+                )
 
         for leaf in tree.get_leaves():
             leaf.id += f"|{leaf.get_time()}"
         dump_newick(tree, f"{tree_filename}.nwk")
 
-        self._generate_one_from_tree(
-            filename=msa_filename, tree_file=f"{tree_filename}.nwk", rng=rng, data=data
-        )
+        self._generate_one_from_tree(msa_filename, f"{tree_filename}.nwk", rng, d)
         if not self.keep_trees:
             os.remove(f"{tree_filename}.nwk")
+
+        return d

phylogenie/generators/dataset.py

@@ -1,16 +1,13 @@
 import os
 from abc import ABC, abstractmethod
 from enum import Enum
-from itertools import product
 from typing import Any
 
 import joblib
-import numpy as np
 import pandas as pd
 from numpy.random import Generator, default_rng
 from tqdm import tqdm
 
-from phylogenie.generators.factories import distribution
 from phylogenie.utils import Distribution, StrictBaseModel
 
 
@@ -31,15 +28,12 @@ class DatasetGenerator(ABC, StrictBaseModel):
     context: dict[str, Distribution] | None = None
 
     @abstractmethod
-    def _generate_one(
-        self, filename: str, rng: Generator, data: dict[str, Any]
-    ) -> None: ...
-
     def generate_one(
-        self, filename: str, data: dict[str, Any] | None = None, seed: int | None = None
-    ) -> None:
-        data = {} if data is None else data
-        self._generate_one(filename=filename, rng=default_rng(seed), data=data)
+        self,
+        filename: str,
+        context: dict[str, Any] | None = None,
+        seed: int | None = None,
+    ) -> dict[str, Any]: ...
 
     def _generate(self, rng: Generator, n_samples: int, output_dir: str) -> None:
         if os.path.exists(output_dir):
@@ -53,24 +47,18 @@ class DatasetGenerator(ABC, StrictBaseModel):
         )
         os.makedirs(data_dir)
 
-        data: list[dict[str, Any]] = [{} for _ in range(n_samples)]
-        if self.context is not None:
-            for d, (k, v) in product(data, self.context.items()):
-                dist = distribution(v, d)
-                d[k] = np.array(getattr(rng, dist.type)(**dist.args)).tolist()
-            df = pd.DataFrame([{"file_id": str(i), **d} for i, d in enumerate(data)])
-            df.to_csv(os.path.join(output_dir, METADATA_FILENAME), index=False)
-
         jobs = joblib.Parallel(n_jobs=self.n_jobs, return_as="generator_unordered")(
             joblib.delayed(self.generate_one)(
-                filename=os.path.join(data_dir, str(i)),
-                data=data[i],
                 seed=int(rng.integers(2**32)),
+                filename=os.path.join(data_dir, str(i)),
+                context=self.context,
             )
             for i in range(n_samples)
         )
-        for _ in tqdm(jobs, total=n_samples, desc=f"Generating {data_dir}..."):
-            pass
+        df = pd.DataFrame(
+            [r for r in tqdm(jobs, total=n_samples, desc=f"Generating {data_dir}...")]
+        )
+        df.to_csv(os.path.join(output_dir, METADATA_FILENAME), index=False)
 
     def generate(self) -> None:
         rng = default_rng(self.seed)

phylogenie/generators/factories.py

@@ -1,6 +1,7 @@
 from typing import Any
 
 import numpy as np
+from numpy.random import Generator
 
 import phylogenie.generators.configs as cfg
 import phylogenie.generators.typeguards as ctg
@@ -30,14 +31,6 @@ def _eval_expression(expression: str, data: dict[str, Any]) -> Any:
     ).tolist()
 
 
-def distribution(x: Distribution, data: dict[str, Any]) -> Distribution:
-    args = x.args
-    for arg_name, arg_value in args.items():
-        if isinstance(arg_value, str):
-            args[arg_name] = _eval_expression(arg_value, data)
-    return Distribution(type=x.type, **args)
-
-
 def integer(x: cfg.Integer, data: dict[str, Any]) -> int:
     if isinstance(x, str):
         e = _eval_expression(x, data)
@@ -209,3 +202,21 @@ def skyline_matrix(
     value = [[[e] * N] * M if isinstance(e, pgt.Scalar) else e for e in value]
 
     return SkylineMatrix(value=value, change_times=change_times)
+
+
+def distribution(x: Distribution, data: dict[str, Any]) -> Distribution:
+    args = x.args
+    for arg_name, arg_value in args.items():
+        if isinstance(arg_value, str):
+            args[arg_name] = _eval_expression(arg_value, data)
+    return Distribution(type=x.type, **args)
+
+
+def data(context: dict[str, Distribution] | None, rng: Generator) -> dict[str, Any]:
+    if context is None:
+        return {}
+    data: dict[str, Any] = {}
+    for k, v in context.items():
+        dist = distribution(v, data)
+        data[k] = np.array(getattr(rng, dist.type)(**dist.args)).tolist()
+    return data

phylogenie/generators/trees.py

@@ -1,14 +1,16 @@
 from abc import abstractmethod
 from enum import Enum
+from pathlib import Path
 from typing import Annotated, Any, Literal
 
 import numpy as np
-from numpy.random import Generator
+from numpy.random import default_rng
 from pydantic import Field
 
 import phylogenie.generators.configs as cfg
 from phylogenie.generators.dataset import DatasetGenerator, DataType
 from phylogenie.generators.factories import (
+    data,
     distribution,
     integer,
     scalar,
@@ -48,19 +50,19 @@ class TreeDatasetGenerator(DatasetGenerator):
     max_time: cfg.Scalar = np.inf
     init_state: str | None = None
     sampling_probability_at_present: cfg.Scalar = 0.0
+    timeout: float = np.inf
 
     @abstractmethod
     def _get_events(self, data: dict[str, Any]) -> list[Event]: ...
 
-    def simulate_one(self, rng: Generator, data: dict[str, Any]) -> Tree | None:
-        events = self._get_events(data)
+    def simulate_one(self, data: dict[str, Any], seed: int | None = None) -> Tree:
         init_state = (
             self.init_state
             if self.init_state is None
             else self.init_state.format(**data)
         )
         return simulate_tree(
-            events=events,
+            events=self._get_events(data),
             min_tips=integer(self.min_tips, data),
             max_tips=integer(self.max_tips, data),
             max_time=scalar(self.max_time, data),
@@ -68,15 +70,27 @@ class TreeDatasetGenerator(DatasetGenerator):
             sampling_probability_at_present=scalar(
                 self.sampling_probability_at_present, data
             ),
-            seed=int(rng.integers(2**32)),
+            seed=seed,
+            timeout=self.timeout,
         )
 
-    def _generate_one(
-        self, filename: str, rng: Generator, data: dict[str, Any]
-    ) -> None:
-        tree = self.simulate_one(rng, data)
-        if tree is not None:
-            dump_newick(tree, f"{filename}.nwk")
+    def generate_one(
+        self,
+        filename: str,
+        context: dict[str, Any] | None = None,
+        seed: int | None = None,
+    ) -> dict[str, Any]:
+        d = {"file_id": Path(filename).stem}
+        rng = default_rng(seed)
+        while True:
+            try:
+                d.update(data(context, rng))
+                tree = self.simulate_one(d, seed)
+                dump_newick(tree, f"{filename}.nwk")
+                break
+            except TimeoutError:
+                print("Simulation timed out, retrying with different parameters...")
+        return d
 
 
 class CanonicalTreeDatasetGenerator(TreeDatasetGenerator):

phylogenie/treesimulator/gillespie.py

@@ -1,4 +1,5 @@
 import os
+import time
 from collections.abc import Sequence
 
 import joblib
@@ -21,7 +22,8 @@ def simulate_tree(
     init_state: str | None = None,
     sampling_probability_at_present: float = 0.0,
     seed: int | None = None,
-) -> Tree | None:
+    timeout: float = np.inf,
+) -> Tree:
     if max_time == np.inf and max_tips == MAX_TIPS:
         raise ValueError("Either max_time or max_tips must be specified.")
 
@@ -41,14 +43,18 @@ def simulate_tree(
         raise ValueError(f"Init state {init_state} not found in event states: {states}")
 
     rng = default_rng(seed)
+    start_clock = time.perf_counter()
     while True:
         model = Model(init_state, events)
         current_time = 0.0
         change_times = sorted(set(t for e in events for t in e.rate.change_times))
         next_change_time = change_times.pop(0) if change_times else np.inf
-
         target_n_tips = rng.integers(min_tips, max_tips) if max_time == np.inf else None
+
         while current_time < max_time:
+            if time.perf_counter() - start_clock > timeout:
+                raise TimeoutError("Simulation timed out.")
+
             events = model.events
             rates = [e.get_propensity(model, current_time) for e in events]
 
@@ -98,26 +104,35 @@ def generate_trees(
     sampling_probability_at_present: float = 0.0,
     seed: int | None = None,
     n_jobs: int = -1,
+    timeout: float = np.inf,
 ) -> None:
+    def _simulate_tree(seed: int) -> Tree:
+        while True:
+            try:
+                return simulate_tree(
+                    events=events,
+                    min_tips=min_tips,
+                    max_tips=max_tips,
+                    max_time=max_time,
+                    init_state=init_state,
+                    sampling_probability_at_present=sampling_probability_at_present,
+                    seed=seed,
+                    timeout=timeout,
+                )
+            except TimeoutError:
+                print("Simulation timed out, retrying with a different seed...")
+            seed += 1
+
     if os.path.exists(output_dir):
         raise FileExistsError(f"Output directory {output_dir} already exists")
     os.mkdir(output_dir)
 
     rng = default_rng(seed)
     jobs = joblib.Parallel(n_jobs=n_jobs, return_as="generator_unordered")(
-        joblib.delayed(simulate_tree)(
-            events=events,
-            min_tips=min_tips,
-            max_tips=max_tips,
-            max_time=max_time,
-            init_state=init_state,
-            sampling_probability_at_present=sampling_probability_at_present,
-            seed=int(rng.integers(2**32)),
-        )
+        joblib.delayed(_simulate_tree)(seed=int(rng.integers(2**32)))
         for _ in range(n_trees)
     )
     for i, tree in tqdm(
         enumerate(jobs), total=n_trees, desc=f"Generating trees in {output_dir}..."
     ):
-        if tree is not None:
-            dump_newick(tree, os.path.join(output_dir, f"{i}.nwk"))
+        dump_newick(tree, os.path.join(output_dir, f"{i}.nwk"))

{phylogenie-2.1.2.dist-info → phylogenie-2.1.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: phylogenie
-Version: 2.1.2
+Version: 2.1.3
 Summary: Generate phylogenetic datasets with minimal setup effort
 Author: Gabriele Marino
 Author-email: gabmarino.8601@gmail.com

{phylogenie-2.1.2.dist-info → phylogenie-2.1.3.dist-info}/RECORD

@@ -1,10 +1,10 @@
 phylogenie/__init__.py,sha256=T2mRLsYtoLlWt8GlxrrUnfXJ9XVioq7hTvVq3uJpwQI,2215
 phylogenie/generators/__init__.py,sha256=zsOxy28-9j9alOQLIgrOAFfmM58NNHO_NEtW-KXQXAY,888
-phylogenie/generators/alisim.py,sha256=dDqlSwLDbRE2u5SZlsq1mArobTBtuk0aeXY3m1N-bWA,2374
+phylogenie/generators/alisim.py,sha256=3mANgyQrlozhslV3_ryt-m4ItkRcKKRLufWf6SNBTnQ,2781
 phylogenie/generators/configs.py,sha256=AiiFS6rpH9BPwDKCkT4SVrRzfLFFrwRCJM4CRj0Srdk,1072
-phylogenie/generators/dataset.py,sha256=wYtb5fxjrM0bD9KVDlgZPpiR4ACezdKKfsmsXyc4__0,2755
-phylogenie/generators/factories.py,sha256=14p1wvJtjdXgM4mPqV9vX34HYcKpemx9TPfemvoB5Wo,7265
-phylogenie/generators/trees.py,sha256=EEj5KPuu7FSOoti_gGTmgWgiX68rEKFh57Hp6BDw8po,10049
+phylogenie/generators/dataset.py,sha256=loVKC_1G7gzkPDN9W3GF-Rj9od8AeOJgIC0aJJa-4KA,2110
+phylogenie/generators/factories.py,sha256=jLwDuq0mrmDz2U5rZM19KJ2hSpamG3r6zb83YCc6snA,7619
+phylogenie/generators/trees.py,sha256=q03WPG82M4ucp-jyjoKEBy7TKMBzD3RkKn8hS0G0-i0,10463
 phylogenie/generators/typeguards.py,sha256=yj4VkhOaUXJ2OrY-6zhOeY9C4yKIQxjZtk2d-vIxttQ,828
 phylogenie/io.py,sha256=y7nQIvLgCvqELsXFKfm1GgKJO_saoQ-7zQpE3Kvajzc,3509
 phylogenie/main.py,sha256=vtvSpQxBNlYABoFQ25czl-l3fIr4QRo3svWVd-jcArw,1170
@@ -20,13 +20,13 @@ phylogenie/treesimulator/events/__init__.py,sha256=UGfvXOVJ_ZAkk_8sBPihjmxciiaEn
 phylogenie/treesimulator/events/contact_tracing.py,sha256=_nJ85yhgGkeruQgMHvGpDYoyhheBf8M4LgZWiWdi5dY,4801
 phylogenie/treesimulator/events/core.py,sha256=JokGmieAv2xEX7KsjBWZr05jHN1jB-XZbpxe9gwdbDA,7953
 phylogenie/treesimulator/events/mutations.py,sha256=xkXUIppbLIWZqKwVf-hi7d-_pS42TG2EPVfJA_grxBg,3443
-phylogenie/treesimulator/gillespie.py,sha256=TMDNKBkFwVyAEhBlbwxCTA61GuGwP-42HxpsAVXiU0s,4275
+phylogenie/treesimulator/gillespie.py,sha256=naoxPyZixWVkd5f7B3KhEtOFiQI4NDIp_589NCLTHKM,4831
 phylogenie/treesimulator/model.py,sha256=0Im6cFTlpMlJrSP4pTTKtvLT9qrQWV8MSTesAsBxT8g,5422
 phylogenie/typeguards.py,sha256=JtqmbEWJZBRHbWgCvcl6nrWm3VcBfzRbklbTBYHItn0,1325
 phylogenie/typings.py,sha256=GknvAFXyiaWeeYJ8Lk5d6E2VHT-xW6ONEojYbtJYiB8,476
 phylogenie/utils.py,sha256=pCg9ob0RpLUHwM49x4knKxL4FNPr3-EU_6zMXsvxtAg,370
-phylogenie-2.1.2.dist-info/LICENSE.txt,sha256=NUrDqElK-eD3I0WqC004CJsy6cs0JgsAoebDv_42-pw,1071
-phylogenie-2.1.2.dist-info/METADATA,sha256=RoxXfMaPI6auRdjONAE8HfcfhZjHX0hTmx_oJlrzaws,5375
-phylogenie-2.1.2.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-phylogenie-2.1.2.dist-info/entry_points.txt,sha256=Rt6_usN0FkBX1ZfiqCirjMN9FKOgFLG8rydcQ8kugeE,51
-phylogenie-2.1.2.dist-info/RECORD,,
+phylogenie-2.1.3.dist-info/LICENSE.txt,sha256=NUrDqElK-eD3I0WqC004CJsy6cs0JgsAoebDv_42-pw,1071
+phylogenie-2.1.3.dist-info/METADATA,sha256=_jiRJ7jNb-CPkhqExE3sWc6RXz8rrbnnxdgXZHbNScc,5375
+phylogenie-2.1.3.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+phylogenie-2.1.3.dist-info/entry_points.txt,sha256=Rt6_usN0FkBX1ZfiqCirjMN9FKOgFLG8rydcQ8kugeE,51
+phylogenie-2.1.3.dist-info/RECORD,,