[v2.3] Pass RNG to evolution steps

posydon/binary_evol/DT/step_detached.py

@@ -163,17 +163,18 @@ class detached_step:
 
     # settings in .ini will override
     DEFAULT_KWARGS = {"dt": None,
-                        "n_o_steps_history": None,
-                        "do_wind_loss": True,
-                        "do_tides": True,
-                        "do_gravitational_radiation": True,
-                        "do_magnetic_braking": True,
-                        "magnetic_braking_mode": "RVJ83",
-                        "do_stellar_evolution_and_spin_from_winds": True,
-                        "RLO_orbit_at_orbit_with_same_am": False,
-                        "metallicity": None,
-                        "track_matcher": None,
-                        "verbose": False}
+                      "n_o_steps_history": None,
+                      "do_wind_loss": True,
+                      "do_tides": True,
+                      "do_gravitational_radiation": True,
+                      "do_magnetic_braking": True,
+                      "magnetic_braking_mode": "RVJ83",
+                      "do_stellar_evolution_and_spin_from_winds": True,
+                      "RLO_orbit_at_orbit_with_same_am": False,
+                      "metallicity": None,
+                      "track_matcher": None,
+                      "RNG": np.random.default_rng(),
+                      "verbose": False}
 
     def __init__(self, **kwargs):
 
@@ -224,12 +225,12 @@ def init_evo_kwargs(self):
             "magnetic_braking_mode": self.magnetic_braking_mode,
             "do_stellar_evolution_and_spin_from_winds": self.do_stellar_evolution_and_spin_from_winds,
             "do_gravitational_radiation": self.do_gravitational_radiation,
-            "verbose": self.verbose,
-        }
+            "verbose": self.verbose}
 
     def __repr__(self):
         """Return the name of evolution step."""
