feature!: MM clauses can access the current simulation state

epymorph/data/mm/centroids.py

@@ -2,6 +2,7 @@
 
 import numpy as np
 from numpy.typing import NDArray
+from typing_extensions import override
 
 from epymorph.attribute import AttributeDef
 from epymorph.data_shape import Shapes
@@ -16,9 +17,6 @@ class CentroidsClause(MovementClause):
     """The clause of the centroids model."""
 
     requirements = (
-        AttributeDef(
-            "population", int, Shapes.N, comment="The total population at each node."
-        ),
         AttributeDef(
             "centroid",
             CentroidType,
@@ -66,16 +64,21 @@ def dispersal_kernel(self) -> NDArray[np.float64]:
         prob = np.exp(-dist_over_phi)
         return row_normalize(prob)
 
-    def evaluate(self, tick: Tick) -> NDArray[np.int64]:
-        pop = self.data("population")
+    @override
+    def evaluate(
+        self,
+        tick: Tick,
+        *,
+        available: NDArray[SimDType],
+    ) -> NDArray[np.int64]:
         comm_prop = self.data("commuter_proportion")
         if comm_prop < 0:
             err = (
                 "Centroids movement model parameter 'commuter_proportion' must be "
                 "greater than or equal to zero."
             )
             raise DataAttributeError(err)
-        n_commuters = np.floor(pop * comm_prop).astype(SimDType)
+        n_commuters = np.floor(available * comm_prop).astype(SimDType)
         return self.rng.multinomial(n_commuters, self.dispersal_kernel)
 
 

epymorph/data/mm/flat.py

@@ -2,6 +2,7 @@
 
 import numpy as np
 from numpy.typing import NDArray
+from typing_extensions import override
 
 from epymorph.attribute import AttributeDef
 from epymorph.data_shape import Shapes
@@ -43,7 +44,13 @@ def dispersal_kernel(self) -> NDArray[np.float64]:
         np.fill_diagonal(ones, 0)
         return row_normalize(ones)
 
-    def evaluate(self, tick: Tick) -> NDArray[SimDType]:
+    @override
+    def evaluate(
+        self,
+        tick: Tick,
+        *,
+        available: NDArray[SimDType],
+    ) -> NDArray[SimDType]:
         pop = self.data("population")
         comm_prop = self.data("commuter_proportion")
         n_commuters = np.floor(pop * comm_prop).astype(SimDType)

epymorph/data/mm/icecube.py

@@ -1,5 +1,6 @@
 import numpy as np
 from numpy.typing import NDArray
+from typing_extensions import override
 
 from epymorph.attribute import AttributeDef
 from epymorph.data_shape import Shapes
@@ -28,7 +29,13 @@ class IcecubeClause(MovementClause):
     leaves = TickIndex(step=0)
     returns = TickDelta(step=1, days=0)
 
-    def evaluate(self, tick: Tick) -> NDArray[np.int64]:
+    @override
+    def evaluate(
+        self,
+        tick: Tick,
+        *,
+        available: NDArray[SimDType],
+    ) -> NDArray[np.int64]:
         N = self.scope.nodes
         pop = self.data("population")
         comm_prop = self.data("commuter_proportion")

epymorph/data/mm/no.py

@@ -1,5 +1,6 @@
 import numpy as np
 from numpy.typing import NDArray
+from typing_extensions import override
 
 from epymorph.data_type import SimDType
 from epymorph.movement_model import EveryDay, MovementClause, MovementModel
@@ -14,7 +15,13 @@ class NoClause(MovementClause):
     leaves = TickIndex(step=0)
     returns = TickDelta(step=0, days=0)
 
-    def evaluate(self, tick: Tick) -> NDArray[np.int64]:
+    @override
+    def evaluate(
+        self,
+        tick: Tick,
+        *,
+        available: NDArray[SimDType],
+    ) -> NDArray[np.int64]:
         N = self.scope.nodes
         return np.zeros((N, N), dtype=SimDType)
 

epymorph/data/mm/pei.py

@@ -2,6 +2,7 @@
 
 import numpy as np
 from numpy.typing import NDArray
+from typing_extensions import override
 
 from epymorph.attribute import AttributeDef
 from epymorph.data_shape import Shapes
@@ -52,7 +53,13 @@ def commuting_probability(self) -> NDArray[np.float64]:
         commuters = self.data("commuters")
         return row_normalize(commuters)
 
-    def evaluate(self, tick: Tick) -> NDArray[np.int64]:
+    @override
+    def evaluate(
+        self,
+        tick: Tick,
+        *,
+        available: NDArray[SimDType],
+    ) -> NDArray[np.int64]:
         move_control = self.data("move_control")
         actual = self.rng.binomial(self.commuters_by_node, move_control)
         return self.rng.multinomial(actual, self.commuting_probability)
@@ -87,7 +94,13 @@ def commuters_average(self) -> NDArray[SimDType]:
         commuters = self.data("commuters")
         return (commuters + commuters.T) // 2
 
-    def evaluate(self, tick: Tick) -> NDArray[SimDType]:
+    @override
+    def evaluate(
+        self,
+        tick: Tick,
+        *,
+        available: NDArray[SimDType],
+    ) -> NDArray[SimDType]:
         theta = self.data("theta")
         return self.rng.poisson(theta * self.commuters_average)
 

epymorph/data/mm/sparsemod.py

@@ -2,6 +2,7 @@
 
 import numpy as np
 from numpy.typing import NDArray
+from typing_extensions import override
 
 from epymorph.attribute import AttributeDef
 from epymorph.data_shape import Shapes
@@ -58,7 +59,13 @@ def dispersal_kernel(self) -> NDArray[np.float64]:
         distance = pairwise_haversine(centroid)
         return row_normalize(1 / np.exp(distance / phi))
 
-    def evaluate(self, tick: Tick) -> NDArray[np.int64]:
+    @override
+    def evaluate(
+        self,
+        tick: Tick,
+        *,
+        available: NDArray[SimDType],
+    ) -> NDArray[np.int64]:
         return self.rng.multinomial(self.commuters_by_node, self.dispersal_kernel)
 
 

epymorph/movement_model.py

@@ -23,7 +23,7 @@
 from epymorph.data_type import SimDType
 from epymorph.simulation import (
     NEVER,
-    SimulationTickFunction,
+    BaseSimulationFunction,
     Tick,
     TickDelta,
     TickIndex,
@@ -129,7 +129,7 @@ def evaluate(self, tick: Tick) -> bool:
 ##################
 
 
-class MovementClause(SimulationTickFunction[NDArray[SimDType]], ABC):
+class MovementClause(BaseSimulationFunction[NDArray[SimDType]], ABC):
     """
     A movement clause is basically a function which calculates _how many_ individuals
     should move between all of the geo nodes.
@@ -179,7 +179,12 @@ def clause_name(self) -> str:
         return self.__class__.__name__
 
     @abstractmethod
-    def evaluate(self, tick: Tick) -> NDArray[SimDType]:
+    def evaluate(
+        self,
+        tick: Tick,
+        *,
+        available: NDArray[SimDType],
+    ) -> NDArray[SimDType]:
         """
         Implement this method to provide logic for the clause.
         Use self methods and properties to access the simulation context or defer
@@ -189,6 +194,12 @@ def evaluate(self, tick: Tick) -> NDArray[SimDType]:
         ----------
         tick :
             The simulation tick being evaluated.
+        available :
+            The number of individuals currently at each location which are available to
+            move, as an N-shaped array. Note: it's not necessary that movement clauses
+            take this into account (e.g., to return requested movement numbers which are
+            less than the available number of individuals.) But it is provided for
+            clauses that wish to take this into account.
 
         Returns
         -------

epymorph/simulation.py

@@ -903,70 +903,3 @@ def defer(
             The result value.
         """
         return self.defer_context(other, scope, time_frame).evaluate()
-
-
-class SimulationTickFunction(BaseSimulationFunction[ResultT]):
-    """
-    A function which runs in the context of a RUME to produce a value
-    (as a numpy array) which is expected to vary over the run of a simulation.
-
-    In typical usage you will not implement a `SimulationTickFunction` directly,
-    but rather one of its more-specific child classes.
-
-    `SimulationTickFunction` is generic in the type of result it produces (`ResultT`).
-
-    See Also
-    --------
-    The only notable child class is [epymorph.movement_model.MovementClause][].
-    """
-
-    @abstractmethod
-    def evaluate(self, tick: Tick) -> ResultT:
-        """
-        Implement this method to provide logic for the function.
-        Use self methods and properties to access the simulation context or defer
-        processing to another function.
-
-        Parameters
-        ----------
-        tick :
-            The simulation tick being evaluated.
-
-        Returns
-        -------
-        :
-            The result value.
-        """
-
-    @final
-    def defer(
-        self,
-        other: "SimulationTickFunction[DeferResultT]",
-        tick: Tick,
-        scope: GeoScope | None = None,
-        time_frame: TimeFrame | None = None,
-    ) -> DeferResultT:
-        """
-        Defer processing to another instance of a `SimulationTickFunction`, returning
-        the result of evaluation.
-
-        This function is generic in the type of result returned by the function
-        to which we are deferring (`DeferResultT`).
-
-        Parameters
-        ----------
-        other :
-            The other function to defer to.
-        tick :
-            The simulation tick being evaluated.
-        scope :
-            Override the geo scope for evaluation; if None, use the same scope.
-        time_frame :
-            Override the time frame for evaluation; if None, use the same time frame.
-
-        Returns
-        -------
-        :
-            The result value.
-        """
-        return self.defer_context(other, scope, time_frame).evaluate(tick)

epymorph/simulator/basic/mm_exec.py

@@ -186,9 +186,13 @@ def apply(self, tick: Tick) -> None:
         for strata, clause in self._clauses:
             if not clause.is_active(tick):
                 continue
+            available_movers = self._world.get_local_array()
 
             try:
-                requested_movers = clause.evaluate(tick)
+                requested_movers = clause.evaluate(
+                    tick,
+                    available=available_movers.sum(axis=1, dtype=SimDType),
+                )
                 np.fill_diagonal(requested_movers, 0)
             except Exception as e:
                 # NOTE: catching exceptions here is necessary to get nice error messages
@@ -200,7 +204,6 @@ def apply(self, tick: Tick) -> None:
                 )
                 raise MMSimError(msg) from e
 
-            available_movers = self._world.get_local_array()
             clause_event = calculate_travelers(
                 clause.clause_name,
                 self._rume.compartment_mobility[strata],

tests/fast/data/mm/centroids_test.py

@@ -42,7 +42,7 @@ def _make_clause(phi: float, commuter_proportion: float = 0.1) -> CentroidsClaus
 def test_evaluate():
     commuter_proportion = 0.1
     clause = _make_clause(phi=40.0, commuter_proportion=commuter_proportion)
-    result = clause.evaluate(_TICK)
+    result = clause.evaluate(_TICK, available=_POPULATION)
 
     assert result.shape == (3, 3)
     assert np.all(result >= 0)
@@ -56,19 +56,19 @@ def test_evaluate():
 def test_phi_zero_error():
     clause = _make_clause(phi=0.0)
     with pytest.raises(DataAttributeError, match="phi"):
-        clause.evaluate(_TICK)
+        clause.evaluate(_TICK, available=_POPULATION)
 
 
 def test_phi_negative_error():
     clause = _make_clause(phi=-5.0)
     with pytest.raises(DataAttributeError, match="phi"):
-        clause.evaluate(_TICK)
+        clause.evaluate(_TICK, available=_POPULATION)
 
 
 def test_commuter_proportion_negative_error():
     clause = _make_clause(phi=40.0, commuter_proportion=-0.1)
     with pytest.raises(DataAttributeError, match="commuter_proportion"):
-        clause.evaluate(_TICK)
+        clause.evaluate(_TICK, available=_POPULATION)
 
 
 def test_small_phi_no_underflow():