MRG: Fix AR estimation

.travis.yml

@@ -6,7 +6,7 @@ sudo: false
 env:
     # Enable python 2 and python 3 builds
     # DEPS=full: build optional dependencies: pandas, nitime, statsmodels,
-    #            scikit-learn, patsy, nibabel pillow;
+    #            scikit-learn, nibabel pillow;
     #            in the case of Python 2, also mayavi, traits, pysurfer
     # DEPS=minimal: don't build optional dependencies; tests that require those
     #               dependencies are supposed to be skipped
@@ -51,7 +51,7 @@ install:
         export MNE_ROOT="${PWD}/minimal_cmds";
         export NEUROMAG2FT_ROOT="${PWD}/minimal_cmds/bin";
         source ${MNE_ROOT}/bin/mne_setup_sh;
-        conda install --yes --quiet $ENSURE_PACKAGES pandas$PANDAS scikit-learn$SKLEARN patsy h5py pillow;
+        conda install --yes --quiet $ENSURE_PACKAGES pandas$PANDAS scikit-learn$SKLEARN h5py pillow;
         pip install -q joblib nibabel;
         if [ "${PYTHON}" == "3.5" ]; then
           conda install --yes --quiet $ENSURE_PACKAGES ipython;

appveyor.yml

@@ -2,7 +2,7 @@ environment:
   global:
       PYTHON: "C:\\conda"
       MINICONDA_VERSION: "latest"
-      CONDA_DEPENDENCIES: "setuptools numpy scipy matplotlib scikit-learn nose mayavi pandas h5py PIL patsy pyside"
+      CONDA_DEPENDENCIES: "setuptools numpy scipy matplotlib scikit-learn nose mayavi pandas h5py PIL pyside"
   matrix:
       - PYTHON_VERSION: "2.7"
         PYTHON_ARCH: "64"

doc/whats_new.rst

@@ -20,7 +20,7 @@ Changelog
 BUG
 ~~~
 
-    - ...
+    - Fix computation of AR coefficients across channels in :func:`mne.time_frequency.fit_iir_model_raw` by `Eric Larson`_
 
 API
 ~~~

examples/time_frequency/plot_temporal_whitening.py

@@ -3,7 +3,7 @@
 Temporal whitening with AR model
 ================================
 
-This script shows how to fit an AR model to data and use it
+Here we fit an AR model to the data and use it
 to temporally whiten the signals.
 
 """
@@ -35,11 +35,11 @@
 picks = mne.pick_types(raw.info, meg='grad', exclude='bads')
 
 order = 5  # define model order
-picks = picks[:5]
+picks = picks[:1]
 
 # Estimate AR models on raw data
 b, a = fit_iir_model_raw(raw, order=order, picks=picks, tmin=60, tmax=180)
-d, times = raw[0, 1e4:2e4]  # look at one channel from now on
+d, times = raw[0, 10000:20000]  # look at one channel from now on
 d = d.ravel()  # make flat vector
 innovation = signal.convolve(d, a, 'valid')
 d_ = signal.lfilter(b, a, innovation)  # regenerate the signal

mne/time_frequency/ar.py

@@ -4,131 +4,39 @@
 # License: BSD (3-clause)
 
 import numpy as np
-from scipy.linalg import toeplitz
+from scipy import linalg
 
-from ..io.pick import pick_types
+from ..defaults import _handle_default
+from ..io.pick import _pick_data_channels, _picks_by_type, pick_info
 from ..utils import verbose
 
 
-# XXX : Back ported from statsmodels
+def _yule_walker(X, order=1):
+    """Compute Yule-Walker (adapted from statsmodels).
 
