Updates to Epstein example (projectmesa#2429)

* rework of epstein

* shift to von neuman grid
* update of perceived risk with round operator
* introduction of enum for state
* cleanup of statistics gathering to use enum
* update of visualization code to be consistent with what is currently supported when drawing new_style discrete grids

* further cleaning

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* cleanup

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* shift to dict for colors

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

mesa/examples/advanced/epstein_civil_violence/agents.py

@@ -1,18 +1,29 @@
 import math
+from enum import Enum
 
 import mesa
 
 
+class CitizenState(Enum):
+    ACTIVE = 1
+    QUIET = 2
+    ARRESTED = 3
+
+
 class EpsteinAgent(mesa.experimental.cell_space.CellAgent):
     def update_neighbors(self):
         """
         Look around and see who my neighbors are
         """
         self.neighborhood = self.cell.get_neighborhood(radius=self.vision)
-
         self.neighbors = self.neighborhood.agents
         self.empty_neighbors = [c for c in self.neighborhood if c.is_empty]
 
+    def move(self):
+        if self.model.movement and self.empty_neighbors:
+            new_pos = self.random.choice(self.empty_neighbors)
+            self.move_to(new_pos)
+
 
 class Citizen(EpsteinAgent):
     """
@@ -38,13 +49,7 @@ class Citizen(EpsteinAgent):
     """
 
     def __init__(
-        self,
-        model,
-        hardship,
-        regime_legitimacy,
-        risk_aversion,
-        threshold,
-        vision,
+        self, model, regime_legitimacy, threshold, vision, arrest_prob_constant
     ):
         """
         Create a new Citizen.
@@ -62,16 +67,21 @@ def __init__(
             model: model instance
         """
         super().__init__(model)
-        self.hardship = hardship
+        self.hardship = self.random.random()
+        self.risk_aversion = self.random.random()
         self.regime_legitimacy = regime_legitimacy
-        self.risk_aversion = risk_aversion
         self.threshold = threshold
-        self.condition = "Quiescent"
+        self.state = CitizenState.QUIET
         self.vision = vision
         self.jail_sentence = 0
         self.grievance = self.hardship * (1 - self.regime_legitimacy)
+        self.arrest_prob_constant = arrest_prob_constant
         self.arrest_probability = None
 
+        self.neighborhood = []
+        self.neighbors = []
+        self.empty_neighbors = []
+
     def step(self):
         """
         Decide whether to activate, then move if applicable.
@@ -81,32 +91,33 @@ def step(self):
             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"
+        if (self.grievance - net_risk) > self.threshold:
+            self.state = CitizenState.ACTIVE
         else:
-            self.condition = "Quiescent"
+            self.state = CitizenState.QUIET
 
-        if self.model.movement and self.empty_neighbors:
-            new_cell = self.random.choice(self.empty_neighbors)
-            self.move_to(new_cell)
+        self.move()
 
     def update_estimated_arrest_probability(self):
         """
         Based on the ratio of cops to actives in my neighborhood, estimate the
         p(Arrest | I go active).
         """
-        cops_in_vision = len([c for c in self.neighbors if isinstance(c, Cop)])
-        actives_in_vision = 1.0  # citizen counts herself
-        for c in self.neighbors:
-            if (
-                isinstance(c, Citizen)
-                and c.condition == "Active"
-                and c.jail_sentence == 0
-            ):
+        cops_in_vision = 0
+        actives_in_vision = 1  # citizen counts herself
+        for neighbor in self.neighbors:
+            if isinstance(neighbor, Cop):
+                cops_in_vision += 1
+            elif neighbor.state == CitizenState.ACTIVE:
                 actives_in_vision += 1
+
+        # there is a body of literature on this equation
+        # the round is not in the pnas paper but without it, its impossible to replicate
+        # the dynamics shown there.
         self.arrest_probability = 1 - math.exp(
-            -1 * self.model.arrest_prob_constant * (cops_in_vision / actives_in_vision)
+            -1 * self.arrest_prob_constant * round(cops_in_vision / actives_in_vision)
         )
 
 
@@ -122,7 +133,7 @@ class Cop(EpsteinAgent):
             able to inspect
     """
 
-    def __init__(self, model, vision):
+    def __init__(self, model, vision, max_jail_term):
         """
         Create a new Cop.
         Args:
@@ -133,6 +144,7 @@ def __init__(self, model, vision):
         """
         super().__init__(model)
         self.vision = vision
+        self.max_jail_term = max_jail_term
 
     def step(self):
         """
@@ -142,17 +154,11 @@ def step(self):
         self.update_neighbors()
         active_neighbors = []
         for agent in self.neighbors:
-            if (
-                isinstance(agent, Citizen)
-                and agent.condition == "Active"
-                and agent.jail_sentence == 0
-            ):
+            if isinstance(agent, Citizen) and agent.state == CitizenState.ACTIVE:
                 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"
-        if self.model.movement and self.empty_neighbors:
-            new_pos = self.random.choice(self.empty_neighbors)
-            self.move_to(new_pos)
+            arrestee.jail_sentence = self.random.randint(0, self.max_jail_term)
+            arrestee.state = CitizenState.ARRESTED
+
+        self.move()

mesa/examples/advanced/epstein_civil_violence/app.py

@@ -1,4 +1,8 @@
-from mesa.examples.advanced.epstein_civil_violence.agents import Citizen, Cop
+from mesa.examples.advanced.epstein_civil_violence.agents import (
+    Citizen,
+    CitizenState,
+    Cop,
+)
 from mesa.examples.advanced.epstein_civil_violence.model import EpsteinCivilViolence
 from mesa.visualization import (
     Slider,
@@ -8,10 +12,12 @@
 )
 
 COP_COLOR = "#000000"
-AGENT_QUIET_COLOR = "#648FFF"
-AGENT_REBEL_COLOR = "#FE6100"
-JAIL_COLOR = "#808080"
-JAIL_SHAPE = "rect"
+
+agent_colors = {
+    CitizenState.ACTIVE: "#FE6100",
+    CitizenState.QUIET: "#648FFF",
+    CitizenState.ARRESTED: "#808080",
+}
 
 
 def citizen_cop_portrayal(agent):
@@ -20,29 +26,12 @@ def citizen_cop_portrayal(agent):
 
     portrayal = {
         "size": 25,
-        "shape": "s",  # square marker
     }
 
     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
-        if shape == "s":
-            portrayal["w"] = 0.9
-            portrayal["h"] = 0.9
-        else:
-            portrayal["r"] = 0.5
-            portrayal["filled"] = False
-        portrayal["layer"] = 0
-
+        portrayal["color"] = agent_colors[agent.state]
     elif isinstance(agent, Cop):
         portrayal["color"] = COP_COLOR
-        portrayal["r"] = 0.9
-        portrayal["layer"] = 1
 
     return portrayal
 
@@ -59,7 +48,7 @@ def citizen_cop_portrayal(agent):
 }
 
 space_component = make_space_matplotlib(citizen_cop_portrayal)
-chart_component = make_plot_measure(["Quiescent", "Active", "Jailed"])
+chart_component = make_plot_measure([state.name.lower() for state in CitizenState])
 
 epstein_model = EpsteinCivilViolence()
 

mesa/examples/advanced/epstein_civil_violence/model.py

@@ -1,13 +1,18 @@
 import mesa
-from mesa.examples.advanced.epstein_civil_violence.agents import Citizen, Cop
+from mesa.examples.advanced.epstein_civil_violence.agents import (
+    Citizen,
+    CitizenState,
+    Cop,
+)
 
 
 class EpsteinCivilViolence(mesa.Model):
     """
     Model 1 from "Modeling civil violence: An agent-based computational
     approach," by Joshua Epstein.
     http://www.pnas.org/content/99/suppl_3/7243.full
-    Attributes:
+
+    Args:
         height: grid height
         width: grid width
         citizen_density: approximate % of cells occupied by citizens.
@@ -45,102 +50,65 @@ def __init__(
         seed=None,
     ):
         super().__init__(seed=seed)
-        self.width = width
-        self.height = height
-        self.citizen_density = citizen_density
-        self.cop_density = cop_density
-        self.citizen_vision = citizen_vision
-        self.cop_vision = cop_vision
-        self.legitimacy = legitimacy
-        self.max_jail_term = max_jail_term
-        self.active_threshold = active_threshold
-        self.arrest_prob_constant = arrest_prob_constant
         self.movement = movement
         self.max_iters = max_iters
-        self.iteration = 0
 
-        self.grid = mesa.experimental.cell_space.OrthogonalMooreGrid(
+        self.grid = mesa.experimental.cell_space.OrthogonalVonNeumannGrid(
             (width, height), capacity=1, torus=True, random=self.random
         )
 
         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,
+            "active": CitizenState.ACTIVE.name,
+            "quiet": CitizenState.QUIET.name,
+            "arrested": CitizenState.ARRESTED.name,
         }
         agent_reporters = {
-            "x": lambda a: a.cell.coordinate[0],
-            "y": lambda a: a.cell.coordinate[1],
-            "breed": lambda a: type(a).__name__,
             "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
         )
-        if self.cop_density + self.citizen_density > 1:
+        if cop_density + citizen_density > 1:
             raise ValueError("Cop density + citizen density must be less than 1")
 
         for cell in self.grid.all_cells:
-            if self.random.random() < self.cop_density:
-                cop = Cop(self, vision=self.cop_vision)
-                cop.move_to(cell)
+            klass = self.random.choices(
+                [Citizen, Cop, None],
+                cum_weights=[citizen_density, citizen_density + cop_density, 1],
+            )[0]
 
-            elif self.random.random() < (self.cop_density + self.citizen_density):
+            if klass == Cop:
+                cop = Cop(self, vision=cop_vision, max_jail_term=max_jail_term)
+                cop.move_to(cell)
+            elif klass == Citizen:
                 citizen = Citizen(
                     self,
-                    hardship=self.random.random(),
-                    regime_legitimacy=self.legitimacy,
-                    risk_aversion=self.random.random(),
-                    threshold=self.active_threshold,
-                    vision=self.citizen_vision,
+                    regime_legitimacy=legitimacy,
+                    threshold=active_threshold,
+                    vision=citizen_vision,
+                    arrest_prob_constant=arrest_prob_constant,
                 )
                 citizen.move_to(cell)
 
         self.running = True
+        self._update_counts()
         self.datacollector.collect(self)
 
     def step(self):
         """
         Advance the model by one step and collect data.
         """
         self.agents.shuffle_do("step")
-        # collect data
+        self._update_counts()
         self.datacollector.collect(self)
-        self.iteration += 1
-        if self.iteration > self.max_iters:
-            self.running = False
-
-    @staticmethod
-    def count_type_citizens(model, condition, exclude_jailed=True):
-        """
-        Helper method to count agents by Quiescent/Active.
-        """
-        citizens = model.agents_by_type[Citizen]
 
-        if exclude_jailed:
-            return len(
-                [
-                    c
-                    for c in citizens
-                    if (c.condition == condition) and (c.jail_sentence == 0)
-                ]
-            )
-        else:
-            return len([c for c in citizens if c.condition == condition])
+        if self.steps > self.max_iters:
+            self.running = False
 
-    @staticmethod
-    def count_jailed(model):
-        """
-        Helper method to count jailed agents.
-        """
-        return len([a for a in model.agents_by_type[Citizen] if a.jail_sentence > 0])
+    def _update_counts(self):
+        """Helper function for counting nr. of citizens in given state."""
+        counts = self.agents_by_type[Citizen].groupby("state").count()
 
-    @staticmethod
-    def count_cops(model):
-        """
-        Helper method to count jailed agents.
-        """
-        return len(model.agents_by_type[Cop])
+        for state in CitizenState:
+            setattr(self, state.name, counts.get(state, 0))