针对pulse-transit的工具

dist/client/mne/time_frequency/ar.py

@@ -0,0 +1,76 @@
+# Authors: The MNE-Python contributors.
+# License: BSD-3-Clause
+# Copyright the MNE-Python contributors.
+
+import numpy as np
+from scipy import linalg
+
+from .._fiff.pick import _picks_by_type, _picks_to_idx, pick_info
+from ..defaults import _handle_default
+from ..utils import _apply_scaling_array, verbose
+
+
+def _yule_walker(X, order=1):
+    """Compute Yule-Walker (adapted from statsmodels).
+
+    Operates in-place.
+    """
+    assert X.ndim == 2
+    denom = X.shape[-1] - np.arange(order + 1)
+    r = np.zeros(order + 1, np.float64)
+    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()
+    return rho, np.sqrt(sigmasq)
+
+
+@verbose
+def fit_iir_model_raw(raw, order=2, picks=None, tmin=None, tmax=None, verbose=None):
+    r"""Fit an AR model to raw data and creates the corresponding IIR filter.
+
+    The computed filter is fitted to data from all of the picked channels,
+    with frequency response given by the standard IIR formula:
+
+    .. math::
+
+        H(e^{jw}) = \frac{1}{a[0] + a[1]e^{-jw} + ... + a[n]e^{-jnw}}
+
+    Parameters
+    ----------
+    raw : Raw object
+        An instance of Raw.
+    order : int
+        Order of the FIR filter.
+    %(picks_good_data)s
+    tmin : float
+        The beginning of time interval in seconds.
+    tmax : float
+        The end of time interval in seconds.
+    %(verbose)s
+
+    Returns
+    -------
+    b : ndarray
+        Numerator filter coefficients.
+    a : ndarray
+        Denominator filter coefficients.
+    """
+    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
+    picks = _picks_to_idx(raw.info, picks)
+    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.0]), np.concatenate(([1.0], -coeffs))