736 Fix mobility coefficients in migration example (#737)

Author: HenrZu

pycode/examples/simulation/migration.py

@@ -20,12 +20,13 @@
 import argparse
 
 import numpy as np
+import matplotlib.pyplot as plt
 
 import memilio.simulation as mio
 import memilio.simulation.secir as secir
 
 
-def parameter_study():
+def run_mobility_example(plot_results=True):
     mio.set_log_level(mio.LogLevel.Warning)
 
     t0 = 0
@@ -52,7 +53,7 @@ def parameter_study():
     model.parameters.CriticalPerSevere[secir.AgeGroup(0)] = 0.25
     model.parameters.DeathsPerCritical[secir.AgeGroup(0)] = 0.3
 
-    # two regions with different populations and with some migration between them
+    # two regions with different populations and with some mobility between them
     graph = secir.MigrationGraph()
     model.populations[secir.AgeGroup(0), secir.InfectionState.Exposed] = 100
     model.populations[secir.AgeGroup(
@@ -88,25 +89,80 @@ def parameter_study():
         2000)
     model.apply_constraints()
     graph.add_node(id=1, model=model, t0=t0)
-    migration_coefficients = 0.1 * np.ones(8)
-    migration_params = mio.MigrationParameters(migration_coefficients)
+    mobility_coefficients = 0.1 * np.ones(model.populations.numel())
+    mobility_params = mio.MigrationParameters(mobility_coefficients)
     # one coefficient per (age group x compartment)
-    graph.add_edge(0, 1, migration_params)
+    graph.add_edge(0, 1, mobility_params)
     # directed graph -> add both directions so coefficients can be different
-    graph.add_edge(1, 0, migration_params)
+    graph.add_edge(1, 0, mobility_params)
 
     # run simulation
     sim = secir.MigrationSimulation(graph, t0, dt=0.5)
     sim.advance(tmax)
 
     # process results
-    region0_result = sim.graph.get_node(0).property.result
-    region1_result = sim.graph.get_node(1).property.result
+    region0_result = secir.interpolate_simulation_result(
+        sim.graph.get_node(0).property.result)
+    region1_result = secir.interpolate_simulation_result(
+        sim.graph.get_node(1).property.result)
+
+    if (plot_results):
+        results = [region0_result.as_ndarray(), region1_result.as_ndarray()]
+        t = results[0][0, :]
+        tick_range = (np.arange(int((len(t) - 1) / 10) + 1) * 10)
+        tick_range[-1] -= 1
+
+        fig, ax = plt.subplots(figsize=(10, 6))
+        for idx, result_region in enumerate(results):
+            region_label = f'Region {idx}'
+            ax.plot(t, result_region[1, :],
+                    label=f'{region_label} - #Susceptible')
+            ax.plot(t, result_region[2, :], label=f'{region_label} - #Exposed')
+            ax.plot(t, result_region[3, :] + result_region[4, :],
+                    label=f'{region_label} - #InfectedNoSymptoms')
+            ax.plot(t, result_region[5, :] + result_region[6, :],
+                    label=f'{region_label} - #InfectedSymptoms')
+            ax.plot(t, result_region[7, :],
+                    label=f'{region_label} - #Hospitalzed')
+            ax.plot(t, result_region[8, :],
+                    label=f'{region_label} - #InfectedCritical')
+            ax.plot(t, result_region[9, :],
+                    label=f'{region_label} - #Recovered')
+            ax.plot(t, result_region[10, :], label=f'{region_label} - #Dead')
+
+        ax.set_title(
+            "SECIR simulation results for both regions (entire population)")
+        ax.set_xticks(tick_range)
+        ax.legend(loc='upper right', bbox_to_anchor=(1, 0.6))
+        plt.yscale('log')
+        fig.tight_layout
+        fig.savefig('Mobility_Secir_by_compartments.pdf')
+
+        fig, ax = plt.subplots(5, 2, figsize=(12, 15))
+        compartments = [
+            'Susceptible', 'Exposed', 'InfectedNoSymptoms',
+            'InfectedNoSymptomsConfirmed', 'InfectedSymptoms',
+            'InfectedSymptomsConfirmed', 'InfectedSevere', 'InfectedCritical',
+            'Recovered', 'Dead']
+        num_compartments = len(compartments)
+
+        for i, title in zip(range(num_compartments), compartments):
+            ax[int(np.floor(i / 2)), int(i % 2)].plot(t,
+                                                      results[0][i+1, :], label="Region 0")
+            ax[int(np.floor(i / 2)), int(i % 2)].plot(t,
+                                                      results[1][i+1, :], label="Region 1")
+            ax[int(np.floor(i / 2)), int(i % 2)].set_title(title, fontsize=10)
+            ax[int(np.floor(i / 2)), int(i % 2)].legend()
+            ax[int(np.floor(i / 2)), int(i % 2)].set_xticks(tick_range)
+
+        plt.subplots_adjust(hspace=0.5, bottom=0.1, top=0.9)
+        fig.suptitle('Simulation results for each region in each compartment')
+        fig.savefig('Region_results_compartments.pdf')
 
 
 if __name__ == "__main__":
     arg_parser = argparse.ArgumentParser(
         'migration',
         description='Example demonstrating the setup and simulation of a geographically resolved SECIR model with travel.')
     args = arg_parser.parse_args()
-    parameter_study()
+    run_mobility_example()