Cosmic integration edits to compute WDWD merger rates

compas_python_utils/cosmic_integration/ClassCOMPAS.py

@@ -2,7 +2,11 @@
 import numpy as np
 import h5py as h5
 import os
-from . import totalMassEvolvedPerZ as MPZ
+import sys
+# Get the COMPAS_ROOT_DIR var, and add the cosmic_integration directory to the path
+compas_root_dir = os.getenv('COMPAS_ROOT_DIR')
+sys.path.append(os.path.join(compas_root_dir, 'compas_python_utils/cosmic_integration'))
+import totalMassEvolvedPerZ as MPZ
 
 
 class COMPASData(object):
@@ -31,12 +35,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,37 +68,43 @@ 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
+        # By default, we mask for BHBHs that merge within a Hubble time, assuming
         # the pessimistic CEE prescription (HG donors cannot survive a CEE) and
         # not allowing immediate RLOF post-CEE
 
         stellar_type_1, stellar_type_2, hubble_flag, dco_seeds = \
             self.get_COMPAS_variables("BSE_Double_Compact_Objects", ["Stellar_Type(1)", "Stellar_Type(2)", "Merges_Hubble_Time", "SEED"])
         dco_seeds = dco_seeds.flatten()
 
-        if types == "CHE_BBH" or types == "NON_CHE_BBH":
+        if types == "CHE_BHBH" or types == "NON_CHE_BHBH":
             stellar_type_1_zams, stellar_type_2_zams, che_ms_1, che_ms_2, sys_seeds = \
                 self.get_COMPAS_variables("BSE_System_Parameters", ["Stellar_Type@ZAMS(1)", "Stellar_Type@ZAMS(2)", "CH_on_MS(1)", "CH_on_MS(2)", "SEED"])
 
             che_mask  = np.logical_and.reduce((stellar_type_1_zams == 16, stellar_type_2_zams == 16, che_ms_1 == True, che_ms_2 == True))
             che_seeds = sys_seeds[()][che_mask]
 
-        self.CHE_mask = np.in1d(dco_seeds, che_seeds) if types == "CHE_BBH" or types == "NON_CHE_BBH" else np.repeat(False, len(dco_seeds))
+        self.CHE_mask = np.in1d(dco_seeds, che_seeds) if types == "CHE_BHBH" or types == "NON_CHE_BHBH" else np.repeat(False, len(dco_seeds))
 
         # if user wants to mask on Hubble time use the flag, otherwise just set all to True, use astype(bool) to set masks to bool type
         hubble_mask = hubble_flag.astype(bool) if withinHubbleTime else np.repeat(True, len(dco_seeds))
 
         # 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),
+            "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])),
+            "BHNS": np.logical_or(np.logical_and(stellar_type_1 == 13, stellar_type_2 == 14),np.logical_and(stellar_type_1 == 14, stellar_type_2 == 13)),
+            "NSWD": np.logical_or(np.logical_and(np.isin(stellar_type_1,[10,11,12]),stellar_type_2 == 13),
+                                  np.logical_and(np.isin(stellar_type_2,[10,11,12]),stellar_type_1 == 13)),
+            "WDBH": np.logical_or(np.logical_and(np.isin(stellar_type_1,[10,11,12]),stellar_type_2 == 14),
+                                  np.logical_and(np.isin(stellar_type_2,[10,11,12]),stellar_type_1 == 14)),
         }
-        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_BHBH"]     = np.logical_and(self.CHE_mask, type_masks["BHBH"]) if types == "CHE_BHBH" else np.repeat(False, len(dco_seeds))
+        type_masks["NON_CHE_BHBH"] = np.logical_and(np.logical_not(self.CHE_mask), type_masks["BHBH"]) if types == "NON_CHE_BHBH" else np.repeat(True, len(dco_seeds))
 
         # if the user wants to make RLOF or optimistic CEs
         if noRLOFafterCEE or pessimistic:
@@ -124,11 +135,14 @@ 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.NSNSmask = type_masks["NSNS"] * hubble_mask * rlof_mask * pessimistic_mask
+        self.WDWDmask = type_masks["WDWD"] * 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.WDWDmask = type_masks["NSWD"] * hubble_mask * rlof_mask * pessimistic_mask
+        self.WDWDmask = type_masks["WDBH"] * hubble_mask * rlof_mask * pessimistic_mask
+        self.CHE_BHBHmask = type_masks["CHE_BHBH"] * hubble_mask * rlof_mask * pessimistic_mask
+        self.NonCHE_BHBHmask = type_masks["NON_CHE_BHBH"] * hubble_mask * rlof_mask * pessimistic_mask
         self.allTypesMask = type_masks["all"] * hubble_mask * rlof_mask * pessimistic_mask
         self.optimisticmask = pessimistic_mask
 
