Refactor forest_fire example to store environment state in a NumPy array

examples/forest_fire/app.py

@@ -9,23 +9,30 @@
 )
 
 COLORS = {"Fine": "#00AA00", "On Fire": "#880000", "Burned Out": "#000000"}
+from forest_fire.model import BURNED_OUT, EMPTY, FINE, ON_FIRE
+
+COLORS_BY_STATE = {
+    EMPTY: "#FFFFFF",
+    FINE: COLORS["Fine"],
+    ON_FIRE: COLORS["On Fire"],
+    BURNED_OUT: COLORS["Burned Out"],
+}
 
 
 def forest_fire_portrayal(tree):
     if tree is None:
         return
+
     portrayal = {"Shape": "rect", "w": 1, "h": 1, "Filled": "true", "Layer": 0}
-    (x, y) = (tree.cell.coordinate[i] for i in (0, 1))
+
+    x, y = tree.cell.coordinate
+    state = tree.model.fire_state[x, y]
+
     portrayal["x"] = x
     portrayal["y"] = y
-    portrayal["color"] = COLORS[tree.condition]
-    return portrayal
+    portrayal["color"] = COLORS_BY_STATE[state]
 
-
-def post_process_space(ax):
-    ax.set_aspect("equal")
-    ax.set_xticks([])
-    ax.set_yticks([])
+    return portrayal
 
 
 def post_process_lines(ax):
@@ -35,7 +42,6 @@ def post_process_lines(ax):
 space_component = make_space_component(
     forest_fire_portrayal,
     draw_grid=False,
-    post_process=post_process_space,
 )
 lineplot_component = make_plot_component(
     COLORS,

examples/forest_fire/forest_fire/agent.py

@@ -2,27 +2,8 @@
 
 
 class TreeCell(FixedAgent):
-    """A tree cell.
-
-    Attributes:
-        condition: Can be "Fine", "On Fire", or "Burned Out"
-
-    """
+    """Tree cell used for visualization and grid coordinates."""
 
     def __init__(self, model, cell):
-        """Create a new tree.
-
-        Args:
-            model: standard model reference for agent.
-        """
         super().__init__(model)
-        self.condition = "Fine"
         self.cell = cell
-
-    def step(self):
-        """If the tree is on fire, spread it to fine trees nearby."""
-        if self.condition == "On Fire":
-            for neighbor in self.cell.neighborhood.agents:
-                if neighbor.condition == "Fine":
-                    neighbor.condition = "On Fire"
-            self.condition = "Burned Out"

examples/forest_fire/forest_fire/model.py

@@ -1,8 +1,18 @@
 import mesa
+import numpy as np
 from mesa.discrete_space import OrthogonalMooreGrid
 
 from .agent import TreeCell
 
+# Environment state stored as NumPy array instead of patch agents.
+# This simplifies the model and avoids creating thousands of objects
+# when cells only store state.
+
+FINE = 0
+ON_FIRE = 1
+BURNED_OUT = 2
+EMPTY = -1
+
 
 class ForestFire(mesa.Model):
     """Simple Forest Fire model."""
@@ -17,39 +27,56 @@ def __init__(self, width=100, height=100, density=0.65, rng=None):
         super().__init__(rng=rng)
 
         # Set up model objects
-
         self.grid = OrthogonalMooreGrid((width, height), capacity=1, random=self.random)
+
+        # Environment state stored in NumPy array instead of patch agents.
+        # TreeCell agents are kept only as lightweight wrappers for visualization.
+        # State array
+        self.fire_state = np.full((width, height), EMPTY, dtype=np.int8)
+
         self.datacollector = mesa.DataCollector(
             {
-                "Fine": lambda m: self.count_type(m, "Fine"),
-                "On Fire": lambda m: self.count_type(m, "On Fire"),
-                "Burned Out": lambda m: self.count_type(m, "Burned Out"),
+                "Fine": lambda m: int(np.sum(m.fire_state == FINE)),
+                "On Fire": lambda m: int(np.sum(m.fire_state == ON_FIRE)),
+                "Burned Out": lambda m: int(np.sum(m.fire_state == BURNED_OUT)),
             }
         )
 
         # Place a tree in each cell with Prob = density
         for cell in self.grid.all_cells:
+            # Create a tree wrapper
+            tree = TreeCell(self, cell)
+            # self.grid.place_agent(tree, cell)  # Not needed/supported in OrthogonalMooreGrid
+
+            x, y = cell.coordinate
             if self.random.random() < density:
-                # Create a tree
-                new_tree = TreeCell(self, cell)
+                self.fire_state[x, y] = FINE
                 # Set all trees in the first column on fire.
-                if cell.coordinate[0] == 0:
-                    new_tree.condition = "On Fire"
+                if x == 0:
+                    self.fire_state[x, y] = ON_FIRE
 
         self.running = True
         self.datacollector.collect(self)
 
     def step(self):
         """Advance the model by one step."""
-        self.agents.shuffle_do("step")
+
+        new_fire_state = self.fire_state.copy()
+
+        for cell in self.grid.all_cells:
+            x, y = cell.coordinate
+            if self.fire_state[x, y] == ON_FIRE:
+                for neighbor in cell.neighborhood:
+                    nx, ny = neighbor.coordinate
+                    if self.fire_state[nx, ny] == FINE:
+                        new_fire_state[nx, ny] = ON_FIRE
+                new_fire_state[x, y] = BURNED_OUT
+
+        self.fire_state = new_fire_state
+
         # collect data
         self.datacollector.collect(self)
 
         # Halt if no more fire
-        if self.count_type(self, "On Fire") == 0:
+        if np.sum(self.fire_state == ON_FIRE) == 0:
             self.running = False
-
-    @staticmethod
-    def count_type(model, tree_condition):
-        """Helper method to count trees in a given condition in a given model."""
-        return len(model.agents.select(lambda x: x.condition == tree_condition))