conven

Author: ahnitz

pycbc/filter/qtransform.py

@@ -58,9 +58,8 @@ def inspiral_qtransform_generator(segments):
         seg /= seg.psd
 
     comb_q_dict = {}
-    for s_num, stilde in enumerate(segments):
+    for stilde in segments:
         # getting q-tiles for segments
-        del s_num
         q_base, q_frange, q_data = inspiral_tiling(stilde)
 
         # getting q-plane for segment
@@ -104,18 +103,13 @@ def inspiral_tiling(seg, frange=(0,np.inf), qrange=(4,64), mismatch=0.2):
         whitened strain of analysis chunk
 
     """
-
-    # retrieve segment psd value
-    seg_psd = seg.psd
-
     # retrieve whitened strain
-    white_strain = (seg / seg_psd ** 0.5 * seg_psd.delta_f)
-    data = white_strain
+    white_strain = (seg / seg.psd ** 0.5 * seg.psd.delta_f)
 
     # perform Q-tiling
-    q_base, frange = qtiling(data, qrange, frange, mismatch)
+    q_base, frange = qtiling(white_strain, qrange, frange, mismatch)
 
-    return q_base, frange, data
+    return q_base, frange, white_strain
 
 def plotter(interp, out_dir, now, frange, tres, fres):
     """Plotting mechanism for pycbc spectrograms
@@ -206,30 +200,30 @@ def qplane(qplane_tile_dict, fseries, frange, normalized=True, tres=1., fres=1.,
     # check for sampling rate
     sampling = fseries.sample_rate
 
-    max_energy = []
+    max_energy, max_key = None, None
     for i, key in enumerate(qplane_tile_dict):
         energies_lst=[]
         for tile in qplane_tile_dict[key]:
-            energies = qtransform(fseries, tile[1], tile[0])
+            normed, raw = qtransform(fseries, tile[1], tile[0])
             if normalized:
-                energies = energies[0]
+                energies = normed
             else:
-                energies = energies[1]
+                energies = raw
+
             energies_lst.append(energies)
+
             if i == 0:
-                max_energy.append(max(energies))
-                max_energy.append(tile)
-                max_energy.append(key)
-            elif max(energies) > max_energy[0]:
-                max_energy[0] = max(energies)
-                max_energy[1] = tile
-                max_energy[2] = key
-                max_energy[3] = energies
+                max_energy = max(energies)
+                max_key = key
+            elif max(energies) > max_energy:
+                max_energy = max(energies)
+                max_key = key
+
         qplane_qtrans_dict[key] = energies_lst
 
     # record q-transform output for peak q
-    result = qplane_qtrans_dict[max_energy[2]]
-    qtile_max = qplane_tile_dict[max_energy[2]]
+    result = qplane_qtrans_dict[max_key]
+    qtile_max = qplane_tile_dict[max_key]
 
     # then interpolate the spectrogram to increase the frequency resolution
     if fres is None:  # unless user tells us not to
@@ -247,7 +241,7 @@ def qplane(qplane_tile_dict, fseries, frange, normalized=True, tres=1., fres=1.,
         # initialize some variables
         frequencies = []
 
-        for i in qplane_tile_dict[max_energy[2]]:
+        for i in qtile_max:
             frequencies.append(i[0])
 
         # 2-D interpolation

pycbc/types/timeseries.py

@@ -427,6 +427,47 @@ def psd(self, segment_duration, **kwds):
                            seg_stride=seg_stride,
                            **kwds)
 
+    def whiten(self, segment_duration, max_filter_duration, trunc_method=None,
+                     remove_corrupted=True, low_frequency_cutoff=None, **kwds):
+        """ Return a whitened time series
+
+        Parameters
+        ----------
+        segment_duration: float
+            Duration in seconds to use for each sample of the spectrum.
+        max_filter_duration : int
+            Maximum length of the time-domain filter in seconds.
+        trunc_method : {None, 'hann'}
+            Function used for truncating the time-domain filter.
+            None produces a hard truncation at `max_filter_len`.
+        remove_corrupted : {True, boolean}
+            If True, the region of the time series corrupted by the whitening
+            is excised before returning. If false, the corrupted regions
+            are not excised and the full time series is returned.
+        kwds : keywords
+            Additional keyword arguments are passed on to the `pycbc.psd.welch` method.
+            
+        Returns
+        -------
+        whitened_data : TimeSeries
+            The whitened time series
+        """
+        from pycbc.psd import inverse_spectrum_truncation, interpolate
+        psd = self.psd(segment_duration, **kwds)
+        psd = interpolate(psd, self.delta_f)
+        max_filter_len = int(max_filter_duration * self.sample_rate)
+        psd = inverse_spectrum_truncation(psd,
+                   max_filter_len=max_filter_len, 
+                   low_frequency_cutoff=low_frequency_cutoff,
+                   trunc_method=trunc_method)
+        white = (self.to_frequencyseries() / psd ** 0.5).to_timeseries()
+
+        if remove_corrupted:
+            white = white[max_filter_len/2:len(self)-max_filter_len/2]
+
+        return white
+        
+
     def save(self, path, group = None):
         """
         Save time series to a Numpy .npy, hdf, or text file. The first column