944 streamline secir python code

pycode/examples/simulation/2020_npis_sarscov2_wildtype_germany.py

@@ -1,13 +1,32 @@
+#############################################################################
+# Copyright (C) 2020-2024 MEmilio
+#
+# Authors: Henrik Zunker
+#
+# Contact: Martin J. Kuehn <Martin.Kuehn@DLR.de>
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#############################################################################
 import numpy as np
 import datetime
 import os
 import memilio.simulation as mio
-import memilio.simulation.secir as secir
+import memilio.simulation.osecir as osecir
 import memilio.plot.createGIF as mp
 
 from enum import Enum
-from memilio.simulation.secir import (Model, Simulation,
-                                      interpolate_simulation_result)
+from memilio.simulation.osecir import (Model, Simulation,
+                                       interpolate_simulation_result)
 
 
 class Location(Enum):
@@ -404,7 +423,7 @@ def get_graph(self, end_date):
         self.set_contact_matrices(model)
         self.set_npis(model.parameters, end_date)
 
-        graph = secir.ModelGraph()
+        graph = osecir.ModelGraph()
 
         scaling_factor_infected = [2.5, 2.5, 2.5, 2.5, 2.5, 2.5]
         scaling_factor_icu = 1.0
@@ -414,7 +433,7 @@ def get_graph(self, end_date):
             self.data_dir, "pydata", "Germany",
             "county_current_population.json")
 
-        mio.secir.set_nodes(
+        mio.osecir.set_nodes(
             model.parameters,
             mio.Date(self.start_date.year,
                      self.start_date.month, self.start_date.day),
@@ -423,7 +442,7 @@ def get_graph(self, end_date):
             path_population_data, True, graph, scaling_factor_infected,
             scaling_factor_icu, tnt_capacity_factor, 0, False)
 
-        mio.secir.set_edges(
+        mio.osecir.set_edges(
             self.data_dir, graph, len(Location))
 
         return graph
@@ -438,9 +457,9 @@ def run(self, num_days_sim, num_runs=10, save_graph=True, create_gif=True):
             path_graph = os.path.join(self.results_dir, "graph")
             if not os.path.exists(path_graph):
                 os.makedirs(path_graph)
-            secir.write_graph(graph, path_graph)
+            osecir.write_graph(graph, path_graph)
 
-        study = secir.ParameterStudy(
+        study = osecir.ParameterStudy(
             graph, 0., num_days_sim, 0.5, num_runs)
         ensemble = study.run()
 
@@ -458,7 +477,7 @@ def run(self, num_days_sim, num_runs=10, save_graph=True, create_gif=True):
         save_percentiles = True
         save_single_runs = False
 
-        secir.save_results(
+        osecir.save_results(
             ensemble_results, ensemble_params, node_ids, self.results_dir,
             save_single_runs, save_percentiles)
         if create_gif:
@@ -474,6 +493,6 @@ def run(self, num_days_sim, num_runs=10, save_graph=True, create_gif=True):
     sim = Simulation(
         data_dir=os.path.join(file_path, "../../../data"),
         start_date=datetime.date(year=2020, month=12, day=12),
-        results_dir=os.path.join(file_path, "../../../results_secir"))
+        results_dir=os.path.join(file_path, "../../../results_osecir"))
     num_days_sim = 50
     sim.run(num_days_sim, num_runs=2)

pycode/examples/simulation/migration_parameter_study.py

@@ -22,14 +22,14 @@
 import numpy as np
 
 import memilio.simulation as mio
-import memilio.simulation.secir as secir
+import memilio.simulation.osecir as osecir
 
 
-def parameter_study():
+def run_migration_parameter_study():
     mio.set_log_level(mio.LogLevel.Warning)
 
     # setup basic parameters
-    model = secir.Model(1)
+    model = osecir.Model(1)
 
     model.parameters.TimeExposed[mio.AgeGroup(0)] = 3.2
     model.parameters.TimeInfectedNoSymptoms[mio.AgeGroup(0)] = 2.
@@ -50,48 +50,48 @@ def parameter_study():
     model.parameters.DeathsPerCritical[mio.AgeGroup(0)] = 0.3
 
     # two regions with different populations and with some migration between them
-    graph = secir.ModelGraph()
-    model.populations[mio.AgeGroup(0), secir.InfectionState.Exposed] = 100
+    graph = osecir.ModelGraph()
+    model.populations[mio.AgeGroup(0), osecir.InfectionState.Exposed] = 100
     model.populations[mio.AgeGroup(
-        0), secir.InfectionState.InfectedNoSymptoms] = 50
+        0), osecir.InfectionState.InfectedNoSymptoms] = 50
     model.populations[mio.AgeGroup(
-        0), secir.InfectionState.InfectedNoSymptomsConfirmed] = 0
+        0), osecir.InfectionState.InfectedNoSymptomsConfirmed] = 0
     model.populations[mio.AgeGroup(
-        0), secir.InfectionState.InfectedSymptoms] = 50
+        0), osecir.InfectionState.InfectedSymptoms] = 50
     model.populations[mio.AgeGroup(
-        0), secir.InfectionState.InfectedSevere] = 20
+        0), osecir.InfectionState.InfectedSevere] = 20
     model.populations[mio.AgeGroup(
-        0), secir.InfectionState.InfectedCritical] = 10
-    model.populations[mio.AgeGroup(0), secir.InfectionState.Recovered] = 10
-    model.populations[mio.AgeGroup(0), secir.InfectionState.Dead] = 0
+        0), osecir.InfectionState.InfectedCritical] = 10
+    model.populations[mio.AgeGroup(0), osecir.InfectionState.Recovered] = 10
+    model.populations[mio.AgeGroup(0), osecir.InfectionState.Dead] = 0
     model.populations.set_difference_from_group_total_AgeGroup((
         mio.AgeGroup(0),
-        secir.InfectionState.Susceptible),
+        osecir.InfectionState.Susceptible),
         10000)
     model.apply_constraints()
     graph.add_node(id=0, model=model)  # copies the model into the graph
