# Homework0: SIR传染病模型

``````"""
Epidemic model
Susceptible 0 Yellow
Infectious 1 Red
Recoverd 2 Blue
"""

import taichi as ti
import numpy as np

ti.init(arch=ti.gpu)

n = 200
beta = 0.04  # rate from S -> I
gamma = 0.01  # rate from I -> R
I0 = 0.01  # initial infectious rate
x = ti.var(dt=ti.i32, shape=(n, n))
c = ti.var(dt=ti.i32, shape=(n, n))

celsize = 4
res = n * celsize
img = ti.Vector(3, dt=ti.i32, shape=(res, res))
color = ti.Vector(3, dt=ti.i32, shape=3)
color[0] = [255, 255, 200]
color[1] = [250, 100, 100]
color[2] = [100, 100, 250]

@ti.func
def count_infected_neighbors(i, j):
return (
x[(i - 1) % n, j] % 2
+ x[(i + 1) % n, j] % 2
+ x[i, (j - 1) % n] % 2
+ x[i, (j + 1) % n] % 2
+ x[(i - 1) % n, (j - 1) % n] % 2
+ x[(i + 1) % n, (j - 1) % n] % 2
+ x[(i - 1) % n, (j + 1) % n] % 2
+ x[(i + 1) % n, (j + 1) % n] % 2
)

@ti.kernel
def run():
for i, j in x:
if x[i, j] == 0:
c[i, j] = count_infected_neighbors(i, j)
for i, j in x:
if x[i, j] == 0:
if ti.random() < c[i, j] / 8 * beta:
x[i, j] = 1
elif x[i, j] == 1:
if ti.random() < gamma:
x[i, j] = 2

@ti.kernel
def init():
for i, j in x:
if ti.random() < I0:
x[i, j] = 1
else:
x[i, j] = 0

@ti.kernel
def render():
for u, v in img:
img[u, v] = color[x[u // celsize, v // celsize]]

init()
gui = ti.GUI("Epidemic SIR model", (res, res))
render()
for t in range(1000):
run()
render()
gui.set_image(img.to_numpy().astype(np.uint8))
gui.show()
``````

5 Likes

Cool!