-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathplot_secir_results.py
More file actions
89 lines (74 loc) · 3.46 KB
/
Copy pathplot_secir_results.py
File metadata and controls
89 lines (74 loc) · 3.46 KB
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
import pandas as pd
import matplotlib.pyplot as plt
# Here we provide a simple script for plotting results of SECIR-type models, working for ODE-based, LCT-based as well
# as IDE-based models. When multiple age groups are used, we plot the results aggregated over all age groups.
# We expect a csv with a name that contains the model name, e.g. "*ode*.csv" or "*lct*.csv".
def plot_secir_results(file):
# In contrast to LCT-based and IDE-based models, the ODE-based SECIR-type model contains compartments for confirmed
# cases of Carrier and Infected compartments. These will not be plotted.
if "ode" in file:
secir_dict = {'Susceptible': 0, 'Exposed': 1, 'Carrier': 2, 'CarrierConfirmed': 3, 'Infected': 3,
'InfectedConfirmed': 4, 'Hospitalized': 5, 'ICU': 6, 'Recovered': 7, 'Dead': 8}
else:
secir_dict = {'Susceptible': 0, 'Exposed': 1, 'Carrier': 2, 'Infected': 3, 'Hospitalized': 4,
'ICU': 5, 'Recovered': 6, 'Dead': 7}
# Read results.
results = pd.read_csv(file)
# Get number of age groups based on shape of results.
num_age_groups = (len(results.columns)-1)//len(secir_dict)
# Define plot.
fig, ax = plt.subplots()
time = results.iloc[:, 0]
# Aggregate results over age groups.
aggregated_susceptibles = 0
aggregated_exposed = 0
aggregated_carrier = 0
aggregated_infected = 0
aggregated_hospitalized = 0
aggregated_icu = 0
aggregated_recovered = 0
aggregated_dead = 0
for group in range(num_age_groups):
aggregated_susceptibles += results.iloc[:,
1 + group*num_age_groups + secir_dict['Susceptible']]
aggregated_exposed += results.iloc[:,
1 + group*num_age_groups + secir_dict['Exposed']]
aggregated_carrier += results.iloc[:,
1 + group*num_age_groups + secir_dict['Carrier']]
aggregated_infected += results.iloc[:,
1 + group*num_age_groups + secir_dict['Infected']]
aggregated_hospitalized += results.iloc[:,
1 + group*num_age_groups + secir_dict['Hospitalized']]
aggregated_icu += results.iloc[:,
1 + group*num_age_groups + secir_dict['ICU']]
aggregated_recovered += results.iloc[:,
1 + group*num_age_groups + secir_dict['Recovered']]
aggregated_dead += results.iloc[:,
1 + group*num_age_groups + secir_dict['Dead']]
# Add results to plot.
ax.plot(
time, aggregated_susceptibles, label='Susceptible')
ax.plot(
time, aggregated_exposed, label='Exposed')
ax.plot(
time, aggregated_carrier, label='Carrier')
ax.plot(
time, aggregated_infected, label='Infected')
ax.plot(
time, aggregated_hospitalized, label='Hospitalized')
ax.plot(
time, aggregated_icu, label='ICU')
ax.plot(
time, aggregated_recovered, label='Recovered')
ax.plot(
time, aggregated_dead, label='Dead')
ax.set_xlabel('Time [days]')
ax.set_ylabel('Individuals [#]')
ax.legend()
plt.show()
def main():
# Example for plotting results of LCT simulation created with tutorial_lct.cpp.
file = "results_lct.csv"
plot_secir_results(file)
if __name__ == '__main__':
main()