-    model.populations[mio.AgeGroup(0), secir.InfectionState.Exposed] = 0
+    model.populations[mio.AgeGroup(0), osecir.InfectionState.Exposed] = 0
     model.populations[mio.AgeGroup(
-        0), secir.InfectionState.InfectedNoSymptoms] = 0
+        0), osecir.InfectionState.InfectedNoSymptoms] = 0
     model.populations[mio.AgeGroup(
-        0), secir.InfectionState.InfectedSymptoms] = 0
+        0), osecir.InfectionState.InfectedSymptoms] = 0
     model.populations[mio.AgeGroup(
-        0), secir.InfectionState.InfectedSevere] = 0
+        0), osecir.InfectionState.InfectedSevere] = 0
     model.populations[mio.AgeGroup(
-        0), secir.InfectionState.InfectedCritical] = 0
-    model.populations[mio.AgeGroup(0), secir.InfectionState.Recovered] = 0
-    model.populations[mio.AgeGroup(0), secir.InfectionState.Dead] = 0
+        0), osecir.InfectionState.InfectedCritical] = 0
+    model.populations[mio.AgeGroup(0), osecir.InfectionState.Recovered] = 0
+    model.populations[mio.AgeGroup(0), osecir.InfectionState.Dead] = 0
     model.populations.set_difference_from_group_total_AgeGroup((
         mio.AgeGroup(0),
-        secir.InfectionState.Susceptible),
+        osecir.InfectionState.Susceptible),
         2000)
     model.apply_constraints()
     graph.add_node(id=1, model=model)
     migration_coefficients = 0.1 * np.ones(model.populations.numel())
+    migration_coefficients[osecir.InfectionState.Dead] = 0
     migration_params = mio.MigrationParameters(migration_coefficients)
     # one coefficient per (age group x compartment)
     graph.add_edge(0, 1, migration_params)
