pole and zero pairing attempt

pairing.py

# -*- coding: utf-8 -*-
"""
Created on Thu Nov 13 21:36:42 2014


"""

from __future__ import division, print_function
import numpy as np
from numpy import zeros, abs, argmin, conj, argsort, concatenate, delete
from scipy import signal
from scipy.signal import zpk2tf

import matplotlib.pyplot as plt
from matplotlib import patches
from collections import defaultdict
from scipy.signal import lti


def pzplot(system, *args, **kwargs):
    """
    Plot the poles and zeros of the given system on the complex plane, with
    real part on horizontal axis and imaginary part on vertical axis.

    `system` is a `scipy.signal.lti` object or a tuple of coefficients in ba
    or zpk form representing a digital or analog system.

    For example, to plot a list of poles p and zeros z, call
    ``pzplot((z, p, 1))``.

    kwargs are passed to `pyplot.plot`, so you can set color, markersize, etc.
    """

    if not isinstance(system, lti):
        system = lti(*system)

    z = system.zeros
    p = system.poles

    # get a figure/plot
    # http://stackoverflow.com/q/13831549/125507#comment19037056_13831816
    ax = kwargs.pop('ax', plt.gca())

    # Get arguments or use defaults
    color = kwargs.pop('color', ax._get_lines.color_cycle.next())
    size = kwargs.pop('markersize', 9)

    # Add unit circle and zero axes
    unit_circle = patches.Circle(xy=(0, 0), radius=1, fill=False,
                                 color='black', ls='solid', alpha=0.1)
    ax.add_patch(unit_circle)
    ax.axvline(x=0, color='black', alpha=0.3)
    ax.axhline(y=0, color='black', alpha=0.3)

    # Plot the poles and set marker properties
    poles = ax.plot(p.real, p.imag, 'x', markersize=size, alpha=0.5,
                    color=color,
                    *args, **kwargs
                    )

    color = poles[0].get_color()

    # Plot the zeros and set marker properties
    zeros = ax.plot(z.real, z.imag, 'o', markersize=size,
                    markerfacecolor='none', alpha=0.5,
                    markeredgecolor=color,  # same as poles
                    # markerfacecolor means the color= argument is ignored,
                    # but set it anyway or else it will use up an entry from
                    # the color cycle and every other color will be skipped
                    color=color,
                    *args, **kwargs
                    )

    # Scale axes to fit
    r = 1.15 * np.amax(np.concatenate((abs(z), abs(p), [1])))
    ax.axis('scaled')
    ax.axis([-r, r, -r, r])

    # 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
    for roots_xy in (zeros_xy, poles_xy):
        d = defaultdict(int)
        coords = defaultdict(tuple)
        for xy in roots_xy:
            key = tuple(np.rint(xy).astype('int'))
            d[key] += 1
            coords[key] = xy
        for key, value in d.iteritems():
            if value > 1:
                print(coords[key])
                x, y = ax.transData.inverted().transform(coords[key])
                ax.text(x, y, r' ${}^{' + str(value) + '}$',
                        fontsize=13, color=color
                        )
    plt.draw()



def _cplxreal(z, tol=None):
    """
    Split into complex and real parts, combining conjugate pairs.

    The 1D input vector `z` is split up into its complex (`zc`) and real (`zr`)
    elements.  Every complex element must be part of a complex-conjugate pair,
    which are combined into a single number (with positive imaginary part) in
    the output.  Two complex numbers are considered a conjugate pair if their
    real and imaginary parts differ in magnitude by less than ``tol * abs(z)``.

    Parameters
    ----------
    z : array_like
        Vector of complex numbers to be sorted and split
    tol : float, optional
        Relative tolerance for testing realness and conjugate equality.
        Default is ``100 * spacing(1)`` of `z`'s data type (i.e. 2e-14 for
        float64)

    Returns
    -------
    zc : ndarray
        Complex elements of `z`, with each pair represented by a single value
        having positive imaginary part, sorted first by real part, and then
        by magnitude of imaginary part.  The pairs are averaged when combined
        to reduce error.
    zr : ndarray
        Real elements of `z` (those having imaginary part less than
        `tol` times their magnitude), sorted by value.

    Raises
    ------
    ValueError
        If there are any complex numbers in `z` for which a conjugate
        cannot be found.

    See Also
    --------
    _cplxpair

    Examples
    --------
    >>> a = [4, 3, 1, 2-2j, 2+2j, 2-1j, 2+1j, 2-1j, 2+1j, 1+1j, 1-1j]
    >>> zc, zr = _cplxreal(a)
    >>> print zc
    [ 1.+1.j  2.+1.j  2.+1.j  2.+2.j]
    >>> print zr
    [ 1.  3.  4.]
    """

    z = atleast_1d(z)
    if z.size == 0:
        return z, z
    elif z.ndim != 1:
        raise ValueError('_cplxreal only accepts 1D input')

    if tol is None:
        # Get tolerance from dtype of input
        tol = 100 * np.finfo((1.0 * z).dtype).eps

    # Sort by real part, magnitude of imaginary part (speed up further sorting)
    z = z[np.lexsort((abs(z.imag), z.real))]

    # Split reals from conjugate pairs
    real_indices = abs(z.imag) <= tol * abs(z)
    zr = z[real_indices].real

    if len(zr) == len(z):
        # Input is entirely real
        return array([]), zr

    # Split positive and negative halves of conjugates
    z = z[~real_indices]
    zp = z[z.imag > 0]
    zn = z[z.imag < 0]

    if len(zp) != len(zn):
        raise ValueError('Array contains complex value with no matching '
                         'conjugate.')

    # Find runs of (approximately) the same real part
    same_real = np.diff(zp.real) <= tol * abs(zp[:-1])
    diffs = numpy.diff(concatenate(([0], same_real, [0])))
    run_starts = numpy.where(diffs > 0)[0]
    run_stops = numpy.where(diffs < 0)[0]

    # Sort each run by their imaginary parts
    for i in range(len(run_starts)):
        start = run_starts[i]
        stop = run_stops[i] + 1
        for chunk in (zp[start:stop], zn[start:stop]):
            chunk[...] = chunk[np.lexsort([abs(chunk.imag)])]

    # Check that negatives match positives
    if any(abs(zp - zn.conj()) > tol * abs(zn)):
        raise ValueError('Array contains complex value with no matching '
                         'conjugate.')

    # Average out numerical inaccuracy in real vs imag parts of pairs
    zc = (zp + zn.conj()) / 2

    return zc, zr


