phylogenie 2.0.8__tar.gz → 2.0.10__tar.gz

{phylogenie-2.0.8 → phylogenie-2.0.10}/PKG-INFO

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

{phylogenie-2.0.8 → phylogenie-2.0.10}/phylogenie/generators/configs.py

@@ -1,6 +1,9 @@
 from pydantic import BaseModel, ConfigDict
 
 import phylogenie.typings as pgt
+from phylogenie.skyline import SkylineMatrix as _SkylineMatrix
+from phylogenie.skyline import SkylineParameter as _SkylineParameter
+from phylogenie.skyline import SkylineVector as _SkylineVector
 
 
 class Distribution(BaseModel):
@@ -34,6 +37,10 @@ class SkylineMatrixModel(StrictBaseModel):
     change_times: ManyScalars
 
 
-SkylineParameter = Scalar | SkylineParameterModel
-SkylineVector = str | pgt.Scalar | list[SkylineParameter] | SkylineVectorModel
-SkylineMatrix = str | pgt.Scalar | list[SkylineVector] | SkylineMatrixModel | None
+SkylineParameter = Scalar | SkylineParameterModel | _SkylineParameter
+SkylineVector = (
+    str | pgt.Scalar | list[SkylineParameter] | SkylineVectorModel | _SkylineVector
+)
+SkylineMatrix = (
+    str | pgt.Scalar | list[SkylineVector] | SkylineMatrixModel | None | _SkylineMatrix
+)

{phylogenie-2.0.8 → phylogenie-2.0.10}/phylogenie/generators/dataset.py

@@ -41,19 +41,21 @@ class DatasetGenerator(ABC, cfg.StrictBaseModel):
         self._generate_one(filename=filename, rng=default_rng(seed), data=data)
 
     def _generate(self, rng: Generator, n_samples: int, output_dir: str) -> None:
-        data_dir = os.path.join(output_dir, DATA_DIRNAME)
-        if os.path.exists(data_dir):
-            print(f"Output directory {data_dir} already exists. Skipping.")
-            return
-        os.makedirs(data_dir)
-
         data: list[dict[str, Any]] = [{} for _ in range(n_samples)]
         if self.context is not None:
+            data_dir = os.path.join(output_dir, DATA_DIRNAME)
             for d, (k, v) in product(data, self.context.items()):
                 args = v.model_extra if v.model_extra is not None else {}
                 d[k] = np.array(getattr(rng, v.type)(**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)
+        else:
+            data_dir = output_dir
+
+        if os.path.exists(data_dir):
+            print(f"Output directory {data_dir} already exists. Skipping.")
+            return
+        os.makedirs(data_dir)
 
         joblib.Parallel(n_jobs=self.n_jobs)(
             joblib.delayed(self.generate_one)(

{phylogenie-2.0.8 → phylogenie-2.0.10}/phylogenie/generators/factories.py

@@ -80,6 +80,8 @@ def one_or_many_scalars(
 def skyline_parameter(
     x: cfg.SkylineParameter, data: dict[str, Any]
 ) -> SkylineParameterLike:
+    if isinstance(x, SkylineParameter):
+        return x
     if isinstance(x, cfg.Scalar):
         return scalar(x, data)
     return SkylineParameter(
@@ -91,6 +93,8 @@ def skyline_parameter(
 def skyline_vector(
     x: cfg.SkylineVector, data: dict[str, Any]
 ) -> SkylineVectorCoercible:
+    if isinstance(x, SkylineVector):
+        return x
     if isinstance(x, str):
         e = _eval_expression(x, data)
         if tg.is_one_or_many_scalars(e):
@@ -151,7 +155,8 @@ def skyline_matrix(
 ) -> SkylineMatrixCoercible | None:
     if x is None:
         return None
-
+    if isinstance(x, SkylineMatrix):
+        return x
     if isinstance(x, str):
         e = _eval_expression(x, data)
         if tg.is_one_or_many_2D_scalars(e):

{phylogenie-2.0.8 → phylogenie-2.0.10}/phylogenie/treesimulator/gillespie.py

@@ -6,7 +6,7 @@ from numpy.random import default_rng
 from phylogenie.skyline import SkylineParameterLike
 from phylogenie.tree import Tree
 from phylogenie.treesimulator.events import Event, get_contact_tracing_events
-from phylogenie.treesimulator.model import Model, is_CT_state
+from phylogenie.treesimulator.model import Model
 
 
 def simulate_tree(
@@ -29,6 +29,8 @@ def simulate_tree(
     if max_tips is None and max_time == np.inf:
         raise ValueError("Either max_tips or max_time must be specified.")
 
+    root_states = list({e.state for e in events})
+
     if notification_probability:
         events = get_contact_tracing_events(
             events,
@@ -38,7 +40,6 @@ def simulate_tree(
         )
 
     n_tries = 0
-    root_states = [e.state for e in events if not is_CT_state(e.state)]
     while max_tries is None or n_tries < max_tries:
         root_state = init_state if init_state is not None else rng.choice(root_states)
         model = Model(root_state, max_notified_contacts, notification_probability, rng)

{phylogenie-2.0.8 → phylogenie-2.0.10}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "phylogenie"
-version = "2.0.8"
+version = "2.0.10"
 description = "Generate phylogenetic datasets with minimal setup effort"
 authors = ["Gabriele Marino <gabmarino.8601@gmail.com>"]
 readme = "README.md"