腳本之家服務(wù)器常用軟件

快捷導(dǎo)航

Python 寫了個新型冠狀病毒疫情傳播模擬程序

更新時間：2020年02月14日 15:57:00 作者：DavyCloud

這篇文章主要介紹了Python 寫了個新型冠狀病毒疫情傳播模擬程序,本文通過實例代碼給大家介紹的非常詳細(xì)，具有一定的參考借鑒價值，需要的朋友可以參考下

病毒擴(kuò)散仿真程序，用 python 也可以。

概述

事情是這樣的，B 站 UP 主 @ele 實驗室，寫了一個簡單的疫情傳播仿真程序，告訴大家在家待著的重要性，視頻相信大家都看過了，并且 UP 主也放出了源碼。

因為是 Java 開發(fā)的，所以開始我并沒有多加關(guān)注。后來看到有人解析代碼，發(fā)現(xiàn)我也能看懂，然后就琢磨用 Python 應(yīng)該怎么實現(xiàn)。

Java 版程序淺析

一個人就是 1 個（x, y）坐標(biāo)點，并且每個人有一個狀態(tài)。

public class Person extends Point {
  private int state = State.NORMAL;
}

在每一輪的迭代中，遍歷每個人，每個人根據(jù)自身的狀態(tài)，做出一定的動作，包括：

移動
狀態(tài)變化
影響他人

這些動作的具體變更，取決于定義的各種系數(shù)。

一輪迭代完成，打印這些點，不同的狀態(tài)對應(yīng)不同的顏色。

繪圖部分直接使用的 Java 繪圖類 Graphics。

Python 版思路

如果我們想用 Python 實現(xiàn)應(yīng)該怎么做呢？

如果完全復(fù)刻 Java 版本，則每次迭代需遍歷所有人，并計算和他人距離，這就是 N^2 次計算。如果是 1000 個人，就需要循環(huán) 1 百萬次。這個 Python 的性能肯定捉急。

不過 Python 有 numpy ，可以快速的操作數(shù)組。結(jié)合 matplotlib 則可以畫出圖形。

import numpy as np
import matplotlib.pyplot as plt

如何模擬人群

為了減少函數(shù)之間互相傳參和使用全局變量，我們也來定義一個類:

class People(object):
  def __init__(self, count=1000, first_infected_count=3):
    self.count = count
    self.first_infected_count = first_infected_count
    self.init()

所有人的坐標(biāo)數(shù)據(jù)就是 N 行 2 列的數(shù)組，同時伴隨一定的狀態(tài)：

  def init(self):
    self._people = np.random.normal(0, 100, (self.count, 2))
    self.reset()

狀態(tài)值和計時器也都是數(shù)組，同時每次隨機(jī)選取指定數(shù)量的人感染：

def reset(self):
    self._round = 0
    self._status = np.array([0] * self.count)
    self._timer = np.array([0] * self.count)
    self.random_people_state(self.first_infected_count, 1)

這里關(guān)鍵的一點是，輔助數(shù)組的大小和人數(shù)保持一致，這樣就能形成一一對應(yīng)的關(guān)系。

狀態(tài)發(fā)生變化的人才順帶記錄時間：

def random_people_state(self, num, state=1):
    """隨機(jī)挑選人設(shè)置狀態(tài)
    """
    assert self.count > num
    # TODO：極端情況下會出現(xiàn)無限循環(huán)
    n = 0
    while n < num:
      i = np.random.randint(0, self.count)
      if self._status[i] == state:
        continue
      else:
        self.set_state(i, state)
        n += 1

  def set_state(self, i, state):
    self._status[i] = state
    # 記錄狀態(tài)改變的時間
    self._timer[i] = self._round

通過狀態(tài)值，就可以過濾出人群，每個人群都是 people 的切片視圖。這里 numpy 的功能相當(dāng)強(qiáng)大，只需要非常簡潔的語法即可實現(xiàn):

 @property
  def healthy(self):
    return self._people[self._status == 0]

  @property
  def infected(self):
    return self._people[self._status == 1]

按照既定的思路，我們先來定義每輪迭代要做的動作:

  def update(self):
    """每一次迭代更新"""
    self.change_state()
    self.affect()
    self.move()
    self._round += 1
    self.report()

順序和開始分析的略有差異，其實并不是十分重要，調(diào)換它們的順序也是可以的。

如何改變狀態(tài)

