Merge pull request #20 from MarinerQ/dev

remove explicit dependency on pycbc
MarinerQ · Sep 28, 2023 · d19c488 · d19c488
2 parents e46c83f + 4d699bf
commit d19c488
Show file tree

Hide file tree

Showing 2 changed files with 71 additions and 33 deletions.
diff --git a/sealgw/calculation/localization.py b/sealgw/calculation/localization.py
@@ -14,7 +14,8 @@
 from matplotlib import figure as Figure
 from matplotlib import pyplot as plt
 import time
-from pycbc.waveform import get_fd_waveform
+import lal
+import lalsimulation
 import bilby
 from gstlal import chirptime
 
@@ -161,48 +162,72 @@ def extract_info_from_xml(
 
 def calculate_template_norms(m1, m2, a1, a2, duration, psd_dict, f_final=1024):
     sigma_dict = {}
-
+    Msun = lal.MSUN_SI
+    Mpc = lal.PC_SI * 1e6
+    f_low = 20
+    f_ref = 50
     mc = (m1 * m2) ** (3 / 5) / (m1 + m2) ** (1 / 5)
-    if mc > 1.73:
-        approximant = 'IMRPhenomD'
-    else:
-        approximant = 'TaylorF2'  # TaylorF2
+    # if mc > 1.73:
+    #    approximant = 'IMRPhenomD'
+    # else:
+    #    approximant = 'TaylorF2'  # TaylorF2
     for detname, psdarray in psd_dict.items():
-        if len(sigma_dict) == 0:
+        if len(sigma_dict) == 0:  # only calculate waveform once
             delta_f = (
                 1 / duration
             )  # psd_dict[detname].index[1]-psd_dict[detname].index[0]
             f_final = min(psd_dict[detname].index[-1], f_final)
             # remove biased PSD from 1000Hz in SPIIR trigger
             if f_final > 972:
                 f_final = 972
-            hp, hc = get_fd_waveform(
-                approximant=approximant,  # SEOBNRv4_ROM TaylorF2
-                mass1=m1,
-                mass2=m2,
-                distance=1,
-                inclination=0,
-                coa_phase=0,
-                spin1x=0,
-                spin1y=0,
-                spin1z=a1,
-                spin2x=0,
-                spin2y=0,
-                spin2z=a2,
-                delta_f=delta_f,
-                f_lower=20,
-                f_final=f_final,
-            )
 
-        hp_farray = hp.sample_frequencies.numpy()
-        mask = hp_farray < f_final
+            if mc < 1.73:
+                hp_complex16series = lalsimulation.SimInspiralTaylorF2(
+                    0,
+                    delta_f,
+                    m1 * Msun,
+                    m2 * Msun,
+                    a1,
+                    a2,
+                    f_low,
+                    f_final,
+                    f_ref,
+                    Mpc,
+                    {},
+                )
+            else:
+                hp_complex16series = lalsimulation.SimIMRPhenomDGenerateFD(
+                    0,
+                    f_ref,
+                    delta_f,
+                    m1 * Msun,
+                    m2 * Msun,
+                    a1,
+                    a2,
+                    f_low,
+                    f_final,
+                    Mpc,
+                    {},
+                    2,
+                )
+
+            hp_farray = np.arange(
+                hp_complex16series.f0,
+                hp_complex16series.f0
+                + hp_complex16series.data.length * hp_complex16series.deltaF,
+                hp_complex16series.deltaF,
+            )
+            mask = hp_farray < f_final
         psd_interp = np.interp(
-            hp.sample_frequencies.numpy()[mask],
+            hp_farray[mask],
             psd_dict[detname].index,
             psd_dict[detname].values,
         )
         sigma = bilby.gw.utils.noise_weighted_inner_product(
-            hp.numpy()[mask], hp.numpy()[mask], psd_interp, 1 / delta_f
+            hp_complex16series.data.data[mask],
+            hp_complex16series.data.data[mask],
+            psd_interp,
+            1 / delta_f,
         )
         sigma = np.sqrt(np.real(sigma))
         sigma_dict[detname] = sigma

diff --git a/sealgw/simulation/generating_data.py b/sealgw/simulation/generating_data.py
@@ -8,10 +8,6 @@
     convert_to_lal_binary_neutron_star_parameters,
 )
 from lal import LIGOTimeGPS
-from pycbc.filter import matched_filter
-from pycbc.types.frequencyseries import FrequencySeries
-from pycbc.types.timeseries import TimeSeries
-from pycbc.waveform import get_fd_waveform, get_td_waveform
 import logging
 
 # fmt: off
@@ -386,6 +382,13 @@ def bns_truncated_fd_bilbypara(
     premerger_time_end,
     **kwargs,
 ):
+    if 'wf_func' in kwargs.keys():  # add in waveform_arguemnt
+        wf_func = kwargs['wf_func']
+    else:
+        from pycbc.waveform import get_fd_waveform
+
+        wf_func = get_fd_waveform
+
     mass_1, mass_2 = chirp_mass_and_mass_ratio_to_component_masses(
         chirp_mass, mass_ratio
     )
@@ -411,7 +414,7 @@ def bns_truncated_fd_bilbypara(
     # The following two return values are pycbc frequency series
     # with frequency stamp from 0 to fhigh.
     # The data between 0 and flow are padded with zeros
-    waveform_polarizations['plus'], waveform_polarizations['cross'] = get_fd_waveform(
+    waveform_polarizations['plus'], waveform_polarizations['cross'] = wf_func(
         approximant=approx,
         mass1=mass_1,
         mass2=mass_2,
@@ -448,6 +451,7 @@ def bns_truncated_fd_bilbypara(
     return waveform_polarizations
 
 
+'''
 def bns_truncated_td_bilbypara(
     tarray,
     chirp_mass,
@@ -635,6 +639,7 @@ def bns_truncated_td(
         )
 
     return waveform_polarizations
+'''
 
 
 def get_wave_gen(source_type, fmin, duration, sampling_frequency):
@@ -852,6 +857,10 @@ def snr_generator(ifos, waveform_generator, injection_parameter, flow=None, fhig
     snr_timeseries_list: a list of snr timeseries (pycbc timeseries)
     sigma_list: a list of sigmas
     """
+    from pycbc.filter import matched_filter
+    from pycbc.types.frequencyseries import FrequencySeries
+    from pycbc.types.timeseries import TimeSeries
+
     injection_parameters_copy = injection_parameter.copy()
     injection_parameters_copy["theta_jn"] = 0
     injection_parameters_copy["luminosity_distance"] = 1
@@ -927,6 +936,10 @@ def snr_generator_fd(
     snr_timeseries_list: a list of snr timeseries (pycbc timeseries)
     sigma_list: a list of sigmas
     """
+    from pycbc.filter import matched_filter
+    from pycbc.types.frequencyseries import FrequencySeries
+    from pycbc.types.timeseries import TimeSeries
+
     injection_parameters_copy = injection_parameter.copy()
     injection_parameters_copy["theta_jn"] = 0
     injection_parameters_copy["luminosity_distance"] = 1