Merge pull request #153 from jaidevd/jd-fix-wv-timestamps

Fix indexing errors and timestamps for Cohen class distributions.

doc/_gallery/plot_4_1_1_pwv_atoms.py

@@ -28,6 +28,7 @@
              [32, .35, 20, 1],
              [96, .35, 20, 1]])
 g = atoms(128, x)
-spec = PseudoWignerVilleDistribution(g)
+t = np.linspace(0, 1, 128)
+spec = PseudoWignerVilleDistribution(g, timestamps=t)
 spec.run()
 spec.plot(kind="contour", scale="log", show_tf=True)

doc/_gallery/plot_4_1_2_sin_gauss_pwv.py

@@ -22,8 +22,10 @@
 
 from tftb.generators import fmconst, amgauss
 from tftb.processing import PseudoWignerVilleDistribution
+import numpy as np
 
+t = np.linspace(0, 1, 128)
 sig = fmconst(128, 0.15)[0] + amgauss(128) * fmconst(128, 0.4)[0]
-tfr = PseudoWignerVilleDistribution(sig)
+tfr = PseudoWignerVilleDistribution(sig, timestamps=t)
 tfr.run()
 tfr.plot(show_tf=True, kind="contour")

tftb/processing/cohen.py

@@ -23,9 +23,9 @@ def run(self):
         lh = (self.fwindow.shape[0] - 1) // 2
         rangemin = min([round(self.n_fbins / 2.0) - 1, lh])
         starts = -np.min(np.c_[rangemin * np.ones(self.ts.shape), self.ts - 1],
-                axis=1).astype(int)
+                         axis=1).astype(int)
         ends = np.min(np.c_[rangemin * np.ones(self.ts.shape),
-            self.signal.shape[0] - self.ts], axis=1).astype(int)
+                      self.signal.shape[0] - self.ts], axis=1).astype(int)
         conj_fwindow = np.conj(self.fwindow)
         for icol in range(self.tfr.shape[1]):
             ti = self.ts[icol]
@@ -85,8 +85,8 @@ def run(self):
         for icol in range(self.ts.shape[0]):
             tau = np.arange(min([self.n_fbins - 1, lh, icol - 1]) + 1)
             indices = np.remainder(self.n_fbins + tau, self.n_fbins) + 1
-            self.tfr[indices, icol] = self.fwindow[lh + tau] * self.signal[icol] * np.conj(
-                    self.signal[icol - tau])
+            self.tfr[indices, icol] = self.fwindow[lh + tau] * \
+                self.signal[icol] * np.conj(self.signal[icol - tau])
         self.tfr = np.real(np.fft.fft(self.tfr, axis=0))
         return self.tfr, self.ts, self.freqs
 
@@ -151,21 +151,20 @@ class WignerVilleDistribution(BaseTFRepresentation):
     def run(self):
         tausec = round(self.n_fbins / 2.0)
         winlength = tausec - 1
-        taulens = np.min(np.c_[self.ts, self.signal.shape[0] - self.ts - 1,
-            winlength * np.ones(self.ts.shape)], axis=1)
+        taulens = np.min(np.c_[np.arange(self.signal.shape[0]),
+                               self.signal.shape[0] - np.arange(self.signal.shape[0]) - 1,
+                         winlength * np.ones(self.ts.shape)], axis=1)
         conj_signal = np.conj(self.signal)
         for icol in range(self.ts.shape[0]):
-            ti = self.ts[icol]
             taumax = taulens[icol]
             tau = np.arange(-taumax, taumax + 1).astype(int)
             indices = np.remainder(self.n_fbins + tau, self.n_fbins).astype(int)
-            self.tfr[indices, icol] = self.signal[ti + tau] * \
-                conj_signal[ti - tau]
-            if (ti <= self.signal.shape[0] - tausec) and (ti >= tausec + 1):
-                self.tfr[tausec, icol] = 0.5 * (self.signal[ti + tausec, 0] *
-                        np.conj(self.signal[ti - tausec, 0])) + \
-                        (self.signal[ti - tausec, 0] *
-                        conj_signal[ti + tausec, 0])
+            self.tfr[indices, icol] = self.signal[icol + tau] * \
+                conj_signal[icol - tau]
+            if (icol <= self.signal.shape[0] - tausec) and (icol >= tausec + 1):
+                self.tfr[tausec, icol] = (self.signal[icol + tausec, 0] *
+                                          np.conj(self.signal[icol - tausec, 0])) + \
+                    (self.signal[icol - tausec, 0] * conj_signal[icol + tausec, 0])
         self.tfr = np.fft.fft(self.tfr, axis=0)
         self.tfr = np.real(self.tfr)
         self.freqs = 0.5 * np.arange(self.n_fbins, dtype=float) / self.n_fbins