-        return "Detached Step."
+        return "detached_step:\n" + \
+            "\n".join([f"{key} = {getattr(self, key)}" for key in self.__dict__])
 
     def __call__(self, binary):
         """
@@ -352,7 +353,7 @@ def __call__(self, binary):
             elif primary.co:
                 mdot_acc = np.atleast_1d(bondi_hoyle(
                     binary, primary, secondary, slice(-len(t), None),
-                    wind_disk_criteria=True, scheme='Kudritzki+2000'))
+                    wind_disk_criteria=True, RNG=self.RNG, scheme='Kudritzki+2000'))
                 primary.lg_mdot = np.log10(mdot_acc.item(-1))
                 primary.lg_mdot_history[len(primary.lg_mdot_history) - len(t) + 1:] = np.log10(mdot_acc[:-1])
             else:

posydon/binary_evol/MESA/step_mesa.py

@@ -139,6 +139,7 @@ class MesaGridStep:
                       'stop_var_name': None,
                       'stop_value': None,
                       'stop_interpolate': True,
+                      'RNG': np.random.default_rng(),
                       'verbose': False}
 
     def __init__(self, **kwargs):
@@ -775,7 +776,7 @@ def update_properties_NN(self, star_1_CO=False, star_2_CO=False,
             key_bh = POSYDON_TO_MESA['star']['lg_mdot']+'_%d' % (k_bh+1)
             tmp_lg_mdot = np.log10(10**cb_bh[key_bh][-1] + cf.bondi_hoyle(
                 binary, accretor, donor, idx=-1,
-                wind_disk_criteria=True, scheme='Kudritzki+2000'))
+                wind_disk_criteria=True, RNG=self.RNG, scheme='Kudritzki+2000'))
             mdot_edd = cf.eddington_limit(binary, idx=-1)[0]
 
             if 10**tmp_lg_mdot > mdot_edd:
@@ -788,7 +789,7 @@ def update_properties_NN(self, star_1_CO=False, star_2_CO=False,
                 history_of_attribute = (np.log10(
                     10**cb_bh[key_bh][0] + cf.bondi_hoyle(
                         binary, accretor, donor, idx=len_binary_hist,
-                        wind_disk_criteria=True, scheme='Kudritzki+2000')))
+                        wind_disk_criteria=True, RNG=self.RNG, scheme='Kudritzki+2000')))
                 if 10**history_of_attribute > edd:
                     history_of_attribute = np.log10(edd)
                 accretor.lg_mdot_history.append(history_of_attribute)
@@ -800,6 +801,7 @@ def update_properties_NN(self, star_1_CO=False, star_2_CO=False,
                 # hence we loop one back range(-N-1,-1)
                 tmp_h = [cf.bondi_hoyle(binary, accretor, donor, idx=i,
                                         wind_disk_criteria=True,
+                                        RNG=self.RNG,
                                         scheme='Kudritzki+2000')
                          for i in range(-length_hist-1, -1)]
                 tmp_edd = [cf.eddington_limit(binary, idx=i)[0]
@@ -966,7 +968,8 @@ def initial_final_interpolation(self, star_1_CO=False, star_2_CO=False):
             tmp_lg_mdot = np.log10(
                 10**fv[key_bh] + cf.bondi_hoyle(
                     binary, accretor, donor, idx=-1,
-                    wind_disk_criteria=True, scheme='Kudritzki+2000'))
+                    wind_disk_criteria=True,
+                    RNG=self.RNG, scheme='Kudritzki+2000'))
 
             mdot_edd = cf.eddington_limit(binary, idx=-1)[0]
             if 10**tmp_lg_mdot > mdot_edd:

posydon/binary_evol/SN/step_SN.py

@@ -249,7 +249,8 @@ class StepSN(object):
         "sigma_kick_ECSN": 20.0,
         "mean_kick_ECSN": None,
         # other
-        "verbose": False,
+        "RNG": np.random.default_rng(),
+        "verbose": False
     }
     # add core collapse physics
     DEFAULT_KWARGS.update(DEFAULT_SN_MODEL)
@@ -1538,13 +1539,13 @@ def orbital_kick(self, binary):
             if not binary.star_1.natal_kick_azimuthal_angle is None:
                 phi = binary.star_1.natal_kick_azimuthal_angle
             else:
-                phi = np.random.uniform(0, 2 * np.pi)
+                phi = self.RNG.uniform(0, 2 * np.pi)
                 binary.star_1.natal_kick_azimuthal_angle = phi
 
             if not binary.star_1.natal_kick_polar_angle is None:
                 cos_theta = np.cos(binary.star_1.natal_kick_polar_angle)
             else:
-                cos_theta = np.random.uniform(-1, 1)
+                cos_theta = self.RNG.uniform(-1, 1)
                 binary.star_1.natal_kick_polar_angle = np.arccos(cos_theta)
 
             # generate random point in the orbit where the kick happens
@@ -1555,7 +1556,7 @@ def orbital_kick(self, binary):
                     raise ValueError("mean_anomaly must be a single float value."
                                      f"\n mean_anomaly = {mean_anomaly}")
             else:
-                mean_anomaly = np.random.uniform(0, 2 * np.pi)
+                mean_anomaly = self.RNG.uniform(0, 2 * np.pi)
                 binary.star_1.natal_kick_mean_anomaly = mean_anomaly
 
         elif binary.event == "CC2":
@@ -1639,13 +1640,13 @@ def orbital_kick(self, binary):
             if not binary.star_2.natal_kick_azimuthal_angle is None:
                 phi = binary.star_2.natal_kick_azimuthal_angle
             else:
-                phi = np.random.uniform(0, 2 * np.pi)
+                phi = self.RNG.uniform(0, 2 * np.pi)
                 binary.star_2.natal_kick_azimuthal_angle = phi
 
             if not binary.star_2.natal_kick_polar_angle is None:
                 cos_theta = np.cos(binary.star_2.natal_kick_polar_angle)
             else:
-                cos_theta = np.random.uniform(-1, 1)
+                cos_theta = self.RNG.uniform(-1, 1)
                 binary.star_2.natal_kick_polar_angle = np.arccos(cos_theta)
 
             # generate random point in the orbit where the kick happens
@@ -1655,7 +1656,7 @@ def orbital_kick(self, binary):
                 if not isinstance(mean_anomaly, float):
                     raise ValueError("mean_anomaly must be a single float value.")
             else:
-                mean_anomaly = np.random.uniform(0, 2 * np.pi)
+                mean_anomaly = self.RNG.uniform(0, 2 * np.pi)
                 binary.star_2.natal_kick_mean_anomaly = mean_anomaly
 
         # update the orbit
@@ -2081,7 +2082,7 @@ def _get_kick_velocity(self, star, sigma=None, mean=None):
             # this is a fallback
             if sigma is None:
                 sigma = 265.0
-            Vkick_ej = sp.stats.maxwell.rvs(loc=0., scale=sigma, size=1)[0]
+            Vkick_ej = sp.stats.maxwell.rvs(loc=0., scale=sigma, size=1, random_state=self.RNG)[0]
 
         elif self.kick_prescription == "log_normal":
             # sigma==None should never be reached, since in that case Vkick=0
@@ -2091,7 +2092,7 @@ def _get_kick_velocity(self, star, sigma=None, mean=None):
                 sigma = 0.68
             if mean is None:
                 mean = np.exp(5.60)
-            Vkick_ej = sp.stats.lognorm.rvs(s=sigma, scale=mean, size=1)[0]
+            Vkick_ej = sp.stats.lognorm.rvs(s=sigma, scale=mean, size=1, random_state=self.RNG)[0]
 
         elif self.kick_prescription == "asym_ej":
             f_kin = 0.1         # Fraction of SN explosion energy that is kinetic energy of the gas

posydon/binary_evol/simulationproperties.py

@@ -17,6 +17,8 @@
 import os
 import time
 
+import numpy as np
+
 from posydon.binary_evol.track_match import TrackMatcher
 from posydon.config import PATH_TO_POSYDON_DATA
 from posydon.interpolation.interpolation import GRIDInterpolator
@@ -229,7 +231,8 @@ def preload_imports(self):
         self._MesaGridStep = MesaGridStep
 
     @classmethod
-    def from_ini(cls, path, metallicity = None, load_steps=False, verbose=False, **override_sim_kwargs):
+    def from_ini(cls, path, metallicity = None, load_steps=False, RNG=np.random.default_rng(),
+                 verbose=False, **override_sim_kwargs):
         """Create a SimulationProperties instance from an inifile.
 
         Parameters
