From aa4607b1db15fa148bd534271ca5be4182add371 Mon Sep 17 00:00:00 2001 From: Ulpu Remes Date: Fri, 28 Jun 2024 00:28:33 +0300 Subject: [PATCH 1/3] make mcmc sample return idata with chains --- elfi/methods/inference/bolfi.py | 6 ++-- elfi/methods/inference/bolfire.py | 18 ++++++------ elfi/methods/results.py | 47 +++++++------------------------ 3 files changed, 22 insertions(+), 49 deletions(-) diff --git a/elfi/methods/inference/bolfi.py b/elfi/methods/inference/bolfi.py index 07f878f9..57be6aa5 100644 --- a/elfi/methods/inference/bolfi.py +++ b/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, diff --git a/elfi/methods/inference/bolfire.py b/elfi/methods/inference/bolfire.py index 4ab51528..2d3f0d1d 100644 --- a/elfi/methods/inference/bolfire.py +++ b/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.""" diff --git a/elfi/methods/results.py b/elfi/methods/results.py index b28b8566..f100cd56 100644 --- a/elfi/methods/results.py +++ b/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.""" From 73eb08634b81c8db038f9e190520b2f9b6f6ad2d Mon Sep 17 00:00:00 2001 From: Ulpu Remes Date: Fri, 28 Jun 2024 00:38:31 +0300 Subject: [PATCH 2/3] update sample class name --- docs/api.rst | 4 ++-- elfi/methods/inference/bolfi.py | 6 +++--- elfi/methods/inference/bolfire.py | 6 +++--- elfi/methods/results.py | 4 ++-- tests/unit/test_results.py | 4 ++-- 5 files changed, 12 insertions(+), 12 deletions(-) diff --git a/docs/api.rst b/docs/api.rst index e58e71d9..2679de62 100644 --- a/docs/api.rst +++ b/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: diff --git a/elfi/methods/inference/bolfi.py b/elfi/methods/inference/bolfi.py index 57be6aa5..5b1b7fe9 100644 --- a/elfi/methods/inference/bolfi.py +++ b/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 MCMCSample, 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 ------- - MCMCSample + 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 MCMCSample( + return McmcSample( method_name='BOLFI', chains=chains, parameter_names=self.target_model.parameter_names, diff --git a/elfi/methods/inference/bolfire.py b/elfi/methods/inference/bolfire.py index 2d3f0d1d..b809c860 100644 --- a/elfi/methods/inference/bolfire.py +++ b/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 MCMCSample +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 ------- - MCMCSample + McmcSample """ # Fit posterior in case not done @@ -282,7 +282,7 @@ def sample(self, self.target_model.is_sampling = False - return MCMCSample(method_name='BOLFIRE', + return McmcSample(method_name='BOLFIRE', chains=chains, parameter_names=self.parameter_names, warmup=warmup, diff --git a/elfi/methods/results.py b/elfi/methods/results.py index f100cd56..9103571f 100644 --- a/elfi/methods/results.py +++ b/elfi/methods/results.py @@ -504,7 +504,7 @@ def plot_pairs(self, selector=None, bins=20, axes=None, all=False, **kwargs): plt.suptitle("Population {}".format(i), fontsize=fontsize) -class MCMCSample(Sample): +class McmcSample(Sample): """Container for MCMC results.""" def __init__(self, method_name, chains, parameter_names, warmup, **kwargs): @@ -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(MCMCSample, self).__init__( + super(McmcSample, self).__init__( method_name=method_name, outputs=outputs, parameter_names=parameter_names, diff --git a/tests/unit/test_results.py b/tests/unit/test_results.py index cd4fd497..f4a7c7dc 100644 --- a/tests/unit/test_results.py +++ b/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, From 563f168818f91d16fad47db149438ac02b925c15 Mon Sep 17 00:00:00 2001 From: Ulpu Remes Date: Fri, 28 Jun 2024 13:02:07 +0300 Subject: [PATCH 3/3] update changelog --- CHANGELOG.rst | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index aa5fbf90..9eb7ffc7 100644 --- a/CHANGELOG.rst +++ b/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