@@ -193,26 +192,27 @@ class PseudoWignerVilleDistribution(WignerVilleDistribution):
     def run(self):
         lh = (self.fwindow.shape[0] - 1) // 2
         for icol in range(self.ts.shape[0]):
-            ti = self.ts[icol]
-            taumaxvals = (ti, self.signal.shape[0] - ti - 1,
+            taumaxvals = (icol, self.signal.shape[0] - icol - 1,
                           np.round(self.n_fbins / 2.0), lh)
             taumax = np.min(taumaxvals)
             tau = np.arange(-taumax, taumax + 1).astype(int)
             indices = np.remainder(self.n_fbins + tau, self.n_fbins).astype(int)
-            self.tfr[indices, icol] = self.fwindow[lh + tau] * self.signal[ti + tau] * \
-                np.conj(self.signal[ti - tau])
+            self.tfr[indices, icol] = self.fwindow[lh + tau] * self.signal[icol + tau] * \
+                np.conj(self.signal[icol - tau])
             tau = np.round(self.n_fbins / 2.0)
-            if (ti <= self.signal.shape[0] - tau) and (ti >= tau + 1) and (tau <= lh):
-                self.tfr[int(tau), icol] = 0.5 * (self.fwindow[lh + tau] * self.signal[ti + tau, 0] *
-                        np.conj(self.signal[ti - tau, 0]) + self.fwindow[lh - tau] *
-                        self.signal[ti - tau, 0] * np.conj(self.signal[ti + tau, 0]))
+            if (icol <= self.signal.shape[0] - tau) and (icol >= tau + 1) and (tau <= lh):
+                self.tfr[int(tau), icol] = self.fwindow[lh + tau] * \
+                    self.signal[icol + tau, 0] * np.conj(self.signal[icol - tau, 0]) + \
+                    self.fwindow[lh - tau] * self.signal[icol - tau, 0] * \
+                    np.conj(self.signal[icol + tau, 0])
+                self.tfr[int(tau), icol] *= 0.5
 
         self.tfr = np.fft.fft(self.tfr, axis=0)
         return np.real(self.tfr), self.ts, self.freqs
 
 
 def smoothed_pseudo_wigner_ville(signal, timestamps=None, freq_bins=None,
-                                twindow=None, fwindow=None):
+                                 twindow=None, fwindow=None):
     """Smoothed Pseudo Wigner-Ville time-frequency distribution.
     :param signal: signal to be analyzed
     :param timestamps: time instants of the signal
@@ -287,6 +287,7 @@ def smoothed_pseudo_wigner_ville(signal, timestamps=None, freq_bins=None,
 if __name__ == '__main__':
     from tftb.generators import anapulse
     sig = anapulse(128)
-    spec = WignerVilleDistribution(sig)
+    t = np.linspace(0, 1, 128)
+    spec = WignerVilleDistribution(sig, timestamps=t)
     spec.run()
     spec.plot(kind="contour", scale="log")

tftb/processing/tests/test_cohen.py

@@ -35,7 +35,7 @@ def test_spectrogram_time_invariance(self):
         shift = 64
         timeshifted_signal = np.roll(signal, shift)
         timeshifted_tfr, _, _ = cohen.Spectrogram(timeshifted_signal, n_fbins=64,
-                                            fwindow=window).run()
+                                                  fwindow=window).run()
         rolled_tfr = np.roll(tfr, shift, axis=1)
         # the time invariance property holds mostly qualitatively. The shifted
         # TFR is not numerically indentical to the rolled TFR, having
@@ -132,5 +132,6 @@ def test_pseudo_wv_reality(self):
         tfr, _, _ = cohen.PseudoWignerVilleDistribution(signal).run()
         self.assertTrue(np.all(np.isreal(tfr)))
 
+
 if __name__ == '__main__':
     unittest.main()