From 1d37b5516fedcd65fb6dc4928f00c3509ce422d4 Mon Sep 17 00:00:00 2001
From: Leo Singer <leo.singer@ligo.org>
Date: Mon, 11 Jun 2018 18:47:43 -0400
Subject: [PATCH] Add vectorize option for Numpy arraywise evaluation of
 like/prior

This behaves exactly the same as the `vectorize` option from
`emcee.EnsembleSampler`.
---
 ptemcee/sampler.py | 35 +++++++++++++++++++++++++----------
 1 file changed, 25 insertions(+), 10 deletions(-)
 mode change 100644 => 100755 ptemcee/sampler.py

diff --git a/ptemcee/sampler.py b/ptemcee/sampler.py
old mode 100644
new mode 100755
index bc81147..320d116
--- a/ptemcee/sampler.py
+++ b/ptemcee/sampler.py
@@ -128,19 +128,22 @@ def __init__(self, logl, logp,
         self.logpkwargs = logpkwargs
 
     def __call__(self, x):
+        s = x.shape
+        x = x.reshape((-1, x.shape[-1]))
+
         lp = self.logp(x, *self.logpargs, **self.logpkwargs)
-        if np.isnan(lp):
+        if np.any(np.isnan(lp)):
             raise ValueError('Prior function returned NaN.')
 
-        if lp == float('-inf'):
-            # Can't return -inf, since this messes with beta=0 behaviour.
-            ll = 0
-        else:
-            ll = self.logl(x, *self.loglargs, **self.loglkwargs)
-            if np.isnan(ll).any():
-                raise ValueError('Log likelihood function returned NaN.')
+        # Can't return -inf, since this messes with beta=0 behaviour.
+        ll = np.empty_like(lp)
+        bad = (lp == -np.inf)
+        ll[bad] = 0
+        ll[~bad] = self.logl(x[~bad], *self.loglargs, **self.loglkwargs)
+        if np.any(np.isnan(ll)):
+            raise ValueError('Log likelihood function returned NaN.')
 
-        return ll, lp
+        return ll.reshape(s[:-1]), lp.reshape(s[:-1])
 
 class Sampler(object):
     """
@@ -198,6 +201,11 @@ class Sampler(object):
     :param adaptation_time: (optional)
         Time-scale for temperature dynamics.  Default: 100.
 
+    :param vectorize: (optional)
+        If ``True``, ``logl`` and ``logp`` are expected to accept a list of
+        position vectors instead of just one. Note that ``pool`` will be
+        ignored if this is ``True``. (default: ``False``)
+
     """
     def __init__(self, nwalkers, dim, logl, logp,
                  ntemps=None, Tmax=None, betas=None,
@@ -205,7 +213,11 @@ def __init__(self, nwalkers, dim, logl, logp,
                  loglargs=[], logpargs=[],
                  loglkwargs={}, logpkwargs={},
                  adaptation_lag=10000, adaptation_time=100,
-                 random=None):
+                 random=None, vectorize=False):
+        self._vectorize = vectorize
+        if vectorize:
+            pool = None
+
         if random is None:
             self._random = np.random.mtrand.RandomState()
         else:
@@ -447,6 +459,9 @@ def _stretch(self, p, logpost, logl):
             self.nprop_accepted[:, jupdate::2] += accepts
 
     def _evaluate(self, ps):
+        if self._vectorize:
+            return self._likeprior(ps)
+
         mapf = map if self.pool is None else self.pool.map
         results = list(mapf(self._likeprior, ps.reshape((-1, self.dim))))