Skip to content

Commit

Permalink
feat: Add agent state management system with discrete & continuous st…
Browse files Browse the repository at this point in the history
…ates

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.
  • Loading branch information
EwoutH committed Dec 13, 2024
1 parent e4e92f4 commit a50dfdc
Showing 1 changed file with 140 additions and 0 deletions.
140 changes: 140 additions & 0 deletions mesa/experimental/states.py
Original file line number Diff line number Diff line change
@@ -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

Check warning on line 17 in mesa/experimental/states.py

View check run for this annotation

Codecov / codecov/patch

mesa/experimental/states.py#L16-L17

Added lines #L16 - L17 were not covered by tests

from mesa import Agent

Check warning on line 19 in mesa/experimental/states.py

View check run for this annotation

Codecov / codecov/patch

mesa/experimental/states.py#L19

Added line #L19 was not covered by tests


class State:

Check warning on line 22 in mesa/experimental/states.py

View check run for this annotation

Codecov / codecov/patch

mesa/experimental/states.py#L22

Added line #L22 was not covered by tests
"""Base class for all states."""

def __init__(self, name: str, initial_value: Any):

Check warning on line 25 in mesa/experimental/states.py

View check run for this annotation

Codecov / codecov/patch

mesa/experimental/states.py#L25

Added line #L25 was not covered by tests
"""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

Check warning on line 30 in mesa/experimental/states.py

View check run for this annotation

Codecov / codecov/patch

mesa/experimental/states.py#L27-L30

Added lines #L27 - L30 were not covered by tests

@property
def value(self) -> Any:

Check warning on line 33 in mesa/experimental/states.py

View check run for this annotation

Codecov / codecov/patch

mesa/experimental/states.py#L32-L33

Added lines #L32 - L33 were not covered by tests
"""Get current state value."""
raise NotImplementedError

Check warning on line 35 in mesa/experimental/states.py

View check run for this annotation

Codecov / codecov/patch

mesa/experimental/states.py#L35

Added line #L35 was not covered by tests

def update(self, time: float) -> None:

Check warning on line 37 in mesa/experimental/states.py

View check run for this annotation

Codecov / codecov/patch

mesa/experimental/states.py#L37

Added line #L37 was not covered by tests
"""Update state to current time."""
raise NotImplementedError

Check warning on line 39 in mesa/experimental/states.py

View check run for this annotation

Codecov / codecov/patch

mesa/experimental/states.py#L39

Added line #L39 was not covered by tests


class DiscreteState(State):

Check warning on line 42 in mesa/experimental/states.py

View check run for this annotation

Codecov / codecov/patch

mesa/experimental/states.py#L42

Added line #L42 was not covered by tests
"""A state with discrete values that change explicitly."""

@property
def value(self) -> Any:

Check warning on line 46 in mesa/experimental/states.py

View check run for this annotation

Codecov / codecov/patch

mesa/experimental/states.py#L45-L46

Added lines #L45 - L46 were not covered by tests
"""Get the current state value."""
return self._value

Check warning on line 48 in mesa/experimental/states.py

View check run for this annotation

Codecov / codecov/patch

mesa/experimental/states.py#L48

Added line #L48 was not covered by tests

@value.setter
def value(self, new_value: Any) -> None:

Check warning on line 51 in mesa/experimental/states.py

View check run for this annotation

Codecov / codecov/patch

mesa/experimental/states.py#L50-L51

Added lines #L50 - L51 were not covered by tests
"""Set the state value."""
self._value = new_value

Check warning on line 53 in mesa/experimental/states.py

View check run for this annotation

Codecov / codecov/patch

mesa/experimental/states.py#L53

Added line #L53 was not covered by tests

def update(self, time: float) -> None:

Check warning on line 55 in mesa/experimental/states.py

View check run for this annotation

Codecov / codecov/patch

mesa/experimental/states.py#L55

Added line #L55 was not covered by tests
"""DiscreteStates only update when value is explicitly changed."""


class ContinuousState(State):

Check warning on line 59 in mesa/experimental/states.py

View check run for this annotation

Codecov / codecov/patch

mesa/experimental/states.py#L59

Added line #L59 was not covered by tests
"""A state that changes continuously over time."""

def __init__(

Check warning on line 62 in mesa/experimental/states.py

View check run for this annotation

Codecov / codecov/patch

mesa/experimental/states.py#L62

Added line #L62 was not covered by tests
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

Check warning on line 70 in mesa/experimental/states.py

View check run for this annotation

Codecov / codecov/patch

mesa/experimental/states.py#L69-L70

Added lines #L69 - L70 were not covered by tests

@property
def value(self) -> float:

Check warning on line 73 in mesa/experimental/states.py

View check run for this annotation

Codecov / codecov/patch

mesa/experimental/states.py#L72-L73

Added lines #L72 - L73 were not covered by tests
"""Calculate and return current value based on elapsed time."""
current_time = self.model.time

Check warning on line 75 in mesa/experimental/states.py

View check run for this annotation

Codecov / codecov/patch

mesa/experimental/states.py#L75

Added line #L75 was not covered by tests
if current_time > self._last_update_time:
self.update(current_time)
return self._value

Check warning on line 78 in mesa/experimental/states.py

View check run for this annotation

Codecov / codecov/patch

mesa/experimental/states.py#L77-L78

Added lines #L77 - L78 were not covered by tests

def update(self, time: float) -> None:

Check warning on line 80 in mesa/experimental/states.py

View check run for this annotation

Codecov / codecov/patch

mesa/experimental/states.py#L80

Added line #L80 was not covered by tests
"""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

