Add multiprocessing option for pocomc

Author: andreicuceu

vega/samplers/pocomc.py

@@ -4,6 +4,7 @@
 import pocomc
 from mpi4py import MPI
 from mpi4py.futures import MPIPoolExecutor
+from multiprocessing import Pool
 from scipy.stats import uniform
 
 from vega.samplers.sampler_interface import Sampler
@@ -25,12 +26,21 @@ def get_sampler_settings(self, sampler_config, num_params, num_derived):
         self.n_evidence = sampler_config.getint('n_evidence', 0)
         self.save_every = sampler_config.getint('save_every', 3)
 
+        self.use_mpi = sampler_config.getboolean('use_mpi', False)
+        self.num_cpu = sampler_config.getint('num_cpu', 64)
+
         self.prior = pocomc.Prior(
             [uniform(self.limits[par][0], self.limits[par][1]-self.limits[par][0])
              for par in self.limits]
         )
 
     def run(self):
+        if self.use_mpi:
+            self._run_mpi()
+        else:
+            self._run_multiprocessing()
+
+    def _run_mpi(self):
         """ Run the PocoMC sampler """
         mpi_comm = MPI.COMM_WORLD
         num_mpi_threads = mpi_comm.Get_size()
@@ -42,6 +52,16 @@ def run(self):
             )
             self.pocomc_sampler.run(self.n_total, self.n_evidence, save_every=self.save_every)
 
+    def _run_multiprocessing(self):
+        """ Run the PocoMC sampler """
+        with Pool(self.num_cpu) as pool:
+            self.pocomc_sampler = pocomc.Sampler(
+                self.prior, self.log_lik, pool=pool, output_dir=self.path,
+                dynamic=self.dynamic, precondition=self.precondition,
+                n_effective=self.n_effective, n_active=self.n_active,
+            )
+            self.pocomc_sampler.run(self.n_total, self.n_evidence, save_every=self.save_every)
+
     def write_chain(self):
         # Get the weighted posterior samples
         samples, weights, logl, logp = self.pocomc_sampler.posterior()