Add multiprocessing basic

Demonstrated to work up to 4 cores with a reduction in time of 2

bilby/core/sampler/ptemcee.py

@@ -39,7 +39,7 @@ class Ptemcee(MCMCSampler):
     # Arguments used by ptemcee
     default_kwargs = dict(
         ntemps=20, nwalkers=200, Tmax=None, betas=None,
-        threads=1, pool=None, a=2.0, loglargs=[], logpargs=[], loglkwargs={},
+        a=2.0, loglargs=[], logpargs=[], loglkwargs={},
         logpkwargs={}, adaptation_lag=10000, adaptation_time=100, random=None,
         iterations=1000, thin=1, storechain=True, adapt=False,
         swap_ratios=False)
@@ -48,7 +48,8 @@ class Ptemcee(MCMCSampler):
                  use_ratio=False, plot=False, skip_import_verification=False,
                  resume=True, nsamples=5000, burn_in_nact=50, thin_by_nact=1,
                  autocorr_c=5, safety=1, ncheck=50, nfrac=5, frac_threshold=0.01,
-                 autocorr_tol=50, min_tau=1, check_point_deltaT=600, **kwargs):
+                 autocorr_tol=50, min_tau=1, check_point_deltaT=600, 
+                 threads=1, **kwargs):
         super(Ptemcee, self).__init__(
             likelihood=likelihood, priors=priors, outdir=outdir,
             label=label, use_ratio=use_ratio, plot=plot,
@@ -71,6 +72,8 @@ class Ptemcee(MCMCSampler):
         self.min_tau = min_tau
         self.check_point_deltaT = check_point_deltaT
 
+        self.threads = threads
+
         self.resume_file = "{}/{}_checkpoint_resume.pickle".format(self.outdir, self.label)
 
     @property
@@ -105,16 +108,28 @@ class Ptemcee(MCMCSampler):
 
         else:
             self.sampler = ptemcee.Sampler(
-                dim=self.ndim, logl=self.log_likelihood, logp=self.log_prior,
-                **self.sampler_init_kwargs)
+                dim=self.ndim, logl=do_nothing_function, logp=do_nothing_function,
+                pool=self.pool, threads=self.threads, **self.sampler_init_kwargs)
+            self.sampler._likeprior = LikePriorEvaluator(
+                self.likelihood, self.priors, self.search_parameter_keys,
+                use_ratio=self.use_ratio)
             pos0 = self.get_pos0_from_prior()
 
         return self.sampler, pos0
 
     def run_sampler(self):
+        import schwimmbad
+        if self.threads > 1:
+            with schwimmbad.MultiPool(self.threads) as pool:
+                self.pool = pool
+                return self.run_sampler_internal()
+        else:
+            self.pool = None
+            return self.run_sampler_internal()
+
+    def run_sampler_internal(self):
         import emcee
         sampler, pos0 = self.get_sampler()
-
         self.time_per_check = []
         self.tau_list = []
         self.tau_list_n = []
@@ -424,3 +439,49 @@ def compute_evidence(sampler, outdir, label, nburn, thin, make_plots=True):
     return lnZ, lnZerr
 
 
+def do_nothing_function():
+    """ This is a do-nothing function, we overwrite the likelihood and prior elsewhere """
+    pass
+
+
+class LikePriorEvaluator(object):
+    """
+    A overwrite of the ptemcee.LikePriorEvaluator to use bilby likelihood and priors
+
+    """
+
+    def __init__(self, likelihood, priors, search_parameter_keys, use_ratio=False):
+        self.likelihood = likelihood
+        self.priors = priors
+        self.search_parameter_keys = search_parameter_keys
+        self.use_ratio = use_ratio
+
+    def logl(self, v_array):
+        parameters = {key: v for key, v in zip(self.search_parameter_keys, v_array)}
+        if self.priors.evaluate_constraints(parameters) > 0:
+            self.likelihood.parameters.update(parameters)
+            if self.use_ratio:
+                return self.likelihood.log_likelihood() - self.likelihood.noise_log_likelihood()
+            else:
+                return self.likelihood.log_likelihood()
+        else:
+            return np.nan_to_num(-np.inf)
+
+    def logp(self, v_array):
+        params = {key: t for key, t in zip(self.search_parameter_keys, v_array)}
+        return self.priors.ln_prob(params)
+
+    def __call__(self, x):
+        lp = self.logp(x)
+        if 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)
+            if np.isnan(ll).any():
+                raise ValueError('Log likelihood function returned NaN.')
+
+        return ll, lp