phylogenie 1.0.7__py3-none-any.whl → 2.0.0__py3-none-any.whl

phylogenie/treesimulator/events.py

@@ -0,0 +1,256 @@
+from abc import ABC, abstractmethod
+
+from numpy.random import Generator
+
+from phylogenie.skyline import (
+    SkylineMatrixCoercible,
+    SkylineParameterLike,
+    SkylineVectorCoercible,
+    skyline_matrix,
+    skyline_parameter,
+    skyline_vector,
+)
+from phylogenie.treesimulator.model import Model
+
+INFECTIOUS_STATE = "I"
+EXPOSED_STATE = "E"
+SUPERSPREADER_STATE = "S"
+
+
+class Event(ABC):
+    def __init__(
+        self,
+        rate: SkylineParameterLike,
+        state: str | None = None,
+    ):
+        self.rate = skyline_parameter(rate)
+        self.state = state
+
+    def get_propensity(self, model: Model, time: float) -> float:
+        return self.rate.get_value_at_time(time) * model.count_leaves(self.state)
+
+    @abstractmethod
+    def apply(self, rng: Generator, model: Model, time: float) -> None: ...
+
+
+class BirthEvent(Event):
+    def __init__(
+        self,
+        rate: SkylineParameterLike,
+        state: str | None = None,
+        child_state: str | None = None,
+    ):
+        super().__init__(rate, state)
+        self.child_state = child_state
+
+    def apply(self, rng: Generator, model: Model, time: float) -> None:
+        node = model.get_random_leaf(self.state, rng)
+        model.add_child(node, time, True, self.child_state)
+
+
+class DeathEvent(Event):
+    def apply(self, rng: Generator, model: Model, time: float) -> None:
+        node = model.get_random_leaf(self.state, rng)
+        model.remove(node)
+
+
+class MigrationEvent(Event):
+    def __init__(self, state: str, target_state: str, rate: SkylineParameterLike):
+        super().__init__(rate, state)
+        self.target_state = target_state
+
+    def apply(self, rng: Generator, model: Model, time: float) -> None:
+        node = model.get_random_leaf(self.state, rng)
+        model.add_child(node, time, False, self.target_state)
+
+
+class SamplingEvent(Event):
+    def __init__(
+        self,
+        rate: SkylineParameterLike,
+        removal_probability: SkylineParameterLike,
+        state: str | None = None,
+    ):
+        super().__init__(rate, state)
+        self.removal_probability = skyline_parameter(removal_probability)
+
+    def apply(self, rng: Generator, model: Model, time: float) -> None:
+        node = model.get_random_leaf(self.state, rng)
+        remove = rng.random() < self.removal_probability.get_value_at_time(time)
+        model.sample(node, time, remove)
+
+
+def get_canonical_events(
+    sampling_rates: SkylineVectorCoercible,
+    birth_rates: SkylineVectorCoercible = 0,
+    death_rates: SkylineVectorCoercible = 0,
+    removal_probabilities: SkylineVectorCoercible = 0,
+    migration_rates: SkylineMatrixCoercible = 0,
+    birth_rates_among_states: SkylineMatrixCoercible = 0,
+    states: list[str] | None = None,
+) -> list[Event]:
+    N = 1 if states is None else len(states)
+
+    birth_rates = skyline_vector(birth_rates, N)
+    death_rates = skyline_vector(death_rates, N)
+    sampling_rates = skyline_vector(sampling_rates, N)
+    removal_probabilities = skyline_vector(removal_probabilities, N)
+    if migration_rates and N == 1:
+        raise ValueError(f"Migration rates require multiple states (got {states}).")
+    if birth_rates_among_states and N == 1:
+        raise ValueError(
+            f"Birth rates among states require multiple states (got {states})."
+        )
+
+    events: list[Event] = []
+    for i in range(N):
+        state = None if states is None else states[i]
+        events.append(BirthEvent(birth_rates[i], state, state))
+        events.append(DeathEvent(death_rates[i], state))
+        events.append(SamplingEvent(sampling_rates[i], removal_probabilities[i], state))
+        if N > 1:
+            migration_rates = skyline_matrix(migration_rates, N, N - 1)
+            birth_rates_among_states = skyline_matrix(
+                birth_rates_among_states, N, N - 1
+            )
+            assert states is not None
+            assert state is not None
+            for j, other_state in enumerate([s for s in states if s != state]):
+                events.append(MigrationEvent(state, other_state, migration_rates[i][j]))
+                events.append(
+                    BirthEvent(birth_rates_among_states[i][j], state, other_state)
+                )
+    return [event for event in events if event.rate]
+
+
+def get_epidemiological_events(
+    sampling_proportions: SkylineVectorCoercible = 1,
+    reproduction_numbers: SkylineVectorCoercible = 0,
+    become_uninfectious_rates: SkylineVectorCoercible = 0,
+    removal_probabilities: SkylineVectorCoercible = 1,
+    migration_rates: SkylineMatrixCoercible = 0,
+    reproduction_numbers_among_states: SkylineMatrixCoercible = 0,
+    states: list[str] | None = None,
+) -> list[Event]:
+    N = 1 if states is None else len(states)
+
+    reproduction_numbers = skyline_vector(reproduction_numbers, N)
+    become_uninfectious_rates = skyline_vector(become_uninfectious_rates, N)
+    sampling_proportions = skyline_vector(sampling_proportions, N)
+    removal_probabilities = skyline_vector(removal_probabilities, N)
+    if N == 1 and reproduction_numbers_among_states:
+        raise ValueError(
+            f"Reproduction numbers among states require multiple states (got {states})."
+        )
+    reproduction_numbers_among_states = (
+        skyline_matrix(reproduction_numbers_among_states, N, N - 1) if N > 1 else 0
+    )
+
+    birth_rates = reproduction_numbers * become_uninfectious_rates
+    sampling_rates = become_uninfectious_rates * sampling_proportions
+    birth_rates_among_states = (
+        reproduction_numbers_among_states * become_uninfectious_rates
+    )
+    death_rates = become_uninfectious_rates - removal_probabilities * sampling_rates
+
+    return get_canonical_events(
+        states=states,
+        birth_rates=birth_rates,
+        death_rates=death_rates,
+        sampling_rates=sampling_rates,
+        removal_probabilities=removal_probabilities,
+        migration_rates=migration_rates,
+        birth_rates_among_states=birth_rates_among_states,
+    )
+
+
+def get_FBD_events(
+    diversification: SkylineVectorCoercible = 0,
+    turnover: SkylineVectorCoercible = 0,
+    sampling_proportions: SkylineVectorCoercible = 1,
+    removal_probabilities: SkylineVectorCoercible = 0,
+    migration_rates: SkylineMatrixCoercible = 0,
+    diversification_between_types: SkylineMatrixCoercible = 0,
+    states: list[str] | None = None,
+):
+    N = 1 if states is None else len(states)
+
+    diversification = skyline_vector(diversification, N)
+    turnover = skyline_vector(turnover, N)
+    sampling_proportions = skyline_vector(sampling_proportions, N)
+    removal_probabilities = skyline_vector(removal_probabilities, N)
+    if N == 1 and diversification_between_types:
+        raise ValueError(
+            f"Diversification rates among states require multiple states (got {states})."
+        )
+    diversification_between_types = (
+        skyline_matrix(diversification_between_types, N, N - 1) if N > 1 else 0
+    )
+
+    birth_rates = diversification / (1 - turnover)
+    death_rates = turnover * birth_rates
+    sampling_rates = (
+        sampling_proportions
+        * death_rates
+        / (1 - removal_probabilities * sampling_proportions)
+    )
+    birth_rates_among_states = diversification_between_types + death_rates
+
+    return get_canonical_events(
+        states=states,
+        birth_rates=birth_rates,
+        death_rates=death_rates,
+        sampling_rates=sampling_rates,
+        removal_probabilities=removal_probabilities,
+        migration_rates=migration_rates,
+        birth_rates_among_states=birth_rates_among_states,
+    )
+
+
+def get_BD_events(
+    reproduction_number: SkylineParameterLike,
+    infectious_period: SkylineParameterLike,
+    sampling_proportion: SkylineParameterLike = 1,
+) -> list[Event]:
+    return get_epidemiological_events(
+        reproduction_numbers=reproduction_number,
+        become_uninfectious_rates=1 / infectious_period,
+        sampling_proportions=sampling_proportion,
+    )
+
+
+def get_BDEI_events(
+    reproduction_number: SkylineParameterLike,
+    infectious_period: SkylineParameterLike,
+    incubation_period: SkylineParameterLike,
+    sampling_proportion: SkylineParameterLike = 1,
+) -> list[Event]:
+    return get_epidemiological_events(
+        states=[EXPOSED_STATE, INFECTIOUS_STATE],
+        sampling_proportions=[0, sampling_proportion],
+        become_uninfectious_rates=[0, 1 / infectious_period],
+        reproduction_numbers_among_states=[[0], [reproduction_number]],
+        migration_rates=[[1 / incubation_period], [0]],
+    )
+
+
+def get_BDSS_events(
+    reproduction_number: SkylineParameterLike,
+    infectious_period: SkylineParameterLike,
+    superspreading_ratio: SkylineParameterLike,
+    superspreaders_proportion: SkylineParameterLike,
+    sampling_proportion: SkylineParameterLike = 1,
+) -> list[Event]:
+    f_SS = superspreaders_proportion
+    r_SS = superspreading_ratio
+    R_0_IS = reproduction_number * f_SS / (1 + r_SS * f_SS - f_SS)
+    R_0_SI = (reproduction_number - r_SS * R_0_IS) * r_SS
+    R_0_S = r_SS * R_0_IS
+    R_0_I = R_0_SI / r_SS
+    return get_epidemiological_events(
+        states=[INFECTIOUS_STATE, SUPERSPREADER_STATE],
+        reproduction_numbers=[R_0_I, R_0_S],
+        reproduction_numbers_among_states=[[R_0_IS], [R_0_SI]],
+        become_uninfectious_rates=1 / infectious_period,
+        sampling_proportions=sampling_proportion,
+    )

