尝试用 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 来模拟病毒感染状况应该是能行得通的。可是其中的影响因子还须要仔细设计。性能也是须要考量的问题。

源码地址

愿疫情能早日过去，武汉加油，中国加油 ?

---

若是本文对你有帮助，请 点赞、分享、关注，谢谢！