<jupyter_start><jupyter_text>Total DCOs in Milky Way                Part of the notebook series for Wagg et al. 2021                This notebook calculates the total number of DCOs in the Milky Way for each model variation and follows appendix B closely.<jupyter_code>import h5py as h5
import numpy as np
from scipy.integrate import quad

import astropy.units as u
import astropy.constants as c

import matplotlib.pyplot as plt
from importlib import reload

import sys
sys.path.append("../src/")
from compas_processing import mask_COMPAS_data, get_COMPAS_vars
from galaxy import simulate_mw, simulate_simple_mw
from variations import variations

%config InlineBackend.figure_format = 'retina'
plt.rc('font', family='serif')
fs = 20<jupyter_output><empty_output><jupyter_text>1 - IMF related functions 1.1 - Normalised IMF functionWe define the initial mass function as follows (using Kroupa)$$    \zeta(m) = \begin{cases}                     \beta_1 m^{-\alpha_1} & m_1 < m \le m_2 \\                    \beta_2 m^{-\alpha_2} & m_2 < m \le m_3 \\                    \beta_3 m^{-\alpha_3} & m_3 < m \le m_4 \\                    0 & \mathrm{else} \\               \end{cases},$$where we define $m_i = [0.01, 0.08, 0.5, 200] \ {\rm M_\odot}$ and $\alpha_i = [0.3, 1.3, 2.3]$. The values of each $\beta$ are defined such that the function is continuous and normalised.<jupyter_code>def IMF(m, m1=0.01, m2=0.08, m3=0.5, m4=200.0, a12=0.3, a23=1.3, a34=2.3):
    """ 
        Calculate the fraction of stellar mass between m and m + dm for a three part broken power law.
        Default values follow Kroupa (2001)
            zeta(m) ~ m^(-a_ij)
        
        Args:
            m       --> [float, list of floats] mass or masses at which to evaluate
            mi      --> [float]                 masses at which to transition the slope
            aij     --> [float]                 slope of the IMF between mi and mj
            
        Returns:
            zeta(m) --> [float, list of floats] value or values of the IMF at m
    """
    # calculate normalisation constants that ensure the IMF is continuous
    b1 = 1 / ( 
                (m2**(1 - a12) - m1**(1 - a12)) / (1 - a12) \
                + m2**(-(a12 - a23)) * (m3**(1 - a23) - m2**(1 - a23)) / (1 - a23) \
                + m2**(-(a12 - a23)) * m3**(-(a23 - a34)) * (m4**(1 - a34) - m3**(1 - a34)) / (1 - a34)
             )
    b2 = b1 * m2**(-(a12 - a23))
    b3 = b2 * m3**(-(a23 - a34))
    
    # evaluate IMF either at a point or for a list of points
    if isinstance(m, float):
        if m < m1:
            return 0
        elif m < m2:
            return b1 * m**(-a12)
        elif m < m3:
            return b2 * m**(-a23)
        elif m < m4:
            return b3 * m**(-a34)
        else:
            return 0
    else:
        imf_vals = np.zeros(len(m))
        imf_vals[np.logical_and(m >= m1, m < m2)] = b1 * m[np.logical_and(m >= m1, m < m2)]**(-a12)
        imf_vals[np.logical_and(m >= m2, m < m3)] = b2 * m[np.logical_and(m >= m2, m < m3)]**(-a23)
        imf_vals[np.logical_and(m >= m3, m < m4)] = b3 * m[np.logical_and(m >= m3, m < m4)]**(-a34)
        return imf_vals<jupyter_output><empty_output><jupyter_text>1.2 - Cumulative density function of the IMFThe CDF is defined as $F(m) = \int_{-\infty}^{m} \zeta(m) \ \mathrm{d}m$ and we so we take the same function as above but integrate it.<jupyter_code>def CDF_IMF(m, m1=0.01, m2=0.08, m3=0.5, m4=200.0, a12=0.3, a23=1.3, a34=2.3):
    """ 
        Calculate the fraction of stellar mass between 0 and m for a three part broken power law.
        Default values follow Kroupa (2001)
            F(m) ~ int_0^m zeta(m) dm
        
        Args:
            m       --> [float, list of floats] mass or masses at which to evaluate
            mi      --> [float]                 masses at which to transition the slope
            aij     --> [float]                 slope of the IMF between mi and mj
            
        Returns:
            zeta(m) --> [float, list of floats] value or values of the IMF at m

        NOTE: this is implemented recursively, probably not the most efficient if you're using this intensively but I'm not so I'm being lazy ¯\_(ツ)_/¯ 
    """

    # calculate normalisation constants that ensure the IMF is continuous
    b1 = 1 / ( 
                (m2**(1 - a12) - m1**(1 - a12)) / (1 - a12) \
                + m2**(-(a12 - a23)) * (m3**(1 - a23) - m2**(1 - a23)) / (1 - a23) \
                + m2**(-(a12 - a23)) * m3**(-(a23 - a34)) * (m4**(1 - a34) - m3**(1 - a34)) / (1 - a34)
             )
    b2 = b1 * m2**(-(a12 - a23))
    b3 = b2 * m3**(-(a23 - a34))
    
    if isinstance(m, float):
        if m <= m1:
            return 0
        elif m <= m2:
            return b1 / (1 - a12) * (m**(1 - a12) - m1**(1 - a12))
        elif m <= m3:
            return CDF_IMF(m2) + b2 / (1 - a23) * (m**(1 - a23) - m2**(1 - a23))
        elif m <= m4:
            return CDF_IMF(m3) + b3 / (1 - a34) * (m**(1 - a34) - m3**(1 - a34))
        else:
            return 0
    else:
        CDF = np.zeros(len(m))
        CDF[np.logical_and(m >= m1, m < m2)] = b1 / (1 - a12) * (m[np.logical_and(m >= m1, m < m2)]**(1 - a12) - m1**(1 - a12))
        CDF[np.logical_and(m >= m2, m < m3)] = CDF_IMF(m2) + b2 / (1 - a23) * (m[np.logical_and(m >= m2, m < m3)]**(1 - a23) - m2**(1 - a23))
        CDF[np.logical_and(m >= m3, m < m4)] = CDF_IMF(m3) + b3 / (1 - a34) * (m[np.logical_and(m >= m3, m < m4)]**(1 - a34) - m3**(1 - a34))
        CDF[m >= m4] = np.ones(len(m[m >= m4]))
        return CDF<jupyter_output><empty_output><jupyter_text>1.3 - Inverse CDF of IMF (for sampling)<jupyter_code>def inverse_CDF_IMF(U, m1=0.01, m2=0.08, m3=0.5, m4=200, a12=0.3, a23=1.3, a34=2.3):
    """ 
        Calculate the inverse CDF for a three part broken power law.
        Default values follow Kroupa (2001)
        
        Args:
            U       --> [float, list of floats] A fraction between 
            mi      --> [float]                 masses at which to transition the slope
            aij     --> [float]                 slope of the IMF between mi and mj
            
        Returns:
            zeta(m) --> [float, list of floats] value or values of the IMF at m

        NOTE: this is implemented recursively, probably not the most efficient if you're using this intensively but I'm not so I'm being lazy ¯\_(ツ)_/¯ 
    """

    # calculate normalisation constants that ensure the IMF is continuous
    b1 = 1 / ( 
                (m2**(1 - a12) - m1**(1 - a12)) / (1 - a12) \
                + m2**(-(a12 - a23)) * (m3**(1 - a23) - m2**(1 - a23)) / (1 - a23) \
                + m2**(-(a12 - a23)) * m3**(-(a23 - a34)) * (m4**(1 - a34) - m3**(1 - a34)) / (1 - a34)
             )
    b2 = b1 * m2**(-(a12 - a23))
    b3 = b2 * m3**(-(a23 - a34))

    # find the probabilities at which the gradient changes
    F1, F2, F3, F4 = CDF_IMF(np.array([m1, m2, m3, m4]), m1=0.01, m2=0.08, m3=0.5, m4=200, a12=0.3, a23=1.3, a34=2.3)

    masses = np.zeros(len(U))
    masses[np.logical_and(U > F1, U <= F2)] = np.power((1 - a12) / b1 * (U[np.logical_and(U > F1, U <= F2)] - F1) + m1**(1 - a12), 1 / (1 - a12))
    masses[np.logical_and(U > F2, U <= F3)] = np.power((1 - a23) / b2 * (U[np.logical_and(U > F2, U <= F3)] - F2) + m2**(1 - a23), 1 / (1 - a23))
    masses[np.logical_and(U > F3, U <= F4)] = np.power((1 - a34) / b3 * (U[np.logical_and(U > F3, U <= F4)] - F3) + m3**(1 - a34), 1 / (1 - a34))
    return masses<jupyter_output><empty_output><jupyter_text>2 - Main normalisation functions 2.1 - Metallicity weights<jupyter_code>def find_metallicity_weights(compas_grid, simple_mw=False):
    # find the bins that have the COMPAS grid at its centre
    inner_bins = np.array([compas_grid[i] + (compas_grid[i+1] - compas_grid[i]) / 2 for i in range(len(compas_grid) - 1)])
    bins = np.concatenate(([compas_grid[0]], inner_bins, [compas_grid[-1]]))
    
    # sample from the chosen MW metallicity distribution
    if simple_mw:
        SAMPLE_SIZE = 2000000
        _, _, sampled_Z = simulate_mw(SAMPLE_SIZE)
    else:
        SAMPLE_SIZE = 20000000
        _, _, sampled_Z = simulate_mw(SAMPLE_SIZE)
    
    # adjust the sample so everything falls inside the compas grid
    sampled_Z = sampled_Z.value
    sampled_Z[sampled_Z > np.max(compas_grid)] = np.max(compas_grid)
    sampled_Z[sampled_Z < np.min(compas_grid)] = np.min(compas_grid)
    
    # create a histogram on the grid and divide by the number of samples to find the weights
    h, _ = np.histogram(sampled_Z, bins=bins)
    w_Z = h / SAMPLE_SIZE
    return w_Z<jupyter_output><empty_output><jupyter_text>2.2 - Find fraction of Universe simulated and average COMPAS mass<jupyter_code>def create_sample_universe(m1_min, m1_max, m2_min, fbin, SAMPLES=20000000):
    # randomly sample a large number of masses, binaries and mass ratios
    primary_mass = inverse_CDF_IMF(np.random.rand(SAMPLES))
    binary = np.random.rand(SAMPLES)
    mass_ratio = np.random.rand(SAMPLES)

    # only fbin fraction of stars have a secondary (in a binary). Assign each a random secondary mass using uniform mass ratio
    secondary_mass = np.zeros(SAMPLES)
    binary_mask = binary < fbin
    secondary_mass[binary_mask] = primary_mass[binary_mask] * mass_ratio[binary_mask]

    # find the total mass of the whole population
    total_mass = np.sum(primary_mass) + np.sum(secondary_mass)

    # apply the COMPAS cuts on primary and secondary mass
    primary_mask = np.logical_and(primary_mass >= m1_min, primary_mass <= m1_max)
    secondary_mask = secondary_mass > m2_min
    full_mask = np.logical_and(primary_mask, secondary_mask)
    total_mass_COMPAS = np.sum(primary_mass[full_mask]) + np.sum(secondary_mass[full_mask])

    fraction = total_mass_COMPAS / total_mass
    average_mass_COMPAS = total_mass_COMPAS / len(primary_mass[full_mask])
    
    return fraction, average_mass_COMPAS<jupyter_output><empty_output><jupyter_text>2.3 - Get star forming mass per Z for each simulation<jupyter_code>def star_forming_mass_per_Z(average_mass_COMPAS, file_path, compas_grid, binary_type, pessimistic=True, hubble_time=True, caseBB_survive=False):
    
    if caseBB_survive and pessimistic:
        pessimistic = False
    
    # open COMPAS file
    with h5.File(file_path, "r") as f:
        # get a mask for the DCOs that you want
        DCO_mask = mask_COMPAS_data(f, binary_type, (hubble_time, True, pessimistic))
        DCO_weights, DCO_Z, DCO_seed = get_COMPAS_vars(f, "doubleCompactObjects", ["weight", "Metallicity1", "seed"], DCO_mask)
        compas_Z_unique = np.unique(DCO_Z)
        
        if caseBB_survive:
            
            to_exclude = np.repeat(False, DCO_weights.shape)
        
            # do it separately by metallicity to ensure seeds are unique
            for Z in compas_Z_unique:                
                ce_path = "/n/holystore01/LABS/berger_lab/Lab/fbroekgaarden/DATA/all_dco_legacy_CEbug_fix/unstableCaseBB/Z_{}/STROOPWAFELcombined/COMPASOutput.h5".format(Z)

                with h5.File(ce_path, "r") as ce_file:
                    ce_seeds, ce_st1, ce_st2 = get_COMPAS_vars(ce_file, "commonEnvelopes", ["randomSeed", "stellarType1", "stellarType2"])                

                # make a mask that just excludes HG but not HeHG
                ce_with_HG = np.logical_or(ce_st1 == 2, ce_st2 == 2)

                # get the corresponding seeds (unique in case there were multiple CE events)
                seeds_to_delete = np.unique(ce_seeds[ce_with_HG])

                # add to the mask
                dco_matching_Z = DCO_Z == Z
                to_exclude[dco_matching_Z] = np.isin(DCO_seed[dco_matching_Z], seeds_to_delete)
                
            # apply the mask
            DCO_weights, DCO_Z = DCO_weights[to_exclude], DCO_Z[to_exclude]

        # sum the weights of the DCOs for each metallicity
        total_BHNS_per_Z = np.array([np.sum(DCO_weights[DCO_Z == Z]) for Z in compas_grid])
        
        all_weights, all_Z = get_COMPAS_vars(f, "systems", ["weight", "Metallicity1"])
        
        # sum the weights of the binaries for each metallicity and multiply by average mass
        MSF_per_Z_COMPAS = np.array([np.sum(all_weights[all_Z == Z]) * average_mass_COMPAS for Z in compas_grid])
        
    return MSF_per_Z_COMPAS, total_BHNS_per_Z<jupyter_output><empty_output><jupyter_text>2.4 - Put it all together to get the total DCOs!<jupyter_code>def get_total_DCO_in_MW(model, binary_type="BHNS", w_Z=None, f_mass_sampled=None, average_mass_COMPAS=None, compas_grid=None, hubble_time=True):
    # find the weight for each metallicity distribution using Frankel Model
    if compas_grid is None:
        compas_grid = np.concatenate((np.round(np.logspace(np.log10(0.0001), np.log10(0.022), 50), 5), [0.0244, 0.02705, 0.03]))
    if w_Z is None:
        w_Z = find_metallicity_weights(compas_grid)
    
    # sample from mass distributions for fraction and average mass
    if f_mass_sampled is None or average_mass_COMPAS is None:
        f_mass_sampled, average_mass_COMPAS = create_sample_universe(5, 150, 0.1, 0.5)
    
    # create a filepath to the particular model
    pessimistic = True
    caseBB_survive = False
    if model == "optimistic":
        model = "fiducial"
        pessimistic = False
    elif model == "unstableCaseBB":
        pessimistic = False
        caseBB_survive = True
    elif model == "unstableCaseBB_opt":
        model = "unstableCaseBB"
        pessimistic = False
        
    file_path = "/n/holystore01/LABS/berger_lab/Lab/fbroekgaarden/DATA/all_dco_legacy_CEbug_fix/{}/COMPASOutputCombined.h5".format(model)
    
    # find the mass and total DCOs
    MSF_per_Z_COMPAS, total_DCO_per_Z = star_forming_mass_per_Z(average_mass_COMPAS, file_path, compas_grid, binary_type,
                                                                pessimistic=pessimistic, hubble_time=hubble_time, caseBB_survive=caseBB_survive)
    MSF_per_Z = MSF_per_Z_COMPAS / f_mass_sampled
    
    M_MW = 10.4e10
    
    total_DCO_in_MW = np.sum(total_DCO_per_Z / MSF_per_Z * w_Z) * M_MW
    
    return total_DCO_in_MW<jupyter_output><empty_output><jupyter_text>3. Actual calculations/execution<jupyter_code>file_path = "/n/holystore01/LABS/berger_lab/Lab/fbroekgaarden/DATA/all_dco_legacy_CEbug_fix/fiducial/COMPASOutputCombined.h5"
