#SIS_ICount.py: Contains functions for displaying plot of the number of times each cell has been infected

#import modules
from numpy import *
from pylab import *
from matplotlib.axes3d import *

#Keep track of what cells are infected intially
def InitICount(CellNum,state):
	ITimes = zeros((CellNum,CellNum),float)
	for i in range(CellNum):
		for j in range(CellNum):
			if state[i][j] == 1:
				ITimes[i][j] = 1
	return ITimes

#if a cell is infected at this update, then add one to ITimes.
def ICountUpdate(CellNum,ITimes,state):
	for i in range(CellNum):
		for j in range(CellNum):
			if state[i][j] == 1:
				ITimes[i][j] += 1
	return ITimes

#plot the number of times a cell has been infected in grayscale
def SIS_Plot(ITimes):
	pcolormesh(ITimes)
	gray()
	return show()