Support tuple and unordered node states

bdms/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "1.0.0"
+__version__ = "1.0.1"
 
 from .tree import Tree as _Tree, TreeError as _TreeError
 

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):

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

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()