Add in additional stubs to mock MPI object and update tests (#35)

PTMCMCSampler/PTMCMCSampler.py

@@ -101,7 +101,7 @@ def __init__(
             self.stream = [np.random.default_rng(s) for s in child_seeds]
         else:
             self.stream = None
-        self.stream = comm.scatter(self.stream, root=0)
+        self.stream = self.comm.scatter(self.stream, root=0)
 
         self.ndim = ndim
         self.logl = _function_wrapper(logl, loglargs, loglkwargs)

PTMCMCSampler/nompi4py.py

@@ -21,6 +21,17 @@ def recv(self, source=1, tag=55):
     def Iprobe(self, source=1, tag=55):
         pass
 
+    def scatter(self, sendobj, **kwargs):
+        if sendobj is not None:
+            return sendobj[0]
+        return None
+
+    def bcast(self, obj, **kwargs):
+        return obj
+
+    def gather(self, sendobj, **kwargs):
+        return [sendobj]
+
 
 # Global object representing no MPI:
 COMM_WORLD = MPIDummy()

tests/test_nuts.py

@@ -4,8 +4,10 @@
 import numpy as np
 import scipy.linalg as sl
 import scipy.optimize as so
+from mpi4py import MPI
 
 from PTMCMCSampler import PTMCMCSampler
+from PTMCMCSampler import nompi4py as MPIDUMMY
 
 
 class GaussianLikelihood(object):
@@ -165,6 +167,9 @@ class TestNuts(TestCase):
     def tearDownClass(cls):
         shutil.rmtree("chains")
 
+    def setUp(self) -> None:
+        self.comm = MPI.COMM_WORLD
+
     def test_nuts(self):
         ndim = 40
         glo = GaussianLikelihood(ndim=ndim, pmin=0.0, pmax=10.0)
@@ -196,6 +201,7 @@ def test_nuts(self):
             logl_grad=gl.lnlikefn_grad,
             logp_grad=gl.lnpriorfn_grad,
             outDir="./chains",
+            comm=self.comm,
         )
 
         sampler.sample(
@@ -213,3 +219,8 @@ def test_nuts(self):
             HMCsteps=100,
             HMCstepsize=0.4,
         )
+
+
+class TestNutsNoMPI(TestNuts):
+    def setUp(self) -> None:
+        self.comm = MPIDUMMY.COMM_WORLD

tests/test_simple.py

@@ -5,8 +5,10 @@
 from unittest import TestCase
 
 import numpy as np
+from mpi4py import MPI
 
 from PTMCMCSampler import PTMCMCSampler
+from PTMCMCSampler import nompi4py as MPIDUMMY
 
 
 class GaussianLikelihood(object):
@@ -65,6 +67,9 @@ class TestSimpleSampler(TestCase):
     def tearDownClass(cls):
         shutil.rmtree("chains")
 
+    def setUp(self) -> None:
+        self.comm = MPI.COMM_WORLD
+
     def test_simple(self):
         # ## Setup Gaussian model class
         ndim = 20
@@ -76,10 +81,25 @@ def test_simple(self):
         p0 = np.random.uniform(pmin, pmax, ndim)
         cov = np.eye(ndim) * 0.1**2
 
-        sampler = PTMCMCSampler.PTSampler(ndim, glo.lnlikefn, glo.lnpriorfn, np.copy(cov), outDir="./chains")
+        sampler = PTMCMCSampler.PTSampler(
+            ndim,
+            glo.lnlikefn,
+            glo.lnpriorfn,
+            np.copy(cov),
+            outDir="./chains",
+            comm=self.comm,
+        )
 
         # add to jump proposal cycle
         ujump = UniformJump(pmin, pmax)
         sampler.addProposalToCycle(ujump.jump, 5)
 
         sampler.sample(p0, 10000, burn=500, thin=1, covUpdate=500, SCAMweight=20, AMweight=20, DEweight=20)
+
+
+class TestSimpleSamplerNoMPI(TestSimpleSampler):
+    def setUp(self) -> None:
+        self.comm = MPIDUMMY.COMM_WORLD
+
+    def test_simple(self):
+        return super().test_simple()