COSMIC-PopSynth · katiebreivik · Apr 24, 2025 · Apr 24, 2025 · Apr 24, 2025
diff --git a/bin/cosmic-pop b/bin/cosmic-pop
@@ -337,7 +337,7 @@ if __name__ == '__main__':
             else:
                 raise ValueError("You must specify either qmin or m2_min in the",
                                  " inifile if you are using the independent sampler")
-            IBT, mass_singles, mass_binaries, n_singles, n_binaries = init_samp_list
+            IBT, mass_singles, mass_binaries, n_singles, n_binaries, n_singles_table = init_samp_list
 
         if sampling['sampling_method'] == 'multidim':
             init_samp_list = InitialBinaryTable.sampler(format_ = sampling['sampling_method'],
@@ -387,6 +387,22 @@ if __name__ == '__main__':
                                                       dtp=dtp,
                                                       timestep_conditions=filters['timestep_conditions'])
 
+        # extract single stars
+
+        if (args.keep_singles==True):
+            bcm_init = bcm.loc[bcm.tphys==0]
+            if sampling['sampling_method'] == 'multidim':
+                bcm_init_singles = bcm_init.iloc[-n_singles:]
+            if sampling['sampling_method'] == 'independent':
+                bcm_init_singles = bcm_init.iloc[-n_singles_table:]
+
+            bcm_singles = bcm.loc[bcm.bin_num.isin(bcm_init_singles.bin_num)]
+            bpp_singles = bpp.loc[bpp.bin_num.isin(bcm_init_singles.bin_num)]
+
+            bpp = bpp.loc[~bpp.bin_num.isin(bcm_init_singles.bin_num)]
+            bcm = bcm.loc[~bcm.bin_num.isin(bcm_init_singles.bin_num)]
+            initCond = initCond.loc[~initCond.bin_num.isin(bcm_init_singles.bin_num)]
+            kick_info = kick_info.loc[~kick_info.bin_num.isin(bcm_init_singles.bin_num)]
         # get any nans and pull them out for now
         nans = np.isnan(bpp.sep)
         if nans.any():
@@ -410,19 +426,27 @@ if __name__ == '__main__':
         # If dtp is not set, filter out first timestep in bcm
         if np.all(dtp == IBT['tphysf'].values):
             bcm = bcm.loc[bcm['tphys'].isin(dtp)]
+            if (args.keep_singles==True):
+                bcm_singles = bcm_singles.loc[bcm_singles['tphys'].isin(dtp)]
 
         # Now get the converging population and bin_nums for conv systems whic
         # satisfy the convergence_limits
         conv_filter, conv_lims_bin_num = utils.conv_select(bcm, bpp,
                                                            kstar1_range, kstar2_range,
                                                            convergence['pop_select'],
                                                            convergence['convergence_limits'])
+        if (args.keep_singles==True):
+            conv_singles_filter = utils.conv_select_singles(bcm_singles, bpp_singles, kstar1_range)
+
         if conv_filter.empty:
             warnings.warn("After filtering for desired convegence systems there were no systems matching your request. It is possible you are suggesting incompatible bin_state choices and pop_select, e.g. bin_state=[0,1], pop_select='disruption'")
             log_file.write("After filtering for desired convegence systems there were no systems matching your request. It is possible you are suggesting incompatible bin_state choices and pop_select, e.g. bin_state=[0,1], pop_select='disruption'")
 
         bcm_filter = bcm.loc[bcm.bin_num.isin(conv_filter.bin_num)]
         bpp_filter = bpp.loc[bpp.bin_num.isin(conv_filter.bin_num)]
+        if (args.keep_singles==True):
+            bpp_singles_filter = bpp_singles.loc[bpp_singles.bin_num.isin(conv_singles_filter.bin_num)]
+            bcm_singles_filter = bcm_singles.loc[bcm_singles.bin_num.isin(conv_singles_filter.bin_num)]
         initC_filter = initCond.loc[initCond.bin_num.isin(conv_filter.bin_num)]
         kick_info_filter = kick_info.loc[kick_info.bin_num.isin(conv_filter.bin_num)]
 
@@ -444,19 +468,15 @@ if __name__ == '__main__':
         conv_filter = conv_filter.loc[conv_filter.bin_num.isin(conv_lims_bin_num)]
 
         # Filter the bcm and bpp arrays according to user specified filters
-        if len(bcm_filter_match) == 0:
-            bcm_filter_match = bcm_filter.copy()
-            bpp_filter_match = bpp_filter.copy()
-            initC_filter_match = initC_filter.copy()
-            conv_filter_match = conv_filter.copy()
-            kick_info_filter_match = kick_info_filter.copy()
-        else:
-            bcm_filter_match = pd.concat([bcm_filter_match,bcm_filter], ignore_index=True)
-            bpp_filter_match = pd.concat([bpp_filter_match,bpp_filter], ignore_index=True)
-            initC_filter_match = pd.concat([initC_filter_match,initC_filter], ignore_index=True)
-            kick_info_filter_match = pd.concat([kick_info_filter_match,kick_info_filter], ignore_index=True)
-            conv_filter_match = pd.concat([conv_filter_match,conv_filter.loc[conv_filter.bin_num.isin(conv_lims_bin_num)]], ignore_index=True)
-
+        bcm_filter_match = bcm_filter.copy()
+        bpp_filter_match = bpp_filter.copy()
+        initC_filter_match = initC_filter.copy()
+        conv_filter_match = conv_filter.copy()
+        kick_info_filter_match = kick_info_filter.copy()
+        if (args.keep_singles==True):
+            conv_singles_filter_match = conv_singles_filter.copy()
+            bpp_singles_filter_match = bpp_singles_filter.copy()
+            bcm_singles_filter_match = bcm_singles_filter.copy()          
 
         if len(conv_filter_match) >= np.min([50, args.Niter]):
             conv_save = pd.concat([conv_save, pd.DataFrame(conv_filter_match)], ignore_index=True)
@@ -478,7 +498,13 @@ if __name__ == '__main__':
             # write the data and the logs!
             mass_list = [total_mass_singles, total_mass_binaries, total_mass_stars]
             n_list = [total_n_singles, total_n_binaries, total_n_stars]
-            utils.pop_write(dat_store, log_file, mass_list, n_list, bcm_filter_match,
+
+            if (args.keep_singles==True):
+                utils.pop_write(dat_store, log_file, mass_list, n_list, bcm_filter_match,
+                            bpp_filter_match, initC_filter_match, conv_filter_match, kick_info_filter_match,
+                            bin_state_nums, match_save, idx, conv_singles=conv_singles_filter_match, bcm_singles=bcm_singles_filter_match, bpp_singles=bpp_singles_filter_match)
+            else:
+                utils.pop_write(dat_store, log_file, mass_list, n_list, bcm_filter_match,
                             bpp_filter_match, initC_filter_match, conv_filter_match, kick_info_filter_match,
                             bin_state_nums, match_save, idx)
 
@@ -488,6 +514,10 @@ if __name__ == '__main__':
             initC_filter_match = []
             conv_filter_match = []
             kick_info_filter_match = []
+            if (args.keep_singles==True):
+                conv_singles_filter_match = []
+                bpp_singles_filter_match = []
+                bcm_singles_filter_match = []
             log_file.write('\n')
         Nstep += args.Nstep
         log_file.flush()

diff --git a/src/cosmic/sample/sampler/independent.py b/src/cosmic/sample/sampler/independent.py
@@ -343,7 +343,7 @@
             np.ones_like(mass1_singles) * 15,       # kstar2 (all massless remnants)
             metallicity_singles,                    # metallicity
         )
-        binary_table = pd.concat([binary_table, singles_table])
+        binary_table = pd.concat([binary_table, singles_table], ignore_index=True)
     else:
         binary_table = InitialBinaryTable.InitialBinaries(
             mass1_binary,
@@ -363,6 +363,7 @@
         m_sampled_binaries,
         n_singles,
         n_binaries,
+        len(mass1_singles)
     )
 
 

diff --git a/src/cosmic/sample/sampler/multidim.py b/src/cosmic/sample/sampler/multidim.py
@@ -25,6 +25,7 @@
 from ... import utils
 
 import numpy as np
+import pandas as pd
 
 __author__ = "Katelyn Breivik <katie.breivik@gmail.com>"
 __credits__ = "Scott Coughlin <scott.coughlin@ligo.org>"
@@ -146,6 +147,7 @@
         mass2_binary,
         porb,
         ecc,
+        single_mass_list,
         mass_singles,
         mass_binaries,
         n_singles,
@@ -187,22 +189,22 @@
             binfrac=binfrac,
         )
         tphysf, metallicity = initconditions.sample_SFH(
-            SF_start=SF_start, SF_duration=SF_duration, met=met, size=mass_singles.size
+            SF_start=SF_start, SF_duration=SF_duration, met=met, size=single_mass_list.size
         )
         metallicity[metallicity < 1e-4] = 1e-4
         metallicity[metallicity > 0.03] = 0.03
-        kstar1 = initconditions.set_kstar(mass_singles)
+        kstar1 = initconditions.set_kstar(single_mass_list)
         singles_table = InitialBinaryTable.InitialBinaries(
-            mass_singles,
-            np.ones_like(mass_singles)*0,
-            np.ones_like(mass_singles)*-1,
-            np.ones_like(mass_singles)*-1,
+            single_mass_list,
+            np.ones_like(single_mass_list)*0,
+            np.ones_like(single_mass_list)*-1,
+            np.ones_like(single_mass_list)*-1,
             tphysf,
             kstar1,
-            np.ones_like(mass_singles)*0,
+            np.ones_like(single_mass_list)*0,
             metallicity,
         )
-        binary_table = binary_table.append(singles_table)
+        binary_table = pd.concat([binary_table, singles_table], ignore_index=True)
     else:
         binary_table = InitialBinaryTable.InitialBinaries(
             mass1_binary,
@@ -321,6 +323,8 @@
             array of orbital periods in days with size=size
         ecc_list : array
             array of eccentricities with size=size
+        single_mass_list : array
+            array of mass of single stars
         mass_singles : `float`
             Total mass in single stars needed to generate population
         mass_binaries : `float`
@@ -360,7 +364,7 @@
             worker = Worker()
             results = list(pool.map(worker, inputs))
 
-        dat_lists = [[], [], [], [], [], [], [], [], []]
+        dat_lists = [[], [], [], [], [], [], [], [], [], []]
 
         for output_list in results:
             ii = 0
@@ -372,17 +376,19 @@
         secondary_mass_list = np.hstack(dat_lists[1])
         porb_list = np.hstack(dat_lists[2])
         ecc_list = np.hstack(dat_lists[3])
-        mass_singles = np.sum(dat_lists[4])
-        mass_binaries = np.sum(dat_lists[5])
-        n_singles = np.sum(dat_lists[6])
-        n_binaries = np.sum(dat_lists[7])
-        binfrac_list = np.hstack(dat_lists[8])
+        single_mass_list = np.hstack(dat_lists[4])
+        mass_singles = np.sum(dat_lists[5])
+        mass_binaries = np.sum(dat_lists[6])
+        n_singles = np.sum(dat_lists[7])
+        n_binaries = np.sum(dat_lists[8])
+        binfrac_list = np.hstack(dat_lists[9])
 
         return (
                 primary_mass_list,
                 secondary_mass_list,
                 porb_list,
                 ecc_list,
+                single_mass_list,
                 mass_singles,
                 mass_binaries,
                 n_singles,
@@ -773,6 +779,7 @@
         n_binaries = 0
         primary_mass_list = []
         secondary_mass_list = []
+        single_mass_list = []
         porb_list = []
         ecc_list = []
         binfrac_list = []
@@ -857,6 +864,7 @@
                 mass_binaries += myq * myM1
                 n_binaries += 1
             else:
+                single_mass_list.append(myM1)
                 mass_singles += myM1
                 n_singles += 1
 
@@ -865,6 +873,7 @@
                 secondary_mass_list,
                 porb_list,
                 ecc_list,
+                single_mass_list,
                 mass_singles,
                 mass_binaries,
                 n_singles,

diff --git a/src/cosmic/tests/test_sample.py b/src/cosmic/tests/test_sample.py
@@ -405,7 +405,7 @@ def test_set_kstar(self):
 
     def test_Moe_sample(self):
         # Test the multidim sampler and system-by-system binary fraction
-        m1, m2, porb, ecc, mass_singles, mass_binaries, n_singles, n_binaries, binfrac = MULTIDIMSAMPLECLASS.initial_sample(rand_seed = 2, size=10, nproc=1, mp_seeds=[0])
+        m1, m2, porb, ecc, single_mass_list, mass_singles, mass_binaries, n_singles, n_binaries, binfrac = MULTIDIMSAMPLECLASS.initial_sample(rand_seed = 2, size=10, nproc=1, mp_seeds=[0])
         self.assertEqual(np.sum(mass_singles), MOE_TOTAL_MASS)
         self.assertAlmostEqual(binfrac.max(), MULTIDIM_BINFRAC_MAX)
         self.assertAlmostEqual(binfrac.min(), MULTIDIM_BINFRAC_MIN)

diff --git a/src/cosmic/utils.py b/src/cosmic/utils.py
@@ -173,6 +173,15 @@
 
     return bcm, bin_state_fraction
 
+def conv_select_singles(bcm_save, bpp_save, final_kstar_1): # fix
+    """Select singles"""
+    conv_save = bpp_save.loc[
+            (bpp_save.kstar_1.isin(final_kstar_1))
+        ]
+    # select the formation parameters
+    conv_save = conv_save.groupby("bin_num").first().reset_index()
+    return conv_save
+
 
 def conv_select(bcm_save, bpp_save, final_kstar_1, final_kstar_2, method, conv_lims):
     """Select bcm data for special convergence cases
@@ -352,6 +361,7 @@
     bin_state_nums,
     match,
     idx,
+    **kwargs,
 ):
     """Writes all the good stuff that you want to save from runFixedPop in a
        single function
@@ -396,6 +406,15 @@
         contains the index of the bcm so we can pick up where we left off
         if runFixedPop hits a wall time
 
+    conv_singles : `pandas.DataFrame`
+        kwargs conv_singles array to write
+
+    bcm_singles : `pandas.DataFrame`
+        kwargs bcm_singles array to write
+
+    bpp_singles : `pandas.DataFrame`
+        kwargs bpp_singles array to write
+
     Returns
     -------
     Nothing!
@@ -433,6 +452,18 @@
 
     # Save the index
     dat_store.append("idx", pd.DataFrame([idx]))
+
+    if "conv_singles" in kwargs.keys():
+
+        # Save the singles dataframe
+        dat_store.append("conv_singles", kwargs["conv_singles"])
+
+        # Save the singles dataframe
+        dat_store.append("bcm_singles", kwargs["bcm_singles"])
+
+        # Save the singles dataframe
+        dat_store.append("bpp_singles", kwargs["bpp_singles"])
+
     return