arma.py 文件源码

def arma_predictor_nonlinear(x, y, m, n, x_hat0=None):
    """
    Calculate the nonlinear fit between the (*m*, *n*) ARMA model and
    the input *x* and output *y*. The optimization starts at *x_hat*
    (a vector with all 0s when `None`). The output is the tuple of the
    *m* AR and *n* MA coefficients.
    """
    if x_hat0 is None:
        x_hat0 = NP.zeros(m + n)
    (x_hat,
     cov_x,
     info,
     mesg,
     ier) = SP.optimize.leastsq(residual,
                                x_hat0,
                                args=(m, x, y),
                                Dfun=Dfun,
                                full_output=True)
    if ier not in [1, 2, 3, 4]:
        raise RuntimeError('optimization failed (ier={}) --- {}'.fomat(ier,
                                                                       mesg))
    a_hat = x_hat[:m]
    b_hat = x_hat[m:]
    a_hat = NP.insert(a_hat, 0, 1)
    return a_hat, b_hat