phylogenie 2.0.14__py3-none-any.whl → 2.1.1__py3-none-any.whl

phylogenie/treesimulator/gillespie.py

@@ -1,86 +1,123 @@
+import os
 from collections.abc import Sequence
 
+import joblib
 import numpy as np
 from numpy.random import default_rng
+from tqdm import tqdm
 
-from phylogenie.skyline import SkylineParameterLike
+from phylogenie.io import dump_newick
 from phylogenie.tree import Tree
-from phylogenie.treesimulator.events import Event, get_contact_tracing_events
-from phylogenie.treesimulator.model import Model
+from phylogenie.treesimulator.model import Event, Model
+
+MAX_TIPS = 2**32
 
 
 def simulate_tree(
     events: Sequence[Event],
     min_tips: int = 1,
-    max_tips: int | None = None,
+    max_tips: int = MAX_TIPS,
     max_time: float = np.inf,
     init_state: str | None = None,
     sampling_probability_at_present: float = 0.0,
-    notification_probability: float = 0,
-    max_notified_contacts: int = 1,
-    samplable_states_after_notification: Sequence[str] | None = None,
-    sampling_rate_after_notification: SkylineParameterLike = np.inf,
-    contacts_removal_probability: SkylineParameterLike = 1,
-    max_tries: int | None = None,
     seed: int | None = None,
 ) -> Tree | None:
-    rng = default_rng(seed)
-
-    if max_tips is None and max_time == np.inf:
-        raise ValueError("Either max_tips or max_time must be specified.")
+    if max_time == np.inf and max_tips == MAX_TIPS:
+        raise ValueError("Either max_time or max_tips must be specified.")
 
-    root_states = list({e.state for e in events})
+    if max_time == np.inf and sampling_probability_at_present:
+        raise ValueError(
+            "sampling_probability_at_present cannot be set when max_time is infinite."
+        )
 
-    if notification_probability:
-        events = get_contact_tracing_events(
-            events,
-            samplable_states_after_notification,
-            sampling_rate_after_notification,
-            contacts_removal_probability,
+    states = {e.state for e in events if e.state}
+    if init_state is None and len(states) > 1:
+        raise ValueError(
+            "Init state must be provided for models with more than one state."
         )
+    elif init_state is None:
+        (init_state,) = states
+    elif init_state not in states:
+        raise ValueError(f"Init state {init_state} not found in event states: {states}")
 
-    n_tries = 0
-    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)
+    rng = default_rng(seed)
+    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
-        n_tips = None if max_tips is None else rng.integers(min_tips, max_tips + 1)
 
-        while current_time < max_time and (n_tips is None or model.n_sampled < n_tips):
+        target_n_tips = rng.integers(min_tips, max_tips) if max_time == np.inf else None
+        while current_time < max_time:
+            events = model.events
             rates = [e.get_propensity(model, current_time) for e in events]
 
             instantaneous_events = [e for e, r in zip(events, rates) if r == np.inf]
             if instantaneous_events:
                 event = instantaneous_events[rng.integers(len(instantaneous_events))]
-                event.apply(model, current_time)
+                event.apply(model, current_time, rng)
                 continue
 
-            if not any(rates):
+            if (
+                not any(rates)
+                or model.n_sampled > max_tips
+                or target_n_tips is not None
+                and model.n_sampled >= target_n_tips
+            ):
                 break
 
             time_step = rng.exponential(1 / sum(rates))
+            if current_time + time_step >= next_change_time:
+                current_time = next_change_time
+                next_change_time = change_times.pop(0) if change_times else np.inf
+                continue
             if current_time + time_step >= max_time:
                 current_time = max_time
                 break
             current_time += time_step
 
-            if current_time >= next_change_time:
-                current_time = next_change_time
-                next_change_time = change_times.pop(0) if change_times else np.inf
-                continue
-
             event_idx = np.searchsorted(np.cumsum(rates) / sum(rates), rng.random())
-            events[int(event_idx)].apply(model, current_time)
+            events[int(event_idx)].apply(model, current_time, rng)
 
         for individual in model.get_population():
             if rng.random() < sampling_probability_at_present:
-                model.sample(individual, current_time, 1)
+                model.sample(individual, current_time, True)
 
