diff --git a/bilby/core/sampler/emcee.py b/bilby/core/sampler/emcee.py index 21b47ef38..0592edc4d 100644 --- a/bilby/core/sampler/emcee.py +++ b/bilby/core/sampler/emcee.py @@ -458,7 +458,10 @@ def _generate_result(self): " Try increasing the number of steps." ) blobs = np.array(self.sampler.blobs) - blobs_trimmed = blobs[self.nburn :, :, :].reshape((-1, 2)) + # Blobs are stored (nsteps, nwalkers, 2) + # Transpose to match the shape of the chain (nwalkers, nsteps, 2) + # and then flatten + blobs_trimmed = blobs[self.nburn :, :, :].transpose(1, 0, 2).reshape((-1, 2)) log_likelihoods, log_priors = blobs_trimmed.T self.result.log_likelihood_evaluations = log_likelihoods self.result.log_prior_evaluations = log_priors diff --git a/test/core/sampler/emcee_test.py b/test/core/sampler/emcee_test.py index ed1ac87da..556adb482 100644 --- a/test/core/sampler/emcee_test.py +++ b/test/core/sampler/emcee_test.py @@ -1,7 +1,9 @@ import unittest +from types import SimpleNamespace import bilby import bilby.core.sampler.emcee +import numpy as np class TestEmcee(unittest.TestCase): @@ -75,13 +77,33 @@ def test_expected_output_files(self): "outdir/emcee_output_test/chain.dat", "outdir/emcee_output_test/sampler.pickle", ] - expected_dirs = [ - "outdir/emcee_output_test" - ] - filenames, dirs = self.sampler.get_expected_outputs(outdir="outdir", label="output_test") + expected_dirs = ["outdir/emcee_output_test"] + filenames, dirs = self.sampler.get_expected_outputs( + outdir="outdir", label="output_test" + ) self.assertListEqual(expected_filenames, filenames) self.assertListEqual(expected_dirs, dirs) + def test_generate_result_flattens_blobs_in_sample_order(self): + chain = np.arange(12).reshape(2, 3, 2) + blobs = np.empty((3, 2, 2)) + blobs[:, :, 0] = chain[:, :, 0].T + blobs[:, :, 1] = chain[:, :, 1].T + self.sampler._sampler = SimpleNamespace(chain=chain, blobs=blobs) + self.sampler.nburn = 1 + self.sampler.result.samples = chain[:, self.sampler.nburn :, :].reshape(-1, 2) + + self.sampler._generate_result() + + np.testing.assert_array_equal( + self.sampler.result.log_likelihood_evaluations, + self.sampler.result.samples[:, 0], + ) + np.testing.assert_array_equal( + self.sampler.result.log_prior_evaluations, + self.sampler.result.samples[:, 1], + ) + if __name__ == "__main__": unittest.main()