-
Notifications
You must be signed in to change notification settings - Fork 0
/
Simulation.py
144 lines (117 loc) · 4.5 KB
/
Simulation.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
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
from Distribution import Distribution
from PIL import Image
import matplotlib.pyplot as plt
import random
from collections import deque
from scipy import stats
import numpy as np
import os
import networkx as nx
import networkx.drawing
class Simulation:
def __init__(self, p, network):
self.p = p
self.network = network
self.network_dicts = nx.convert.to_dict_of_dicts(self.network) #convert to dictionaryform
self.N = len(set(self.network_dicts.keys()))
self.distr = Distribution(stats.bernoulli(self.p))
print("transmission-probability: {}, populationsize: {}".format(self.p, self.N))
def simulate(self, verbose=False, makeGIF=False):
simres = SimulationResults(self.network)
susceptible = set(self.network_dicts.keys())
infected = set()
recovered = set()
first_infect = random.sample(susceptible, 1)[0]
susceptible.remove(first_infect)
infected.add(first_infect)
t=0
while len(infected) > 0:
assert len(susceptible) + len(infected) + len(recovered)==self.N, "incorrect populationsize"
simres.registerState(susceptible, infected, recovered, drawGIF=makeGIF)
#new infects
new_infects = set()
for person in infected:
for friend in self.network_dicts[person].keys():
if friend in susceptible:
if self.distr.rvs():
new_infects.add(friend)
#move infected to susceptible
recovered = recovered.union(infected)
susceptible = susceptible.difference(new_infects)
infected = new_infects
t+=1
if verbose:
print("not-infected: {}, recovered: {}, timesteps: {}".format(len(susceptible), len(recovered), t))
simres.plotStates()
simres.registerFinalAffected(recovered)
return simres
class SimulationResults:
"""
Class for storing results from Simulation, also return type
"""
def __init__(self, network):
self.sState = deque()
self.iState = deque()
self.rState = deque()
self.tState = deque([0,])
self.network = network
self.pos = nx.spring_layout(network)
def registerState(self, s, i, r, drawGIF=False):
self.sState.append(len(s))
self.iState.append(len(i))
self.rState.append(len(r))
self.tState.append(self.tState[-1]+1)
if drawGIF:
self.makeGIFFrame(s, i, r)
def registerFinalAffected(self, recovered):
self.affected = recovered
def getSusceptipleDevelopment(self):
return self.sState
def getRecoveredDevelopment(self):
return self.rState
def getTotalTimesteps(self):
return self.tState[-1]
def makeGIFFrame(self, s, i, r):
colormap = []
for node in self.network:
if node in r:
colormap.append("Orange")
elif node in i:
colormap.append("Red")
elif node in s:
colormap.append("Blue")
else:
print("shouldnt happen")
colormap.append("Blue")
plt.figure(figsize=(30,20))
nx.draw(self.network, node_color=colormap, with_labels=True, pos=self.pos, node_size=600)
if not os.path.isdir("./GIF"):
os.mkdir("./GIF")
plt.savefig(os.getcwd() + "\\GIF\\frame_"+str(self.tState[-1]).zfill(3)+".png")
plt.clf()
def makeGIF(self, fp):
# Create the frames
frames = []
imgs = [os.getcwd() + "\\GIF\\"+x for x in os.listdir(os.getcwd() + "\\GIF\\")]
for i in imgs:
new_frame = Image.open(i)
frames.append(new_frame)
# Save into a GIF file that loops forever
frames[0].save(fp, format='GIF',
append_images=frames[1:],
save_all=True,
duration=300, loop=1)
#remove saved frames and the folder
for img in imgs:
os.remove(img)
os.rmdir(os.getcwd() + "\\GIF\\")
def plotStates(self):
'''make plot of evoltuion of system states'''
plt.figure(figsize=(8,6))
plt.plot(self.sState, color='blue')
plt.plot(self.iState, color='red')
plt.plot(self.rState, color='green')
#plt.title("first infect: " + str(self.firstInfect), size=10)
plt.legend(['suspectible', 'infected', 'recovered'])
plt.grid(True)
plt.show()