Use random number generator seeding in SMC

Author: ricardoV94

pymc/smc/sample_smc.py

@@ -42,7 +42,7 @@ def sample_smc(
     *,
     start=None,
     model=None,
-    random_seed=-1,
+    random_seed=None,
     chains=None,
     cores=None,
     compute_convergence_checks=True,
@@ -191,15 +191,19 @@ def sample_smc(
         cores = min(chains, cores)
 
     if random_seed == -1:
+        raise FutureWarning(
+            f"random_seed should be a non-negative integer or None, got: {random_seed}"
+            "This will raise a ValueError in the Future"
+        )
         random_seed = None
-    if chains == 1 and isinstance(random_seed, int):
-        random_seed = [random_seed]
-    if random_seed is None or isinstance(random_seed, int):
-        if random_seed is not None:
-            np.random.seed(random_seed)
-        random_seed = [np.random.randint(2 ** 30) for _ in range(chains)]
-    if not isinstance(random_seed, Iterable):
-        raise TypeError("Invalid value for `random_seed`. Must be tuple, list or int")
+    if isinstance(random_seed, int) or random_seed is None:
+        rng = np.random.default_rng(seed=random_seed)
+        random_seed = list(rng.integers(2 ** 30, size=chains))
+    elif isinstance(random_seed, Iterable):
+        if len(random_seed) != chains:
+            raise ValueError(f"Length of seeds ({len(seeds)}) must match number of chains {chains}")
+    else:
+        raise TypeError("Invalid value for `random_seed`. Must be tuple, list, int or None")
 
     model = modelcontext(model)
 

pymc/smc/smc.py

@@ -126,7 +126,7 @@ def __init__(
         draws=2000,
         start=None,
         model=None,
-        random_seed=-1,
+        random_seed=None,
         threshold=0.5,
     ):
         """
@@ -140,6 +140,8 @@ def __init__(
             Starting point in parameter space. It should be a list of dict with length `chains`.
             When None (default) the starting point is sampled from the prior distribution.
         model: Model (optional if in ``with`` context)).
+        random_seed: int
+            Value used to initialize the random number generator.
         threshold: float
             Determines the change of beta from stage to stage, i.e.indirectly the number of stages,
             the higher the value of `threshold` the higher the number of stages. Defaults to 0.5.
@@ -151,10 +153,7 @@ def __init__(
         self.start = start
         self.threshold = threshold
         self.model = model
-        self.random_seed = random_seed
-
-        if self.random_seed != -1:
-            np.random.seed(self.random_seed)
+        self.rng = np.random.default_rng(seed=random_seed)
 
         self.model = modelcontext(model)
         self.variables = inputvars(self.model.value_vars)
@@ -262,7 +261,7 @@ def update_beta_and_weights(self):
 
     def resample(self):
         """Resample particles based on importance weights"""
-        self.resampling_indexes = np.random.choice(
+        self.resampling_indexes = self.rng.choice(
             np.arange(self.draws), size=self.draws, p=self.weights
         )
 
@@ -382,11 +381,11 @@ def mutate(self):
         ac_ = np.empty((self.n_steps, self.draws))
 
         cov = self.proposal_dist.cov
-        log_R = np.log(np.random.rand(self.n_steps, self.draws))
+        log_R = np.log(self.rng.random((self.n_steps, self.draws)))
         for n_step in range(self.n_steps):
             # The proposal is independent from the current point.
             # We have to take that into account to compute the Metropolis-Hastings acceptance
-            proposal = floatX(self.proposal_dist.rvs(size=self.draws))
+            proposal = floatX(self.proposal_dist.rvs(size=self.draws, random_state=self.rng))
             proposal = proposal.reshape(len(proposal), -1)
             # To do that we compute the logp of moving to a new point
             forward = self.proposal_dist.logpdf(proposal)
@@ -497,11 +496,12 @@ def mutate(self):
         """Metropolis-Hastings perturbation."""
         ac_ = np.empty((self.n_steps, self.draws))
 
-        log_R = np.log(np.random.rand(self.n_steps, self.draws))
+        log_R = np.log(self.rng.random((self.n_steps, self.draws)))
         for n_step in range(self.n_steps):
             proposal = floatX(
                 self.tempered_posterior
-                + self.proposal_dist(num_draws=self.draws) * self.proposal_scales[:, None]
+                + self.proposal_dist(num_draws=self.draws, rng=self.rng)
+                * self.proposal_scales[:, None]
             )
             ll = np.array([self.likelihood_logp_func(prop) for prop in proposal])
             pl = np.array([self.prior_logp_func(prop) for prop in proposal])

pymc/tests/helpers.py

@@ -17,6 +17,7 @@
 from logging.handlers import BufferingHandler
 
 import aesara
+import numpy as np
 import numpy.random as nr
 
 from aesara.gradient import verify_grad as at_verify_grad
@@ -123,3 +124,11 @@ def verify_grad(op, pt, n_tests=2, rng=None, *args, **kwargs):
     if rng is None:
         rng = nr.RandomState(411342)
     at_verify_grad(op, pt, n_tests, rng, *args, **kwargs)
+
+
+def assert_random_state_equal(state1, state2):
+    for field1, field2 in zip(state1, state2):
+        if isinstance(field1, np.ndarray):
+            np.testing.assert_array_equal(field1, field2)
+        else:
+            assert field1 == field2

pymc/tests/test_smc.py

@@ -32,7 +32,7 @@
 from pymc.aesaraf import floatX
 from pymc.backends.base import MultiTrace
 from pymc.smc.smc import IMH
-from pymc.tests.helpers import SeededTest
+from pymc.tests.helpers import SeededTest, assert_random_state_equal
 
 
 class TestSMC(SeededTest):
@@ -77,8 +77,10 @@ def two_gaussians(x):
             y = pm.Normal("y", x, 1, observed=0)
 
     def test_sample(self):
+        initial_rng_state = np.random.get_state()
         with self.SMC_test:
             mtrace = pm.sample_smc(draws=self.samples, return_inferencedata=False)
+        assert_random_state_equal(initial_rng_state, np.random.get_state())
 
         x = mtrace["X"]
         mu1d = np.abs(x).mean(axis=0)
@@ -531,11 +533,14 @@ def test_named_model(self):
 class TestMHKernel(SeededTest):
     def test_normal_model(self):
         data = st.norm(10, 0.5).rvs(1000, random_state=self.get_random_state())
+
+        initial_rng_state = np.random.get_state()
         with pm.Model() as m:
             mu = pm.Normal("mu", 0, 3)
             sigma = pm.HalfNormal("sigma", 1)
             y = pm.Normal("y", mu, sigma, observed=data)
             idata = pm.sample_smc(draws=2000, kernel=pm.smc.MH)
+        assert_random_state_equal(initial_rng_state, np.random.get_state())
 
         post = idata.posterior.stack(sample=("chain", "draw"))
         assert np.abs(post["mu"].mean() - 10) < 0.1