FEAT: add option to crop whitened time series

bilby/core/series.py

@@ -128,3 +128,7 @@ def start_time(self):
     def start_time(self, start_time):
         self._start_time = start_time
         self._time_array_updated = False
+
+    @property
+    def end_time(self):
+        return self.start_time + self.duration

bilby/gw/detector/interferometer.py

@@ -44,6 +44,10 @@ class Interferometer(object):
     minimum_frequency = PropertyAccessor('strain_data', 'minimum_frequency')
     maximum_frequency = PropertyAccessor('strain_data', 'maximum_frequency')
     frequency_mask = PropertyAccessor('strain_data', 'frequency_mask')
+    time_mask = PropertyAccessor('strain_data', 'time_mask')
+    crop_duration = PropertyAccessor('strain_data', 'crop_duration')
+    cropped_duration = PropertyAccessor('strain_data', 'cropped_duration')
+    cropped_frequency_mask = PropertyAccessor('strain_data', 'cropped_frequency_mask')
     frequency_domain_strain = PropertyAccessor('strain_data', 'frequency_domain_strain')
     time_domain_strain = PropertyAccessor('strain_data', 'time_domain_strain')
 
@@ -614,12 +618,9 @@ def optimal_snr_squared(self, signal):
 
         Returns
         =======
-        float: The optimal signal to noise ratio possible squared
+        float: The optimal signal-to-noise ratio squared of the signal
         """
-        return gwutils.optimal_snr_squared(
-            signal=signal[self.strain_data.frequency_mask],
-            power_spectral_density=self.power_spectral_density_array[self.strain_data.frequency_mask],
-            duration=self.strain_data.duration)
+        return (abs(self.whiten_frequency_series(signal))**2).sum()
 
     def inner_product(self, signal):
         """
@@ -631,13 +632,13 @@ def inner_product(self, signal):
 
         Returns
         =======
-        float: The optimal signal to noise ratio possible squared
+        float:
+            The noise-weighted inner product between the passed signal
+            and the data stored in the :code:`Interferometer`.
         """
-        return gwutils.noise_weighted_inner_product(
-            aa=signal[self.strain_data.frequency_mask],
-            bb=self.strain_data.frequency_domain_strain[self.strain_data.frequency_mask],
-            power_spectral_density=self.power_spectral_density_array[self.strain_data.frequency_mask],
-            duration=self.strain_data.duration)
+        whitened_signal = self.whiten_frequency_series(signal)
+        whitened_data = self.whitened_frequency_domain_strain
+        return (whitened_signal.T * whitened_data.conj()).sum()
 
     def template_template_inner_product(self, signal_1, signal_2):
         """A noise weighted inner product between two templates, using this ifo's PSD.
@@ -653,11 +654,9 @@ def template_template_inner_product(self, signal_1, signal_2):
         =======
         float: The noise weighted inner product of the two templates
         """
-        return gwutils.noise_weighted_inner_product(
-            aa=signal_1[self.strain_data.frequency_mask],
-            bb=signal_2[self.strain_data.frequency_mask],
-            power_spectral_density=self.power_spectral_density_array[self.strain_data.frequency_mask],
-            duration=self.strain_data.duration)
+        whitened_1 = self.whiten_frequency_series(signal_1)
+        whitened_2 = self.whiten_frequency_series(signal_2)
+        return (whitened_1 * whitened_2.conj()).sum()
 
     def matched_filter_snr(self, signal):
         """
@@ -672,13 +671,9 @@ def matched_filter_snr(self, signal):
         complex: The matched filter signal to noise ratio
 
         """
-        return gwutils.matched_filter_snr(
-            signal=signal[self.strain_data.frequency_mask],
-            frequency_domain_strain=self.strain_data.frequency_domain_strain[self.strain_data.frequency_mask],
-            power_spectral_density=self.power_spectral_density_array[self.strain_data.frequency_mask],
-            duration=self.strain_data.duration)
+        return self.inner_product(signal) / self.optimal_snr_squared(signal)**0.5
 
-    def whiten_frequency_series(self, frequency_series : np.array) -> np.array:
+    def whiten_frequency_series(self, frequency_series: np.array) -> np.array:
         """Whitens a frequency series with the noise properties of the detector
 
         .. math::
