ongoing work

Author: quaquel

mesa/experimental/datacollection/collectors.py

@@ -8,7 +8,7 @@
 
 
 class BaseCollector:
-    def __init__(self, name:str, obj: Any, attributes: str|List[str]=None, callable: Callable=None):
+    def __init__(self, name:str, obj: Any, attributes: str|List[str]=None, fn: Callable=None):
         """
 
         Args
@@ -31,7 +31,7 @@ def __init__(self, name:str, obj: Any, attributes: str|List[str]=None, callable:
         self.name = name
         self.obj = obj
         self.attributes = attributes
-        self.callable = callable
+        self.callable = fn
         self.data_over_time = {}
 
     def collect(self, time):
@@ -67,8 +67,8 @@ def to_dataframe(self):
 
 
 class AgentSetCollector(BaseCollector):
-    def __init__(self, name, obj, attributes=None, callable=None):
-        super().__init__(name, obj, attributes=attributes, callable=callable)
+    def __init__(self, name, obj, attributes=None, fn=None):
+        super().__init__(name, obj, attributes=attributes, fn=fn)
         self.attributes.append("unique_id")
 
     def collect(self, time):
@@ -165,24 +165,24 @@ def collect_all(self):
             collector.collect(time)
 
 
-def collect(name:str, obj:Any, attributes: str|List[str]=None, callable: Callable=None):
+def collect(name:str, obj:Any, attributes: str|List[str]=None, fn: 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.
+        fn : 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?
 
     """
 
     if isinstance(obj, AgentSet):
-        return AgentSetCollector(name, obj, attributes, callable)
+        return AgentSetCollector(name, obj, attributes, fn)
     else:
-        return BaseCollector(name, obj, attributes, callable)
+        return BaseCollector(name, obj, attributes, fn)
 
 
 
@@ -193,13 +193,13 @@ class Measure:
     # 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):
+    def __init__(self, model, obj: Any, callable: Callable):
         super().__init__()
         self.obj = obj
         self.callable = callable
-        self.model = model
         self._update_step = -1
         self._cached_value = None
+        self.model = model
 
     def get_value(self, force_update: bool = False):
         """
@@ -213,3 +213,20 @@ def get_value(self, force_update: bool = False):
             self._cached_value = self.callable(self.obj)
         return self._cached_value
 
+
+
+class MeasureDescriptor:
+
+    def some_test_method(self, obj, *args, **kwargs):
+        print("blaat")
+
+    def __set_name__(self, owner, name):
+        self.public_name = name
+        self.private_name = "_" + name
+    def __get__(self, obj, owner):
+        return getattr(obj, self.private_name).get_value()
+
+    def __set__(self, obj, value):
+        setattr(obj, self.private_name, value)
+
+

mesa/experimental/datacollection/examples/boltzmann.py

@@ -5,17 +5,18 @@
 
 from mesa.experimental.datacollection.mesa_classes import ObservableModel, ObservableAgent
 from mesa.experimental.datacollection.pubsub import Events, ObservableNumber
-from mesa.experimental.datacollection.collectors import collect, DataCollector
+from mesa.experimental.datacollection.collectors import collect, DataCollector, Measure, MeasureDescriptor
 from mesa.experimental.datacollection.pubsub import AgentSetObserver
 
-def compute_gini(model):
-    agent_wealths = [agent.wealth for agent in model.schedule.agents]
+def compute_gini(agents):
+    agent_wealths = [agent.wealth for agent in agents]
     x = sorted(agent_wealths)
     N = model.num_agents
     B = sum(xi * (N - i) for i, xi in enumerate(x)) / (N * sum(x))
     return 1 + (1 / N) - 2 * B
 
 
+
 class BoltzmannWealthModel(ObservableModel):
     """A simple model of an economy where agents exchange currency at random.
 
@@ -24,14 +25,13 @@ class BoltzmannWealthModel(ObservableModel):
     highly skewed distribution of wealth.
     """
 
+    # gini = MeasureDescriptor()
+
     def __init__(self, N=100, width=10, height=10):
         super().__init__()
         self.num_agents = N
         self.grid = MultiGrid(width, height, True)
         self.schedule = RandomActivation(self)
-        # self.datacollector = DataCollector(
-        #     model_reporters={"Gini": compute_gini}, agent_reporters={"Wealth": "wealth"}
-        # )
 
         # Create agents
         for i in range(self.num_agents):
@@ -43,21 +43,16 @@ def __init__(self, N=100, width=10, height=10):
             self.grid.place_agent(a, (x, y))
 
         self.running = True
-        # self.datacollector.collect(self)
+
+        self.gini.some_test_method()
+        self.gini = Measure(self, self.agents, compute_gini)
 
     def step(self):
         self.schedule.step()
-        # collect data
-        # self.datacollector.collect(self)
-
-    def run_model(self, n):
-        for i in range(n):
-            self.step()
 
 
 class MoneyAgent(ObservableAgent):
     """An agent with fixed initial wealth."""
-    wealth = ObservableNumber()
 
     def __init__(self, unique_id, model):
         super().__init__(unique_id, model)
@@ -90,14 +85,20 @@ def handler(subject, state):
     if state == "wealth":
         return getattr(subject, state)
 
+
+def some_func(obj):
+    return obj.get_value()
+
 if __name__ == '__main__':
     model = BoltzmannWealthModel()
 
-    datacollector = DataCollector(model, [collect("wealth", model.agents)])
+    datacollector = DataCollector(model, [collect("wealth", model.agents),
+                                          collect("gini", model)])
 
 
     for _ in range(10):
         model.step()
         datacollector.collect_all()
 
-    print(datacollector.wealth.to_dataframe().head())
+    print(datacollector.wealth.to_dataframe().head())
+    print(datacollector.gini.to_dataframe().head())

mesa/experimental/datacollection/mesa_classes.py

@@ -16,6 +16,7 @@ def get(self, attr_names: str | list[str]) -> list[Any]:
 
 
 class ObservableModel(Model):
+
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         super().__init__(*args, **kwargs)
         self.event_producer = EventProducer(self)

mesa/experimental/datacollection/scratchpad.py

@@ -0,0 +1,44 @@
+class Measure:
+
+    def __init__(self, model, identifier, *args, **kwargs):
+        self.model = model
+        self.identifier = identifier
+
+    def get_value(self):
+        return getattr(self.model, self.identifier)
+
+
+class MeasureDescriptor:
+    def __set_name__(self, owner, name):
+        self.public_name = name
+        self.private_name = "_" + name
+
+    def __get__(self, obj, owner):
+        return getattr(obj, self.private_name).get_value()
+
+    def __set__(self, obj, value):
+        setattr(obj, self.private_name, value)
+
+
+class Model:
+
+    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, identifier, *args, **kwargs):
+        self.gini = Measure(self, "identifier")
+        self.identifier = identifier
+
+
+if __name__ == '__main__':
+    model1 = Model(1)
+    model2 = Model(2)
+    print(model1.gini)
+    print(model2.gini)

mesa/model.py

@@ -21,7 +21,17 @@
 TimeT = Union[float, int]
 
 
-class Model:
+class MetaClass(type):
+    def __new__(cls, name, bases, attrs):
+        return super().__new__(cls, name, bases, attrs)
+
+    def __init__(self, name, bases, attrs):
+        # perform any additional initialization here...
+        super().__init__(name, bases, attrs)
+        print("blaat")
+
+
+class Model(metaclass=MetaClass):
     """Base class for models in the Mesa ABM library.
 
     This class serves as a foundational structure for creating agent-based models.