def _cplxpair(z, tol=None):
    """
    Sort into pairs of complex conjugates.

    Complex conjugates in `z` are sorted by increasing real part.  In each
    pair, the number with negative imaginary part appears first.

    If pairs have identical real parts, they are sorted by increasing
    imaginary magnitude.

    Two complex numbers are considered a conjugate pair if their real and
    imaginary parts differ in magnitude by less than ``tol * abs(z)``.  The
    pairs are forced to be exact complex conjugates by averaging the positive
    and negative values.

    Purely real numbers are also sorted, but placed after the complex
    conjugate pairs.  A number is considered real if its imaginary part is
    smaller than `tol` times the magnitude of the number.

    Parameters
    ----------
    z : array_like
        1-dimensional input array to be sorted.
    tol : float, optional
        Relative tolerance for testing realness and conjugate equality.
        Default is ``100 * spacing(1)`` of `z`'s data type (i.e. 2e-14 for
        float64)

    Returns
    -------
    y : ndarray
        Complex conjugate pairs followed by real numbers.

    Raises
    ------
    ValueError
        If there are any complex numbers in `z` for which a conjugate
        cannot be found.

    See Also
    --------
    _cplxreal

    Examples
    --------
    >>> a = [4, 3, 1, 2-2j, 2+2j, 2-1j, 2+1j, 2-1j, 2+1j, 1+1j, 1-1j]
    >>> z = _cplxpair(a)
    >>> print(z)
    [ 1.-1.j  1.+1.j  2.-1.j  2.+1.j  2.-1.j  2.+1.j  2.-2.j  2.+2.j  1.+0.j
      3.+0.j  4.+0.j]
    """

    z = atleast_1d(z)
    if z.size == 0 or np.isrealobj(z):
        return np.sort(z)

    if z.ndim != 1:
        raise ValueError('z must be 1-dimensional')

    zc, zr = _cplxreal(z, tol)

    # Interleave complex values and their conjugates, with negative imaginary
    # parts first in each pair
    zc = np.dstack((zc.conj(), zc)).flatten()
    z = np.append(zc, zr)
    return z





































z, p, k = signal.ellip(7, 1, 40, [0.2, 0.7], 'bp', False, 'zpk') # 2 real zeros
# z, p, k = signal.ellip(7, 1, 40, [0.2, 0.7], 'bs', False, 'zpk') # 2 real poles
# z, p, k = signal.ellip(7, 1, 40, 0.2, 'lp', False, 'zpk') # 1 real pole, 1 real zero
# z, p, k = signal.ellip(7, 1, 40, 0.2, 'hp', False, 'zpk') # 1 real pole, 1 real zero
# z, p, k = signal.ellip(6, 1, 40, [0.2, 0.7], 'bp', False, 'zpk')
# z, p, k = signal.ellip(6, 1, 40, [0.2, 0.7], 'bs', False, 'zpk')
# z, p, k = signal.ellip(3, 1, 40, [0.2, 0.7], 'bp', False, 'zpk')














