Adds 'CHE_BBH' and 'NON_CHE_BBH' to selectable DCO types for ClassCOMPAS DCOmask (see setCOMPASDCOmask() in ClassCOMPAS.py)

Author: jeffriley

postProcessing/Folders/CosmicIntegration/PythonScripts/ClassCOMPAS.py

@@ -38,6 +38,8 @@ def __init__(
         self.BBHmask = None
         self.DNSmask = None
         self.BHNSmask = None
+        self.CHE_BBHmask = None
+        self.NonCHE_BBHmask = None
         self.initialZ = None
         self.sw_weights = None
         self.n_systems = None
@@ -69,9 +71,26 @@ def setCOMPASDCOmask(
         # By default, we mask for BBHs that merge within a Hubble time, assumming
         # 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"])
 
+        if types == "CHE_BBH":
+            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_1_mask = np.logical_and(np.logical_and(stellar_type_1_zams, 16), (che_ms_1 == True))
+            che_2_mask = np.logical_and(np.logical_and(stellar_type_2_zams, 16), (che_ms_2 == True))
+            che_seeds  = sys_seeds[()][che_1_mask & che_2_mask]
+
+        if types == "NON_CHE_BBH":
+            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"])
+
+            non_che_1_mask = np.logical_or(np.logical_not(np.logical_and(stellar_type_1_zams, 16)), np.logical_and(np.logical_and(stellar_type_1_zams, 16), (che_ms_1 == False)))
+            non_che_2_mask = np.logical_or(np.logical_not(np.logical_and(stellar_type_1_zams, 16)), np.logical_and(np.logical_and(stellar_type_2_zams, 16), (che_ms_2 == False)))
+            non_che_seeds  = sys_seeds[()][non_che_1_mask & non_che_2_mask]
+
         # if user wants to mask on Hubble time use the flag, otherwise just set all to True
         hubble_mask = hubble_flag.astype(bool) if withinHubbleTime else np.repeat(True, len(dco_seeds))
 
@@ -81,6 +100,8 @@ def setCOMPASDCOmask(
             "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),
+            "CHE_BBH": np.logical_and(np.in1d(dco_seeds, che_seeds), np.logical_and(stellar_type_1 == 14, stellar_type_2 == 14)) if types == "CHE_BBH" else np.repeat(False, len(dco_seeds)),
+            "NON_CHE_BBH": np.logical_and(np.in1d(dco_seeds, non_che_seeds), np.logical_and(stellar_type_1 == 14, stellar_type_2 == 14)) if types == "NON_CHE_BBH" else np.repeat(True, len(dco_seeds))
         }
 
         # if the user wants to make RLOF or optimistic CEs
@@ -115,6 +136,8 @@ def setCOMPASDCOmask(
         self.BBHmask = type_masks["BBH"] * 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.allTypesMask = type_masks["all"] * hubble_mask * rlof_mask * pessimistic_mask
         self.optimisticmask = pessimistic_mask