import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import matplotlib.colors as cc
from matplotlib.colors import ListedColormap
# creates the four colors for our simulation, EMPTY is white,
# TREE is green, FIRE is red and BURNT is black
colors = cc.ColorConverter.colors
cols = [colors['w'],
colors['g'],
colors['r'],
colors['k']]
ff1 = ListedColormap(cols)
# TODO revise this method to include a boolean parameter called torus
# when true, the function should wrap around the edges, when false
# return the function as written.
def von_neuman_neighbors(i, j, size):
n = []
if i > 0:
n.append((i - 1, j))
if j > 0:
n.append((i, j - 1))
if i < size - 1:
n.append((i + 1, j))
if j < size - 1:
n.append((i, j + 1))
return n
# TODO write this method to include all eight neighbors. Use the
# torus parameter as described above
def moore_neighbors(i, j, size, torus):
pass
class Forest():
states = {"EMPTY":0, "TREE":1, "FIRE":2, "BURNT":3}
def __init__(self, size):
self.size = size
self.x = np.arange(size)
self.y = np.arange(size)
self.cells = np.zeros((size, size), dtype="int")
def set_fire(self, locs):
for p in locs:
self.cells[p[0] + self.size // 2, p[1] + self.size // 2] = Forest.states["FIRE"]
def random_setup(self, d):
for i in range(self.size):
for j in range(self.size):
if (np.random.random() < d):
self.cells[i, j] = Forest.states["TREE"]
else:
self.cells[i, j] = Forest.states["EMPTY"]
def image_setup(self):
plt.title("Forest Fire Simulation")
self.plt = plt.imshow(self.cells, interpolation='nearest',
origin='bottom',
vmin=Forest.states["EMPTY"],
vmax=Forest.states["BURNT"],
cmap=ff1)
def update(self):
newcells = np.zeros((self.size, self.size), dtype="int")
for i in range(self.size):
for j in range(self.size):
neighbors = von_neuman_neighbors(i, j, self.size)
# TODO Add in code to make the fire spread. If a tree is near
# a fire cell, it catches. All trees on fire burn out completely in
# one timestep.
if self.cells[i,j] == Forest.states["FIRE"]:
newcells[i,j] = self.cells[i,j]
elif self.cells[i,j] == Forest.states["TREE"]:
newcells[i,j] = self.cells[i,j]
elif self.cells[i,j] == Forest.states["BURNT"]:
newcells[i,j] = self.cells[i,j]
elif self.cells[i,j] == Forest.states["EMPTY"]:
newcells[i,j] = self.cells[i,j]
# TODO Revise to have burnt trees turn into empty cells randomly
# TODO Revise to have empty cells grow new trees randomly if they
# they are adjacent to a living tree
# TODO Add in a probability of a lightning strike each round. If it hits a
# tree, the tree is set on fire.
self.cells = newcells
def plot(self):
self.plt.set_data(self.cells)
return self.plt
size = 60
density = 0.35
fig, ax = plt.subplots()
ax.set_ylim(-1, size)
ax.set_xlim(-1, size)
middle = ((0, 0), )
ff = Forest(size)
ff.random_setup(density)
ff.set_fire(middle)
################
# ANIMATION
ff.image_setup()
def update(data):
ff.update()
return ff.plot(),
def data_gen():
while True: yield 1
ani = animation.FuncAnimation(fig, update, data_gen, blit=False, interval=50)
plt.show()