-def yule_walker(X, order=1, method="unbiased", df=None, inv=False,
-                demean=True):
+    Operates in-place.
     """
-    Estimate AR(p) parameters from a sequence X using Yule-Walker equation.
-
-    Unbiased or maximum-likelihood estimator (mle)
-
-    See, for example:
-
-    http://en.wikipedia.org/wiki/Autoregressive_moving_average_model
-
-    Parameters
-    ----------
-    X : array-like
-        1d array
-    order : integer, optional
-        The order of the autoregressive process.  Default is 1.
-    method : string, optional
-       Method can be "unbiased" or "mle" and this determines denominator in
-       estimate of autocorrelation function (ACF) at lag k. If "mle", the
-       denominator is n=X.shape[0], if "unbiased" the denominator is n-k.
-       The default is unbiased.
-    df : integer, optional
-       Specifies the degrees of freedom. If `df` is supplied, then it is
-       assumed the X has `df` degrees of freedom rather than `n`.  Default is
-       None.
-    inv : bool
-        If inv is True the inverse of R is also returned.  Default is False.
-    demean : bool
-        True, the mean is subtracted from `X` before estimation.
-
-    Returns
-    -------
-    rho
-        The autoregressive coefficients
-    sigma
-        TODO
-
-    """
-    # TODO: define R better, look back at notes and technical notes on YW.
-    # First link here is useful
-    # http://www-stat.wharton.upenn.edu/~steele/Courses/956/ResourceDetails/YuleWalkerAndMore.htm  # noqa
-    method = str(method).lower()
-    if method not in ["unbiased", "mle"]:
-        raise ValueError("ACF estimation method must be 'unbiased' or 'MLE'")
-    X = np.array(X, float)
-    if demean:
-        X -= X.mean()                  # automatically demean's X
-    n = df or X.shape[0]
-
-    if method == "unbiased":        # this is df_resid ie., n - p
-        def denom(k):
-            return n - k
-    else:
-        def denom(k):
-            return n
-    if X.ndim > 1 and X.shape[1] != 1:
-        raise ValueError("expecting a vector to estimate AR parameters")
+    assert X.ndim == 2
+    denom = X.shape[-1] - np.arange(order + 1)
     r = np.zeros(order + 1, np.float64)
-    r[0] = (X ** 2).sum() / denom(0)
-    for k in range(1, order + 1):
-        r[k] = (X[0:-k] * X[k:]).sum() / denom(k)
-    R = toeplitz(r[:-1])
-
-    rho = np.linalg.solve(R, r[1:])
+    for di, d in enumerate(X):
+        d -= d.mean()
+        r[0] += np.dot(d, d)
+        for k in range(1, order + 1):
+            r[k] += np.dot(d[0:-k], d[k:])
+    r /= denom * len(X)
+    rho = linalg.solve(linalg.toeplitz(r[:-1]), r[1:])
     sigmasq = r[0] - (r[1:] * rho).sum()
-    if inv:
-        return rho, np.sqrt(sigmasq), np.linalg.inv(R)
-    else:
-        return rho, np.sqrt(sigmasq)
-
-
-def ar_raw(raw, order, picks, tmin=None, tmax=None):
-    """Fit AR model on raw data
-
-    Fit AR models for each channels and returns the models
-    coefficients for each of them.
-
-    Parameters
-    ----------
-    raw : Raw instance
-        The raw data
-    order : int
-        The AR model order
-    picks : array-like of int
-        The channels indices to include
-    tmin : float
-        The beginning of time interval in seconds.
-    tmax : float
-        The end of time interval in seconds.
-
-    Returns
-    -------
-    coefs : array
-        Sets of coefficients for each channel
-    """
-    start, stop = None, None
-    if tmin is not None:
-        start = raw.time_as_index(tmin)[0]
-    if tmax is not None:
-        stop = raw.time_as_index(tmax)[0] + 1
-    data, times = raw[picks, start:stop]
-
-    coefs = np.empty((len(data), order))
-    for k, d in enumerate(data):
-        this_coefs, _ = yule_walker(d, order=order)
-        coefs[k, :] = this_coefs
-    return coefs
+    return rho, np.sqrt(sigmasq)
 
 
 @verbose
 def fit_iir_model_raw(raw, order=2, picks=None, tmin=None, tmax=None,
                       verbose=None):
     """Fits an AR model to raw data and creates the corresponding IIR filter
 
-    The computed filter is the average filter for all the picked channels.
-    The frequency response is given by:
+    The computed filter is fitted to data from all of the picked channels,
+    with frequency response given by the standard IIR formula:
 
     .. math::
 
@@ -157,9 +65,19 @@ def fit_iir_model_raw(raw, order=2, picks=None, tmin=None, tmax=None,
     a : ndarray
         Denominator filter coefficients
     """
+    from ..cov import _apply_scaling_array
+    start, stop = None, None
+    if tmin is not None:
+        start = raw.time_as_index(tmin)[0]
+    if tmax is not None:
+        stop = raw.time_as_index(tmax)[0] + 1
     if picks is None:
