Source code for mmf.signal.spectrum

r"""Tools for computing the spectrum of a time-series."""
from __future__ import division

import numpy as np

import scikits.talkbox

def extend(f, order, steps):
    r"""Use a `order+1` Linear Prediction Coefficients (LPC) to extend the
    data.

    Note
    ----
    The sign convention is opposite that of Numerical Recipes
    """
    a, _e, _k = scikits.talkbox.lpc(f, order=order)
    fnew = np.empty(len(f) + steps, dtype=f.dtype)
    fnew[:len(f)] = f
    for n in xrange(steps):
        fnew[n-steps] = -sum(fnew[n - steps - order:n - steps]*np.flipud(a[1:]))
    return fnew
    
def spectrum(f_t, dt, ws, order):
    r"""
    Parameters
    ----------
    f_t : array
        Tabulated `f(t)` on an equally spaced grid of times `t` with spacing
        `dt`
    `dt` : float
        Time-step
    `ws` : array
        Angular frequencies at which the results are returned. 
    `order` : int
        Order.
    """
    z = np.exp(1j*ws*dt)
    a, _e, _k = scikits.talkbox.lpc(f_t, order=order)
    c = np.flipud(1*a)
    c[-1] = 1.0
    return np.sqrt(a[0]+0j)/np.polyval(c, z)

def compare(f_t, dt, order):
    r"""Compare the spectrum with the FFT."""
    from matplotlib import pyplot as plt
    Nt = len(f_t)
    T = dt * Nt
    ws = 2*np.pi*np.arange(Nt//2)/T
    plt.plot(ws, abs(np.fft.fft(f_t)[:Nt//2]),':')
    plt.plot(ws, np.fft.fft(f_t)[:Nt//2].real,':')
    plt.plot(ws, np.fft.fft(f_t)[:Nt//2].imag,':')
    ws = np.linspace(0,10,1000)
    plt.plot(ws, abs(spectrum(f_t=f_t, dt=dt, ws=ws, order=order)))
    plt.plot(ws, spectrum(f_t=f_t, dt=dt, ws=ws, order=order).real)
    plt.plot(ws, spectrum(f_t=f_t, dt=dt, ws=ws, order=order).imag)