set_injection_parameters.py

import sys
# cf https://docs.python.org/3/tutorial/modules.html
sys.path.append('../')
import numpy as np
from configparser import ConfigParser

import bilby.gw.conversion as conv

from Headers.PN_Coefficients import * # Headers.Constants already imported in PN_Coeff
import Library.python_utils as pu
import Library.param_utils as paru

import Library.CONST as CONST
import Library.roq_utils as roqu


LAL_PI = 3.141592653589793238462643383279502884
LAL_MSUN_SI = 1.988546954961461467461011951140572744e30
LAL_MTSUN_SI = 4.925491025543575903411922162094833998e-6

def Compute_fISCO(Rmin, m1, m2):
    """
    Compute the frequency at the Innermost Stable Circular Orbit
    """

    vISCO = np.sqrt(1.0 / Rmin)
    Mt = m1 + m2
    fmax = vISCO**3 / (PI * GEOM * Mt)      # fmax signal
    return fmax

def Compute_LAL_fISCO(m1_SI, m2_SI):
    m1 = m1_SI / LAL_MSUN_SI
    m2 = m2_SI / LAL_MSUN_SI
    m_sec = (m1 + m2) * LAL_MTSUN_SI
    piM = LAL_PI * m_sec
    vISCO = 1 / np.sqrt(6)
    fISCO = vISCO**3 / piM

    return fISCO


def ComputeChirpTime3p5PN(f_low, m1, m2):
    """
    Computes chirp times at 3.5 PN.

    Inputs:
    -------
        f_low: [float] low frequency cutoff of the detector
        sigpar: [array_like] parameters of the system
    Outputs:
    --------

    """
    Mt = m1 + m2
    Mt_sec = Mt * GEOM           # Total mass in s
    etaM = m1 * m2 / Mt**2
    etaM2 = etaM*etaM

    Tau_0 = tf1 * Mt_sec / etaM
    Tau_2 = tf2 + tf3 * etaM
    Tau_3 = tf4 * PI
    Tau_4 = tf5 + tf6 * etaM + tf7 * etaM2
    Tau_5 = PI*(tf9 + tf8*etaM)
    Tau_6 = tf11 + tf12*PI2 + tf10*(EULER + np.log(4.0)) + ( tf13 + tf14*PI2 )*etaM + tf15*etaM2 + tf16*etaM2*etaM
    Tau_6_log = tf10
    Tau_7 = PI*(tf19 + tf18*etaM + tf17*etaM2 )
    v_0 = pow(PI*Mt_sec*f_low,1.0/3.0)
    v_0_2 = v_0*v_0
    v_0_3 = v_0_2*v_0
    v_0_4 = v_0_3*v_0
    v_0_5 = v_0_4*v_0
    v_0_6 = v_0_5*v_0
    v_0_7 = v_0_6*v_0

    Tau_chirp = Tau_0*pow(v_0,-8.0)*(1.0 + Tau_2*v_0_2 + Tau_3*v_0_3 + Tau_4 *v_0_4 +  Tau_5 *v_0_5 +  (Tau_6 + Tau_6_log*np.log(v_0))*v_0_6 +  Tau_7 *v_0_7 )
    return Tau_chirp

def compute_time_from_flow_to_fhigh_3p5PN(f_low, f_high, mass_1, mass_2):
    """
    Computes the time it takes for the binary to go from f_low to f_high, using 3.5 PN approximation.
    """

    # Compute the time it takes from flow to get to merger
    tc_flow = ComputeChirpTime3p5PN(f_low, mass_1, mass_2)

    # Compute the time it takes from fhigh to get to merger
    tc_fhigh = ComputeChirpTime3p5PN(f_high, mass_1, mass_2)

    # Substracting the two gives the time it take for the binary to go from f_low to f_high

    return tc_flow - tc_fhigh

def ini_file_to_dict(inj_file_path='../examples/injection_files/GW170817.ini'):
    config = ConfigParser()
    config.read(inj_file_path)
    injection_parameters = pu.config_parser_to_dict(config)['parameters']

    injection_parameters = conv.generate_mass_parameters(injection_parameters)
    injection_parameters['reduced_mass'] = paru.component_masses_to_reduced_mass(injection_parameters['mass_1'], injection_parameters['mass_2'])

    # These two lines have been added to be able to plot the injected value on phi_c on the corner plots. Otherwise the `phi_c` and `tc_3p5PN` keys are not used in the code.
    # injection_parameters['phi_c'] = (injection_parameters['phase'] * 2) % (2*np.pi)
    # injection_parameters['tc_3p5PN'] = tc_3p5PN
    return injection_parameters

