Add seed argument to EnsembleSampler

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):

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

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

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):

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)

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)

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

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)

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)

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)
 

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())

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))
+    )