Example #1
0
def arma_estimate(X, P, Q, lag):
    """Autoregressive and moving average estimators.

    This function provides an estimate of the autoregressive 
    parameters, the moving average parameters, and the driving
    white noise variance of  an ARMA(P,Q) for a complex or real data sequence.

    The parameters are estimated using three steps:    

        * Estimate the AR parameters from the original data based on a least 
          squares modified Yule-Walker technique,  
        * Produce a residual time sequence by filtering the original data 
          with a filter based on the AR parameters, 
        * Estimate the MA parameters from the residual time sequence.

    :param array X: Array of data samples (length N)
    :param int P: Desired number of AR parameters
    :param int Q: Desired number of MA parameters
    :param int lag: Maximum lag to use for autocorrelation estimates

    :return:
        * A     - Array of complex P AR parameter estimates
        * B     - Array of complex Q MA parameter estimates
        * RHO   - White noise variance estimate

    .. note::
      *  lag must be >= Q (MA order)

    **dependencies**: 
        * :meth:`spectrum.correlation.CORRELATION`
        * :meth:`spectrum.covar.arcovar`
        * :meth:`spectrum.arma.ma`

    .. plot::
        :width: 80%
        :include-source:

        from spectrum import *
        from pylab import *
        a,b, rho = arma_estimate(marple_data, 15, 15, 30)
        psd = arma2psd(A=a, B=b, rho=rho, sides='centerdc', norm=True)
        plot(10 * log10(psd))
        ylim([-50,0])

    :reference: [Marple]_
    """
    R = CORRELATION(X, maxlags=lag, norm='unbiased')
    R0 = R[0]
    # C   Estimate the AR parameters (no error weighting is used).
    # C   Number of equation errors is M-Q .
    MPQ = lag - Q + P

    N = len(X)
    Y = zeros(N - P, dtype=complex)

    for K in range(0, MPQ):
        KPQ = K + Q - P + 1
        if KPQ < 0:
            Y[K] = R[-KPQ].conjugate()
        if KPQ == 0:
            Y[K] = R0
        if KPQ > 0:
            Y[K] = R[KPQ]

    # The resize is very important for the normalissation.
    Y.resize(lag)
    if P <= 4:
        res = arcovar_marple(Y.copy(), P)  # ! Eq. (10.12)
        ar_params = res[0]
    else:
        res = arcovar(Y.copy(), P)  # ! Eq. (10.12)
        ar_params = res[0]

    # the .copy is used to prevent a reference somewhere. this is a bug
    # to be tracked down.
    Y.resize(N - P)

    # C   Filter the original time series
    for k in range(P, N):
        SUM = X[k]
        # SUM += sum([ar_params[j]*X[k-j-1] for j in range(0,P)])
        for j in range(0, P):
            SUM = SUM + ar_params[j] * X[k - j - 1]  # ! Eq. (10.17)
        Y[k - P] = SUM

    # Estimate the MA parameters (a "long" AR of order at least 2*IQ
    # C   is suggested)
    # Y.resize(N-P)
    ma_params, rho = ma(Y, Q, 2 * Q)  # ! Eq. (10.3)

    return ar_params, ma_params, rho
