From e3cb5379dc5ba7a92cd05b1d0c5fdd0be5fcd7c4 Mon Sep 17 00:00:00 2001
From: Timothy Willard <9395586+TimothyWillard@users.noreply.github.com>
Date: Thu, 21 Nov 2024 17:20:21 -0500
Subject: [PATCH] Bug fix `compute_all_multioutcomes`

Added a bug fix to `compute_all_multioutcomes` to look for if a sourced
outcome exists and has not been processed yet and if so skip that
outcome for now and re-add it to the end of the outcomes to process.
Required for the `test_outcomes_pcomp` and
`test_outcomes_pcomp_read_write` tests to pass. Also minor edits to
`config_mc_selection.yml` to add the required 'compartments'/'seir'
sections.
---
 flepimop/gempyor_pkg/src/gempyor/outcomes.py  |  6 ++-
 .../tests/outcomes/config_mc_selection.yml    | 45 +++++++++++++++++++
 2 files changed, 50 insertions(+), 1 deletion(-)

diff --git a/flepimop/gempyor_pkg/src/gempyor/outcomes.py b/flepimop/gempyor_pkg/src/gempyor/outcomes.py
index edf0b6cc5..7303fffee 100644
--- a/flepimop/gempyor_pkg/src/gempyor/outcomes.py
+++ b/flepimop/gempyor_pkg/src/gempyor/outcomes.py
@@ -369,7 +369,8 @@ def compute_all_multioutcomes(
     else:
         seir_sim = bypass_seir_df
 
-    for new_comp in parameters:
+    parameters_keys = list(parameters.keys())
+    for new_comp in parameters_keys:
         if "source" in parameters[new_comp]:
             # Read the config for this compartment: if a source is specified, we
             # 1. compute incidence from binomial draw
@@ -400,6 +401,9 @@ def compute_all_multioutcomes(
             else:  # already defined outcomes
                 if source_name in all_data:
                     source_array = all_data[source_name]
+                elif source_name in parameters_keys:
+                    parameters_keys.append(new_comp)
+                    continue
                 else:
                     raise ValueError(
                         f"ERROR with outcome {new_comp}: the specified source {source_name} is not a dictionnary (for seir outcome) nor an existing pre-identified outcomes."
diff --git a/flepimop/gempyor_pkg/tests/outcomes/config_mc_selection.yml b/flepimop/gempyor_pkg/tests/outcomes/config_mc_selection.yml
index 717135674..68d0d8764 100644
--- a/flepimop/gempyor_pkg/tests/outcomes/config_mc_selection.yml
+++ b/flepimop/gempyor_pkg/tests/outcomes/config_mc_selection.yml
@@ -8,6 +8,51 @@ subpop_setup:
   geodata: data/geodata.csv
 
 
+compartments:
+  infection_stage:
+    - S
+    - I1
+    - R
+  vaccination_stage:
+    - unvaccinated
+    - first_dose
+
+
+seir:
+  integration:
+    method: rk4
+    dt: 1
+  transitions:
+    - source: [
+        ["S"],
+        ["unvaccinated", "first_dose"]
+      ]
+      destination: [
+        ["I1"],
+        ["unvaccinated", "first_dose"]
+      ]
+      proportional_to: [
+        "source",
+        [
+          [
+            ["I1"]
+          ],
+          [
+            ["unvaccinated", "first_dose"],
+            ["unvaccinated"],
+          ],
+        ]
+      ]
+      rate: [
+        [0.85],
+        [1, 0.5]
+      ]
+      proportion_exponent: [
+        [1, [1, 1]],
+        [1, [1, 1]]
+      ]
+
+
 outcomes:
   method: delayframe
   param_from_file: False