improve tau(f)

MarinerQ · May 22, 2024 · 212560e · 212560e
1 parent 7906296
commit 212560e
Show file tree

Hide file tree

Showing 4 changed files with 247 additions and 52 deletions.
diff --git a/sealgw/simulation/__init__.py b/sealgw/simulation/__init__.py
@@ -1 +1 @@
-from . import generating_data, prior_fitting, sealinterferometers
+from . import generating_data, prior_fitting, sealinterferometers, antenna
diff --git a/sealgw/simulation/antenna.py b/sealgw/simulation/antenna.py
@@ -0,0 +1,103 @@
+import numpy as np
+import bilby 
+
+
+
+class GWAntennaOnCPU():
+    def __init__(self, detector_tensor, vertex):
+        self.detector_tensor = detector_tensor
+        self.vertex = vertex
+
+    def getgha(self, gpstime, ra):
+        # Greenwich hour angle of source (radians).
+        gha = np.zeros_like(gpstime) - ra
+        for i,gpst in enumerate(gpstime):
+            gha[i] += bilby.gw.utils.greenwich_mean_sidereal_time(gpst)
+        return gha
+
+    def response(self, ra, dec, psi, gpstime):
+        bs = ra.shape[0]
+
+        X = np.zeros((bs, 3))
+        Y = np.zeros((bs, 3))
+
+
+        gha = self.getgha(gpstime, ra)
+
+        cosgha = np.cos(gha)
+        singha = np.sin(gha)
+        cosdec = np.cos(dec)
+        sindec = np.sin(dec)
+        cospsi = np.cos(psi)
+        sinpsi = np.sin(psi)
+
+        X[:,0] = -cospsi * singha - sinpsi * cosgha * sindec
+        X[:,1] = -cospsi * cosgha + sinpsi * singha * sindec
+        X[:,2] =  sinpsi * cosdec
+
+        Y[:,0] =  sinpsi * singha - cospsi * cosgha * sindec
+        Y[:,1] =  sinpsi * cosgha + cospsi * singha * sindec
+        Y[:,2] =  cospsi * cosdec
+
+
+        D = self.detector_tensor
+        fp = np.einsum('ij,jk,ik->i', X, D, X) - np.einsum('ij,jk,ik->i', Y, D, Y)
+        fc = np.einsum('ij,jk,ik->i', X, D, Y) + np.einsum('ij,jk,ik->i', Y, D, X)
+
+        return (fp, fc)
+
+    def time_delay_from_geocenter(self, ra, dec, gpstime):
+        bs = ra.shape[0]
+        gha = self.getgha(gpstime, ra)
+
+        cosgha = np.cos(gha)
+        singha = np.sin(gha)
+        cosdec = np.cos(dec)
+        sindec = np.sin(dec)
+
+        wavevector = np.zeros((bs, 3))
+        wavevector[:,0],wavevector[:,1],wavevector[:,2]  = \
+            -cosgha*cosdec, cosdec*singha, -sindec
+
+        loc = self.vertex
+        dt = np.einsum('ij,j->i', wavevector, loc) / 299792458
+
+        return dt
+
+    def resp_and_dt(self, ra, dec, gpstime, psi):
+        bs = ra.shape[0]        
+        X = np.zeros((bs, 3))
+        Y = np.zeros((bs, 3))
+
+        gha = self.getgha(gpstime, ra)
+
+        cosgha = np.cos(gha)
+        singha = np.sin(gha)
+        cosdec = np.cos(dec)
+        sindec = np.sin(dec)
+        cospsi = np.cos(psi)
+        sinpsi = np.sin(psi)
+
+        X[:,0] = -cospsi * singha - sinpsi * cosgha * sindec
+        X[:,1] = -cospsi * cosgha + sinpsi * singha * sindec
+        X[:,2] =  sinpsi * cosdec
+
+        Y[:,0] =  sinpsi * singha - cospsi * cosgha * sindec
+        Y[:,1] =  sinpsi * cosgha + cospsi * singha * sindec
+        Y[:,2] =  cospsi * cosdec
+
+
+
+        wavevector = np.zeros((bs, 3))
+        wavevector[:,0],wavevector[:,1],wavevector[:,2]  = \
+            -cosgha*cosdec, cosdec*singha, -sindec
+
+
+        loc = self.vertex
+        D = self.detector_tensor
+
+        fp = np.einsum('ij,jk,ik->i', X, D, X) - np.einsum('ij,jk,ik->i', Y, D, Y)
+        fc = np.einsum('ij,jk,ik->i', X, D, Y) + np.einsum('ij,jk,ik->i', Y, D, X)
+        dt = np.einsum('ij,j->i', wavevector, loc) / 299792458
+
+        return (fp,fc,dt)
diff --git a/sealgw/simulation/generating_data.py b/sealgw/simulation/generating_data.py
@@ -9,6 +9,12 @@
 )
 from lal import LIGOTimeGPS
 import logging
