稍微改了一下例子里面生命游戏的代码做了一个最简单的传染病模型。黄色是易感者，红色是感染者，蓝色是痊愈者。
"""
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()
