adding moves interface

emcee/ensemble.py

@@ -13,11 +13,13 @@
 
 __all__ = ["EnsembleSampler"]
 
+import logging
+
 import numpy as np
 
 from . import autocorr
 from .sampler import Sampler
-from .interruptible_pool import InterruptiblePool
+from .moves import StretchMove
 
 
 class EnsembleSampler(Sampler):
@@ -72,22 +74,39 @@ class EnsembleSampler(Sampler):
         can be any object with a ``map`` method that follows the same
         calling sequence as the built-in ``map`` function.
 
-    :param runtime_sortingfn: (optional)
+    :param runtime_sortingfn: (deprecated; ignored)
         A function implementing custom runtime load-balancing. See
         :ref:`loadbalance` for more information.
 
     """
     def __init__(self, nwalkers, dim, lnpostfn, a=2.0, args=[], kwargs={},
                  postargs=None, threads=None, pool=None, live_dangerously=False,
-                 runtime_sortingfn=None):
+                 runtime_sortingfn=None, moves=None):
         if threads is not None:
             logging.warn("the 'threads' argument is deprecated; "
                          "use 'pool' instead")
 
+        if runtime_sortingfn is not None:
+            logging.warn("the 'runtime_sortingfn' argument is deprecated")
+
+        if moves is None:
+            self._moves = [StretchMove()]
+            self._weights = [1.0]
+        elif isinstance(moves, Iterable):
+            try:
+                self._moves, self._weights = zip(*moves)
+            except TypeError:
+                self._moves = moves
+                self._weights = np.ones(len(moves))
+        else:
+            self._moves = [moves]
+            self._weights = [1.0]
+        self._weights = np.atleast_1d(self._weights).astype(float)
+        self._weights /= np.sum(self._weights)
+
         self.k = nwalkers
         self.a = a
         self.pool = pool
-        self.runtime_sortingfn = runtime_sortingfn
 
         if postargs is not None:
             args = postargs
@@ -195,20 +214,26 @@ def sample(self, p0, lnprob0=None, rstate0=None, blobs0=None,
         lnprob = lnprob0
         blobs = blobs0
         if lnprob is None:
-            lnprob, blobs = self._get_lnprob(p)
+            lnprob, blobs = self.compute_log_prob(p)
+            # lnprob, blobs = self._get_lnprob(p)
 
         # Check to make sure that the probability function didn't return
         # ``np.nan``.
         if np.any(np.isnan(lnprob)):
             raise ValueError("The initial lnprob was NaN.")
 
         # Store the initial size of the stored chain.
-        i0 = self._chain.shape[1]
+        initial_size = self._chain.shape[1]
+
+        # Check that the thin keyword is reasonable.
+        thin = int(thin)
+        if thin <= 0:
+            raise ValueError("Invalid thinning argument")
 
         # Here, we resize chain in advance for performance. This actually
         # makes a pretty big difference.
         if storechain:
-            N = int(iterations / thin)
+            N = iterations // thin
             self._chain = np.concatenate((self._chain,
                                           np.zeros((self.k, N, self.dim))),
                                          axis=1)
@@ -218,64 +243,17 @@ def sample(self, p0, lnprob0=None, rstate0=None, blobs0=None,
         for i in range(int(iterations)):
             self.iterations += 1
 
-            # If we were passed a Metropolis-Hastings proposal
-            # function, use it.
-            if mh_proposal is not None:
-                # Draw proposed positions & evaluate lnprob there
-                q = mh_proposal(p)
-                newlnp, blob = self._get_lnprob(q)
-
-                # Accept if newlnp is better; and ...
-                acc = (newlnp > lnprob)
-
-                # ... sometimes accept for steps that got worse
-                worse = np.flatnonzero(~acc)
-                acc[worse] = ((newlnp[worse] - lnprob[worse]) >
-                              np.log(self._random.rand(len(worse))))
-                del worse
-
-                # Update the accepted walkers.
-                lnprob[acc] = newlnp[acc]
-                p[acc] = q[acc]
-                self.naccepted[acc] += 1
-
-                if blob is not None:
-                    assert blobs is not None, (
-                        "If you start sampling with a given lnprob, you also "
-                        "need to provide the current list of blobs at that "
-                        "position.")
-                    ind = np.arange(self.k)[acc]
-                    for j in ind:
-                        blobs[j] = blob[j]
+            # Choose a random move
+            move = self._random.choice(self._moves, p=self._weights)
 
-            else:
-                # Loop over the two ensembles, calculating the proposed
-                # positions.
-
-                # Slices for the first and second halves
-                first, second = slice(halfk), slice(halfk, self.k)
-                for S0, S1 in [(first, second), (second, first)]:
-                    q, newlnp, acc, blob = self._propose_stretch(p[S0], p[S1],
-                                                                 lnprob[S0])
-                    if np.any(acc):
-                        # Update the positions, log probabilities and
-                        # acceptance counts.
-                        lnprob[S0][acc] = newlnp[acc]
-                        p[S0][acc] = q[acc]
-                        self.naccepted[S0][acc] += 1
-
-                        if blob is not None:
-                            assert blobs is not None, (
-                                "If you start sampling with a given lnprob, "
-                                "you also need to provide the current list of "
-                                "blobs at that position.")
-                            ind = np.arange(len(acc))[acc]
-                            indfull = np.arange(self.k)[S0][acc]
-                            for j in range(len(ind)):
-                                blobs[indfull[j]] = blob[ind[j]]
-
-            if storechain and (i+1) % thin == 0:
-                ind = i0 + int(i / thin)
+            # Propose
+            p, lnprob, blobs, accepted = move.propose(
+                p, lnprob, blobs, self.compute_log_prob, self._random)
+            self.naccepted += accepted
+
+            # Save the results
+            if storechain and (i + 1) % thin == 0:
+                ind = initial_size + (i + 1) // thin - 1
                 self._chain[:, ind, :] = p
                 self._lnprob[:, ind] = lnprob
                 if blobs is not None:
@@ -289,55 +267,7 @@ def sample(self, p0, lnprob0=None, rstate0=None, blobs0=None,
             else:
                 yield p, lnprob, self.random_state
 
-    def _propose_stretch(self, p0, p1, lnprob0):
-        """
-        Propose a new position for one sub-ensemble given the positions of
-        another.
-
-        :param p0:
-            The positions from which to jump.
-
-        :param p1:
-            The positions of the other ensemble.
-
-        :param lnprob0:
-            The log-probabilities at ``p0``.
-
-        This method returns:
-
-        * ``q`` - The new proposed positions for the walkers in ``ensemble``.
-
-        * ``newlnprob`` - The vector of log-probabilities at the positions
-          given by ``q``.
-
-        * ``accept`` - A vector of type ``bool`` indicating whether or not
-          the proposed position for each walker should be accepted.
-
-        * ``blob`` - The new meta data blobs or ``None`` if nothing was
-          returned by ``lnprobfn``.
-
-        """
-        s = np.atleast_2d(p0)
-        Ns = len(s)
-        c = np.atleast_2d(p1)
-        Nc = len(c)
-
-        # Generate the vectors of random numbers that will produce the
-        # proposal.
-        zz = ((self.a - 1.) * self._random.rand(Ns) + 1) ** 2. / self.a
-        rint = self._random.randint(Nc, size=(Ns,))
-
-        # Calculate the proposed positions and the log-probability there.
-        q = c[rint] - zz[:, np.newaxis] * (c[rint] - s)
-        newlnprob, blob = self._get_lnprob(q)
-
-        # Decide whether or not the proposals should be accepted.
-        lnpdiff = (self.dim - 1.) * np.log(zz) + newlnprob - lnprob0
-        accept = (lnpdiff > np.log(self._random.rand(len(lnpdiff))))
-
-        return q, newlnprob, accept, blob
-
-    def _get_lnprob(self, pos=None):
+    def compute_log_prob(self, coords=None):
         """
         Calculate the vector of log-probability for the walkers.
 
