Fix caching_and_replay example for Mesa 3.x compatibility

examples/caching_and_replay/README.md

@@ -2,56 +2,70 @@
 
 ## Summary
 
-This example applies caching on the Mesa [Schelling example](https://github.com/mesa/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.
+This example demonstrates how to record and replay Mesa simulations using the [Mesa-Replay](https://github.com/Logende/mesa-replay) library. It wraps the standard [Schelling segregation model](https://github.com/mesa/mesa-examples/tree/main/examples/schelling) with `CacheableModel`, allowing you to:
 
-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`.
-From the user's perspective, the new model behaves the same way as the original Schelling model, but additionally supports caching.
+- **Record** simulations by saving the model state at each step to a cache file
+- **Replay** previously recorded simulations to examine specific runs
 
-Note that the main purpose of this example is to demonstrate that caching and replaying simulation runs is possible.
-The example is designed to be accessible.
-In practice, someone who wants to replay their simulation might not necessarily embed a replay button into the web view, but instead have a dedicated script to run a simulation that is being cached, separate from a script to replay a simulation run from a given cache file.
-More examples of caching and replay can be found in the [Mesa-Replay Repository](https://github.com/Logende/mesa-replay/tree/main/examples).
+The cacheable model behaves identically to a regular Mesa model but with added recording/replay capabilities. More examples can be found in the [Mesa-Replay repository](https://github.com/Logende/mesa-replay/tree/main/examples).
 
 ## Installation
 
-To install the dependencies use pip and the requirements.txt in this directory. e.g.
+Install dependencies:
 
-```
-    $ pip install -r requirements.txt
+```bash
+pip install -r requirements.txt
 ```
 
 ## How to Run
 
-To run the model interactively, run ``mesa runserver`` in this directory. e.g.
+### Interactive Visualization
+
+Run the cacheable version with replay controls:
 
+```bash
+solara run run.py
 ```
-    $ mesa runserver
+
+Or run the standard (non-cacheable) version:
+
+```bash
+solara run server.py
 ```
 
-or
+Then open [http://localhost:8765](http://localhost:8765) in your browser.
 
-Directly run the file ``run.py`` in the terminal. e.g.
+### Recording a Simulation
 
-```
-    $ python run.py
-```
+1. Uncheck the **'Replay cached run?'** checkbox
+2. Adjust the **Cache File Path** if desired (default: `./my_cache_file_path.cache`)
+3. Click **Reset**, then **Run**
+4. The cache is saved automatically when the simulation completes
+
+### Replaying a Simulation
 
-Then open your browser to [http://127.0.0.1:8521/](http://127.0.0.1:8521/) and press Reset, then Run.
+1. Check the **'Replay cached run?'** checkbox
+2. Ensure the cache file path matches your recorded file
+3. Click **Reset**, then **Run**
+4. The simulation will replay the exact recorded sequence
 
-First, run the **simulation** with the 'Replay' switch disabled.
-When the simulation run is finished (e.g. all agents are happy, no more new steps are simulated), the run will automatically be stored in a cache file.
+## Notes
 
-Next, **replay** your latest cached simulation run by enabling the Replay switch and then pressing Reset.
+- The cache file is written when the simulation completes (when `model.running` becomes `False`)
+- During recording, model states are held in memory and written once at the end
+- For large simulations, consider the memory requirements of caching all steps
+- By default, every step is cached. For large runs, you can adjust `cache_step_rate` to cache fewer steps
 
 ## Files
 
-* ``run.py``: Launches a model visualization server and uses `CacheableModelSchelling` as simulation model
-* ``cacheablemodel.py``: Implements `CacheableModelSchelling` to make the original Schelling model cacheable
-* ``model.py``: Taken from the original Mesa Schelling example
-* ``server.py``: Taken from the original Mesa Schelling example
+* `run.py` - Cacheable model visualization with replay controls
+* `cacheablemodel.py` - CacheableSchelling implementation with Mesa 3.x support
+* `model.py` - Standard Schelling segregation model
+* `server.py` - Standard (non-cacheable) visualization
+* `requirements.txt` - Required packages
 
 ## Further Reading
 
 * [Mesa-Replay library](https://github.com/Logende/mesa-replay)
-* [More caching and replay examples](https://github.com/Logende/mesa-replay/tree/main/examples)
+* [Mesa-Replay examples](https://github.com/Logende/mesa-replay/tree/main/examples)
+* [Original Schelling model](https://github.com/mesa/mesa-examples/tree/main/examples/schelling)

examples/caching_and_replay/cacheablemodel.py

@@ -1,17 +1,19 @@
+import random
+from pathlib import Path
+
+import dill
+from mesa.discrete_space import OrthogonalMooreGrid
 from mesa_replay import CacheableModel, CacheState
-from model import Schelling
+from model import Schelling, SchellingAgent
 
 
 class CacheableSchelling(CacheableModel):
-    """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.
+    """Schelling model with caching and replay capabilities.
+
+    This wraps the standard Schelling model with Mesa-Replay's CacheableModel,
+    allowing simulations to be recorded and replayed. The model behaves like
+    a regular Mesa model, but can save its state history to a cache file or
+    replay from a previously saved cache.
     """
 
     def __init__(
@@ -23,10 +25,22 @@ def __init__(
         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,
+        verbose=False,
     ):
+        """Create a new cacheable Schelling model.
+
+        Args:
+            width: Grid width
+            height: Grid height
+            density: Initial population density
+            minority_pc: Fraction of minority agents
+            homophily: Minimum similar neighbors needed to be happy
+            radius: Neighborhood radius for similarity check
+            cache_file_path: Where to save/load the cache file
+            replay: If True, replay from cache; if False, record new simulation
+            verbose: If True, print status messages
+        """
         actual_model = Schelling(
             width=width,
             height=height,
@@ -35,9 +49,184 @@ def __init__(
             homophily=homophily,
             radius=radius,
         )
-        cache_state = CacheState.REPLAY if replay else CacheState.RECORD
+
+        self.verbose = verbose
+
+        cache_path = Path(cache_file_path)
+        effective_replay = replay and cache_path.exists() and cache_path.is_file()
+
+        if replay and not effective_replay:
+            print(
+                f"Cache file not found at {cache_file_path}. "
+                "Running in record mode instead."
+            )
+
+        cache_state = CacheState.REPLAY if effective_replay else CacheState.RECORD
+
+        if verbose:
+            mode = "REPLAY" if cache_state == CacheState.REPLAY else "RECORD"
+            print(f"Initializing in {mode} mode")
+            print(f"Cache file: {cache_file_path}")
+
         super().__init__(
             model=actual_model,
             cache_file_path=cache_file_path,
             cache_state=cache_state,
         )
+
+    def step(self):
+        """Run one model step. Automatically records or replays based on mode."""
+        super().step()
+
+        # Write cache after each step in record mode
+        if self._cache_state == CacheState.RECORD:
+            self._write_cache_file()
+
+        if self.verbose:
+            mode = "REPLAY" if self._cache_state == CacheState.REPLAY else "RECORD"
+            print(f"Step {self.step_count} ({mode})")
+
+    def __setattr__(self, key, value):
+        """Handle attribute setting. Finalizes cache when model stops running."""
+        if key == "running":
+            was_running = (
+                getattr(self.model, "running", True) if hasattr(self, "model") else True
+            )
+            if hasattr(self, "model"):
+                self.model.__setattr__(key, value)
+            else:
+                super().__setattr__(key, value)
+
+            if (
+                hasattr(self, "_cache_state")
+                and was_running
+                and not value
+                and self._cache_state == CacheState.RECORD
+                and not self.run_finished
+            ):
+                if self.verbose:
+                    print(f"Finalizing cache at step {self.step_count}...")
+                self.finish_run()
+        else:
+            super().__setattr__(key, value)
+
+    def _serialize_state(self) -> bytes:
+        """Serialize model state for caching.
+
+        Custom implementation needed for Mesa 3.x CellAgent positioning.
+        """
+        state_dict = self.model.__dict__.copy()
+
+        # Save random number generator state
+        state_dict["random_state"] = self.model.random.getstate()
+        state_dict.pop("random", None)
+
+        # Serialize agents with cell coordinates (not pos, which is None for CellAgents)
+        agents_data = []
+        for agent in self.model.agents:
+            coord = (
+                agent.cell.coordinate if hasattr(agent, "cell") and agent.cell else None
+            )
+            agents_data.append(
+                {
+                    "unique_id": agent.unique_id,
+                    "type": agent.type,
+                    "coord": coord,
+                }
+            )
+        state_dict["agents_data"] = agents_data
+        state_dict.pop("agents", None)
+        state_dict.pop("_agents", None)
+        state_dict.pop("_agents_by_type", None)
+
+        # Save grid contents
+        grid_data = {}
+        for coord, cell in self.model.grid._cells.items():
+            grid_data[coord] = [a.unique_id for a in cell.agents]
+        state_dict["grid_data"] = grid_data
+        state_dict.pop("grid", None)
+
+        # Save data collector state
+        if hasattr(self.model, "datacollector"):
+            dc = self.model.datacollector
+            state_dict["datacollector_data"] = {
+                "model_vars": dc.model_vars.copy(),
+                "agent_records": dc._agent_records.copy()
+                if hasattr(dc, "_agent_records")
+                else {},
+            }
+
+        # Save agent ID counter
+        state_dict["_next_id"] = (
+            self.model.agent_id_counter
+            if hasattr(self.model, "agent_id_counter")
+            else getattr(self.model, "_next_id", 0)
+        )
+
+        return dill.dumps(state_dict)
+
+    def _deserialize_state(self, state: bytes) -> None:
+        """Restore model state from cache.
+
+        Custom implementation needed for Mesa 3.x agent management.
+        """
+        if self.verbose:
+            print(f"Restoring state at step {self.step_count}")
+        state_dict = dill.loads(state)  # noqa: S301
+
+        # Restore basic attributes first
+        for k, v in state_dict.items():
+            if k not in [
+                "random_state",
+                "agents_data",
+                "grid_data",
+                "datacollector_data",
+                "_next_id",
+            ]:
+                setattr(self.model, k, v)
+
+        # Random
+        self.model.random = random.Random()
+        self.model.random.setstate(state_dict["random_state"])
+
+        # Grid
+        self.model.grid = OrthogonalMooreGrid(
+            (self.model.width, self.model.height), torus=True, random=self.model.random
+        )
+
+        # Clear existing agents to avoid duplicates
+        self.model._agents.clear()
+        if hasattr(self.model, "_agents_by_type"):
+            self.model._agents_by_type.clear()
+        if hasattr(self.model, "_all_agents"):
+            self.model._all_agents.clear()
+
+        # Restore agents from cached data
+        agent_map = {}
+        agents_to_restore = state_dict.get("agents_data", [])
+
+        for a_data in agents_to_restore:
+            agent = SchellingAgent(self.model, a_data["type"])
+            agent.unique_id = a_data["unique_id"]
+            coord = a_data.get("coord")
+
+            # Place agent in its cell
+            if coord is not None and coord in self.model.grid._cells:
+                cell = self.model.grid._cells[coord]
+                agent.cell = cell
+
+            agent_map[a_data["unique_id"]] = agent
+
+        # Restore data collector state
+        if "datacollector_data" in state_dict and hasattr(self.model, "datacollector"):
+            dc_data = state_dict["datacollector_data"]
+            self.model.datacollector.model_vars = dc_data["model_vars"].copy()
+            if "agent_records" in dc_data and hasattr(
+                self.model.datacollector, "_agent_records"
+            ):
+                self.model.datacollector._agent_records = dc_data[
+                    "agent_records"
+                ].copy()
+
+        # Restore agent ID counter
+        self.model.agent_id_counter = state_dict.get("_next_id", 0)

examples/caching_and_replay/model.py

@@ -1,4 +1,6 @@
-"""This file was copied over from the original Schelling mesa example."""
+"""Schelling segregation model."""
+
+from contextlib import suppress
 
 import mesa
 from mesa.discrete_space import CellAgent, OrthogonalMooreGrid
@@ -25,7 +27,9 @@ def step(self):
 
         # If unhappy, move:
         if similar < self.model.homophily:
-            self.cell = self.model.grid.select_random_empty_cell()
+            with suppress(IndexError):
+                # Try to move; if no empty cells, agent stays in place
+                self.cell = self.model.grid.select_random_empty_cell()
         else:
             self.model.happy += 1
 
@@ -69,9 +73,7 @@ def __init__(
         )
 
         # Set up agents
-        # We use a grid iterator that returns
-        # the coordinates of a cell as well as
-        # its contents. (coord_iter)
+        # We use a grid iterator that returns the coordinates of a cell as well as its contents. (coord_iter)
         for cell in self.grid.all_cells:
             if self.random.random() < self.density:
                 agent_type = 1 if self.random.random() < self.minority_pc else 0