@@ -697,7 +692,21 @@ def whiten_frequency_series(self, frequency_series : np.array) -> np.array:
         frequency_series : np.array
             The frequency series, whitened by the ASD
         """
-        return frequency_series / (self.amplitude_spectral_density_array * (self.duration / 4)**0.5)
+        if self.crop_duration == 0:
+            return gwutils.frequency_domain_whiten(
+                frequency_series=frequency_series,
+                amplitude_spectral_density=self.amplitude_spectral_density_array,
+                frequency_mask=self.frequency_mask,
+                duration=self.duration,
+            )
+        else:
+            return gwutils.whiten_and_crop(
+                frequency_series=frequency_series,
+                amplitude_spectral_density=self.amplitude_spectral_density_array,
+                frequency_mask=self.frequency_mask,
+                time_mask=self.time_mask,
+                duration=self.duration,
+            )
 
     def get_whitened_time_series_from_whitened_frequency_series(
         self,
@@ -734,7 +743,9 @@ def get_whitened_time_series_from_whitened_frequency_series(
 
         whitened_time_series = (
             xp.fft.irfft(whitened_frequency_series)
-            * self.frequency_mask.sum()**0.5 / frequency_window_factor
+            * self.frequency_mask.sum()**0.5
+            / frequency_window_factor
+            * self.time_mask.mean()**0.5
         )
 
         return whitened_time_series

bilby/gw/detector/strain_data.py

@@ -5,6 +5,7 @@
 from ...core import utils
 from ...core.series import CoupledTimeAndFrequencySeries
 from ...core.utils import logger, PropertyAccessor
+from ...core.utils.series import create_frequency_series
 from .. import utils as gwutils
 
 
@@ -14,11 +15,12 @@ class InterferometerStrainData(object):
     duration = PropertyAccessor('_times_and_frequencies', 'duration')
     sampling_frequency = PropertyAccessor('_times_and_frequencies', 'sampling_frequency')
     start_time = PropertyAccessor('_times_and_frequencies', 'start_time')
+    end_time = PropertyAccessor('_times_and_frequencies', 'end_time')
     frequency_array = PropertyAccessor('_times_and_frequencies', 'frequency_array')
     time_array = PropertyAccessor('_times_and_frequencies', 'time_array')
 
     def __init__(self, minimum_frequency=0, maximum_frequency=np.inf,
-                 roll_off=0.2, notch_list=None):
+                 roll_off=0.2, notch_list=None, crop_duration=0):
         """ Initiate an InterferometerStrainData object
 
         The initialised object contains no data, this should be added using one
