Python学习笔记#9：病毒传播模拟程序

概述

本章借助matplotlib包，模拟疫情传播，有两种思路：

1. 使用面向对象的方式，每个人为一个独立对象
2. 使用矩阵的方式

方法1. 面向对象方式

面向对象的方式，每个人为一个独立对象，算法：

1. 遍历所有对象：
   1. 更新当前对象位置（移动）
   2. 计算当前对象到其他所有人的相对距离
   3. 根据距离，感染率等因素，改变对象邻居当前的状态。

算法逻辑非常简单，但是可以看到会嵌套多次循环，时间复杂度为：$O(n^2)$

1. 导入包

import numpy as np
import matplotlib.pyplot as plt

2. 全局参数

# 地图宽度
width = 100

# 总人口
pop = 2000

# 初始病人数量
n = 10

# 感染半径
sd = 10

# 感染几率 50%
sr = 0.5

3. 定义人

1. 坐标：随机分布
2. 颜色：绿色代表健康，红色代表感染
3. 运动函数：人随机运动

# 人
class People(object):
    def __init__(self):
    	# 随机坐标
        self.x = np.random.rand() * width
        self.y = np.random.rand() * width
        self.loc = np.array([self.x, self.y])
        self.color = 'g'

    # 随机运动
    def move(self):
        self.x += np.random.randn()
        self.y += np.random.randn()

4. 环境

1. 按照人口生成不同的人口对象
2. 随机选择病源

# 人群
all_people = np.array([People() for i in range(pop)])

# 初始化患病人群
sick_people = np.random.choice(all_people, size=n, replace=False)
for p in sick_people:
    p.color = "r"

5. 病毒感染

1. 遍历所有人群，区分患者与普通人
2. 遍历患者：
   1. 遍历所有人，计算与当前患者的距离
      1. 当目标与患者距离小于传播半径，则有一定几率感染。
3. 返回患病熟料

# 病毒感染函数
def affect(all_people):
    sick_people = []
    healthy_people = []
    n = 0
    for p in all_people:
        if p.color == "r":
            sick_people.append(p)
            n += 1
        if p.color == "g":
            healthy_people.append(p)

    for sp in sick_people:
        for hp in healthy_people:
            dist = np.linalg.norm(sp.loc - hp.loc)
            rad = np.random.rand()
            if hp.color == "g" and dist <= sd and rad < sr:
                hp.color = "r"
                n += 1
    return n

6. 动态显示

使用matplotlib互交模式，动态显示

plt.ion()

# 模拟
while n < pop:
    plt.clf()
    update(all_people)
    n = affect(all_people)
    plt.scatter([p.x for p in all_people], [p.y for p in all_people], c=[p.color for p in all_people], s=3)
    plt.axis([0, width, 0, width])
    plt.pause(0.5)
    print("总人数：{}，传染人数：{}".format(pop, n))

plt.ioff()
plt.show()

7. 结果

总人数：1000，传染人数：150
总人数：1000，传染人数：585
总人数：1000，传染人数：875
总人数：1000，传染人数：984
总人数：1000，传染人数：1000

方法2. 矩阵运算方式

借助numpy，可以非常快速的处理矩阵运算，特点：

1. 速度快
2. 处理数据量大
3. 比较抽象

时间复杂度为：$O(n)$

1. 代码

import numpy as np
import matplotlib.pyplot as plt

