[fix] RNG arguments for distribution samplers

conda/meta.yaml

@@ -1,5 +1,5 @@
 {% set name = "posydon" %}
-{% set version = "2.2.6" %}
+{% set version = "2.2.7" %}
 
 package:
   name: "{{ name|lower }}"

posydon/popsyn/Moes_distributions.py

@@ -114,6 +114,8 @@ def __init__(self, n_M1=101, n_logP=158, n_q=91, n_e=200,
         self.numq = n_q
         self.nume = n_e
 
+        self.RNG = kwargs.get('RNG', np.random.default_rng())
+
         # ranges where M+D17 has statistics corrected for selection effects:
         # 0.8 < M1/Msun < 40 with self.numM1 steps in log space
         self.M1v = 10**(np.linspace(np.log10(0.8), np.log10(40.0), self.numM1))
@@ -429,7 +431,7 @@ def __call__(self, M1, M_min=0.08, M_max=150.0, all_binaries=True):
             else:
                 mybinfrac = np.max(mycumPbindist())
             # Generate random number myrand between 0 and 1
-            myrand = np.random.rand()
+            myrand = self.RNG.random()
             # If random number < binary star fraction, generate a binary
             if(myrand < mybinfrac):
                 # Given myrand, select P and corresponding index in logPv
@@ -438,7 +440,7 @@ def __call__(self, M1, M_min=0.08, M_max=150.0, all_binaries=True):
                                          == min(abs(mylogP-self.logPv)))
                 indlogP = indlogP[0]
                 # Given M1 & P, select e from eccentricity distribution
-                mye = np.interp(np.random.rand(),
+                mye = np.interp(self.RNG.random(),
                                 self.cumedist[:,indlogP,indM1].flatten(), self.ev)
                 # Given M1 & P, determine mass ratio distribution.
                 # If M1 < 0.8 Msun, truncate q distribution and consider only mass
@@ -454,7 +456,7 @@ def __call__(self, M1, M_min=0.08, M_max=150.0, all_binaries=True):
                     # Set probability = 0 where q < q_min
                     mycumqdist[self.qv <= q_min] = 0.0
                 # Given M1 & P, select q from cumulative mass ratio distribution
-                myq = np.interp(np.random.rand(), mycumqdist, self.qv)
+                myq = np.interp(self.RNG.random(), mycumqdist, self.qv)
             else:
                 # If instead random number > binary star fraction, generate single
                 # star
@@ -463,7 +465,7 @@ def __call__(self, M1, M_min=0.08, M_max=150.0, all_binaries=True):
                 mye = np.nan
             # Get metallicity
             Zsun = 0.02
-            logZ = -2.3 + (0.176-(-2.3)) * np.random.rand()
+            logZ = -2.3 + (0.176-(-2.3)) * self.RNG.random()
             Z = Zsun*10**logZ
             M2s.append(M1*myq)
             logPs.append(mylogP)

posydon/popsyn/independent_sample.py

@@ -164,7 +164,8 @@ def pdf(logp):
         orbital_periods_M_gt_15 = 10**(rejection_sampler(
             size=number_of_binaries,
             x_lim=[np.log10(orbital_period_min), np.log10(orbital_period_max)],
-            pdf=pdf))
+            pdf=pdf,
+            rng=RNG))
 
         orbital_periods = np.where(primary_masses <= 15.0,
                                    orbital_periods_M_lt_15,

posydon/popsyn/star_formation_history.py

@@ -1102,15 +1102,15 @@ def get_formation_times(N_binaries, star_formation="constant", **kwargs): # prag
     if scenario in ["custom_linear", "custom_log10"]:
         custom_ages_file = kwargs.get("custom_ages_file")
         x, y = np.loadtxt(custom_ages_file, unpack=True)
-        current_binary_ages = rejection_sampler(x, y, N_binaries)
+        current_binary_ages = rejection_sampler(x, y, N_binaries, rng=RNG)
         if "log10" in scenario:
             current_binary_ages = 10.0**current_binary_ages
         return max_time - current_binary_ages
 
     if scenario in ["custom_linear_histogram", "custom_log10_histogram"]:
         custom_ages_file = kwargs.get("custom_ages_file")
         x, y = read_histogram_from_file(custom_ages_file)
-        current_binary_ages = histogram_sampler(x, y)
+        current_binary_ages = histogram_sampler(x, y, rng=RNG)
         if "log10" in scenario:
             current_binary_ages = 10.0**current_binary_ages
         return max_time - current_binary_ages

posydon/unit_tests/utils/test_common_functions.py

@@ -754,8 +754,10 @@ def mock_rand2(shape):
                approx(5.13970075150e-8, abs=6e-20)
 
     def test_rejection_sampler(self, monkeypatch):
-        def mock_uniform(low=0.0, high=1.0, size=1):
-            return np.linspace(low, high, num=size)
+        class MockRNG:
+            def uniform(self, low=0.0, high=1.0, size=1):
+                return np.linspace(low, high, num=size)
+        mock_rng = MockRNG()
         def mock_interp1d(x, y):
             if x[0] > x[-1]:
                 raise ValueError
@@ -786,7 +788,6 @@ def mock_pdf(x):
                                       +" function in rejection sampling"):
             totest.rejection_sampler(pdf=mock_pdf)
         # examples:
-        monkeypatch.setattr(np.random, "uniform", mock_uniform)
         monkeypatch.setattr(totest, "interp1d", mock_interp1d)
         tests = [(np.array([0.0, 1.0]), np.array([0.4, 0.6]), 5,\
                   np.array([0.0, 0.25, 0.5, 0.75, 1.0])),\
@@ -795,22 +796,24 @@ def mock_pdf(x):
                  (np.array([1.0, 0.0]), np.array([0.2, 0.8]), 6,\
                   np.array([0.0, 0.2, 0.4, 0.0, 0.5, 0.0]))]
         for (x, y, s, r) in tests:
-            assert np.allclose(totest.rejection_sampler(x=x, y=y, size=s),\
+            assert np.allclose(totest.rejection_sampler(x=x, y=y, size=s, rng=mock_rng),\
                                r)
         assert np.allclose(totest.rejection_sampler(x_lim=np.array([0.0,\
                                                                      1.0]),\
-                                                     pdf=mock_pdf, size=5),\
+                                                     pdf=mock_pdf, size=5, rng=mock_rng),\
                            np.array([0.0, 0.25, 0.0, 0.0, 0.0]))
         assert np.allclose(totest.rejection_sampler(x=np.array([1.0, 0.0]),\
                                                        y=np.array([0.2, 0.8]),\
                                                        x_lim=np.array([0.0,\
                                                                        1.0]),\
-                                                       pdf=mock_pdf, size=5),\
+                                                       pdf=mock_pdf, size=5, rng=mock_rng),\
                               np.array([0.0, 0.25, 0.0, 0.0, 0.0]))
 
     def test_inverse_sampler(self, monkeypatch):
-        def mock_uniform(low=0.0, high=1.0, size=1):
-            return np.linspace(low, high, num=size)
+        class MockRNG:
+            def uniform(self, low=0.0, high=1.0, size=1):
+                return np.linspace(low, high, num=size)
+        mock_rng = MockRNG()
         # missing argument
         with raises(TypeError, match="missing 2 required positional "\
                                      +"arguments: 'x' and 'y'"):
@@ -826,26 +829,37 @@ def mock_uniform(low=0.0, high=1.0, size=1):
             totest.inverse_sampler(x=np.array([0.0, 1.0]),\
                                    y=np.array([-0.4, 0.6]))
         # examples:
-        monkeypatch.setattr(np.random, "uniform", mock_uniform)
         assert np.allclose(totest.inverse_sampler(x=np.array([0.0, 1.0]),\
                                                   y=np.array([0.4, 0.6]),\
-                                                  size=5),\
+                                                  size=5, rng=mock_rng),\
                            np.array([0.0, 0.29128785, 0.54950976, 0.78388218,\
                                      1.0]))
         assert np.allclose(totest.inverse_sampler(x=np.array([0.0, 1.0]),\
                                                   y=np.array([0.6, 0.4]),\
-                                                  size=4),\
+                                                  size=4, rng=mock_rng),\
                            np.array([0.0, 0.2919872, 0.61952386, 1.0]))
         with warns(RuntimeWarning, match="invalid value encountered in "\
                                          +"divide"):
             assert np.allclose(totest.inverse_sampler(x=np.array([0.0, 1.0]),\
                                                       y=np.array([0.5, 0.5]),\
-                                                      size=5),\
+                                                      size=5, rng=mock_rng),\
                                np.array([0.0, 0.25, 0.5, 0.75, 1.0]))
 
     def test_histogram_sampler(self, monkeypatch):
