-
Notifications
You must be signed in to change notification settings - Fork 0
/
sim.py
123 lines (111 loc) · 5.18 KB
/
sim.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from states import States
class Draw:
def __init__(self, df, path):
self.path = path
self.fig, self.axs = plt.subplots(2, figsize=(10, 8))
self.fig.suptitle('Covid-19 Epidemic Sample Model', fontsize=16)
self.ld = ['Healthy',
'Covid-19(+)',
'Hospitalized',
'Cured',
'Self Cured',
'Death Toll',
'Limbo',
'Total Cured']
self.df = df[0]
self.stateOfDay = df[1]
self.MoversList = df[2]
self.state = States()
self.dayLimit = self.state.returnConfig('SampleDefaults', 'dayLimit', int)
self.x = self.state.returnConfig('SampleDefaults', 'xLimit', int)
self.y = self.state.returnConfig('SampleDefaults', 'yLimit', int)
self.groupsize = self.state.returnConfig('SampleDefaults', 'groupSize', int)
self.sampleSize = self.state.returnConfig('SampleDefaults', 'sampleSize', float)
self.distance = self.state.returnConfig('SampleDefaults', 'distance', float)
self.beds = self.state.returnConfig('SampleDefaults', 'numberOfBeds', float)
self.distanceLimit = self.state.returnConfig("SampleDefaults", "distance", float)
def setup_plot(self, day):
self.axs[0].cla()
self.scat = self.axs[0].scatter(self.df['X'], self.df['Y'], s=1, c=self.plt1color())
for i in self.MoversList:
self.axs[0].scatter(self.df.loc[i,'X'], self.df.loc[i, 'Y'], s=6,
facecolors='none', edgecolors='black')
self.axs[0].text(self.df.loc[i, 'X'] + 0.02, self.df.loc[i, 'Y'] +
0.02, str(i), fontsize=5)
sDay = str(day)
title = f"Day={sDay}, Group Size={self.groupsize}, Movers={self.sampleSize * 100}%, Infection Distance={self.distance}"
self.axs[0].set_title(title, loc='center')
title2 = f"Healthy={len(self.df[self.df['State'] == 0])}, Infected={len(self.df[self.df['State'] == 1])}, Dead={len(self.df[self.df['State'] == 3])}"
self.axs[0].set_xlabel(title2)
self.axs[0].set_yticklabels([])
self.axs[0].set_xticklabels([])
self.axs[0].tick_params(
which='both',
bottom=False,
top=False,
right=False,
left=False,
labelbottom=True)
self.axs[1].set_yticklabels([])
self.axs[1].set_xticklabels([])
self.axs[1].cla()
self.axs[1].plot(self.stateOfDay.Healthy, label=self.ld[0], color=self.plt2color()[0])
self.axs[1].plot(self.stateOfDay['Covid-19(+)'], label=self.ld[1], color=self.plt2color()[1])
self.axs[1].plot(self.stateOfDay.Hospitalized, label=self.ld[2], color=self.plt2color()[2])
self.axs[1].plot(self.stateOfDay.Cured, label=self.ld[3], color=self.plt2color()[3])
self.axs[1].plot(self.stateOfDay['Self Cured'], label=self.ld[4], color=self.plt2color()[4])
self.axs[1].plot(self.stateOfDay.Dead, label=self.ld[5], color=self.plt2color()[5])
self.axs[1].plot(self.stateOfDay.Limbo, label=self.ld[6], color=self.plt2color()[6])
self.axs[1].plot(self.stateOfDay['Total Cured'], label=self.ld[7], color=self.plt2color()[7])
self.axs[1].set_prop_cycle(color=self.plt2color())
self.axs[1].legend(bbox_to_anchor=(0, 1), loc='upper left', borderaxespad=0.)
plt.tight_layout(rect=[0, 0.03, 1, 0.95])
plt.ion()
plt.pause(0.5)
plt.xlabel("Days")
if day < 10:
sDay = '0' + sDay
title = f"{self.path}/days/Day{sDay}.png"
plt.savefig(title)
return
def plt1color(self):
cols = []
for l in self.df.index:
if self.df.loc[l, 'State'] == 1: # Infected
cols.append(self.state.stateColors(1))
elif self.df.loc[l, 'State'] == 3: # Dead
cols.append(self.state.stateColors(3))
elif self.df.loc[l, 'State'] == 2: # Hospitalized
cols.append(self.state.stateColors(2))
elif self.df.loc[l, 'State'] == 4: # Cured
cols.append(self.state.stateColors(4))
elif self.df.loc[l, 'State'] == 5: # Limbo
cols.append(self.state.stateColors(5))
elif self.df.loc[l, 'State'] == 100: #Total Cured
cols.append(self.state.stateColors(100))
else:
cols.append('blue') # Healthy
return cols
def plt2color(self):
cols = []
for i in self.stateOfDay.columns:
if i == 'Covid-19(+)': # Infected
cols.append('red')
elif i == 'Dead': # Dead
cols.append('black')
elif i == 'Hospitalized': # Hospitalized
cols.append('orange')
elif i == 'Cured': # Cured
cols.append('green')
elif i == 'Self Cured': # limbo
cols.append('cyan')
elif i == 'Limbo': # limbo
cols.append('yellow')
elif i == 'Total Cured': # limbo
cols.append('lime')
else:
cols.append('blue') # Healthy
return cols