Merge pull request #4326 from tacaswell/fix_mlab_psd

jenshnielsen · jenshnielsen · commit 063e3f1d2a62 · 2015-04-12T17:52:39.000+01:00
BUG : fix scaling issue with mlab.psd and odd NFFT
diff --git a/lib/matplotlib/mlab.py b/lib/matplotlib/mlab.py
@@ -782,7 +782,15 @@ def _spectral_helper(x, y=None, NFFT=None, Fs=None, detrend_func=None,
         # Also include scaling factors for one-sided densities and dividing by
         # the sampling frequency, if desired. Scale everything, except the DC
         # component and the NFFT/2 component:
-        result[1:-1] *= scaling_factor
+
+        # if we have a even number of frequencies, don't scale NFFT/2
+        if not NFFT % 2:
+            slc = slice(1, -1, None)
+        # if we have an odd number, just don't scale DC
+        else:
+            slc = slice(1, None, None)
+
+        result[slc] *= scaling_factor
 
         # MATLAB divides by the sampling frequency so that density function
         # has units of dB/Hz and can be integrated by the plotted frequency
diff --git a/lib/matplotlib/tests/test_mlab.py b/lib/matplotlib/tests/test_mlab.py
@@ -2960,8 +2960,17 @@ def test_evaluate_equal_dim_and_num_lt(self):
         np.testing.assert_array_almost_equal(y, y_expected, 7)
 
 
-#*****************************************************************
-#*****************************************************************
+def test_psd_onesided_norm():
+    u = np.array([0, 1, 2, 3, 1, 2, 1])
+    dt = 1.0
+    Su = np.abs(np.fft.fft(u) * dt)**2 / float(dt * u.size)
+    P, f = mlab.psd(u, NFFT=u.size, Fs=1/dt, window=mlab.window_none,
+                    detrend=mlab.detrend_none, noverlap=0, pad_to=None,
+                    scale_by_freq=None,
+                    sides='onesided')
+    Su_1side = np.append([Su[0]], Su[1:4] + Su[4:][::-1])
+    assert_allclose(P, Su_1side, atol=1e-06)
+
 
 if __name__ == '__main__':
     import nose
