update to boltzman

Author: quaquel

mesa/experimental/datacollection/examples/boltzmann.py

@@ -95,7 +95,7 @@ def handler(subject, state):
 
     observer = AgentSetObserver(model.agents, Events.STATE_CHANGE, handler)
 
-    datacollector = DataCollector()
+    datacollector = DataCollector(model)
     datacollector.add_collector(AgentSetCollector("wealth", model.agents, "wealth"))
 
 

mesa/experimental/datacollection/mesa_classes.py

@@ -1,9 +1,19 @@
 from typing import Any
 
+import itertools
+
 from ...agent import Agent, AgentSet
 from ...model import Model
 from .pubsub import EventProducer, Events
 
+class UpdatedAgentSet(AgentSet):
+
+    def get(self, attr_names: str | list[str]) -> list[Any]:
+        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
+
 
 class ObservableModel(Model):
     def __init__(self, *args: Any, **kwargs: Any) -> None:
@@ -26,6 +36,14 @@ def get_agents_of_type(self, agenttype: type[Agent]) -> AgentSet:
         """Retrieves an AgentSet containing all agents of the specified type."""
         return UpdatedAgentSet(self.agents_[agenttype].keys(), self)
 
+    @property
+    def agents(self) -> AgentSet:
+        if hasattr(self, "_agents"):
+            return self._agents
+        else:
+            all_agents = itertools.chain.from_iterable(self.agents_.values())
+            return UpdatedAgentSet(all_agents, self)
+
     def subscribe(self, event: str, event_handler: callable):
         self.event_producer.subscribe(event, event_handler)
 
@@ -62,13 +80,7 @@ def unsubscribe(self, event: str, event_handler: callable):
         self.event_producer.unsubscribe(event, event_handler)
 
 
-class UpdatedAgentSet(AgentSet):
 
-    def get(self, attr_names: str | list[str]) -> list[Any]:
-        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