ongoing work

Author: quaquel

mesa/experimental/datacollection/collectors.py

@@ -4,21 +4,26 @@
 import pandas as pd
 
 from mesa.agent import AgentSet
+from mesa import Model
 
 
 class BaseCollector:
-    def __init__(self, name:str, obj: Any, attributes: str|List[str]=None, callable: Callable=None, column_names: List[str]=None):
+    def __init__(self, name:str, obj: Any, attributes: str|List[str]=None, callable: Callable=None):
         """
 
         Args
         name : name of the collector
         obj : object
         attributes :
         callable : callable
-        column_names : list of column names, optional, defaults to attributes
+
+        Note::
+        if a callable is passed, it is assumed that there will only be a single return value
 
         """
         super().__init__()
+        if attributes is None:
+            attributes = [name,]
 
         if isinstance(attributes, str):
             attributes = [attributes,]
@@ -27,7 +32,6 @@ def __init__(self, name:str, obj: Any, attributes: str|List[str]=None, callable:
         self.obj = obj
         self.attributes = attributes
         self.callable = callable
-        self.column_names = column_names
         self.data_over_time = {}
 
     def collect(self, time):
@@ -54,9 +58,7 @@ def to_dataframe(self):
         # or name, or even something else if callable does funky stuff
         # so we need meaningful defaults and a way to override them
 
-        if self.column_names is not None:
-            columns = self.column_names
-        elif self.callable is not None:
+        if self.callable is not None:
             columns = [self.name]
         else:
             columns = self.attributes
@@ -65,8 +67,8 @@ def to_dataframe(self):
 
 
 class AgentSetCollector(BaseCollector):
-    def __init__(self, name, obj, attributes=None, callable=None, column_names=None):
-        super().__init__(name, obj, attributes, callable, column_names)
+    def __init__(self, name, obj, attributes=None, callable=None):
+        super().__init__(name, obj, attributes=attributes, callable=callable)
         self.attributes.append("unique_id")
 
     def collect(self, time):
@@ -163,9 +165,51 @@ def collect_all(self):
             collector.collect(time)
 
 
+def collect(name:str, obj:Any, attributes: str|List[str]=None, callable: Callable=None):
+    """
+
+    Args
+        name : name of the collector
+        obj : object form which to collect information
+        attributes : attributes to collect, option. If not provided, attributes defaults to name
+        callable : callable to apply to collected data.
+
+    FIXME:: what about callable to object directly? or simply not allow for it and solve this
+    FIXME:: through measures?
+
+    """
 
-def collect_from(name, object, attributes, callable=None):
-    if isinstance(object, AgentSet):
-        return AgentSetCollector(name, object, attributes, callable)
+    if isinstance(obj, AgentSet):
+        return AgentSetCollector(name, obj, attributes, callable)
     else:
-        return BaseCollector(name, object, attributes, callable)
+        return BaseCollector(name, obj, attributes, callable)
+
+
+
+
+class Measure:
+    # FIXME:: do we want AgentSet based measures?
+    # FIXME:: can we play some property trick to enable attribute retrieval
+    # FIXME:: doing so would turn measure into a descriptor
+    # FIXME:: what about callable vs. attribute based Measures?
+
+    def __init__(self, model:Model, obj: Any, callable: Callable):
+        super().__init__()
+        self.obj = obj
+        self.callable = callable
+        self.model = model
+        self._update_step = -1
+        self._cached_value = None
+
+    def get_value(self, force_update: bool = False):
+        """
+
+        Args:
+            force_update (bool): force recalculation of measure.
+
+        """
+
+        if force_update or (self.model.step != self._update_step):
+            self._cached_value = self.callable(self.obj)
+        return self._cached_value
+

mesa/experimental/datacollection/examples/boltzmann.py

@@ -5,7 +5,7 @@
 
 from mesa.experimental.datacollection.mesa_classes import ObservableModel, ObservableAgent
 from mesa.experimental.datacollection.pubsub import Events, ObservableNumber
-from mesa.experimental.datacollection.collectors import BaseCollector, AgentSetCollector, DataCollector
+from mesa.experimental.datacollection.collectors import collect, DataCollector
 from mesa.experimental.datacollection.pubsub import AgentSetObserver
 
 def compute_gini(model):
@@ -93,10 +93,7 @@ def handler(subject, state):
 if __name__ == '__main__':
     model = BoltzmannWealthModel()
 
-    observer = AgentSetObserver(model.agents, Events.STATE_CHANGE, handler)
-
-    datacollector = DataCollector(model)
-    datacollector.add_collector(AgentSetCollector("wealth", model.agents, "wealth"))
+    datacollector = DataCollector(model, [collect("wealth", model.agents)])
 
 
     for _ in range(10):

mesa/experimental/datacollection/examples/epstein.py

@@ -4,8 +4,7 @@
 from mesa.space import SingleGrid
 
 from mesa.experimental.datacollection.mesa_classes import ObservableModel, ObservableAgent
-from mesa.experimental.datacollection.pubsub import Events, ObservableNumber
-from mesa.experimental.datacollection.collectors import DataCollector, collect_from
+from mesa.experimental.datacollection.collectors import DataCollector, collect, Measure
 
 
 class EpsteinAgent(ObservableAgent):
@@ -224,10 +223,15 @@ def __init__(
         self.arrest_prob_constant = arrest_prob_constant
         self.movement = movement
         self.max_iters = max_iters
-        self.iteration = 0
         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():
@@ -256,31 +260,31 @@ def step(self):
         """
         self.schedule.step()
         # collect data
-        self.iteration += 1
-        if self.iteration > self.max_iters:
+        if self._steps > self.max_iters:
             self.running = False
 
 
 if __name__ == '__main__':
     model = EpsteinCivilViolence()
 
     citizens = model.get_agents_of_type(Citizen)
-    cops = model.get_agents_of_type(Citizen)
+    cops = model.get_agents_of_type(Cop)
 
     dc = DataCollector(model, [
-        collect_from("n_quiescent", citizens, "condition",
-                     lambda d: len([e for e in d if e["condition"] == "Quiescent"])),
-        collect_from("n_active", cops, "condition",
-                     lambda d: len([e for e in d if e["condition"] == "Active"])),
-        collect_from("jailed", citizens, "jail_sentence",
-                     lambda d: len([e for e in d if e["jail_sentence"] > 0])),
-        collect_from("data", citizens, ["jail_sentence", "arrest_probability"])
+        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"])
     ])
 
     dc.collect_all()
     for _ in range(10):
         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())