"""
pseudocode:

Sort complex poles by Q (or distance from unit circle)
For each complex pole
    find the nearest zero (complex or real)
    If the nearest zero is complex pair
        pair them
        remove zero pair from future consideration
        move on to the next complex pole
    if the nearest zero is real
        if there are no zeros left:
            pair with zeros at origin(?)
        else:
            find the second-nearest real zero
            pair 2 real zeros with 1 complex pole pair
            remove 2 real zeros from future consideration
            move on to the next complex pole
    if all zeros (both complex and real) have been used up
        pair the remaining complex poles with zeros at origin?
otherwise all the complex poles have been used up. start pairing real poles:

Sort real poles by their distance from unit circle (all have the same Q)
for all real poles:
    find the nearest zero (complex or real)
    If the nearest zero is complex pair
        pair them
        remove zero pair from future consideration
        move on to the next complex pole
    if the nearest zero is real
        if there are no zeros left:
            pair with zeros at origin
        else:
            find the second-nearest real zero
            pair 2 real zeros with 1 complex pole pair
            remove 2 real zeros from future consideration
            move on to the next complex pole
    if all zeros (both complex and real) have been used up
        pair the remaining real poles with zeros at origin
    if there are still zeros left:
        ???
"""









#############################################
# def zpk2sos


p = np.array(p)
z = np.array(z)
if len(z) > len(p):
    raise ValueError('Cannot have more zeros than poles')
n_sections = (len(p) + 1) // 2
sos = zeros((n_sections, 6))

# Separate complex/real poles/zeros
p_c, p_r = _cplxreal(p)
z_c, z_r = _cplxreal(z)

current_section = 0

plt.subplot(1, 2, 1)
pzplot((z, p, k))
plt.subplot(1, 2, 2)


# Sort poles by distance from unit circle (approximation of Q factor)
p_c = p_c[argsort(abs(1 - abs(p_c)))]
p_r = p_r[argsort(    1 - abs(p_r)) ]

# TODO: could sort by actual Q factor, but what's the formula in the Z plane?
# Is that even actually better?
# Q = -sqrt(log(abs(p_c))**2 + angle(p_c)**2)/(2*log(abs(p_c))) ???


# Match up complex poles
for pole in p_c:
    # find closest zero, whether complex or real
    closest_i = argmin(abs(pole - concatenate((z_c, z_r))))

    if closest_i < len(z_c):
        # Complex zero is closer than real zero
        zero = z_c[closest_i]
        section = zpk2tf([zero, conj(zero)], [pole, conj(pole)], 1)
        z_c = delete(z_c, closest_i)
    else:
        # Real zero is closer
        closest_i = closest_i - len(z_c)  # Undo concatenation of z_c, z_r
        zero_1 = z_r[closest_i]

        if len(z_r) == 1:
            # There are no more real zeros to pair with
            zero_2 = 0
        else:
            z_r = delete(z_r, closest_i)
            closest_i = argmin(abs(pole - z_r))
            zero_2 = z_r[closest_i]

        section = zpk2tf([zero_1, zero_2], [pole, conj(pole)], 1)
        z_r = delete(z_r, closest_i)

    sos[current_section] = concatenate(section)
    current_section += 1

    print(pole, zero)
    pzplot(section, markeredgewidth=2)

    if len(z_r) == 0 and len(z_c) == 0:
        # Ran out of zeros to pair with poles
        break
        for pole in p_c:
            section = zpk2tf([0, 0], [pole, conj(pole)], 1)
            sos[current_section] = concatenate(section)
            current_section += 1

# If odd number of real poles, add another at origin:
if len(p_r) % 2:
    p_r = np.append(p_r, 0)

# Match up pairs of real poles
for pole_1, pole_2 in p_r.reshape(-1, 2):
    # find closest zero, whether complex or real
    closest_i = argmin(abs(pole_1 - concatenate((z_c, z_r))))

    if closest_i < len(z_c):
        # Complex zero is closer than real zero
        zero = z_c[closest_i]
        section = zpk2tf([zero, conj(zero)], [pole_1, pole_2], 1)
        z_c = delete(z_c, closest_i)
    else:
        # Real zero is closer
        closest_i = closest_i - len(z_c)  # Undo concatenation of z_c, z_r
        zero_1 = z_r[closest_i]

        if len(z_r) == 1:
            # There are no more real zeros to pair with
            zero_2 = 0
        else:
            z_r = delete(z_r, closest_i)
            closest_i = argmin(abs(pole_1 - z_r))
            zero_2 = z_r[closest_i]

        section = zpk2tf([zero_1, zero_2], [pole_1, pole_2], 1)
        z_r = delete(z_r, closest_i)

    sos[current_section] = concatenate(section)
    current_section += 1

    print(pole, zero)
    pzplot(section, markeredgewidth=2)

    if len(z_r) == 0 and len(z_c) == 0:
        # Ran out of zeros to pair with poles
        break
        for pole in p_c:
            section = zpk2tf([0, 0], [pole_1, pole_2], 1)
            sos[current_section] = concatenate(section)
            current_section += 1

assert current_section == n_sections
# TODO HANDLE K
# TODO combine these into one loop that uses pole_2 = conj(pole_1)