phylogenie/treesimulator/gillespie.py

@@ -0,0 +1,66 @@
+from collections.abc import Sequence
+
+import numpy as np
+from numpy.random import default_rng
+
+from phylogenie.tree import Tree
+from phylogenie.treesimulator.events import Event
+from phylogenie.treesimulator.model import Model
+
+
+def simulate_tree(
+    events: Sequence[Event],
+    min_tips: int = 1,
+    max_tips: int | None = None,
+    max_time: float = np.inf,
+    init_state: str | None = None,
+    sampling_probability_at_present: float = 0.0,
+    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.")
+
+    n_tries = 0
+    states = [e.state for e in events if e.state is not None]
+    init_state = (
+        init_state
+        if init_state is not None
+        else str(rng.choice(states)) if states else None
+    )
+    while max_tries is None or n_tries < max_tries:
+        model = Model(init_state)
+        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):
+            rates = [e.get_propensity(model, current_time) for e in events]
+            if not any(rates):
+                break
+
+            current_time += rng.exponential(1 / sum(rates))
+            if current_time >= max_time:
+                break
+
+            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(rng, model, current_time)
+
+        for leaf in model.get_leaves():
+            if rng.random() < sampling_probability_at_present:
+                model.sample(leaf, current_time, True)
+
+        if model.n_sampled >= min_tips and (
+            max_tips is None or model.n_sampled <= max_tips
+        ):
+            return model.get_sampled_tree()
+
+        n_tries += 1