Example #2
0
def CORRELOGRAMPSD(X, Y=None, lag=-1, window='hamming',
                    norm='unbiased', NFFT=4096, window_params={},
                    correlation_method='xcorr'):
    """PSD estimate using correlogram method.


    :param array X: complex or real data samples X(1) to X(N)
    :param array Y: complex data samples Y(1) to Y(N). If provided, computes
        the cross PSD, otherwise the PSD is returned
    :param int lag: highest lag index to compute. Must be less than N
    :param str window_name: see :mod:`window` for list of valid names
    :param str norm: one of the valid normalisation of :func:`xcorr` (biased, 
        unbiased, coeff, None)
    :param int NFFT: total length of the final data sets (padded with zero 
        if needed; default is 4096)
    :param str correlation_method: either `xcorr` or `CORRELATION`.
        CORRELATION should be removed in the future.

    :return:
        * Array of real (cross) power spectral density estimate values. This is
          a two sided array with negative values following the positive ones
          whatever is the input data (real or complex).

    .. rubric:: Description:

    The exact power spectral density is the Fourier transform of the
    autocorrelation sequence:

    .. math:: P_{xx}(f) = T \sum_{m=-\infty}^{\infty} r_{xx}[m] exp^{-j2\pi fmT}

    The correlogram method of PSD estimation substitutes a finite sequence of
    autocorrelation estimates :math:`\hat{r}_{xx}` in place of :math:`r_{xx}`.
    This estimation can be computed with :func:`xcorr` or :func:`CORRELATION` by
    chosing a proprer lag `L`. The estimated PSD is then

    .. math:: \hat{P}_{xx}(f) = T \sum_{m=-L}^{L} \hat{r}_{xx}[m] exp^{-j2\pi fmT}

    The lag index must be less than the number of data samples `N`. Ideally, it
    should be around `L/10` [Marple]_ so as to avoid greater statistical
    variance associated with higher lags.

    To reduce the leakage of the implicit rectangular window and therefore to
    reduce the bias in the estimate, a tapering window is normally used and lead
    to the so-called Blackman and Tukey correlogram:

    .. math:: \hat{P}_{BT}(f) = T \sum_{m=-L}^{L} w[m] \hat{r}_{xx}[m] exp^{-j2\pi fmT}

    The correlogram for the cross power spectral estimate is

    .. math:: \hat{P}_{xx}(f) = T \sum_{m=-L}^{L} \hat{r}_{xx}[m] exp^{-j2\pi fmT}

    which is computed if :attr:`Y` is not provide. In such case,
    :math:`r_{yx} = r_{xy}` so we compute the correlation only once.

    .. plot::
        :width: 80%
        :include-source:

        from spectrum import CORRELOGRAMPSD, marple_data
        from spectrum.tools import cshift
        from pylab import log10, axis, grid, plot,linspace

        psd = CORRELOGRAMPSD(marple_data, marple_data, lag=15)
        f = linspace(-0.5, 0.5, len(psd))
        psd = cshift(psd, len(psd)/2)
        plot(f, 10*log10(psd/max(psd)))
        axis([-0.5,0.5,-50,0])
        grid(True)

    .. seealso:: :func:`create_window`, :func:`CORRELATION`, :func:`xcorr`,
        :class:`pcorrelogram`.
    """
    N = len(X)
    assert lag<N, 'lag must be < size of input data'
    assert correlation_method in ['CORRELATION', 'xcorr']
    if Y is None:
        Y = numpy.array(X)
        crosscorrelation = False
    else:
        crosscorrelation = True

    if NFFT is None:
        NFFT = N
    psd = numpy.zeros(NFFT, dtype=complex)

    # Window should be centered around zero. Moreover, we want only the
    # positive values. So, we need to use 2*lag + 1 window and keep values on
    # the right side.
    w = Window(2.*lag+1, window, **window_params)
    w = w.data[lag+1:]

    # compute the cross correlation
    if correlation_method == 'CORRELATION':
        rxy = CORRELATION (X, Y, maxlags=lag, norm=norm)
    elif correlation_method == 'xcorr':
        rxy, _l = xcorr (X, Y, maxlags=lag, norm=norm)
        rxy = rxy[lag:]

    # keep track of the first elt.
    psd[0] = rxy[0]

    # create the first part of the PSD
    psd[1:lag+1] = rxy[1:] * w

    # create the second part.
    # First, we need to compute the auto or cross correlation ryx
    if crosscorrelation is True:
        # compute the cross correlation
        if correlation_method == 'CORRELATION':
            ryx = CORRELATION(Y, X, maxlags=lag, norm=norm)
        elif correlation_method == 'xcorr':
            ryx, _l = xcorr(Y, X, maxlags=lag, norm=norm)
            ryx = ryx[lag:]
        #print len(ryx), len(psd[-1:NPSD-lag-1:-1])

        psd[-1:NFFT-lag-1:-1] = ryx[1:].conjugate() * w
    else: #autocorrelation no additional correlation call required
        psd[-1:NFFT-lag-1:-1] = rxy[1:].conjugate() * w

    psd = numpy.real(fft(psd))

    return psd