diff --git a/.jules/bolt.md b/.jules/bolt.md new file mode 100644 index 00000000..acd90371 --- /dev/null +++ b/.jules/bolt.md @@ -0,0 +1,5 @@ +## 2025-05-15 - Vectorized Diagonal Scaling in Covariance Operations + +**Learning:** Matrix operations of the form $V \cdot \text{diag}(D) \cdot V^T$ (common in matrix functions like logm, expm, sqrtm) are significantly more efficient when implemented using NumPy broadcasting (`V * D[..., np.newaxis, :]`) instead of explicit diagonal matrices. This avoids $O(N^3)$ matrix multiplication for the scaling step and eliminates large $O(N^2)$ memory allocations for the diagonal matrices. Additionally, `diag_nd` (creating a batch of diagonal matrices) is much faster using advanced indexing (`res[..., i, i] = M`) than looping and concatenating. + +**Action:** Prefer broadcasting for diagonal scaling and advanced indexing for diagonal matrix creation in all performance-sensitive matrix code. diff --git a/src/eegprep/plugins/clean_rawdata/private/covariance.py b/src/eegprep/plugins/clean_rawdata/private/covariance.py index cd646640..de9286eb 100644 --- a/src/eegprep/plugins/clean_rawdata/private/covariance.py +++ b/src/eegprep/plugins/clean_rawdata/private/covariance.py @@ -35,8 +35,10 @@ 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]) + res = np.zeros((*dims, N, N), dtype=M.dtype) + i = np.arange(N) + res[..., i, i] = M + return res else: return np.diag(M) @@ -44,31 +46,32 @@ def diag_nd(M): 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)) + # Using broadcasting (V * logD) is equivalent to V @ diag(logD) but faster + 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)) + 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)) + 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)) + 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)) + return finite_matmul(V * (1.0 / np.sqrt(D))[..., np.newaxis, :], V.swapaxes(-2, -1)) def cov_sqrtm2(C): @@ -76,8 +79,8 @@ def cov_sqrtm2(C): D, V = np.linalg.eigh(C) sqrtD = np.sqrt(D) 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, :], V.swapaxes(-2, -1)), + finite_matmul(V * (1.0 / sqrtD)[..., np.newaxis, :], V.swapaxes(-2, -1)), )