Skip to content

Refactor forest_fire example to store environment state in a NumPy array#375

Open
ShreyasN707 wants to merge 3 commits intomesa:mainfrom
ShreyasN707:forest-fire-example-refactor
Open

Refactor forest_fire example to store environment state in a NumPy array#375
ShreyasN707 wants to merge 3 commits intomesa:mainfrom
ShreyasN707:forest-fire-example-refactor

Conversation

@ShreyasN707
Copy link
Contributor

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.

Screenshot from 2026-03-11 23-36-00 Screenshot from 2026-03-11 23-36-38 Screenshot from 2026-03-11 23-37-05

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.

@ShreyasN707
Copy link
Contributor Author

ShreyasN707 commented Mar 12, 2026

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant