diff --git a/.jules/bolt.md b/.jules/bolt.md new file mode 100644 index 00000000..eebce40f --- /dev/null +++ b/.jules/bolt.md @@ -0,0 +1,3 @@ +## 2025-05-15 - [Optimization of batched covariance functions] +**Learning:** Replacing intermediate diagonal matrix creation and full matrix multiplication with broadcasting scaling (`V * D[..., np.newaxis, :]`) significantly improves performance by reducing memory allocation and operation complexity from $O(N^3)$ to $O(N^2)$ for the scaling step. Batched diagonal matrix creation can also be optimized using advanced indexing instead of loops. +**Action:** Always look for diagonal matrix multiplications and replace them with vectorized broadcasting when applicable. For batched operations, avoid loops and `np.concatenate` in favor of pre-allocated arrays and advanced indexing. diff --git a/src/eegprep/plugins/clean_rawdata/private/covariance.py b/src/eegprep/plugins/clean_rawdata/private/covariance.py index cd646640..4b21c27a 100644 --- a/src/eegprep/plugins/clean_rawdata/private/covariance.py +++ b/src/eegprep/plugins/clean_rawdata/private/covariance.py @@ -34,41 +34,42 @@ 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: + if not dims: 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)) + 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 +77,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)), )