Merge pull request #108 from jaidevd/jd-fix-spectra

Fix spectrum plots

doc/Makefile

@@ -6,6 +6,7 @@ SPHINXOPTS    =
 SPHINXBUILD   = sphinx-build
 PAPER         =
 BUILDDIR      = _build
+EXAMPLEDIR    = auto_examples
 
 # User-friendly check for sphinx-build
 ifeq ($(shell which $(SPHINXBUILD) >/dev/null 2>&1; echo $$?), 1)
@@ -48,6 +49,7 @@ help:
 
 clean:
 	rm -rf $(BUILDDIR)/*
+	rm -rf $(EXAMPLEDIR)/*
 
 html:
 	$(SPHINXBUILD) -b html $(ALLSPHINXOPTS) $(BUILDDIR)/html

doc/_gallery/plot_4_1_2_sin_gauss_pwv.py

@@ -22,11 +22,8 @@
 
 from tftb.generators import fmconst, amgauss
 from tftb.processing import PseudoWignerVilleDistribution
-import numpy as np
 
 sig = fmconst(128, 0.15)[0] + amgauss(128) * fmconst(128, 0.4)[0]
 tfr = PseudoWignerVilleDistribution(sig)
 tfr.run()
-tfr.plot(show_tf=True, kind="contour",
-        freq_x=(abs(np.fft.fftshift(np.fft.fft(sig))) ** 2)[::-1][:64],
-        freq_y=np.arange(sig.shape[0] / 2))
+tfr.plot(show_tf=True, kind="contour")

doc/_gallery/plot_4_1_2_sin_gauss_wv.py

@@ -22,11 +22,8 @@
 
 from tftb.generators import fmconst, amgauss
 from tftb.processing import WignerVilleDistribution
-import numpy as np
 
 sig = fmconst(128, 0.15)[0] + amgauss(128) * fmconst(128, 0.4)[0]
 tfr = WignerVilleDistribution(sig)
 tfr.run()
-tfr.plot(show_tf=True, kind='contour',
-        freq_x=(abs(np.fft.fftshift(np.fft.fft(sig))) ** 2)[::-1][:64],
-        freq_y=np.arange(sig.shape[0] / 2))
+tfr.plot(show_tf=True, kind='contour')

tftb/processing/affine.py

@@ -99,9 +99,10 @@ def _mellin_transform(self, S1, S2):
         return beta, mellin1, mellin2
 
     def plot(self, kind="contour", show_tf=True, **kwargs):
+        freq_y = np.linspace(self.fmin, self.fmax, self.signal.shape[0] / 2)
         super(AffineDistribution, self).plot(kind=kind, show_tf=show_tf,
                                              contour_y=self.freqs,
-                                             freq_x=self.freqs, **kwargs)
+                                             freq_y=freq_y, **kwargs)
 
     def run(self):
         raise NotImplementedError

tftb/processing/base.py

@@ -15,6 +15,9 @@
 
 
 class BaseTFRepresentation(object):
+
+    isaffine = False
+
     def __init__(self, signal, **kwargs):
         """Create a base time-frequency representation object.
 
@@ -48,6 +51,14 @@ def __init__(self, signal, **kwargs):
         self.freqs = freqs.astype(float) / self.n_fbins
         self.tfr = np.zeros((self.n_fbins, self.ts.shape[0]), dtype=complex)
 
+    def _get_spectrum(self):
+        if not self.isaffine:
+            return np.fft.fftshift(np.abs(np.fft.fft(self.signal)) ** 2)
+        nf2 = self.tfr.shape[0]
+        spec = np.abs(np.fft.fft(self.signal[self.ts.min():(self.ts.max() + 1)],
+                                 2 * nf2)) ** 2
+        return spec[:nf2]
+
     def _make_window(self):
         """Make a Hamming window function.
 
@@ -103,7 +114,7 @@ def plot(self, ax=None, kind='cmap', show=True, default_annotation=True,
                 if levels is not None:
                     axTF.contour(t, f, self.tfr, levels, **kwargs)
                 else:
-                    if getattr(self, "isaffine", False):
+                    if self.isaffine:
                         maxi = np.amax(self.tfr)
                         mini = max(np.amin(self.tfr), maxi * threshold)
                         levels = np.linspace(mini, maxi, 65)
@@ -116,10 +127,11 @@ def plot(self, ax=None, kind='cmap', show=True, default_annotation=True,
             axTime.plot(np.real(self.signal))
             axFreq = divider.append_axes("left", 1.2, pad=0.5)
             k = int(np.floor(self.signal.shape[0] / 2.0))
-            freq_x = kwargs.get('freq_x',
-                    abs(np.fft.fftshift(np.fft.fft(self.signal))) ** 2)[::-1][:k]
-            freq_y = kwargs.get('freq_y',
-                    np.arange(k))
+            freq_x = kwargs.get('freq_x', self._get_spectrum()[::-1][:k])
+            if self.isaffine:
+                freq_y = kwargs.get('freq_y', self.freqs)
+            else:
+                freq_y = kwargs.get('freq_y', np.arange(k))
             axFreq.plot(freq_x, freq_y)
             if default_annotation:
                 axTF.grid(True)
@@ -132,12 +144,17 @@ def plot(self, ax=None, kind='cmap', show=True, default_annotation=True,
                 axTime.set_ylabel('Real part')
                 axTime.set_title('Signal in time')
                 axTime.grid(True)
-                axFreq.set_ylim(0, freq_y.shape[0] - 1)
+                if not self.isaffine:
+                    axFreq.set_ylim(0, freq_y.shape[0] - 1)
+                else:
+                    axFreq.set_ylim(freq_y[0], freq_y[-1])
                 axFreq.set_ylabel('Spectrum')
                 axFreq.set_yticklabels([])
                 axFreq.set_xticklabels([])
                 axFreq.grid(True)
-                axFreq.invert_xaxis()
+                if not self.isaffine:
+                    axFreq.invert_xaxis()
+                    axFreq.invert_yaxis()
         else:
             if (ax is None) and (kind != "surf"):
                 fig = plt.figure()

tftb/processing/cohen.py

@@ -106,7 +106,7 @@ def plot(self, kind='cmap', threshold=0.05, sqmod=True, **kwargs):
 
 class PseudoMargenauHillDistribution(MargenauHillDistribution):
 
-    name = "pseudo margenau hill"
+    name = "pseudo margenau-hill"
 
     def _make_window(self):
         hlength = np.floor(self.n_fbins / 4.0)

tftb/processing/utils.py

@@ -11,6 +11,10 @@
 import numpy as np
 
 
+def get_spectrum(signal):
+    return np.fft.fftshift(np.abs(np.fft.fft(signal)) ** 2)
+
+
 def integrate_2d(mat, x=None, y=None):
     """integrate_2d