Cosmic integration edits to compute WDWD merger rates

compas_python_utils/cosmic_integration/ClassCOMPAS.py

@@ -2,7 +2,8 @@
 import numpy as np
 import h5py as h5
 import os
-from . import totalMassEvolvedPerZ as MPZ
+# from . import totalMassEvolvedPerZ as MPZ
+import totalMassEvolvedPerZ as MPZ
 
 
 class COMPASData(object):
@@ -31,12 +32,13 @@ def __init__(
         self.mass2 = None  # Msun
         self.DCOmask = None
         self.allTypesMask = None
-        self.BBHmask = None
-        self.DNSmask = None
+        self.BHBHmask = None
+        self.NSNSmask = None
         self.BHNSmask = None
+        self.WDWDmask = None
         self.CHE_mask = None
-        self.CHE_BBHmask = None
-        self.NonCHE_BBHmask = None
+        self.CHE_BHBHmask = None
+        self.NonCHE_BHBHmask = None
         self.initialZ = None
         self.sw_weights = None
         self.n_systems = None
@@ -63,7 +65,7 @@ def __init__(
             print("          and optionally self.setGridAndMassEvolved() if using a metallicity grid")
 
     def setCOMPASDCOmask(
-        self, types="BBH", withinHubbleTime=True, pessimistic=True, noRLOFafterCEE=True
+        self, types="BHBH", withinHubbleTime=True, pessimistic=True, noRLOFafterCEE=True
     ):
         # By default, we mask for BBHs that merge within a Hubble time, assuming
         # the pessimistic CEE prescription (HG donors cannot survive a CEE) and
@@ -88,12 +90,14 @@ def setCOMPASDCOmask(
         # mask on stellar types (where 14=BH and 13=NS), BHNS can be BHNS or NSBH
         type_masks = {
             "all": np.repeat(True, len(dco_seeds)),
-            "BBH": np.logical_and(stellar_type_1 == 14, stellar_type_2 == 14),
-            "BHNS": np.logical_or(np.logical_and(stellar_type_1 == 14, stellar_type_2 == 13), np.logical_and(stellar_type_1 == 13, stellar_type_2 == 14)),
-            "BNS": np.logical_and(stellar_type_1 == 13, stellar_type_2 == 13),
+            "BHBH": np.logical_and(stellar_type_1 == 14, stellar_type_2 == 14),
+            "BHNS": np.logical_and(np.isin(stellar_type_1,[13,14]),np.isin(stellar_type_2,[13,14])),
+            "NSNS": np.logical_and(stellar_type_1 == 13, stellar_type_2 == 13),  
+            "WDWD": np.logical_and(np.isin(stellar_type_1,[10,11,12]),np.isin(stellar_type_2,[10,11,12]))
         }
-        type_masks["CHE_BBH"]     = np.logical_and(self.CHE_mask, type_masks["BBH"]) if types == "CHE_BBH" else np.repeat(False, len(dco_seeds))
-        type_masks["NON_CHE_BBH"] = np.logical_and(np.logical_not(self.CHE_mask), type_masks["BBH"]) if types == "NON_CHE_BBH" else np.repeat(True, len(dco_seeds))
+
+        type_masks["CHE_BBH"]     = np.logical_and(self.CHE_mask, type_masks["BHBH"]) if types == "CHE_BBH" else np.repeat(False, len(dco_seeds))
+        type_masks["NON_CHE_BBH"] = np.logical_and(np.logical_not(self.CHE_mask), type_masks["BHBH"]) if types == "NON_CHE_BBH" else np.repeat(True, len(dco_seeds))
 
         # if the user wants to make RLOF or optimistic CEs
         if noRLOFafterCEE or pessimistic:
@@ -124,11 +128,12 @@ def setCOMPASDCOmask(
 
         # create a mask for each dco type supplied
         self.DCOmask = type_masks[types] * hubble_mask * rlof_mask * pessimistic_mask
-        self.BBHmask = type_masks["BBH"] * hubble_mask * rlof_mask * pessimistic_mask
+        self.BHBHmask = type_masks["BHBH"] * hubble_mask * rlof_mask * pessimistic_mask
         self.BHNSmask = type_masks["BHNS"] * hubble_mask * rlof_mask * pessimistic_mask
-        self.DNSmask = type_masks["BNS"] * hubble_mask * rlof_mask * pessimistic_mask
-        self.CHE_BBHmask = type_masks["CHE_BBH"] * hubble_mask * rlof_mask * pessimistic_mask
-        self.NonCHE_BBHmask = type_masks["NON_CHE_BBH"] * hubble_mask * rlof_mask * pessimistic_mask
+        self.NSNSmask = type_masks["NSNS"] * hubble_mask * rlof_mask * pessimistic_mask
+        self.WDWDmask = type_masks["WDWD"] * hubble_mask * rlof_mask * pessimistic_mask
+        self.CHE_BHBHmask = type_masks["CHE_BBH"] * hubble_mask * rlof_mask * pessimistic_mask
+        self.NonCHE_BHBHmask = type_masks["NON_CHE_BBH"] * hubble_mask * rlof_mask * pessimistic_mask
         self.allTypesMask = type_masks["all"] * hubble_mask * rlof_mask * pessimistic_mask
         self.optimisticmask = pessimistic_mask
 

compas_python_utils/cosmic_integration/FastCosmicIntegration.py

@@ -6,13 +6,16 @@
 import scipy
 from scipy.interpolate import interp1d
 from scipy.stats import norm as NormDist
-from compas_python_utils.cosmic_integration import ClassCOMPAS
-from compas_python_utils.cosmic_integration import selection_effects
+# from compas_python_utils.cosmic_integration import ClassCOMPAS
+import ClassCOMPAS
+# from compas_python_utils.cosmic_integration import selection_effects
+import selection_effects
 import warnings
 import astropy.units as u
 import argparse
 import importlib
-from compas_python_utils.cosmic_integration.cosmology import get_cosmology
+# from compas_python_utils.cosmic_integration.cosmology import get_cosmology
+from cosmology import get_cosmology
 
 def calculate_redshift_related_params(max_redshift=10.0, max_redshift_detection=1.0, redshift_step=0.001, z_first_SF = 10.0, cosmology=None):
     """ 
@@ -310,7 +313,7 @@ def find_detection_probability(Mc, eta, redshifts, distances, n_redshifts_detect
 
     return detection_probability
 
-def find_detection_rate(path, dco_type="BBH", merger_output_filename=None, weight_column=None,
+def find_detection_rate(path, dco_type="BHBH", merger_output_filename=None, weight_column=None,
                         merges_hubble_time=True, pessimistic_CEE=True, no_RLOF_after_CEE=True,
                         max_redshift=10.0, max_redshift_detection=1.0, redshift_step=0.001, z_first_SF = 10,
                         use_sampled_mass_ranges=True, m1_min=5 * u.Msun, m1_max=150 * u.Msun, m2_min=0.1 * u.Msun, fbin=0.7,
@@ -332,7 +335,7 @@ def find_detection_rate(path, dco_type="BBH", merger_output_filename=None, weigh
             == Arguments for finding and masking COMPAS file ==
             ===================================================
             path                   --> [string] Path to the COMPAS data file that contains the output
-            dco_type               --> [string] Which DCO type to calculate rates for: one of ["all", "BBH", "BHNS", "BNS"]
+            dco_type               --> [string] Which DCO type to calculate rates for: one of ["all", "BHBH", "BHNS", "NSNS", "WDWD"]
             merger_output_filename --> [string] Optional name of output file to store merging DCOs (do not create the extra output if None)
             weight_column          --> [string] Name of column in "DoubleCompactObjects" file that contains adaptive sampling weights
                                                     (Leave this as None if you have unweighted samples)
@@ -393,7 +396,7 @@ def find_detection_rate(path, dco_type="BBH", merger_output_filename=None, weigh
     # assert that input will not produce errors
     assert max_redshift_detection <= max_redshift, "Maximum detection redshift cannot be below maximum redshift"
     assert m1_min <= m1_max, "Minimum sampled primary mass cannot be above maximum sampled primary mass"
-    assert np.logical_and(fbin >= 0.0, fbin <= 1.0), "Binary fraction must be between 0 and 1"
+    assert fbin is None or (0.0 <= fbin <= 1.0), "Binary fraction must be between 0 and 1, or if None will vary with mass"
     assert Mc_step < Mc_max, "Chirp mass step size must be less than maximum chirp mass"
     assert eta_step < eta_max, "Symmetric mass ratio step size must be less than maximum symmetric mass ratio"
     assert snr_step < snr_max, "SNR step size must be less than maximum SNR"
@@ -529,6 +532,8 @@ def append_rates(path, detection_rate, formation_rate, merger_rate, redshifts, C
     print('shape redshifts', np.shape(redshifts))
     print('shape COMPAS.sw_weights', np.shape(COMPAS.sw_weights) )
     print('COMPAS.DCOmask', COMPAS.DCOmask, ' was set for dco_type', dco_type)
+    if dco_type=='all':
+        print('Note that rates are calculated for ALL systems in the DCO table, this could include WDWD')
     print('shape COMPAS COMPAS.DCOmask', np.shape(COMPAS.DCOmask) )
 
     #################################################
@@ -550,7 +555,6 @@ def append_rates(path, detection_rate, formation_rate, merger_rate, redshifts, C
         else:
             print(new_rate_group, 'exists, we will overrwrite the data')
 
-
         #################################################
         # Bin rates by redshifts
         #################################################
@@ -579,7 +583,7 @@ def append_rates(path, detection_rate, formation_rate, merger_rate, redshifts, C
             N_dco_in_z_bin      = (merger_rate[:,:] * fine_shell_volumes[:])
             print('fine_shell_volumes', fine_shell_volumes)
 
-            # The number of merging BBHs that need a weight
+            # The number of merging BHBHs that need a weight
             N_dco  = len(merger_rate[:,0])
 
             ####################
@@ -609,7 +613,7 @@ def append_rates(path, detection_rate, formation_rate, merger_rate, redshifts, C
             detection_index = z_index if z_index < n_redshifts_detection else n_redshifts_detection
 
             print('You will only save data up to redshift ', maxz, ', i.e. index', z_index)
-            save_redshifts        = redshifts
+            save_redshifts        = redshifts[:z_index]
             save_merger_rate      = merger_rate[:,:z_index]
             save_detection_rate   = detection_rate[:,:detection_index]
 
@@ -619,7 +623,8 @@ def append_rates(path, detection_rate, formation_rate, merger_rate, redshifts, C
         # Write the rates as a separate dataset
         # re-arrange your list of rate parameters
         DCO_to_rate_mask     = COMPAS.DCOmask #save this bool for easy conversion between BSE_Double_Compact_Objects, and CI weights
-        rate_data_list       = [DCO['SEED'][DCO_to_rate_mask], DCO_to_rate_mask , save_redshifts,  save_merger_rate, merger_rate[:,0], save_detection_rate]
+        DCO_seeds            = h_new['BSE_Double_Compact_Objects']['SEED'][DCO_to_rate_mask] # Get DCO seed
+        rate_data_list       = [DCO_seeds, DCO_to_rate_mask , save_redshifts,  save_merger_rate, merger_rate[:,0], save_detection_rate]
         rate_list_names      = ['SEED', 'DCOmask', 'redshifts',  'merger_rate','merger_rate_z0', 'detection_rate'+sensitivity]
         for i, data in enumerate(rate_data_list):
             print('Adding rate info of shape', np.shape(data))
@@ -760,20 +765,29 @@ def plot_rates(save_dir, formation_rate, merger_rate, detection_rate, redshifts,
     else:
         plt.close()
 
-
+# To allow f_binary to be None or Float
+def none_or_float(value):
+    return None if value.lower() == "none" else float(value)
 
 def parse_cli_args():
     parser = argparse.ArgumentParser()
     parser.add_argument("--path", dest='path', help="Path to the COMPAS file that contains the output", type=str,
                         default="COMPAS_Output.h5")
-    # For what DCO would you like the rate?  options: ALL, BHBH, BHNS NSNS
+
+    # For what DCO would you like the rate?  options: ALL, BHBH, BHNS NSNS, WDWD
     parser.add_argument("--dco_type", dest='dco_type',
-                        help="Which DCO type you used to calculate rates, one of: ['all', 'BBH', 'BHNS', 'BNS'] ",
-                        type=str, default="BBH")
+                        help="Which DCO type you used to calculate rates, one of: ['all', 'BHBH', 'BHNS', 'NSNS', 'WDWD'] ",
+                        type=str, default="BHBH")
     parser.add_argument("--weight", dest='weight_column',
                         help="Name of column w AIS sampling weights, i.e. 'mixture_weight'(leave as None for unweighted samples) ",
                         type=str, default=None)
-
+    parser.add_argument("--keep_pessimistic_CEE", dest='remove_pessimistic_CEE',
+                    help="keep_pessimistic_CEE will set remove_pessimistic_CEE to false. The default behaviour (remove_pessimistic_CEE == True), will mask binaries that binaries that experience a CEE while on the HG", 
+                    action='store_false', default=True)
+    parser.add_argument("--keepRLOF_postCE", dest='remove_RLOF_after_CEE',
+                        help="keepRLOF_postCE will set remove_RLOF_after_CEE to false. The default behaviour (remove_RLOF_after_CEE == True), will mask binaries that have immediate RLOF after a CCE", 
+                        action='store_false', default=True)
+
     # Options for the redshift evolution and detector sensitivity
     parser.add_argument("--maxz", dest='max_redshift', help="Maximum redshift to use in array", type=float, default=10)
     parser.add_argument("--zSF", dest='z_first_SF', help="redshift of first star formation", type=float, default=10)
@@ -792,7 +806,7 @@ def parse_cli_args():
                         default=150.)
     parser.add_argument("--m2min", dest='m2_min', help="Minimum secondary mass sampled by COMPAS", type=float,
                         default=0.1)
-    parser.add_argument("--fbin", dest='fbin', help="Binary fraction used by COMPAS", type=float, default=0.7)
+    parser.add_argument("--fbin", dest='fbin', help="Binary fraction used by COMPAS, if None f_bin will be changing with mass", type=none_or_float, default=0.7)
 
     # Parameters determining dP/dZ and SFR(z), default options from Neijssel 2019
     parser.add_argument("--mu0", dest='mu0', help="mean metallicity at redshhift 0", type=float, default=0.035)
@@ -847,6 +861,8 @@ def main():
         args.path,
         dco_type=args.dco_type,
         weight_column=args.weight_column,
+        pessimistic_CEE=args.remove_pessimistic_CEE,
+        no_RLOF_after_CEE=args.remove_RLOF_after_CEE,
         max_redshift=args.max_redshift,
         max_redshift_detection=args.max_redshift_detection,
         redshift_step=args.redshift_step,