Advanced Computational Coursework.py

from random import random
import matplotlib.pyplot as pyplot
import matplotlib as plt
from matplotlib import animation
#import time
#import numba

#Simulation Variables
EMPTY = 0
TREE = 1
BURNING = 2
probTree = 0.8
probBurning = 0.01
probImmune = 0.3
probLightning = 0.001
GridSizes = [100,400,800,1000,1200,2000]
timeStep =500
SequentialTimeList = []
ParallelTimeList = []
barWidth = 0.25

#Initialization of the the initial forest grid
def InitializeForest(n,probTree,probBurning):
    forestGrid = []
    for x in range(n):
        rowList = []
        for y in range(n):
            if (random()<probTree):
                if(random()<probBurning):
                    cell = BURNING
                else:
                    cell = TREE
            else:
                cell = EMPTY
            rowList.append(cell)
        forestGrid.append(rowList)    
    return forestGrid

#Parallel copy of InitializeForest Function
#@numba.jit(nopython=False, parallel=True)
def InitializeForestParallel(n,probTree,probBurning):   
    forestGrid = []
    for x in range(n):
        rowList = []
        for y in range(n):
            if (random()<probTree):
                if(random()<probBurning):
                    cell = BURNING
                else:
                    cell = TREE
            else:
                cell = EMPTY
            rowList.append(cell)
        forestGrid.append(rowList)    
    return forestGrid

def ReflectiveBoundaryCondition(forestGrid, n):
    ExtendedGrid = []
    GridNS = forestGrid  
    GridNorth = forestGrid[0]
    GridSouth = forestGrid[n-1] 
    GridNS.insert(0,GridNorth)
    GridNS.insert(n,GridSouth)
    for x in range(n+2):
        rowList = []
        for y in range(n+2):
            if (y == 0):
                cell = GridNS[x][0]
            elif(y == n+1):
                cell = GridNS[x][y-2]
            else:
                cell = GridNS[x][y-1]
            rowList.append(cell)
        ExtendedGrid.append(rowList)
    return ExtendedGrid

#@numba.jit(nopython=False, parallel=True)
def ReflectiveBoundaryConditionParallel(forestGrid, n):
    ExtendedGrid = []
    GridNS = forestGrid  
    GridNorth = forestGrid[0]
    GridSouth = forestGrid[n-1] 
    GridNS.insert(0,GridNorth)
    GridNS.insert(n,GridSouth)
    for x in range(n+2):
        rowList = []
        for y in range(n+2):
            if (y == 0):
                cell = GridNS[x][0]
            elif(y == n+1):
                cell = GridNS[x][y-2]
            else:
                cell = GridNS[x][y-1]
            rowList.append(cell)
        ExtendedGrid.append(rowList)
    return ExtendedGrid

def MooreNeighborhood(ExtForestGrid,x,y):
    if (ExtForestGrid[x][y+1]==2):
        BURNING = True
    elif (ExtForestGrid[x][y-1]==2):
        BURNING = True
    elif (ExtForestGrid[x+1][y]==2):
        BURNING = True
    elif (ExtForestGrid[x-1][y]==2):
        BURNING = True
    elif (ExtForestGrid[x-1][y-1]==2):
        BURNING = True
    elif (ExtForestGrid[x+1][y-1]==2):
        BURNING = True
    elif (ExtForestGrid[x-1][y+1]==2):
        BURNING = True
    elif (ExtForestGrid[x+1][y+1]==2):
        BURNING = True
    else:
        BURNING = False
    return BURNING

#@numba.jit(nopython=False, parallel=True)
def MooreNeighborhoodParallel(ExtForestGrid,x,y):
    if (ExtForestGrid[x][y+1]==2):
        BURNING = True
    elif (ExtForestGrid[x][y-1]==2):
        BURNING = True
    elif (ExtForestGrid[x+1][y]==2):
        BURNING = True
    elif (ExtForestGrid[x-1][y]==2):
        BURNING = True
    elif (ExtForestGrid[x-1][y-1]==2):
        BURNING = True
    elif (ExtForestGrid[x+1][y-1]==2):
        BURNING = True
    elif (ExtForestGrid[x-1][y+1]==2):
        BURNING = True
    elif (ExtForestGrid[x+1][y+1]==2):
        BURNING = True
    else:
        BURNING = False
    return BURNING

