update to pubsub and ConditionalAgentSet

Author: quaquel

mesa/experimental/datacollection/collectors.py

@@ -193,10 +193,10 @@ class Measure:
     # FIXME:: doing so would turn measure into a descriptor
     # FIXME:: what about callable vs. attribute based Measures?
 
-    def __init__(self, model, obj: Any, callable: Callable):
+    def __init__(self, model, obj: Any, fn: Callable):
         super().__init__()
         self.obj = obj
-        self.callable = callable
+        self.fn = fn
         self._update_step = -1
         self._cached_value = None
         self.model = model
@@ -210,15 +210,15 @@ def get_value(self, force_update: bool = False):
         """
 
         if force_update or (self.model.step != self._update_step):
-            self._cached_value = self.callable(self.obj)
+            self._cached_value = self.fn(self.obj)
         return self._cached_value
 
 
 
 class MeasureDescriptor:
 
     def some_test_method(self, obj, *args, **kwargs):
-        print("blaat")
+        print(obj)
 
     def __set_name__(self, owner, name):
         self.public_name = name

mesa/experimental/datacollection/examples/boltzmann.py

@@ -4,7 +4,7 @@
 from mesa.time import RandomActivation
 
 from mesa.experimental.datacollection.mesa_classes import ObservableModel, ObservableAgent
-from mesa.experimental.datacollection.pubsub import Events, ObservableNumber
+from mesa.experimental.datacollection.pubsub import MessageType, ObservableState
 from mesa.experimental.datacollection.collectors import collect, DataCollector, Measure, MeasureDescriptor
 from mesa.experimental.datacollection.pubsub import AgentSetObserver
 
@@ -44,7 +44,6 @@ def __init__(self, N=100, width=10, height=10):
 
         self.running = True
 
-        self.gini.some_test_method()
         self.gini = Measure(self, self.agents, compute_gini)
 
     def step(self):

mesa/experimental/datacollection/examples/epstein.py

@@ -1,10 +1,13 @@
 import math
+import enum
 
 from mesa.time import RandomActivation
 from mesa.space import SingleGrid
 
 from mesa.experimental.datacollection.mesa_classes import ObservableModel, ObservableAgent
 from mesa.experimental.datacollection.collectors import DataCollector, collect, Measure
+from mesa.experimental.datacollection.mesa_classes import ConditionalAgentSet
+from mesa.experimental.datacollection.pubsub import ObservableState
 
 
 class EpsteinAgent(ObservableAgent):
@@ -13,6 +16,12 @@ def __init__(self, unique_id, model, vision):
         self.vision = vision
 
 
+class CitizenState(enum.StrEnum):
+    ACTIVE = "active"
+    QUIESCENT = "quiescent"
+    JAILED = "jailed"
+
+
 class Citizen(EpsteinAgent):
     """
     A member of the general population, may or may not be in active rebellion.
@@ -37,6 +46,7 @@ class Citizen(EpsteinAgent):
         arrest_probability: agent's assessment of arrest probability, given
             rebellion
     """
+    condition = ObservableState()
 
     def __init__(
             self,
@@ -68,7 +78,7 @@ def __init__(
         self.regime_legitimacy = regime_legitimacy
         self.risk_aversion = risk_aversion
         self.threshold = threshold
-        self.condition = "Quiescent"
+        self.condition = CitizenState.QUIESCENT
         self.jail_sentence = 0
         self.grievance = self.hardship * (1 - self.regime_legitimacy)
         self.arrest_probability = None
@@ -77,16 +87,20 @@ def step(self):
         """
         Decide whether to activate, then move if applicable.
         """
-        if self.jail_sentence:
+        if self.condition == CitizenState.JAILED:
             self.jail_sentence -= 1
             return  # no other changes or movements if agent is in jail.
+
         self.update_neighbors()
         self.update_estimated_arrest_probability()
+
         net_risk = self.risk_aversion * self.arrest_probability
-        if self.grievance - net_risk > self.threshold:
-            self.condition = "Active"
-        else:
-            self.condition = "Quiescent"
+        if (self.grievance - net_risk > self.threshold) and (self.condition != CitizenState.ACTIVE):
+            self.condition = CitizenState.ACTIVE
+        elif self.condition == CitizenState.ACTIVE:
+            self.condition = CitizenState.QUIESCENT
+        # else, agent is quiescent and stays that way
+
         if self.model.movement and self.empty_neighbors:
             new_pos = self.random.choice(self.empty_neighbors)
             self.model.grid.move_agent(self, new_pos)
@@ -113,14 +127,17 @@ def update_estimated_arrest_probability(self):
         for c in self.neighbors:
             if (
                     isinstance(c, Citizen)
-                    and c.condition == "Active"
-                    and c.jail_sentence == 0
+                    and c.condition == CitizenState.ACTIVE
             ):
                 actives_in_vision += 1
         self.arrest_probability = 1 - math.exp(
             -1 * self.model.arrest_prob_constant * (cops_in_vision / actives_in_vision)
         )
 
+    def sent_to_jail(self, sentence):
+        self.condition = CitizenState.JAILED
+        self.jail_sentence = sentence
+
 
 class Cop(EpsteinAgent):
     """
