Pole-zero matching scratch work

pz_match.py

import numpy as np
from scipy.signal.filter_design import _cplxreal


def pop_real(p):
    """Pop the first real value in the list p."""
    k = 0
    while k < len(p) and p[k].imag != 0:
        k += 1
    if k == len(p):
        raise ValueError("no real values in p")
    pr = p.pop(k)
    return pr


def pop_pole_pair(p):
    """Pop a single complex value or a pair of real values from the list p."""
    p1 = p.pop(0)
    pole_pair = [p1]
    if p1.imag == 0:
        p2 = pop_real(p)
        pole_pair.append(p2)
    return pole_pair


def find_closest_zero_index(zeros, pole_pair, real=False):
    """
    Find the index in `zeros` of the value closest to `pole_pair`.

    If `real` is True, find the index in `zeros` of the *real* value
    closest to `pole_pair`.

    `pole_pair` must be a 1-d sequence of values.
    """
    z = np.array(zeros)
    if real:
        real_indices = np.where(z.imag == 0)[0]
        z = z[real_indices]
    dist = np.abs(z - np.array(pole_pair).reshape(-1, 1))
    k = np.unravel_index(dist.argmin(), dist.shape)[1]
    if real:
        k = real_indices[k]
    return k


def pole_zero_pairs(p, z):
    """
    `p` holds all the poles.
    `z` holds all the zeros.

    Assumptions:
    * len(z) <= len(p)
    * If there is a complex value in either sequence, the complex conjugate of
      that value must also be in the sequence.

    Returns a list of tuples.  Each tuple has length 2; it holds a matching
    of poles and zeros.  Each element in the tuple is a list.  If the list has
    length 2, it contains two real values (i.e. the imaginary part is 0).
    If the list has length 1, it is a complex value, and it means that value
    and its complex conjugate are part of the match.
    """
    if len(p) % 2:
        # Make the number of poles even.
        p = np.r_[p, 0]

    if len(z) < len(p):
        # Make the number of zeros equal to the number of poles.
        z = np.r_[z, np.zeros(len(p) - len(z))]

    p_c, p_r = _cplxreal(p)
    p = np.r_[p_c, p_r]
    # Sort p in order of increasing distance from the unit circle. 
    p = p[np.argsort(np.abs(np.abs(p) - 1))]

    z_c, z_r = _cplxreal(z)
    z = np.r_[z_c, z_r]

    # Use python lists for the convenience of the `pop` method.
    poles = p.tolist()
    zeros = z.tolist()

    pole_pairs = []
    zero_pairs = []
    while len(poles) > 0:
        pole_pair = pop_pole_pair(poles)
        pole_pairs.append(pole_pair)

        z1index = find_closest_zero_index(zeros, pole_pair)
        z1 = zeros.pop(z1index)
        zero_pair = [z1]
        if z1.imag == 0:
            z2index = find_closest_zero_index(zeros, pole_pair, real=True)
            z2 = zeros.pop(z2index)
            zero_pair.append(z2)
        zero_pairs.append(zero_pair)

    return zip(pole_pairs, zero_pairs)


if __name__ == "__main__":
    from scipy.signal import ellip, butter
    import matplotlib.pyplot as plt

    #p = np.array([0.1, 1+1j, 0.75, 0.5+0.25j, 0.25, 0])
    #z = np.array([0.25, 0.9, -.5+0.5j, 0, 0.25+0.25j, -0.5])

    #p = np.array([ 0.99+0.01j,  0.99-0.01j,  0.10+0.9j ,  0.10-0.9j ])
    #z = np.array([ 1.+0.j ,  1.+0.j ,  0.+0.9j,  0.-0.9j])

    order = 7
    z, p, k = butter(order, 0.2, output='zpk')

    #order = 13
    #Wn = 0.125
    #z, p, k = ellip(order, 0.087, 40, Wn, output='zpk')

    np.set_printoptions(precision=4)
    print "poles:", p
    print "zeros:", z
    pairs = pole_zero_pairs(p, z)
    print
    for pair in pairs:
        print "%25s  %25s" % (np.array(pair[0]), np.array(pair[1]))

    doplot = True
    if doplot:
        colors = ['b', 'g', 'r', 'c', 'm', 'y']
        alphas = [1.0, 0.5, 0.25]
        for i, (pole_pair, zero_pair) in enumerate(pairs):
            if len(pole_pair) == 1:
                p1, p2 = pole_pair[0], pole_pair[0].conjugate()
            else:
                p1, p2 = pole_pair
            if len(zero_pair) == 1:
                z1, z2 = zero_pair[0], zero_pair[0].conjugate()
            else:
                z1, z2 = zero_pair
            clr = i % len(colors)
            a = i // len(colors)
            plt.plot(p1.real, p1.imag, 'x', color=colors[clr], alpha=alphas[a], ms=10, mew=2)
            plt.plot(p2.real, p2.imag, 'x', color=colors[clr], alpha=alphas[a], ms=10, mew=2)
            plt.plot(z1.real, z1.imag, 'o', color=colors[clr], alpha=alphas[a], ms=10)
            plt.plot(z2.real, z2.imag, 'o', color=colors[clr], alpha=alphas[a], ms=10)

        t = np.linspace(0, 2*np.pi, 401)
        ux = np.cos(t)
        uy = np.sin(t)
        plt.plot(ux, uy, 'k--')
        plt.axis('equal')
        plt.grid()
        plt.show()