MAINT: ensure dynesty is seeded when global RNG is not set

bilby/core/sampler/base_sampler.py

@@ -538,7 +538,7 @@ def get_random_draw_from_prior(self):
         self.check_draw(draw)
         return draw
 
-    def get_initial_points_from_prior(self, npoints=1):
+    def get_initial_points_from_prior(self, npoints=1, *, random_state=None):
         """Method to draw a set of live points from the prior
 
         This iterates over draws from the prior until all the samples have a
@@ -548,6 +548,9 @@ def get_initial_points_from_prior(self, npoints=1):
         ==========
         npoints: int
             The number of values to return
+        random_state: None, int, numpy.random.Generator
+            Random state used to draw the initial points. If None, use Bilby's
+            global random number generator.
 
         Returns
         =======
@@ -560,12 +563,13 @@ def get_initial_points_from_prior(self, npoints=1):
         """
         from ..utils import random
 
+        rng = random.resolve_random_state(random_state)
         logger.info("Generating initial points from the prior")
         unit_cube = []
         parameters = []
         likelihood = []
         while len(unit_cube) < npoints:
-            unit = random.rng.uniform(0, 1, self.ndim)
+            unit = rng.uniform(0, 1, self.ndim)
             theta = self.prior_transform(unit)
             if self.check_draw(theta, warning=False):
                 unit_cube.append(unit)

bilby/core/sampler/dynesty.py