-        def mock_uniform(low=0.0, high=1.0, size=1):
-            return np.linspace(low, high, num=size)
+        class MockRNG:
+            def uniform(self, low=0.0, high=1.0, size=1):
+                return np.linspace(low, high, num=size)
+            def choice(self, a, size=None, replace=True, p=None):
+                if isinstance(a, int):
+                    a = np.arange(a)
+                sample = []
+                for (v, q) in zip(a, p):
+                    sample += round(size*q) * [v]
+                if len(sample) < size:
+                    sample += (size-len(sample)) * [a[-1]]
+                return np.array(sample[:size])
+
+        mock_rng = MockRNG()
         def mock_choice(a, size=None, replace=True, p=None):
             if isinstance(a, int):
                 a=np.arange(a)
@@ -870,20 +884,18 @@ def mock_choice(a, size=None, replace=True, p=None):
             totest.histogram_sampler(x_edges=np.array([0.0, 1.0]),\
                                      y=np.array([0.4, 0.6]))
         # examples:
-        monkeypatch.setattr(np.random, "uniform", mock_uniform)
-        monkeypatch.setattr(np.random, "choice", mock_choice)
         assert np.allclose(totest.histogram_sampler(x_edges=np.array([0.0,\
                                                                       0.5,\
                                                                       1.0]),\
                                                     y=np.array([0.2, 0.8]),\
-                                                    size=5),\
+                                                    size=5, rng=mock_rng),\
                            np.array([0.0, 0.5, 0.66666667, 0.83333333, 1.0]))
         assert np.allclose(totest.histogram_sampler(x_edges=np.array([0.0,\
                                                                          0.5,\
                                                                          1.0]\
                                                        ), y=np.array([0.2,\
                                                                       0.8]),\
-                                                       size=4),\
+                                                       size=4, rng=mock_rng),\
                               np.array([0.0, 0.5, 0.75, 1.0]))
 
     def test_read_histogram_from_file(self, csv_path_failing_3_data_lines,\

posydon/utils/common_functions.py

@@ -669,7 +669,7 @@ def bondi_hoyle(binary, accretor, donor, idx=-1, wind_disk_criteria=True,
     return np.squeeze(mdot_acc)
 
 
-def rejection_sampler(x=None, y=None, size=1, x_lim=None, pdf=None):
+def rejection_sampler(x=None, y=None, size=1, x_lim=None, pdf=None, rng=None):
     """Generate a sample from a 1d PDF using the acceptance-rejection method.
 
     Parameters
@@ -684,13 +684,17 @@ def rejection_sampler(x=None, y=None, size=1, x_lim=None, pdf=None):
         The boundary where the pdf function is defined if passed as pdf.
     pdf : func
         The pdf function
+    rng : numpy.random.Generator, optional
+        Random number generator. If None, uses np.random.default_rng().
 
     Returns
     -------
     ndarray
         An array with `size` random numbers generated from the PDF.
 
     """
+    if rng is None:
+        rng = np.random.default_rng()
     if pdf is None:
         if ((x is None) or (y is None)):
             raise ValueError("x and y PDF values must be specified if no PDF"
@@ -702,32 +706,32 @@ def rejection_sampler(x=None, y=None, size=1, x_lim=None, pdf=None):
             idxs = np.argsort(x)
             pdf = interp1d(x.take(idxs), y.take(idxs))
 
-        x_rand = np.random.uniform(x.min(), x.max(), size)
-        y_rand = np.random.uniform(0, y.max(), size)
+        x_rand = rng.uniform(x.min(), x.max(), size)
+        y_rand = rng.uniform(0, y.max(), size)
         values = x_rand[y_rand <= pdf(x_rand)]
         while values.shape[0] < size:
             n = size - values.shape[0]
-            x_rand = np.random.uniform(x.min(), x.max(), n)
-            y_rand = np.random.uniform(0, y.max(), n)
+            x_rand = rng.uniform(x.min(), x.max(), n)
+            y_rand = rng.uniform(0, y.max(), n)
             values = np.hstack([values, x_rand[y_rand <= pdf(x_rand)]])
     else:
         if x_lim is None:
             raise ValueError("x_lim must be specified for passed PDF function"
                              " in rejection sampling")
-        x_rand = np.random.uniform(x_lim[0], x_lim[1], size)
-        pdf_max = max(pdf(np.random.uniform(x_lim[0], x_lim[1], 50000)))
-        y_rand = np.random.uniform(0, pdf_max, size)
+        x_rand = rng.uniform(x_lim[0], x_lim[1], size)
+        pdf_max = max(pdf(rng.uniform(x_lim[0], x_lim[1], 50000)))
+        y_rand = rng.uniform(0, pdf_max, size)
         values = x_rand[y_rand <= pdf(x_rand)]
         while values.shape[0] < size:
             n = size - values.shape[0]
-            x_rand = np.random.uniform(x_lim[0], x_lim[1], n)
-            y_rand = np.random.uniform(0, pdf_max, n)
+            x_rand = rng.uniform(x_lim[0], x_lim[1], n)
+            y_rand = rng.uniform(0, pdf_max, n)
             values = np.hstack([values, x_rand[y_rand <= pdf(x_rand)]])
 
     return values
 
 
-def inverse_sampler(x, y, size=1):
+def inverse_sampler(x, y, size=1, rng=None):
     """Sample from a PDF using the inverse-sampling method.
 
     Parameters
@@ -738,13 +742,17 @@ def inverse_sampler(x, y, size=1):
         The probablity density at `x` (or a scaled version of it).
     size : int
         Number of samples to generate.
+    rng : numpy.random.Generator, optional
+        Random number generator. If None, uses np.random.default_rng().
 
     Returns
     -------
     array
         The sample drawn from the PDF.
 
     """
+    if rng is None:
+        rng = np.random.default_rng()
     # x should be sorted
     assert np.all(np.diff(x) > 0)
     # y should be above 0
@@ -757,7 +765,7 @@ def inverse_sampler(x, y, size=1):
     total_area = cumsum_areas[-1]
 
     # start the inverse sampling
-    u = np.random.uniform(size=size)
+    u = rng.uniform(size=size)
     # index of "bin" where each sampled value corresponds too
     u_indices = np.searchsorted(cumsum_areas, u * total_area)
     # the area that should be covered from the end of the previous bin
@@ -779,14 +787,14 @@ def inverse_sampler(x, y, size=1):
     if n_where_nan:
         assert np.all(dy_bins[where_nan] == 0)
         sample[where_nan] = x[u_indices][where_nan] + (
-            dx_bins[where_nan] * np.random.uniform(size=n_where_nan))
+            dx_bins[where_nan] * rng.uniform(size=n_where_nan))
 
     # make sure that everything worked as expected
     assert np.all(np.isfinite(sample))
     return sample
 
 
-def histogram_sampler(x_edges, y, size=1):
+def histogram_sampler(x_edges, y, size=1, rng=None):
     """Sample from an empirical PDF represented by a histogram.
 
     Parameters
@@ -797,28 +805,32 @@ def histogram_sampler(x_edges, y, size=1):
         The counts (or a scaled version of them) of the histogram.
     size : int
         Number of random values to produce.
+    rng : numpy.random.Generator, optional
+        Random number generator. If None, uses np.random.default_rng().
 
     Returns
     -------
     array
         The random sample.
 
     """
+    if rng is None:
+        rng = np.random.default_rng()
     assert np.all(y >= 0.0)
 
     # make sure that the lengths of the input arrays are correct
     n_bins = len(y)
     assert n_bins > 0 and len(x_edges) == n_bins + 1
     # first decide which will be the bin of each element in the sample
-    bin_sample = np.random.choice(n_bins, replace=True, p=y/sum(y), size=size)
+    bin_sample = rng.choice(n_bins, replace=True, p=y/sum(y), size=size)
 
     sample = np.ones(size) * np.nan
 
     # select each bin and based on its uniform distribution, decide the sample
     bins_found = set(bin_sample)
     for bin_index in bins_found:
         in_this_bin = np.argwhere(bin_sample == bin_index)[:, 0]
-        sample[in_this_bin] = np.random.uniform(
+        sample[in_this_bin] = rng.uniform(
             x_edges[bin_index], x_edges[bin_index+1], size=len(in_this_bin))
 
     assert(np.all(np.isfinite(sample)))