Python 写了个新型冠状病毒疫情传播模拟程序

2020-02-14 18:01:25 来源:易采站长站 作者:王旭

病毒扩散仿真程序,用 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
            




              

微信扫一扫

易采站长站微信账号