Epstein civil violence: Switch to experimental datacollector

examples/epstein_civil_violence/app.py

@@ -0,0 +1,58 @@
+from mesa.visualization import Slider, SolaraViz
+
+from epstein_civil_violence.agent import Citizen, Cop
+from epstein_civil_violence.model import EpsteinCivilViolence
+
+
+COP_COLOR = "#000000"
+AGENT_QUIET_COLOR = "#648FFF"
+AGENT_REBEL_COLOR = "#FE6100"
+JAIL_COLOR = "#808080"
+JAIL_SHAPE = "rect"
+
+
+def citizen_cop_portrayal(agent):
+    portrayal = {
+        "Shape": "circle",
+        "x": agent.pos[0],
+        "y": agent.pos[1],
+        "Filled": "true",
+    }
+
+    if isinstance(agent, Citizen):
+        color = (
+            AGENT_QUIET_COLOR if agent.condition == "Quiescent" else AGENT_REBEL_COLOR
+        )
+        color = JAIL_COLOR if agent.jail_sentence else color
+        shape = JAIL_SHAPE if agent.jail_sentence else "circle"
+        portrayal["color"] = color
+        portrayal["Shape"] = shape
+        # TODO add marker to SolaraViz to display rectangle.
+        portrayal["size"] = 0.5
+        portrayal["Filled"] = "false"
+
+    else:
+        assert isinstance(agent, Cop)
+        portrayal["color"] = COP_COLOR
+        portrayal["size"] = 0.9
+
+    return portrayal
+
+
+model_params = {
+    "height": 40,
+    "width": 40,
+    "citizen_density": Slider("Initial Agent Density", 0.7, 0.0, 0.9, 0.1),
+    "cop_density": Slider("Initial Cop Density", 0.04, 0.0, 0.1, 0.01),
+    "citizen_vision": Slider("Citizen Vision", 7, 1, 10, 1),
+    "cop_vision": Slider("Cop Vision", 7, 1, 10, 1),
+    "legitimacy": Slider("Government Legitimacy", 0.82, 0.0, 1, 0.01),
+    "max_jail_term": Slider("Max Jail Term", 30, 0, 50, 1),
+}
+page = SolaraViz(
+    EpsteinCivilViolence,
+    model_params,
+    measures=[{"Quiescent": "#648FFF", "Active": "#FE6100", "Jailed": "#808080"}],
+    agent_portrayal=citizen_cop_portrayal,
+)
+page  # noqa

examples/epstein_civil_violence/epstein_civil_violence/model.py

@@ -1,4 +1,5 @@
 import mesa
+from mesa.experimental import DataCollector
 
 from .agent import Citizen, Cop
 
@@ -61,23 +62,6 @@ def __init__(
         self.schedule = mesa.time.RandomActivation(self)
         self.grid = mesa.space.SingleGrid(width, height, torus=True)
 
-        model_reporters = {
-            "Quiescent": lambda m: self.count_type_citizens(m, "Quiescent"),
-            "Active": lambda m: self.count_type_citizens(m, "Active"),
-            "Jailed": self.count_jailed,
-            "Cops": self.count_cops,
-        }
-        agent_reporters = {
-            "x": lambda a: a.pos[0],
-            "y": lambda a: a.pos[1],
-            "breed": lambda a: a.breed,
-            "jail_sentence": lambda a: getattr(a, "jail_sentence", None),
-            "condition": lambda a: getattr(a, "condition", None),
-            "arrest_probability": lambda a: getattr(a, "arrest_probability", None),
-        }
-        self.datacollector = mesa.DataCollector(
-            model_reporters=model_reporters, agent_reporters=agent_reporters
-        )
         unique_id = 0
         if self.cop_density + self.citizen_density > 1:
             raise ValueError("Cop density + citizen density must be less than 1")
@@ -102,16 +86,40 @@ def __init__(
                 self.grid[x][y] = citizen
                 self.schedule.add(citizen)
 
-        self.running = True
-        self.datacollector.collect(self)
+        model_reporters = {
+            "Quiescent": lambda m: self.count_type_citizens(m, "Quiescent"),
+            "Active": lambda m: self.count_type_citizens(m, "Active"),
+            "Jailed": self.count_jailed,
+            "Cops": self.count_cops,
+        }
+        agent_reporters = {
+            "x": lambda agents: [pos[0] for pos in agents.get("pos")],
+            "y": lambda agents: [pos[1] for pos in agents.get("pos")],
+            "breed": "breed",
+        }
+        self.citizens = self.get_agents_of_type(Citizen)
+        self.datacollector = DataCollector(
+            self,
+            {
+                "model": model_reporters,
+                "agents": agent_reporters,
+                "citizens": {
+                    "jail_sentence": "jail_sentence",
+                    "condition": "condition",
+                    "arrest_probability": "arrest_probability",
+                },
+            },
+        )
+
+        self.datacollector.collect()
 
     def step(self):
         """
         Advance the model by one step and collect data.
         """
         self.schedule.step()
         # collect data
-        self.datacollector.collect(self)
+        self.datacollector.collect()
         self.iteration += 1
         if self.iteration > self.max_iters:
             self.running = False