#This is the spread function that determines the state of the individual cell at the next iteration.
def Spread(ExtForestGrid,x,y,probImmune,probLightning):
    if (ExtForestGrid[x][y] == 0):
        cell = EMPTY
    elif (ExtForestGrid[x][y] == 1):
        if(random()<probImmune):
            cell = TREE
        elif (MooreNeighborhood(ExtForestGrid,x,y) or random()<probLightning):
            cell = BURNING
        else:
            cell = TREE
    else:
        cell = EMPTY
    return cell

#Parallel copy of FireSpread Function
#@numba.jit(nopython=False, parallel=True)
def SpreadParallel(ExtForestGrid,x,y,probImmune,probLightning):
    if (ExtForestGrid[x][y] == 0):
        cell = EMPTY
    elif (ExtForestGrid[x][y] == 1):
        if(random()<probImmune):
            cell = TREE
        elif (MooreNeighborhoodParallel(ExtForestGrid,x,y) or random()<probLightning):
            cell = BURNING
        else:
            cell = TREE
    else:
        cell = EMPTY
    return cell

#This function applies the function FireSpread to the inner cells in the extended grid.
def ApplySpread(ExtGrid,probImmune,probLightning):
    newGrid = []
    n = len(ExtGrid)
    for x in range(1,n-1):
        rowList = []
        for y in range(1,n-1):
            cell = Spread(ExtGrid,x,y,probImmune,probLightning)
            rowList.append(cell)
        newGrid.append(rowList.copy())
    return newGrid

#Parallel copy of ApplyFireSpread Function
#@numba.jit(nopython=False, parallel=True)
def ApplySpreadParallel(ExtGrid,probImmune,probLightning):
    newGrid = []
    n = len(ExtGrid)
    for x in range(1,n-1):
        rowList = []
        for y in range(1,n-1):
            cell = SpreadParallel(ExtGrid,x,y,probImmune,probLightning)
            rowList.append(cell)
        newGrid.append(rowList.copy())
    return newGrid

#Function to run simulation
def Simulation(n,probTree,probBurning,probLightning,probImmune,timeStep):     
    initialGrid = InitializeForest(n,probTree,probBurning)
    itemCount = 0
    gridList = []
    gridList.append(initialGrid)
    for t in range(timeStep):   
        ExtendedGrid = ReflectiveBoundaryCondition(gridList[itemCount],n)
        newGrid = ApplySpread(ExtendedGrid,probImmune,probLightning)  
        gridList.append(newGrid)
        itemCount = itemCount + 1   
    return gridList

#Parallel copy of RunSimulation Function
#@numba.jit(nopython=False, parallel=True)
def SimulationParallel(n,probTree,probBurning,probLightning,probImmune,timeStep):     
    initialGrid = InitializeForestParallel(n,probTree,probBurning)
    itemCount = 0
    gridList = []
    gridList.append(initialGrid)
    for t in range(timeStep):   
        ExtendedGrid = ReflectiveBoundaryConditionParallel(gridList[itemCount],n)
        newGrid = ApplySpreadParallel(ExtendedGrid,probImmune,probLightning)  
        gridList.append(newGrid)
        itemCount = itemCount + 1   
    return gridList

result = Simulation(500, probTree, probBurning, probLightning, probImmune, timeStep)

cmap = plt.colors.ListedColormap(['yellow','green','red'])
bounds=[-0.5,0.5,1.5,2.5]
norm = plt.colors.BoundaryNorm(bounds, cmap.N)


fig = pyplot.figure(figsize=(25/3, 6.25))
ax = fig.add_subplot(111)
ax.set_axis_off()
im = ax.imshow(result[0], cmap=cmap, norm=norm)#, interpolation='nearest')

def animate(i):
    im.set_data(result[i])
    

anim = animation.FuncAnimation(fig, animate, interval=200, frames=len(result))
pyplot.show()