phylogenie/treesimulator/model.py

@@ -0,0 +1,100 @@
+from collections import defaultdict
+
+from numpy.random import Generator, default_rng
+
+from phylogenie.tree import Tree
+
+
+class Model:
+    def __init__(self, init_state: str | None = None):
+        self._next_id = 0
+        self._n_sampled = 0
+        self._leaves: dict[str, Tree] = {}
+        self._leaf2state: dict[str, str | None] = {}
+        self._state2leaves: dict[str | None, set[str]] = defaultdict(set)
+        self._tree = self._get_new_node(init_state, None)
+
+    @property
+    def next_id(self) -> int:
+        self._next_id += 1
+        return self._next_id
+
+    @property
+    def n_sampled(self) -> int:
+        return self._n_sampled
+
+    def _get_new_node(self, state: str | None, branch_length: float | None) -> Tree:
+        id = str(self.next_id) if state is None else f"{self.next_id}|{state}"
+        node = Tree(id, branch_length)
+        if branch_length is None:
+            self._leaves[id] = node
+            self._leaf2state[id] = state
+            self._state2leaves[state].add(id)
+        return node
+
+    def remove(self, node_id: str) -> None:
+        self._state2leaves[self._leaf2state[node_id]].remove(node_id)
+        self._leaf2state.pop(node_id, None)
+        self._leaves.pop(node_id)
+
+    def add_child(
+        self,
+        node_id: str,
+        time: float,
+        stem: bool,
+        state: str | None,
+        branch_length: float | None = None,
+    ) -> None:
+        node = self._leaves[node_id]
+        if node.branch_length is not None:
+            raise ValueError("Cannot add a child to a node with a set branch length.")
+        node.add_child(self._get_new_node(state, branch_length))
+        if stem:
+            node.add_child(self._get_new_node(self._leaf2state[node.id], None))
+        node.branch_length = (
+            time if node.parent is None else time - node.parent.get_time()
+        )
+        self.remove(node_id)
+
+    def sample(self, node_id: str, time: float, remove: bool) -> None:
+        self.add_child(node_id, time, not remove, self._leaf2state[node_id], 0.0)
+        self._n_sampled += 1
+
+    def get_sampled_tree(self) -> Tree:
+        tree = self._tree.copy()
+        for node in list(tree.postorder_traversal()):
+            if node.branch_length is None or (
+                node.branch_length > 0 and not node.children
+            ):
+                if node.parent is None:
+                    raise ValueError("No samples in the tree.")
+                else:
+                    node.parent.children.remove(node)
+            elif len(node.children) == 1:
+                (child,) = node.children
+                child.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)
+        return tree
+
+    def get_random_leaf(
+        self, state: str | None = None, rng: int | Generator | None = None
+    ) -> str:
+        rng = rng if isinstance(rng, Generator) else default_rng(rng)
+        if state is None:
+            return rng.choice(list(self._leaves))
+        return rng.choice(list(self._state2leaves[state]))
+
+    def get_leaves(self) -> list[str]:
+        return list(self._leaves)
+
+    def count_leaves(self, state: str | None = None) -> int:
+        if state is None:
+            return len(self._leaves)
+        return len(self._state2leaves[state])

