嘗試用 Python 寫的病毒擴散模擬程序

病毒擴散仿真程序，用 python 也能夠。

概述

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

由於是 Java 開發的，因此開始我並無多加關注。後來看到有人解析代碼，發現我也能看懂，而後就琢磨用 Python 應該怎麼實現。

Java 版程序淺析

一我的就是 1 個（x, y）座標點，而且每一個人有一個狀態。

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

在每一輪的迭代中，遍歷每一個人，每一個人根據自身的狀態，作出必定的動做，包括：

• 移動
• 狀態變化
• 影響他人

這些動做的具體變動，取決於定義的各類係數。

一輪迭代完成，打印這些點，不一樣的狀態對應不一樣的顏色。

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

Python 版思路

若是咱們想用 Python 實現應該怎麼作呢？

若是徹底復刻 Java 版本，則每次迭代需遍歷全部人，並計算和他人距離，這就是 N^2 次計算。若是是 1000 我的，就須要循環 1 百萬次。這個 Python 的性能確定捉急。

不過 Python 有 numpy ，能夠快速的操做數組。結合 matplotlib 則能夠畫出圖形。

import numpy as np
import matplotlib.pyplot as plt

如何模擬人羣

爲了減小函數之間互相傳參和使用全局變量，咱們也來定義一個類:

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

全部人的座標數據就是 N 行 2 列的數組，同時伴隨必定的狀態：

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

狀態值和計時器也都是數組，同時每次隨機選取指定數量的人感染：

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)

這裏關鍵的一點是，輔助數組的大小和人數保持一致，這樣就能造成一一對應的關係。

狀態發生變化的人才順帶記錄時間：

def random_people_state(self, num, state=1):
        """隨機挑選人設置狀態
        """
        assert self.count > num
        # TODO：極端狀況下會出現無限循環
        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
        # 記錄狀態改變的時間
        self._timer[i] = self._round

經過狀態值，就能夠過濾出人羣，每一個人羣都是 people 的切片視圖。這裏 numpy 的功能至關強大，只須要很是簡潔的語法便可實現:

@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()

順序和開始分析的略有差別，其實並非十分重要，調換它們的順序也是能夠的。

如何改變狀態

這一步就是更新狀態數組 self._status 和 計時器數組 self._timer:

def change_state(self):
        dt = self._round - self._timer
        # 必須先更新時鐘再更新狀態
        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

仍然是經過切片過濾出要更改的目標，而後所有更新。

這裏具體的實現我寫的很是簡單，沒有引入太多的變量：

在必定週期內的 感染者（infected），狀態置爲 確診（confirmed）。 我這裏簡單假設了確診者就被醫院收治，因此失去了繼續感染他人的機會（見下面）。若是要搞複雜點，能夠引入病牀，治癒，死亡等狀態。

如何影響他人

影響別人是整個程序的性能瓶頸，由於須要計算每一個人之間的距離。

這裏繼續作了簡化，只處理感染者：

def infect_possible(self, x=0., safe_distance=3.0):
        """按機率感染接近的健康人
        x 的取值參考正態分佈機率表，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
                # 記錄狀態改變的時間
                self._timer[i] = self._round

能夠看到，距離的計算仍然是經過 numpy 的矩陣操做。可是須要對每個感染者單獨計算，因此若是感染者較多，python 的處理效率感人。

如何移動

_people 是一個座標矩陣，只要生成移動距離矩陣 dt，而後它相加便可。咱們能夠設置一個可移動的範圍 width，把移動距離控制在必定範圍內。

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

這裏還須要增長一個控制移動意向的選項，仍然是利用了正態分佈機率。考慮到這種場景有可能會重用，因此特意把這個方法提取了出來，生成一個只包含 0 1 的數組充當開關。

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

        x 大體取值範圍 -1.99 - 1.99；
        對應正態分佈的機率， 取值 0 的時候對應機率是 50%
        :param x: 控制開關比例
        :return:
        """
        normal = np.random.normal(0, 1, self.count)
        switch = np.where(normal < x, 1, 0)
        return switch

輸出結果

有了一切數據和變化以後，接下來最重要的事情天然就是圖形化顯示結果了。直接使用 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 主要是我的用來練手，目前一些參數沒有徹底抽出來。有須要的只能直接改源碼。

後記

從屢次實驗的結果，經過調整人員的流動意願，流動距離等因素，是能夠獲得直觀的結論的。

本人也是初次使用 numpy 和 matplotlib，現學現賣，如有使用不當之處請指正。其中的機率參數設置 基本沒有科學依據，僅供 Python 愛好者參考。

總得來講，用 numpy 來模擬病毒感染狀況應該是能行得通的。可是其中的影響因子還須要仔細設計。性能也是須要考量的問題。

源碼地址

願疫情能早日過去，武漢加油，中國加油 ?

---

若是本文對你有幫助，請 點贊、分享、關注，謝謝！