/
model.py
72 lines (60 loc) · 2.13 KB
/
model.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
import random
from mesa import Model
from mesa.datacollection import DataCollector
from mesa.space import Grid
from mesa.time import RandomActivation
from .agent import TreeCell
class ForestFire(Model):
"""
Simple Forest Fire model.
"""
def __init__(self, height=100, width=100, density=0.65):
"""
Create a new forest fire model.
Args:
height, width: The size of the grid to model
density: What fraction of grid cells have a tree in them.
"""
# Initialize model parameters
self.height = height
self.width = width
self.density = density
# Set up model objects
self.schedule = RandomActivation(self)
self.grid = Grid(height, width, torus=False)
self.datacollector = 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")})
# Place a tree in each cell with Prob = density
for (contents, x, y) in self.grid.coord_iter():
if random.random() < self.density:
# Create a tree
new_tree = TreeCell((x, y), self)
# Set all trees in the first column on fire.
if x == 0:
new_tree.condition = "On Fire"
self.grid._place_agent((x, y), new_tree)
self.schedule.add(new_tree)
self.running = True
self.datacollector.collect(self)
def step(self):
"""
Advance the model by one step.
"""
self.schedule.step()
# collect data
self.datacollector.collect(self)
# Halt if no more fire
if self.count_type(self, "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.
"""
count = 0
for tree in model.schedule.agents:
if tree.condition == tree_condition:
count += 1
return count