-        if model.n_sampled >= min_tips and (
-            max_tips is None or model.n_sampled <= max_tips
-        ):
+        if min_tips <= model.n_sampled <= max_tips:
             return model.get_sampled_tree()
 
-        n_tries += 1
+
+def generate_trees(
+    output_dir: str,
+    n_trees: int,
+    events: list[Event],
+    min_tips: int = 1,
+    max_tips: int = 2**32,
+    max_time: float = np.inf,
+    init_state: str | None = None,
+    sampling_probability_at_present: float = 0.0,
+    seed: int | None = None,
+    n_jobs: int = -1,
+) -> None:
+    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)),
+        )
+        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"))

phylogenie/treesimulator/model.py

@@ -1,62 +1,74 @@
+from abc import ABC, abstractmethod
 from collections import defaultdict
-from dataclasses import dataclass, field
-from typing import ClassVar
+from collections.abc import Sequence
+from copy import deepcopy
+from dataclasses import dataclass
+from typing import Any
 
-from numpy.random import Generator, default_rng
+import numpy as np
+from numpy.random import Generator
 
+from phylogenie.skyline import SkylineParameterLike, skyline_parameter
 from phylogenie.tree import Tree
 
-CT_POSTFIX = "-CT"
 
+class Event(ABC):
+    def __init__(self, state: str, rate: SkylineParameterLike):
+        self.state = state
+        self.rate = skyline_parameter(rate)
 
-def get_CT_state(state: str) -> str:
-    return f"{state}{CT_POSTFIX}"
+    def draw_individual(self, model: "Model", rng: Generator) -> int:
+        return rng.choice(model.get_population(self.state))
 
+    def get_propensity(self, model: "Model", time: float) -> float:
+        n_individuals = model.count_individuals(self.state)
+        rate = self.rate.get_value_at_time(time)
+        if rate == np.inf and not n_individuals:
+            return 0
+        return rate * n_individuals
 
-def is_CT_state(state: str) -> bool:
-    return state.endswith(CT_POSTFIX)
+    @abstractmethod
+    def apply(self, model: "Model", time: float, rng: Generator) -> None: ...
 
 
 @dataclass
 class Individual:
+    id: int
     node: Tree
     state: str
-    id: int = field(init=False)
-    _id_counter: ClassVar[int] = 0
-
-    def __post_init__(self):
-        Individual._id_counter += 1
-        self.id = Individual._id_counter
 
 
 class Model:
-    def __init__(
-        self,
-        init_state: str,
-        max_notified_contacts: int = 1,
-        notification_probability: float = 0,
-        rng: int | Generator | None = None,
-    ):
+    def __init__(self, init_state: str, events: Sequence[Event]):
         self._next_node_id = 0
+        self._next_individual_id = 0
         self._population: dict[int, Individual] = {}
         self._states: dict[str, set[int]] = defaultdict(set)
-        self._contacts: dict[int, list[Individual]] = defaultdict(list)
         self._sampled: set[str] = set()
         self._tree = self._get_new_individual(init_state).node
-        self._max_notified_contacts = max_notified_contacts
-        self._notification_probability = notification_probability
-        self._rng = rng if isinstance(rng, Generator) else default_rng(rng)
+        self._events = list(events)
+        self.context: dict[str, Any] = {}
 
     @property
     def n_sampled(self) -> int:
         return len(self._sampled)
 
+    @property
+    def events(self) -> tuple[Event, ...]:
+        return tuple(self._events)
+
+    def add_event(self, event: Event) -> None:
+        self._events.append(event)
+
     def _get_new_node(self, state: str) -> Tree:
         self._next_node_id += 1
         return Tree(f"{self._next_node_id}|{state}")
 
     def _get_new_individual(self, state: str) -> Individual:
-        individual = Individual(self._get_new_node(state), state)
+        self._next_individual_id += 1
+        individual = Individual(
+            self._next_individual_id, self._get_new_node(state), state
+        )
         self._population[individual.id] = individual
         self._states[state].add(individual.id)
         return individual
@@ -77,9 +89,8 @@ class Model:
     def remove(self, id: int, time: float) -> None:
         individual = self._population[id]
         self._set_branch_length(individual.node, time)
-        state = individual.state
         self._population.pop(id)
-        self._states[state].remove(id)
+        self._states[individual.state].remove(id)
 
     def migrate(self, id: int, state: str, time: float) -> None:
         individual = self._population[id]
@@ -88,17 +99,16 @@ class Model:
         self._states[state].add(id)
         self._stem(individual, time)
 
-    def birth_from(self, id: int, state: str, time: float) -> None:
+    def birth_from(self, id: int, state: str, time: float) -> int:
         individual = self._population[id]
         new_individual = self._get_new_individual(state)
         individual.node.add_child(new_individual.node)
         self._stem(individual, time)
-        self._contacts[id].append(new_individual)
-        self._contacts[new_individual.id].append(individual)
+        return new_individual.id
 
-    def sample(self, id: int, time: float, removal_probability: float) -> None:
+    def sample(self, id: int, time: float, removal: bool) -> None:
         individual = self._population[id]
-        if self._rng.random() < removal_probability:
+        if removal:
             self._sampled.add(individual.node.id)
             self.remove(id, time)
         else:
@@ -108,47 +118,38 @@ class Model:
             individual.node.add_child(sample_node)
             self._stem(individual, time)
 
-        for contact in self._contacts[id][-self._max_notified_contacts :]:
-            if (
-                contact.id in self._population
-                and not is_CT_state(contact.state)
-                and self._rng.random() < self._notification_probability
-            ):
-                self.migrate(contact.id, get_CT_state(contact.state), time)
+    def get_state(self, id: int) -> str:
+        return self._population[id].state
 
     def get_sampled_tree(self) -> Tree:
-        tree = self._tree.copy()
+        tree = deepcopy(self._tree)
         for node in list(tree.postorder_traversal()):
             if node.id not in self._sampled and not node.children:
                 if node.parent is None:
                     raise ValueError("No samples in the tree.")
                 else:
-                    node.parent.children.remove(node)
+                    node.parent.remove_child(node)
             elif len(node.children) == 1:
                 (child,) = node.children
-                child.parent = node.parent
+                child.set_parent(node.parent)
                 assert child.branch_length is not None
                 assert node.branch_length is not None
                 child.branch_length += node.branch_length
                 if node.parent is None:
                     return child
                 else:
-                    node.parent.children.append(child)
-                    node.parent.children.remove(node)
+                    node.parent.remove_child(node)
         return tree
 
     def get_full_tree(self) -> Tree:
-        return self._tree.copy()
-
-    def get_random_individual(self, state: str | None = None) -> int:
-        if state is None:
-            return self._rng.choice(list(self._population))
-        return self._rng.choice(list(self._states[state]))
+        return deepcopy(self._tree)
 
-    def get_population(self) -> list[int]:
-        return list(self._population)
+    def get_population(self, states: str | None = None) -> list[int]:
+        if states is None:
+            return list(self._population)
+        return list(self._states[states])
 
-    def count_individuals(self, state: str | None = None) -> int:
-        if state is None:
+    def count_individuals(self, states: str | None = None) -> int:
+        if states is None:
             return len(self._population)
-        return len(self._states[state])
+        return len(self._states[states])

phylogenie/typings.py

@@ -16,6 +16,5 @@ Many2DScalars = Many2D[Scalar]
 Many3DScalars = Many3D[Scalar]
 
 Vector1D = list[Scalar]
-IntVector1D = list[int]
 Vector2D = list[Vector1D]
 Vector3D = list[Vector2D]

phylogenie/utils.py

