From 6b31ec19bd27df24be5e17e15f1f73d50dd02bfa Mon Sep 17 00:00:00 2001 From: "Michael J. Williams" Date: Mon, 19 Oct 2026 17:13:42 +0200 Subject: [PATCH 1/2] MAINT: add `random_state` to sampler utils --- bilby/core/sampler/base_sampler.py | 8 ++++++-- bilby/core/sampler/dynesty.py | 16 +++++++++++----- 2 files changed, 17 insertions(+), 7 deletions(-) diff --git a/bilby/core/sampler/base_sampler.py b/bilby/core/sampler/base_sampler.py index e7aa9cbb7..a0a943406 100644 --- a/bilby/core/sampler/base_sampler.py +++ b/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) diff --git a/bilby/core/sampler/dynesty.py b/bilby/core/sampler/dynesty.py index 563cc24e2..7a9bb554d 100644 --- a/bilby/core/sampler/dynesty.py +++ b/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,12 +543,16 @@ 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) @@ -556,14 +560,16 @@ def _generate_result(self, out): 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, From 809e526ce249ae855492ce4dcc0c9bb79059e4e9 Mon Sep 17 00:00:00 2001 From: "Michael J. Williams" Date: Mon, 19 Oct 2026 17:14:56 +0200 Subject: [PATCH 2/2] TST: update tests for sampler rng --- test/core/sampler/base_sampler_test.py | 22 ++++++++++++++ test/core/sampler/dynesty_test.py | 40 +++++++++++++++++++------- 2 files changed, 52 insertions(+), 10 deletions(-) diff --git a/test/core/sampler/base_sampler_test.py b/test/core/sampler/base_sampler_test.py index 2a52a24c3..7ef8757f3 100644 --- a/test/core/sampler/base_sampler_test.py +++ b/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() diff --git a/test/core/sampler/dynesty_test.py b/test/core/sampler/dynesty_test.py index ca15eb021..2ffe3dfdc 100644 --- a/test/core/sampler/dynesty_test.py +++ b/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__":