Merge branch 'add-time-domain-injection-into-data' into 'master'

Adds helper method to inject a signal into time-domain data

Closes #77

See merge request !444

bilby/gw/detector.py

@@ -13,6 +13,7 @@ from ..core import utils
 from ..core.utils import logger
 from ..core.series import CoupledTimeAndFrequencySeries
 from .calibration import Recalibrate
+from . import conversion
 
 try:
     import gwpy
@@ -23,6 +24,7 @@ except ImportError:
 
 try:
     import lal
+    import lalsimulation as lalsim
 except ImportError:
     logger.warning("You do not have lalsuite installed currently. You will"
                    " not be able to use some of the prebuilt functions.")
@@ -2409,3 +2411,104 @@ def load_data_from_cache_file(
         raise ValueError('Data not loaded for {}'.format(ifo.name))
     elif not psd_set:
         raise ValueError('PSD not created for {}'.format(ifo.name))
+
+
+def get_safe_signal_duration(mass_1, mass_2, a_1, a_2, tilt_1, tilt_2, flow=10,
+                             **kwargs):
+    """ Calculate the safe signal duration, given the parameters
+
+    Parameters
+    ----------
+    mass_1, mass_2, a_1, a_2, tilt_1, tilt_2: float
+        The signal parameters
+    flow: float
+        The lower frequency cutoff from which to calculate the signal duration
+
+     Returns
+     -------
+     safe_signal_duration: float
+        The shortest safe signal duration (i.e., the signal duration rounded up
+        to the nearest power of 2)
+
+    """
+    chirp_time = lalsim.SimInspiralChirpTimeBound(
+        flow, mass_1 * lal.MSUN_SI, mass_2 * lal.MSUN_SI,
+        a_1 * np.cos(tilt_1), a_2 * np.cos(tilt_2))
+    return max(2**(int(np.log2(chirp_time)) + 1), 4)
+
+
+def inject_signal_into_gwpy_timeseries(
+        data, waveform_generator, parameters, det, outdir=None, label=None):
+    """ Inject a signal into a gwpy timeseries
+
+    Parameters
+    ----------
+    data: gwpy.timeseries.TimeSeries
+        The time-series data into which we want to inject the signal
+    waveform_generator: bilby.gw.waveform_generator.WaveformGenerator
+        An initialised waveform_generator
+    parameters: dict
+        A dictionary of the signal-parameters to inject
+    ifo: bilby.gw.detector.Interferometer
+        The interferometer for which the data refers too
+    outdir, label: str
+        If given, the outdir and label used to generate a plot
+
+    Returns
+    -------
+    data_and_signal: gwpy.timeseries.TimeSeries
+        The data with the time-domain signal added
+    meta_data: dict
+        A dictionary of meta data about the injection
+
+    """
+    ifo = get_empty_interferometer(det)
+    ifo.strain_data.set_from_gwpy_timeseries(data)
+
+    parameters_check, _ = conversion.convert_to_lal_binary_black_hole_parameters(parameters)
+    safe_time = get_safe_signal_duration(**parameters_check)
+    if data.duration.value < safe_time:
+        ValueError(
+            "Injecting a signal with safe-duration {} longer than the data {}"
+            .format(safe_time, data.duration.value))
+
+    waveform_polarizations = waveform_generator.time_domain_strain(parameters)
+
+    signal = np.zeros(len(data))
+
+    for mode in waveform_polarizations.keys():
+        det_response = ifo.antenna_response(
+            parameters['ra'], parameters['dec'], parameters['geocent_time'],
+            parameters['psi'], mode)
+        signal += waveform_polarizations[mode] * det_response
+    time_shift = ifo.time_delay_from_geocenter(
+        parameters['ra'], parameters['dec'], parameters['geocent_time'])
+
+    dt = parameters['geocent_time'] + time_shift - data.times[0].value
+    n_roll = dt * data.sample_rate.value
+    n_roll = int(np.round(n_roll))
+    signal_shifted = gwpy.timeseries.TimeSeries(
+        data=np.roll(signal, n_roll), times=data.times, unit=data.unit)
+
+    signal_and_data = data.inject(signal_shifted)
+
+    if outdir is not None and label is not None:
+        fig = gwpy.plot.Plot(signal_shifted)
+        fig.savefig('{}/{}_{}_time_domain_injected_signal'.format(
+            outdir, ifo.name, label))
+
+    meta_data = dict()
+    frequency_domain_signal, _ = utils.nfft(signal, waveform_generator.sampling_frequency)
+    meta_data['optimal_SNR'] = (
+        np.sqrt(ifo.optimal_snr_squared(signal=frequency_domain_signal)).real)
+    meta_data['matched_filter_SNR'] = (
+        ifo.matched_filter_snr(signal=frequency_domain_signal))
+    meta_data['parameters'] = parameters
+
+    logger.info("Injected signal in {}:".format(ifo.name))
+    logger.info("  optimal SNR = {:.2f}".format(meta_data['optimal_SNR']))
+    logger.info("  matched filter SNR = {:.2f}".format(meta_data['matched_filter_SNR']))
+    for key in parameters:
+        logger.info('  {} = {}'.format(key, parameters[key]))
+
+    return signal_and_data, meta_data