当前位置:首页 >> 脚本专栏

用Python制作在地图上模拟瘟疫扩散的Gif图

受杰森的《Almost Looks Like Work》启发,我来展示一些病毒传播模型。需要注意的是这个模型并不反映现实情况,因此不要误以为是西非可怕的传染病。相反,它更应该被看做是某种虚构的僵尸爆发现象。那么,让我们进入主题。

用Python制作在地图上模拟瘟疫扩散的Gif图

这就是SIR模型,其中字母S、I和R反映的是在僵尸疫情中,个体可能处于的不同状态。

  •     S 代表易感群体,即健康个体中潜在的可能转变的数量。
  •     I 代表染病群体,即僵尸数量。
  •     R 代表移除量,即因死亡而退出游戏的僵尸数量,或者感染后又转回人类的数量。但对与僵尸不存在治愈者,所以我们就不要自我愚弄了(如果要把SIR模型应用到流感传染中,还是有治愈者的)。
  • 至于β(beta)和γ(gamma):
  •     β(beta)表示疾病的传染性程度,只要被咬就会感染。
  •     γ(gamma)表示从僵尸走向死亡的速率,取决于僵尸猎人的平均工作速率,当然,这不是一个完美的模型,请对我保持耐心。
  • S′="htmlcode">
    import numpy as np
    import math
    import matplotlib.pyplot as plt  
    %matplotlib inline
    from matplotlib import rcParams
    import matplotlib.image as mpimg
    rcParams['font.family'] = 'serif'
    rcParams['font.size'] = 16
    rcParams['figure.figsize'] = 12, 8
    from PIL import Image
    

    适当的beta和gamma值就能够摧毁大半江山
     

    beta = 0.010
    gamma = 1
    

    还记得导数的定义么?当导数已知,假设Δt很小的情况下,经过重新整理,它可以用来近似预测函数的下一个取值,我们已经声明过u′(t)。

    用Python制作在地图上模拟瘟疫扩散的Gif图

    初始化一些东东。
     

    import numpy as np
    import math
    import matplotlib.pyplot as plt  
    %matplotlib inline
    from matplotlib import rcParams
    import matplotlib.image as mpimg
    rcParams['font.family'] = 'serif'
    rcParams['font.size'] = 16
    rcParams['figure.figsize'] = 12, 8
    from PIL import Image
    

    适当的beta和gamma值就能够摧毁大半江山
     

    beta = 0.010
    gamma = 1
    

    还记得导数的定义么?当导数已知,假设Δt很小的情况下,经过重新整理,它可以用来近似预测函数的下一个取值,我们已经声明过u′(t)。

    用Python制作在地图上模拟瘟疫扩散的Gif图

    这种方法叫做欧拉法,代码如下:
     

    def euler_step(u, f, dt):
      return u + dt * f(u)
    

    我们需要函数f(u)。友好的numpy提供了简洁的数组操作。我可能会在另一篇文章中回顾它,因为它们太强大了,需要更多的解释,但现在这样就能达到效果:

     
    def f(u):
      S = u[0]
      I = u[1]
      R = u[2]
     
      new = np.array([-beta*(S[1:-1, 1:-1]*I[1:-1, 1:-1] +
                  S[0:-2, 1:-1]*I[0:-2, 1:-1] +
                  S[2:, 1:-1]*I[2:, 1:-1] +
                  S[1:-1, 0:-2]*I[1:-1, 0:-2] +
                  S[1:-1, 2:]*I[1:-1, 2:]),
               beta*(S[1:-1, 1:-1]*I[1:-1, 1:-1] +
                  S[0:-2, 1:-1]*I[0:-2, 1:-1] +
                  S[2:, 1:-1]*I[2:, 1:-1] +
                  S[1:-1, 0:-2]*I[1:-1, 0:-2] +
                  S[1:-1, 2:]*I[1:-1, 2:]) - gamma*I[1:-1, 1:-1],
               gamma*I[1:-1, 1:-1]
              ])
     
      padding = np.zeros_like(u)
      padding[:,1:-1,1:-1] = new
      padding[0][padding[0] < 0] = 0
      padding[0][padding[0] > 255] = 255
      padding[1][padding[1] < 0] = 0
      padding[1][padding[1] > 255] = 255
      padding[2][padding[2] < 0] = 0
      padding[2][padding[2] > 255] = 255
     
      return padding
    

    导入北欧国家的人口密度图并进行下采样,以便较快地得到结果
     

    from PIL import Image
    img = Image.open('popdens2.png')
    img = img.resize((img.size[0]/2,img.size[1]/2))
    img = 255 - np.asarray(img)
    imgplot = plt.imshow(img)
    imgplot.set_interpolation('nearest')
    
    

    用Python制作在地图上模拟瘟疫扩散的Gif图

    北欧国家的人口密度图(未包含丹麦)

    S矩阵,也就是易感个体,应该近似于人口密度。感染者初始值是0,我们把斯德哥尔摩作为第一感染源。
     

    S_0 = img[:,:,1]
    I_0 = np.zeros_like(S_0)
    I_0[309,170] = 1 # patient zero
    

    因为还没人死亡,所以把矩阵也置为0.
     

    R_0 = np.zeros_like(S_0)
    

    接着初始化模拟时长等。
     

    T = 900             # final time
    dt = 1             # time increment
    N = int(T/dt) + 1        # number of time-steps
    t = np.linspace(0.0, T, N)   # time discretization
     
    # initialize the array containing the solution for each time-step
    u = np.empty((N, 3, S_0.shape[0], S_0.shape[1]))
    u[0][0] = S_0
    u[0][1] = I_0
    u[0][2] = R_0
    

    我们需要自定义一个颜色表,这样才能将感染矩阵显示在地图上。
     

    import matplotlib.cm as cm
    theCM = cm.get_cmap("Reds")
    theCM._init()
    alphas = np.abs(np.linspace(0, 1, theCM.N))
    theCM._lut[:-3,-1] = alphas
    

    下面坐下来欣赏吧…

     
    for n in range(N-1):
      u[n+1] = euler_step(u[n], f, dt)
    

    让我们再做一下图像渲染,把它做成gif,每个人都喜欢gifs!
     

    from images2gif import writeGif
     
    keyFrames = []
    frames = 60.0
     
    for i in range(0, N-1, int(N/frames)):
      imgplot = plt.imshow(img, vmin=0, vmax=255)
      imgplot.set_interpolation("nearest")
      imgplot = plt.imshow(u[i][1], vmin=0, cmap=theCM)
      imgplot.set_interpolation("nearest")
      filename = "outbreak" + str(i) + ".png"
      plt.savefig(filename)
      keyFrames.append(filename)
     
    images = [Image.open(fn) for fn in keyFrames]
    gifFilename = "outbreak.gif"
    writeGif(gifFilename, images, duration=0.3)
    plt.clf()