Merge pull request #6843 from afvincent/fix_psd_demo_complex

tacaswell · web-flow · commit e6efa8ec6ae0 · 2016-07-27T23:22:17.000-04:00
DOC: Fix `psd_demo_complex` similarly to `psd_demo3`
diff --git a/examples/pylab_examples/psd_demo_complex.py b/examples/pylab_examples/psd_demo_complex.py
@@ -1,38 +1,46 @@
-# This is a ported version of a MATLAB example from the signal processing
-# toolbox that showed some difference at one time between Matplotlib's and
-# MATLAB's scaling of the PSD.  This differs from psd_demo3.py in that
-# this uses a complex signal, so we can see that complex PSD's work properly
+"""This is a ported version of a MATLAB example from the signal
+processing toolbox that showed some difference at one time between
+Matplotlib's and MATLAB's scaling of the PSD.
+
+This differs from psd_demo3.py in that this uses a complex signal,
+so we can see that complex PSD's work properly
+
+"""
+
 import numpy as np
 import matplotlib.pyplot as plt
 import matplotlib.mlab as mlab
 
+prng = np.random.RandomState(123456)  # to ensure reproducibility
+
 fs = 1000
 t = np.linspace(0, 0.3, 301)
 A = np.array([2, 8]).reshape(-1, 1)
 f = np.array([150, 140]).reshape(-1, 1)
-xn = (A * np.exp(2j * np.pi * f * t)).sum(axis=0) + 5 * np.random.randn(*t.shape)
+xn = (A * np.exp(2j * np.pi * f * t)).sum(axis=0) + 5 * prng.randn(*t.shape)
+
+fig, (ax0, ax1) = plt.subplots(ncols=2)
 
+fig.subplots_adjust(hspace=0.45, wspace=0.3)
 yticks = np.arange(-50, 30, 10)
-xticks = np.arange(-500, 550, 100)
-plt.subplots_adjust(hspace=0.45, wspace=0.3)
-ax = plt.subplot(1, 2, 1)
+yrange = (yticks[0], yticks[-1])
+xticks = np.arange(-500, 550, 200)
 
-plt.psd(xn, NFFT=301, Fs=fs, window=mlab.window_none, pad_to=1024,
+ax0.psd(xn, NFFT=301, Fs=fs, window=mlab.window_none, pad_to=1024,
         scale_by_freq=True)
-plt.title('Periodogram')
-plt.yticks(yticks)
-plt.xticks(xticks)
-plt.grid(True)
-plt.xlim(-500, 500)
-
-plt.subplot(1, 2, 2, sharex=ax, sharey=ax)
-plt.psd(xn, NFFT=150, Fs=fs, window=mlab.window_none, noverlap=75, pad_to=512,
+ax0.set_title('Periodogram')
+ax0.set_yticks(yticks)
+ax0.set_xticks(xticks)
+ax0.grid(True)
+ax0.set_ylim(yrange)
+
+ax1.psd(xn, NFFT=150, Fs=fs, window=mlab.window_none, pad_to=512, noverlap=75,
         scale_by_freq=True)
-plt.title('Welch')
-plt.xticks(xticks)
-plt.yticks(yticks)
-plt.ylabel('')
-plt.grid(True)
-plt.xlim(-500, 500)
+ax1.set_title('Welch')
+ax1.set_xticks(xticks)
+ax1.set_yticks(yticks)
+ax1.set_ylabel('')  # overwrite the y-label added by `psd`
+ax1.grid(True)
+ax1.set_ylim(yrange)
 
 plt.show()
