cellular_automata_epidemics.jl

using StatsBase
using Plots

nt = 500
nx = 200
ny = 200

# state
x = zeros(Int8,nx,ny,nt) 

p_infection = zeros(Float16,nx,ny)
p_infection[:,:] .= 0.5
p_infection[:,begin:2:end] .= 0.025

p_removal = 0.2

# initial infection
x[45:55,45:55,1] .= 1

for t in 1:nt-1
     for j in 2:ny-1
        for i in 2:nx-1
            
            # propagate current state to next time step
            x[i,j,t+1] = x[i,j,t]
            
             # S -> I
            if x[i,j,t] == 0
                for ni in -1:1 
                    for nj in -1:1
                        if x[i+ni,j+nj,t] == 1 && x[i,j,t+1] == 0
                            x[i,j,t+1] = sample([0, 1], Weights([1-p_infection[i,j], p_infection[i,j]]))
                        end
                    end 
                end
            end
            
            # I -> R
            if x[i,j,t] == 1
                x[i,j,t+1] = sample([1, 2], Weights([1-p_removal, p_removal]))
            end
        end
    end
end

@gif for t ∈ 1:nt
    heatmap(x[:,:,t], clims=(0, 2), color=cgrad(:PiYG_4, rev=true))
end