def compute_extended_analysis_dict(mass_1, mass_2, geocent_time, chirp_mass, **analysis_kwargs):

    approximant = analysis_kwargs['approximant']
    roq = analysis_kwargs['roq']
    roq_directory = analysis_kwargs['roq_b_matrix_directory']
    minimum_frequency = analysis_kwargs['minimum_frequency']
    reference_frequency = analysis_kwargs['reference_frequency']
    maximum_frequency = analysis_kwargs['maximum_frequency']

    ext_analysis_dict = {}
    ifo_chosen = analysis_kwargs['ifos'].split(',')
    if set.intersection(set(CONST.IFOS_POSSIBLE), set(ifo_chosen)) != set(ifo_chosen):
        raise ValueError("IFOs wrongly chosen: you must choose between {}. Example: '--ifos=H1,V1'. ".format(CONST.IFOS_POSSIBLE))
    else:
        ext_analysis_dict['ifos'] = ifo_chosen
    # approximant = 'TaylorF2'
    # approximant = 'SpinTaylorF2'
    # approximant = 'IMRPhenomD'
    # approximant = 'IMRPhenomPv2'
    # approximant = 'IMRPhenomPv2_NRTidal'

    # Set the duration and sampling frequency of the data segment that we're going
    # to inject the signal into. For the
    # TaylorF2 waveform, we cut the signal close to the isco frequency
    tc_3p5PN = ComputeChirpTime3p5PN(minimum_frequency, mass_1, mass_2)
    # We integrate the signal up to the frequency of the "Innermost stable circular orbit (ISCO)"
    R_isco = 6.      # Orbital separation at ISCO, in geometric units. 6M for PN ISCO; 2.8M for EOB
    m1_min = 1
    m2_min = 1
    f_ISCO = Compute_fISCO(R_isco, mass_1, mass_2)
    f_ISCO_run = Compute_fISCO(R_isco, m1_min, m2_min)
    duration = None
    t_end_minus_tc = 2
    # f_ISCO = Compute_LAL_fISCO(mass_1 * LAL_MSUN_SI, mass_2 * LAL_MSUN_SI)
    # f_ISCO_run = Compute_LAL_fISCO(m1_min * LAL_MSUN_SI, m2_min * LAL_MSUN_SI)
    if approximant == 'TaylorF2':
        f_high = f_ISCO
        f_high_run = f_ISCO_run
        # maximum_frequency_injected_waveform = 0
        # maximum_frequency_search_waveform = 0
        # # maximum_frequency_ifo = min(4096 * 4, f_high_run)
        # maximum_frequency_ifo = f_high_run
        # sampling_frequency = 2 * f_high_run
        # duration = int(tc_3p5PN) + 1 + 2
        # reference_frequency = 0
        maximum_frequency_ifo = min(maximum_frequency, f_high_run)
        maximum_frequency_injected_waveform = maximum_frequency_ifo
        maximum_frequency_search_waveform = maximum_frequency_ifo

        sampling_frequency = 2 * maximum_frequency_ifo
        # duration = 64
        duration = int(tc_3p5PN) + 1 + t_end_minus_tc
    elif approximant == 'IMRPhenomPv2' and roq:
        # cf:
        # - https://lscsoft.docs.ligo.org/lalsuite/lalsimulation/group___l_a_l_sim_i_m_r_phenom__c.html#gad3f98acfe9527259a7f73a8ef69a2f7b
        # - line 103 of https://lscsoft.docs.ligo.org/lalsuite/lalsimulation/_l_a_l_sim_i_m_r_phenom_d_8c_source.html#l00103
        params_best_basis, scale_factor = roqu.get_best_segment_params_from_chirp_mass(chirp_mass, roq_directory=roq_directory)

        ext_analysis_dict['roq'] = {}
        ext_analysis_dict['roq']['directory'] = roq_directory + str(int(params_best_basis['seglen'])) + 's/'
        ext_analysis_dict['roq']['params_path'] = ext_analysis_dict['roq']['directory'] + 'params.dat'

        rescaled_params = roqu.rescale_params(params_best_basis, scale_factor)

        # minimum_frequency = max(minimum_frequency, rescaled_params['flow'])
        # maximum_frequency_ifo = min(int(maximum_frequency), int(rescaled_params['fhigh']))
        # sampling_frequency = 2 * maximum_frequency_ifo
        # tc_3p5PN = ComputeChirpTime3p5PN(minimum_frequency, mass_1, mass_2)
        # # int() function because we need to make sure that sampling_freq * duration = integer to the 14 decimal precision
        # duration = min(int(tc_3p5PN + 2 + 1), int(rescaled_params['seglen']))
        # maximum_frequency_injected_waveform = rescaled_params['fhigh']
        # maximum_frequency_search_waveform = maximum_frequency_ifo
        # reference_frequency = reference_frequency * scale_factor


        # Settings which give me consistent results between the classic and the roq compute logL
        minimum_frequency = rescaled_params['flow']
        maximum_frequency_injected_waveform = rescaled_params['fhigh']
        # maximum_frequency_injected_waveform = min(2048, rescaled_params['fhigh'])
        maximum_frequency_ifo = min(maximum_frequency, rescaled_params['fhigh'])
        maximum_frequency_search_waveform = maximum_frequency_injected_waveform
        # maximum_frequency_search_waveform = maximum_frequency_ifo
        # maximum_frequency_ifo = rescaled_params['fhigh']
        sampling_frequency = 2 * max(maximum_frequency, rescaled_params['fhigh'])
        duration = rescaled_params['seglen']
        reference_frequency = reference_frequency * scale_factor





        ext_analysis_dict['roq']['params'] = params_best_basis
        ext_analysis_dict['roq']['rescaled_params'] = rescaled_params
        ext_analysis_dict['roq']['scale_factor'] = scale_factor

        tc_3p5PN = ComputeChirpTime3p5PN(minimum_frequency, mass_1, mass_2)

        LAL_MTSUN_SI = 4.925491025543575903411922162094833998e-6
        f_CUT = 0.2
        M_sec = (mass_1 + mass_2) * LAL_MTSUN_SI
        M_min_sec = (m1_min + m2_min) * LAL_MTSUN_SI
        f_high = f_CUT / M_sec
        f_high_run = f_CUT / M_min_sec
    else:
        # cf:
        # - https://lscsoft.docs.ligo.org/lalsuite/lalsimulation/group___l_a_l_sim_i_m_r_phenom__c.html#gad3f98acfe9527259a7f73a8ef69a2f7b
        # - line 103 of https://lscsoft.docs.ligo.org/lalsuite/lalsimulation/_l_a_l_sim_i_m_r_phenom_d_8c_source.html#l00103
        LAL_MTSUN_SI = 4.925491025543575903411922162094833998e-6
        f_CUT = 0.2
        M_sec = (mass_1 + mass_2) * LAL_MTSUN_SI
        M_min_sec = (m1_min + m2_min) * LAL_MTSUN_SI
        f_high = f_CUT / M_sec
        f_high_run = f_CUT / M_min_sec

        maximum_frequency_ifo = min(maximum_frequency, f_high_run)
        maximum_frequency_injected_waveform = maximum_frequency_ifo
        maximum_frequency_search_waveform = maximum_frequency_ifo

        sampling_frequency = 2 * maximum_frequency_ifo
        # duration = 64
        duration = int(tc_3p5PN) + 1 + t_end_minus_tc # formula for 59 sec
        # duration = next_power_of_two(tc_3p5PN + t_end_minus_tc) # formula for 64 sec


    # Use an integer because frame files have an integer start_time
    # hence when extracting a subpart of it, the index will be exact
    # start_time = geocent_time + t_end_minus_tc - duration # new formula which should be used
    start_time = geocent_time  - tc_3p5PN # formula used for reruns of chap 8
    # start_time = int(geocent_time - tc_3p5PN)
    # if duration is None:
    #     duration = int(geocent_time - start_time) + 1 + 2

    ext_analysis_dict.update(
        minimum_frequency = minimum_frequency,
        maximum_frequency_injected_waveform = maximum_frequency_injected_waveform,
        maximum_frequency_search_waveform = maximum_frequency_search_waveform,
        maximum_frequency_ifo = maximum_frequency_ifo,
        maximum_frequency_generated_waveform = None,
        sampling_frequency = sampling_frequency,
        duration = duration,
        tc_3p5PN = tc_3p5PN,
        start_time = start_time,
        f_high = f_high,
        f_high_run = f_high_run,
        approximant = approximant,
        reference_frequency = reference_frequency
    )

    return ext_analysis_dict

