Skip to content

Commit 369edec

Browse files
MaxBetzDLRMaxBetzDLR
authored
944 streamline secir python code (#1041)
1 parent 575a146 commit 369edec

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

43 files changed

+461
-291
lines changed

‎pycode/examples/simulation/2020_npis_sarscov2_wildtype_germany.py‎

Lines changed: 29 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,32 @@
1+
#############################################################################
2+
# Copyright (C) 2020-2024 MEmilio
3+
#
4+
# Authors: Henrik Zunker
5+
#
6+
# Contact: Martin J. Kuehn <[email protected]>
7+
#
8+
# Licensed under the Apache License, Version 2.0 (the "License");
9+
# you may not use this file except in compliance with the License.
10+
# You may obtain a copy of the License at
11+
#
12+
# http://www.apache.org/licenses/LICENSE-2.0
13+
#
14+
# Unless required by applicable law or agreed to in writing, software
15+
# distributed under the License is distributed on an "AS IS" BASIS,
16+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
17+
# See the License for the specific language governing permissions and
18+
# limitations under the License.
19+
#############################################################################
120
import numpy as np
221
import datetime
322
import os
423
import memilio.simulation as mio
5-
import memilio.simulation.secir as secir
24+
import memilio.simulation.osecir as osecir
625
import memilio.plot.createGIF as mp
726

827
from enum import Enum
9-
from memilio.simulation.secir import (Model, Simulation,
10-
interpolate_simulation_result)
28+
from memilio.simulation.osecir import (Model, Simulation,
29+
interpolate_simulation_result)
1130

1231

1332
class Location(Enum):
@@ -404,7 +423,7 @@ def get_graph(self, end_date):
404423
self.set_contact_matrices(model)
405424
self.set_npis(model.parameters, end_date)
406425

407-
graph = secir.ModelGraph()
426+
graph = osecir.ModelGraph()
408427

409428
scaling_factor_infected = [2.5, 2.5, 2.5, 2.5, 2.5, 2.5]
410429
scaling_factor_icu = 1.0
@@ -414,7 +433,7 @@ def get_graph(self, end_date):
414433
self.data_dir, "pydata", "Germany",
415434
"county_current_population.json")
416435

