BUG(emcee): fix blob ordering

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

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()