Improve Caching and Replay Example (#102)

* Improve Caching and Replay

* Update to benchmark

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

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

* Update model.py

* Update run.py

* Update cacheablemodel.py

---------

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

Author: coderbeta1

examples/caching_and_replay/README.md

@@ -2,7 +2,7 @@
 
 ## Summary
 
-This example applies caching on the Mesa [Schelling example](https://github.com/projectmesa/mesa-examples/tree/main/examples/Schelling).
+This example applies caching on the Mesa [Schelling example](https://github.com/projectmesa/mesa-examples/tree/main/examples/schelling).
 It enables a simulation run to be "cached" or in other words recorded. The recorded simulation run is persisted on the local file system and can be replayed at any later point.
 
 It uses the [Mesa-Replay](https://github.com/Logende/mesa-replay) library and puts the Schelling model inside a so-called `CacheableModel` wrapper that we name `CacheableSchelling`.
@@ -29,6 +29,14 @@ To run the model interactively, run ``mesa runserver`` in this directory. e.g.
     $ mesa runserver
 ```
 
+or
+
+Directly run the file ``run.py`` in the terminal. e.g.
+
+```
+    $ python run.py
+```
+
 Then open your browser to [http://127.0.0.1:8521/](http://127.0.0.1:8521/) and press Reset, then Run.
 
 First, run the **simulation** with the 'Replay' switch disabled.

examples/caching_and_replay/cacheablemodel.py

@@ -3,15 +3,17 @@
 
 
 class CacheableSchelling(CacheableModel):
-    """A wrapper around the original Schelling model to make the simulation cacheable
+    """
+    A wrapper around the original Schelling model to make the simulation cacheable
     and replay-able.  Uses CacheableModel from the Mesa-Replay library,
     which is a wrapper that can be put around any regular mesa model to make it
     "cacheable".
     From outside, a CacheableSchelling instance can be treated like any
     regular Mesa model.
     The only difference is that the model will write the state of every simulation step
     to a cache file or when in replay mode use a given cache file to replay that cached
-    simulation run."""
+    simulation run.
+    """
 
     def __init__(
         self,
@@ -20,14 +22,23 @@ def __init__(
         density=0.8,
         minority_pc=0.2,
         homophily=3,
+        radius=1,
+        cache_file_path="./my_cache_file_path.cache",
         # Note that this is an additional parameter we add to our model,
         # which decides whether to simulate or replay
         replay=False,
     ):
-        actual_model = Schelling(width, height, density, minority_pc, homophily)
+        actual_model = Schelling(
+            width=width,
+            height=height,
+            density=density,
+            minority_pc=minority_pc,
+            homophily=homophily,
+            radius=radius,
+        )
         cache_state = CacheState.REPLAY if replay else CacheState.RECORD
         super().__init__(
-            actual_model,
-            cache_file_path="my_cache_file_path.cache",
+            model=actual_model,
+            cache_file_path=cache_file_path,
             cache_state=cache_state,
         )

examples/caching_and_replay/model.py

@@ -8,7 +8,7 @@ class SchellingAgent(mesa.Agent):
     Schelling segregation agent
     """
 
-    def __init__(self, pos, model, agent_type):
+    def __init__(self, unique_id, model, agent_type):
         """
         Create a new Schelling agent.
 
@@ -17,13 +17,14 @@ def __init__(self, pos, model, agent_type):
            x, y: Agent initial location.
            agent_type: Indicator for the agent's type (minority=1, majority=0)
         """
-        super().__init__(pos, model)
-        self.pos = pos
+        super().__init__(unique_id, model)
         self.type = agent_type
 
     def step(self):
         similar = 0
-        for neighbor in self.model.grid.iter_neighbors(self.pos, True):
+        for neighbor in self.model.grid.iter_neighbors(
+            self.pos, moore=True, radius=self.model.radius
+        ):
             if neighbor.type == self.type:
                 similar += 1
 
@@ -39,48 +40,64 @@ class Schelling(mesa.Model):
     Model class for the Schelling segregation model.
     """
 
-    def __init__(self, width=20, height=20, density=0.8, minority_pc=0.2, homophily=3):
-        """ """
-        super().__init__()
-        self.width = width
+    def __init__(
+        self,
+        height=20,
+        width=20,
+        homophily=3,
+        radius=1,
+        density=0.8,
+        minority_pc=0.3,
+        seed=None,
+    ):
+        """
+        Create a new Schelling model.
+
+        Args:
+            width, height: Size of the space.
+            density: Initial Chance for a cell to populated
+            minority_pc: Chances for an agent to be in minority class
+            homophily: Minimum number of agents of same class needed to be happy
+            radius: Search radius for checking similarity
+            seed: Seed for Reproducibility
+        """
+
+        super().__init__(seed=seed)
         self.height = height
+        self.width = width
         self.density = density
         self.minority_pc = minority_pc
         self.homophily = homophily
+        self.radius = radius
 
         self.schedule = mesa.time.RandomActivation(self)
         self.grid = mesa.space.SingleGrid(width, height, torus=True)
 
         self.happy = 0
         self.datacollector = mesa.DataCollector(
-            {"happy": "happy"},  # Model-level count of happy agents
-            # For testing purposes, agent's individual x and y
-            {"x": lambda a: a.pos[0], "y": lambda a: a.pos[1]},
+            model_reporters={"happy": "happy"},  # Model-level count of happy agents
         )
 
         # Set up agents
         # We use a grid iterator that returns
         # the coordinates of a cell as well as
         # its contents. (coord_iter)
-        for cell in self.grid.coord_iter():
-            x, y = cell[1]
+        for _, pos in self.grid.coord_iter():
             if self.random.random() < self.density:
                 agent_type = 1 if self.random.random() < self.minority_pc else 0
-
-                agent = SchellingAgent((x, y), self, agent_type)
-                self.grid.place_agent(agent, (x, y))
+                agent = SchellingAgent(self.next_id(), self, agent_type)
+                self.grid.place_agent(agent, pos)
                 self.schedule.add(agent)
 
-        self.running = True
         self.datacollector.collect(self)
 
     def step(self):
         """
-        Run one step of the model. If All agents are happy, halt the model.
+        Run one step of the model.
         """
         self.happy = 0  # Reset counter of happy agents
         self.schedule.step()
-        # collect data
+
         self.datacollector.collect(self)
 
         if self.happy == self.schedule.get_agent_count():

examples/caching_and_replay/run.py

@@ -6,25 +6,30 @@
 
 # As 'replay' is a simulation model parameter in this example, we need to make it available as such
 model_params["replay"] = mesa.visualization.Checkbox("Replay cached run?", False)
+model_params["cache_file_path"] = "./my_cache_file_path.cache"
 
 
 def get_cache_file_status(_):
     """
     Display an informational text about caching and the status of the cache file (existing versus not existing)
     """
-    cache_file = Path("./my_cache_file_path.cache")
+    cache_file = Path(model_params["cache_file_path"])
     return (
         f"Only activate the 'Replay cached run?' switch when a cache file already exists, otherwise it will fail. "
         f"Cache file existing: '{cache_file.exists()}'."
     )
 
 
 server = mesa.visualization.ModularServer(
-    # Note that Schelling was replaced by CacheableSchelling here
-    CacheableSchelling,
-    [get_cache_file_status, canvas_element, get_happy_agents, happy_chart],
-    "Schelling",
-    model_params,
+    model_cls=CacheableSchelling,  # Note that Schelling was replaced by CacheableSchelling here
+    visualization_elements=[
+        get_cache_file_status,
+        canvas_element,
+        get_happy_agents,
+        happy_chart,
+    ],
+    name="Schelling Segregation Model",
+    model_params=model_params,
 )
 
 server.launch()

examples/caching_and_replay/server.py

@@ -28,20 +28,35 @@ def schelling_draw(agent):
     return portrayal
 
 
-canvas_element = mesa.visualization.CanvasGrid(schelling_draw, 20, 20, 500, 500)
+canvas_element = mesa.visualization.CanvasGrid(
+    portrayal_method=schelling_draw,
+    grid_width=20,
+    grid_height=20,
+    canvas_width=500,
+    canvas_height=500,
+)
 happy_chart = mesa.visualization.ChartModule([{"Label": "happy", "Color": "Black"}])
 
 model_params = {
     "height": 20,
     "width": 20,
-    "density": mesa.visualization.Slider("Agent density", 0.6, 0.1, 1.0, 0.1),
-    "minority_pc": mesa.visualization.Slider("Fraction minority", 0.2, 0.00, 1.0, 0.05),
-    "homophily": mesa.visualization.Slider("Homophily", 2, 0, 8, 1),
+    "density": mesa.visualization.Slider(
+        name="Agent density", value=0.8, min_value=0.1, max_value=1.0, step=0.1
+    ),
+    "minority_pc": mesa.visualization.Slider(
+        name="Fraction minority", value=0.2, min_value=0.00, max_value=1.0, step=0.05
+    ),
+    "homophily": mesa.visualization.Slider(
+        name="Homophily", value=3, min_value=0, max_value=8, step=1
+    ),
+    "radius": mesa.visualization.Slider(
+        name="Search Radius", value=1, min_value=1, max_value=5, step=1
+    ),
 }
 
 server = mesa.visualization.ModularServer(
-    Schelling,
-    [canvas_element, get_happy_agents, happy_chart],
-    "Schelling",
-    model_params,
+    model_cls=Schelling,
+    visualization_elements=[canvas_element, get_happy_agents, happy_chart],
+    name="Schelling Segregation Model",
+    model_params=model_params,
 )