@@ -35,6 +37,11 @@ def __init__(self, minimum_frequency=0, maximum_frequency=np.inf,
             This corresponds to alpha * duration / 2 for scipy tukey window.
         notch_list: bilby.gw.detector.strain_data.NotchList
             A list of notches
+        crop_duration: float | tuple
+            The duration of data to crop at the beginning/end of the segment
+            to avoid whitening artifacts. If a float, that duration is excluded
+            at each end, if a tuple, this specifies the truncation duration
+            at the beginning and end.
 
         """
 
@@ -43,11 +50,14 @@ def __init__(self, minimum_frequency=0, maximum_frequency=np.inf,
         self.notch_list = notch_list
         self.roll_off = roll_off
         self.window_factor = 1
+        self._crop_duration = crop_duration
 
         self._times_and_frequencies = CoupledTimeAndFrequencySeries()
 
         self._frequency_mask_updated = False
         self._frequency_mask = None
+        self._time_mask_updated = False
+        self._time_mask = None
         self._frequency_domain_strain = None
         self._time_domain_strain = None
         self._channel = None
@@ -139,6 +149,33 @@ def notch_list(self, notch_list):
             raise ValueError("notch_list {} not understood".format(notch_list))
         self._frequency_mask_updated = False
 
+    @property
+    def crop_duration(self):
+        """
+        The duration of data to crop at the beginning/end of the segment
+        to avoid conditioning artifacts. If a float, that duration is
+        excluded at each end, if a tuple, this specifies the truncation
+        duration at the beginning and end.
+        """
+        return self._crop_duration
+
+    @crop_duration.setter
+    def crop_duration(self, crop_duration):
+        if not isinstance(self.crop_duration, (float, int, list, tuple)):
+            raise TypeError(f"Invalid crop specification {self.crop_duration}")
+        self._crop_duration = crop_duration
+        self._time_mask_updated = False
+
+    @property
+    def cropped_duration(self):
+        """
+        The duration after applying the time-domain mask.
+        """
+        if isinstance(self.crop_duration, (float, int)):
+            return self.duration - 2 * self.crop_duration
+        else:
+            return self.duration - sum(self.crop_duration[:2])
+
     @property
     def frequency_mask(self):
         """ Masking array for limiting the frequency band.
@@ -149,20 +186,69 @@ def frequency_mask(self):
             An array of boolean values
         """
         if not self._frequency_mask_updated:
-            frequency_array = self._times_and_frequencies.frequency_array
+            self._update_frequency_mask()
+        return self._frequency_mask
+
+    def _update_frequency_mask(self):
+        def calculate_frequency_mask(frequency_array):
             mask = ((frequency_array >= self.minimum_frequency) &
                     (frequency_array <= self.maximum_frequency))
             for notch in self.notch_list:
                 mask[notch.get_idxs(frequency_array)] = False
-            self._frequency_mask = mask
-            self._frequency_mask_updated = True
-        return self._frequency_mask
+            return mask
+
+        self._frequency_mask = calculate_frequency_mask(
+            self._times_and_frequencies.frequency_array
+        )
+
+        cropped_frequencies = create_frequency_series(
+            duration=self.cropped_duration,
+            sampling_frequency=self.sampling_frequency
+        )
+        self._cropped_frequency_mask = calculate_frequency_mask(
+            cropped_frequencies
+        )
+
+        self._frequency_mask_updated = True
 
     @frequency_mask.setter
     def frequency_mask(self, mask):
         self._frequency_mask = mask
         self._frequency_mask_updated = True
 
+    @property
+    def cropped_frequency_mask(self):
+        if not self._frequency_mask_updated:
+            self._update_frequency_mask
+        return self._cropped_frequency_mask
+
+    @property
+    def time_mask(self):
+        """ Masking array for cropping corrupted data at the edges.
+
+        Returns
+        =======
+        mask: np.ndarray
+            An array of boolean values
+        """
+        if not self._time_mask_updated:
+            if isinstance(self.crop_duration, (tuple, list)):
+                crop_start, crop_end = self.crop_duration
+            elif isinstance(self.crop_duration, (float, int)):
+                crop_start = crop_end = self.crop_duration
+
+            time_array = self._times_and_frequencies.time_array
+            mask = ((time_array > self.start_time + crop_start) &
+                    (time_array <= self.end_time - crop_end))
+            self._time_mask = mask
+            self._time_mask_updated = True
+        return self._time_mask
+
+    @time_mask.setter
+    def time_mask(self, mask):
+        self._time_mask = mask
+        self._time_mask_updated = True
+
     @property
     def alpha(self):
         return 2 * self.roll_off / self.duration

bilby/gw/likelihood/base.py

@@ -12,7 +12,7 @@
 from ...core.prior import Interped, Prior, Uniform, DeltaFunction
 from ..detector import InterferometerList, get_empty_interferometer, calibration
 from ..prior import BBHPriorDict, Cosmological
-from ..utils import noise_weighted_inner_product, zenith_azimuth_to_ra_dec, ln_i0
+from ..utils import zenith_azimuth_to_ra_dec, ln_i0
 
 
 class GravitationalWaveTransient(Likelihood):
@@ -284,64 +284,56 @@ def calculate_snrs(self, waveform_polarizations, interferometer, *, return_array
             interferometer=interferometer,
             parameters=parameters,
         )
-        _mask = interferometer.frequency_mask
 
         if 'recalib_index' in parameters:
-            signal[_mask] *= self.calibration_draws[interferometer.name][int(parameters['recalib_index'])]
+            signal *= self.calibration_draws[interferometer.name][int(parameters['recalib_index'])]
 
-        d_inner_h = interferometer.inner_product(signal=signal)
-        optimal_snr_squared = interferometer.optimal_snr_squared(signal=signal)
+        whitened_signal = interferometer.whiten_frequency_series(signal)
+        xp = aac.array_namespace(whitened_signal)
+
+        d_inner_h = (xp.conjugate(interferometer.whitened_frequency_domain_strain) * whitened_signal).sum()
+        optimal_snr_squared = (abs(whitened_signal)**2).sum()
         complex_matched_filter_snr = d_inner_h / (optimal_snr_squared**0.5)
 
         d_inner_h_array = None
         optimal_snr_squared_array = None
 
-        normalization = 4 / self.waveform_generator.duration
-        xp = aac.array_namespace(signal)
-
         if return_array is False:
             d_inner_h_array = None
             optimal_snr_squared_array = None
         elif self.time_marginalization and self.calibration_marginalization:
 
             d_inner_h_integrand = xp.tile(
-                interferometer.frequency_domain_strain.conj() * signal /
-                interferometer.power_spectral_density_array, (self.number_of_response_curves, 1)).T
+                xp.conjugate(interferometer.whitened_frequency_domain_strain) * whitened_signal,
+                (self.number_of_response_curves, 1)
+            ).T
 
-            d_inner_h_integrand[_mask] *= self.calibration_draws[interferometer.name].T
+            d_inner_h_integrand[interferometer.frequency_mask] *= self.calibration_draws[interferometer.name].T
 
-            d_inner_h_array = 4 / self.waveform_generator.duration * xp.fft.fft(
-                d_inner_h_integrand[0:-1], axis=0
-            ).T
+            d_inner_h_array = xp.fft.fft(d_inner_h_integrand[:-1], axis=0).T
 
-            optimal_snr_squared_integrand = (
-                normalization * xp.abs(signal)**2 / interferometer.power_spectral_density_array
-            )
+            optimal_snr_squared_integrand = xp.abs(whitened_signal)**2
             optimal_snr_squared_array = xp.dot(
-                optimal_snr_squared_integrand[_mask],
+                optimal_snr_squared_integrand[interferometer.frequency_mask],
                 self.calibration_abs_draws[interferometer.name].T
             )
 
         elif self.time_marginalization and not self.calibration_marginalization:
-            d_inner_h_array = normalization * xp.fft.fft(
-                signal[0:-1]
-                * interferometer.frequency_domain_strain.conj()[0:-1]
-                / interferometer.power_spectral_density_array[0:-1]
+            d_inner_h_integrand = (
+                whitened_signal
+                * xp.conjugate(interferometer.whitened_frequency_domain_strain)
             )
+            d_inner_h_array = xp.fft.fft(d_inner_h_integrand[:-1])
 
         elif self.calibration_marginalization and ('recalib_index' not in parameters):
             d_inner_h_integrand = (
-                normalization *
-                interferometer.frequency_domain_strain.conj() * signal
-                / interferometer.power_spectral_density_array
-            )
-            d_inner_h_array = xp.dot(d_inner_h_integrand[_mask], self.calibration_draws[interferometer.name].T)
+                xp.conjugate(interferometer.whitened_frequency_domain_strain) * whitened_signal
+            )[interferometer.frequency_mask]
+            d_inner_h_array = xp.dot(d_inner_h_integrand, self.calibration_draws[interferometer.name].T)
 
-            optimal_snr_squared_integrand = (
-                normalization * xp.abs(signal)**2 / interferometer.power_spectral_density_array
-            )
+            optimal_snr_squared_integrand = xp.abs(whitened_signal)**2
             optimal_snr_squared_array = xp.dot(
-                optimal_snr_squared_integrand[_mask],
+                optimal_snr_squared_integrand[interferometer.frequency_mask],
                 self.calibration_abs_draws[interferometer.name].T
             )
 
@@ -402,13 +394,10 @@ def priors(self, priors):
     def _calculate_noise_log_likelihood(self):
         log_l = 0
         for interferometer in self.interferometers:
-            mask = interferometer.frequency_mask
-            log_l -= abs(noise_weighted_inner_product(
-                interferometer.frequency_domain_strain[mask],
-                interferometer.frequency_domain_strain[mask],
-                interferometer.power_spectral_density_array[mask],
-                self.waveform_generator.duration) / 2)
-        return log_l
+            log_l -= (
+                abs(interferometer.whitened_frequency_domain_strain)**2
+            ).sum() / 2
+        return log_l.real
 
     def noise_log_likelihood(self):
         # only compute likelihood if called for the 1st time
@@ -472,7 +461,7 @@ def compute_log_likelihood_from_snrs(self, total_snrs, parameters):
                 d_inner_h=total_snrs.d_inner_h, h_inner_h=total_snrs.optimal_snr_squared)
 
         else:
-            log_l = np.real(total_snrs.d_inner_h) - total_snrs.optimal_snr_squared / 2
+            log_l = total_snrs.d_inner_h.real - total_snrs.optimal_snr_squared / 2
 
         return log_l