diff --git a/epymorph/data/mm/centroids.py b/epymorph/data/mm/centroids.py index 413423a9..af35f53f 100644 --- a/epymorph/data/mm/centroids.py +++ b/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,8 +64,8 @@ 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 = ( @@ -75,7 +73,7 @@ def evaluate(self, tick: Tick) -> NDArray[np.int64]: "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) diff --git a/epymorph/data/mm/flat.py b/epymorph/data/mm/flat.py index 1c386d60..eac1a750 100644 --- a/epymorph/data/mm/flat.py +++ b/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,8 @@ 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) diff --git a/epymorph/data/mm/icecube.py b/epymorph/data/mm/icecube.py index e0e4d600..bf89f0a5 100644 --- a/epymorph/data/mm/icecube.py +++ b/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,8 @@ 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") diff --git a/epymorph/data/mm/no.py b/epymorph/data/mm/no.py index 2b8fc451..ac7464cd 100644 --- a/epymorph/data/mm/no.py +++ b/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,8 @@ 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) diff --git a/epymorph/data/mm/pei.py b/epymorph/data/mm/pei.py index 669f4ccb..e2e92a81 100644 --- a/epymorph/data/mm/pei.py +++ b/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,8 @@ 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 +89,8 @@ 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) diff --git a/epymorph/data/mm/sparsemod.py b/epymorph/data/mm/sparsemod.py index f7aee2ec..5054ae12 100644 --- a/epymorph/data/mm/sparsemod.py +++ b/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,8 @@ 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) diff --git a/epymorph/movement_model.py b/epymorph/movement_model.py index 55b88094..4d0550a2 100644 --- a/epymorph/movement_model.py +++ b/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,7 @@ 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 +189,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 ------- diff --git a/epymorph/simulation.py b/epymorph/simulation.py index 5937949f..94e7cade 100644 --- a/epymorph/simulation.py +++ b/epymorph/simulation.py @@ -552,7 +552,7 @@ class BaseSimulationFunction(ABC, Generic[ResultT]): -------- Refer to [epymorph.simulation.SimulationFunction][] and - [epymorph.simulation.SimulationTickFunction][] for more concrete subclasses. + [epymorph.movement_model.MovementClause][] for more concrete subclasses. """ @classmethod @@ -561,7 +561,7 @@ def __init_subclass__(cls, **kwargs) -> None: _validate_simulation_function(cls) # NOTE: this base class exists so that we don't limit the signature of `evaluate`. - # `SimulationTickFunction` evaluates using the current tick, while + # `MovementClause` evaluates using the current tick and available population, while # `SimulationFunction` does not require any parameters. requirements: Sequence[AttributeDef] | property = () @@ -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) diff --git a/epymorph/simulator/basic/mm_exec.py b/epymorph/simulator/basic/mm_exec.py index 7d25b03b..135d14bc 100644 --- a/epymorph/simulator/basic/mm_exec.py +++ b/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=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], diff --git a/tests/fast/data/mm/centroids_test.py b/tests/fast/data/mm/centroids_test.py index c3161442..a0680a93 100644 --- a/tests/fast/data/mm/centroids_test.py +++ b/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, _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, _POPULATION) def test_phi_negative_error(): clause = _make_clause(phi=-5.0) with pytest.raises(DataAttributeError, match="phi"): - clause.evaluate(_TICK) + clause.evaluate(_TICK, _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, _POPULATION) def test_small_phi_no_underflow(): diff --git a/tests/fast/data_test.py b/tests/fast/data_test.py index bd5867f7..c8c69573 100644 --- a/tests/fast/data_test.py +++ b/tests/fast/data_test.py @@ -1,11 +1,12 @@ # ruff: noqa: PT009,PT027 import math -import unittest +from functools import cached_property import numpy as np -import numpy.testing as npt +import pytest import sympy from numpy.typing import NDArray +from typing_extensions import override from epymorph.attribute import ( AbsoluteName, @@ -14,12 +15,12 @@ ) from epymorph.compartment_model import MultiStrataModelSymbols, edge from epymorph.data.ipm.sirs import SIRS -from epymorph.data.mm.centroids import Centroids from epymorph.data_shape import Shapes -from epymorph.data_type import AttributeArray, CentroidDType +from epymorph.data_type import AttributeArray, CentroidDType, CentroidType, SimDType from epymorph.error import DataAttributeError from epymorph.geography.us_census import StateScope from epymorph.initializer import SingleLocation +from epymorph.movement_model import EveryDay, MovementClause, MovementModel from epymorph.params import ( ParamFunctionNode, ParamFunctionNumpy, @@ -28,326 +29,352 @@ simulation_symbols, ) from epymorph.rume import GPM, RUME, MultiStrataRUME -from epymorph.simulation import ParamValue +from epymorph.simulation import ParamValue, Tick, TickDelta, TickIndex from epymorph.time import TimeFrame - - -class EvaluateParamsTest(unittest.TestCase): - def assert_db( - self, - db: dict[AbsoluteName, AttributeArray], - key: str, - value: AttributeArray, - ) -> None: - matched = db.get(AbsoluteName.parse(key)) - if matched is None: - self.fail(f"Database did not contain the expected key: {key}") +from epymorph.util import pairwise_haversine, row_normalize + + +# Define the movement model to use for the purposes of these tests. +class TestCentroidsClause(MovementClause): + requirements = ( + AttributeDef("population", int, Shapes.N), + AttributeDef("centroid", CentroidType, Shapes.N), + AttributeDef("phi", float, Shapes.Scalar, default_value=40.0), + AttributeDef("commuter_proportion", float, Shapes.Scalar, default_value=0.1), + ) + + predicate = EveryDay() + leaves = TickIndex(step=0) + returns = TickDelta(step=1, days=0) + + @cached_property + def dispersal_kernel(self) -> NDArray[np.float64]: + centroid = self.data("centroid") + phi = self.data("phi") + distance = pairwise_haversine(centroid) + dist_over_phi = np.clip(distance / phi, a_min=None, a_max=700.0) + prob = np.exp(-dist_over_phi) + return row_normalize(prob) + + @override + def evaluate(self, tick: Tick, available: NDArray[SimDType]) -> NDArray[np.int64]: + comm_prop = self.data("commuter_proportion") + n_commuters = np.floor(available * comm_prop).astype(SimDType) + return self.rng.multinomial(n_commuters, self.dispersal_kernel) + + +class TestCentroids(MovementModel): + steps = (1 / 3, 2 / 3) + clauses = (TestCentroidsClause(),) + + +def assert_db( + db: dict[AbsoluteName, AttributeArray], + key: str, + value: AttributeArray, +) -> None: + matched = db.get(AbsoluteName.parse(key)) + if matched is None: + raise AssertionError(f"Database did not contain the expected key: {key}") + else: + msg = f"Database value at key {key} did not match expected." + if value.dtype == np.float64: + np.testing.assert_array_almost_equal(matched, value, err_msg=msg) # type: ignore else: - msg = f"Database value at key {key} did not match expected." - if value.dtype == np.float64: - npt.assert_array_almost_equal(matched, value, err_msg=msg) # type: ignore - else: - npt.assert_array_equal(matched, value, err_msg=msg) - - def _default_params(self) -> dict[str, ParamValue]: - return { - "gpm:aaa::ipm::beta": 0.4, - "gpm:bbb::ipm::beta": 0.3, - "gamma": 1 / 10, # gamma for all strata will be the same - "gpm:aaa::ipm::xi": 0, - "gpm:bbb::ipm::xi": 1 / 90, - "ipm::beta_bbb_aaa": 0.2, - # use the same population values for init and mm modules - # test input as lists and np arrays - "gpm:aaa::*::population": [100, 200], - "gpm:bbb::*::population": np.array([300, 400], dtype=np.int64), - # param names can also include leading stars explicitly - "*::*::centroid": np.array([(1.0, 1.0), (2.0, 2.0)], dtype=CentroidDType), - } - - def _create_rume(self, rume_params: dict[str, ParamValue] | None = None) -> RUME: - meta_requirements = [ - AttributeDef("beta_bbb_aaa", float, Shapes.TxN), + np.testing.assert_array_equal(matched, value, err_msg=msg) + + +def default_params() -> dict[str, ParamValue]: + return { + "gpm:aaa::ipm::beta": 0.4, + "gpm:bbb::ipm::beta": 0.3, + "gamma": 1 / 10, # gamma for all strata will be the same + "gpm:aaa::ipm::xi": 0, + "gpm:bbb::ipm::xi": 1 / 90, + "ipm::beta_bbb_aaa": 0.2, + # use the same population values for init and mm modules + # test input as lists and np arrays + "gpm:aaa::*::population": [100, 200], + "gpm:bbb::*::population": np.array([300, 400], dtype=np.int64), + # param names can also include leading stars explicitly + "*::*::centroid": np.array([(1.0, 1.0), (2.0, 2.0)], dtype=CentroidDType), + } + + +def create_rume(rume_params: dict[str, ParamValue] | None = None) -> RUME: + meta_requirements = [ + AttributeDef("beta_bbb_aaa", float, Shapes.TxN), + ] + + def meta_edges(s: MultiStrataModelSymbols): + [S_aaa, I_aaa, R_aaa] = s.strata_compartments("aaa") # noqa: N806 + [S_bbb, I_bbb, R_bbb] = s.strata_compartments("bbb") # noqa: N806 + [beta_bbb_aaa] = s.all_meta_requirements + N_aaa = sympy.Max(1, S_aaa + I_aaa + R_aaa) # noqa: N806 + return [ + edge(S_bbb, I_bbb, beta_bbb_aaa * S_bbb * I_aaa / N_aaa), ] - def meta_edges(s: MultiStrataModelSymbols): - [S_aaa, I_aaa, R_aaa] = s.strata_compartments("aaa") # noqa: N806 - [S_bbb, I_bbb, R_bbb] = s.strata_compartments("bbb") # noqa: N806 - [beta_bbb_aaa] = s.all_meta_requirements - N_aaa = sympy.Max(1, S_aaa + I_aaa + R_aaa) # noqa: N806 - return [ - edge(S_bbb, I_bbb, beta_bbb_aaa * S_bbb * I_aaa / N_aaa), - ] - - return MultiStrataRUME.build( - strata=[ - GPM( - name="aaa", - ipm=SIRS(), - mm=Centroids(), - init=SingleLocation(location=0, seed_size=100), - params={ - # leave phi unspecified to test default value resolution - }, - ), - GPM( - name="bbb", - ipm=SIRS(), - mm=Centroids(), - init=SingleLocation(location=0, seed_size=100), - params={ - ModuleNamePattern.parse(k): v - for k, v in { - "beta": 99.0, # we'll override this value to test shadowing - "phi": 33.0, # test GPM value resolution - }.items() - }, - ), - ], - meta_requirements=meta_requirements, - meta_edges=meta_edges, - scope=StateScope.in_states(["04", "35"], year=2020), - time_frame=TimeFrame.of("2021-01-01", 180), - params=rume_params or self._default_params(), - ) - - def test_eval_1(self): - rume = self._create_rume() - - db = rume.evaluate_params(rng=np.random.default_rng(1)).to_dict() + return MultiStrataRUME.build( + strata=[ + GPM( + name="aaa", + ipm=SIRS(), + mm=TestCentroids(), + init=SingleLocation(location=0, seed_size=100), + params={ + # leave phi unspecified to test default value resolution + }, + ), + GPM( + name="bbb", + ipm=SIRS(), + mm=TestCentroids(), + init=SingleLocation(location=0, seed_size=100), + params={ + ModuleNamePattern.parse(k): v + for k, v in { + "beta": 99.0, # we'll override this value to test shadowing + "phi": 33.0, # test GPM value resolution + }.items() + }, + ), + ], + meta_requirements=meta_requirements, + meta_edges=meta_edges, + scope=StateScope.in_states(["04", "35"], year=2020), + time_frame=TimeFrame.of("2021-01-01", 180), + params=rume_params or default_params(), + ) + + +@pytest.fixture +def rume() -> RUME: + return create_rume() + + +def test_eval_1(rume): + db = rume.evaluate_params(rng=np.random.default_rng(1)).to_dict() + + # We should have as many entries in our DB as we have attributes in the RUME. + assert len(db) == len(rume.requirements) + + assert_db(db, "gpm:aaa::ipm::beta", np.array(0.4, dtype=np.float64)) + assert_db(db, "gpm:bbb::ipm::beta", np.array(0.3, dtype=np.float64)) + assert_db(db, "gpm:aaa::ipm::gamma", np.array(0.1, dtype=np.float64)) + assert_db(db, "gpm:bbb::ipm::gamma", np.array(0.1, dtype=np.float64)) + assert_db(db, "gpm:aaa::ipm::xi", np.array(0.0, dtype=np.float64)) + assert_db(db, "gpm:bbb::ipm::xi", np.array(1 / 90, dtype=np.float64)) + assert_db(db, "meta::ipm::beta_bbb_aaa", np.array(0.2, dtype=np.float64)) + + assert_db(db, "gpm:aaa::init::population", np.array([100, 200], dtype=np.int64)) + assert_db(db, "gpm:bbb::init::population", np.array([300, 400], dtype=np.int64)) + + assert_db(db, "gpm:aaa::mm::population", np.array([100, 200], dtype=np.int64)) + assert_db(db, "gpm:bbb::mm::population", np.array([300, 400], dtype=np.int64)) + + assert_db(db, "gpm:aaa::mm::phi", np.array(40.0, dtype=np.float64)) + assert_db(db, "gpm:bbb::mm::phi", np.array(33.0, dtype=np.float64)) + + assert_db( + db, + "gpm:aaa::mm::centroid", + np.array([(1.0, 1.0), (2.0, 2.0)], dtype=CentroidDType), + ) + assert_db( + db, + "gpm:bbb::mm::centroid", + np.array([(1.0, 1.0), (2.0, 2.0)], dtype=CentroidDType), + ) + + # When params are provided as the same literal value, + # they should evaluate to the same object. + x1 = db[AbsoluteName.parse("gpm:aaa::ipm::gamma")] + x2 = db[AbsoluteName.parse("gpm:bbb::ipm::gamma")] + assert x1 is x2 + + +def test_eval_2(rume): + # Test with override values. + db = rume.evaluate_params( + override_params={"*::*::beta": 0.5}, + rng=np.random.default_rng(1), + ).to_dict() + + # Beta should be overridden from test case 1, + assert_db(db, "gpm:aaa::ipm::beta", np.array(0.5, dtype=np.float64)) + assert_db(db, "gpm:bbb::ipm::beta", np.array(0.5, dtype=np.float64)) + # the rest should be the same. + assert_db(db, "gpm:aaa::ipm::gamma", np.array(0.1, dtype=np.float64)) + assert_db(db, "gpm:bbb::ipm::gamma", np.array(0.1, dtype=np.float64)) + assert_db(db, "gpm:aaa::ipm::xi", np.array(0.0, dtype=np.float64)) + assert_db(db, "gpm:bbb::ipm::xi", np.array(1 / 90, dtype=np.float64)) + + +def test_eval_3(): + # Test for missing attribute. + # Use the default params but delete one of the attributes. + params = default_params() + del params["gamma"] + + rume = create_rume(params) + + with pytest.raises(DataAttributeError) as exc: + rume.evaluate_params(rng=np.random.default_rng(1)) + + err = str(exc.value).lower() + assert "there are missing values" in err + assert "gpm:aaa::ipm::gamma" in err + assert "gpm:bbb::ipm::gamma" in err + + +def test_eval_sympy_expression(rume): + # Test param as sympy expression + t, T, n = simulation_symbols("day", "duration_days", "node_index") + beta_expr = 0.04 * sympy.sin(8 * sympy.pi * t / T) + 0.34 + (0.02 * n) + + db = rume.evaluate_params( + override_params={"gpm:aaa::ipm::beta": beta_expr}, + rng=np.random.default_rng(1), + ).to_dict() + + expected = np.stack( + [ + 0.04 * np.sin(8 * np.pi * np.arange(180) / 180) + 0.34, + 0.04 * np.sin(8 * np.pi * np.arange(180) / 180) + 0.36, + ], + axis=1, + dtype=np.float64, + ) + assert_db(db, "gpm:aaa::ipm::beta", expected) + + +def test_eval_param_function_1(rume): + # Test param as shaped function + class Beta(ParamFunctionTimeAndNode): + GAMMA = AttributeDef("gamma", float, Shapes.TxN) + + requirements = [GAMMA] + + r_0: float + + def __init__(self, r_0: float): + self.r_0 = r_0 + + def evaluate1(self, day: int, node_index: int) -> float: + T = self.time_frame.days + gamma = self.data(self.GAMMA)[day, node_index] + magnitude = self.r_0 * gamma + return ( + 0.1 * magnitude * math.sin(8 * math.pi * day / T) + + (0.85 * magnitude) + + (0.05 * magnitude * node_index) + ) - # We should have as many entries in our DB as we have attributes in the RUME. - self.assertEqual(len(db), len(rume.requirements)) + db = rume.evaluate_params( + override_params={"gpm:aaa::ipm::beta": Beta(4.0)}, + rng=np.random.default_rng(1), + ).to_dict() - self.assert_db(db, "gpm:aaa::ipm::beta", np.array(0.4, dtype=np.float64)) - self.assert_db(db, "gpm:bbb::ipm::beta", np.array(0.3, dtype=np.float64)) - self.assert_db(db, "gpm:aaa::ipm::gamma", np.array(0.1, dtype=np.float64)) - self.assert_db(db, "gpm:bbb::ipm::gamma", np.array(0.1, dtype=np.float64)) - self.assert_db(db, "gpm:aaa::ipm::xi", np.array(0.0, dtype=np.float64)) - self.assert_db(db, "gpm:bbb::ipm::xi", np.array(1 / 90, dtype=np.float64)) - self.assert_db(db, "meta::ipm::beta_bbb_aaa", np.array(0.2, dtype=np.float64)) + expected = np.stack( + [ + 0.04 * np.sin(8 * np.pi * np.arange(180) / 180) + 0.34, + 0.04 * np.sin(8 * np.pi * np.arange(180) / 180) + 0.36, + ], + axis=1, + dtype=np.float64, + ) + assert_db(db, "gpm:aaa::ipm::beta", expected) - self.assert_db( - db, "gpm:aaa::init::population", np.array([100, 200], dtype=np.int64) - ) - self.assert_db( - db, "gpm:bbb::init::population", np.array([300, 400], dtype=np.int64) - ) - self.assert_db( - db, "gpm:aaa::mm::population", np.array([100, 200], dtype=np.int64) - ) - self.assert_db( - db, "gpm:bbb::mm::population", np.array([300, 400], dtype=np.int64) - ) +def test_eval_param_function_2(rume): + # Test param as shaped function, with difference between strata + class Xi(ParamFunctionNode): + BETA = AttributeDef("beta", float, Shapes.TxN) - self.assert_db(db, "gpm:aaa::mm::phi", np.array(40.0, dtype=np.float64)) - self.assert_db(db, "gpm:bbb::mm::phi", np.array(33.0, dtype=np.float64)) + requirements = [BETA] - self.assert_db( - db, - "gpm:aaa::mm::centroid", - np.array([(1.0, 1.0), (2.0, 2.0)], dtype=CentroidDType), - ) - self.assert_db( - db, - "gpm:bbb::mm::centroid", - np.array([(1.0, 1.0), (2.0, 2.0)], dtype=CentroidDType), - ) + def evaluate1(self, node_index: int) -> float: + beta = self.data(self.BETA)[0, node_index] + return beta / (5 * (node_index + 1)) - # When params are provided as the same literal value, - # they should evaluate to the same object. - x1 = db[AbsoluteName.parse("gpm:aaa::ipm::gamma")] - x2 = db[AbsoluteName.parse("gpm:bbb::ipm::gamma")] - self.assertIs(x1, x2) + db = rume.evaluate_params( + override_params={"ipm::xi": Xi()}, + rng=np.random.default_rng(1), + ).to_dict() - def test_eval_2(self): - # Test with override values. - rume = self._create_rume() + expected_aaa = np.array([(0.4 / 5), (0.4 / 10)], dtype=np.float64) + expected_bbb = np.array([(0.3 / 5), (0.3 / 10)], dtype=np.float64) + assert_db(db, "gpm:aaa::ipm::xi", expected_aaa) + assert_db(db, "gpm:bbb::ipm::xi", expected_bbb) - db = rume.evaluate_params( - override_params={"*::*::beta": 0.5}, - rng=np.random.default_rng(1), - ).to_dict() - - # Beta should be overridden from test case 1, - self.assert_db(db, "gpm:aaa::ipm::beta", np.array(0.5, dtype=np.float64)) - self.assert_db(db, "gpm:bbb::ipm::beta", np.array(0.5, dtype=np.float64)) - # the rest should be the same. - self.assert_db(db, "gpm:aaa::ipm::gamma", np.array(0.1, dtype=np.float64)) - self.assert_db(db, "gpm:bbb::ipm::gamma", np.array(0.1, dtype=np.float64)) - self.assert_db(db, "gpm:aaa::ipm::xi", np.array(0.0, dtype=np.float64)) - self.assert_db(db, "gpm:bbb::ipm::xi", np.array(1 / 90, dtype=np.float64)) - - def test_eval_3(self): - # Test for missing attribute. - # Use the default params but delete one of the attributes. - params = self._default_params() - del params["gamma"] - - rume = self._create_rume(params) - - with self.assertRaises(DataAttributeError) as ctx: - rume.evaluate_params(rng=np.random.default_rng(1)) - - err = str(ctx.exception).lower() - self.assertIn("there are missing values", err) - self.assertIn("gpm:aaa::ipm::gamma", err) - self.assertIn("gpm:bbb::ipm::gamma", err) - - def test_eval_sympy_expression(self): - # Test param as sympy expression - t, T, n = simulation_symbols("day", "duration_days", "node_index") - beta_expr = 0.04 * sympy.sin(8 * sympy.pi * t / T) + 0.34 + (0.02 * n) - - rume = self._create_rume() - - db = rume.evaluate_params( - override_params={"gpm:aaa::ipm::beta": beta_expr}, - rng=np.random.default_rng(1), - ).to_dict() - - expected = np.stack( - [ - 0.04 * np.sin(8 * np.pi * np.arange(180) / 180) + 0.34, - 0.04 * np.sin(8 * np.pi * np.arange(180) / 180) + 0.36, - ], - axis=1, - dtype=np.float64, - ) - self.assert_db(db, "gpm:aaa::ipm::beta", expected) - def test_eval_param_function_1(self): - # Test param as shaped function - class Beta(ParamFunctionTimeAndNode): - GAMMA = AttributeDef("gamma", float, Shapes.TxN) +def test_eval_param_function_chained(rume): + class Gamma(ParamFunctionScalar): + BETA = AttributeDef("beta", float, Shapes.Scalar) - requirements = [GAMMA] + requirements = [BETA] - r_0: float + def evaluate1(self) -> float: + beta = self.data(self.BETA) + return float(beta) / 4.0 - def __init__(self, r_0: float): - self.r_0 = r_0 + class Xi(ParamFunctionNumpy): + ALPHA = AttributeDef("alpha", float, Shapes.Scalar) + GAMMA = AttributeDef("gamma", float, Shapes.Scalar) - def evaluate1(self, day: int, node_index: int) -> float: - T = self.time_frame.days - gamma = self.data(self.GAMMA)[day, node_index] - magnitude = self.r_0 * gamma - return ( - 0.1 * magnitude * math.sin(8 * math.pi * day / T) - + (0.85 * magnitude) - + (0.05 * magnitude * node_index) - ) + requirements = [ALPHA, GAMMA] - rume = self._create_rume() + def evaluate(self) -> NDArray[np.float64]: + # alpha and gamma are both scalars, + # but I'm using ParamFunctionNumpy + # so it's on me to make sure my result is an NDArray + alpha = self.data(self.ALPHA) + gamma = self.data(self.GAMMA) + return np.asarray(gamma / alpha) - db = rume.evaluate_params( - override_params={"gpm:aaa::ipm::beta": Beta(4.0)}, - rng=np.random.default_rng(1), - ).to_dict() - - expected = np.stack( - [ - 0.04 * np.sin(8 * np.pi * np.arange(180) / 180) + 0.34, - 0.04 * np.sin(8 * np.pi * np.arange(180) / 180) + 0.36, - ], - axis=1, - dtype=np.float64, - ) - self.assert_db(db, "gpm:aaa::ipm::beta", expected) - - def test_eval_param_function_2(self): - # Test param as shaped function, with difference between strata - class Xi(ParamFunctionNode): - BETA = AttributeDef("beta", float, Shapes.TxN) - - requirements = [BETA] - - def evaluate1(self, node_index: int) -> float: - beta = self.data(self.BETA)[0, node_index] - return beta / (5 * (node_index + 1)) - - rume = self._create_rume() - - db = rume.evaluate_params( - override_params={"ipm::xi": Xi()}, - rng=np.random.default_rng(1), - ).to_dict() + db = rume.evaluate_params( + override_params={ + "gpm:aaa::ipm::alpha": 9, + "gpm:aaa::ipm::beta": 0.4, + "gpm:aaa::ipm::gamma": Gamma(), + "gpm:aaa::ipm::xi": Xi(), + }, + rng=np.random.default_rng(1), + ).to_dict() - expected_aaa = np.array([(0.4 / 5), (0.4 / 10)], dtype=np.float64) - expected_bbb = np.array([(0.3 / 5), (0.3 / 10)], dtype=np.float64) - self.assert_db(db, "gpm:aaa::ipm::xi", expected_aaa) - self.assert_db(db, "gpm:bbb::ipm::xi", expected_bbb) + assert_db(db, "gpm:aaa::ipm::alpha", np.array(9)) + assert_db(db, "gpm:aaa::ipm::beta", np.array(0.4)) + assert_db(db, "gpm:aaa::ipm::gamma", np.array(0.1)) + assert_db(db, "gpm:aaa::ipm::xi", np.array(1 / 90)) - def test_eval_param_function_chained(self): - class Gamma(ParamFunctionScalar): - BETA = AttributeDef("beta", float, Shapes.Scalar) - requirements = [BETA] +def test_eval_param_function_circular(rume): + class Gamma(ParamFunctionNumpy): + XI = AttributeDef("xi", float, Shapes.Scalar) - def evaluate1(self) -> float: - beta = self.data(self.BETA) - return float(beta) / 4.0 + requirements = [XI] - class Xi(ParamFunctionNumpy): - ALPHA = AttributeDef("alpha", float, Shapes.Scalar) - GAMMA = AttributeDef("gamma", float, Shapes.Scalar) + def evaluate(self) -> NDArray[np.float64]: + return np.array(0) - requirements = [ALPHA, GAMMA] + class Xi(ParamFunctionNumpy): + GAMMA = AttributeDef("gamma", float, Shapes.Scalar) - def evaluate(self) -> NDArray[np.float64]: - # alpha and gamma are both scalars, - # but I'm using ParamFunctionNumpy - # so it's on me to make sure my result is an NDArray - alpha = self.data(self.ALPHA) - gamma = self.data(self.GAMMA) - return np.asarray(gamma / alpha) + requirements = [GAMMA] - rume = self._create_rume() + def evaluate(self) -> NDArray[np.float64]: + return np.array(0) - db = rume.evaluate_params( + with pytest.raises(DataAttributeError) as exc: + rume.evaluate_params( override_params={ - "gpm:aaa::ipm::alpha": 9, - "gpm:aaa::ipm::beta": 0.4, "gpm:aaa::ipm::gamma": Gamma(), "gpm:aaa::ipm::xi": Xi(), }, rng=np.random.default_rng(1), - ).to_dict() - - self.assert_db(db, "gpm:aaa::ipm::alpha", np.array(9)) - self.assert_db(db, "gpm:aaa::ipm::beta", np.array(0.4)) - self.assert_db(db, "gpm:aaa::ipm::gamma", np.array(0.1)) - self.assert_db(db, "gpm:aaa::ipm::xi", np.array(1 / 90)) - - def test_eval_param_function_circular(self): - class Gamma(ParamFunctionNumpy): - XI = AttributeDef("xi", float, Shapes.Scalar) - - requirements = [XI] - - def evaluate(self) -> NDArray[np.float64]: - return np.array(0) - - class Xi(ParamFunctionNumpy): - GAMMA = AttributeDef("gamma", float, Shapes.Scalar) - - requirements = [GAMMA] - - def evaluate(self) -> NDArray[np.float64]: - return np.array(0) - - rume = self._create_rume() - - with self.assertRaises(DataAttributeError) as ctx: - rume.evaluate_params( - override_params={ - "gpm:aaa::ipm::gamma": Gamma(), - "gpm:aaa::ipm::xi": Xi(), - }, - rng=np.random.default_rng(1), - ) + ) - err = str(ctx.exception).lower() - self.assertIn("circular dependency", err) - self.assertIn("gpm:aaa::ipm::gamma", err) + err = str(exc.value).lower() + assert "circular dependency" in err + assert "gpm:aaa::ipm::gamma" in err diff --git a/tests/fast/movement_model_test.py b/tests/fast/movement_model_test.py index 4f45e0df..2e49f09e 100644 --- a/tests/fast/movement_model_test.py +++ b/tests/fast/movement_model_test.py @@ -41,7 +41,9 @@ class MyClause(MovementClause): returns = TickDelta(days=0, step=1) predicate = EveryDay() - def evaluate(self, tick: Tick) -> NDArray[SimDType]: + def evaluate( + self, tick: Tick, available: NDArray[SimDType] + ) -> NDArray[SimDType]: return np.array([0]) clause = MyClause() @@ -57,7 +59,9 @@ class MyClause(MovementClause): returns = TickDelta(days=0, step=1) predicate = EveryDay() - def evaluate(self, tick: Tick) -> NDArray[SimDType]: + def evaluate( + self, tick: Tick, available: NDArray[SimDType] + ) -> NDArray[SimDType]: return np.array([0]) @@ -69,7 +73,9 @@ class MyClause(MovementClause): # returns = TickDelta(days=0, step=1) # noqa: ERA001 predicate = EveryDay() - def evaluate(self, tick: Tick) -> NDArray[SimDType]: + def evaluate( + self, tick: Tick, available: NDArray[SimDType] + ) -> NDArray[SimDType]: return np.array([0]) @@ -81,7 +87,9 @@ class MyClause(MovementClause): returns = TickDelta(days=0, step=1) # predicate = EveryDay() # noqa: ERA001 - def evaluate(self, tick: Tick) -> NDArray[SimDType]: + def evaluate( + self, tick: Tick, available: NDArray[SimDType] + ) -> NDArray[SimDType]: return np.array([0]) @@ -93,7 +101,9 @@ class MyClause(MovementClause): returns = TickDelta(days=0, step=1) predicate = EveryDay() - def evaluate(self, tick: Tick) -> NDArray[SimDType]: + def evaluate( + self, tick: Tick, available: NDArray[SimDType] + ) -> NDArray[SimDType]: return np.array([0]) @@ -105,7 +115,9 @@ class MyClause(MovementClause): returns = TickDelta(days=0, step=-23) predicate = EveryDay() - def evaluate(self, tick: Tick) -> NDArray[SimDType]: + def evaluate( + self, tick: Tick, available: NDArray[SimDType] + ) -> NDArray[SimDType]: return np.array([0]) @@ -115,7 +127,9 @@ class MyClause(MovementClause): returns = NEVER predicate = EveryDay() - def evaluate(self, tick: Tick) -> NDArray[SimDType]: + def evaluate( + self, tick: Tick, available: NDArray[SimDType] + ) -> NDArray[SimDType]: return np.array([0]) clause = MyClause() @@ -128,7 +142,7 @@ class MyClause(MovementClause): returns = TickDelta(days=0, step=1) predicate = EveryDay() - def evaluate(self, tick: Tick) -> NDArray[SimDType]: + def evaluate(self, tick: Tick, available: NDArray[SimDType]) -> NDArray[SimDType]: return np.array([0]) @@ -187,7 +201,9 @@ class MyClauseLocal(MovementClause): returns = TickDelta(days=0, step=9) predicate = EveryDay() - def evaluate(self, tick: Tick) -> NDArray[SimDType]: + def evaluate( + self, tick: Tick, available: NDArray[SimDType] + ) -> NDArray[SimDType]: return np.array([0]) with pytest.raises( diff --git a/tests/fast/rume_test.py b/tests/fast/rume_test.py index 26280272..02305d0d 100644 --- a/tests/fast/rume_test.py +++ b/tests/fast/rume_test.py @@ -161,7 +161,9 @@ class Clause1(MovementClause): returns = TickDelta(days=0, step=1) predicate = EveryDay() - def evaluate(self, tick: Tick) -> NDArray[np.int64]: + def evaluate( + self, tick: Tick, available: NDArray[np.int64] + ) -> NDArray[np.int64]: return np.array([]) class Model1(MovementModel): @@ -173,7 +175,9 @@ class Clause2(MovementClause): returns = TickDelta(days=0, step=1) predicate = EveryDay() - def evaluate(self, tick: Tick) -> NDArray[np.int64]: + def evaluate( + self, tick: Tick, available: NDArray[np.int64] + ) -> NDArray[np.int64]: return np.array([]) class Model2(MovementModel): @@ -200,7 +204,9 @@ class Clause1(MovementClause): returns = TickDelta(days=0, step=1) predicate = EveryDay() - def evaluate(self, tick: Tick) -> NDArray[np.int64]: + def evaluate( + self, tick: Tick, available: NDArray[np.int64] + ) -> NDArray[np.int64]: return np.array([]) class Model1(MovementModel): @@ -212,7 +218,9 @@ class Clause2(MovementClause): returns = NEVER predicate = EveryDay() - def evaluate(self, tick: Tick) -> NDArray[np.int64]: + def evaluate( + self, tick: Tick, available: NDArray[np.int64] + ) -> NDArray[np.int64]: return np.array([]) class Model2(MovementModel): @@ -384,12 +392,6 @@ def test_create_multistrata_2(self): AbsoluteName("gpm:aaa", "ipm", "gamma"): AttributeDef( "gamma", float, Shapes.TxN ), - AbsoluteName("gpm:aaa", "mm", "population"): AttributeDef( - "population", - int, - Shapes.N, - comment="The total population at each node.", - ), AbsoluteName("gpm:aaa", "mm", "centroid"): AttributeDef( "centroid", (("longitude", float), ("latitude", float)),