diff --git a/.jules/bolt.md b/.jules/bolt.md new file mode 100644 index 00000000..a9e0c7fd --- /dev/null +++ b/.jules/bolt.md @@ -0,0 +1,5 @@ +## 2025-05-14 - Optimized Covariance Matrix Operations + +**Learning:** Matrix power functions (logm, expm, powm, sqrtm) implemented via eigen-decomposition can be significantly optimized by replacing diagonal matrix creation and full matrix multiplication with NumPy broadcasting. Scaling eigenvectors by eigenvalues is $O(N^2)$ vs $O(N^3)$ for matrix multiplication, and avoids $O(N^2)$ space for the diagonal matrix. `diag_nd` was also a bottleneck due to loop-based concatenation. + +**Action:** Prefer broadcasting for scaling columns/rows of matrices by vectors. Use advanced indexing for batch diagonal matrix creation. diff --git a/src/eegprep/plugins/clean_rawdata/private/covariance.py b/src/eegprep/plugins/clean_rawdata/private/covariance.py index cd646640..5c006244 100644 --- a/src/eegprep/plugins/clean_rawdata/private/covariance.py +++ b/src/eegprep/plugins/clean_rawdata/private/covariance.py @@ -34,50 +34,56 @@ def diag_nd(M): """Like np.diag, but in case of a ...,N, returns a ...,N,N array of diag matrices.""" *dims, N = M.shape - if dims: - cat = np.concatenate([np.diag(d) for d in M.reshape((-1, N))]) - return np.reshape(cat, dims + [N, N]) - else: - return np.diag(M) + res = np.zeros((*dims, N, N), dtype=M.dtype) + idx = np.arange(N) + res[..., idx, idx] = M + return res def cov_logm(C): """Calculate the matrix logarithm of a covariance matrix or ...,N,N array.""" D, V = np.linalg.eigh(C) - return finite_matmul(finite_matmul(V, diag_nd(np.log(D))), V.swapaxes(-2, -1)) + # Optimized: V * log(D) @ V^T + return finite_matmul(V * np.log(D)[..., np.newaxis, :], V.swapaxes(-2, -1)) def cov_expm(C): """Calculate the matrix exponent of a covariance matrix or ...,N,N array.""" D, V = np.linalg.eigh(C) - return finite_matmul(finite_matmul(V, diag_nd(np.exp(D))), V.swapaxes(-2, -1)) + # Optimized: V * exp(D) @ V^T + return finite_matmul(V * np.exp(D)[..., np.newaxis, :], V.swapaxes(-2, -1)) def cov_powm(C, exp): """Calculate a matrix power of a covariance matrix or ...,N,N array.""" D, V = np.linalg.eigh(C) - return finite_matmul(finite_matmul(V, diag_nd(D**exp)), V.swapaxes(-2, -1)) + # Optimized: V * (D^exp) @ V^T + return finite_matmul(V * (D**exp)[..., np.newaxis, :], V.swapaxes(-2, -1)) def cov_sqrtm(C): """Calculate the matrix square root of a covariance matrix or ...,N,N array.""" D, V = np.linalg.eigh(C) - return finite_matmul(finite_matmul(V, diag_nd(np.sqrt(D))), V.swapaxes(-2, -1)) + # Optimized: V * sqrt(D) @ V^T + return finite_matmul(V * np.sqrt(D)[..., np.newaxis, :], V.swapaxes(-2, -1)) def cov_rsqrtm(C): """Calculate the matrix reciprocal square root of a covariance matrix or ...,N,N array.""" D, V = np.linalg.eigh(C) - return finite_matmul(finite_matmul(V, diag_nd(1.0 / np.sqrt(D))), V.swapaxes(-2, -1)) + # Optimized: V * (1/sqrt(D)) @ V^T + return finite_matmul(V * (1.0 / np.sqrt(D))[..., np.newaxis, :], V.swapaxes(-2, -1)) def cov_sqrtm2(C): """Calculate the matrix square root, and its reciprocal, for a covariance matrix or ...,N,N array.""" D, V = np.linalg.eigh(C) sqrtD = np.sqrt(D) + # Optimized: V * sqrt(D) @ V^T and V * (1/sqrt(D)) @ V^T + VT = V.swapaxes(-2, -1) return ( - finite_matmul(finite_matmul(V, diag_nd(sqrtD)), V.swapaxes(-2, -1)), - finite_matmul(finite_matmul(V, diag_nd(1.0 / sqrtD)), V.swapaxes(-2, -1)), + finite_matmul(V * sqrtD[..., np.newaxis, :], VT), + finite_matmul(V * (1.0 / sqrtD)[..., np.newaxis, :], VT), )