+import lal 
+
+LAL_MTSUN_SI = lal.MTSUN_SI
+LAL_PI = lal.PI
+LAL_GAMMA = lal.GAMMA
+Pi_p2 = LAL_PI**2
 
 # fmt: off
 _PARAMETERS = [
@@ -305,13 +311,80 @@ def f_of_tau(tau, m1=None, m2=None, mc=None):
     return f
 
 
-def tau_of_f(f, m1=None, m2=None, mc=None):
+def tau_of_f( f, m1=None, m2=None, mc=None, chi=0):
     '''
-    Use 0PN if mc is provided. Otherwise use 3.5PN (TaylorF2).
+    Use 0PN if mc is provided. Otherwise use 3.5PN (TaylorF2), based on XLALSimInspiralTaylorF2ReducedSpinChirpTime
     '''
-    if mc is None:
-        tau = bilby.gw.utils.calculate_time_to_merger(f, m1, m2)
-    else:
+
+    if isinstance(f, list):
+        f = np.array(f)
+
+    if mc is None: # use 3.5PN
+        if isinstance(f, (float, int, np.float64)):
+            tau = bilby.gw.utils.calculate_time_to_merger(f, m1, m2, safety=1)
+        elif isinstance(f, (np.ndarray)):
+            if isinstance(m1, (float, int, np.float64)):
+                m = m1 + m2
+                eta = m1 * m2 / (m * m)
+                eta2 = eta * eta
+                chi2 = chi * chi
+                sigma0 = (-12769 * (-81. + 4. * eta)) / (16. * (-113. + 76. * eta) * (-113. + 76. * eta))
+                gamma0 = (565 * (-146597. + 135856. * eta + 17136. * eta2)) / (2268. * (-113. + 76. * eta))
+
+                v = (LAL_PI * m * LAL_MTSUN_SI * f)**(1/3)
+                tk = np.zeros((8, len(f)))  # chirp time coefficients up to 3.5 PN
+
+                # chirp time coefficients up to 3.5PN
+                tk[0] = (5. * m * LAL_MTSUN_SI) / (256. * np.power(v, 8) * eta)
+                tk[1] = 0.
+                tk[2] = 2.9484126984126986 + (11 * eta) / 3.
+                tk[3] = (-32 * LAL_PI) / 5. + (226. * chi) / 15.
+                tk[4] = 6.020630590199042 - 2 * sigma0 * chi2 + (5429 * eta) / 504. + (617 * eta2) / 72.
+                tk[5] = (3 * gamma0 * chi) / 5. - (7729 * LAL_PI) / 252. + (13 * LAL_PI * eta) / 3.
+                tk[6] = -428.291776175525 + (128 * Pi_p2) / 3. + (6848 * LAL_GAMMA) / 105. + (3147553127 * eta) / 3.048192e6 - \
+                        (451 * Pi_p2 * eta) / 12. - (15211 * eta2) / 1728. + (25565 * eta2 * eta) / 1296. + (6848 * np.log(4 * v)) / 105.
+                tk[7] = (-15419335 * LAL_PI) / 127008. - (75703 * LAL_PI * eta) / 756. + (14809 * LAL_PI * eta2) / 378.
+
+                vk = v.reshape((len(f), 1)) ** np.arange(8)  # v^k
+                tau = (1+np.sum(tk[2:,:] * vk.T[2:,:], axis=0) ) * tk[0]
+            elif isinstance(m1, np.ndarray):
+                m = m1 + m2
+                eta = m1 * m2 / (m * m)
+                eta2 = eta * eta
+                chi2 = chi * chi
+                sigma0 = (-12769 * (-81. + 4. * eta)) / (16. * (-113. + 76. * eta) * (-113. + 76. * eta))
+                gamma0 = (565 * (-146597. + 135856. * eta + 17136. * eta2)) / (2268. * (-113. + 76. * eta))
+
+                # Expand dimensions for batch processing
+                eta = eta[:, np.newaxis]
+                eta2 = eta2[:, np.newaxis]
+                chi = chi[:, np.newaxis]
+                chi2 = chi2[:, np.newaxis]
+                sigma0 = sigma0[:, np.newaxis]
+                gamma0 = gamma0[:, np.newaxis]
+                m_expanded = m[:, np.newaxis]
+                f_expanded = f[np.newaxis, :]
+
+                v = (LAL_PI * m_expanded * LAL_MTSUN_SI * f_expanded)**(1/3)
+                tk = np.zeros((len(m1), 8, len(f)))  # chirp time coefficients up to 3.5 PN for each batch
+
+                # chirp time coefficients up to 3.5PN
+                tk[:, 0, :] = (5. * m_expanded * LAL_MTSUN_SI) / (256. * np.power(v, 8) * eta)
+                tk[:, 1, :] = 0.
+                tk[:, 2, :] = 2.9484126984126986 + (11 * eta) / 3.
+                tk[:, 3, :] = (-32 * LAL_PI) / 5. + (226. * chi) / 15.
+                tk[:, 4, :] = 6.020630590199042 - 2 * sigma0 * chi2 + (5429 * eta) / 504. + (617 * eta2) / 72.
+                tk[:, 5, :] = (3 * gamma0 * chi) / 5. - (7729 * LAL_PI) / 252. + (13 * LAL_PI * eta) / 3.
+                tk[:, 6, :] = -428.291776175525 + (128 * Pi_p2) / 3. + (6848 * LAL_GAMMA) / 105. + (3147553127 * eta) / 3.048192e6 - \
+                              (451 * Pi_p2 * eta) / 12. - (15211 * eta2) / 1728. + (25565 * eta2 * eta) / 1296. + (6848 * np.log(4 * v)) / 105.
+                tk[:, 7, :] = (-15419335 * LAL_PI) / 127008. - (75703 * LAL_PI * eta) / 756. + (14809 * LAL_PI * eta2) / 378.
+
+                vk = np.power(v[:,:, np.newaxis], np.arange(8)).swapaxes(1,2)  # v^k with correct broadcasting
+                tau = (1 + np.sum(tk[:, 2:, :] * vk[:, 2:, :], axis=1)) * tk[:, 0, :]
+        else:
+            print(f, type(f))
+            raise Exception("Unsupported type of f")
+    else: # use 0PN
         tau = 2.18 * (1.21 / mc) ** (5 / 3) * (100 / f) ** (8 / 3)
     return tau
 

diff --git a/sealgw/simulation/sealinterferometers.py b/sealgw/simulation/sealinterferometers.py
@@ -25,7 +25,7 @@
 # from ..calculation.localization import lal_et_response_function, lal_ce_response_function, lal_dt_function
 from ..calculation.localization import lal_response_function, lal_dt_function
 from .generating_data import f_of_tau, tau_of_f, segmentize_tau
-
+from .antenna import GWAntennaOnCPU
 
 class SealInterferometer(Interferometer):
     """Class for the Interferometer for SealGW. ET response function aligned with LAL's."""
@@ -130,6 +130,7 @@ def __init__(
         )
         self.antenna_response_change = antenna_response_change
         self.antenna_response_change_timescale = antenna_response_change_timescale
+        self.antenna_func = GWAntennaOnCPU(self.detector_tensor, self.vertex)
 
     def antenna_response(self, ra, dec, time, psi, mode):
         """
@@ -194,7 +195,6 @@ def time_delay_from_geocenter(self, ra, dec, time):
         if self.name in ['ET1', 'ET2', 'ET3', 'CE', 'CEL', 'M1', 'M2', 'M3']:
             return lal_dt_function(ra, dec, time, self.name)
         else:
-            # print('!!!!!!')
             return time_delay_from_geocenter(self.geometry.vertex, ra, dec, time)
 
     def get_detector_response(
@@ -220,14 +220,16 @@ def get_detector_response(
         if frequencies is None:
             frequencies = self.frequency_array[self.frequency_mask]
             mask = self.frequency_mask
+            # !!Assume only low freqs are masked!!
+            # So that masked_index + masked_length = unmasked_index
+            masked_length = len(self.frequency_array) - len(frequencies)
+
         else:
             mask = np.ones(len(frequencies), dtype=bool)
+            masked_length = 0
 
-        # !!Assume only low freqs are masked!!
-        # So that masked_index + masked_length = unmasked_index
-        masked_length = len(self.frequency_array) - len(frequencies)
+
         signal = {}
-
         if self.antenna_response_change:
             try:
                 tau = tau_of_f(
@@ -237,61 +239,78 @@ def get_detector_response(
                 print("Using 0PN order to calculate time to merger.")
                 tau = tau_of_f(frequencies, mc=parameters['chirp_mass'])
             times = parameters['geocent_time'] - tau
-            segment_starts = segmentize_tau(tau, self.antenna_response_change_timescale)
 
-            antenna_response_array_dict = dict()
-            for mode in waveform_polarizations.keys():
-                antenna_response_array_dict[mode] = np.zeros(
-                    len(waveform_polarizations[mode])
-                )
+
+            # start calculating fp fc dt
+            if self.antenna_response_change_timescale==0: # use parallel calculation for all time(freq) points
+                print('use parallel calculation for all time(freq) points')
+                L = len(times)
+                ra_array = np.zeros(L) + parameters['ra']
+                dec_array = np.zeros(L)+ parameters['dec']
+                psi_array = np.zeros(L)+ parameters['psi']
+                fp, fc, dt = self.antenna_func.resp_and_dt(ra_array, dec_array, times, psi_array)
+                antenna_response_array_dict = {'plus':fp, 'cross':fc}
+                time_shift = dt
+            else: # use segment calculation. assume earth is fixed within each segment
+                print('use segment calculation. assume earth is fixed within each segment')
+                segment_starts = segmentize_tau(tau, self.antenna_response_change_timescale)
+
+                antenna_response_array_dict = dict()
+                for mode in waveform_polarizations.keys():
+                    antenna_response_array_dict[mode] = np.zeros(
+                        len(waveform_polarizations[mode])
+                    )
+                    for iseg, segment_start in enumerate(segment_starts):
+                        time = times[segment_start]
+                        if iseg < len(segment_starts) - 1:
+                            this_seg_start = segment_start + masked_length
+                            next_seg_start = segment_starts[iseg + 1] + masked_length
+                            length_to_fill = next_seg_start - this_seg_start
+                        else:
+                            this_seg_start = segment_start + masked_length
+                            next_seg_start = None
+                            length_to_fill = (
+                                len(waveform_polarizations[mode]) - this_seg_start
+                            )
+                        antenna_response_array_dict[mode][
+                            this_seg_start:next_seg_start
+                        ] = np.full(
+                            length_to_fill,
+                            self.antenna_response(
+                                parameters['ra'],
+                                parameters['dec'],
+                                time,
+                                parameters['psi'],
+                                mode,
+                            ),
+                        )
+                time_shift = np.zeros(len(frequencies))
                 for iseg, segment_start in enumerate(segment_starts):
                     time = times[segment_start]
                     if iseg < len(segment_starts) - 1:
-                        this_seg_start = segment_start + masked_length
-                        next_seg_start = segment_starts[iseg + 1] + masked_length
+                        this_seg_start = segment_start
+                        next_seg_start = segment_starts[iseg + 1]
                         length_to_fill = next_seg_start - this_seg_start
                     else:
-                        this_seg_start = segment_start + masked_length
+                        this_seg_start = segment_start
                         next_seg_start = None
-                        length_to_fill = (
-                            len(waveform_polarizations[mode]) - this_seg_start
-                        )
-                    antenna_response_array_dict[mode][
-                        this_seg_start:next_seg_start
-                    ] = np.full(
+                        length_to_fill = len(frequencies) - this_seg_start
+                    time_shift[this_seg_start:next_seg_start] = np.full(
                         length_to_fill,
-                        self.antenna_response(
-                            parameters['ra'],
-                            parameters['dec'],
-                            time,
-                            parameters['psi'],
-                            mode,
+                        self.time_delay_from_geocenter(
+                            parameters['ra'], parameters['dec'], time
                         ),
                     )
+            # end of calculating fp, fc, dt
 
+            for mode in waveform_polarizations.keys():
                 signal[mode] = (
                     waveform_polarizations[mode] * antenna_response_array_dict[mode]
                 )
 
             signal_ifo = sum(signal.values()) * mask
 
-            time_shift = np.zeros(len(frequencies))
-            for iseg, segment_start in enumerate(segment_starts):
-                time = times[segment_start]
-                if iseg < len(segment_starts) - 1:
-                    this_seg_start = segment_start
-                    next_seg_start = segment_starts[iseg + 1]
-                    length_to_fill = next_seg_start - this_seg_start
-                else:
-                    this_seg_start = segment_start
-                    next_seg_start = None
-                    length_to_fill = len(frequencies) - this_seg_start
-                time_shift[this_seg_start:next_seg_start] = np.full(
-                    length_to_fill,
-                    self.time_delay_from_geocenter(
-                        parameters['ra'], parameters['dec'], time
-                    ),
-                )
+
 
             # Be careful to first subtract the two GPS times which are ~1e9 sec.
             # And then add the time_shift which varies at ~1e-5 sec
@@ -301,7 +320,8 @@ def get_detector_response(
 
             dt = dt_geocent + time_shift
 
-        else:
+        else: #  not include earth rotation
+            print('not including the earth rotation')
             for mode in waveform_polarizations.keys():
                 det_response = self.antenna_response(
                     parameters['ra'],
@@ -316,7 +336,6 @@ def get_detector_response(
             time_shift = self.time_delay_from_geocenter(
                 parameters['ra'], parameters['dec'], parameters['geocent_time']
             )
-
             # Be careful to first subtract the two GPS times which are ~1e9 sec.
             # And then add the time_shift which varies at ~1e-5 sec
             dt_geocent = parameters['geocent_time'] - self.strain_data.start_time
Original file line number	Diff line number	Diff line change
		@@ -1 +1 @@
		from . import generating_data, prior_fitting, sealinterferometers
		from . import generating_data, prior_fitting, sealinterferometers, antenna