-    # directed graph -> add both directions so coefficients can be different
     graph.add_edge(1, 0, migration_params)
 
     # process the result of one run
@@ -107,7 +107,7 @@ def handle_result(graph, run_idx):
                 "  initial InfectedNoSymptoms count {}.".format(
                     model.populations
                     [mio.AgeGroup(0),
-                     secir.InfectionState.InfectedNoSymptoms].value))
+                     osecir.InfectionState.InfectedNoSymptoms].value))
             print(f"  compartments at t = {result.get_time(0)}:")
             print("  ", result.get_value(0))
             print(f"  compartments at t = {result.get_last_time()}:")
@@ -117,17 +117,17 @@ def handle_result(graph, run_idx):
     # study with unknown number of undetected InfectedNoSymptoms
     carrier_distribution = mio.ParameterDistributionNormal(50, 2000, 200, 100)
     graph.get_node(0).property.populations[mio.AgeGroup(
-        0), secir.InfectionState.InfectedNoSymptoms].set_distribution(carrier_distribution)
+        0), osecir.InfectionState.InfectedNoSymptoms].set_distribution(carrier_distribution)
 
     t0 = 0
     tmax = 50
-    study = secir.ParameterStudy(graph, t0, tmax, dt=1.0, num_runs=3)
+    study = osecir.ParameterStudy(graph, t0, tmax, dt=1.0, num_runs=3)
     study.run(handle_result)
 
 
 if __name__ == "__main__":
     arg_parser = argparse.ArgumentParser(
         'migration_parameter_study',
-        description='Example demonstrating setup of ensemble runs of a geographically resolved SECIR model with travel.')
+        description='Example demonstrating setup of ensemble runs of a geographically resolved ODE SECIHURD model with mobility.')
     args = arg_parser.parse_args()
-    parameter_study()
+    run_migration_parameter_study()

pycode/examples/simulation/secir_ageres.py → pycode/examples/simulation/ode_secir_ageres.py

@@ -26,14 +26,14 @@
 import pandas as pd
 
 from memilio.simulation import ContactMatrix, Damping, UncertainContactMatrix, AgeGroup
-from memilio.simulation.secir import Index_InfectionState
-from memilio.simulation.secir import InfectionState as State
-from memilio.simulation.secir import Model, Simulation, simulate
+from memilio.simulation.osecir import Index_InfectionState
+from memilio.simulation.osecir import InfectionState as State
+from memilio.simulation.osecir import Model, Simulation, simulate
 
 
-def run_secir_simulation(show_plot=True):
+def run_ode_secir_ageres_simulation(show_plot=True):
     """
-    Runs the c++ secir model using multiple age groups
+    Runs the c++ ODE SECIHURD model using multiple age groups
     and plots the results
     """
 
@@ -129,12 +129,12 @@ def run_secir_simulation(show_plot=True):
         ax.plot(t, data[:, 5], label='#InfectedCritical')
         ax.plot(t, data[:, 6], label='#Recovered')
         ax.plot(t, data[:, 7], label='#Died')
-        ax.set_title("SECIR model simulation")
+        ax.set_title("ODE SECIR model simulation")
         ax.set_xticks(tick_range)
         ax.set_xticklabels(datelist[tick_range], rotation=45)
         ax.legend()
         fig.tight_layout
-        fig.savefig('Secir_simple.pdf')
+        fig.savefig('osecir_simple.pdf')
 
         plt.show()
         plt.close()
@@ -145,9 +145,9 @@ def run_secir_simulation(show_plot=True):
 
 if __name__ == "__main__":
     arg_parser = argparse.ArgumentParser(
-        'secir_ageres',
-        description='Example demonstrating the setup and simulation of the SECIR model with multiple age groups.')
+        'ode_secir_ageres',
+        description='Example demonstrating the setup and simulation of the ODE SECIHURD model with multiple age groups.')
     arg_parser.add_argument('-p', '--show_plot',
                             action='store_const', const=True, default=False)
     args = arg_parser.parse_args()
-    run_secir_simulation(**args.__dict__)
+    run_ode_secir_ageres_simulation(**args.__dict__)

pycode/examples/simulation/secir_groups.py → pycode/examples/simulation/ode_secir_groups.py

@@ -26,15 +26,15 @@
 import pandas as pd
 
 from memilio.simulation import AgeGroup, ContactMatrix, Damping, UncertainContactMatrix
-from memilio.simulation.secir import Index_InfectionState
-from memilio.simulation.secir import InfectionState as State
-from memilio.simulation.secir import (Model, Simulation,
-                                      interpolate_simulation_result, simulate)
+from memilio.simulation.osecir import Index_InfectionState
+from memilio.simulation.osecir import InfectionState as State
+from memilio.simulation.osecir import (Model, Simulation,
+                                       interpolate_simulation_result, simulate)
 
 
-def run_secir_groups_simulation(show_plot=True):
+def run_ode_secir_groups_simulation(show_plot=True):
     """
-    Runs the c++ secir model using mulitple age groups 
+    Runs the c++ ODE SECIHURD model using mulitple age groups 
     and plots the results
     """
 
@@ -159,12 +159,12 @@ def run_secir_groups_simulation(show_plot=True):
     ax.plot(t, data[:, 5], label='#InfectedCritical')
     ax.plot(t, data[:, 6], label='#Recovered')
     ax.plot(t, data[:, 7], label='#Dead')
-    ax.set_title("SECIR simulation results (entire population)")
+    ax.set_title("ODE SECIR simulation results (entire population)")
     ax.set_xticks(tick_range)
     ax.set_xticklabels(datelist[tick_range], rotation=45)
     ax.legend()
     fig.tight_layout
-    fig.savefig('Secir_by_compartments.pdf')
+    fig.savefig('osecir_by_compartments.pdf')
 
     # plot dynamics in each comparment by age group
     fig, ax = plt.subplots(4, 2, figsize=(12, 15))
@@ -182,8 +182,9 @@ def run_secir_groups_simulation(show_plot=True):
         ax[int(np.floor(i / 2)), int(i % 2)
            ].set_xticklabels(datelist[tick_range], rotation=45)
     plt.subplots_adjust(hspace=0.5, bottom=0.1, top=0.9)
-    fig.suptitle('SECIR simulation results by age group in each compartment')
-    fig.savefig('Secir_age_groups_in_compartments.pdf')
+    fig.suptitle(
+        'ODE SECIR simulation results by age group in each compartment')
+    fig.savefig('osecir_age_groups_in_compartments.pdf')
 
     fig, ax = plt.subplots(4, 2, figsize=(12, 15))
     for i, title in zip(range(num_compartments), compartments):
@@ -194,8 +195,9 @@ def run_secir_groups_simulation(show_plot=True):
         ax[int(np.floor(i / 2)), int(i % 2)
            ].set_xticklabels(datelist[tick_range], rotation=45)
     plt.subplots_adjust(hspace=0.5, bottom=0.1, top=0.9)
-    fig.suptitle('SECIR simulation results by compartment (entire population)')
-    fig.savefig('Secir_all_parts.pdf')
+    fig.suptitle(
+        'ODE SECIR simulation results by compartment (entire population)')
+    fig.savefig('osecir_all_parts.pdf')
 
     if show_plot:
         plt.show()
@@ -206,9 +208,9 @@ def run_secir_groups_simulation(show_plot=True):
 
 if __name__ == "__main__":
     arg_parser = argparse.ArgumentParser(
-        'secir_groups',
-        description='Simple example demonstrating the setup and simulation of the SECIR model with multiple age groups.')
+        'ode_secir_groups',
+        description='Simple example demonstrating the setup and simulation of the ODE SECIHURD model with multiple age groups.')
     arg_parser.add_argument('-p', '--show_plot',
                             action='store_const', const=True, default=False)
     args = arg_parser.parse_args()
-    run_secir_groups_simulation(**args.__dict__)
+    run_ode_secir_groups_simulation(**args.__dict__)