class VirusSimluator(object):
    def __init__(self):

        # 地图宽度
        self.width = 1000

        # 总人口
        self.pop = 2800

        # 初始病人数量
        self.first_patients = 10

        # 感染半径
        self.infection_radius = 5

        # 感染几率 10%
        self.infection_potential = 0.5

        # 人群：横坐标、纵坐标
        self.locations = np.random.rand(self.pop, 2) * self.width

        # 状态：（g-正常、r-感染）
        self.status = np.array(["g"] * self.pop)

        # 初始化感染群
        self.initPatients()

    # 更新人群位置
    # 限制人群移动范围在 [0, self.width, 0, self.width]
    def move(self):
        self.locations += np.random.randn(self.pop, 2) * self.width
        # 所有超过边界的人，需要边界反弹
        # 左下角两条边
        n = self.locations[self.locations < 0].size
        self.locations[self.locations < 0] = np.random.rand(n) * self.width
        # 右上角两条边
        n = self.locations[self.locations > self.width].size
        self.locations[self.locations > self.width] = self.width - np.random.rand(n) * self.width

    # 初始化患病人群
    def initPatients(self):
        self.status[np.random.choice(self.pop, size=self.first_patients, replace=False)] = "r"

    # 统计感染人群
    @property
    def patients(self):
        return self.locations[self.status == "r"]

    # 统计感染人数
    @property
    def patients_num(self):
        return self.status[self.status == "r"].size

    # 传染函数
    def affect(self):
        for ip in self.patients:
            # 人与人的距离矩阵
            d = np.sqrt(np.sum(np.asarray(ip - self.locations)**2, axis=1))
            # 传染几率的矩阵
            p = np.random.rand(len(d))
            # 既小于传染距离又达到传染几率的人群
            self.status[(d < self.infection_radius) & (p < self.infection_potential)] = "r"

    # 统计人数
    def reportStatistic(self, n):
        current_patient_num = self.patients_num
        print("第【{}】轮感染，总人数：{}，传染人数：{}".format(n, self.pop, current_patient_num))
        n += 1
        return current_patient_num, n

    # 显示函数
    def display(self):
        k, n = self.reportStatistic(1)
        plt.ion()
        fig = plt.figure(figsize=(10, 5))

        ax1 = fig.add_subplot(1, 2, 1)
        ax2 = fig.add_subplot(1, 2, 2)

        x = np.array(n)
        y = np.array(k)

        while k < self.pop:
            ax1.cla()
            # 画离散点
            ax1.scatter(self.locations[:, 0], self.locations[:, 1], c=self.status, s=1)
            # 限制图像范围
            ax1.axis([0, self.width, 0, self.width])
            self.move()
            self.affect()
            k, n = self.reportStatistic(n)

            # 统计曲线
            ax2.cla()
            x = np.append(x, n)
            y = np.append(y, k)
            ax2.plot(x, y)
            ax2.axis([0, 100, 0, self.pop])

            plt.pause(0.1)

        plt.ioff()
        plt.show()

if __name__ == "__main__":
    vs = VirusSimluator()
    vs.display()

2. 重难点

运动

• 简单的生成一个随机向量矩阵，大小和原来的坐标矩阵一致，然后相加，即可达到随机运动的效果。
• 所有超过边界的人，需要边界反弹，限制其移动范围。

    # 更新人群位置
    # 限制人群移动范围在 [0, self.width, 0, self.width]
    def move(self):
        self.locations += np.random.randn(self.pop, 2) * self.width
        # 所有超过边界的人，需要边界反弹
        # 左下角两条边
        n1 = self.locations[self.locations < 0].size
        self.locations[self.locations < 0] = np.random.randn(n1) * self.width
        # 右上角两条边
        n2 = self.locations[self.locations > self.width].size
        self.locations[self.locations > self.width] = self.width - np.random.rand(n2) * self.width

装饰器
Python内置的@property装饰器就是负责把一个方法变成属性调用

    # 统计感染人群
    @property
    def patients(self):
        return self.locations[self.status == "r"]

    # 统计感染人数
    @property
    def patients_num(self):
        return self.status[self.status == "r"].size

感染函数

1. 得到所有的感染者坐标
2. 遍历感染者：
   1. 计算感染者与所有人的距离矩阵d
   2. 计算传染概率矩阵p
   3. 将所有符合条件的邻居标记为感染

    # 传染函数
    def affect(self):
        for ip in self.patients:
            # 人与人的距离矩阵
            d = np.sqrt(np.sum(np.asarray(ip - self.locations)**2, axis=1))
            # 传染几率的矩阵
            p = np.random.rand(len(d))
            # 既小于传染距离又达到传染几率的人群
            self.status[(d < self.infection_radius) & (p < self.infection_potential)] = "r"

3. 单点到多点的距离

距离算法：

d = np.sqrt(np.sum(np.asarray(ip - self.locations)**2, axis=1))

实际上可以分解为

l = np.asarray(ip - self.locations)

其中ip为单一的点，self.locations为数组，结果也为数组，大小和self.locations一样。

更进一步，得到每个点的差的乘积

l = np.asarray(ip - self.locations)**2

乘积相加，但是这样只有一个数字，

s = np.sum(np.array(ip - self.locations)**2)

想要得到矩阵，需要控制参数axis

# axis=1，表示沿着x轴相加，横向
l = np.sum(np.array(p - locations)**2,axis=1)

# axis=1，表示沿着y轴相加，竖向
l = np.sum(np.array(p - locations)**2,axis=0)

最后再对每一个元素开方，得到结果

4. 结果