Track random process AUC

autumn/models/sm_covid2/model.py

@@ -526,6 +526,7 @@ def build_model(params: dict, build_options: dict = None, ret_builder=False) ->
 
     if params.activate_random_process:
         outputs_builder.request_random_process_outputs()
+        outputs_builder.request_random_process_auc()
 
     # request extra output to store the number of students*weeks of school missed
     outputs_builder.request_student_weeks_missed_output(student_weeks_missed)

autumn/models/sm_covid2/outputs.py

@@ -405,6 +405,21 @@ def array_max(x):
             func=peak_func
         )
 
+    def request_random_process_auc(self):
+        """
+        Create an output to calculate the area between the (transformed) random process and the horizontal line y=1.
+        """
+
+        def sum_diffs_to_one(x):
+            return jnp.repeat(jnp.sum(jnp.abs(1 - x)), jnp.size(x))
+
+        sum_diffs_to_one_func = Function(sum_diffs_to_one, [DerivedOutput("transformed_random_process")])
+
+        self.model.request_function_output(
+            "random_process_auc",
+            func=sum_diffs_to_one_func
+        )
+
     # def request_icu_outputs(
     #     self,
     #     prop_icu_among_hospitalised: float,

autumn/projects/sm_covid2/common_school/runner_tools.py

@@ -273,14 +273,15 @@ def get_quantile_outputs(outputs_df, diff_outputs_df, quantiles=[.025, .25, .5,
     times = sorted(outputs_df.index.unique())
     scenarios = outputs_df["scenario"].unique()
     unc_output_names = [
-        "infection_deaths_ma7", "prop_ever_infected_age_matched", "prop_ever_infected", "transformed_random_process", "cumulative_incidence", "cumulative_infection_deaths",
-        "peak_hospital_occupancy", "hospital_occupancy"
+        "infection_deaths_ma7", "prop_ever_infected_age_matched", "prop_ever_infected", "transformed_random_process", "random_process_auc",
+        "cumulative_incidence", "cumulative_infection_deaths", "peak_hospital_occupancy", "hospital_occupancy"
     ]
 
     uncertainty_data = []
     for scenario in scenarios:
         scenario_mask = outputs_df["scenario"] == scenario
         scenario_df = outputs_df[scenario_mask]
+
         for time in times:
             masked_df = scenario_df.loc[time]
             if masked_df.empty: