From 19da4f22380407acf201305359c734117cd5e126 Mon Sep 17 00:00:00 2001
From: Romain Ragonnet <romain.ragonnet@monash.edu>
Date: Wed, 1 Nov 2023 14:52:54 +1100
Subject: [PATCH] Track incidence and cum-incidence by strain

---
 autumn/models/sm_covid2/model.py   |  5 +++-
 autumn/models/sm_covid2/outputs.py | 44 ++++++++++++++++++++++++++++--
 2 files changed, 46 insertions(+), 3 deletions(-)

diff --git a/autumn/models/sm_covid2/model.py b/autumn/models/sm_covid2/model.py
index 2156497aa..0daa14c97 100644
--- a/autumn/models/sm_covid2/model.py
+++ b/autumn/models/sm_covid2/model.py
@@ -523,9 +523,12 @@ def build_model(params: dict, build_options: dict = None, ret_builder=False) ->
     )
 
     outputs_builder.request_cumulative_outputs(
-        params.requested_cumulative_outputs, params.cumulative_start_time
+        params.requested_cumulative_outputs, params.cumulative_start_time, strain_strata
     )
 
+    if "incidence" in params.requested_cumulative_outputs:
+        outputs_builder.request_cumulative_incidence_prop_by_strain(strain_strata)
+
     if params.activate_random_process:
         outputs_builder.request_random_process_outputs()
         outputs_builder.request_random_process_auc()
diff --git a/autumn/models/sm_covid2/outputs.py b/autumn/models/sm_covid2/outputs.py
index d78605dd1..f452144b0 100644
--- a/autumn/models/sm_covid2/outputs.py
+++ b/autumn/models/sm_covid2/outputs.py
@@ -73,6 +73,7 @@ def request_incidence(
         strain_strata: List[str],
         incidence_flow: str,
         request_incidence_by_age: bool,
+        request_incidence_by_strain: bool = True,
     ):
         """
         Calculate incident disease cases. This is associated with the transition to infectiousness if there is only one
@@ -92,6 +93,7 @@ def request_incidence(
         self.model.request_output_for_flow(name="incidence", flow_name=incidence_flow)
 
         # Stratified
+        strain_incidence_sources = {strain: [] for strain in strain_strata}
         for agegroup in age_groups:
             agegroup_string = f"Xagegroup_{agegroup}"
             age_incidence_sources = []
@@ -108,6 +110,7 @@ def request_incidence(
 
                     output_name = f"incidence{agegroup_string}{immunity_string}{strain_string}"
                     age_incidence_sources.append(output_name)
+                    strain_incidence_sources[strain].append(output_name)
 
                     self.model.request_output_for_flow(
                         name=output_name,
@@ -124,6 +127,21 @@ def request_incidence(
                     save_results=True,
                 )
 
+
+        # Aggregated incidence by strain
+        if request_incidence_by_strain:
+
+            for strain in strain_strata:                 
+                strain_string = f"Xstrain_{strain}" if strain else ""
+                output_name = f"incidence{strain_string}"
+
+                self.model.request_aggregate_output(
+                    name=output_name,
+                    sources=strain_incidence_sources[strain],
+                    save_results=False,
+                )       
+            
+
     def request_elderly_incidence_prop(
         self,
         age_groups: List[str],
@@ -136,6 +154,21 @@ def request_elderly_incidence_prop(
             save_results=True
         )
 
+    def request_cumulative_incidence_prop_by_strain(
+        self,
+        strain_strata: List[str],
+    ):
+        for strain in strain_strata: 
+            strain_string = f"Xstrain_{strain}" if strain else ""
+            cum_output_name = f"cumulative_incidence{strain_string}"
+
+            self.model.request_function_output(
+                name=f"cumulative_incidence_prop{strain_string}",
+                func=DerivedOutput(cum_output_name) / DerivedOutput("cumulative_incidence"),
+                save_results=True
+            )
+
+
     def request_infection_deaths(
         self,
         model_times: np.ndarray,
@@ -647,7 +680,7 @@ def request_immunity_props(self, immunity_strata, age_pops, request_immune_prop_
                         func=make_age_immune_prop_func(popsize)(DerivedOutput(n_age_immune_name))
                     )
 
-    def request_cumulative_outputs(self, requested_cumulative_outputs, cumulative_start_time):
+    def request_cumulative_outputs(self, requested_cumulative_outputs, cumulative_start_time, strain_strata):
         """
         Compute cumulative outputs for requested outputs.
 
@@ -655,8 +688,15 @@ def request_cumulative_outputs(self, requested_cumulative_outputs, cumulative_st
             requested_cumulative_outputs: List of requested derived outputs to accumulate
             cumulative_start_time: reference time for cumulative output calculation
         """
+        computed_cumulative_outputs = requested_cumulative_outputs
+
+        # also request cumulative incidence by strain
+        for strain in strain_strata:                 
+            strain_string = f"Xstrain_{strain}" if strain else ""
+            output_name = f"incidence{strain_string}"
+            computed_cumulative_outputs.append(output_name)
 
-        for output in requested_cumulative_outputs:
+        for output in computed_cumulative_outputs:
             self.model.request_cumulative_output(
                 name=f"cumulative_{output}", source=output, start_time=cumulative_start_time
             )