Merge pull request #4593 from e-q/fix4328

tacaswell · tacaswell · commit cefab28635aa · 2015-07-15T23:20:20.000-04:00
FIX: Correct output of mlab._spectral_helper when scale_by_freq=False
diff --git a/lib/matplotlib/mlab.py b/lib/matplotlib/mlab.py
@@ -765,9 +765,6 @@ def _spectral_helper(x, y=None, NFFT=None, Fs=None, detrend_func=None,
         pass
 
     if mode == 'psd':
-        # Scale the spectrum by the norm of the window to compensate for
-        # windowing loss; see Bendat & Piersol Sec 11.5.2.
-        result /= (np.abs(windowVals)**2).sum()
 
         # Also include scaling factors for one-sided densities and dividing by
         # the sampling frequency, if desired. Scale everything, except the DC
@@ -787,6 +784,12 @@ def _spectral_helper(x, y=None, NFFT=None, Fs=None, detrend_func=None,
         # values. Perform the same scaling here.
         if scale_by_freq:
             result /= Fs
+            # Scale the spectrum by the norm of the window to compensate for
+            # windowing loss; see Bendat & Piersol Sec 11.5.2.
+            result /= (np.abs(windowVals)**2).sum()
+        else:
+            # In this case, preserve power in the segment, not amplitude
+            result /= np.abs(windowVals).sum()**2
 
     t = np.arange(NFFT/2, len(x) - NFFT/2 + 1, NFFT - noverlap)/Fs
 
diff --git a/lib/matplotlib/tests/test_mlab.py b/lib/matplotlib/tests/test_mlab.py
@@ -1805,31 +1805,37 @@ def test_psd_windowarray(self):
 
     def test_psd_windowarray_scale_by_freq(self):
         freqs = self.freqs_density
+        win = mlab.window_hanning(np.ones(self.NFFT_density_real))
+
         spec, fsp = mlab.psd(x=self.y,
                              NFFT=self.NFFT_density,
                              Fs=self.Fs,
                              noverlap=self.nover_density,
                              pad_to=self.pad_to_density,
-                             sides=self.sides)
+                             sides=self.sides,
+                             window=mlab.window_hanning)
         spec_s, fsp_s = mlab.psd(x=self.y,
                                  NFFT=self.NFFT_density,
                                  Fs=self.Fs,
                                  noverlap=self.nover_density,
                                  pad_to=self.pad_to_density,
                                  sides=self.sides,
+                                 window=mlab.window_hanning,
                                  scale_by_freq=True)
         spec_n, fsp_n = mlab.psd(x=self.y,
                                  NFFT=self.NFFT_density,
                                  Fs=self.Fs,
                                  noverlap=self.nover_density,
                                  pad_to=self.pad_to_density,
                                  sides=self.sides,
+                                 window=mlab.window_hanning,
                                  scale_by_freq=False)
-
         assert_array_equal(fsp, fsp_s)
         assert_array_equal(fsp, fsp_n)
         assert_array_equal(spec, spec_s)
-        assert_allclose(spec_s, spec_n/self.Fs, atol=1e-08)
+        assert_allclose(spec_s*(win**2).sum(),
+                        spec_n/self.Fs*win.sum()**2,
+                        atol=1e-08)
 
     def test_complex_spectrum(self):
         freqs = self.freqs_spectrum
