740 Fix ParameterStudy examples in Python (#741)

Author: HenrZu

pycode/examples/simulation/migration_parameter_study.py

@@ -85,15 +85,15 @@ def parameter_study():
         2000)
     model.apply_constraints()
     graph.add_node(id=1, model=model)
-    migration_coefficients = 0.1 * np.ones(8)
+    migration_coefficients = 0.1 * np.ones(model.populations.numel())
     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
-    def handle_result(graph):
+    def handle_result(graph, run_idx):
         print(f'run {handle_result.c}')
         handle_result.c = handle_result.c + 1
         for node_idx in range(graph.num_nodes):

pycode/examples/simulation/parameter_studies.py

@@ -73,7 +73,7 @@ def parameter_study():
     print(model.parameters.ContactPatterns.cont_freq_mat[1].baseline)
     # process the result of one run
 
-    def handle_result(graph):
+    def handle_result(graph, run_idx):
         group = secir.AgeGroup(0)
         print("run {} with infection rate {:.2G}".format(handle_result.c, graph.get_node(
             0).property.model.parameters.TransmissionProbabilityOnContact[group].value))
@@ -93,8 +93,8 @@ def handle_result(graph):
     graph = secir.ModelGraph()
     graph.add_node(0, model)
     graph.add_node(1, model)
-    graph.add_edge(0, 1, 0.01 * np.ones(8*num_groups))
-    graph.add_edge(1, 0, 0.01 * np.ones(8*num_groups))
+    graph.add_edge(0, 1, 0.01 * np.ones(model.populations.numel()*num_groups))
+    graph.add_edge(1, 0, 0.01 * np.ones(model.populations.numel()*num_groups))
 
     study = secir.ParameterStudy(graph, t0=1, tmax=10, dt=0.5, num_runs=3)
     study.run(handle_result)