@@ -145,14 +162,12 @@ def step(self):
             if (
                     isinstance(agent, Citizen)
                     and agent.condition == "Active"
-                    and agent.jail_sentence == 0
             ):
                 active_neighbors.append(agent)
         if active_neighbors:
             arrestee = self.random.choice(active_neighbors)
             sentence = self.random.randint(0, self.model.max_jail_term)
-            arrestee.jail_sentence = sentence
-            arrestee.condition = "Quiescent"
+            arrestee.sent_to_jail(sentence)
         if self.model.movement and self.empty_neighbors:
             new_pos = self.random.choice(self.empty_neighbors)
             self.model.grid.move_agent(self, new_pos)
@@ -226,12 +241,6 @@ def __init__(
         self.schedule = RandomActivation(self)
         self.grid = SingleGrid(width, height, torus=True)
 
-
-        citizens = m.get_agents_of_type(Citizen)
-
-        self.quiescent = Measure(citizens, lambda agent_set:
-                                 agent_set.select(lambda agent: agent.condition == "quiescent"))
-
         if self.cop_density + self.citizen_density > 1:
             raise ValueError("Cop density + citizen density must be less than 1")
         for contents, pos in self.grid.coord_iter():
@@ -252,6 +261,22 @@ def __init__(
             self.grid.place_agent(agent, pos)
             self.schedule.add(agent)
 
+        # static groups
+        citizens = self.get_agents_of_type(Citizen)
+        cops = self.get_agents_of_type(Cop)
+
+        # conditional groups
+        self.quiescent = ConditionalAgentSet(citizens, self,
+                                             condition=lambda agent: agent.condition == CitizenState.QUIESCENT)
+        self.active = ConditionalAgentSet(citizens, self,
+                                          condition=lambda agent: agent.condition == CitizenState.ACTIVE)
+        self.jailed = ConditionalAgentSet(citizens, self, condition=lambda agent: agent.jail_sentence > 0)
+
+        # measures
+        self.n_quiescent = Measure(self, self.quiescent, lambda obj: len(obj))
+        self.n_active = Measure(self, self.active, lambda obj: len(obj))
+        self.n_jailed = Measure(self, self.jailed, lambda obj: len(obj))
+
         self.running = True
 
     def step(self):
@@ -265,26 +290,17 @@ def step(self):
 
 
 if __name__ == '__main__':
-    model = EpsteinCivilViolence()
-
-    citizens = model.get_agents_of_type(Citizen)
-    cops = model.get_agents_of_type(Cop)
+    model = EpsteinCivilViolence(seed=15)
 
     dc = DataCollector(model, [
-        collect("n_quiescent", citizens, attributes="condition",
-                     callable=lambda d: len([e for e in d if e["condition"] == "Quiescent"])),
-        collect("n_active", citizens, attributes="condition",
-                     callable=lambda d: len([e for e in d if e["condition"] == "Active"])),
-        collect("jail_sentence", citizens,
-                     callable=lambda d: len([e for e in d if e["jail_sentence"] > 0])),
-        collect("data", citizens, ["jail_sentence", "arrest_probability"])
+        collect("model_data", model, attributes=["n_quiescent", "n_active", "n_jailed"]),
+        collect("jail_sentence", model.jailed),
+        collect("citizen_data", model.get_agents_of_type(Citizen), ["jail_sentence", "arrest_probability"])
     ])
 
     dc.collect_all()
-    for _ in range(10):
+    for _ in range(50):
         model.step()
         dc.collect_all()
 
-    print(dc.jail_sentence.to_dataframe().head())
-    print(dc.n_quiescent.to_dataframe().head())
-    print(dc.data.to_dataframe().head())
+    print(dc.data.to_dataframe())

mesa/experimental/datacollection/examples/scratchpad.py

@@ -0,0 +1,8 @@
+from mesa.experimental.datacollection.mesa_classes import ObservableModel
+from mesa.experimental.datacollection.pubsub import MessageType
+
+if __name__ == '__main__':
+    model = ObservableModel()
+
+
+    print(MessageType.message_types)

mesa/experimental/datacollection/mesa_classes.py

@@ -1,37 +1,120 @@
-from typing import Any
+from typing import Any, Iterable, Callable, List
 
 import itertools
 
 from ...agent import Agent, AgentSet
 from ...model import Model
-from .pubsub import EventProducer, Events
+from .pubsub import MessageProducer, MessageType
+from .collectors import Measure, MeasureDescriptor
+
 
 class UpdatedAgentSet(AgentSet):
 
-    def get(self, attr_names: str | list[str]) -> list[Any]:
+    def get(self, attr_names: str | List[str]) -> list[Any]:
+        """
+        Retrieve the specified attribute(s) from each agent in the AgentSet.
+
+        Args:
+            attr_names (str | List[str]): The name(s) of the attribute(s) to retrieve from each agent.
+
+        Returns:
+            list[Any]: A list of attribute values for each agent in the set.
+        """
+
         if isinstance(attr_names, str):
-            attr_names = [attr_names]
-        a = [[getattr(agent, attr_name) for attr_name in attr_names] for agent in self._agents]
-        return a
+            return [getattr(agent, attr_names) for agent in self._agents]
+        else:
+            return [[getattr(agent, attr_name) for attr_name in attr_names] for agent in self._agents]
+
+
+class ConditionalAgentSet(UpdatedAgentSet):
+    """This is a dynamic agent set where membership depends on a specified condition
+
+    For this agent set, memberships depends on a user specified condition, and it can change over the course of the
+    simulation. Agent membership is evaluated everytime the agent sends a STATE_CHANGE message.
+
+    If the user passes an initial set of agents, only these agents are considered to be potentially part of the
+    agent set
+
+    If the suer does not pass an initial set of agents, it defaults to model.agents
+
+
+    """
+
+    def __init__(self, agents: Iterable[Agent] | None, model: Model, condition: Callable[[Agent], bool]) -> None:
+        """
+
+        Args:
+            agents (Iterable[Agent]): An iterable of agents. These form the basis of the agents in the set
+            model (Model): A model instance
+            condition (Callable[[Agent], bool]): a function that takes an agent and returns boolean. If true, the agent
+            is considered part of the agent set, otherwise the is currently not part of the agent set.
+
+
+        """
+
+        super().__init__({}, model)
+        self._condition = condition
+
+        if agents is None:
+            agents = model.agents
+            model.subscribe(model.AGENT_ADDED.name)
+
+        for agent in agents:
+            self.add_permanently(agent)
+
+    def add_permanently(self, agent: Agent) -> None:
+        agent.subscribe(agent.STATE_CHANGE.name, self.state_change_handler)
+        self._apply_condition(agent)
+    def remove_permanently(self, agent):
+        self.remove(agent)
+        agent.unsubscribe(agent.STATE_CHANGE.name, self.state_change_handler)
+
+    def _apply_condition(self, agent: Agent):
+        if self._condition(agent):
+            self.add(agent)
+        else:
+            self.discard(agent)
+
+    def state_change_handler(self, message):
+        self._apply_condition(message.sender)
+
+    def agent_added_handler(self, message):
+        agent = message.agent
+        agent.subscribe(agent.STATE_CHANGE)
+        self._apply_condition(agent)
 
 
 class ObservableModel(Model):
+    AGENT_ADDED = MessageType("agent")
+    AGENT_REMOVED = MessageType("agent")
+    STATE_CHANGED = MessageType("state")
+
+    def __setattr__(self, name, value):
+        if isinstance(value, Measure) and not name.startswith("_"):
+            klass = type(self)
+            descr = MeasureDescriptor()
+            descr.__set_name__(klass, name)
+            setattr(klass, name, descr)
+            descr.__set__(self, value)
+        else:
+            super().__setattr__(name, value)
 
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         super().__init__(*args, **kwargs)
-        self.event_producer = EventProducer(self)
+        self.event_producer = MessageProducer(self)
 
     @property
     def time(self):
         return self._time
 
-    def add_agent(self, agent:Agent) -> None:
+    def add_agent(self, agent: Agent) -> None:
         self.agents_[type(agent)][agent] = None
-        self.event_producer.fire_event(Events.AGENT_ADDED, agent)
+        self.event_producer.send_message(self.AGENT_ADDED, agent=agent)
 
-    def remove_agent(self, agent:Agent) -> None:
+    def remove_agent(self, agent: Agent) -> None:
         self.agents_[type(agent)].pop(agent, default=None)
-        self.event_producer.fire_event(Events.AGENT_REMOVED, agent)
+        self.event_producer.send_message(self.AGENT_REMOVED, agent=agent)
 
     def get_agents_of_type(self, agenttype: type[Agent]) -> AgentSet:
         """Retrieves an AgentSet containing all agents of the specified type."""
@@ -54,6 +137,8 @@ def unsubscribe(self, event: str, event_handler: callable):
 
 
 class ObservableAgent(Agent):
+    STATE_CHANGE = MessageType("state")
+
     def __init__(self, unique_id: int, model: ObservableModel) -> None:
         """
         Create a new agent.
@@ -62,10 +147,11 @@ def __init__(self, unique_id: int, model: ObservableModel) -> None:
             unique_id (int): A unique identifier for this agent.
             model (Model): The model instance in which the agent exists.
         """
+        super().__init__(unique_id, model)
         self.unique_id = unique_id
         self.model = model
         self.pos = None
-        self.event_producer = EventProducer(self)
+        self.event_producer = MessageProducer(self)
 
         # register agent
         self.model.add_agent(self)
@@ -79,9 +165,3 @@ def subscribe(self, event: str, event_handler: callable):
     def unsubscribe(self, event: str, event_handler: callable):
         # or try except pass, which is slightly faster
         self.event_producer.unsubscribe(event, event_handler)
-
-
-
-
-
-