diff --git a/bdms/__init__.py b/bdms/__init__.py index d5cd2a2..ed39a45 100644 --- a/bdms/__init__.py +++ b/bdms/__init__.py @@ -1,4 +1,4 @@ -__version__ = "1.0.0" +__version__ = "1.0.1" from .tree import Tree as _Tree, TreeError as _TreeError diff --git a/bdms/poisson.py b/bdms/poisson.py index d9bcc39..42970a5 100644 --- a/bdms/poisson.py +++ b/bdms/poisson.py @@ -181,7 +181,9 @@ def Λ(self, x: Hashable, t: float, Δt: float) -> NDArray[np.floating]: def Λ_inv(self, x: Hashable, t: float, τ: float) -> NDArray[np.floating]: rate = self.λ_homogeneous(x) if rate == 0: - return np.full_like(x, np.inf, dtype=float) + # NOTE: not np.full_like(x, np.inf), which has the wrong shape if the state + # is a sequence, such as a tuple + return np.asarray(np.inf) return τ / rate @@ -204,7 +206,10 @@ def λ_homogeneous( ) -> NDArray[np.floating]: if isinstance(x, Hashable): # type: ignore return float(self.value) # type: ignore - return self.value * np.ones_like(x, dtype=float) + # NOTE: one rate per state, so not np.ones_like(x) unless x is an array, since a + # sequence of states may be a sequence of tuples + shape = x.shape if isinstance(x, np.ndarray) else len(x) + return np.full(shape, float(self.value)) class DiscreteProcess(HomogeneousProcess): diff --git a/bdms/tree.py b/bdms/tree.py index 9471746..fa2f1f0 100644 --- a/bdms/tree.py +++ b/bdms/tree.py @@ -379,21 +379,25 @@ def mutation_event_fn(node: Tree, rng: np.random.Generator) -> bool: raise TreeError(f"{capacity=} exceeded at time={current_time}") else: raise ValueError(f"{capacity_method=} not recognized") + # NOTE: only waiting times are compared, so states need not be orderable waiting_time, event, state = min( ( - processes[event].waiting_time_rv( + ( + processes[event].waiting_time_rv( + state, + current_time, + rate_multiplier=rate_multipliers[event] + * len(state_names_active[state]), + seed=rng, + ), + event, state, - current_time, - rate_multiplier=rate_multipliers[event] - * len(state_names_active[state]), - seed=rng, - ), - event, - state, - ) - for event in processes - for state in state_names_active - if len(state_names_active[state]) > 0 + ) + for event in processes + for state in state_names_active + if len(state_names_active[state]) > 0 + ), + key=lambda item: item[0], ) Δt = float(np.minimum(waiting_time, end_time - current_time)) current_time += Δt diff --git a/tests/test_bdms.py b/tests/test_bdms.py index 3850513..03f8ddb 100644 --- a/tests/test_bdms.py +++ b/tests/test_bdms.py @@ -1,3 +1,4 @@ +import dataclasses import numpy as np import ete4 import bdms @@ -131,5 +132,45 @@ def test_copy(self): self.assertFalse(self.tree._pruned) +class TestStates(unittest.TestCase): + def test_tuple_states(self): + # the default death and mutation processes have rate zero + tree = bdms.Tree(state=("a", 0)) + tree.evolve( + 2.0, birth_process=bdms.poisson.DiscreteProcess({("a", 0): 1.0}), seed=0 + ) + scalar_tree = bdms.Tree(state="a0") + scalar_tree.evolve( + 2.0, birth_process=bdms.poisson.DiscreteProcess({"a0": 1.0}), seed=0 + ) + self.assertEqual( + [node.t for node in tree.traverse()], + [node.t for node in scalar_tree.traverse()], + ) + states = [("a", 0), ("a", 1), ("b", 0)] + rates = bdms.poisson.ConstantProcess(1.0).λ_homogeneous(states) + self.assertEqual(rates.shape, (3,)) + + def test_unordered_states(self): + # states that can't be compared, two of which can't give birth + @dataclasses.dataclass(frozen=True) + class State: + kind: str + + a, b, c = State("a"), State("b"), State("c") + tree = bdms.Tree(state=a) + tree.evolve( + 6.0, + birth_process=bdms.poisson.DiscreteProcess({a: 0.0, b: 0.0, c: 2.0}), + mutation_process=bdms.poisson.ConstantProcess(1.0), + mutator=bdms.mutators.DiscreteMutator( + (a, b, c), np.array([[0, 1, 0], [0, 0, 1], [1, 0, 0]]) + ), + init_population=3, + seed=0, + ) + self.assertTrue(any(node.event == "birth" for node in tree.traverse())) + + if __name__ == "__main__": unittest.main()