Support show argument in plot_iir_filter_analysis().

applications/lofar2/model/pfb_os/dsp.py

@@ -339,19 +339,29 @@ def plot_time_response(h, name='', markers=False):
     plt.grid(True)
 
 
-def plot_iir_filter_analysis(b, a, fs=1, Ntime=100, step=False):
-    """Plot iir filter analysis results.
+def plot_iir_filter_analysis(b, a, fs=1, whole=False, Ntime=100, step=False, show=[]):
+    """Plot and print iir filter analysis results.
 
     Input:
     . b, a: IIR filter coefficients in same format as for scipy.signal.freqz
     . fs: sample frequency
+    . whole: False to have only positive frequencies, so 0 to fs / 2 in
+      spectrum plots, True to have 0 to fs.
     . Ntime: number of timesamples for impulse response
     . step: False for impulse response, True for step response
+    . show[]: which plots to show from: 'zplane', 'power spectrum',
+      'phase spectrum', 'time response'
+
+    Return: z, p, k
     """
 
     # Plot poles / zeros diagram in z-plane
+    if 'zplane' in show:
         fig1, ax1 = plt.subplots(1)
+        ax1.set_title('Zeros (o) and poles (x) in z-plane')
         z, p, k = dsp_fpga_lib.zplane(b, a, plt_ax=ax1)  # uses np.roots(a), np.roots(b)
+    else:
+        z, p, k = dsp_fpga_lib.zplane(b, a)  # no plot, only calculate z, p, k
     print('Zeros, poles and gain from b, a coefficients:')
     if len(z) > 0:
         print('. zeros:')
@@ -370,24 +380,38 @@ def plot_iir_filter_analysis(b, a, fs=1, Ntime=100, step=False):
 
     # Plot transfer function H(f), is H(z) for z = exp(j w), so along the unit circle
     # . 0 Hz at 1 + 0j, fs / 4 at 0 + 1j, fNyquist = fs / 2 at -1 + 0j
-    # . use whole=False to have only positve frequencies, so 0 to fs / 2
+    # . use whole=False to have only positive frequencies, so 0 to fs / 2
     # . use Nfreq frequency points
-    fig2, ax2 = plt.subplots(1)
+    if 'power spectrum' in show or 'phase spectrum' in show:
         Nfreq = 1024
-    [f, HF] = signal.freqz(b, a, Nfreq, whole=False, fs=fs)
+        [f, HF] = signal.freqz(b, a, Nfreq, whole=whole, fs=fs)
+        if 'power spectrum' in show:
+            fig2, ax2 = plt.subplots(1)
             ax2.plot(f, pow_db(HF))
+            ax2.set_title('Power spectrum')
             ax2.set_xlabel('frequency [fs = %f]' % fs)
             ax2.set_ylabel('HF power [dB]')
 
+        if 'phase spectrum' in show:
+            fig3, ax3 = plt.subplots(1)
+            ax3.plot(f, np.unwrap(np.angle(HF)))
+            ax3.set_title('Phase spectrum')
+            ax3.set_xlabel('frequency [fs = %f]' % fs)
+            ax3.set_ylabel('HF phase [rad]')
+
     # Plot impulse response
+    if 'time response' in show:
         Ts = 1 / fs
-    fig3, ax3 = plt.subplots(1)
-    step = True  # step response (makes impz use np.cumsum(h))
-    step = False
-    [h, t] = dsp_fpga_lib.impz(b, a, FS=fs, N=Ntime, step=step)  # uses signal.lfilter()
-    (ml, sl, bl) = ax3.stem(t, h, linefmt='b-', markerfmt='ro', basefmt='k')
-    ax3.set_xlabel('time [Ts = %f]' % Ts)
-    ax3.set_ylabel('h[n]')
+        fig4, ax4 = plt.subplots(1)
+        # impz uses signal.lfilter(), and cumsum(h) for step response
+        [h, t] = dsp_fpga_lib.impz(b, a, FS=fs, N=Ntime, step=step)
+        (ml, sl, bl) = ax4.stem(t, h, linefmt='b-', markerfmt='ro', basefmt='k')
+        if step:
+            ax4.set_title('Step response')
+        else:
+            ax4.set_title('Impulse response')
+        ax4.set_xlabel('time [Ts = %f]' % Ts)
+        ax4.set_ylabel('h[n]')
     return z, p, k