diff --git a/PTMCMCSampler/PTMCMCSampler.py b/PTMCMCSampler/PTMCMCSampler.py index 4aa1fee..5550ea6 100755 --- a/PTMCMCSampler/PTMCMCSampler.py +++ b/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) diff --git a/PTMCMCSampler/nompi4py.py b/PTMCMCSampler/nompi4py.py index 944024c..c73160d 100644 --- a/PTMCMCSampler/nompi4py.py +++ b/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() diff --git a/tests/test_nuts.py b/tests/test_nuts.py index 1586948..e9547a6 100644 --- a/tests/test_nuts.py +++ b/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 diff --git a/tests/test_simple.py b/tests/test_simple.py index e5719a4..e4d689c 100644 --- a/tests/test_simple.py +++ b/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()