Convert logic of get_open_strain to methods of strain_data

Adds functions to set the strain data from open data, generate a PSD,
filter, and window.

tupak/gw/detector.py

@@ -6,7 +6,6 @@ import os
 import matplotlib.pyplot as plt
 import numpy as np
 import gwpy
-from gwpy.signal import filter_design
 from scipy import signal
 from scipy.interpolate import interp1d
 
@@ -99,6 +98,7 @@ class InterferometerStrainData(object):
         self.sampling_frequency = None
         self.duration = None
         self.start_time = None
+        self.buffer_time = 0
 
         self._frequency_domain_strain = None
         self._frequency_array = None
@@ -265,8 +265,8 @@ class InterferometerStrainData(object):
             raise ValueError("Input timeseries is not a gwpy TimeSeries")
         self.start_time = timeseries.epoch.value
         self.sampling_frequency = timeseries.sample_rate.value
-        self.duration = timeseries.value
-        self.time_domain_strain = timeseries.value
+        self.duration = timeseries.duration.value
+        self._time_domain_strain = timeseries.value
 
     def set_from_time_domain_strain(
             self, time_domain_strain, sampling_frequency=None, duration=None,
@@ -283,6 +283,85 @@ class InterferometerStrainData(object):
         else:
             raise ValueError("Data times do not match time array")
 
+    def set_from_open_data(
+            self, name, start_time, duration=4, outdir='outdir', cache=True,
+            buffer_time=1, **kwargs):
+        """
+
+        Parameters
+        ----------
+        name: str
+            Detector name, e.g., 'H1'.
+        start_time: float
+            Start GPS time of segment.
+        duration: float, optional
+            The total time (in seconds) to analyse. Defaults to 4s.
+        outdir: str
+            Directory where the psd files are saved
+        cache: bool, optional
+            Whether or not to store/use the acquired data.
+        buffer_time: float
+            Read in data with `start_time-buffer_time` and
+            `start_time+duration+buffer_time`
+        **kwargs:
+            All keyword arguments are passed to
+            `gwpy.timeseries.TimeSeries.fetch_open_data()`.
+
+        """
+
+        self.buffer_time = buffer_time
+
+        timeseries = tupak.gw.utils.get_open_strain_data(
+            name, start_time, start_time+duration, outdir=outdir, cache=cache,
+            buffer_time=buffer_time, **kwargs)
+
+        self.set_from_gwpy_timeseries(timeseries)
+
+    def low_pass_filter(self, filter_freq):
+        """ Low pass filter the data """
+        bp = gwpy.signal.filter_design.lowpass(
+            filter_freq, self.sampling_frequency)
+        strain = gwpy.timeseries.TimeSeries(
+            self.time_domain_strain, sample_rate=self.sampling_frequency)
+        strain = strain.filter(bp, filtfilt=True)
+        strain = strain.crop(*strain.span.contract(self.buffer_time))
+        self._time_domain_strain = strain.value
+
+    def apply_tukey_window(self, alpha):
+        N = len(self.time_domain_strain)
+        self._time_domain_strain *= signal.windows.tukey(N, alpha=alpha)
+
+    def create_power_spectral_density(self, alpha, name, outdir='outdir'):
+        """ Use the time domain strain to generate a power spectral density
+
+        This create a Tukey-windowed power spectral density and writes it to a
+        PSD file.
+
+        Parameters
+        ----------
+        alpha: float
+            The tukey window shape parameter passed to `scipy.signal.tukey`.
+        outdir: str
+            The output directory to write the PSD file too
+
+        Returns
+        -------
+        psd_file_name: str
+            The path to the PSD file
+
+        """
+        NFFT = int(self.sampling_frequency * self.duration)
+        window = signal.windows.tukey(NFFT, alpha=alpha)
+        strain = gwpy.timeseries.TimeSeries(
+            self.time_domain_strain, sample_rate=self.sampling_frequency)
+        psd = strain.psd(fftlength=self.duration, window=window)
+        psd_file = '{}/{}_PSD_{}_{}.txt'.format(
+            outdir, name, self.start_time, self.duration)
+        with open('{}'.format(psd_file), 'w+') as file:
+            for f, p in zip(psd.frequencies.value, psd.value):
+                file.write('{} {}\n'.format(f, p))
+        return psd_file
+
     def _infer_frequency_domain_dependence(
             self, sampling_frequency, duration, frequency_array):
         """ Helper function to figure out if the frequency_array, or
@@ -1181,7 +1260,7 @@ class PowerSpectralDensity(object):
             sampling_frequency = int(strain.sample_rate.value)
 
             # Low pass filter
-            bp = filter_design.lowpass(filter_freq, strain.sample_rate)
+            bp = gwpy.signal.filter_design.lowpass(filter_freq, strain.sample_rate)
             strain = strain.filter(bp, filtfilt=True)
             strain = strain.crop(*strain.span.contract(1))
 
@@ -1404,50 +1483,30 @@ def get_interferometer_with_open_data(
     if start_time is None:
         start_time = trigger_time + 2 - time_duration
 
-    strain = tupak.gw.utils.get_open_strain_data(
-        name, start_time - 1, start_time + time_duration + 1, outdir=outdir, cache=cache,
-        raw_data_file=raw_data_file, **kwargs)
-
-    strain_psd = tupak.gw.utils.get_open_strain_data(
-        name, start_time + time_duration + psd_offset,
-        start_time + time_duration + psd_offset + psd_duration,
-        raw_data_file=raw_data_file,
+    strain = InterferometerStrainData()
+    strain.set_from_open_data(
+        name=name, start_time=start_time, duration=time_duration,
         outdir=outdir, cache=cache, **kwargs)
 
-    sampling_frequency = int(strain.sample_rate.value)
+    strain_psd = InterferometerStrainData()
+    strain_psd.set_from_open_data(
+        name=name, start_time=start_time + time_duration + psd_offset,
+        duration=psd_duration, outdir=outdir, cache=cache, **kwargs)
 
     # Low pass filter
-    bp = filter_design.lowpass(filter_freq, strain.sample_rate)
-    strain = strain.filter(bp, filtfilt=True)
-    strain = strain.crop(*strain.span.contract(1))
-    strain_psd = strain_psd.filter(bp, filtfilt=True)
-    strain_psd = strain_psd.crop(*strain_psd.span.contract(1))
+    strain.low_pass_filter(filter_freq)
+    strain_psd.low_pass_filter(filter_freq)
 
     # Create and save PSDs
-    NFFT = int(sampling_frequency * time_duration)
-    window = signal.windows.tukey(NFFT, alpha=alpha)
-    psd = strain_psd.psd(fftlength=time_duration, window=window)
-    psd_file = '{}/{}_PSD_{}_{}.txt'.format(
-        outdir, name, start_time + time_duration + psd_offset, psd_duration)
-    with open('{}'.format(psd_file), 'w+') as file:
-        for f, p in zip(psd.frequencies.value, psd.value):
-            file.write('{} {}\n'.format(f, p))
-
-    time_series = strain.times.value
-    time_duration = time_series[-1] - time_series[0]
-
-    # Apply Tukey window
-    N = len(time_series)
-    strain = strain * signal.windows.tukey(N, alpha=alpha)
+    psd_file = strain_psd.create_power_spectral_density(
+        name=name, alpha=alpha, outdir=outdir)
+
+    strain.apply_tukey_window(alpha=alpha)
 
     interferometer = get_empty_interferometer(name)
     interferometer.power_spectral_density = PowerSpectralDensity(
         psd_file=psd_file)
-    interferometer.set_strain_data_from_frequency_domain_strain(
-        frequency_domain_strain=utils.nfft(
-            strain.value, sampling_frequency)[0],
-        sampling_frequency=sampling_frequency, duration=time_duration,
-        start_time=strain.epoch.value)
+    interferometer.strain_data = strain
 
     if plot:
         interferometer.plot_data(outdir=outdir, label=label)
@@ -1593,7 +1652,8 @@ def get_event_data(
                 psd_offset=psd_offset, psd_duration=psd_duration, cache=cache,
                 outdir=outdir, label=label, plot=plot, filter_freq=filter_freq,
                 raw_data_file=raw_data_file, **kwargs))
-        except ValueError:
+        except ValueError as e:
+            logging.debug("Error raised {}".format(e))
             logging.warning('No data found for {}.'.format(name))
 
     return InterferometerSet(interferometers)

tupak/gw/utils.py

@@ -291,7 +291,7 @@ def get_event_time(event):
 
 
 def get_open_strain_data(
-        name, t1, t2, outdir, cache=False, raw_data_file=None, **kwargs):
+        name, t1, t2, outdir, cache=False, buffer_time=0, **kwargs):
     """ A function which accesses the open strain data
 
     This uses `gwpy` to download the open data and then saves a cached copy for
@@ -307,9 +307,10 @@ def get_open_strain_data(
         The output directory to place data in
     cache: bool
         If true, cache the data
+    buffer_time: float
+        Time to add to the begining and end of the segment.
     **kwargs:
         Passed to `gwpy.timeseries.TimeSeries.fetch_open_data`
-    raw_data_file
 
     Returns
     -----------
@@ -317,19 +318,18 @@ def get_open_strain_data(
 
     """
     filename = '{}/{}_{}_{}.txt'.format(outdir, name, t1, t2)
-    if raw_data_file is not None:
-        logging.info('Using raw_data_file {}'.format(raw_data_file))
-        strain = TimeSeries.read(raw_data_file)
-        if (t1 > strain.times[0].value) and (t2 < strain.times[-1].value):
-            logging.info('Using supplied raw data file')
-            strain = strain.crop(t1, t2)
-        else:
-            raise ValueError('Supplied file does not contain requested data')
-    elif os.path.isfile(filename) and cache:
+
+    if buffer_time < 0:
+        raise ValueError("buffer_time < 0")
+    t1 = t1 - buffer_time
+    t2 = t2 + buffer_time
+
+    if os.path.isfile(filename) and cache:
         logging.info('Using cached data from {}'.format(filename))
         strain = TimeSeries.read(filename)
     else:
-        logging.info('Fetching open data ...')
+        logging.info('Fetching open data from {} to {} with buffer time {}'
+                     .format(t1, t2, buffer_time))
         strain = TimeSeries.fetch_open_data(name, t1, t2, **kwargs)
         logging.info('Saving data to {}'.format(filename))
         strain.write(filename)