Skip to content

Instantly share code, notes, and snippets.

@endolith
Last active March 16, 2023 14:38
Show Gist options
  • Save endolith/4625838 to your computer and use it in GitHub Desktop.
Save endolith/4625838 to your computer and use it in GitHub Desktop.
Plot poles and zeros in Z plane for a transfer function
# -*- coding: utf-8 -*-
"""
Combination of
http://scipy-central.org/item/52/1/zplane-function
and
http://www.dsprelated.com/showcode/244.php
with my own modifications
"""
# Copyright (c) 2011 Christopher Felton
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.
# The following is derived from the slides presented by
# Alexander Kain for CS506/606 "Special Topics: Speech Signal Processing"
# CSLU / OHSU, Spring Term 2011.
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import patches
from matplotlib.pyplot import axvline, axhline
from collections import defaultdict
def zplane(z, p, filename=None):
"""Plot the complex z-plane given zeros and poles.
"""
# get a figure/plot
ax = plt.subplot(2, 2, 1)
# TODO: should just inherit whatever subplot it's called in?
# Add unit circle and zero axes
unit_circle = patches.Circle((0,0), radius=1, fill=False,
color='black', ls='solid', alpha=0.1)
ax.add_patch(unit_circle)
axvline(0, color='0.7')
axhline(0, color='0.7')
# Plot the poles and set marker properties
poles = plt.plot(p.real, p.imag, 'x', markersize=9, alpha=0.5)
# Plot the zeros and set marker properties
zeros = plt.plot(z.real, z.imag, 'o', markersize=9,
color='none', alpha=0.5,
markeredgecolor=poles[0].get_color(), # same color as poles
)
# Scale axes to fit
r = 1.5 * np.amax(np.concatenate((abs(z), abs(p), [1])))
plt.axis('scaled')
plt.axis([-r, r, -r, r])
# ticks = [-1, -.5, .5, 1]
# plt.xticks(ticks)
# plt.yticks(ticks)
"""
If there are multiple poles or zeros at the same point, put a
superscript next to them.
TODO: can this be made to self-update when zoomed?
"""
# Finding duplicates by same pixel coordinates (hacky for now):
poles_xy = ax.transData.transform(np.vstack(poles[0].get_data()).T)
zeros_xy = ax.transData.transform(np.vstack(zeros[0].get_data()).T)
# dict keys should be ints for matching, but coords should be floats for
# keeping location of text accurate while zooming
# TODO make less hacky, reduce duplication of code
d = defaultdict(int)
coords = defaultdict(tuple)
for xy in poles_xy:
key = tuple(np.rint(xy).astype('int'))
d[key] += 1
coords[key] = xy
for key, value in d.iteritems():
if value > 1:
x, y = ax.transData.inverted().transform(coords[key])
plt.text(x, y,
r' ${}^{' + str(value) + '}$',
fontsize=13,
)
d = defaultdict(int)
coords = defaultdict(tuple)
for xy in zeros_xy:
key = tuple(np.rint(xy).astype('int'))
d[key] += 1
coords[key] = xy
for key, value in d.iteritems():
if value > 1:
x, y = ax.transData.inverted().transform(coords[key])
plt.text(x, y,
r' ${}^{' + str(value) + '}$',
fontsize=13,
)
if filename is None:
plt.show()
else:
plt.savefig(filename)
print 'Pole-zero plot saved to ' + str(filename)
if __name__ == "__main__":
from scipy.signal import (freqz, butter, bessel, cheby1, cheby2, ellip,
tf2zpk, zpk2tf, lfilter, buttap, bilinear, cheb2ord, cheb2ap
)
from numpy import asarray, tan, array, pi, arange, cos, log10, unwrap, angle
from matplotlib.pyplot import (stem, title, grid, show, plot, xlabel,
ylabel, subplot, xscale, figure, xlim,
margins)
# # Cosine function
# omega = pi/4
# b = array([1.0, -cos(omega)])
# a = array([1, -2*cos(omega), 1.0])
b, a = butter(2, [0.06, 0.7], 'bandpass')
# Get the poles and zeros
z, p, k = tf2zpk(b, a)
# Create zero-pole plot
figure(figsize=(16, 9))
subplot(2, 2, 1)
zplane(z, p)
grid(True, color='0.9', linestyle='-', which='both', axis='both')
title('Poles and zeros')
# Display zeros, poles and gain
print str(len(z)) + " zeros: " + str(z)
print str(len(p)) + " poles: " + str(p)
print "gain: " + str(k)
# Impulse response
index = arange(0,20)
u = 1.0*(index==0)
y = lfilter(b, a, u)
subplot(2, 2, 3)
stem(index,y)
title('Impulse response')
margins(0, 0.1)
grid(True, color='0.9', linestyle='-', which='both', axis='both')
show()
# Frequency response
w, h = freqz(b, a)
subplot(2, 2, 2)
plot(w/pi, 20*log10(abs(h)))
xscale('log')
title('Frequency response')
xlabel('Normalized frequency')
ylabel('Amplitude [dB]')
margins(0, 0.1)
grid(True, color = '0.7', linestyle='-', which='major', axis='both')
grid(True, color = '0.9', linestyle='-', which='minor', axis='both')
show()
# Phase
subplot(2, 2, 4)
plot(w/pi, 180/pi * unwrap(angle(h)))
xscale('log')
xlabel('Normalized frequency')
ylabel('Phase [degrees]')
grid(True, color = '0.7', linestyle='-', which='major')
grid(True, color = '0.9', linestyle='-', which='minor')
show()
@endolith
Copy link
Author

bug: alpha only applies to X not O

@heetbeet
Copy link

Thanks for this gist. I was looking for an easy Matlab replacement and this will do just fine.

@rdpoor
Copy link

rdpoor commented Mar 13, 2020

@endolith: Thanks for the inspiration on superscripting multiple poles and zeros. Since the computations that create the poles and zeros tend to produce exact duplicates, I found it was sufficient to do the comparison on the direct poles and zeros without translating to and from pixel coordinates:

def plot_zpk(zeros, poles, k):
    plot_unit_circle()
    t1 = plt.plot(zeros.real, zeros.imag, 'o', markersize=10.0, alpha=0.5)
    t2 = plt.plot(poles.real, poles.imag, 'x', markersize=10.0, alpha=0.5)
    mark_overlapping(zeros)
    mark_overlapping(poles)

def mark_overlapping(items):
    """
    Given `items` as a list of complex coordinates, make a tally of identical 
    values, and, if there is more than one, plot a superscript on the graph.
    """
    d = defaultdict(int)
    for i in items:
        d[i] += 1
    for item, count in d.items():
        if count > 1:
            plt.text(item.real, item.imag, r' ${}^{' + str(count) + '}$', fontsize=13)

@rdpoor
Copy link

rdpoor commented Mar 13, 2020

e.g.:

plot_zpk(*scipy.signal.cheby1(8, 0.1, [5000*0.7, 5000/0.7], btype='bandstop', output='zpk', fs=48000))

image

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment