ENH: Add whitened TD strain as a property of the Interferometer

bilby/gw/detector/interferometer.py

@@ -616,19 +616,95 @@ class Interferometer(object):
             power_spectral_density=self.power_spectral_density_array[self.strain_data.frequency_mask],
             duration=self.strain_data.duration)
 
+    def whiten_frequency_series(self, frequency_series : np.array) -> np.array:
+        """Whitens a frequency series with the noise properties of the detector
+
+        .. math::
+            \\tilde{a}_w(f) = \\tilde{a}(f) \\sqrt{\\frac{4}{T S_n(f)}}
+
+        Such that
+
+        .. math::
+            Var(n) = \\frac{1}{N} \\sum_k=0^N n_W(f_k)n_W^*(f_k) = 2
+
+        Where the factor of two is due to the independent real and imaginary
+        components.
+
+        Parameters
+        ==========
+        frequency_series : np.array
+            The frequency series, whitened by the ASD
+        """
+        return frequency_series / (self.amplitude_spectral_density_array * np.sqrt(self.duration / 4))
+
+    def get_whitened_time_series_from_whitened_frequency_series(
+        self,
+        whitened_frequency_series : np.array
+    ) -> np.array:
+        """Gets the whitened time series from a whitened frequency series.
+
+        This ifft's and also applies a windowing factor,
+        since when f_min and f_max are set bilby applies a mask to the series.
+
+        Per 6.2a-b in https://arxiv.org/pdf/gr-qc/0509116 since our window
+        is just a band pass,
+        this coefficient is :math:`w/W` where
+
+        .. math::
+
+            W = \\frac{1}{N} \\sum_{k=0}^N w^2[j]
+
+        Since our window :math:`w` is simply 1 or 0, depending on the mask, we get
+
+        .. math::
+
+            W = \\frac{1}{N} \\sum_{k=0}^N \\Theta(f_{max} - f_k)\\Theta(f_k - f_{min})}
+
+        and accordingly the termwise window factor is
+
+        .. math::
+            w = \\sqrt{N W} = \\sqrt{\\sum_{k=0}^N \\Theta(f_{max} - f_k)\\Theta(f_k - f_{min})}}
+
+        """
+        frequency_window_factor = (
+            np.sum(self.frequency_mask)
+            / len(self.frequency_mask)
+        )
+
+        whitened_time_series = (
+            np.fft.irfft(whitened_frequency_series)
+            * np.sqrt(np.sum(self.frequency_mask)) / frequency_window_factor
+        )
+
+        return whitened_time_series
+
     @property
     def whitened_frequency_domain_strain(self):
-        """ Calculates the whitened data by dividing the frequency domain data by
-        ((amplitude spectral density) * (duration / 4) ** 0.5). The resulting
-        data will have unit variance.
+        r"""Whitens the frequency domain data by dividing through by ASD,
+        with appropriate normalization.
+
+        See `whiten_frequency_series()` for details.
 
         Returns
         =======
         array_like: The whitened data
         """
-        return self.strain_data.frequency_domain_strain / (
-            self.amplitude_spectral_density_array * np.sqrt(self.duration / 4)
-        )
+        return self.whiten_frequency_series(self.strain_data.frequency_domain_strain)
+
+    @property
+    def whitened_time_domain_strain(self) -> np.array:
+        """Calculates the whitened time domain strain
+        by iffting the whitened frequency domain strain,
+        with the appropriate normalization.
+
+        See `get_whitened_time_series_from_whitened_frequency_series()` for details
+
+        Returns
+        =======
+        array_like
+            The whitened data in the time domain
+        """
+        return self.get_whitened_time_series_from_whitened_frequency_series(self.whitened_frequency_domain_strain)
 
     def save_data(self, outdir, label=None):
         """ Creates save files for interferometer data in plain text format.

bilby/gw/result.py

@@ -377,10 +377,6 @@ class CompactBinaryCoalescenceResult(CoreResult):
         logger.debug("Downsampling frequency mask to {} values".format(
             len(frequency_idxs))
         )
-        frequency_window_factor = (
-            np.sum(interferometer.frequency_mask)
-            / len(interferometer.frequency_mask)
-        )
         plot_times = interferometer.time_array[time_idxs]
         plot_times -= interferometer.strain_data.start_time
         start_time -= interferometer.strain_data.start_time
@@ -451,11 +447,7 @@ class CompactBinaryCoalescenceResult(CoreResult):
                 fig.add_trace(
                     go.Scatter(
                         x=plot_times,
-                        y=np.fft.irfft(
-                            interferometer.whitened_frequency_domain_strain
-                            * np.sqrt(np.sum(interferometer.frequency_mask))
-                            / frequency_window_factor
-                        )[time_idxs],
+                        y=interferometer.whitened_time_domain_strain[time_idxs],
                         fill=None,
                         mode='lines', line_color=DATA_COLOR,
                         opacity=0.5,
@@ -478,11 +470,7 @@ class CompactBinaryCoalescenceResult(CoreResult):
                     interferometer.amplitude_spectral_density_array[frequency_idxs],
                     color=DATA_COLOR, label='ASD')
                 axs[1].plot(
-                    plot_times, np.fft.irfft(
-                        interferometer.whitened_frequency_domain_strain
-                        * np.sqrt(np.sum(interferometer.frequency_mask))
-                        / frequency_window_factor
-                    )[time_idxs],
+                    plot_times, interferometer.whitened_time_domain_strain[time_idxs],
                     color=DATA_COLOR, alpha=0.3)
             logger.debug('Plotted interferometer data.')
 
@@ -493,10 +481,10 @@ class CompactBinaryCoalescenceResult(CoreResult):
             wf_pols = waveform_generator.frequency_domain_strain(params)
             fd_waveform = interferometer.get_detector_response(wf_pols, params)
             fd_waveforms.append(fd_waveform[frequency_idxs])
-            td_waveform = infft(
-                fd_waveform * np.sqrt(2. / interferometer.sampling_frequency) /
-                interferometer.amplitude_spectral_density_array,
-                self.sampling_frequency)[time_idxs]
+            whitened_fd_waveform = interferometer.whiten_frequency_series(fd_waveform)
+            td_waveform = interferometer.get_whitened_time_series_from_whitened_frequency_series(
+                whitened_fd_waveform
+            )[time_idxs]
             td_waveforms.append(td_waveform)
         fd_waveforms = asd_from_freq_series(
             fd_waveforms,

test/gw/detector/interferometer_test.py

@@ -560,21 +560,92 @@ class TestInterferometerAntennaPatternAgainstLAL(unittest.TestCase):
 
 class TestInterferometerWhitenedStrain(unittest.TestCase):
     def setUp(self):
+        self.duration = 64
+        self.sampling_frequency = 4096
         self.ifo = bilby.gw.detector.get_empty_interferometer('H1')
         self.ifo.set_strain_data_from_power_spectral_density(
-            sampling_frequency=4096, duration=64)
+            sampling_frequency=self.sampling_frequency, duration=self.duration)
+        self.waveform_generator = bilby.gw.waveform_generator.WaveformGenerator(
+            duration=self.duration,
+            sampling_frequency=self.sampling_frequency,
+            frequency_domain_source_model=bilby.gw.source.lal_binary_black_hole,
+            waveform_arguments={
+                "waveform_approximant": "IMRPhenomXP"
+            })
+
+        self.parameters = {
+            'mass_1': 10,
+            'mass_2': 10,
+            'a_1': 0,
+            'a_2': 0,
+            'tilt_1': 0,
+            'tilt_2': 0,
+            'phi_12': 0,
+            'phi_jl': 0,
+            'theta_jn': 0,
+            'luminosity_distance': 40,
+            'phase': 0,
+            'ra': 0,
+            'dec': 0,
+            'geocent_time': 62,
+            'psi': 0
+        }
 
     def tearDown(self):
         del self.ifo
+        del self.waveform_generator
+        del self.parameters
+        del self.duration
+        del self.sampling_frequency
 
-    def test_whitened_strain(self):
-        mask = self.ifo.frequency_mask
-        white = self.ifo.whitened_frequency_domain_strain[mask]
-        std_real = np.std(white.real)
-        std_imag = np.std(white.imag)
+    def _check_frequency_series_whiteness(self, frequency_series):
+        std_real = np.std(frequency_series.real)
+        std_imag = np.std(frequency_series.imag)
         self.assertAlmostEqual(std_real, 1, places=2)
         self.assertAlmostEqual(std_imag, 1, places=2)
 
+    def _check_time_series_whiteness(self, time_series):
+        std = np.std(time_series)
+        self.assertAlmostEqual(std, 1, places=2)
 
-if __name__ == "__main__":
-    unittest.main()
+    def test_frequency_domain_whitened_strain(self):
+        mask = self.ifo.frequency_mask
+        white = self.ifo.whitened_frequency_domain_strain[mask]
+        self._check_frequency_series_whiteness(white)
+
+    def test_time_domain_whitened_strain(self):
+        whitened_td = self.ifo.whitened_time_domain_strain
+        self._check_time_series_whiteness(whitened_td)
+
+    def test_frequency_domain_noise_and_signal_whitening(self):
+        # Inject some (loud) signal
+        self.ifo.inject_signal(waveform_generator=self.waveform_generator, parameters=self.parameters)
+        # Make the template separately
+        waveform_polarizations = self.waveform_generator.frequency_domain_strain(parameters=self.parameters)
+        signal_ifo = self.ifo.get_detector_response(
+            waveform_polarizations=waveform_polarizations,
+            parameters=self.parameters
+        )
+        # Whiten the template
+        whitened_signal_ifo = self.ifo.whiten_frequency_series(signal_ifo)
+        mask = self.ifo.frequency_mask
+        white = self.ifo.whitened_frequency_domain_strain[mask] - whitened_signal_ifo[mask]
+        self._check_frequency_series_whiteness(white)
+
+    def test_time_domain_noise_and_signal_whitening(self):
+        # Inject some (loud) signal
+        self.ifo.inject_signal(waveform_generator=self.waveform_generator, parameters=self.parameters)
+        # Make the template separately
+        waveform_polarizations = self.waveform_generator.frequency_domain_strain(parameters=self.parameters)
+        signal_ifo = self.ifo.get_detector_response(
+            waveform_polarizations=waveform_polarizations,
+            parameters=self.parameters
+        )
+        # Whiten the template in FD
+        whitened_signal_ifo_fd = self.ifo.whiten_frequency_series(signal_ifo)
+        # Get whitened template in TD
+        whitened_signal_ifo_td = self.ifo.get_whitened_time_series_from_whitened_frequency_series(
+            whitened_signal_ifo_fd
+        )
+        whitened_td = self.ifo.whitened_time_domain_strain - whitened_signal_ifo_td
+        self._check_time_series_whiteness(whitened_td)