-
Notifications
You must be signed in to change notification settings - Fork 1
/
SIS.py
94 lines (78 loc) · 2.9 KB
/
SIS.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
import numpy as np
from scipy.integrate import odeint
import matplotlib.pyplot as plt
class SIS(object):
"""docstring"""
name = "SIS"
initial = { "S0" : "Suceptible",
"I0" : "Infecté"}
vars = {"beta" : "taux d'infection",
"gamma" : "taux de guérison"}
def get(self, var):
try:
return eval("self."+var)
except:
return None
def set(self, var, val):
exec("self."+var+" = "+str(val))
def __init__(self, nbSscptbl0=999, nbInfctd0=1,\
infectiousRate=0.3, recoveryRate=0.05, \
timeStart=0, timeStop=1000, nbSteps=1001):
self.S0 = nbSscptbl0
self.I0 = nbInfctd0
self.N = self.S0 + self.I0
self.S, self.I = None, None
self.beta = infectiousRate
self.gamma = recoveryRate
self.timeParam = [timeStart, timeStop, nbSteps]
self.timeVector = np.linspace(timeStart, timeStop, nbSteps)
def differentialEq(self, y, t):
pass
def solveDifferential(self):
# vecteur initial
list0 = []
for elem in self.initial.keys():
list0.append(self.get(elem))
y0 = tuple(list0)
ret = odeint(self.differentialEq, y0, self.timeVector)
for index, elem in enumerate(self.initial.keys()):
self.set(elem[0], list(ret.T[index]))
return ret.T
def createGraph(self, duration=None, Color=['b','y','r','g']):
"""
Susceptible in blue, Exposed in yellow, Infected in red, recovered in green
"""
if duration == None:
duration = self.timeVector[-1]
self.solveDifferential()
fig = plt.figure()
ax = fig.add_subplot()
for index, elem in enumerate(self.initial.keys()):
var = self.get(elem[0])
name = self.initial[elem]
name = '('+name[0].upper()+')'+name[1:]
ax.plot(self.timeVector, var, Color[index], label=name)
ax.set_xlabel('Time (in days)')
ax.set_ylabel('Populaton (in person)')
ax.set_xlim(0, duration)
legend = ax.legend() #ET?
def plot(self, duration=None):
self.createGraph(duration)
plt.show()
def export(self, filename=None, duration=None, d=":"):
if filename ==None:
filename = "Images/"+self.name
for elem in self.initial.keys():
filename += self.get(elem)
filename += d
for val in self.vars.keys():
filename += self.get(val)
filename += d
filename += str(self.timeParam[0]) + d + str(self.timeParam[1]) + d + str(self.timeParam[2])
filename += ".png"
self.createGraph(duration)
plt.savefig(filename)
if __name__ == "__main__":
model = SIS()
model.set("S0", 10)
print(model.get("S0"))