adding some blobs tests

docs/user/upgrade.rst

@@ -0,0 +1,9 @@
+.. _upgrade:
+
+.. module:: emcee
+
+Upgrading From Pre-0.3 Versions
+===============================
+
+The v0.3 release of emcee is the biggest update in years.
+

emcee/ensemble.py

@@ -364,6 +364,12 @@ def compute_log_prob(self, coords=None):
             log_prob = np.array([float(l[0]) for l in results])
             blob = [l[1:] for l in results]
             blob = np.array(blob, dtype=np.atleast_1d(blob[0]).dtype)
+
+            # Deal with single blobs in a better way
+            if blob.shape[1] == 1:
+                m = [slice(None) for i in range(len(blob.shape))]
+                m[1] = 0
+                blob = blob[m]
         except (IndexError, TypeError):
             log_prob = np.array([float(l) for l in results])
             blob = None

tests/unit/test_blobs.py

@@ -0,0 +1,59 @@
+# -*- coding: utf-8 -*-
+
+from __future__ import division, print_function
+
+import pytest
+import numpy as np
+
+from emcee import backends, EnsembleSampler
+
+__all__ = ["test_blob_shape"]
+
+
+class BlobLogProb(object):
+
+    def __init__(self, blob_function):
+        self.blob_function = blob_function
+
+    def __call__(self, params):
+        return -0.5 * np.sum(params**2), self.blob_function(params)
+
+
+@pytest.mark.parametrize(
+    "backend", [backends.Backend, backends.hdf.TempHDFBackend],
+)
+def test_blob_shape(backend):
+    with backend() as be:
+        np.random.seed(42)
+
+        nblobs = 5
+        model = BlobLogProb(lambda x: np.random.randn(nblobs))
+
+        coords = np.random.randn(32, 3)
+        nwalkers, ndim = coords.shape
+        sampler = EnsembleSampler(nwalkers, ndim, model, backend=be)
+        nsteps = 10
+        sampler.run_mcmc(coords, nsteps)
+        assert sampler.get_blobs().shape == (nsteps, nwalkers, nblobs)
+
+        model = BlobLogProb(lambda x: np.random.randn())
+        be.reset(nwalkers, ndim)
+        sampler = EnsembleSampler(nwalkers, ndim, model, backend=be)
+        sampler.run_mcmc(coords, nsteps)
+        assert sampler.get_blobs().shape == (nsteps, nwalkers)
+
+        # HDF backend doesn't support the object type
+        if backend == backends.hdf.TempHDFBackend:
+            return
+
+        model = BlobLogProb(lambda x: "face")
+        be.reset(nwalkers, ndim)
+        sampler = EnsembleSampler(nwalkers, ndim, model, backend=be)
+        sampler.run_mcmc(coords, nsteps)
+        assert sampler.get_blobs().shape == (nsteps, nwalkers)
+
+        model = BlobLogProb(lambda x: (np.random.randn(nblobs), "face"))
+        be.reset(nwalkers, ndim)
+        sampler = EnsembleSampler(nwalkers, ndim, model, backend=be)
+        sampler.run_mcmc(coords, nsteps)
+        assert sampler.get_blobs().shape == (nsteps, nwalkers, 2)