├── README.md └── SEIR.py /README.md: -------------------------------------------------------------------------------- 1 | # CA_SEIR 2 | 元胞自动机模拟病毒传染（SEIR模型）可视化 3 | #### 功能描述 4 | 设置一定大小的人数、四种人群（四种状态）、传染概率、潜伏时间、治愈时间、免疫时间，即可对其进行模拟也可以更改规则和人数、人群，实现自己想要的模拟 5 | 6 | #### 环境配置&需要的包 7 | * Python 3.x+Anaconda+Pycharm 8 | * ramdom、numpy、sys 9 | * pygame 10 | * matplotlib 11 | 12 | #### 如何更改参数或制定自己想要的规则 13 | - 1.更改传染病模型人群种类（状态）代码2-6行注释 14 | - 2.更改规则，对每一个类别的人群，确定相应规则 代码9-12行，98-128行 15 | - 3.更改人群总数 代码226行（后两个参数的乘积）+246行（退出循环条件） 16 | - 4.更改概率，代码20-26行 17 | 18 | #### 相关示例 19 | 20 | 21 | 22 | ![fig.1](https://github.com/Windxy/CA_SEIR/blob/image/1.png) 23 | ![fig.2](https://github.com/Windxy/CA_SEIR/blob/image/2.png) 24 | ![fig.3](https://github.com/Windxy/CA_SEIR/blob/image/3.png) 25 | ![fig.4](https://github.com/Windxy/CA_SEIR/blob/image/4.png) 26 | ![fig.5](https://github.com/Windxy/CA_SEIR/blob/image/5.png) 27 | -------------------------------------------------------------------------------- /SEIR.py: -------------------------------------------------------------------------------- 1 | # coding:utf-8 2 | # 状态： 3 | # S0表示易感者S 4 | # S1表示感染者I 5 | # S2表示治愈者R 6 | # S3表示潜伏者E 7 | 8 | # 规则： 9 | # 1. 当S=S0时，为易感者，被感染率为p=k*(上下左右邻居的感染者)/4+l*(左上左下右上右下邻居的感染者)/4,概率c1变为潜伏者,概率1-c1变为感染者 10 | # 2. 当S=S1时，为感染者，具有一定的感染能力，等待t_max时间，会自动变为S2治愈者，染病时间初始化为0 11 | # 3. 当S=S2时，为治愈者，[具有一定的免疫能力，等待T_max时间，会自动变为S0易感者，免疫时间初始化为0],改为永久免疫 12 | # 4. 当S=S3时,为潜伏者,等待t1_max时间(潜伏时间),会自动变为S1感染者,潜伏时间初始化为0, 13 | from pylab import * 14 | import random 15 | import numpy as np 16 | import pygame 17 | import sys 18 | import matplotlib.pyplot as plt 19 | 20 | # 初始化相关数据 21 | k=0.85 #上下左右概率 22 | l=0.55 #其它概率 23 | c1=0.4 #c1概率变为潜伏者 24 | t1_max=15 #潜伏时间 25 | t_max=50 #治愈时间 26 | T_max=75 #免疫时间 #暂时不用 27 | 28 | # p = np.array([0.6, 0.4]) 29 | # probability = np.random.choice([True, False], 30 | # p=p.ravel()) # p(感染or潜伏)=(self.s1_0*k+self.s1_1*l) p(易感)=1-(self.s1_0*k+self.s1_1*l) 31 | 32 | def probablity_fun(): 33 | np.random.seed(0) 34 | # p = np.array([(self.s1_0*k+self.s1_1*l),1-(self.s1_0*k+self.s1_1*l)]) 35 | p = np.array([0.6, 0.4]) 36 | probability = np.random.choice([True, False], 37 | p=p.ravel()) # p(感染or潜伏)=(self.s1_0*k+self.s1_1*l) p(易感)=1-(self.s1_0*k+self.s1_1*l) 38 | return probability 39 | 40 | RED = (255, 0, 0) 41 | GREY = (127, 127, 127) 42 | Green = (0, 255, 0) 43 | BLACK = (0, 0, 0) 44 | 45 | """细胞类，单个细胞""" 46 | class Cell: 47 | # 初始化 48 | stage = 0 49 | 50 | def __init__(self, ix, iy, stage): 51 | self.ix = ix 52 | self.iy = iy 53 | self.stage = stage #状态，初始化默认为0，易感者 54 | # self.neighbour_count = 0 #周围细胞的数量 55 | self.s1_0 = 0 #上下左右为感染者的数量 56 | self.s1_1 = 0 #左上左下右上右下感染者的数量 57 | self.T_ = 0 #免疫时间 58 | self.t_ = 0 #患病时间 59 | self.t1_ = 0 #潜伏时间 60 | 61 | # 计算周围有多少个感染者 62 | def calc_neighbour_count(self): 63 | count_0 = 0 64 | count_1 = 0 65 | pre_x = self.ix - 1 if self.ix > 0 else 0 66 | for i in range(pre_x, self.ix+1+1): 67 | pre_y = self.iy - 1 if self.iy > 0 else 0 68 | for j in range(pre_y, self.iy+1+1): 69 | if i == self.ix and j == self.iy: # 判断是否为自身 70 | continue 71 | if self.invalidate(i, j): # 判断是否越界 72 | continue 73 | if CellGrid.cells[i][j].stage == 1 or CellGrid.cells[i][j] == 3 : #此时这个邻居是感染者 74 | #如果是在上下左右 75 | if (i==self.ix and j==self.iy-1) or \ 76 | (i==self.ix and j==self.iy+1) or \ 77 | (i==self.ix-1 and j==self.iy) or \ 78 | (i==self.ix+1 and j==self.iy): 79 | count_0+=1 80 | else: 81 | count_1+=1 82 | # print(count_0) 83 | self.s1_0 = count_0 84 | # if self.s1_1!=0: 85 | # print(count_1,count_0,self.ix,self.iy) 86 | self.s1_1 = count_1 87 | 88 | # 判断是否越界 89 | def invalidate(self, x, y): 90 | if x >= CellGrid.cx or y >= CellGrid.cy: 91 | return True 92 | if x < 0 or y < 0: 93 | return True 94 | return False 95 | 96 | # 定义规则 97 | def next_iter(self): 98 | # 规则1，易感者 99 | if self.stage==0: 100 | probability=random.random()#生成0到1的随机数 101 | s1_01 = self.s1_0 * k + self.s1_1 * l 102 | 103 | if (s1_01>probability) and (s1_01!=0): 104 | p1 = random.random() 105 | if p1>c1: 106 | self.stage=1 107 | else: 108 | self.stage=3 109 | else: 110 | self.stage = 0 111 | # 规则2，感染者 112 | elif self.stage == 1: 113 | if self.t_ >= t_max: 114 | self.stage = 2 115 | else: 116 | self.t_ = self.t_ + 1 117 | # 规则3，治愈者(永久免疫规则) 118 | elif self.stage == 2: 119 | if self.T_ >= T_max: 120 | self.stage = 0 121 | else: 122 | self.T_ = self.T_ + 1 123 | # 规则4,潜伏者 124 | elif self.stage == 3: 125 | if self.t1_ >= t1_max: 126 | self.stage = 1 # 转变为感染者 127 | else: 128 | self.t1_ += 1 129 | 130 | """细胞网格类，处在一个长cx,宽cy的网格中""" 131 | class CellGrid: 132 | 133 | cells = [] 134 | cx = 0 135 | cy = 0 136 | 137 | # 初始化 138 | def __init__(self, cx, cy): 139 | CellGrid.cx = cx 140 | CellGrid.cy = cy 141 | for i in range(cx): 142 | cell_list = [] 143 | for j in range(cy): 144 | cell = Cell(i, j, 0) #首先默认为全是易感者 145 | if (i == cx/2 and j ==cy/2) or (i==cx/2+1 and j==cy/2) or (i==cx/2+1 and j==cy/2+1):#看26行就可以了 146 | cell_list.append(Cell(i,j,1)) 147 | else: 148 | cell_list.append(cell) 149 | CellGrid.cells.append(cell_list) 150 | 151 | def next_iter(self): 152 | for cell_list in CellGrid.cells: 153 | for item in cell_list: 154 | item.next_iter() 155 | 156 | def calc_neighbour_count(self): 157 | for cell_list in CellGrid.cells: 158 | for item in cell_list: 159 | item.calc_neighbour_count() 160 | 161 | 162 | def num_of_nonstage(self): 163 | # global count0_,count1_,count2_ 164 | count0 = 0 165 | count1 = 0 166 | count2 = 0 167 | count3 = 0 168 | for i in range(self.cx): 169 | for j in range(self.cy): 170 | # 计算全部的方格数 171 | cell = self.cells[i][j].stage 172 | if cell == 0: 173 | count0 += 1 174 | elif cell == 1: 175 | count1 += 1 176 | elif cell == 2: 177 | count2 += 1 178 | elif cell == 3: 179 | count3 += 1 180 | return count0, count1, count2, count3 181 | 182 | '''界面类''' 183 | class Game: 184 | screen = None 185 | count0 = 0 186 | count1 = 9 187 | count2 = 0 188 | count3 = 0 189 | def __init__(self, width, height, cx, cy):#屏幕宽高，细胞生活区域空间大小 190 | self.width = width 191 | self.height = height 192 | self.cx_rate = int(width / cx) 193 | self.cy_rate = int(height / cy) 194 | self.screen = pygame.display.set_mode([width, height])# 195 | self.cells = CellGrid(cx, cy) 196 | 197 | def show_life(self): 198 | for cell_list in self.cells.cells: 199 | for item in cell_list: 200 | x = item.ix 201 | y = item.iy 202 | if item.stage == 0: 203 | pygame.draw.rect(self.screen, GREY, 204 | [x * self.cx_rate, y * self.cy_rate, self.cx_rate, self.cy_rate]) 205 | elif item.stage == 2: 206 | pygame.draw.rect(self.screen, Green, 207 | [x * self.cx_rate, y * self.cy_rate, self.cx_rate, self.cy_rate]) 208 | elif item.stage == 1: 209 | pygame.draw.rect(self.screen, RED, 210 | [x * self.cx_rate, y * self.cy_rate, self.cx_rate, self.cy_rate]) 211 | elif item.stage == 3: 212 | pygame.draw.rect(self.screen, BLACK, 213 | [x * self.cx_rate, y * self.cy_rate, self.cx_rate, self.cy_rate]) 214 | 215 | # def count_num(self): 216 | # self.count0, self.count1, self.count2,self.count3 = self.cells.num_of_nonstage() 217 | mpl.rcParams['font.sans-serif'] = ['FangSong'] # 指定默认字体 218 | mpl.rcParams['axes.unicode_minus'] = False # 解决保存图像是负号'-'显示为方块的问题 219 | if __name__ == '__main__': 220 | count0_ = [] 221 | count1_ = [] 222 | count2_ = [] 223 | count3_ = [] 224 | pygame.init() 225 | pygame.display.set_caption("传染病模型") 226 | game = Game(800, 800, 200, 200) 227 | 228 | clock = pygame.time.Clock() 229 | k1 = 0 230 | while True: 231 | k1 += 1 232 | print(k1) 233 | 234 | # game.screen.fill(GREY)#底部全置灰 235 | clock.tick(100) # 每秒循环10次 236 | for event in pygame.event.get(): 237 | if event.type == pygame.QUIT: 238 | sys.exit() 239 | game.cells.calc_neighbour_count() 240 | count0, count1, count2,count3 = game.cells.num_of_nonstage() 241 | # count0,count1,count2 = game.count_num() 242 | count0_.append(count0) 243 | count1_.append(count1) 244 | count2_.append(count2) 245 | count3_.append(count3) 246 | if count2 > 200*190: # 退出条件 247 | break 248 | 249 | plt.plot(count0_, color='y', label='易感者') 250 | plt.plot(count3_, color='b', label='潜伏者') 251 | plt.plot(count1_, color='r', label='感染者') 252 | plt.plot(count2_, color='g', label='治愈者') 253 | # plt.ylim([0,80000]) 254 | plt.legend() 255 | plt.xlabel('时间单位') 256 | plt.ylabel('人数单位') 257 | plt.pause(0.1)#0.1秒停一次 258 | plt.clf()#清除 259 | 260 | # plt.close()#退出 261 | game.show_life() 262 | pygame.display.flip() 263 | game.cells.next_iter() 264 | 265 | # plt.show()#显示 --------------------------------------------------------------------------------