From a50dfdc604520efc5bc02505927beedfcb24e847 Mon Sep 17 00:00:00 2001 From: Ewout ter Hoeven Date: Fri, 13 Dec 2024 09:47:52 +0100 Subject: [PATCH 1/5] feat: Add agent state management system with discrete & continuous states Adds a minimal state management system to Mesa that supports: - Discrete states with explicit value changes - Continuous states that update based on elapsed time - Composite states derived from other states - Integration with Mesa's Agent class via StateAgent This provides a clean API for managing agent states in simulations while handling both discrete events and continuous changes efficiently. --- mesa/experimental/states.py | 140 ++++++++++++++++++++++++++++++++++++ 1 file changed, 140 insertions(+) create mode 100644 mesa/experimental/states.py diff --git a/mesa/experimental/states.py b/mesa/experimental/states.py new file mode 100644 index 00000000000..232f5a05a89 --- /dev/null +++ b/mesa/experimental/states.py @@ -0,0 +1,140 @@ +"""State management system for Mesa agents. + +This module provides a flexible state management system for Mesa agents, supporting +both discrete and continuous state changes. It enables agents to maintain multiple +states that can change either explicitly (discrete) or based on time (continuous), +and includes support for composite states derived from other states. + +Core Classes: + State: Base class defining the state interface + DiscreteState: States with explicit value changes + ContinuousState: States that change over time + CompositeState: States computed from other states + StateAgent: Mesa Agent subclass with state management +""" + +from collections.abc import Callable +from typing import Any + +from mesa import Agent + + +class State: + """Base class for all states.""" + + def __init__(self, name: str, initial_value: Any): + """Create a new state.""" + self.name = name + self._value = initial_value + self._last_update_time = 0 + self.model = None # Set when state is added to agent + + @property + def value(self) -> Any: + """Get current state value.""" + raise NotImplementedError + + def update(self, time: float) -> None: + """Update state to current time.""" + raise NotImplementedError + + +class DiscreteState(State): + """A state with discrete values that change explicitly.""" + + @property + def value(self) -> Any: + """Get the current state value.""" + return self._value + + @value.setter + def value(self, new_value: Any) -> None: + """Set the state value.""" + self._value = new_value + + def update(self, time: float) -> None: + """DiscreteStates only update when value is explicitly changed.""" + + +class ContinuousState(State): + """A state that changes continuously over time.""" + + def __init__( + self, + name: str, + initial_value: float, + rate_function: Callable[[float, float], float], + ): + """Create a new continuous state.""" + super().__init__(name, initial_value) + self.rate_function = rate_function + + @property + def value(self) -> float: + """Calculate and return current value based on elapsed time.""" + current_time = self.model.time + if current_time > self._last_update_time: + self.update(current_time) + return self._value + + def update(self, time: float) -> None: + """Update state value based on elapsed time.""" + elapsed = time - self._last_update_time + self._value = self._value + self.rate_function(self._value, elapsed) + self._last_update_time = time + + +class CompositeState(State): + """A state derived from other states.""" + + def __init__( + self, + name: str, + dependent_states: list[State], + computation_function: Callable[..., Any], + ): + """Create a new composite state.""" + self.dependent_states = dependent_states + self.computation_function = computation_function + super().__init__(name, None) # Value computed on first access + + @property + def value(self) -> Any: + """Compute value based on dependent states.""" + state_values = [state.value for state in self.dependent_states] + return self.computation_function(*state_values) + + def update(self, time: float) -> None: + """Update all dependent states.""" + for state in self.dependent_states: + state.update(time) + + +class StateAgent(Agent): + """An agent with integrated state management.""" + + def __init__(self, model) -> None: + """Create a new agent with state management.""" + super().__init__(model) + self.states = {} # name -> State mapping + + def add_state(self, state: State) -> None: + """Add a new state to the agent.""" + state.model = self.model + self.states[state.name] = state + + def get_state(self, name: str) -> Any: + """Get the current value of a state.""" + return self.states[name].value + + def set_state(self, name: str, value: Any) -> None: + """Set value for a discrete state.""" + if isinstance(self.states[name], DiscreteState): + self.states[name].value = value + else: + raise ValueError("Cannot directly set value of non-discrete state") + + def update_states(self) -> None: + """Update all states to current time.""" + for state in self.states.values(): + state.update(self.model.time) From 19245c2a42cf29368c0a006e88f5f7e70637f03b Mon Sep 17 00:00:00 2001 From: Ewout ter Hoeven Date: Fri, 13 Dec 2024 11:39:52 +0100 Subject: [PATCH 2/5] StateAgent: allows direct attribute-based state access --- mesa/experimental/states.py | 56 ++++++++++++++++++++++++------------- 1 file changed, 36 insertions(+), 20 deletions(-) diff --git a/mesa/experimental/states.py b/mesa/experimental/states.py index 232f5a05a89..8ead177d5c0 100644 --- a/mesa/experimental/states.py +++ b/mesa/experimental/states.py @@ -80,7 +80,7 @@ def value(self) -> float: def update(self, time: float) -> None: """Update state value based on elapsed time.""" elapsed = time - self._last_update_time - self._value = self._value + self.rate_function(self._value, elapsed) + self._value += self.rate_function(self._value, elapsed) self._last_update_time = time @@ -111,30 +111,46 @@ def update(self, time: float) -> None: class StateAgent(Agent): - """An agent with integrated state management.""" + """An agent with integrated state management that allows direct attribute-based state access.""" - def __init__(self, model) -> None: + def __init__(self, model): """Create a new agent with state management.""" super().__init__(model) - self.states = {} # name -> State mapping - - def add_state(self, state: State) -> None: - """Add a new state to the agent.""" - state.model = self.model - self.states[state.name] = state - - def get_state(self, name: str) -> Any: - """Get the current value of a state.""" - return self.states[name].value - - def set_state(self, name: str, value: Any) -> None: - """Set value for a discrete state.""" - if isinstance(self.states[name], DiscreteState): - self.states[name].value = value - else: - raise ValueError("Cannot directly set value of non-discrete state") + object.__setattr__(self, "states", {}) # Use object.__setattr__ to avoid recursion def update_states(self) -> None: """Update all states to current time.""" for state in self.states.values(): state.update(self.model.time) + + def __getattribute__(self, name: str) -> Any: + """Get an attribute, routing state access to its value.""" + states = object.__getattribute__(self, "states") + if name in states: + # If it's a known state, return its current value + return states[name].value + else: + # Otherwise, return the attribute normally + return object.__getattribute__(self, name) + + def __setattr__(self, name: str, value: Any) -> None: + """Set an attribute, allowing direct state assignment or updates.""" + states = object.__getattribute__(self, "states") + # If setting a State object, add or update the states dictionary + if isinstance(value, State): + # The state's name should match the attribute name + if value.name != name: + raise ValueError(f"State name '{value.name}' does not match attribute name '{name}'") + states[name] = value + value.model = self.model + else: + # If we're setting a non-state value and a corresponding state exists + if name in states: + # The state must be discrete to allow direct setting + if isinstance(states[name], DiscreteState): + states[name].value = value + else: + raise ValueError("Cannot directly set value of non-discrete state") + else: + # Otherwise, it's just a normal attribute + object.__setattr__(self, name, value) From 3cd6859e109b65ce0391ea3c227243418782c2a8 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 13 Dec 2024 10:40:33 +0000 Subject: [PATCH 3/5] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- mesa/experimental/states.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/mesa/experimental/states.py b/mesa/experimental/states.py index 8ead177d5c0..59fe9073a3b 100644 --- a/mesa/experimental/states.py +++ b/mesa/experimental/states.py @@ -116,7 +116,9 @@ class StateAgent(Agent): def __init__(self, model): """Create a new agent with state management.""" super().__init__(model) - object.__setattr__(self, "states", {}) # Use object.__setattr__ to avoid recursion + object.__setattr__( + self, "states", {} + ) # Use object.__setattr__ to avoid recursion def update_states(self) -> None: """Update all states to current time.""" @@ -140,7 +142,9 @@ def __setattr__(self, name: str, value: Any) -> None: if isinstance(value, State): # The state's name should match the attribute name if value.name != name: - raise ValueError(f"State name '{value.name}' does not match attribute name '{name}'") + raise ValueError( + f"State name '{value.name}' does not match attribute name '{name}'" + ) states[name] = value value.model = self.model else: From badc3a3d4bf9e54eb65f918e6367d927ee37d1ed Mon Sep 17 00:00:00 2001 From: Ewout ter Hoeven Date: Fri, 13 Dec 2024 11:51:31 +0100 Subject: [PATCH 4/5] Make sure the state value is up to date before returning it --- mesa/experimental/states.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/mesa/experimental/states.py b/mesa/experimental/states.py index 59fe9073a3b..a5fcca4c619 100644 --- a/mesa/experimental/states.py +++ b/mesa/experimental/states.py @@ -129,8 +129,12 @@ def __getattribute__(self, name: str) -> Any: """Get an attribute, routing state access to its value.""" states = object.__getattribute__(self, "states") if name in states: - # If it's a known state, return its current value - return states[name].value + # If it's a known state, ensure it is updated before returning its value + state = states[name] + current_time = object.__getattribute__(self.model, "time") + if current_time > state._last_update_time: + state.update(current_time) + return state.value else: # Otherwise, return the attribute normally return object.__getattribute__(self, name) From 102dd7087f72b50f2e01287ac60604adb8d41589 Mon Sep 17 00:00:00 2001 From: Ewout ter Hoeven Date: Fri, 13 Dec 2024 11:52:49 +0100 Subject: [PATCH 5/5] states: Don't require a name simplifies the API --- mesa/experimental/states.py | 14 +++----------- 1 file changed, 3 insertions(+), 11 deletions(-) diff --git a/mesa/experimental/states.py b/mesa/experimental/states.py index a5fcca4c619..476c053942c 100644 --- a/mesa/experimental/states.py +++ b/mesa/experimental/states.py @@ -22,9 +22,8 @@ class State: """Base class for all states.""" - def __init__(self, name: str, initial_value: Any): + def __init__(self, initial_value: Any): """Create a new state.""" - self.name = name self._value = initial_value self._last_update_time = 0 self.model = None # Set when state is added to agent @@ -61,12 +60,11 @@ class ContinuousState(State): def __init__( self, - name: str, initial_value: float, rate_function: Callable[[float, float], float], ): """Create a new continuous state.""" - super().__init__(name, initial_value) + super().__init__(initial_value) self.rate_function = rate_function @property @@ -89,14 +87,13 @@ class CompositeState(State): def __init__( self, - name: str, dependent_states: list[State], computation_function: Callable[..., Any], ): """Create a new composite state.""" self.dependent_states = dependent_states self.computation_function = computation_function - super().__init__(name, None) # Value computed on first access + super().__init__(None) # Value computed on first access @property def value(self) -> Any: @@ -144,11 +141,6 @@ def __setattr__(self, name: str, value: Any) -> None: states = object.__getattribute__(self, "states") # If setting a State object, add or update the states dictionary if isinstance(value, State): - # The state's name should match the attribute name - if value.name != name: - raise ValueError( - f"State name '{value.name}' does not match attribute name '{name}'" - ) states[name] = value value.model = self.model else: