diff --git a/src/emcee/ensemble.py b/src/emcee/ensemble.py index 45cc4ba9..f38f4042 100644 --- a/src/emcee/ensemble.py +++ b/src/emcee/ensemble.py @@ -10,7 +10,7 @@ from .moves import StretchMove from .pbar import get_progress_bar from .state import State -from .utils import deprecated, deprecation_warning +from .utils import deprecated, deprecation_warning, _check_random_state __all__ = ["EnsembleSampler", "walkers_independent"] @@ -61,6 +61,15 @@ class EnsembleSampler(object): to accept a list of position vectors instead of just one. Note that ``pool`` will be ignored if this is ``True``. (default: ``False``) + seed (Union[int, np.random.RandomState, np.random.Generator, None]): If + `seed` is not specified the `np.RandomState` singleton is used. + If `seed` is an int, a new `np.random.RandomState` instance is used, + seeded with seed. + If `seed` is already a `np.random.RandomState` or a + `np.random.Generator` instance, then that `RandomState` or + `Generator` instance is used, omitting the stored state if + re-using a backend. + Specify `seed` for reproducable minimizations. """ @@ -76,6 +85,7 @@ def __init__( backend=None, vectorize=False, blobs_dtype=None, + seed=None, # Deprecated... a=None, postargs=None, @@ -124,6 +134,7 @@ def __init__( self.backend = Backend() if backend is None else backend # Deal with re-used backends + state = None if not self.backend.initialized: self._previous_state = None self.reset() @@ -150,8 +161,7 @@ def __init__( # This is a random number generator that we can easily set the state # of without affecting the numpy-wide generator - self._random = np.random.mtrand.RandomState() - self._random.set_state(state) + self._random = _check_random_state(seed, state) # Do a little bit of _magic_ to make the likelihood call with # ``args`` and ``kwargs`` pickleable. @@ -167,7 +177,10 @@ def random_state(self): so silently. """ - return self._random.get_state() + try: + return self._random.get_state() + except AttributeError: + return self._random.bit_generator.state @random_state.setter # NOQA def random_state(self, state): diff --git a/src/emcee/moves/de.py b/src/emcee/moves/de.py index 32e60cfa..e3314af0 100644 --- a/src/emcee/moves/de.py +++ b/src/emcee/moves/de.py @@ -3,6 +3,7 @@ import numpy as np from .red_blue import RedBlueMove +from ..utils import rng_integers __all__ = ["DEMove"] @@ -42,9 +43,9 @@ def get_proposal(self, s, c, random): Nc = list(map(len, c)) ndim = s.shape[1] q = np.empty((Ns, ndim), dtype=np.float64) - f = self.sigma * random.randn(Ns) + f = self.sigma * random.standard_normal(Ns) for i in range(Ns): - w = np.array([c[j][random.randint(Nc[j])] for j in range(2)]) + w = np.array([c[j][rng_integers(random, Nc[j])] for j in range(2)]) random.shuffle(w) g = np.diff(w, axis=0) * self.g0 + f[i] q[i] = s[i] + g diff --git a/src/emcee/moves/de_snooker.py b/src/emcee/moves/de_snooker.py index 0b53efa4..6ca38b44 100644 --- a/src/emcee/moves/de_snooker.py +++ b/src/emcee/moves/de_snooker.py @@ -3,6 +3,7 @@ import numpy as np from .red_blue import RedBlueMove +from ..utils import rng_integers __all__ = ["DESnookerMove"] @@ -35,7 +36,7 @@ def get_proposal(self, s, c, random): q = np.empty((Ns, ndim), dtype=np.float64) metropolis = np.empty(Ns, dtype=np.float64) for i in range(Ns): - w = np.array([c[j][random.randint(Nc[j])] for j in range(3)]) + w = np.array([c[j][rng_integers(random, Nc[j])] for j in range(3)]) random.shuffle(w) z, z1, z2 = w delta = s[i] - z diff --git a/src/emcee/moves/gaussian.py b/src/emcee/moves/gaussian.py index eeafe7ef..e95225c5 100644 --- a/src/emcee/moves/gaussian.py +++ b/src/emcee/moves/gaussian.py @@ -3,6 +3,7 @@ import numpy as np from .mh import MHMove +from ..utils import rng_integers __all__ = ["GaussianMove"] @@ -88,13 +89,13 @@ def get_factor(self, rng): return np.exp(rng.uniform(-self._log_factor, self._log_factor)) def get_updated_vector(self, rng, x0): - return x0 + self.get_factor(rng) * self.scale * rng.randn(*(x0.shape)) + return x0 + self.get_factor(rng) * self.scale * rng.standard_normal((x0.shape)) def __call__(self, x0, rng): nw, nd = x0.shape xnew = self.get_updated_vector(rng, x0) if self.mode == "random": - m = (range(nw), rng.randint(x0.shape[-1], size=nw)) + m = (range(nw), rng_integers(rng, x0.shape[-1], size=nw)) elif self.mode == "sequential": m = (range(nw), self.index % nd + np.zeros(nw, dtype=int)) self.index = (self.index + 1) % nd @@ -107,7 +108,7 @@ def __call__(self, x0, rng): class _diagonal_proposal(_isotropic_proposal): def get_updated_vector(self, rng, x0): - return x0 + self.get_factor(rng) * self.scale * rng.randn(*(x0.shape)) + return x0 + self.get_factor(rng) * self.scale * rng.standard_normal((x0.shape)) class _proposal(_isotropic_proposal): diff --git a/src/emcee/moves/mh.py b/src/emcee/moves/mh.py index 4b190498..253362cc 100644 --- a/src/emcee/moves/mh.py +++ b/src/emcee/moves/mh.py @@ -56,7 +56,7 @@ def propose(self, model, state): # Loop over the walkers and update them accordingly. lnpdiff = new_log_probs - state.log_prob + factors - accepted = np.log(model.random.rand(nwalkers)) < lnpdiff + accepted = np.log(model.random.uniform(size=nwalkers)) < lnpdiff # Update the parameters new_state = State(q, log_prob=new_log_probs, blobs=new_blobs) diff --git a/src/emcee/moves/red_blue.py b/src/emcee/moves/red_blue.py index e6ab59d5..3d32e55f 100644 --- a/src/emcee/moves/red_blue.py +++ b/src/emcee/moves/red_blue.py @@ -97,7 +97,7 @@ def propose(self, model, state): zip(all_inds[S1], factors, new_log_probs) ): lnpdiff = f + nlp - state.log_prob[j] - if lnpdiff > np.log(model.random.rand()): + if lnpdiff > np.log(model.random.uniform()): accepted[j] = True new_state = State(q, log_prob=new_log_probs, blobs=new_blobs) diff --git a/src/emcee/moves/stretch.py b/src/emcee/moves/stretch.py index 40d8e537..b20fca7d 100644 --- a/src/emcee/moves/stretch.py +++ b/src/emcee/moves/stretch.py @@ -3,6 +3,7 @@ import numpy as np from .red_blue import RedBlueMove +from ..utils import rng_integers __all__ = ["StretchMove"] @@ -27,7 +28,7 @@ def get_proposal(self, s, c, random): c = np.concatenate(c, axis=0) Ns, Nc = len(s), len(c) ndim = s.shape[1] - zz = ((self.a - 1.0) * random.rand(Ns) + 1) ** 2.0 / self.a + zz = random.uniform(low=1, high=self.a, size=Ns) ** 2.0 / self.a factors = (ndim - 1.0) * np.log(zz) - rint = random.randint(Nc, size=(Ns,)) + rint = rng_integers(random, Nc, size=(Ns,)) return c[rint] - (c[rint] - s) * zz[:, None], factors diff --git a/src/emcee/tests/integration/test_de.py b/src/emcee/tests/integration/test_de.py index d9c8ec68..c7f7fa75 100644 --- a/src/emcee/tests/integration/test_de.py +++ b/src/emcee/tests/integration/test_de.py @@ -1,19 +1,32 @@ # -*- coding: utf-8 -*- +import pytest +import packaging +import numpy as np from emcee import moves from .test_proposal import _test_normal, _test_uniform -__all__ = ["test_normal_de", "test_normal_de_no_gamma", "test_uniform_de"] - -def test_normal_de(**kwargs): - _test_normal(moves.DEMove(), **kwargs) - - -def test_normal_de_no_gamma(**kwargs): - _test_normal(moves.DEMove(gamma0=1.0), **kwargs) - - -def test_uniform_de(**kwargs): - _test_uniform(moves.DEMove(), **kwargs) +@pytest.mark.parametrize( + "generator", + [ + False, + pytest.param( + True, + marks=pytest.mark.skipif( + packaging.version.parse(np.__version__) < packaging.version.parse("1.17.0"), + reason="requires numpy 1.17.0 or higher", + ) + ) + ] +) +class TestDE: + def test_normal_de(self, generator): + _test_normal(moves.DEMove(), generator=generator) + + def test_normal_de_no_gamma(self, generator): + _test_normal(moves.DEMove(gamma0=1.0), generator=generator) + + def test_uniform_de(self, generator): + _test_uniform(moves.DEMove(), generator=generator) diff --git a/src/emcee/tests/integration/test_longdouble.py b/src/emcee/tests/integration/test_longdouble.py index b8d2c7a8..53cdb89d 100644 --- a/src/emcee/tests/integration/test_longdouble.py +++ b/src/emcee/tests/integration/test_longdouble.py @@ -12,7 +12,7 @@ def log_prob(x, ivar): ndim, nwalkers = 5, 20 ivar = 1. / np.random.rand(ndim).astype(np.longdouble) - p0 = np.random.randn(nwalkers, ndim).astype(np.longdouble) + p0 = np.random.standard_normal((nwalkers, ndim)).astype(np.longdouble) sampler = emcee.EnsembleSampler(nwalkers, ndim, log_prob, args=[ivar]) sampler.run_mcmc(p0, 100) diff --git a/src/emcee/tests/integration/test_proposal.py b/src/emcee/tests/integration/test_proposal.py index a590be0a..37d71f89 100644 --- a/src/emcee/tests/integration/test_proposal.py +++ b/src/emcee/tests/integration/test_proposal.py @@ -32,6 +32,7 @@ def _test_normal( nwalkers=32, nsteps=2000, seed=1234, + generator=False, check_acceptance=True, pool=None, blobs=False, @@ -47,8 +48,12 @@ def _test_normal( else: lp = normal_log_prob + sampler_kwargs = {} + if generator: + sampler_kwargs["seed"] = np.random.default_rng(seed) + sampler = emcee.EnsembleSampler( - nwalkers, ndim, lp, moves=proposal, pool=pool + nwalkers, ndim, lp, moves=proposal, pool=pool, **sampler_kwargs ) if hasattr(proposal, "ntune") and proposal.ntune > 0: coords = sampler.run_mcmc(coords, proposal.ntune, tune=True) @@ -74,15 +79,19 @@ def _test_normal( assert ks < 0.05, "The K-S test failed" -def _test_uniform(proposal, nwalkers=32, nsteps=2000, seed=1234): +def _test_uniform(proposal, nwalkers=32, nsteps=2000, seed=1234, generator=False): # Set up the random number generator. np.random.seed(seed) # Initialize the ensemble and proposal. coords = np.random.rand(nwalkers, 1) + sampler_kwargs = {} + if generator: + sampler_kwargs["seed"] = np.random.default_rng(seed) + sampler = emcee.EnsembleSampler( - nwalkers, 1, normal_log_prob, moves=proposal + nwalkers, 1, normal_log_prob, moves=proposal, **sampler_kwargs ) sampler.run_mcmc(coords, nsteps) diff --git a/src/emcee/tests/unit/test_sampler.py b/src/emcee/tests/unit/test_sampler.py index 8c629074..da8fff75 100644 --- a/src/emcee/tests/unit/test_sampler.py +++ b/src/emcee/tests/unit/test_sampler.py @@ -2,9 +2,11 @@ import pickle from itertools import product, islice +from copy import deepcopy import numpy as np import pytest +import packaging from emcee import EnsembleSampler, backends, moves, walkers_independent @@ -42,6 +44,7 @@ def test_shapes(backend, moves, nwalkers=32, ndim=3, nsteps=10, seed=1234): # Run the sampler. sampler.run_mcmc(coords, nsteps) + chain = sampler.get_chain() assert len(chain) == nsteps, "wrong number of steps" @@ -137,7 +140,10 @@ def run_sampler( ): np.random.seed(seed) coords = np.random.randn(nwalkers, ndim) - sampler = EnsembleSampler(nwalkers, ndim, normal_log_prob, backend=backend) + np.random.seed(None) + sampler = EnsembleSampler( + nwalkers, ndim, normal_log_prob, backend=backend, seed=seed + ) sampler.run_mcmc( coords, nsteps, @@ -333,3 +339,52 @@ def test_infinite_iterations(backend, nwalkers=32, ndim=3): coords = np.random.randn(nwalkers, ndim) for state in islice(EnsembleSampler(nwalkers, ndim, normal_log_prob, backend=be).sample(coords, iterations=None, store=False), 10): pass + +def test_sampler_seed(): + nwalkers = 32 + ndim = 3 + nsteps = 25 + np.random.seed(456) + coords = np.random.randn(nwalkers, ndim) + sampler1 = EnsembleSampler(nwalkers, ndim, normal_log_prob, seed=1234) + sampler2 = EnsembleSampler(nwalkers, ndim, normal_log_prob, seed=2) + sampler3 = EnsembleSampler(nwalkers, ndim, normal_log_prob, seed=1234) + sampler4 = EnsembleSampler( + nwalkers, ndim, normal_log_prob, seed=deepcopy(sampler1._random) + ) + for sampler in (sampler1, sampler2, sampler3, sampler4): + sampler.run_mcmc(coords, nsteps) + for k in ["get_chain", "get_log_prob"]: + attr1 = getattr(sampler1, k)() + attr2 = getattr(sampler2, k)() + attr3 = getattr(sampler3, k)() + attr4 = getattr(sampler4, k)() + assert not np.allclose(attr1, attr2), "inconsistent {0}".format(k) + np.testing.assert_allclose(attr1, attr3, err_msg="inconsistent {0}".format(k)) + np.testing.assert_allclose(attr1, attr4, err_msg="inconsistent {0}".format(k)) + + +def test_sampler_bad_seed(): + nwalkers = 32 + ndim = 3 + with pytest.raises(TypeError, match="seed must be"): + EnsembleSampler(nwalkers, ndim, normal_log_prob, seed="bad_seed") + +@pytest.mark.skipif( + packaging.version.parse(np.__version__) < packaging.version.parse("1.17.0"), + reason="requires numpy 1.17.0 or higher", +) +def test_sampler_generator(): + nwalkers = 32 + ndim = 3 + nsteps = 5 + np.random.seed(456) + coords = np.random.randn(nwalkers, ndim) + seed1 = np.random.default_rng(1) + sampler1 = EnsembleSampler(nwalkers, ndim, normal_log_prob, seed=seed1) + sampler1.run_mcmc(coords, nsteps) + seed2 = np.random.default_rng(1) + sampler2 = EnsembleSampler(nwalkers, ndim, normal_log_prob, seed=seed2) + sampler2.run_mcmc(coords, nsteps) + np.testing.assert_allclose(sampler1.get_chain(), sampler2.get_chain()) + np.testing.assert_allclose(sampler1.get_log_prob(), sampler2.get_log_prob()) diff --git a/src/emcee/utils.py b/src/emcee/utils.py index 42bbcfd4..4f77d903 100644 --- a/src/emcee/utils.py +++ b/src/emcee/utils.py @@ -60,3 +60,47 @@ def sample_ellipsoid(p0, covmat, size=1): return np.random.multivariate_normal( np.atleast_1d(p0), np.atleast_2d(covmat), size=size ) + +try: + Generator = np.random.Generator +except AttributeError: + Generator = None + +def rng_integers(gen, low, **kwargs): + """ + Generate integers from a Generator or RandomState. + + :param gen: Generator or RandomState object + :param low: Passed as first argument to gen.integers() or to + gen.randint() depending of its type. + :param kwargs: Passed to gen.integers() or to gen.randint() + depending on its type. + + """ + if Generator is not None and isinstance(gen, Generator): + return gen.integers(low, **kwargs) + return gen.randint(low, **kwargs) + +def _check_random_state(seed, state): + """Check seed argument and set RandomState. + + Based on scikit-learn utils/validation.py. + """ + if isinstance(seed, (int, np.integer)) or seed is None: + random = np.random.mtrand.RandomState() + random.set_state(state) + if seed is not None: + random.seed(seed) + return random + if isinstance(seed, np.random.RandomState): + return seed + try: + # Generator is only available in numpy >= 1.17 + if isinstance(seed, np.random.Generator): + return seed + except AttributeError: + pass + raise TypeError( + "seed must be an int, np.random.RandomState, np.random.Generator or " + "None of type {}".format(type(seed)) + )