Merge e5bea9b into 143170d

scikit-signal · Jul 27, 2017 · 26d5d3b · 26d5d3b
2 parents 143170d + e5bea9b
commit 26d5d3b
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/tftb/processing/cohen.py b/tftb/processing/cohen.py
@@ -179,6 +179,8 @@ def plot(self, kind='cmap', threshold=0.05, sqmod=False, **kwargs):
             kwargs['levels'] = levels
         if sqmod:
             self.tfr = np.abs(self.tfr) ** 2
+        else:
+            self.tfr = np.abs(self.tfr)
         _threshold = np.amax(self.tfr) * threshold
         self.tfr[self.tfr <= _threshold] = 0.0
         super(WignerVilleDistribution, self).plot(kind=kind, threshold=threshold,
@@ -211,10 +213,8 @@ def run(self):
         return np.real(self.tfr), self.ts, self.freqs
 
     def plot(self, **kwargs):
-        if kwargs.get("kind", "") == "cmap":
-            import warnings
-            warnings.warn("cmap is not supported for PseudoWignerVille distributions.")
-        kwargs['kind'] = "contour"
+        if "kind" not in kwargs:
+            kwargs['kind'] = "contour"
         super(PseudoWignerVilleDistribution, self).plot(**kwargs)