Aluode/PerceptionLabPortable
0
1"""TimeDelayingRidge class."""2 3# Authors: The MNE-Python contributors.4# License: BSD-3-Clause5# Copyright the MNE-Python contributors.6 7import numpy as np8from scipy import linalg9from scipy.signal import fftconvolve10from scipy.sparse.csgraph import laplacian11from sklearn.base import BaseEstimator, RegressorMixin12from sklearn.utils.validation import check_is_fitted13 14from ..cuda import _setup_cuda_fft_multiply_repeated15from ..filter import next_fast_len16from ..fixes import jit17from ..utils import ProgressBar, _check_option, logger, warn18from ._fixes import _check_n_features_3d, validate_data19 20 21def _compute_corrs(22 X, y, smin, smax, n_jobs=None, fit_intercept=False, edge_correction=True23):24 """Compute auto- and cross-correlations."""25 if fit_intercept:26 # We could do this in the Fourier domain, too, but it should27 # be a bit cleaner numerically to do it here.28 X_offset = np.mean(X, axis=0)29 y_offset = np.mean(y, axis=0)30 if X.ndim == 3:31 X_offset = X_offset.mean(axis=0)32 y_offset = np.mean(y_offset, axis=0)33 X = X - X_offset34 y = y - y_offset35 else:36 X_offset = y_offset = 0.037 if X.ndim == 2:38 assert y.ndim == 239 X = X[:, np.newaxis, :]40 y = y[:, np.newaxis, :]41 assert X.shape[:2] == y.shape[:2]42 len_trf = smax - smin43 len_x, n_epochs, n_ch_x = X.shape44 len_y, n_epochs_y, n_ch_y = y.shape45 assert len_x == len_y46 assert n_epochs == n_epochs_y47 48 n_fft = next_fast_len(2 * X.shape[0] - 1)49 50 _, cuda_dict = _setup_cuda_fft_multiply_repeated(51 n_jobs, [1.0], n_fft, "correlation calculations"52 )53 del n_jobs # only used to set as CUDA54 55 # create our Toeplitz indexer56 ij = np.empty((len_trf, len_trf), int)57 for ii in range(len_trf):58 ij[ii, ii:] = np.arange(len_trf - ii)59 x = np.arange(n_fft - 1, n_fft - len_trf + ii, -1)60 ij[ii + 1 :, ii] = x61 62 x_xt = np.zeros([n_ch_x * len_trf] * 2)63 x_y = np.zeros((len_trf, n_ch_x, n_ch_y), order="F")64 n = n_epochs * (n_ch_x * (n_ch_x + 1) // 2 + n_ch_x)65 logger.info(f"Fitting {n_epochs} epochs, {n_ch_x} channels")66 pb = ProgressBar(n, mesg="Sample")67 count = 068 pb.update(count)69 for ei in range(n_epochs):70 this_X = X[:, ei, :]71 # XXX maybe this is what we should parallelize over CPUs at some point72 X_fft = cuda_dict["rfft"](this_X, n=n_fft, axis=0)73 X_fft_conj = X_fft.conj()74 y_fft = cuda_dict["rfft"](y[:, ei, :], n=n_fft, axis=0)75 76 for ch0 in range(n_ch_x):77 for oi, ch1 in enumerate(range(ch0, n_ch_x)):78 this_result = cuda_dict["irfft"](79 X_fft[:, ch0] * X_fft_conj[:, ch1], n=n_fft, axis=080 )81 # Our autocorrelation structure is a Toeplitz matrix, but82 # it's faster to create the Toeplitz ourselves than use83 # linalg.toeplitz.84 this_result = this_result[ij]85 # However, we need to adjust for coeffs that are cut off,86 # i.e. the non-zero delays should not have the same AC value87 # as the zero-delay ones (because they actually have fewer88 # coefficients).89 #90 # These adjustments also follow a Toeplitz structure, so we91 # construct a matrix of what has been left off, compute their92 # inner products, and remove them.93 if edge_correction:94 _edge_correct(this_result, this_X, smax, smin, ch0, ch1)95 96 # Store the results in our output matrix97 x_xt[98 ch0 * len_trf : (ch0 + 1) * len_trf,99 ch1 * len_trf : (ch1 + 1) * len_trf,100 ] += this_result101 if ch0 != ch1:102 x_xt[103 ch1 * len_trf : (ch1 + 1) * len_trf,104 ch0 * len_trf : (ch0 + 1) * len_trf,105 ] += this_result.T106 count += 1107 pb.update(count)108 109 # compute the crosscorrelations110 cc_temp = cuda_dict["irfft"](111 y_fft * X_fft_conj[:, slice(ch0, ch0 + 1)], n=n_fft, axis=0112 )113 if smin < 0 and smax >= 0:114 x_y[:-smin, ch0] += cc_temp[smin:]115 x_y[len_trf - smax :, ch0] += cc_temp[:smax]116 else:117 x_y[:, ch0] += cc_temp[smin:smax]118 count += 1119 pb.update(count)120 121 x_y = np.reshape(x_y, (n_ch_x * len_trf, n_ch_y), order="F")122 return x_xt, x_y, n_ch_x, X_offset, y_offset123 124 125@jit()126def _edge_correct(this_result, this_X, smax, smin, ch0, ch1):127 if smax > 0:128 tail = _toeplitz_dot(this_X[-1:-smax:-1, ch0], this_X[-1:-smax:-1, ch1])129 if smin > 0:130 tail = tail[smin - 1 :, smin - 1 :]131 this_result[max(-smin + 1, 0) :, max(-smin + 1, 0) :] -= tail132 if smin < 0:133 head = _toeplitz_dot(this_X[:-smin, ch0], this_X[:-smin, ch1])[::-1, ::-1]134 if smax < 0:135 head = head[:smax, :smax]136 this_result[:-smin, :-smin] -= head137 138 139@jit()140def _toeplitz_dot(a, b):141 """Create upper triangular Toeplitz matrices & compute the dot product."""142 # This is equivalent to:143 # a = linalg.toeplitz(a)144 # b = linalg.toeplitz(b)145 # a[np.triu_indices(len(a), 1)] = 0146 # b[np.triu_indices(len(a), 1)] = 0147 # out = np.dot(a.T, b)148 assert a.shape == b.shape and a.ndim == 1149 out = np.outer(a, b)150 for ii in range(1, len(a)):151 out[ii, ii:] += out[ii - 1, ii - 1 : -1]152 out[ii + 1 :, ii] += out[ii:-1, ii - 1]153 return out154 155 156def _compute_reg_neighbors(n_ch_x, n_delays, reg_type, method="direct", normed=False):157 """Compute regularization parameter from neighbors."""158 known_types = ("ridge", "laplacian")159 if isinstance(reg_type, str):160 reg_type = (reg_type,) * 2161 if len(reg_type) != 2:162 raise ValueError(f"reg_type must have two elements, got {len(reg_type)}")163 for r in reg_type:164 if r not in known_types:165 raise ValueError(f"reg_type entries must be one of {known_types}, got {r}")166 reg_time = reg_type[0] == "laplacian" and n_delays > 1167 reg_chs = reg_type[1] == "laplacian" and n_ch_x > 1168 if not reg_time and not reg_chs:169 return np.eye(n_ch_x * n_delays)170 # regularize time171 if reg_time:172 reg = np.eye(n_delays)173 stride = n_delays + 1174 reg.flat[1::stride] += -1175 reg.flat[n_delays::stride] += -1176 reg.flat[n_delays + 1 : -n_delays - 1 : stride] += 1177 args = [reg] * n_ch_x178 reg = linalg.block_diag(*args)179 else:180 reg = np.zeros((n_delays * n_ch_x,) * 2)181 182 # regularize features183 if reg_chs:184 block = n_delays * n_delays185 row_offset = block * n_ch_x186 stride = n_delays * n_ch_x + 1187 reg.flat[n_delays:-row_offset:stride] += -1188 reg.flat[n_delays + row_offset :: stride] += 1189 reg.flat[row_offset:-n_delays:stride] += -1190 reg.flat[: -(n_delays + row_offset) : stride] += 1191 assert np.array_equal(reg[::-1, ::-1], reg)192 193 if method == "direct":194 if normed:195 norm = np.sqrt(np.diag(reg))196 reg /= norm197 reg /= norm[:, np.newaxis]198 return reg199 else:200 # Use csgraph. Note that our -1's above are really the neighbors!201 # If we ever want to allow arbitrary adjacency matrices, this is how202 # we'd want to do it.203 reg = laplacian(-reg, normed=normed)204 return reg205 206 207def _fit_corrs(x_xt, x_y, n_ch_x, reg_type, alpha, n_ch_in):208 """Fit the model using correlation matrices."""209 # do the regularized solving210 n_ch_out = x_y.shape[1]211 assert x_y.shape[0] % n_ch_x == 0212 n_delays = x_y.shape[0] // n_ch_x213 reg = _compute_reg_neighbors(n_ch_x, n_delays, reg_type)214 mat = x_xt + alpha * reg215 # From sklearn216 try:217 # Note: we must use overwrite_a=False in order to be able to218 # use the fall-back solution below in case a LinAlgError219 # is raised220 w = linalg.solve(mat, x_y, overwrite_a=False, assume_a="pos")221 except np.linalg.LinAlgError:222 warn(223 "Singular matrix in solving dual problem. Using "224 "least-squares solution instead."225 )226 w = linalg.lstsq(mat, x_y, lapack_driver="gelsy")[0]227 w = w.T.reshape([n_ch_out, n_ch_in, n_delays])228 return w229 230 231class TimeDelayingRidge(RegressorMixin, BaseEstimator):232 """Ridge regression of data with time delays.233 234 Parameters235 ----------236 tmin : int | float237 The starting lag, in seconds (or samples if ``sfreq`` == 1).238 Negative values correspond to times in the past.239 tmax : int | float240 The ending lag, in seconds (or samples if ``sfreq`` == 1).241 Positive values correspond to times in the future.242 Must be >= tmin.243 sfreq : float244 The sampling frequency used to convert times into samples.245 alpha : float246 The ridge (or laplacian) regularization factor.247 reg_type : str | list248 Can be ``"ridge"`` (default) or ``"laplacian"``.249 Can also be a 2-element list specifying how to regularize in time250 and across adjacent features.251 fit_intercept : bool252 If True (default), the sample mean is removed before fitting.253 n_jobs : int | str254 The number of jobs to use. Can be an int (default 1) or ``'cuda'``.255 256 .. versionadded:: 0.18257 edge_correction : bool258 If True (default), correct the autocorrelation coefficients for259 non-zero delays for the fact that fewer samples are available.260 Disabling this speeds up performance at the cost of accuracy261 depending on the relationship between epoch length and model262 duration. Only used if ``estimator`` is float or None.263 264 .. versionadded:: 0.18265 266 See Also267 --------268 mne.decoding.ReceptiveField269 270 Notes271 -----272 This class is meant to be used with :class:`mne.decoding.ReceptiveField`273 by only implicitly doing the time delaying. For reasonable receptive274 field and input signal sizes, it should be more CPU and memory275 efficient by using frequency-domain methods (FFTs) to compute the276 auto- and cross-correlations.277 """278 279 def __init__(280 self,281 tmin,282 tmax,283 sfreq,284 alpha=0.0,285 reg_type="ridge",286 fit_intercept=True,287 n_jobs=None,288 edge_correction=True,289 ):290 self.tmin = tmin291 self.tmax = tmax292 self.sfreq = sfreq293 self.alpha = alpha294 self.reg_type = reg_type295 self.fit_intercept = fit_intercept296 self.edge_correction = edge_correction297 self.n_jobs = n_jobs298 299 def __sklearn_tags__(self):300 """..."""301 tags = super().__sklearn_tags__()302 tags.input_tags.three_d_array = True303 tags.target_tags.one_d_labels = True304 tags.target_tags.two_d_labels = True305 tags.target_tags.multi_output = True306 return tags307 308 @property309 def _smin(self):310 return int(round(self.tmin_ * self.sfreq_))311 312 @property313 def _smax(self):314 return int(round(self.tmax_ * self.sfreq_)) + 1315 316 def _check_data(self, X, y=None, reset=False):317 if reset:318 X, y = validate_data(319 self,320 X=X,321 y=y,322 reset=reset,323 validate_separately=( # to take care of 3D y324 dict(allow_nd=True),325 dict(allow_nd=True, ensure_2d=False),326 ),327 )328 if X.ndim == 3:329 assert y.ndim == 3330 assert X.shape[:2] == y.shape[:2]331 else:332 if y.ndim == 1:333 y = y[:, np.newaxis]334 assert y.ndim == 2335 _check_option("y.shape[0]", y.shape[0], (X.shape[0],))336 else:337 X = validate_data(self, X=X, allow_nd=True, ensure_2d=False, reset=reset)338 # Because when ensure_2d=True, sklearn takes n_features from X.shape[1],339 # when we need X.shape[-1]. So we ensure 2D and check features ourselves.340 if X.ndim == 1:341 raise ValueError(342 "Reshape your data either using array.reshape(-1, 1) if "343 "your data has a single feature or array.reshape(1, -1) "344 "if it contains a single sample."345 )346 _check_n_features_3d(self, X, reset)347 348 return X, y349 350 def _validate_params(self, X):351 self.tmin_ = float(self.tmin)352 self.tmax_ = float(self.tmax)353 self.sfreq_ = float(self.sfreq)354 self.alpha_ = float(self.alpha)355 if self.tmin_ > self.tmax_:356 raise ValueError(f"tmin must be <= tmax, got {self.tmin_} and {self.tmax_}")357 n_delays = self._smax - self._smin358 min_samples = (n_delays + 1) // 2359 if X.shape[0] < min_samples:360 raise ValueError(361 f"Got n_samples = {X.shape[0]}, but at least {min_samples} "362 f"are required to estimate {n_delays} delays."363 )364 365 def fit(self, X, y):366 """Estimate the coefficients of the linear model.367 368 Parameters369 ----------370 X : array, shape (n_samples[, n_epochs], n_features)371 The training input samples to estimate the linear coefficients.372 y : array, shape (n_samples[, n_epochs], n_outputs)373 The target values.374 375 Returns376 -------377 self : instance of TimeDelayingRidge378 Returns the modified instance.379 """380 X, y = self._check_data(X, y, reset=True)381 self._validate_params(X)382 # These are split into two functions because it's possible that we383 # might want to allow people to do them separately (e.g., to test384 # different regularization parameters).385 self.cov_, x_y_, n_ch_x, X_offset, y_offset = _compute_corrs(386 X,387 y,388 self._smin,389 self._smax,390 self.n_jobs,391 self.fit_intercept,392 self.edge_correction,393 )394 self.coef_ = _fit_corrs(395 self.cov_, x_y_, n_ch_x, self.reg_type, self.alpha_, n_ch_x396 )397 # This is the sklearn formula from LinearModel (will be 0. for no fit)398 if self.fit_intercept:399 self.intercept_ = y_offset - np.dot(X_offset, self.coef_.sum(-1).T)400 else:401 self.intercept_ = 0.0402 return self403 404 def predict(self, X):405 """Predict the output.406 407 Parameters408 ----------409 X : array, shape (n_samples[, n_epochs], n_features)410 The data.411 412 Returns413 -------414 X : ndarray415 The predicted response.416 """417 check_is_fitted(self)418 X, _ = self._check_data(X)419 if X.ndim == 2:420 X = X[:, np.newaxis, :]421 singleton = True422 else:423 singleton = False424 out = np.zeros(X.shape[:2] + (self.coef_.shape[0],))425 smin = self._smin426 offset = max(smin, 0)427 for ei in range(X.shape[1]):428 for oi in range(self.coef_.shape[0]):429 for fi in range(self.coef_.shape[1]):430 temp = fftconvolve(X[:, ei, fi], self.coef_[oi, fi])431 temp = temp[max(-smin, 0) :][: len(out) - offset]432 out[offset : len(temp) + offset, ei, oi] += temp433 out += self.intercept_434 if singleton:435 out = out[:, 0, :]436 out = out.squeeze()437 return out438