417-
mio.secir.set_nodes(
436+
mio.osecir.set_nodes(
418437
model.parameters,
419438
mio.Date(self.start_date.year,
420439
self.start_date.month, self.start_date.day),
@@ -423,7 +442,7 @@ def get_graph(self, end_date):
423442
path_population_data, True, graph, scaling_factor_infected,
424443
scaling_factor_icu, tnt_capacity_factor, 0, False)
425444

426-
mio.secir.set_edges(
445+
mio.osecir.set_edges(
427446
self.data_dir, graph, len(Location))
428447

429448
return graph
@@ -438,9 +457,9 @@ def run(self, num_days_sim, num_runs=10, save_graph=True, create_gif=True):
438457
path_graph = os.path.join(self.results_dir, "graph")
439458
if not os.path.exists(path_graph):
440459
os.makedirs(path_graph)
441-
secir.write_graph(graph, path_graph)
460+
osecir.write_graph(graph, path_graph)
442461

443-
study = secir.ParameterStudy(
462+
study = osecir.ParameterStudy(
444463
graph, 0., num_days_sim, 0.5, num_runs)
445464
ensemble = study.run()
446465

@@ -458,7 +477,7 @@ def run(self, num_days_sim, num_runs=10, save_graph=True, create_gif=True):
458477
save_percentiles = True
459478
save_single_runs = False
460479

461-
secir.save_results(
480+
osecir.save_results(
462481
ensemble_results, ensemble_params, node_ids, self.results_dir,
463482
save_single_runs, save_percentiles)
464483
if create_gif:
@@ -474,6 +493,6 @@ def run(self, num_days_sim, num_runs=10, save_graph=True, create_gif=True):
474493
sim = Simulation(
475494
data_dir=os.path.join(file_path, "../../../data"),
476495
start_date=datetime.date(year=2020, month=12, day=12),
477-
results_dir=os.path.join(file_path, "../../../results_secir"))
496+
results_dir=os.path.join(file_path, "../../../results_osecir"))
478497
num_days_sim = 50
479498
sim.run(num_days_sim, num_runs=2)

‎pycode/examples/simulation/migration_parameter_study.py‎

Lines changed: 27 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -22,14 +22,14 @@
2222
import numpy as np
2323

2424
import memilio.simulation as mio
25-
import memilio.simulation.secir as secir
25+
import memilio.simulation.osecir as osecir
2626

2727

28-
def parameter_study():
28+
def run_migration_parameter_study():
2929
mio.set_log_level(mio.LogLevel.Warning)
3030

3131
# setup basic parameters
32-
model = secir.Model(1)
32+
model = osecir.Model(1)
3333

3434
model.parameters.TimeExposed[mio.AgeGroup(0)] = 3.2
3535
model.parameters.TimeInfectedNoSymptoms[mio.AgeGroup(0)] = 2.
@@ -50,48 +50,48 @@ def parameter_study():
5050
model.parameters.DeathsPerCritical[mio.AgeGroup(0)] = 0.3
5151

5252
# two regions with different populations and with some migration between them
53-
graph = secir.ModelGraph()
54-
model.populations[mio.AgeGroup(0), secir.InfectionState.Exposed] = 100
53+
graph = osecir.ModelGraph()
54+
model.populations[mio.AgeGroup(0), osecir.InfectionState.Exposed] = 100
5555
model.populations[mio.AgeGroup(
56-
0), secir.InfectionState.InfectedNoSymptoms] = 50
56+
0), osecir.InfectionState.InfectedNoSymptoms] = 50
5757
model.populations[mio.AgeGroup(
58-
0), secir.InfectionState.InfectedNoSymptomsConfirmed] = 0
58+
0), osecir.InfectionState.InfectedNoSymptomsConfirmed] = 0
5959
model.populations[mio.AgeGroup(
60-
0), secir.InfectionState.InfectedSymptoms] = 50
60+
0), osecir.InfectionState.InfectedSymptoms] = 50
6161
model.populations[mio.AgeGroup(
62-
0), secir.InfectionState.InfectedSevere] = 20
62+
0), osecir.InfectionState.InfectedSevere] = 20
6363
model.populations[mio.AgeGroup(
64-
0), secir.InfectionState.InfectedCritical] = 10
65-
model.populations[mio.AgeGroup(0), secir.InfectionState.Recovered] = 10
66-
model.populations[mio.AgeGroup(0), secir.InfectionState.Dead] = 0
64+
0), osecir.InfectionState.InfectedCritical] = 10
65+
model.populations[mio.AgeGroup(0), osecir.InfectionState.Recovered] = 10
66+
model.populations[mio.AgeGroup(0), osecir.InfectionState.Dead] = 0
6767
model.populations.set_difference_from_group_total_AgeGroup((
6868
mio.AgeGroup(0),
69-
secir.InfectionState.Susceptible),
69+
osecir.InfectionState.Susceptible),
7070
10000)
7171
model.apply_constraints()
7272
graph.add_node(id=0, model=model) # copies the model into the graph
73-
model.populations[mio.AgeGroup(0), secir.InfectionState.Exposed] = 0
73+
model.populations[mio.AgeGroup(0), osecir.InfectionState.Exposed] = 0
7474
model.populations[mio.AgeGroup(
75-
0), secir.InfectionState.InfectedNoSymptoms] = 0
75+
0), osecir.InfectionState.InfectedNoSymptoms] = 0
7676
model.populations[mio.AgeGroup(
77-
0), secir.InfectionState.InfectedSymptoms] = 0
77+
0), osecir.InfectionState.InfectedSymptoms] = 0
7878
model.populations[mio.AgeGroup(
79-
0), secir.InfectionState.InfectedSevere] = 0
79+
0), osecir.InfectionState.InfectedSevere] = 0
8080
model.populations[mio.AgeGroup(
81-
0), secir.InfectionState.InfectedCritical] = 0
82-
model.populations[mio.AgeGroup(0), secir.InfectionState.Recovered] = 0
83-
model.populations[mio.AgeGroup(0), secir.InfectionState.Dead] = 0
81+
0), osecir.InfectionState.InfectedCritical] = 0
82+
model.populations[mio.AgeGroup(0), osecir.InfectionState.Recovered] = 0
83+
model.populations[mio.AgeGroup(0), osecir.InfectionState.Dead] = 0
8484
model.populations.set_difference_from_group_total_AgeGroup((
8585
mio.AgeGroup(0),
86-
secir.InfectionState.Susceptible),
86+
osecir.InfectionState.Susceptible),
8787
2000)
8888
model.apply_constraints()
8989
graph.add_node(id=1, model=model)
9090
migration_coefficients = 0.1 * np.ones(model.populations.numel())
91+
migration_coefficients[osecir.InfectionState.Dead] = 0
9192
migration_params = mio.MigrationParameters(migration_coefficients)
9293
# one coefficient per (age group x compartment)
9394
graph.add_edge(0, 1, migration_params)
94-
# directed graph -> add both directions so coefficients can be different
9595
graph.add_edge(1, 0, migration_params)
9696

