include AdaptiveAggregatedDistance in exercises as well

Author: charlie0614

exercises/exercise04.py

@@ -161,13 +161,13 @@ def _(mo):
 
 @app.cell
 def _(contact_frequency, np, osecir, t0, tmax):
-    def run_simulation(parameters, tmax = tmax):
+    def run_simulation(parameters, tmax=tmax):
         # Create model and set parameters
         influenza_model = osecir.Model(1)
         # TODO: We need to run the set_population and set_parameters functions:
-    
-    
-        influenza_model.parameters.ContactPatterns.cont_freq_mat[0].baseline = np.ones((1,1)) * contact_frequency
+
+        influenza_model.parameters.ContactPatterns.cont_freq_mat[0].baseline = np.ones(
+            (1, 1)) * contact_frequency
         # Check that the parameters are meaningful, i.e., no negative dwelling times
         influenza_model.check_constraints()
 
@@ -263,7 +263,7 @@ def _(mo):
     mo.md(r"""
     ## Defining the objective function
 
-    The last step before running the fitting is the defintion of an objective (or distance) function. Here, we are given data for the ICU cases and deaths per day. Thus an obvious choice for the distance function is to calculate the difference between the simulated and the observed numbers per day and adding them up.
+    The last step before running the fitting is the defintion of an objective (or distance) function. Here, we are given data for the ICU cases and deaths per day. Thus an obvious choice for the distance function is to calculate the difference between the simulated and the observed numbers per day and then add them up. As they live on different scales, we define the distance on ICU cases and deaths seperately and use `pyabc.AdaptiveAggregatedDistance` to scale them and aggregate.
 
     We need a function that takes a `data` dictionary provided by our `run_simulation` function and an `observation` dictionary, given by our input data. As in the plotting sections of the previous tutorials, we have to access the correct columns of our simulation results by indexing as there is no name provided.
     """)
@@ -272,15 +272,23 @@ def _(mo):
 
 app._unparsable_cell(
     r"""
-    def distance_function(simulation, real_data):
+    def distance_ICU(simulation, real_data):
         # TODO: Extract the correct columns from the data dataframe:
         real_ICU = 
-        real_Deaths = 
         sim = simulation['data']
         sim_ICU = sim[1+ int(osecir.InfectionState.InfectedCritical), :]
+        # TODO: Add the code to calculate the absolute differences between data and simulation 
+        return np.sum()
+
+    def distance_Deaths(simulation, real_data):
+        # TODO: Extract the correct columns from the data dataframe:
+        real_Deaths =
+        sim = simulation['data']
         sim_Death = sim[1 + int(osecir.InfectionState.Dead), :]
         # TODO: Add the code to calculate the absolute differences between data and simulation 
-        return (np.sum( ) + np.sum( )) / 1000
+        return np.sum()
+
+    distance = pyabc.AdaptiveAggregatedDistance([distance_ICU, distance_Deaths], adaptive=False, scale_function=pyabc.distance.median)
     """,
     name="_"
 )
@@ -325,7 +333,7 @@ def _(example_results):
 
 @app.cell
 def _():
-    # TODO: Test the distance function on the example_results and the observation_data
+    # TODO: Test the distance function on the example_results and the observation_data. You need to specify a time t=-1 to avoid an error later.
 
     return
 
@@ -362,15 +370,15 @@ def _(mo):
 
 @app.cell
 def _(
-    distance_function,
+    distance,
     observation_data,
     os,
     prior,
     pyabc,
     run_simulation,
     tempfile,
 ):
-    abc = pyabc.ABCSMC(run_simulation, prior, distance_function, population_size=400)
+    abc = pyabc.ABCSMC(run_simulation, prior, distance, population_size=400)
     db_path = "sqlite:///" + os.path.join(tempfile.gettempdir(), "tmp.db")
     abc.new(db_path, observation_data)
     return
@@ -424,14 +432,16 @@ def _(mo):
 @app.cell
 def _(np, osecir):
     def plot_critical_data(sum_stat, weight, ax, **kwargs):
-        ax.plot(range(0, 31), sum_stat['data'][1+ int(osecir.InfectionState.InfectedCritical), :], color = 'grey', alpha = 0.1)
+        ax.plot(range(0, 31), sum_stat['data'][1 + int(
+            osecir.InfectionState.InfectedCritical), :], color='grey', alpha=0.1)
 
     def plot_critical_mean(sum_stats, weights, ax, **kwargs):
         weights = np.array(weights)
         weights /= weights.sum()
-        data = np.array([sum_stat['data'][1+ int(osecir.InfectionState.InfectedCritical), :] for sum_stat in sum_stats])
+        data = np.array([sum_stat['data'][1 + int(osecir.InfectionState.InfectedCritical), :]
+                        for sum_stat in sum_stats])
         mean = (data * weights.reshape((-1, 1))).sum(axis=0)