phylogenie/typings.py

@@ -19,5 +19,3 @@ Vector1D = list[Scalar]
 IntVector1D = list[int]
 Vector2D = list[Vector1D]
 Vector3D = list[Vector2D]
-
-Data = dict[str, str | Scalar | Vector1D | Vector2D | Vector3D]

{phylogenie-1.0.7.dist-info → phylogenie-2.0.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: phylogenie
-Version: 1.0.7
+Version: 2.0.0
 Summary: Generate phylogenetic datasets with minimal setup effort
 Author: Gabriele Marino
 Author-email: gabmarino.8601@gmail.com
@@ -23,9 +23,9 @@ Description-Content-Type: text/markdown
 
 ---
 
-[![TreeSimulator](https://img.shields.io/badge/Powered%20by-TreeSimulator-green?style=flat-square)](https://github.com/evolbioinfo/treesimulator)
-[![Remaster](https://img.shields.io/badge/Powered%20by-Remaster-blue?style=flat-square)](https://tgvaughan.github.io/remaster/)
 [![AliSim](https://img.shields.io/badge/Powered%20by-AliSim-orange?style=flat-square)](https://iqtree.github.io/doc/AliSim)
+[![PyPI package](https://badge.fury.io/py/phylogenie.svg)](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!
 
@@ -46,7 +46,7 @@ Phylogenie comes packed with useful features, including:
   Simply specify the number of cores to use, and Phylogenie handles multiprocessing automatically.
 
 - **Pre-implemented parameterizations** 🎯  
-  Include canonical, fossilized birth-death, epidemiological, birth-death with exposed-infectious (BDEI), birth-death with superspreading (BDSS), contact-tracing (CT), and more.
+  Include canonical, fossilized birth-death, epidemiological, birth-death with exposed-infectious (BDEI), birth-death with superspreading (BDSS), and more.
 
 - **Skyline parameter support** 🪜  
   Support for piece-wise constant parameters.
@@ -54,9 +54,6 @@ Phylogenie comes packed with useful features, including:
 - **Arithmetic operations on parameters** 🧮  
   Perform flexible arithmetic operations between parameters directly within the config file.
 
-- **Support for common phylogenetic simulation tools** 🛠️  
-  Compatible backends include ReMASTER, TreeSimulator, and AliSim.
-
 - **Modular and extendible architecture** 🧩  
   Easily add new simulation backends as needed.
 
@@ -76,18 +73,9 @@ cd phylogenie
 pip install .
 ```
 
-## 🛠 Backend dependencies
-
-Phylogenie works with the following simulation backends:
-
-- **[TreeSimulator](https://github.com/evolbioinfo/treesimulator)**  
-  A [Python](https://www.python.org/) package for simulating phylogenetic trees. It is automatically installed with Phylogenie, so you can use it right away.
-
-- **[ReMASTER](https://tgvaughan.github.io/remaster/)**  
-  A [BEAST2](https://www.beast2.org/) package designed for tree simulation. To use ReMASTER as a backend, you need to install it separately.
+## 🛠 Backend dependency
 
-- **[AliSim](https://iqtree.github.io/doc/AliSim)**  
-  A tool for simulating multiple sequence alignments (MSAs). It is distributed with [IQ-TREE](https://iqtree.github.io/) and also requires separate installation if you wish to use it as a backend.
+Phylogenie relies on [AliSim](https://iqtree.github.io/doc/AliSim) for simulating multiple sequence alignments (MSAs). AliSim is a powerful MSAs simulation tool distributed with [IQ-TREE](https://iqtree.github.io/), and requires separate installation to use it as a simulation backend.
 
 ## 🚀 Quick Start
 

phylogenie-2.0.0.dist-info/RECORD

@@ -0,0 +1,28 @@
+phylogenie/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+phylogenie/generators/__init__.py,sha256=kHDCQUSCl2O-wtPxcYEc1d4NTh5wjt-Oci9bYVNbhIY,440
+phylogenie/generators/alisim.py,sha256=Psh_fAHHrJBCd1kqVtz_BdoVVmNvieHJRMmapQncJhM,2288
+phylogenie/generators/configs.py,sha256=2NctdpKhOK7uuTrau18QCG1OqAAmmTl_c6QYwFnRSU0,1108
+phylogenie/generators/dataset.py,sha256=Q8GsO6b4792VACvwle26QAsbeiyOGUqID2Q3-_urVS4,2543
+phylogenie/generators/factories.py,sha256=AxLNuvZKnrSSUltr--kJ1zr-JjObD-KIip31ETSN0wk,6952
+phylogenie/generators/trees.py,sha256=CXFXeYLYjgJe7tDwCaYrtHBVLl8Wn83in_pDM_2oH6M,8717
+phylogenie/generators/typeguards.py,sha256=6WIQIy6huHok4vYqA9ym6g2kvyyXPggBONWwbiHtDiY,925
+phylogenie/io.py,sha256=lDLWz-7fsw85HCpVjtpGpw2oWxf5Aurd3_jTx2Gt_pU,2865
+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=S9CTsiK8hh4hXEnbsnPFnIXVNUFL9Oe3MMBKWRNEsSE,8951
+phylogenie/skyline/parameter.py,sha256=CJ5OEyRQG2Tg1WJWQ1IpfX-6hjJv80Zj8lMoRke5nnQ,4648
+phylogenie/skyline/vector.py,sha256=becZlWLeT12mvpDC5MLGXZPKlgmbkY-DEcuhdQpnsSc,7135
+phylogenie/tree.py,sha256=nNUQrTtUqgc-Z6CHflqROZ8W2ZL1YWGG7xeWmcKwZUU,1684
+phylogenie/treesimulator/__init__.py,sha256=koyVZxKUeMtY1S2WUgqbeIthoGbdUmx9vP4svFi4a6U,459
+phylogenie/treesimulator/events.py,sha256=4TJpf9v9Cx-8g886Ilw2fbjSCuSkIoQ_xfOkm_a4ex8,9356
+phylogenie/treesimulator/gillespie.py,sha256=p2bjwc6iyaf4mMI-xregJwOnXkjT2tIhxXbcGncQ5b0,2293
+phylogenie/treesimulator/model.py,sha256=pPl4JYnBIHjWeSgpIypaDOes11tA4bFAGuMRArXbsVo,3577
+phylogenie/typeguards.py,sha256=WBOSJSaOC8VDtrYoA2w_AYEXTpyKdCfmsM29KaKXl3A,1350
+phylogenie/typings.py,sha256=O1X6lGKTjJ2YJz3ApQ-rYb_tEJNUIcHdUIeYlSM4s5o,500
+phylogenie-2.0.0.dist-info/LICENSE.txt,sha256=NUrDqElK-eD3I0WqC004CJsy6cs0JgsAoebDv_42-pw,1071
+phylogenie-2.0.0.dist-info/METADATA,sha256=n2Rf9bfLtz_ctQrAZaXc-nQCe1Z3xqoUbiNrpwLOD94,5602
+phylogenie-2.0.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+phylogenie-2.0.0.dist-info/entry_points.txt,sha256=Rt6_usN0FkBX1ZfiqCirjMN9FKOgFLG8rydcQ8kugeE,51
+phylogenie-2.0.0.dist-info/RECORD,,

phylogenie/backend/remaster/__init__.py

@@ -1,21 +0,0 @@
-from phylogenie.backend.remaster.generate import generate_trees
-from phylogenie.backend.remaster.reactions import (
-    DEFAULT_POPULATION,
-    SAMPLE_POPULATION,
-    PunctualReaction,
-    Reaction,
-    get_canonical_reactions,
-    get_epidemiological_reactions,
-    get_FBD_reactions,
-)
-
-__all__ = [
-    "DEFAULT_POPULATION",
-    "SAMPLE_POPULATION",
-    "PunctualReaction",
-    "Reaction",
-    "get_canonical_reactions",
-    "get_epidemiological_reactions",
-    "get_FBD_reactions",
-    "generate_trees",
-]