9797
# process the result of one run
@@ -107,7 +107,7 @@ def handle_result(graph, run_idx):
107107
" initial InfectedNoSymptoms count {}.".format(
108108
model.populations
109109
[mio.AgeGroup(0),
110-
secir.InfectionState.InfectedNoSymptoms].value))
110+
osecir.InfectionState.InfectedNoSymptoms].value))
111111
print(f" compartments at t = {result.get_time(0)}:")
112112
print(" ", result.get_value(0))
113113
print(f" compartments at t = {result.get_last_time()}:")
@@ -117,17 +117,17 @@ def handle_result(graph, run_idx):
117117
# study with unknown number of undetected InfectedNoSymptoms
118118
carrier_distribution = mio.ParameterDistributionNormal(50, 2000, 200, 100)
119119
graph.get_node(0).property.populations[mio.AgeGroup(
120-
0), secir.InfectionState.InfectedNoSymptoms].set_distribution(carrier_distribution)
120+
0), osecir.InfectionState.InfectedNoSymptoms].set_distribution(carrier_distribution)
121121

122122
t0 = 0
123123
tmax = 50
124-
study = secir.ParameterStudy(graph, t0, tmax, dt=1.0, num_runs=3)
124+
study = osecir.ParameterStudy(graph, t0, tmax, dt=1.0, num_runs=3)
125125
study.run(handle_result)
126126

127127

128128
if __name__ == "__main__":
129129
arg_parser = argparse.ArgumentParser(
130130
'migration_parameter_study',
131-
description='Example demonstrating setup of ensemble runs of a geographically resolved SECIR model with travel.')
131+
description='Example demonstrating setup of ensemble runs of a geographically resolved ODE SECIHURD model with mobility.')
132132
args = arg_parser.parse_args()
133-
parameter_study()
133+
run_migration_parameter_study()

pycode/examples/simulation/secir_ageres.py renamed to pycode/examples/simulation/ode_secir_ageres.py

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -26,14 +26,14 @@
2626
import pandas as pd
2727

2828
from memilio.simulation import ContactMatrix, Damping, UncertainContactMatrix, AgeGroup
29-
from memilio.simulation.secir import Index_InfectionState
30-
from memilio.simulation.secir import InfectionState as State
31-
from memilio.simulation.secir import Model, Simulation, simulate
29+
from memilio.simulation.osecir import Index_InfectionState
30+
from memilio.simulation.osecir import InfectionState as State
31+
from memilio.simulation.osecir import Model, Simulation, simulate
3232

3333

34-
def run_secir_simulation(show_plot=True):
34+
def run_ode_secir_ageres_simulation(show_plot=True):
3535
"""
36-
Runs the c++ secir model using multiple age groups
36+
Runs the c++ ODE SECIHURD model using multiple age groups
3737
and plots the results
3838
"""
3939

@@ -129,12 +129,12 @@ def run_secir_simulation(show_plot=True):
129129
ax.plot(t, data[:, 5], label='#InfectedCritical')
130130
ax.plot(t, data[:, 6], label='#Recovered')
131131
ax.plot(t, data[:, 7], label='#Died')
132-
ax.set_title("SECIR model simulation")
132+
ax.set_title("ODE SECIR model simulation")
133133
ax.set_xticks(tick_range)
134134
ax.set_xticklabels(datelist[tick_range], rotation=45)
135135
ax.legend()
136136
fig.tight_layout
137-
fig.savefig('Secir_simple.pdf')
137+
fig.savefig('osecir_simple.pdf')
138138

139139
plt.show()
140140
plt.close()
@@ -145,9 +145,9 @@ def run_secir_simulation(show_plot=True):
145145

146146
if __name__ == "__main__":
147147
arg_parser = argparse.ArgumentParser(
148-
'secir_ageres',
149-
description='Example demonstrating the setup and simulation of the SECIR model with multiple age groups.')
148+
'ode_secir_ageres',
149+
description='Example demonstrating the setup and simulation of the ODE SECIHURD model with multiple age groups.')
150150
arg_parser.add_argument('-p', '--show_plot',
151151
action='store_const', const=True, default=False)
152152
args = arg_parser.parse_args()
153-
run_secir_simulation(**args.__dict__)
153+
run_ode_secir_ageres_simulation(**args.__dict__)

