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!