@@ -246,9 +249,19 @@ def from_ini(cls, path, metallicity = None, load_steps=False, verbose=False, **o
         load_steps : bool
             Whether or not evolution steps should be automatically loaded.
 
+        RNG : numpy.random.Generator, optional
+            Random number generator used for any stochastic components of
+            the simulation. Defaults to a new NumPy Generator instance
+            created via ``np.random.default_rng()``.
+
         verbose : bool
             Print useful info.
 
+        **override_sim_kwargs
+            Additional keyword arguments that override values specified
+            in the .ini file when constructing the SimulationProperties
+            instance.
+
         Returns
         -------
         SimulationProperties
@@ -264,11 +277,12 @@ def from_ini(cls, path, metallicity = None, load_steps=False, verbose=False, **o
         if load_steps:
             # Load the steps and required data
             new_instance.load_steps(metallicity=metallicity,
+                                    RNG=RNG,
                                     verbose=verbose)
 
         return new_instance
 
-    def load_steps(self, metallicity=None, verbose=False):
+    def load_steps(self, metallicity=None, RNG=np.random.default_rng(), verbose=False):
         """Instantiate all step classes and set as instance attributes.
 
         Parameters
@@ -292,15 +306,16 @@ def load_steps(self, metallicity=None, verbose=False):
             if isinstance(tup, tuple):
                 step_kwargs = tup[1]
                 metallicity = step_kwargs.get('metallicity', metallicity)
-                self.load_a_step(name, tup, metallicity=metallicity, verbose=verbose)
+                self.load_a_step(name, tup, metallicity=metallicity, RNG=RNG, verbose=verbose)
 
         if verbose:
             if self.steps_loaded:
                 print("All steps loaded successfully.")
             else:
                 print("Not all steps were loaded successfully. Check warnings for details.")
 
-    def load_a_step(self, step_name, step_tup=(NullStep, {}), metallicity=None, from_ini='', verbose=False):
+    def load_a_step(self, step_name, step_tup=(NullStep, {}), metallicity=None,
+                    RNG=np.random.default_rng(), from_ini='', verbose=False):
         """
         Instantiate and attach a simulation step to this object.
 
@@ -366,8 +381,8 @@ def load_a_step(self, step_name, step_tup=(NullStep, {}), metallicity=None, from
             # check to make sure the step has a...
             # 1) metallicity assigned (if needed)
             # 2) TrackMatcher assigned (if needed)
-            step_tup = self.check_step(metallicity, step_name,
-                                    step_tup, verbose)
+            step_tup = self.check_step(metallicity, RNG, step_name,
+                                       step_tup, verbose)
 
         step_func, step_kwargs = step_tup
 
@@ -392,7 +407,7 @@ def load_a_step(self, step_name, step_tup=(NullStep, {}), metallicity=None, from
                                 for name, tup in self.kwargs.items()
                                 if isinstance(tup, tuple))
 
-    def check_step(self, metallicity, step_name, step_tup, verbose=False):
+    def check_step(self, metallicity, RNG, step_name, step_tup, verbose=False):
         """
         Validate and update configuration for an evolution step.
 
@@ -457,6 +472,9 @@ def check_step(self, metallicity, step_name, step_tup, verbose=False):
                 print(f"matcher_kwargs: \n" + "\n".join(kw_list))
             step_kwargs['track_matcher'] = self.track_matchers[matcher_key]
 
+        if "RNG" in step_func.DEFAULT_KWARGS:
+            step_kwargs['RNG'] = RNG
+
         return step_tup
 
     def create_track_matcher(self, metallicity, step_name, matcher_kwargs):

posydon/popsyn/binarypopulation.py

@@ -304,7 +304,7 @@ def _safe_evolve(self, **kwargs):
                     modified_tup = (step_function, step_kwargs)
                     self.population_properties.kwargs[step_name] = modified_tup
 
-            self.population_properties.load_steps()
+            self.population_properties.load_steps(RNG=self.RNG)
 
         indices = kwargs.get('indices', list(range(self.number_of_binaries)))
 

posydon/unit_tests/utils/test_common_functions.py

@@ -696,10 +696,12 @@ def test_beaming(self, binary):
             assert totest.beaming(binary) == r
 
     def test_bondi_hoyle(self, binary, monkeypatch):
-        def mock_rand(shape):
-            return np.zeros(shape)
-        def mock_rand2(shape):
-            return np.full(shape, 0.1)
+        class MockRNG:
+            def random(self, shape):
+                return np.zeros(shape)
+        class MockRNG2:
+            def random(self, shape):
+                return np.full(shape, 0.1)
         # missing argument
         with raises(TypeError, match="missing 3 required positional "\
                                      +"arguments: 'binary', 'accretor', and "\
@@ -725,32 +727,32 @@ def mock_rand2(shape):
                                              +"associated with a value"):
             # undefined scheme
             totest.bondi_hoyle(binary, binary.star_1, binary.star_2, scheme='')
-        monkeypatch.setattr(np.random, "rand", mock_rand)
-        assert totest.bondi_hoyle(binary, binary.star_1, binary.star_2) ==\
+        rng = MockRNG()
+        assert totest.bondi_hoyle(binary, binary.star_1, binary.star_2, RNG=rng) ==\
                approx(3.92668160462e-17, abs=6e-29)
         assert totest.bondi_hoyle(binary, binary.star_1, binary.star_2,\
-                                  scheme='Kudritzki+2000') ==\
+                                  RNG=rng, scheme='Kudritzki+2000') ==\
                approx(3.92668160462e-17, abs=6e-29)
         binary.star_2.log_R = 1.5          #donor's radius is 10^{1.5}Rsun
         assert totest.bondi_hoyle(binary, binary.star_1, binary.star_2,\
-                                  scheme='Kudritzki+2000') ==\
+                                  RNG=rng, scheme='Kudritzki+2000') ==\
                approx(3.92668160462e-17, abs=6e-29)
         binary.star_2.log_R = -1.5         #donor's radius is 10^{-1.5}Rsun
         assert totest.bondi_hoyle(binary, binary.star_1, binary.star_2,\
-                                  scheme='Kudritzki+2000') == 1e-99
+                                  RNG=rng, scheme='Kudritzki+2000') == 1e-99
         binary.star_2.surface_h1 = 0.25    #donor's X_surf=0.25
-        assert totest.bondi_hoyle(binary, binary.star_1, binary.star_2) ==\
+        assert totest.bondi_hoyle(binary, binary.star_1, binary.star_2, RNG=rng) ==\
                1e-99
         binary.star_2.lg_wind_mdot = -4.0  #donor's wind is 10^{-4}Msun/yr
-        assert totest.bondi_hoyle(binary, binary.star_1, binary.star_2) ==\
+        assert totest.bondi_hoyle(binary, binary.star_1, binary.star_2, RNG=rng) ==\
                1e-99
         assert totest.bondi_hoyle(binary, binary.star_1, binary.star_2,\
-                                  wind_disk_criteria=False) ==\
+                                  RNG=rng, wind_disk_criteria=False) ==\
                approx(5.34028698228e-17, abs=6e-29) # form always a disk
-        monkeypatch.setattr(np.random, "rand", mock_rand2) # other angle
+        rng = MockRNG2() # other angle
         binary.star_1.state = 'BH'         #accretor is BH
         assert totest.bondi_hoyle(binary, binary.star_1, binary.star_2,\
-                                  wind_disk_criteria=False) ==\
+                                  wind_disk_criteria=False, RNG=rng) ==\
                approx(5.13970075150e-8, abs=6e-20)
 
     def test_rejection_sampler(self, monkeypatch):

posydon/utils/common_functions.py

@@ -507,7 +507,7 @@ def beaming(binary):
 
 
 def bondi_hoyle(binary, accretor, donor, idx=-1, wind_disk_criteria=True,
-                scheme='Hurley+2002'):
+                RNG=np.random.default_rng(), scheme='Hurley+2002'):
     """Calculate the Bondi-Hoyle accretion rate of a binary [1]_.
 
     Parameters
@@ -629,7 +629,7 @@ def bondi_hoyle(binary, accretor, donor, idx=-1, wind_disk_criteria=True,
             pass
 
     n = np.sqrt((G * (m_acc + m) * Msun) / ((radius * Rsun)**3))
-    t0 = np.random.rand(len(sep)) * 2 * np.pi / n
+    t0 = RNG.random(len(sep)) * 2 * np.pi / n
     E = newton(lambda x: x - ecc * np.sin(x) - n * t0,
                np.ones_like(sep) * np.pi / 2,
                maxiter=100)