Update arviz data extracted from MCMC results

CHANGELOG.rst

@@ -1,6 +1,7 @@
 Changelog
 =========
 
+- Update BOLFI and BOLFIRE to use a shared sample class that returns individual chains in the arviz inference data
 - Use kernel copy to avoid pickle issue and allow BOLFI parallelisation with non-default kernel
 - Restrict matplotlib version < 3.9 for compatibility with GPy
 - Add option to use additive or multiplicative adjustment in any acquisition method

docs/api.rst

@@ -65,7 +65,7 @@ Below is a list of inference methods included in ELFI.
    OptimizationResult
    Sample
    SmcSample
-   BolfiSample
+   McmcSample
 
 
 **Post-processing**
@@ -244,7 +244,7 @@ Inference API classes
    :members:
    :inherited-members:
 
-.. autoclass:: BolfiSample
+.. autoclass:: McmcSample
    :members:
    :inherited-members:
 

elfi/methods/inference/bolfi.py

@@ -16,7 +16,7 @@
 from elfi.methods.bo.utils import stochastic_optimization
 from elfi.methods.inference.parameter_inference import ParameterInference
 from elfi.methods.posteriors import BolfiPosterior
-from elfi.methods.results import BolfiSample, OptimizationResult
+from elfi.methods.results import McmcSample, OptimizationResult
 from elfi.methods.utils import arr2d_to_batch, batch_to_arr2d, ceil_to_batch_size, resolve_sigmas
 from elfi.model.extensions import ModelPrior
 
@@ -507,7 +507,7 @@ def sample(self,
 
         Returns
         -------
-        BolfiSample
+        McmcSample
 
         """
         if self.state['n_batches'] == 0:
@@ -588,7 +588,7 @@ def sample(self,
                   mcmc.gelman_rubin_statistic(chains[:, :, ii]))
         self.target_model.is_sampling = False
 
-        return BolfiSample(
+        return McmcSample(
             method_name='BOLFI',
             chains=chains,
             parameter_names=self.target_model.parameter_names,

elfi/methods/inference/bolfire.py

@@ -14,7 +14,7 @@
 from elfi.methods.classifier import Classifier, LogisticRegression
 from elfi.methods.inference.parameter_inference import ModelBased
 from elfi.methods.posteriors import BOLFIREPosterior
-from elfi.methods.results import BOLFIRESample
+from elfi.methods.results import McmcSample
 from elfi.methods.utils import batch_to_arr2d, resolve_sigmas
 from elfi.model.extensions import ModelPrior
 
@@ -201,7 +201,7 @@ def sample(self,
 
         Returns
         -------
-        BOLFIRESample
+        McmcSample
 
         """
         # Fit posterior in case not done
@@ -282,13 +282,13 @@ def sample(self,
 
         self.target_model.is_sampling = False
 
-        return BOLFIRESample(method_name='BOLFIRE',
-                             chains=chains,
-                             parameter_names=self.parameter_names,
-                             warmup=warmup,
-                             n_sim=self.state['n_sim'],
-                             seed=self.seed,
-                             *args, **kwargs)
+        return McmcSample(method_name='BOLFIRE',
+                          chains=chains,
+                          parameter_names=self.parameter_names,
+                          warmup=warmup,
+                          n_sim=self.state['n_sim'],
+                          seed=self.seed,
+                          *args, **kwargs)
 
     def _resolve_marginal(self, marginal, seed_marginal=None):
         """Resolve marginal data."""

elfi/methods/results.py

@@ -504,8 +504,8 @@ def plot_pairs(self, selector=None, bins=20, axes=None, all=False, **kwargs):
             plt.suptitle("Population {}".format(i), fontsize=fontsize)
 
 
-class BolfiSample(Sample):
-    """Container for results from BOLFI."""
+class McmcSample(Sample):
+    """Container for MCMC results."""
 
     def __init__(self, method_name, chains, parameter_names, warmup, **kwargs):
         """Initialize result.
@@ -529,7 +529,7 @@ def __init__(self, method_name, chains, parameter_names, warmup, **kwargs):
         concatenated = warmed_up.reshape((-1,) + shape[2:])
         outputs = dict(zip(parameter_names, concatenated.T))
 
-        super(BolfiSample, self).__init__(
+        super(McmcSample, self).__init__(
             method_name=method_name,
             outputs=outputs,
             parameter_names=parameter_names,
@@ -538,6 +538,13 @@ def __init__(self, method_name, chains, parameter_names, warmup, **kwargs):
             warmup=warmup,
             **kwargs)
 
+    @property
+    def idata(self):
+        """Convert MCMC chains to arviz InferenceData object."""
+        warmed_up = self.chains[:, self.warmup:]
+        sample_chains = dict(zip(self.parameter_names, np.transpose(warmed_up, (2, 0, 1))))
+        return az.from_dict(sample_chains)
+
     def plot_traces(self, selector=None, axes=None, **kwargs):
         """Plot MCMC traces."""
         return vis.plot_traces(self, selector, axes, **kwargs)
@@ -605,40 +612,6 @@ def compute_ess(self):
         return {p: eff_sample_size(self.samples[p]) for p in self.parameter_names}
 
 
-class BOLFIRESample(Sample):
-    """Container for results from BOLFIRE."""
-
-    def __init__(self, method_name, chains, parameter_names, warmup, *args, **kwargs):
-        """Initialize BOLFIRE result.
-
-        Parameters
-        ----------
-        method_name: str
-            Name of the inference method.
-        chains: np.ndarray (n_chains, n_samples, n_parameters)
-            Chains from sampling, warmup included.
-        parameter_names: list
-            List of names in the outputs dict that refer to model parameters.
-        warmup: int
-            Number of warmup iterations in chains.
-
-        """
-        n_chains = chains.shape[0]
-        warmed_up = chains[:, warmup:, :]
-        concatenated = warmed_up.reshape((-1,) + chains.shape[2:])
-        outputs = dict(zip(parameter_names, concatenated.T))
-
-        super(BOLFIRESample, self).__init__(
-            method_name=method_name,
-            outputs=outputs,
-            parameter_names=parameter_names,
-            chains=chains,
-            n_chains=n_chains,
-            warmup=warmed_up,
-            *args, **kwargs
-        )
-
-
 class RomcSample(Sample):
     """Container for results from ROMC."""
 

tests/unit/test_results.py

@@ -43,14 +43,14 @@ def test_sample():
     sample.summary()
 
 
-def test_bolfi_sample():
+def test_mcmc_sample():
     n_chains = 3
     n_iters = 10
     warmup = 5
     parameter_names = ['a', 'b']
     chains = np.random.random((n_chains, n_iters, len(parameter_names)))
 
-    result = elfi.methods.results.BolfiSample(
+    result = elfi.methods.results.McmcSample(
         method_name="TestRes",
         chains=chains,
         parameter_names=parameter_names,