Virus diffusion emulator, also available in Python.

An overview of the

The thing is, B station UP @ele lab wrote a simple epidemic transmission simulation program, telling everyone the importance of staying at home, I believe you have seen the video, and UP also released the source code.

Because it was developed in Java, I didn’t pay much attention at first. Then I saw someone parse the code, realized THAT I could understand it, and wondered how to do it in Python.

Java version of the program analysis

A person is one (x, y) coordinate point, and each person has a state.

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

In each iteration, everyone is traversed, and everyone makes certain actions according to their own state, including:

  • mobile
  • Change of state
  • Influence others

The specific changes to these actions depend on the various coefficients defined.

An iteration is completed, and the dots are printed, with different states corresponding to different colors.

The drawing part uses the Java drawing class Graphics directly.

Python version of thinking

What should we do if we want to implement it in Python?

If the Java version was completely replicated, each iteration would iterate over everyone and calculate the distance to everyone else, which is N^2 calculations. If it’s 1,000 people, it’s a million cycles. The Python performance is definitely in a hurry.

But Python has Numpy, which allows you to quickly manipulate arrays. With matplotlib you can draw a graph.

import numpy as np
import matplotlib.pyplot as plt
Copy the code

How to simulate a crowd

To reduce functions passing arguments to each other and using global variables, let’s also define a class:

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

Everyone’s coordinate data is an array of N rows and 2 columns, accompanied by certain states:

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

Status values and timers are also arrays, and a specified number of people are randomly selected at a time:

    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)
Copy the code

The key point here is that the size of the auxiliary array is the same as the number of people, so that there is a one-to-one correspondence.

Record the time of the talent whose status has changed:

    def random_people_state(self, num, state=1):
        """Randomly selected person setting status"""
        assert self.count > num
        # TODO: An infinite loop can occur in extreme cases
        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
        Record the time when the state changes
        self._timer[i] = self._round
Copy the code

Through the status value, the crowd can be filtered, and each crowd is a slice view of people. Here numpy is quite powerful and requires only a very compact syntax:

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

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

Following the established thinking, we first define the actions to be done in each iteration:

    def update(self):
        """Update every iteration"""
        self.change_state()
        self.affect()
        self.move()
        self._round += 1
        self.report()
Copy the code

The slight difference between the order and the initial analysis is not really important, and it is ok to reverse the order.

How to change the state

This step is to update the status array self._status and the timer array self._timer:

    def change_state(self):
        dt = self._round - self._timer
        The clock must be updated before the status is updated
        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
Copy the code

Again, the target to change is filtered through slices and then all updated.

The implementation here is very simple, without introducing too many variables:

Infected patients within a specified period of time will become confirmed. I have simply assumed that those who are diagnosed are admitted to the hospital, thus losing the chance of further infection (see below). If you want to complicate things, you can introduce sickbeds, cures, death, etc.

How to Influence Others

Affecting others is a performance bottleneck for the entire program because you need to calculate the distance between everyone.

Here we continue to simplify and only deal with infected people:

Def infect_possible (self, x = 0., safe_distance = 3.0) :"""If x=0, the infection probability is 50%."""
        for inf inSelf. 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  # Out of range, forget it
                if self._status[i] > 0:
                    continue
                if np.random.normal() > x:
                    continue
                self._status[i] = 1
                Record the time when the state changes
                self._timer[i] = self._round
Copy the code

As you can see, the distance is still computed by numpy’s matrix operation. However, each infected person needs to be counted individually, so if there are many infected people, Python’s processing efficiency is impressive.

How to move

The _people is a coordinate matrix, you just generate the moving distance matrix dt, and then you add them. We can set a moveable range width to control the moving distance within a certain range.

    def move(self, width=1, x=.0):
        movement = self.random_movement(width=width)
        # Restrict the movement of people in a particular state
        switch = self.random_switch(x=x)
        movement[switch == 0] = 0
        self._people = self._people + movement
Copy the code

There is also an option to control the intention to move, again using normally distributed probability. Given the potential for reuse in this scenario, this method was specifically extracted to generate an array of 0’s and 1’s to act as switches.

    def random_switch(self, x=0.):
        """Randomly generated switch, 0-off, 1-on x approximate value range -1.99-1.99; Param x: control switch ratio :return:"""
        normal = np.random.normal(0, 1, self.count)
        switch = np.where(normal < x, 1, 0)
        return switch
Copy the code

The output

With all the data and changes, the next most important thing is to graphically display the results. Just use matplotlib’s scatter diagram:

    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)
Copy the code

The actual effect

Start.

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)
Copy the code

Because this small demo is mainly for personal practice, some parameters have not been completely extracted. There is a need to directly change the source code.

Afterword.

From the results of many experiments, intuitive conclusions can be obtained by adjusting factors such as personnel’s willingness to flow and flow distance.

I am also the first time to use numpy and matplotlib, now learn now sell, if there is any improper use please correct. The probability parameter Settings are largely unscientific and are for Python enthusiasts only.

In general, using NUMpy to simulate a virus infection should work. But the impact factors need to be carefully designed. Performance is also an issue to consider.

The source address

May the epidemic pass as soon as possible. Come on, Wuhan, China 💪

If this article has helped you, pleasegive a like,share,Focus on, thank you!