@@ -355,10 +285,10 @@ def _get_lnprob(self, pos=None):
           this position or ``None`` if nothing was returned.
 
         """
-        if pos is None:
+        if coords is None:
             p = self.pos
         else:
-            p = pos
+            p = coords
 
         # Check that the parameters are in physical ranges.
         if np.any(np.isinf(p)):
@@ -374,10 +304,6 @@ def _get_lnprob(self, pos=None):
         else:
             M = map
 
-        # sort the tasks according to (user-defined) some runtime guess
-        if self.runtime_sortingfn is not None:
-            p, idx = self.runtime_sortingfn(p)
-
         # Run the log-probability calculations (optionally in parallel).
         results = list(M(self.lnprobfn, [p[i] for i in range(len(p))]))
 
@@ -388,15 +314,6 @@ def _get_lnprob(self, pos=None):
             lnprob = np.array([float(l) for l in results])
             blob = None
 
-        # sort it back according to the original order - get the same
-        # chain irrespective of the runtime sorting fn
-        if self.runtime_sortingfn is not None:
-            orig_idx = np.argsort(idx)
-            lnprob = lnprob[orig_idx]
-            p = [p[i] for i in orig_idx]
-            if blob is not None:
-                blob = [blob[i] for i in orig_idx]
-
         # Check for lnprob returning NaN.
         if np.any(np.isnan(lnprob)):
             # Print some debugging stuff.

emcee/moves/__init__.py

@@ -0,0 +1,9 @@
+# -*- coding: utf-8 -*-
+
+from __future__ import division, print_function
+
+from .move import Move
+from .stretch import StretchMove
+from .red_blue import RedBlueMove
+
+__all__ = ["Move", "RedBlueMove", "StretchMove"]

emcee/moves/gaussian.py

@@ -0,0 +1,118 @@
+# -*- coding: utf-8 -*-
+
+from __future__ import division, print_function
+
+import numpy as np
+
+from .mh import MHMove
+
+__all__ = ["GaussianMove"]
+
+
+class GaussianMove(MHMove):
+    """A Metropolis step with a Gaussian proposal function.
+
+    Args:
+        cov: The covariance of the proposal function. This can be a scalar,
+            vector, or matrix and the proposal will be assumed isotropic,
+            axis-aligned, or general respectively.
+        mode (Optional): Select the method used for updating parameters. This
+            can be one of ``"vector"``, ``"random"``, or ``"sequential"``. The
+            ``"vector"`` mode updates all dimensions simultaneously,
+            ``"random"`` randomly selects a dimension and only updates that
+            one, and ``"sequential"`` loops over dimensions and updates each
+            one in turn.
+        factor (Optional[float]): If provided the proposal will be made with a
+            standard deviation uniformly selected from the range
+            ``exp(U(-log(factor), log(factor))) * cov``. This is invalid for
+            the ``"vector"`` mode.
+
+    Raises:
+        ValueError: If the proposal dimensions are invalid or if any of any of
+            the other arguments are inconsistent.
+
+    """
+    def __init__(self, cov, mode="vector", factor=None):
+        # Parse the proposal type.
+        try:
+            float(cov)
+
+        except TypeError:
+            cov = np.atleast_1d(cov)
+            if len(cov.shape) == 1:
+                # A diagonal proposal was given.
+                ndim = len(cov)
+                proposal = _diagonal_proposal(np.sqrt(cov), factor, mode)
+
+            elif len(cov.shape) == 2 and cov.shape[0] == cov.shape[1]:
+                # The full, square covariance matrix was given.
+                ndim = cov.shape[0]
+                proposal = _proposal(cov, factor, mode)
+
+            else:
+                raise ValueError("Invalid proposal scale dimensions")
+
+        else:
+            # This was a scalar proposal.
+            ndim = None
+            proposal = _isotropic_proposal(np.sqrt(cov), factor, mode)
+
+        super(GaussianMove, self).__init__(proposal, ndim=ndim)
+
+
+class _isotropic_proposal(object):
+
+    allowed_modes = ["vector", "random", "sequential"]
+
+    def __init__(self, scale, factor, mode):
+        self.index = 0
+        self.scale = scale
+        if factor is None:
+            self._log_factor = None
+        else:
+            if factor < 1.0:
+                raise ValueError("'factor' must be >= 1.0")
+            self._log_factor = np.log(factor)
+
+        if mode not in self.allowed_modes:
+            raise ValueError(("'{0}' is not a recognized mode. "
+                              "Please select from: {1}")
+                             .format(mode, self.allowed_modes))
+        self.mode = mode
+
+    def get_factor(self, rng):
+        if self._log_factor is None:
+            return 1.0
+        return np.exp(rng.uniform(-self._log_factor, self._log_factor))
+
+    def get_updated_vector(self, rng, x0):
+        return x0 + self.get_factor(rng) * self.scale * rng.randn(*(x0.shape))
+
+    def __call__(self, x0, rng):
+        nw, nd = x0.shape
+        xnew = self.get_updated_vector(rng, x0)
+        if self.mode == "random":
+            m = (range(nw), rng.randint(x0.shape[-1], size=nw))
+        elif self.mode == "sequential":
+            m = (range(nw), self.index % nd + np.zeros(nw, dtype=int))
+            self.index = (self.index + 1) % nd
+        else:
+            return xnew, np.zeros(nw)
+        x = np.array(x0)
+        x[m] = xnew[m]
+        return x, np.zeros(nw)
+
+
+class _diagonal_proposal(_isotropic_proposal):
+
+    def get_updated_vector(self, rng, x0):
+        return x0 + self.get_factor(rng) * self.scale * rng.randn(*(x0.shape))
+
+
+class _proposal(_isotropic_proposal):
+
+    allowed_modes = ["vector"]
+
+    def get_updated_vector(self, rng, x0):
+        return x0 + self.get_factor(rng) * rng.multivariate_normal(
+            np.zeros(len(self.scale)), self.scale)