@@ -492,7 +492,7 @@ def run_sampler(self):
         else:
             if self.kwargs["live_points"] is None:
                 self.kwargs["live_points"] = self.get_initial_points_from_prior(
-                    self.nlive
+                    self.nlive, random_state=self.kwargs["rstate"]
                 )
                 self.kwargs["live_points"] = (*self.kwargs["live_points"], None)
             self.sampler = self.sampler_init(
@@ -528,12 +528,12 @@ def run_sampler(self):
         dynesty_result = f"{self.outdir}/{self.label}_dynesty.pickle"
         safe_file_dump(out, dynesty_result, "dill")
 
-        self._generate_result(out)
+        self._generate_result(out, random_state=self.sampler.rstate)
         self.result.sampling_time = self.sampling_time
 
         return self.result
 
-    def _generate_result(self, out):
+    def _generate_result(self, out, *, random_state=None):
         """
         Extract the information we need from the dynesty output. This includes
         the evidence, nested samples, run statistics. In addition, we generate
@@ -543,27 +543,33 @@ def _generate_result(self, out):
         ==========
         out: dynesty.result.Result
             The dynesty output.
+        random_state: None, int, numpy.random.Generator
+            Random state used to generate posterior samples. If None, use
+            Bilby's global random number generator.
         """
         import dynesty
         from scipy.special import logsumexp
 
         from ..utils import random
 
+        rng = random.resolve_random_state(random_state)
         logwts = out["logwt"]
         weights = np.exp(logwts - out["logz"][-1])
         nested_samples = DataFrame(out.samples, columns=self.search_parameter_keys)
         nested_samples["weights"] = weights
         nested_samples["log_likelihood"] = out.logl
         self.result.nested_samples = nested_samples
         if self.rejection_sample_posterior:
-            keep = weights > random.rng.uniform(0, max(weights), len(weights))
+            keep = weights > rng.uniform(0, max(weights), len(weights))
             self.result.samples = out.samples[keep]
             self.result.log_likelihood_evaluations = out.logl[keep]
             logger.info(
                 f"Rejection sampling nested samples to obtain {sum(keep)} posterior samples"
             )
         else:
-            self.result.samples = dynesty.utils.resample_equal(out.samples, weights)
+            self.result.samples = dynesty.utils.resample_equal(
+                out.samples, weights, rstate=rng
+            )
             self.result.log_likelihood_evaluations = self.reorder_loglikelihoods(
                 unsorted_loglikelihoods=out.logl,
                 unsorted_samples=out.samples,

test/core/sampler/base_sampler_test.py

@@ -113,6 +113,28 @@ def test_log_likelihood_without_use_ratio(self):
     def test_get_random_draw(self):
         self.assertEqual(self.sampler.get_random_draw_from_prior(), np.array([0.5]))
 
+    def test_initial_points_with_random_state(self):
+        global_rng = bilby.core.utils.random.rng
+        global_state = copy.deepcopy(global_rng.bit_generator.state)
+        with patch.object(self.sampler, "check_draw", return_value=True), patch.object(
+            self.sampler, "prior_transform", side_effect=lambda unit: unit
+        ):
+            points, _, _ = self.sampler.get_initial_points_from_prior(
+                5, random_state=np.random.default_rng(1234)
+            )
+        expected = np.random.default_rng(1234).uniform(0, 1, (5, self.sampler.ndim))
+        np.testing.assert_array_equal(points, expected)
+        self.assertEqual(global_rng.bit_generator.state, global_state)
+
+    def test_initial_points_default_random_state(self):
+        rng = np.random.default_rng(1234)
+        with patch.object(bilby.core.utils.random.Generator, "rng", rng), patch.object(
+            self.sampler, "check_draw", return_value=True
+        ), patch.object(self.sampler, "prior_transform", side_effect=lambda unit: unit):
+            points, _, _ = self.sampler.get_initial_points_from_prior(5)
+        expected = np.random.default_rng(1234).uniform(0, 1, (5, self.sampler.ndim))
+        np.testing.assert_array_equal(points, expected)
+
     def test_base_run_sampler(self):
         sampler_copy = copy.copy(self.sampler)
         self.sampler.run_sampler()

test/core/sampler/dynesty_test.py

@@ -486,7 +486,6 @@ def _remove_tree(self):
             pass
 
     def _run_sampler(self, **kwargs):
-        bilby.core.utils.random.seed(42)
         return bilby.run_sampler(
             likelihood=self.likelihood,
             priors=self.priors,
@@ -496,26 +495,47 @@ def _run_sampler(self, **kwargs):
             dlogz=1.0,
             nlive=20,
             sample="acceptance-walk",
+            check_point=False,
+            check_point_plot=False,
+            print_progress=False,
             **kwargs,
         )
 
-    def test_reproducibility_seed(self):
-        res0 = self._run_sampler(seed=1234)
-        res1 = self._run_sampler(seed=1234)
+    @parameterized.parameterized.expand([True, False])
+    def test_reproducibility_seed(self, rejection_sample_posterior):
+        res0 = self._run_sampler(
+            seed=1234, rejection_sample_posterior=rejection_sample_posterior
+        )
+        res1 = self._run_sampler(
+            seed=1234, rejection_sample_posterior=rejection_sample_posterior
+        )
         assert res0.log_evidence == res1.log_evidence
+        np.testing.assert_array_equal(res0.samples, res1.samples)
 
-    def test_reproducibility_state(self):
+    @parameterized.parameterized.expand([True, False])
+    def test_reproducibility_state(self, rejection_sample_posterior):
         rstate = np.random.default_rng(1234)
-        res0 = self._run_sampler(rstate=rstate)
+        res0 = self._run_sampler(
+            rstate=rstate, rejection_sample_posterior=rejection_sample_posterior
+        )
         rstate = np.random.default_rng(1234)
-        res1 = self._run_sampler(rstate=rstate)
+        res1 = self._run_sampler(
+            rstate=rstate, rejection_sample_posterior=rejection_sample_posterior
+        )
         assert res0.log_evidence == res1.log_evidence
+        np.testing.assert_array_equal(res0.samples, res1.samples)
 
-    def test_reproducibility_state_and_seed(self):
+    @parameterized.parameterized.expand([True, False])
+    def test_reproducibility_state_and_seed(self, rejection_sample_posterior):
         rstate = np.random.default_rng(1234)
-        res0 = self._run_sampler(rstate=rstate)
-        res1 = self._run_sampler(seed=1234)
+        res0 = self._run_sampler(
+            rstate=rstate, rejection_sample_posterior=rejection_sample_posterior
+        )
+        res1 = self._run_sampler(
+            seed=1234, rejection_sample_posterior=rejection_sample_posterior
+        )
         assert res0.log_evidence == res1.log_evidence
+        np.testing.assert_array_equal(res0.samples, res1.samples)
 
 
 if __name__ == "__main__":