Refactor forest_fire example to store environment state in a NumPy array

[ShreyasN707]

This PR refactors the forest_fire example to store the environment state using a NumPy array instead of representing each grid cell as a stateful patch agent.

Previously, each grid cell was implemented as a TreeCell agent that stored its condition (Fine, On Fire, Burned Out). In this refactor, the cell state is stored in a 2D NumPy array (fire_state), while TreeCell agents are kept only as lightweight wrappers for visualization and grid positioning.

The goal is to simplify the model logic and avoid creating thousands of objects when the environment state can be represented more efficiently.

Motivation

In the original implementation, each grid cell was represented by a TreeCell agent whose primary purpose was to hold the state of the environment. For large grids this means creating thousands of agents just to store simple state values.

For example, the previous implementation stored state directly on the agent:

class TreeCell(FixedAgent):
    def __init__(self, model, cell):
        super().__init__(model)
        self.condition = "Fine"

While this works, the agents themselves do not make decisions — they only store environment data.

This PR moves the environment state into a NumPy grid, which:

• avoids unnecessary agent objects

• simplifies the simulation logic

• makes the example clearer for new users

aligns the example with patterns where environment state is separate from agents

Implementation

The main change is the introduction of a NumPy-based state grid inside the model.

Environment state

A NumPy array now stores the fire state of each grid cell:

FINE = 0
ON_FIRE = 1
BURNED_OUT = 2
EMPTY = -1

self.fire_state = np.full((width, height), EMPTY, dtype=np.int8)

This array replaces the previous TreeCell.condition attribute.

Fire propagation logic

The fire spread logic is now handled in model.step() using the array instead of agent state:

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

Using a copied array ensures that updates occur synchronously, preserving the behavior of the original cellular automaton.

Data collection

The DataCollector was updated to read directly from the NumPy grid:

"Fine": lambda m: np.sum(m.fire_state == FINE),
"On Fire": lambda m: np.sum(m.fire_state == ON_FIRE),
"Burned Out": lambda m: np.sum(m.fire_state == BURNED_OUT),

Visualization

TreeCell agents are still created so the existing visualization system can iterate over cells. However, they now act only as wrappers for location, while the actual state is read from the NumPy array:

x, y = tree.cell.coordinate
state = tree.model.fire_state[x, y]

portrayal["color"] = COLORS_BY_STATE[state]

Result

The simulation behavior remains identical to the original example, but the implementation becomes:

• simpler

• more explicit about environment state

• less dependent on large numbers of state-holding agents

The example still runs using:

solara run examples/forest_fire/app.py

and the visualization, fire propagation, and density behavior remain unchanged.

Additional Notes

TreeCell agents are retained only for visualization compatibility.

The refactor focuses solely on internal model structure and does not change the user-facing interface.

The example was tested manually to confirm that fire propagation and visualization behave the same as before.

Does it belong?

• No significant overlap with an existing example, or if there is, the overlap is justified
• The model is well-scoped: the simplest model that demonstrates its idea
• It showcases Mesa features not already well-covered by other examples
• It showcases interesting ABM mechanics, dynamics, or phenomena

Is it correct and current?

• Uses current Mesa APIs (no deprecated schedulers, manual unique_id, dict portrayals, etc.)
• Runs and visualizes out of the box
• Deterministic with a fixed rng seed

Is it clean?

• Minimal: no dead code, unused imports, or unnecessary complexity
• Clear naming; logic readable without deep domain knowledge
• README explains what it does, what it demonstrates, and how to run it (concise, links to references for theory)
• [] PR follows the template, commits are clean, and there are no unrelated changes

[ShreyasN707]

@EwoutH @quaquel is there any improvements that can be made in this PR?

[EwoutH]

Thanks for the PR, looks like a cool model.

Could you:

• Checkout the test failure
• Request a peer-review (and maybe do one or two yourself)

[ShreyasN707]

@EwoutH I’ve self-reviewed the changes in this PR. It completes the refactor to a NumPy array for this example and does not change the existing model logic.

• Refactor preserves original model behavior (no logic change)
• Environment state moved from agents to model (NumPy array)
• Clear separation between agent representation and environment state
• Uses current Mesa patterns and APIs
• No unnecessary complexity introduced
• No dead code or unused variables
• Changes are focused and limited to the refactor