Check warning on line 84 in mesa/experimental/states.py

View check run for this annotation

Codecov / codecov/patch

mesa/experimental/states.py#L82-L84

Added lines #L82 - L84 were not covered by tests


class CompositeState(State):

Check warning on line 87 in mesa/experimental/states.py

View check run for this annotation

Codecov / codecov/patch

mesa/experimental/states.py#L87

Added line #L87 was not covered by tests
"""A state derived from other states."""

def __init__(

Check warning on line 90 in mesa/experimental/states.py

View check run for this annotation

Codecov / codecov/patch

mesa/experimental/states.py#L90

Added line #L90 was not covered by tests
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

Check warning on line 99 in mesa/experimental/states.py

View check run for this annotation

Codecov / codecov/patch

mesa/experimental/states.py#L97-L99

Added lines #L97 - L99 were not covered by tests

@property
def value(self) -> Any:

Check warning on line 102 in mesa/experimental/states.py

View check run for this annotation

Codecov / codecov/patch

mesa/experimental/states.py#L101-L102

Added lines #L101 - L102 were not covered by tests
"""Compute value based on dependent states."""
state_values = [state.value for state in self.dependent_states]
return self.computation_function(*state_values)

Check warning on line 105 in mesa/experimental/states.py

View check run for this annotation

Codecov / codecov/patch

mesa/experimental/states.py#L104-L105

Added lines #L104 - L105 were not covered by tests

def update(self, time: float) -> None:

Check warning on line 107 in mesa/experimental/states.py

View check run for this annotation

Codecov / codecov/patch

mesa/experimental/states.py#L107

Added line #L107 was not covered by tests
"""Update all dependent states."""
for state in self.dependent_states:
state.update(time)

Check warning on line 110 in mesa/experimental/states.py

View check run for this annotation

Codecov / codecov/patch

mesa/experimental/states.py#L110

Added line #L110 was not covered by tests


class StateAgent(Agent):

Check warning on line 113 in mesa/experimental/states.py

View check run for this annotation

Codecov / codecov/patch

mesa/experimental/states.py#L113

Added line #L113 was not covered by tests
"""An agent with integrated state management."""

def __init__(self, model) -> None:

Check warning on line 116 in mesa/experimental/states.py

View check run for this annotation

Codecov / codecov/patch

mesa/experimental/states.py#L116

Added line #L116 was not covered by tests
"""Create a new agent with state management."""
super().__init__(model)
self.states = {} # name -> State mapping

Check warning on line 119 in mesa/experimental/states.py

View check run for this annotation

Codecov / codecov/patch

mesa/experimental/states.py#L118-L119

Added lines #L118 - L119 were not covered by tests

def add_state(self, state: State) -> None:

Check warning on line 121 in mesa/experimental/states.py

View check run for this annotation

Codecov / codecov/patch

mesa/experimental/states.py#L121

Added line #L121 was not covered by tests
"""Add a new state to the agent."""
state.model = self.model
self.states[state.name] = state

Check warning on line 124 in mesa/experimental/states.py

View check run for this annotation

Codecov / codecov/patch

mesa/experimental/states.py#L123-L124

Added lines #L123 - L124 were not covered by tests

def get_state(self, name: str) -> Any:

Check warning on line 126 in mesa/experimental/states.py

View check run for this annotation

Codecov / codecov/patch

mesa/experimental/states.py#L126

Added line #L126 was not covered by tests
"""Get the current value of a state."""
return self.states[name].value

Check warning on line 128 in mesa/experimental/states.py

View check run for this annotation

Codecov / codecov/patch

mesa/experimental/states.py#L128

Added line #L128 was not covered by tests

def set_state(self, name: str, value: Any) -> None:

Check warning on line 130 in mesa/experimental/states.py

View check run for this annotation

Codecov / codecov/patch

mesa/experimental/states.py#L130

Added line #L130 was not covered by tests
"""Set value for a discrete state."""
if isinstance(self.states[name], DiscreteState):
self.states[name].value = value

Check warning on line 133 in mesa/experimental/states.py

View check run for this annotation

Codecov / codecov/patch

mesa/experimental/states.py#L133

Added line #L133 was not covered by tests
else:
raise ValueError("Cannot directly set value of non-discrete state")

Check warning on line 135 in mesa/experimental/states.py

View check run for this annotation

Codecov / codecov/patch

mesa/experimental/states.py#L135

Added line #L135 was not covered by tests

def update_states(self) -> None:

Check warning on line 137 in mesa/experimental/states.py

View check run for this annotation

Codecov / codecov/patch

mesa/experimental/states.py#L137

Added line #L137 was not covered by tests
"""Update all states to current time."""
for state in self.states.values():
state.update(self.model.time)

Check warning on line 140 in mesa/experimental/states.py

View check run for this annotation

Codecov / codecov/patch

mesa/experimental/states.py#L140

Added line #L140 was not covered by tests

0 comments on commit a50dfdc

Please sign in to comment.