這一步就是更新狀態(tài)數(shù)組 self._status 和計時器數(shù)組 self._timer:

  def change_state(self):
    dt = self._round - self._timer
    # 必須先更新時鐘再更新狀態(tài)
    d = np.random.randint(3, 5)
    self._timer[(self._status == 1) & ((dt == d) | (dt > 14))] = self._round
    self._status[(self._status == 1) & ((dt == d) | (dt > 14))] += 1

仍然是通過切片過濾出要更改的目標(biāo)，然后全部更新。

這里具體的實現(xiàn)我寫的非常簡單，沒有引入太多的變量：

在一定周期內(nèi)的感染者（infected），狀態(tài)置為確診（confirmed）。我這里簡單假設(shè)了確診者就被醫(yī)院收治，所以失去了繼續(xù)感染他人的機(jī)會（見下面）。如果要搞復(fù)雜點，可以引入病床，治愈，死亡等狀態(tài)。

如何影響他人

影響別人是整個程序的性能瓶頸，因為需要計算每個人之間的距離。

這里繼續(xù)做了簡化，只處理感染者：

  def infect_possible(self, x=0., safe_distance=3.0):
    """按概率感染接近的健康人
    x 的取值參考正態(tài)分布概率表，x=0 時感染概率是 50%
    """
    for inf in self.infected:
      dm = (self._people - inf) ** 2
      d = dm.sum(axis=1) ** 0.5
      sorted_index = d.argsort()
      for i in sorted_index:
        if d[i] >= safe_distance:
          break # 超出范圍，不用管了
        if self._status[i] > 0:
          continue
        if np.random.normal() > x:
          continue
        self._status[i] = 1
        # 記錄狀態(tài)改變的時間
        self._timer[i] = self._round

可以看到，距離的計算仍然是通過 numpy 的矩陣操作。但是需要對每一個感染者單獨(dú)計算，所以如果感染者較多，python 的處理效率感人。

如何移動

_people 是一個坐標(biāo)矩陣，只要生成移動距離矩陣 dt，然后它相加即可。我們可以設(shè)置一個可移動的范圍 width，把移動距離控制在一定范圍內(nèi)。

  def move(self, width=1, x=.0):
    movement = self.random_movement(width=width)
    # 限定特定狀態(tài)的人員移動
    switch = self.random_switch(x=x)
    movement[switch == 0] = 0
    self._people = self._people + movement

這里還需要增加一個控制移動意向的選項，仍然是利用了正態(tài)分布概率。考慮到這種場景有可能會重用，所以特地把這個方法提取了出來，生成一個只包含 0 1 的數(shù)組充當(dāng)開關(guān)。

  def random_switch(self, x=0.):
    """隨機(jī)生成開關(guān)，0 - 關(guān)，1 - 開

    x 大致取值范圍 -1.99 - 1.99；
    對應(yīng)正態(tài)分布的概率， 取值 0 的時候?qū)?yīng)概率是 50%
    :param x: 控制開關(guān)比例
    :return:
    """
    normal = np.random.normal(0, 1, self.count)
    switch = np.where(normal < x, 1, 0)
    return switch

輸出結(jié)果

有了一切數(shù)據(jù)和變化之后，接下來最重要的事情自然就是圖形化顯示結(jié)果了。直接使用 matplotlib 的散點圖就可以了：

 def report(self):
    plt.cla()
    # plt.grid(False)
    p1 = plt.scatter(self.healthy[:, 0], self.healthy[:, 1], s=1)
    p2 = plt.scatter(self.infected[:, 0], self.infected[:, 1], s=1, c='pink')
    p3 = plt.scatter(self.confirmed[:, 0], self.confirmed[:, 1], s=1, c='red')

    plt.legend([p1, p2, p3], ['healthy', 'infected', 'confirmed'], loc='upper right', scatterpoints=1)
    t = "Round: %s, Healthy: %s, Infected: %s, Confirmed: %s" % \
      (self._round, len(self.healthy), len(self.infected), len(self.confirmed))
    plt.text(-200, 400, t, ha='left', wrap=True)

實際效果
啟動。

if __name__ == '__main__':
  np.random.seed(0)
  plt.figure(figsize=(16, 16), dpi=100)
  plt.ion()
  p = People(5000, 3)
  for i in range(100):
    p.update()
    p.report()
    plt.pause(.1)
  plt.pause(3)

因為這個小 demo 主要是個人用來練手，目前一些參數(shù)沒有完全抽出來。有需要的只能直接改源碼。