def next_power_of_two(val):
    return 2**int(np.log(val)/np.log(2) + 1)

def get_inj_parameters_and_analysis_dict(inj_file_path='../examples/injection_files/GW170817.ini', **analysis_kwargs):

    injection_parameters = ini_file_to_dict(inj_file_path)

    ext_analysis_dict = compute_extended_analysis_dict(injection_parameters['mass_1'], injection_parameters['mass_2'], injection_parameters['geocent_time'], injection_parameters['chirp_mass'], **analysis_kwargs)

    # These two lines have been added to be able to plot the injected value on phi_c on the corner plots. Otherwise the `phi_c` and `tc_3p5PN` keys are not used in the code.
    injection_parameters['phi_c'] = (injection_parameters['phase'] * 2) % (2*np.pi)
    injection_parameters['tc_3p5PN'] = ext_analysis_dict['tc_3p5PN']

    return injection_parameters, ext_analysis_dict


if __name__=='__main__':
    from optparse import OptionParser
    usage = """%prog [options]
    Plots the three waveforms in the time domain"""
    parser=OptionParser(usage)

    parser.add_option("--inj_file", default='../examples/injection_files/GW170817.ini', action="store", type="string", help="""Injection file path leading to the `.ini` file specifying the parameters of the injection.""")

    (opts,args)=parser.parse_args()

    injection_parameters = ini_file_to_dict(opts.inj_file)

    pu.print_dict(injection_parameters, indent=3, align_keys=True)

    breakpoint()