pycode/examples/simulation/secir_groups.py renamed to pycode/examples/simulation/ode_secir_groups.py

Lines changed: 17 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -26,15 +26,15 @@
2626
import pandas as pd
2727

2828
from memilio.simulation import AgeGroup, ContactMatrix, Damping, UncertainContactMatrix
29-
from memilio.simulation.secir import Index_InfectionState
30-
from memilio.simulation.secir import InfectionState as State
31-
from memilio.simulation.secir import (Model, Simulation,
32-
interpolate_simulation_result, simulate)
29+
from memilio.simulation.osecir import Index_InfectionState
30+
from memilio.simulation.osecir import InfectionState as State
31+
from memilio.simulation.osecir import (Model, Simulation,
32+
interpolate_simulation_result, simulate)
3333

3434

35-
def run_secir_groups_simulation(show_plot=True):
35+
def run_ode_secir_groups_simulation(show_plot=True):
3636
"""
37-
Runs the c++ secir model using mulitple age groups
37+
Runs the c++ ODE SECIHURD model using mulitple age groups
3838
and plots the results
3939
"""
4040

@@ -159,12 +159,12 @@ def run_secir_groups_simulation(show_plot=True):
159159
ax.plot(t, data[:, 5], label='#InfectedCritical')
160160
ax.plot(t, data[:, 6], label='#Recovered')
161161
ax.plot(t, data[:, 7], label='#Dead')
162-
ax.set_title("SECIR simulation results (entire population)")
162+
ax.set_title("ODE SECIR simulation results (entire population)")
163163
ax.set_xticks(tick_range)
164164
ax.set_xticklabels(datelist[tick_range], rotation=45)
165165
ax.legend()
166166
fig.tight_layout
167-
fig.savefig('Secir_by_compartments.pdf')
167+
fig.savefig('osecir_by_compartments.pdf')
168168

169169
# plot dynamics in each comparment by age group
170170
fig, ax = plt.subplots(4, 2, figsize=(12, 15))
@@ -182,8 +182,9 @@ def run_secir_groups_simulation(show_plot=True):
182182
ax[int(np.floor(i / 2)), int(i % 2)
183183
].set_xticklabels(datelist[tick_range], rotation=45)
184184
plt.subplots_adjust(hspace=0.5, bottom=0.1, top=0.9)
185-
fig.suptitle('SECIR simulation results by age group in each compartment')
186-
fig.savefig('Secir_age_groups_in_compartments.pdf')
185+
fig.suptitle(
186+
'ODE SECIR simulation results by age group in each compartment')
187+
fig.savefig('osecir_age_groups_in_compartments.pdf')
187188

188189
fig, ax = plt.subplots(4, 2, figsize=(12, 15))
189190
for i, title in zip(range(num_compartments), compartments):
@@ -194,8 +195,9 @@ def run_secir_groups_simulation(show_plot=True):
194195
ax[int(np.floor(i / 2)), int(i % 2)
195196
].set_xticklabels(datelist[tick_range], rotation=45)
196197
plt.subplots_adjust(hspace=0.5, bottom=0.1, top=0.9)
197-
fig.suptitle('SECIR simulation results by compartment (entire population)')
198-
fig.savefig('Secir_all_parts.pdf')
198+
fig.suptitle(
199+
'ODE SECIR simulation results by compartment (entire population)')
200+
fig.savefig('osecir_all_parts.pdf')
199201

200202
if show_plot:
201203
plt.show()
@@ -206,9 +208,9 @@ def run_secir_groups_simulation(show_plot=True):
206208

207209
if __name__ == "__main__":
208210
arg_parser = argparse.ArgumentParser(
209-
'secir_groups',
210-
description='Simple example demonstrating the setup and simulation of the SECIR model with multiple age groups.')
211+
'ode_secir_groups',
212+
description='Simple example demonstrating the setup and simulation of the ODE SECIHURD model with multiple age groups.')
211213
arg_parser.add_argument('-p', '--show_plot',
212214
action='store_const', const=True, default=False)
213215
args = arg_parser.parse_args()
214-
run_secir_groups_simulation(**args.__dict__)
216+
run_ode_secir_groups_simulation(**args.__dict__)

0 commit comments

Comments
 (0)