with h5.File(file_path, "r") as test:
    Z = get_COMPAS_vars(test, "doubleCompactObjects", "Metallicity1")
    compas_grid = np.unique(Z)
w_Z = find_metallicity_weights(compas_grid)
f_mass_sampled, average_mass_COMPAS = create_sample_universe(5, 150, 0.1, 0.5)
btypes = ["BHBH", "BHNS", "NSNS"]
models = [v["file"] for v in variations]

totals = np.zeros(shape=(len(btypes), len(models)))

for i in range(len(btypes)):
    for j in range(len(models)):
        totals[i][j] = get_total_DCO_in_MW(model=models[j], binary_type=btypes[i], w_Z=w_Z,
                                           f_mass_sampled=f_mass_sampled,
                                           average_mass_COMPAS=average_mass_COMPAS, compas_grid=compas_grid)
        print("DCO: {}, Model: {}, total: {}".format(btypes[i], models[j], totals[i][j]))
np.save("../data/total_DCO_in_MW.npy", totals)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.12642473470457224, 'fraction_numerical': 0.043495349141883924, 'mean_word_length': 2.43387314439946, 'pattern_counts': {'":': 3, '<': 69, '<?xml version=': 0, '>': 72, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 2, 'word_list_counts': {'cursed_substrings.json': 9, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '9', 'dir': 'jupyter-structured-clean-dedup', 'id': '639553', 'license': 'MIT', 'n_tokens_mistral': 5918, 'n_tokens_neox': 5488, 'n_words': 1404, 'path': 'simulation/postprocessing_notebooks/get_total_DCOs_in_MW.ipynb', 'repo_name': 'TomWagg/detecting-DCOs-in-LISA', 'size': '15266'}