-
Notifications
You must be signed in to change notification settings - Fork 2
/
Copy pathaggregate_visualization.py
118 lines (85 loc) · 4.15 KB
/
aggregate_visualization.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
"""This module contains the aggregate visualization of the simulation.
"""
import sys
import pyqtgraph as pg # pylint: disable=E0401
from pyqtgraph.Qt import QtCore, QtGui # pylint: disable=E0401
import numpy as np
class AggregatePlot(object):
'''This class renders aggregate variables from a disease simulation.
'''
def __init__(self, disease, kennel):
self.disease = disease
self.kennel = kennel
self.app = QtGui.QApplication([])
self.win = pg.GraphicsWindow(title="Aggregate Measures")
self.win.resize(1500, 600)
self.win.setWindowTitle('Aggregate Measures')
self.win.keyPressEvent = self.keyPressEvent
pg.setConfigOptions(antialias=True)
plots = []
curves = []
plots.append(self.win.addPlot(title="Total Population"))
curves.append(plots[-1].plot(pen='b'))
plots.append(self.win.addPlot(title="Current Population"))
curves.append(plots[-1].plot(pen='b'))
plots.append(self.win.addPlot(title="Total Infected"))
curves.append(plots[-1].plot(pen='y'))
plots.append(self.win.addPlot(title="Current Infected"))
curves.append(plots[-1].plot(pen='y'))
plots.append(self.win.addPlot(title="Infection Rate"))
curves.append(plots[-1].plot(pen=(255, 165, 0)))
self.win.nextRow()
plots.append(self.win.addPlot(title="Total Survived/Immune"))
curves.append(plots[-1].plot(pen='g'))
plots.append(self.win.addPlot(title="Current Survived/Immune"))
curves.append(plots[-1].plot(pen='g'))
plots.append(self.win.addPlot(title="Total Died"))
curves.append(plots[-1].plot(pen='r'))
plots.append(self.win.addPlot(title="Current Died"))
curves.append(plots[-1].plot(pen='r'))
plots.append(self.win.addPlot(title="Survival Rate"))
curves.append(plots[-1].plot(pen='w'))
self.curves = curves
self.plots = plots
self.data = {'pop': [0], 'inf': [0], 'imm': [0], 'die': [0],
'tpop': [0], 'tinf': [0], 'timm': [0], 'tdie': [0],
'ir': [0], 'sr': [0]}
def _key_press_event(self, event):
if event.key() == QtCore.Qt.Key_Escape:
sys.exit(0)
keyPressEvent = _key_press_event
def update(self):
'''Update the graph.
'''
self.app.processEvents()
self.data['tpop'].append(self.disease.total_intake)
self.data['tinf'].append(self.disease.total_infected)
self.data['timm'].append(self.disease.total_discharged)
self.data['tdie'].append(self.disease.total_died)
susceptible_nodes = len(self.disease.get_state_node('S')['members'])
survived_nodes = len(self.disease.get_state_node('IS')['members'])
infected_nodes = len(self.disease.get_state_node('I')['members'])
symptomatic_nodes = len(self.disease.get_state_node('SY')['members'])
died_nodes = len(self.disease.get_state_node('D')['members'])
self.data['pop'].append(susceptible_nodes +
survived_nodes +
infected_nodes +
symptomatic_nodes +
died_nodes)
self.data['inf'].append(infected_nodes + symptomatic_nodes)
self.data['imm'].append(survived_nodes)
self.data['die'].append(died_nodes)
if self.disease.total_intake != 0:
self.data['ir'].append(self.disease.total_infected/self.disease.total_intake)
self.data['sr'].append(self.disease.total_discharged/self.disease.total_intake)
else:
self.data['ir'].append(np.nan)
self.data['sr'].append(np.nan)
for curve, data in zip(self.curves, [self.data['tpop'], self.data['pop'],
self.data['tinf'], self.data['inf'], self.data['ir'],
self.data['timm'], self.data['imm'],
self.data['tdie'], self.data['die'], self.data['sr']]):
curve.setData(data)
if __name__ == '__main__':
from main import main
main()