⚡ Bolt: optimize covariance matrix operations

.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.

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),
     )