python绘制动态图，模拟病毒传播

前几天在b站看到一个模拟病毒传播的视频，当时想着说python能不能做，今天闲着没事，试了一下，还可以。直接上代码

import numpy as np
from matplotlib import pyplot as plt
from matplotlib import animation
from sklearn.neighbors import NearestNeighbors

n=10#初始病人数量
n1=1#每变化一次，每个病人传播人数

class point():#点类，用来模拟人
    def __init__(self):
        self.x=np.random.normal()
        self.y=np.random.normal()
        self.color='g'

    def get_loc(self):
        return self.x,self.y

    def get_color(self):
        return self.color

    def set_color(self,color):
        self.color=color

    def move(self):
        random=np.random.random()
        if random<0.25:
            self.x+=0.1
        elif 0.25<=random<0.5:
            self.x-=0.1
        elif 0.5<=random<0.75:
            self.y+=0.1
        else:
            self.y-=0.1

l=[]
for i in range(500):#人数
    s=point()
    l.append(s)

l1=[]
for i in np.random.random(n):#选出n个病人，改变颜色
    l[int(len(l)*i)].set_color('r')
    l1.append(l[int(len(l)*i)])

fig = plt.figure()

scat = plt.scatter([x.get_loc()[0] for x in l],[x.get_loc()[1] for x in l],c=[x.get_color() for x in l])

def animate(i):
    print(i)
    for j in l:
        j.move()

    #传播机制，每次从距离最近的五个人中传播给一个人
    l2=[]
    for j in set(l1):
        nnarray = NearestNeighbors(n_neighbors=5).fit(np.array([x.get_loc() for x in l])).kneighbors(np.array(j.get_loc()).reshape(1, -1), return_distance=False)[0]
        nn=nnarray[int(np.random.random()*len(nnarray))]
        l[nn].set_color('r')
        l2.append(l[nn])
    for j in l2:
        l1.append(j)
    scat.set_offsets([[x.get_loc()[0],x.get_loc()[1]] for x in l])
    scat.set_color([x.get_color() for x in l])
    return scat,

anim = animation.FuncAnimation(fig, animate, frames=100, interval=60, blit=False)
# plt.show()
#图片保存，保存格式为.gif文件，需要使用一个额外的软件imagemagick，7.0以后的版本不用配置，可以直接使用
anim.save('1.gif',writer='imagemagick')

做的比较简单，自己也是小白，感觉python还是好强大。效果是这样的，红的是病人，绿的是健康人。