add inspection mode

forecasttools/sbc.py

@@ -18,6 +18,7 @@ def __init__(
         num_simulations=10,
         sample_kwargs=None,
         seed=None,
+        inspection_mode=False,
         **kwargs,
     ) -> None:
         """
@@ -71,13 +72,18 @@ def __init__(
 
         self.num_simulations = num_simulations
         self.sample_kwargs = sample_kwargs
-
+        # Initialize the simulations and random seeds
         self.simulations = {}
         self._simulations_complete = 0
         prior_pred_rng, sampler_rng = random.split(seed)
         self._prior_pred_rng = prior_pred_rng
         self._sampler_rng = sampler_rng
         self.num_samples = None
+        # Set the inspection mode
+        # if in inspection mode, store all idata objects from fitting
+        self.inspection_mode = inspection_mode
+        if inspection_mode:
+            self.idatas = []
 
     def _get_prior_predictive_samples(
         self,
@@ -185,6 +191,9 @@ def run_simulations(self) -> None:
             initial=self._simulations_complete,
             total=self.num_simulations,
         )
+        if self.inspection_mode:
+            self.prior = prior
+            self.prior_pred = prior_pred
         try:
             while self._simulations_complete < self.num_simulations:
                 idx = self._simulations_complete
@@ -195,6 +204,8 @@ def run_simulations(self) -> None:
                 idata = self._get_posterior_samples(
                     sampler_seeds[idx], prior_predictive_draw
                 )
+                if self.inspection_mode:
+                    self.idatas.append(idata)
                 self._increment_rank_statistics(prior_draw, idata["posterior"])
                 self._simulations_complete += 1
                 progress.update()

tests/test_sbc.py

@@ -41,6 +41,13 @@ def sbc_instance(mcmc_kernel, observed_vars):
     return SBC(mcmc_kernel, y=None, observed_vars=observed_vars)
 
 
+@pytest.fixture
+def sbc_instance_inspection_on(mcmc_kernel, observed_vars):
+    return SBC(
+        mcmc_kernel, y=None, observed_vars=observed_vars, inspection_mode=True
+    )
+
+
 def test_sbc_initialization(sbc_instance, mcmc_kernel, observed_vars):
     """
     Test that the SBC class is initialized correctly.
@@ -98,6 +105,16 @@ def test_run_simulations(sbc_instance):
     assert "mu" in sbc_instance.simulations
 
 
+def test_run_simulations_with_inspection(sbc_instance_inspection_on):
+    """
+    Test that the simulations for SBC are run correctly.
+    """
+    sbc_instance_inspection_on.run_simulations()
+    assert isinstance(sbc_instance_inspection_on.idatas, list)
+    assert isinstance(sbc_instance_inspection_on.prior, dict)
+    assert isinstance(sbc_instance_inspection_on.prior_pred, dict)
+
+
 def test_plot_results(sbc_instance):
     """
     Test that the results are plotted.