⚡ Bolt: optimize covariance matrix operations and diag_nd

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

src/eegprep/plugins/clean_rawdata/private/covariance.py

@@ -35,49 +35,52 @@ 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)
 
 
 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):
     """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)
     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)),
     )