@@ -0,0 +1,17 @@
+from typing import Any
+
+from pydantic import BaseModel, ConfigDict
+
+
+class StrictBaseModel(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+
+class Distribution(BaseModel):
+    type: str
+    model_config = ConfigDict(extra="allow")
+
+    @property
+    def args(self) -> dict[str, Any]:
+        assert self.model_extra is not None
+        return self.model_extra

{phylogenie-2.0.14.dist-info → phylogenie-2.1.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: phylogenie
-Version: 2.0.14
+Version: 2.1.1
 Summary: Generate phylogenetic datasets with minimal setup effort
 Author: Gabriele Marino
 Author-email: gabmarino.8601@gmail.com
@@ -24,7 +24,6 @@ Description-Content-Type: text/markdown
 
 [![AliSim](https://img.shields.io/badge/Powered%20by-AliSim-orange?style=flat-square)](https://iqtree.github.io/doc/AliSim)
 [![PyPI version](https://img.shields.io/pypi/v/phylogenie)](https://pypi.org/project/phylogenie/)
-[![PyPI downloads](https://shields.io/pypi/dm/phylogenie)](https://pypi.org/project/phylogenie/)
 
 Phylogenie is a [Python](https://www.python.org/) package designed to easily simulate phylogenetic datasets—such as trees and multiple sequence alignments (MSAs)—with minimal setup effort. Simply specify the distributions from which your parameters should be sampled, and Phylogenie will handle the rest!
 

phylogenie-2.1.1.dist-info/RECORD

@@ -0,0 +1,32 @@
+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/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/typeguards.py,sha256=yj4VkhOaUXJ2OrY-6zhOeY9C4yKIQxjZtk2d-vIxttQ,828
+phylogenie/io.py,sha256=y7nQIvLgCvqELsXFKfm1GgKJO_saoQ-7zQpE3Kvajzc,3509
+phylogenie/main.py,sha256=vtvSpQxBNlYABoFQ25czl-l3fIr4QRo3svWVd-jcArw,1170
+phylogenie/msa.py,sha256=JDGyZUsAq6-m-SQjoCDjAkAZIxfgyl_PDIhdYn5HOow,2064
+phylogenie/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+phylogenie/skyline/__init__.py,sha256=7pF4CUb4ZCLzNYJNhOjpuTOLTRhlK7L6ugfccNqjIGo,620
+phylogenie/skyline/matrix.py,sha256=Gl8OgKjtieG0NwPYiPimKI36gefV8fm_OeorjdXxPTs,9146
+phylogenie/skyline/parameter.py,sha256=EM9qlPt0JhMBy3TbztM0dj24BaGNEy8KWKdTObDKhbI,4644
+phylogenie/skyline/vector.py,sha256=bJP7_FNX_Klt6wXqsyfj0KX3VNj6-dIhzCKSJuQcOV0,7115
+phylogenie/tree.py,sha256=fmCLDNIls3VjY413b-khzj-2jlaQ9rJo7AFN1cz4zXk,2411
+phylogenie/treesimulator/__init__.py,sha256=XG_xwETKWgDmCihqNUFCcMHtFg4WvZu5qbqWn9Dndt8,879
+phylogenie/treesimulator/events/__init__.py,sha256=UGfvXOVJ_ZAkk_8sBPihjmxciiaEnXZEPFIY53sttWI,940
+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/model.py,sha256=bWoFiO-99tTjHtueKPvGV9RSHQKYJ625XP1dvYbP1_Q,5454
+phylogenie/typeguards.py,sha256=JtqmbEWJZBRHbWgCvcl6nrWm3VcBfzRbklbTBYHItn0,1325
+phylogenie/typings.py,sha256=GknvAFXyiaWeeYJ8Lk5d6E2VHT-xW6ONEojYbtJYiB8,476
+phylogenie/utils.py,sha256=pCg9ob0RpLUHwM49x4knKxL4FNPr3-EU_6zMXsvxtAg,370
+phylogenie-2.1.1.dist-info/LICENSE.txt,sha256=NUrDqElK-eD3I0WqC004CJsy6cs0JgsAoebDv_42-pw,1071
+phylogenie-2.1.1.dist-info/METADATA,sha256=Ay--Y2G4F1KwK9GgJFhtiTwGyKJ815JUCLYbXMyek1o,5375
+phylogenie-2.1.1.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+phylogenie-2.1.1.dist-info/entry_points.txt,sha256=Rt6_usN0FkBX1ZfiqCirjMN9FKOgFLG8rydcQ8kugeE,51
+phylogenie-2.1.1.dist-info/RECORD,,

phylogenie-2.0.14.dist-info/RECORD

@@ -1,28 +0,0 @@
-phylogenie/__init__.py,sha256=4BytT42_M1K6T3W9eqHQCrKc6g0Lh5LTQxP8dGIJTsk,1915
-phylogenie/generators/__init__.py,sha256=VCpuvmOoY_N6p1h_Q0peYgGIIUIFLsvZT3T7vbHG6w0,1090
-phylogenie/generators/alisim.py,sha256=dDqlSwLDbRE2u5SZlsq1mArobTBtuk0aeXY3m1N-bWA,2374
-phylogenie/generators/configs.py,sha256=4jSBUZiFo2GacXWed5dy7lUEkaOWZkZG-KY9vHfhqGU,993
-phylogenie/generators/dataset.py,sha256=hbkN5McM4BKY7D0hLNaxdoAGsLHac6O-D4sgnZ0wFX4,2618
-phylogenie/generators/factories.py,sha256=0ckeAsKnPy69Vbdoi1rIyf6zRcqamz9VfSi0mAiTzds,6938
-phylogenie/generators/trees.py,sha256=jukaVXGcPGzDBEYMGJ1MKqWt4XbAB5EEfuHXDpwKTqM,9173
-phylogenie/generators/typeguards.py,sha256=yj4VkhOaUXJ2OrY-6zhOeY9C4yKIQxjZtk2d-vIxttQ,828
-phylogenie/io.py,sha256=ZXlofnSh7FX5UJiP0svRHrTraMSNgKa1GiAv0bMz7jU,2854
-phylogenie/main.py,sha256=vtvSpQxBNlYABoFQ25czl-l3fIr4QRo3svWVd-jcArw,1170
-phylogenie/msa.py,sha256=JDGyZUsAq6-m-SQjoCDjAkAZIxfgyl_PDIhdYn5HOow,2064
-phylogenie/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-phylogenie/skyline/__init__.py,sha256=7pF4CUb4ZCLzNYJNhOjpuTOLTRhlK7L6ugfccNqjIGo,620
-phylogenie/skyline/matrix.py,sha256=Gl8OgKjtieG0NwPYiPimKI36gefV8fm_OeorjdXxPTs,9146
-phylogenie/skyline/parameter.py,sha256=EM9qlPt0JhMBy3TbztM0dj24BaGNEy8KWKdTObDKhbI,4644
-phylogenie/skyline/vector.py,sha256=bJP7_FNX_Klt6wXqsyfj0KX3VNj6-dIhzCKSJuQcOV0,7115
-phylogenie/tree.py,sha256=dk8Sj1tqyGOunVO2crtIqb0LH-ws-PXqA8SuNcYfVHI,1738
-phylogenie/treesimulator/__init__.py,sha256=INPU9LrPdUmt3dYGzWDRoRKrPR9xENcHu44pJVUbyNA,525
-phylogenie/treesimulator/events.py,sha256=X3_0U9qqMpYgh6-7TwQEnlUipANkHz6QTCXlm-qXFQk,9524
-phylogenie/treesimulator/gillespie.py,sha256=NHiBLVlpKdcTPohMr6ehywbPeB8gnn6erEv2bkgvf4c,3268
-phylogenie/treesimulator/model.py,sha256=XpzAicmg2O6K0Trk5YolH-B_HJZxoSauF2wZOMqp-Iw,5559
-phylogenie/typeguards.py,sha256=JtqmbEWJZBRHbWgCvcl6nrWm3VcBfzRbklbTBYHItn0,1325
-phylogenie/typings.py,sha256=O1X6lGKTjJ2YJz3ApQ-rYb_tEJNUIcHdUIeYlSM4s5o,500
-phylogenie-2.0.14.dist-info/LICENSE.txt,sha256=NUrDqElK-eD3I0WqC004CJsy6cs0JgsAoebDv_42-pw,1071
-phylogenie-2.0.14.dist-info/METADATA,sha256=NXOQmoUHSTA2AjrXjt4VKSaM2KO6H2rBF7nrRuGK2f4,5473
-phylogenie-2.0.14.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-phylogenie-2.0.14.dist-info/entry_points.txt,sha256=Rt6_usN0FkBX1ZfiqCirjMN9FKOgFLG8rydcQ8kugeE,51
-phylogenie-2.0.14.dist-info/RECORD,,