-        picks = pick_types(raw.info, meg=True, eeg=True)
-    coefs = ar_raw(raw, order=order, picks=picks, tmin=tmin, tmax=tmax)
-    mean_coefs = np.mean(coefs, axis=0)  # mean model across channels
-    a = np.concatenate(([1.], -mean_coefs))  # filter coefficients
-    return np.array([1.]), a
+        picks = _pick_data_channels(raw.info, exclude='bads')
+    data = raw[picks, start:stop][0]
+    # rescale data to similar levels
+    picks_list = _picks_by_type(pick_info(raw.info, picks))
+    scalings = _handle_default('scalings_cov_rank', None)
+    _apply_scaling_array(data, picks_list=picks_list, scalings=scalings)
+    # do the fitting
+    coeffs, _ = _yule_walker(data, order=order)
+    return np.array([1.]), np.concatenate(([1.], -coeffs))

mne/time_frequency/tests/test_ar.py

@@ -1,36 +1,51 @@
 import os.path as op
+
 import numpy as np
-from numpy.testing import assert_array_almost_equal
-from nose.tools import assert_true, assert_equal
+from numpy.testing import assert_array_almost_equal, assert_allclose
+from nose.tools import assert_equal
+from scipy.signal import lfilter
 
-from mne import io, pick_types
-from mne.time_frequency.ar import yule_walker, fit_iir_model_raw
-from mne.utils import requires_statsmodels, requires_patsy
+from mne import io
+from mne.time_frequency.ar import _yule_walker, fit_iir_model_raw
+from mne.utils import requires_statsmodels, run_tests_if_main
 
 
 raw_fname = op.join(op.dirname(__file__), '..', '..', 'io', 'tests', 'data',
                     'test_raw.fif')
 
 
-@requires_patsy
 @requires_statsmodels
 def test_yule_walker():
     """Test Yule-Walker against statsmodels."""
     from statsmodels.regression.linear_model import yule_walker as sm_yw
     d = np.random.randn(100)
     sm_rho, sm_sigma = sm_yw(d, order=2)
-    rho, sigma = yule_walker(d, order=2)
+    rho, sigma = _yule_walker(d[np.newaxis], order=2)
     assert_array_almost_equal(sm_sigma, sigma)
     assert_array_almost_equal(sm_rho, rho)
 
 
 def test_ar_raw():
     """Test fitting AR model on raw data."""
-    raw = io.read_raw_fif(raw_fname)
+    raw = io.read_raw_fif(raw_fname).crop(0, 2).load_data()
+    raw.pick_types(meg='grad')
     # pick MEG gradiometers
-    picks = pick_types(raw.info, meg='grad', exclude='bads')
-    picks = picks[:2]
-    tmin, tmax, order = 0, 10, 2
-    coefs = fit_iir_model_raw(raw, order, picks, tmin, tmax)[1][1:]
-    assert_equal(coefs.shape, (order,))
-    assert_true(0.9 < -coefs[0] < 1.1)
+    for order in (2, 5, 10):
+        coeffs = fit_iir_model_raw(raw, order)[1][1:]
+        assert_equal(coeffs.shape, (order,))
+        assert_allclose(-coeffs[0], 1., atol=0.5)
+    # let's make sure we're doing something reasonable: first, white noise
+    rng = np.random.RandomState(0)
+    raw._data = rng.randn(*raw._data.shape)
+    raw._data *= 1e-15
+    for order in (2, 5, 10):
+        coeffs = fit_iir_model_raw(raw, order)[1]
+        assert_allclose(coeffs, [1.] + [0.] * order, atol=2e-2)
+    # Now let's try pink noise
+    iir = [1, -1, 0.2]
+    raw._data = lfilter([1.], iir, raw._data)
+    for order in (2, 5, 10):
+        coeffs = fit_iir_model_raw(raw, order)[1]
+        assert_allclose(coeffs, iir + [0.] * (order - 2), atol=5e-2)
+
+run_tests_if_main()

mne/utils.py

@@ -1043,7 +1043,6 @@ def wrapper(func):
 requires_tvtk = partial(requires_module, name='TVTK',
                         call='from tvtk.api import tvtk')
 requires_statsmodels = partial(requires_module, name='statsmodels')
-requires_patsy = partial(requires_module, name='patsy')
 requires_pysurfer = partial(requires_module, name='PySurfer',
                             call='from surfer import Brain')
 requires_PIL = partial(requires_module, name='PIL',