diff --git a/pyeeg/models.py b/pyeeg/models.py index da0ce24..177bc24 100644 --- a/pyeeg/models.py +++ b/pyeeg/models.py @@ -192,8 +192,9 @@ def fromArray(arr, tmin, tmax, fs): return trf def __init__(self, times=(0.,), tmin=None, tmax=None, srate=1., - alpha=0., fit_intercept=True, verbose=True): - + alpha=0., fit_intercept=True, verbose=True, + quadratic_reg=None): + # if tmin and tmax: # LOGGER.info("Will use lags spanning form tmin to tmax.\nTo use individual lags, use the `times` argument...") # self.lags = lag_span(tmin, tmax, srate=srate)[::-1] #pylint: disable=invalid-unary-operand-type @@ -202,7 +203,7 @@ def __init__(self, times=(0.,), tmin=None, tmax=None, srate=1., # else: # self.times = np.asarray(times) # self.lags = lag_sparse(self.times, srate)[::-1] - + self.tmin = tmin self.tmax = tmax self.times = times @@ -227,6 +228,8 @@ def __init__(self, times=(0.,), tmin=None, tmax=None, srate=1., # The two following are only defined if simple least-square (no reg.) is used self.tvals_ = None self.pvals_ = None + # Quadratic regularization + self.quadratic_reg = quadratic_reg # Can be None, a matrix, or a string ('smoothness', 'laplacian') def fill_lags(self): """Fill the lags attributes. @@ -333,9 +336,26 @@ def fit(self, X, y, lagged=False, drop=True, feat_names=(), rotations=()): # Solving with svd or least square: if self.verbose: LOGGER.info("Computing coefficients..") + + # Build quadratic regularization matrix M if specified + M = None + if self.quadratic_reg is not None: + if isinstance(self.quadratic_reg, str): + # Create matrix from string spec (e.g., 'smoothness', 'laplacian') + from pyeeg.solvers import create_quadratic_regularizer + n_lags = len(self.lags) + n_feats = self.n_feats_ + # Create block-diagonal matrix for per-feature smoothness + L_single = create_quadratic_regularizer(self.quadratic_reg, n_lags) + # Repeat for each feature (block diagonal) + M = np.kron(np.eye(n_feats), L_single) + else: + # Assume it's already a matrix + M = self.quadratic_reg + if self.use_regularisation or np.ndim(y) == 3: # svd method: - betas = _svd_regress(X, y, self.alpha, self.verbose) + betas = _svd_regress(X, y, self.alpha, M=M, verbose=self.verbose) self.all_betas = betas # Storing only the first as the main betas = betas[..., 0] @@ -359,6 +379,22 @@ def fit(self, X, y, lagged=False, drop=True, feat_names=(), rotations=()): self.coef_ = self.coef_[::-1, :, :] # need to flip the first axis of array to get correct lag order self.fitted = True + # Compute standardized coefficients (beta * std(X) / std(y)) + # This gives the change in y (in SDs) for a 1 SD change in X + if not lagged and self.fitted: + try: + X_std = np.std(X[:, self.fit_intercept:], axis=0, keepdims=True) # (1, n_lags*n_feats) + y_std = np.std(y, axis=0, keepdims=True) # (1, n_chans) + # Reshape X_std to match coef_ shape: (n_lags, n_feats, n_chans) + X_std_reshaped = X_std.reshape(len(self.lags), self.n_feats_, 1) + self.standardized_coef_ = self.coef_ * X_std_reshaped / y_std[None, None, :] + except Exception as e: + if self.verbose: + LOGGER.warning("Could not compute standardized coefficients: %s", e) + self.standardized_coef_ = None + else: + self.standardized_coef_ = None + # Get t-statistic and p-vals if regularization is ommited if not self.use_regularisation: if self.verbose: LOGGER.info("Computing statistics...") @@ -433,12 +469,24 @@ def _fitlists(self, X, y, drop=True, feat_names=(), lagged=False, verbose=True): assert len(feat_names) == X.shape[1], err_msg self.feat_names_ = feat_names + # Build quadratic regularization matrix M if specified + M = None + if self.quadratic_reg is not None: + if isinstance(self.quadratic_reg, str): + from pyeeg.solvers import create_quadratic_regularizer + n_lags = len(self.lags) + n_feats = self.n_feats_ + L_single = create_quadratic_regularizer(self.quadratic_reg, n_lags) + M = np.kron(np.eye(n_feats), L_single) + else: + M = self.quadratic_reg + if lagged: - betas = _svd_regress([np.hstack([np.ones((len(x), 1)), x])for x in X] if self.fit_intercept else X, [yy[s] for s,yy in zip(valid_samples, y)], self.alpha, self.verbose) + betas = _svd_regress([np.hstack([np.ones((len(x), 1)), x])for x in X] if self.fit_intercept else X, [yy[s] for s,yy in zip(valid_samples, y)], self.alpha, M=M, verbose=self.verbose) else: filling = np.nan if drop else 0. betas = _svd_regress([np.hstack([np.ones((sum(s), 1)), lag_matrix(x, self.lags, filling=filling, drop_missing=drop)]) if self.fit_intercept else lag_matrix(x, self.lags, filling=filling, drop_missing=drop) - for s,x in zip(valid_samples, X)], [yy[s] for s,yy in zip(valid_samples, y)], self.alpha, self.verbose) + for s,x in zip(valid_samples, X)], [yy[s] for s,yy in zip(valid_samples, y)], self.alpha, M=M, verbose=self.verbose) # Storing all alpha's betas self.all_betas = betas # Storing only the first as the main diff --git a/pyeeg/solvers.py b/pyeeg/solvers.py index baa9249..699b830 100644 --- a/pyeeg/solvers.py +++ b/pyeeg/solvers.py @@ -4,17 +4,82 @@ from functools import reduce from tqdm import tqdm import logging -from typing import Union, List +from typing import Union, List, Optional LOGGER = logging.getLogger(__name__) -def svd_solver(A, b, lambda_=0., truncated_svd=False, verbose=False): +def create_laplacian_matrix(n_lags: int, alpha: float = 1.0) -> np.ndarray: + """ + Create a Laplacian matrix for smoothness constraints in quadratic regularization. + + The Laplacian matrix approximates the second derivative, promoting smoothness + in the TRF coefficients across time lags. + + Parameters: + ---------- + n_lags : int + Number of time lags (dimension of the TRF). + alpha : float, optional + Scaling factor for the Laplacian. Default is 1.0. + + Returns: + ------- + L : ndarray (n_lags, n_lags) + Laplacian matrix for smoothness regularization. + + Notes: + ----- + The Laplacian matrix L has the form: + L[i, i-1] = -1 + L[i, i] = 2 + L[i, i+1] = -1 + for interior points, with appropriate boundary conditions. + """ + L = np.zeros((n_lags, n_lags)) + for i in range(n_lags): + L[i, i] = 2.0 + if i > 0: + L[i, i-1] = -1.0 + L[i-1, i] = -1.0 + return alpha * L + + +def create_quadratic_regularizer(reg_type: str, n_lags: int, alpha: float = 1.0) -> np.ndarray: + """ + Factory function to create quadratic regularization matrices. + + Parameters: + ---------- + reg_type : str + Type of regularization: 'smoothness' or 'laplacian'. + n_lags : int + Number of time lags. + alpha : float, optional + Regularization strength. Default is 1.0. + + Returns: + ------- + M : ndarray + Quadratic regularization matrix. + + Raises: + ------ + ValueError + If reg_type is not recognized. + """ + if reg_type in ('smoothness', 'laplacian'): + return create_laplacian_matrix(n_lags, alpha) + else: + raise ValueError(f"Unknown regularization type: {reg_type}. Use 'smoothness' or 'laplacian'.") + + +def svd_solver(A, b, lambda_=0., M=None, truncated_svd=False, verbose=False): """ Solve the linear system Ax = b using the SVD method. - This method assunes that we are solving the normal equation: - (X^T X + lambda I) x = X^T y + This method assumes that we are solving the normal equation: + (X^T X + lambda I + M) x = X^T y Thus, A = X^T X and b = X^T y. Parameters: @@ -23,7 +88,10 @@ def svd_solver(A, b, lambda_=0., truncated_svd=False, verbose=False): b : ndarray Right-hand side vector. Typically of shape (n_features * n_lags, n_outputs) in the context of TRF. lambda_ : float, optional - Regularization parameter. + Regularization parameter (Tikhonov/L2 regularization). + M : ndarray, optional + Quadratic regularization matrix. If provided, solves (A + M) x = b instead of (A + lambda I) x = b. + Useful for smoothness constraints (e.g., Laplacian matrix). truncated_svd : bool, optional Whether to use the truncated SVD method. If True, lambda_ must be between 0 and 1; it represents the fraction of the total variance to keep. @@ -34,6 +102,15 @@ def svd_solver(A, b, lambda_=0., truncated_svd=False, verbose=False): """ # Check symmetricity of A assert np.allclose(A, A.T), 'Matrix A must be symmetric' + + if M is not None: + # Quadratic regularization: solve (A + M) x = b + # Use eigendecomposition for (A + M) + C = A + M + eigvals, eigvecs = np.linalg.eigh(C) + s_inv = np.diag(1 / eigvals) + return eigvecs @ s_inv @ eigvecs.T @ b + U, s, Vt = np.linalg.svd(A, full_matrices=False, hermitian=True) if truncated_svd: assert 0 < lambda_ < 1 @@ -211,12 +288,13 @@ def _lstsq_regress(x: Union[np.ndarray, List[np.ndarray]], else: betas = np.linalg.lstsq(x, y)[0] return betas - def _svd_regress(x: Union[np.ndarray, List[np.ndarray]], y: Union[np.ndarray, List[np.ndarray]], alpha: Union[float, np.ndarray], + M: np.ndarray = None, verbose: bool = False) -> np.ndarray: - """Linear regression using svd. + """ + Linear regression using svd. Parameters ---------- @@ -229,7 +307,13 @@ def _svd_regress(x: Union[np.ndarray, List[np.ndarray]], list is treated as an individual subject, the resulting `betas` coefficients are thus computed on the averaged covariance matrices. alpha : float or array-like - If array, will compute betas for every regularisation parameters at once + If array, will compute betas for every regularisation parameters at once. + Used for Tikhonov/L2 regularization. + M : ndarray, optional + Quadratic regularization matrix. If provided, used in conjunction with alpha. + The solution becomes: betas = (XᵀX + M + alpha*I)⁻¹ Xᵀy + verbose : bool, optional + Whether to print progress information. Returns ------- @@ -266,6 +350,11 @@ def _svd_regress(x: Union[np.ndarray, List[np.ndarray]], if (len(x) == len(y)) and np.ndim(x[0])==2: # will accumulate covariances assert all([xtr.shape[0] == ytr.shape[0] for xtr, ytr in zip(x, y)]), "Inconsistent trial lengths!" XtX = reduce(lambda x, y: x + y, [xx.T @ xx for xx in x]) + + # Add quadratic regularization matrix M if provided + if M is not None: + XtX = XtX + M + [U, s, V] = np.linalg.svd(XtX, full_matrices=False) # here V = U.T XtY = np.zeros((XtX.shape[0], y[0].shape[1]), dtype=y[0].dtype) count = 1 @@ -282,13 +371,22 @@ def _svd_regress(x: Union[np.ndarray, List[np.ndarray]], #betas = U @ np.diag(1/(s + alpha)) @ U.T @ XtY - eigenvals_scaled = np.zeros((*V.shape, np.size(alpha))) - eigenvals_scaled[range(len(V)), range(len(V)), :] = 1 / \ + eigvals_scaled = np.zeros((*V.shape, np.size(alpha))) + eigvals_scaled[range(len(V)), range(len(V)), :] = 1 / \ (np.repeat(s[:, None], np.size(alpha), axis=1) + np.repeat(alpha[:, None].T, len(s), axis=0)) - Vsreg = np.dot(V.T, eigenvals_scaled) # np.diag(1/(s + alpha)) + Vsreg = np.dot(V.T, eigvals_scaled) # np.diag(1/(s + alpha)) betas = np.einsum('...jk, jl -> ...lk', Vsreg, U.T @ XtY) #Vsreg @ Ut else: [U, s, V] = np.linalg.svd(x, full_matrices=False) + if M is not None: + # For single matrix case, we need to handle M properly + # XᵀX + M regularization + if np.ndim(x) == 2: + XtX = x.T @ x + XtX = XtX + M + # Recompute SVD with regularization + [U, s, V] = np.linalg.svd(XtX, full_matrices=False) + if np.ndim(y) == 3: Uty = np.zeros((U.shape[1], y.shape[2])) for Y in y: @@ -299,11 +397,11 @@ def _svd_regress(x: Union[np.ndarray, List[np.ndarray]], # Broadcast all alphas (regularization param) in a 3D matrix, # each slice being a diagonal matrix of s/(s**2+lambda) - eigenvals_scaled = np.zeros((*V.shape, np.size(alpha))) - eigenvals_scaled[range(len(V)), range(len(V)), :] = np.repeat(s[:, None], np.size(alpha), axis=1) / \ + eigvals_scaled = np.zeros((*V.shape, np.size(alpha))) + eigvals_scaled[range(len(V)), range(len(V)), :] = np.repeat(s[:, None], np.size(alpha), axis=1) / \ (np.repeat(s[:, None]**2, np.size(alpha), axis=1) + np.repeat(alpha[:, None].T, len(s), axis=0)) # A dot product instead of matmul allows to repeat multiplication alike across third dimension (alphas) - Vsreg = np.dot(V.T, eigenvals_scaled) # np.diag(s/(s**2 + alpha)) + Vsreg = np.dot(V.T, eigvals_scaled) # np.diag(s/(s**2 + alpha)) # Using einsum to control which access get multiplied, again leaving alpha's dimension "untouched" betas = np.einsum('...jk, jl -> ...lk', Vsreg, Uty) #Vsreg @ Uty