From e39f961e5d9d702f827517b47dc441c03860f817 Mon Sep 17 00:00:00 2001
From: Timothy Willard <9395586+TimothyWillard@users.noreply.github.com>
Date: Tue, 17 Dec 2024 12:01:57 -0500
Subject: [PATCH] Add `Parameters.reinitialize_distributions` method

Added a method to reinit distribution parameters. When the `Parameters`
class is initialized it creates this distribution parameters but also
captures the random seed at the time of initialization which makes the
draws non-random across processes. Added a method to reinitialize
distribution parameters when reseeding. Not a perfect solution, the
ambigious random seed settings is rather confusing for development.
---
 flepimop/gempyor_pkg/src/gempyor/parameters.py       | 12 ++++++++++++
 flepimop/gempyor_pkg/src/gempyor/seir.py             |  2 ++
 .../test_random_seir_parameter_draw_per_slot.py      |  7 +++++--
 3 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/flepimop/gempyor_pkg/src/gempyor/parameters.py b/flepimop/gempyor_pkg/src/gempyor/parameters.py
index 7ab0de151..780a55119 100644
--- a/flepimop/gempyor_pkg/src/gempyor/parameters.py
+++ b/flepimop/gempyor_pkg/src/gempyor/parameters.py
@@ -165,6 +165,18 @@ def __init__(
         logging.debug(f"Index in arrays are: {self.pnames2pindex}")
         logging.debug(f"NPI overlap operation is {self.stacked_modifier_method} ")
 
+    def reinitialize_distributions(self) -> None:
+        """
+        Reinitialize all random distributions.
+
+        This method reinitializes all random distributions for the parameters contained
+        within this class. The random seed for each distribution is captured on
+        initialization so this method will change those seeds.
+        """
+        for pn in self.pnames:
+            if "dist" in self.pdata[pn]:
+                self.pdata[pn]["dist"] = self.pconfig[pn]["value"].as_random_distribution()
+
     def picklable_lamda_alpha(self):
         """
         Read the `alpha_val` attribute.
diff --git a/flepimop/gempyor_pkg/src/gempyor/seir.py b/flepimop/gempyor_pkg/src/gempyor/seir.py
index 5b76860ab..851b5f023 100644
--- a/flepimop/gempyor_pkg/src/gempyor/seir.py
+++ b/flepimop/gempyor_pkg/src/gempyor/seir.py
@@ -258,6 +258,8 @@ def onerun_SEIR(
     config=None,
 ):
     np.random.seed()
+    modinf.parameters.reinitialize_distributions()
+
     npi = None
     if modinf.npi_config_seir:
         npi = build_npi_SEIR(
diff --git a/flepimop/gempyor_pkg/tests/simulate/test_random_seir_parameter_draw_per_slot.py b/flepimop/gempyor_pkg/tests/simulate/test_random_seir_parameter_draw_per_slot.py
index c12235195..d9864c6ef 100644
--- a/flepimop/gempyor_pkg/tests/simulate/test_random_seir_parameter_draw_per_slot.py
+++ b/flepimop/gempyor_pkg/tests/simulate/test_random_seir_parameter_draw_per_slot.py
@@ -25,15 +25,18 @@ def setup_sample_2pop_vaccine_scenarios(tmp_path: Path) -> Path:
     return tmp_path
 
 
+@pytest.mark.parametrize("n_jobs", [1, 2])
 def test_random_seir_parameter_draw_per_slot(
-    monkeypatch: pytest.MonkeyPatch, setup_sample_2pop_vaccine_scenarios: Path
+    monkeypatch: pytest.MonkeyPatch, setup_sample_2pop_vaccine_scenarios: Path, n_jobs: int
 ) -> None:
     # Test setup
     monkeypatch.chdir(setup_sample_2pop_vaccine_scenarios)
 
     # Test execution of `gempyor-simulate`
     runner = CliRunner()
-    result = runner.invoke(_click_simulate, ["config_sample_2pop_vaccine_scenarios.yml"])
+    result = runner.invoke(
+        _click_simulate, ["config_sample_2pop_vaccine_scenarios.yml", "--jobs", str(n_jobs)]
+    )
     assert result.exit_code == 0
 
     # Get the contents of 'spar' and 'hpar' directories as DataFrames