@@ -158,6 +172,9 @@ def setCOMPASData(self):
 
         primary_masses, secondary_masses, formation_times, coalescence_times, dco_seeds = \
             self.get_COMPAS_variables("BSE_Double_Compact_Objects", ["Mass(1)", "Mass(2)", "Time", "Coalescence_Time", "SEED"])
+        # Raise an error if DCO table is empty
+        if len(primary_masses) == 0:
+            raise ValueError("BSE_Double_Compact_Objects is empty!")
 
         initial_seeds, initial_Z = self.get_COMPAS_variables("BSE_System_Parameters", ["SEED", "Metallicity@ZAMS(1)"])
 
@@ -173,6 +190,10 @@ def setCOMPASData(self):
         self.mass1 = primary_masses[self.DCOmask]
         self.mass2 = secondary_masses[self.DCOmask]
 
+        #Check that you have some systems of interest in your DCO table (i.e.  len(primary_masses[self.DCOmask])>0 )
+        if len(self.mass1) == 0:
+            raise ValueError("No DCOs found with the current mask. Please check your DCO table, or change your mask settings.")
+
         # Stuff of data I dont need for integral
         # but I might be to laze to read in myself
         # and often use. Might turn it of for memory efficiency

compas_python_utils/cosmic_integration/FastCosmicIntegration.py

@@ -1,18 +1,24 @@
 import numpy as np
 import h5py  as h5
 import os
+import sys
 import time
 import matplotlib.pyplot as plt
 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
 import warnings
 import astropy.units as u
 import argparse
 import importlib
-from compas_python_utils.cosmic_integration.cosmology import get_cosmology
+
+# Get the COMPAS_ROOT_DIR var, and add the cosmic_integration directory to the path
+compas_root_dir = os.getenv('COMPAS_ROOT_DIR')
+sys.path.append(os.path.join(compas_root_dir, 'compas_python_utils/cosmic_integration'))
+import ClassCOMPAS
+from cosmology import get_cosmology
+import selection_effects
+
 
 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):
     """ 
@@ -212,7 +218,7 @@ def find_formation_and_merger_rates(n_binaries, redshifts, times, time_first_SF,
             merger_rate[i, :first_too_early_index - 1] = formation_rate[i, z_of_formation_index]
     return formation_rate, merger_rate
 
-def compute_snr_and_detection_grids(sensitivity="O1", snr_threshold=8.0, Mc_max=300.0, Mc_step=0.1,
+def compute_snr_and_detection_grids(dco_type, sensitivity="O1", snr_threshold=8.0, Mc_max=300.0, Mc_step=0.1,
                                     eta_max=0.25, eta_step=0.01, snr_max=1000.0, snr_step=0.1):
     """
         Compute a grid of SNRs and detection probabilities for a range of masses and SNRs
@@ -238,6 +244,10 @@ def compute_snr_and_detection_grids(sensitivity="O1", snr_threshold=8.0, Mc_max=
             snr_grid_at_1Mpc               --> [2D float array] The snr of a binary with masses (Mc, eta) at a distance of 1 Mpc
             detection_probability_from_snr --> [list of floats] A list of detection probabilities for different SNRs
     """
+    # If DCO type includes a WD, return empty arrays since we currently only support LVK sensitivity
+    if dco_type in ["WDWD", "NSWD", "WDBH"]:
+        warnings.warn("!! Detected rate is not computed since DCO type {} doesnt work with LVK sensitivity {}".format(dco_type, sensitivity))
+
     # get interpolator given sensitivity
     interpolator = selection_effects.SNRinterpolator(sensitivity)
 
@@ -310,7 +320,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 +342,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", "NSNS", "WDWD", "BHNS", "NSWD", "WDBH"]
             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 +403,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"
@@ -469,7 +479,7 @@ def find_detection_rate(path, dco_type="BBH", merger_output_filename=None, weigh
                                                                     COMPAS.delayTimes, COMPAS.sw_weights)
 
     # create lookup tables for the SNR at 1Mpc as a function of the masses and the probability of detection as a function of SNR
-    snr_grid_at_1Mpc, detection_probability_from_snr = compute_snr_and_detection_grids(sensitivity, snr_threshold, Mc_max, Mc_step,
+    snr_grid_at_1Mpc, detection_probability_from_snr = compute_snr_and_detection_grids(dco_type, sensitivity, snr_threshold, Mc_max, Mc_step,
                                                                                     eta_max, eta_step, snr_max, snr_step)
 
     # use lookup tables to find the probability of detecting each binary at each redshift
@@ -529,6 +539,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 +562,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 +590,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 DCO systems that need a weight
             N_dco  = len(merger_rate[:,0])
 
             ####################
@@ -609,7 +620,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 +630,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 +772,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', 'NSNS', 'WDWD', 'BHNS', 'NSWD', 'WDBH'] ",
+                        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 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 +813,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 +868,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,