-        ax.plot(range(0, 31), mean, color='C2', label = "Simulation mean")
+        ax.plot(range(0, 31), mean, color='C2', label="Simulation mean")
 
     return plot_critical_data, plot_critical_mean
 
@@ -446,9 +456,11 @@ def _(
     pyabc,
 ):
     fig, ax = plt.subplots()
-    ax = pyabc.visualization.plot_data_callback(history, plot_critical_data, plot_critical_mean, ax=ax)
+    ax = pyabc.visualization.plot_data_callback(
+        history, plot_critical_data, plot_critical_mean, ax=ax)
 
-    plt.scatter(range(0, 31), observation_data["Critical"], color = "C1", label = "Data", zorder = 2)
+    plt.scatter(
+        range(0, 31), observation_data["Critical"], color="C1", label="Data", zorder=2)
     plt.xlabel("Time")
     plt.ylabel("# Cases")
     plt.title("Number of ICU patients")
@@ -467,9 +479,11 @@ def _():
 @app.cell
 def _(history, observation_data, plot_dead_data, plot_dead_mean, plt, pyabc):
     fig_dead, ax_dead = plt.subplots()
-    ax_dead = pyabc.visualization.plot_data_callback(history, plot_dead_data, plot_dead_mean, ax=ax_dead)
+    ax_dead = pyabc.visualization.plot_data_callback(
+        history, plot_dead_data, plot_dead_mean, ax=ax_dead)
 
-    plt.scatter(range(0, 31), observation_data["Deaths"], color = "C1", label = "Data", zorder = 2)
+    plt.scatter(
+        range(0, 31), observation_data["Deaths"], color="C1", label="Data", zorder=2)
     plt.xlabel("Time")
     plt.ylabel("# Cases")
     plt.title("Cumulative number of dead patients")
@@ -504,34 +518,40 @@ def _():
 
 @app.cell
 def _(df, run_simulation_with_params):
-    predictions = df.apply(run_simulation_with_params, axis = 1)
+    predictions = df.apply(run_simulation_with_params, axis=1)
     return (predictions,)
 
 
 @app.cell
 def _(np, observation_data, osecir, plt, predictions, w):
-    fig_prog, ax_prog = plt.subplots(1, 2, figsize = (12,4))
-    ax_prog[0].scatter(range(0, 31), observation_data["Deaths"], color = "C1", label = "Data")
+    fig_prog, ax_prog = plt.subplots(1, 2, figsize=(12, 4))
+    ax_prog[0].scatter(
+        range(0, 31), observation_data["Deaths"], color="C1", label="Data")
     for prediction in predictions:
-        ax_prog[0].plot(range(61), prediction["data"][1+int(osecir.InfectionState.Dead)], color = "grey", alpha = 0.2)
+        ax_prog[0].plot(range(61), prediction["data"]
+                        [1+int(osecir.InfectionState.Dead)], color="grey", alpha=0.2)
     weights = np.array(w)
     weights /= weights.sum()
-    data = np.array([prediction['data'][1+ int(osecir.InfectionState.Dead), :] for prediction in predictions])
+    data = np.array([prediction['data'][1 + int(osecir.InfectionState.Dead), :]
+                    for prediction in predictions])
     mean = (data * weights.reshape((-1, 1))).sum(axis=0)
-    ax_prog[0].plot(range(0, 61), mean, color='C2', label = "Simulation mean")
+    ax_prog[0].plot(range(0, 61), mean, color='C2', label="Simulation mean")
     ax_prog[0].legend()
     ax_prog[0].set_title("Number of Deaths")
     ax_prog[0].set_ylabel("Number")
     ax_prog[0].set_xlabel("Time [d]")
 
-    ax_prog[1].scatter(range(0, 31), observation_data["Critical"], color = "C1", label = "Data")
+    ax_prog[1].scatter(
+        range(0, 31), observation_data["Critical"], color="C1", label="Data")
     for prediction in predictions:
-        ax_prog[1].plot(range(61), prediction["data"][1+int(osecir.InfectionState.InfectedCritical)], color = "grey", alpha = 0.2)
+        ax_prog[1].plot(range(61), prediction["data"]
+                        [1+int(osecir.InfectionState.InfectedCritical)], color="grey", alpha=0.2)
     weights = np.array(w)
     weights /= weights.sum()
-    data = np.array([prediction['data'][1+ int(osecir.InfectionState.InfectedCritical), :] for prediction in predictions])
+    data = np.array([prediction['data'][1 + int(osecir.InfectionState.InfectedCritical), :]
+                    for prediction in predictions])
     mean = (data * weights.reshape((-1, 1))).sum(axis=0)
-    ax_prog[1].plot(range(0, 61), mean, color='C2', label = "Simulation mean")
+    ax_prog[1].plot(range(0, 61), mean, color='C2', label="Simulation mean")
     ax_prog[1].legend()
     ax_prog[1].set_title("Number of ICU cases")
     ax_prog[1].set_ylabel("Number")