offner22 primary mass dependent binary fraction (#644)

MarkGM02 · katiebreivik · web-flow · commit 29c2eca3e484 · 2024-06-21T16:54:01.000-04:00
* added offner22 primary mass dependent binary fraction * pinning numpy version because of new numpy version released on june 16 (https://stackoverflow.com/questions/78634235/numpy-dtype-size-changed-may-indicate-binary-incompatibility-expected-96-from) * pinned numpy version so 3.7 works --------- Co-authored-by: katiebreivik <kbreivik@flatironinstitute.org>
diff --git a/bin/cosmic-pop b/bin/cosmic-pop
@@ -101,7 +101,7 @@ def parse_commandline():
     parser.add_argument("--binary_state", nargs='+', type=int)
     parser.add_argument("--sampling_method")
     parser.add_argument("--primary_model", help="Chooses the initial primary mass function from: salpeter55, kroupa93, kroupa01", type=str)
-    parser.add_argument("--binfrac_model", help="Chooses the binary fraction model from: a float between [0,1] and vanHaaften", type=float)
+    parser.add_argument("--binfrac_model", help="Chooses the binary fraction model from: a float between [0,1], vanHaaften, and offner22", type=float)
     parser.add_argument("--ecc_model", help="Chooses the initial eccentricity distribution model from: thermal, uniform, and sana12", type=str)
     parser.add_argument("--porb_model", help="Chooses the initial orbital period distribution model from: log_uniform and sana12", type=str)
     parser.add_argument("--SF_start", help="Sets the time in the past when star formation initiates in Myr", type=float)
diff --git a/requirements.txt b/requirements.txt
@@ -1,5 +1,5 @@
 scipy >= 0.12.1
-numpy >= 1.16
+numpy == 1.21.6
 astropy >= 1.1.1
 configparser
 tqdm >= 4.0
diff --git a/src/cosmic/sample/sampler/cmc.py b/src/cosmic/sample/sampler/cmc.py
@@ -83,7 +83,7 @@ def get_cmc_sampler(
         Sets the pairing of stars M>msort only with stars with M>msort
 
     binfrac_model : `str or float`
-        Model for binary fraction; choices include: vanHaaften or a fraction where 1.0 is 100% binaries
+        Model for binary fraction; choices include: vanHaaften, offner22, or a fraction where 1.0 is 100% binaries
 
     binfrac_model_msort : `str or float`
         Same as binfrac_model for M>msort
diff --git a/src/cosmic/sample/sampler/independent.py b/src/cosmic/sample/sampler/independent.py
@@ -114,7 +114,7 @@ def get_independent_sampler(
         Duration of constant star formation beginning from SF_Start in Myr
 
     binfrac_model : `str or float`
-        Model for binary fraction; choices include: vanHaaften or a fraction where 1.0 is 100% binaries
+        Model for binary fraction; choices include: vanHaaften, offner22, or a fraction where 1.0 is 100% binaries
 
     binfrac_model_msort : `str or float`
         Same as binfrac_model for M>msort
@@ -601,6 +601,7 @@ def binary_select(self, primary_mass, binfrac_model=0.5, **kwargs):
         either a binary fraction specified by a float or a
         primary-mass dependent binary fraction following
         `van Haaften et al.(2009) <http://adsabs.harvard.edu/abs/2013A%26A...552A..69V>`_ in appdx
+        or `Offner et al.(2022) <https://arxiv.org/abs/2203.10066>`_ in fig 1
 
         Parameters
         ----------
@@ -609,6 +610,9 @@ def binary_select(self, primary_mass, binfrac_model=0.5, **kwargs):
 
             binfrac_model : str or float
                 vanHaaften - primary mass dependent and ONLY VALID up to 100 Msun
+                
+                offner22 - primary mass dependent
+                
                 float - fraction of binaries; 0.5 means 2 in 3 stars are a binary pair while 1
                 means every star is in a binary pair
 
@@ -650,13 +654,35 @@ def binary_select(self, primary_mass, binfrac_model=0.5, **kwargs):
                 binary_choose_low = np.random.uniform(
                     0, 1.0, primary_mass[lowmassIdx].size)
 
+                (singleIdx_low,) = np.where(
+                    binary_fraction_low < binary_choose_low)
+                (binaryIdx_low,) = np.where(
+                    binary_fraction_low >= binary_choose_low)
+            elif binfrac_model == "offner22":
+                from scipy.interpolate import BSpline
+                t = [0.0331963853, 0.0331963853, 0.0331963853, 0.0331963853, 0.106066017,
+                     0.212132034, 0.424264069, 0.866025404, 1.03077641, 1.11803399,
+                     1.95959179, 3.87298335, 6.32455532, 11.6619038, 29.1547595,
+                     40, 50, 60, 70, 80, 90, 100, 110, 120, 130, 150, 150, 150, 150]
+                c = [0.08, 0.15812003, 0.20314101, 0.23842953, 0.33154153, 0.39131739,
+                     0.46020725, 0.59009569, 0.75306454, 0.81652502, 0.93518422, 0.92030594,
+                     0.96, 0.96, 0.96, 0.96, 0.96, 0.96, 0.96, 0.96, 0.96, 0.96, 0.96, 0.96, 0.96]
+                k = 3
+                def offner_curve(x):
+                    a = -0.16465041
+                    b = -0.11616329
+                    return np.piecewise(x, [x < 6.4, x >= 6.4], [BSpline(t,c,k), lambda x : a * np.exp(b * x) + 0.97])
+                binary_fraction_low = offner_curve(primary_mass[lowmassIdx])
+                binary_choose_low = np.random.uniform(
+                    0, 1.0, primary_mass[lowmassIdx].size)
+
                 (singleIdx_low,) = np.where(
                     binary_fraction_low < binary_choose_low)
                 (binaryIdx_low,) = np.where(
                     binary_fraction_low >= binary_choose_low)
             else:
                 raise ValueError(
-                    "You have supplied a non-supported binary fraction model. Please choose vanHaaften or a float"
+                    "You have supplied a non-supported binary fraction model. Please choose vanHaaften, offner22, or a float"
                 )
         elif type(binfrac_model) == float:
             if (binfrac_model <= 1.0) & (binfrac_model >= 0.0):
@@ -675,7 +701,7 @@ def binary_select(self, primary_mass, binfrac_model=0.5, **kwargs):
                 )
         else:
             raise ValueError(
-                "You have not supplied a model or a fraction. Please choose either vanHaaften or a float"
+                "You have not supplied a model or a fraction. Please choose either vanHaaften, offner22, or a float"
             )
 
         # --- if using a different binary fraction for high-mass systems
@@ -686,13 +712,35 @@ def binary_select(self, primary_mass, binfrac_model=0.5, **kwargs):
                 binary_choose_high = np.random.uniform(
                     0, 1.0, primary_mass[highmassIdx].size)
 
+                (singleIdx_high,) = np.where(
+                    binary_fraction_high < binary_choose_high)
+                (binaryIdx_high,) = np.where(
+                    binary_fraction_high >= binary_choose_high)
+            elif binfrac_model_msort == "offner22":
+                from scipy.interpolate import BSpline
+                t = [0.0331963853, 0.0331963853, 0.0331963853, 0.0331963853, 0.106066017,
+                     0.212132034, 0.424264069, 0.866025404, 1.03077641, 1.11803399,
+                     1.95959179, 3.87298335, 6.32455532, 11.6619038, 29.1547595,
+                     40, 50, 60, 70, 80, 90, 100, 110, 120, 130, 150, 150, 150, 150]
+                c = [0.08, 0.15812003, 0.20314101, 0.23842953, 0.33154153, 0.39131739,
+                     0.46020725, 0.59009569, 0.75306454, 0.81652502, 0.93518422, 0.92030594,
+                     0.96, 0.96, 0.96, 0.96, 0.96, 0.96, 0.96, 0.96, 0.96, 0.96, 0.96, 0.96, 0.96]
+                k = 3
+                def offner_curve(x):
+                    a = -0.16465041
+                    b = -0.11616329
+                    return np.piecewise(x, [x < 6.4, x >= 6.4], [BSpline(t,c,k), lambda x : a * np.exp(b * x) + 0.97])
+                binary_fraction_high = offner_curve(primary_mass[highmassIdx])
+                binary_choose_high = np.random.uniform(
+                    0, 1.0, primary_mass[highmassIdx].size)
+
                 (singleIdx_high,) = np.where(
                     binary_fraction_high < binary_choose_high)
                 (binaryIdx_high,) = np.where(
                     binary_fraction_high >= binary_choose_high)
             else:
                 raise ValueError(
-                    "You have supplied a non-supported binary fraction model. Please choose vanHaaften or a float"
+                    "You have supplied a non-supported binary fraction model. Please choose vanHaaften, offner22, or a float"
                 )
         elif (binfrac_model_msort is not None) and (type(binfrac_model_msort) == float):
             if (binfrac_model_msort <= 1.0) & (binfrac_model_msort >= 0.0):
@@ -711,7 +759,7 @@ def binary_select(self, primary_mass, binfrac_model=0.5, **kwargs):
                 )
         elif (binfrac_model_msort is not None):
             raise ValueError(
-                "You have not supplied a model or a fraction. Please choose either vanHaaften or a float"
+                "You have not supplied a model or a fraction. Please choose either vanHaaften, offner22, or a float"
             )
 
 
diff --git a/src/cosmic/tests/test_sample.py b/src/cosmic/tests/test_sample.py
@@ -42,6 +42,9 @@
 N_BINARY_SELECT = 85
 VANHAAFTEN_BINFRAC_MAX = 0.9989087986493874
 VANHAAFTEN_BINFRAC_MIN = 0.6192803136799157
+OFFNER_MASS_RANGES = [(0.075,0.15), (0.15,0.30), (0.3,0.6), (0.75,1.00), (0.85,1.25), (1.00,1.25), (1.6,2.4), (3,5), (5,8), (8,17), (17,50)]
+OFFNER_DATA = [0.19, 0.23, 0.30, 0.42, 0.47, 0.50, 0.68, 0.81, 0.89, 0.93, 0.96]
+OFFNER_ERRORS = [0.03, 0.02, 0.02, 0.03, 0.03, 0.04, 0.07, 0.06, 0.05, 0.04, 0.04]
 MULTIDIM_BINFRAC_MAX = 0.6146916774140262
 MULTIDIM_BINFRAC_MIN = 0.13786300908773025
 CONST_SFR_SUM = 460028.2453521937
@@ -207,6 +210,20 @@ def test_binary_fraction(self):
         m1_b, m1_s, binfrac, bin_index = SAMPLECLASS.binary_select(primary_mass=np.arange(1,100), binfrac_model='vanHaaften')
         self.assertEqual(binfrac.max(), VANHAAFTEN_BINFRAC_MAX)
         self.assertEqual(binfrac.min(), VANHAAFTEN_BINFRAC_MIN)
+        test_fracs = []
+        test_errs = []
+        primary_mass = np.array([float(x) for x in np.logspace(np.log10(0.08), np.log10(150), num=100000)])
+        m1_b, m1_s, binfrac, bin_index = SAMPLECLASS.binary_select(primary_mass=primary_mass, binfrac_model='offner22')
+        for i in range(len(OFFNER_MASS_RANGES)):
+            low, high = OFFNER_MASS_RANGES[i][0], OFFNER_MASS_RANGES[i][1]
+            offner_value = OFFNER_DATA[i]
+            offner_error = OFFNER_ERRORS[i]
+            bins_count = len(m1_b[(m1_b >= low) & (m1_b <= high)])
+            singles_count = len(m1_s[(m1_s >= low) & (m1_s <= high)])
+            bin_frac = bins_count / (bins_count + singles_count)
+            error = abs(offner_value - bin_frac)
+            self.assertLess(error, offner_error)
+        
 
     def test_msort(self):
         np.random.seed(2)
