Zombie simulation.py

import numpy as np
import random
import matplotlib.pyplot as plt

##==============================================================================

def SI_simulation(tmax, zom = 0.4, a = 0.2, d = 0.1, 
                  S0 = 998, Z0 = 2, D0 = 0, da = 0, dz = 0, dd = 0):
    '''
    tmax: [Only required variable] How many days the simulation is run for
    zom: initial chance a zombiie will infect a human
    a: chance a human will succefully permanently kill a zombie
    d: chance a zombie completely decays
    S0: Initial suscceptable population
    D0: Initial dead population
    Z0: Initial zombie population
    da: rate at which a changes
    dz: rate at which zom changes
    dd: rate at which d changes
    '''
    Nt = tmax + 1
    t = np.linspace(0,tmax,Nt)
    
    (S, Z, D) = (np.zeros(Nt, dtype = 'int'), 
                 np.zeros(Nt, dtype = 'int'), 
                 np.zeros(Nt, dtype = 'int'))
    
    S[0], Z[0], D[0] = S0, Z0, D0 # Initial variables
    
    for i in range(1, Nt):
        S[i] = S[i-1]
        Z[i] = Z[i-1]
        D[i] = D[i-1]

        sp = random.randint(0, S[i-1])
        zp = random.randint(0, Z[i-1])
        ref  = min(sp,zp)
        
        if i > 5:
            '''
            Everyday after day 5, A bunch of zombies will decay completely and
            permaanently die
            '''
            for j in range(Z[i-1]):
                p = random.random()
                if p <= d:
                    Z[i] -= 1
                    D[i] += 1
        
        for j in range(ref):
            '''
            Everyday, a couple of survivors will interact with a couple of zombies
            annd they will fight.
            '''
            p = random.random()
            q = random.random()
            if p <= zom:
                S[i] -= 1
                Z[i] += 1
            if q <= a:
                Z[i] -= 1
                D[i] += 1
        
        '''
        Everyday after the first five days, the survivors
        get better at killing zombies and avoiding death. The zommbies also decay
        faster
        '''
        if all([i>= 5, a < 0.75, d < 0.10, zom > 0.20]): #Limits how muuch the variables change
            a += da
            zom -= dz
            d += dd
            
        
    return t, S, Z, D

##==============================================================================

tmax = 250 # How many days the simualtions are run for 
nsims = 20 # How many simulations are run

Z0 = 1 
S0 = 1_000 - Z0
D0 = 0

# Important: S0 + Z0 + D0 should *always* be constant

Zrate = 0.60
Attack = 0.15
Decay = 0.001

dZrate = 0.001
dAttack = 0.001
dDecay = 0.0001

##==============================================================================

#Used for testing indivdual simulations.

'''data = SI_simulation(tmax, zom = Zrate, a = Attack, d = Decay, 
                         dz = dZrate, da = dAttack, dd = dDecay)

plt.plot(t, data[1], lw = 3, c = 'b', label = 'Susceptable')
plt.plot(t, data[2], lw = 3, c = 'r', label = 'Zombies')
plt.plot(t, data[3], lw = 3, c = 'k', label = 'Dead')
plt.plot(t, data[3]+data[2]+data[1], lw = 3, c = 'g', label = 'total')'''


##==============================================================================

(Ssims, Zsims, Dsims) = (np.zeros((tmax+1, nsims)), np.zeros((tmax+1, nsims)), 
                         np.zeros((tmax+1, nsims)))

# Runs the simmuation Nsims amount of times and plots alll of them

for i in range(nsims):
    data = SI_simulation(tmax, zom = Zrate, a = Attack, d = Decay, 
                         dz = dZrate, da = dAttack, dd = dDecay)
    Ssims[:,i] = data[1]
    Zsims[:,i] = data[2]
    Dsims[:,i] = data[3]

S = np.transpose(Ssims)
Z = np.transpose(Zsims)
D = np.transpose(Dsims)
t = np.linspace(0, tmax, tmax+1)

width = 0.2

for i in range(nsims):
    plt.plot(t, S[i], c = 'b', lw=width)
    plt.plot(t, Z[i], c = 'r', lw=width)
    plt.plot(t, D[i], c = 'k', lw=width)
    #plt.plot(t, (D[i]+Z[i]+S[i]), c = 'g', lw=width)

def average_out(X):
    ans = []
    for i in X:
        ans.append(np.mean(i))
    return np.array(ans)

# The averages of the simulations

S1, Z1, D1 = average_out(Ssims), average_out(Zsims), average_out(Dsims)

lwidth = 3

plt.plot(t, S1, lw = lwidth, c = 'b', label = 'Susceptable')
plt.plot(t, Z1, lw = lwidth, c = 'r', label = 'Zombies')
plt.plot(t, D1, lw = lwidth, c = 'k', label = 'Dead')
#plt.plot(t, (S1 + Z1 + D1), c = 'g', lw = lwidth)

##==============================================================================

plt.legend()
plt.grid()

if dZrate+dAttack+dDecay == 0:
    title = f'Test Zombie apocalypse model with {nsims} simulations. \n\
first {tmax} days [static variables]'
else:
    title = f'Test Zombie apocalypse model with {nsims} simulations. \n\
first {tmax} days [dynamic variables]'

plt.title(title)

s = f'Time [Days] \n {Zrate*100}% initial chance of human getting infected on interaction. \n\
{Attack * 100}% initial chance of a human killing a zombie on interaction \n\
{Decay*100}% initial chance of a zombie decaying completely (on a given day)'

plt.xlabel(s)

plt.ylabel('Number of People')

plt.show()