-
Notifications
You must be signed in to change notification settings - Fork 1.3k
/
ar.py
78 lines (66 loc) · 2.3 KB
/
ar.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
# Authors: Alexandre Gramfort <alexandre.gramfort@inria.fr>
# The statsmodels folks for AR yule_walker
#
# License: BSD-3-Clause
import numpy as np
from ..defaults import _handle_default
from ..io.pick import _picks_to_idx, _picks_by_type, pick_info
from ..utils import verbose, _apply_scaling_array
def _yule_walker(X, order=1):
"""Compute Yule-Walker (adapted from statsmodels).
Operates in-place.
"""
from scipy import linalg
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.]), np.concatenate(([1.], -coeffs))