⚡ Vectorize nested loops in spectral computations for performance

benchmark_quadratic.py

@@ -0,0 +1,28 @@
+import numpy as np
+import time
+from general_python.physics.spectral.quadratic.spectral_backend_quadratic import greens_function_quadratic, greens_function_quadratic_finite_T
+
+def main():
+    np.random.seed(42)
+    N = 500
+    omega = 0.5
+    eta = 0.01
+    beta = 1.0
+
+    eigenvalues = np.sort(np.random.randn(N))
+    operator_a = np.random.randn(N, N) + 1j * np.random.randn(N, N)
+    operator_b = np.random.randn(N, N) + 1j * np.random.randn(N, N)
+    occ = np.random.randint(0, 2, size=N)
+
+    start = time.perf_counter()
+    res1 = greens_function_quadratic(omega, eigenvalues, operator_a=operator_a, operator_b=operator_b, occupations=occ, eta=eta, backend="numpy", basis_transform=False)
+    end = time.perf_counter()
+    print(f"greens_function_quadratic: {end - start:.4f}s")
+
+    start = time.perf_counter()
+    res2 = greens_function_quadratic_finite_T(omega, eigenvalues, operator_a=operator_a, operator_b=operator_b, beta=beta, eta=eta, backend="numpy", basis_transform=False)
+    end = time.perf_counter()
+    print(f"greens_function_quadratic_finite_T: {end - start:.4f}s")
+
+if __name__ == "__main__":
+    main()

benchmark_spectral.py

@@ -0,0 +1,32 @@
+import numpy as np
+import time
+from general_python.physics.spectral.spectral_backend import operator_spectral_function_lehmann, susceptibility_bubble
+
+def main():
+    np.random.seed(42)
+    N = 500
+    omega = 0.5
+    eta = 0.01
+    temperature = 1.0
+
+    eigenvalues = np.sort(np.random.randn(N))
+    eigenvectors = np.linalg.qr(np.random.randn(N, N) + 1j * np.random.randn(N, N))[0]
+    operator = np.random.randn(N, N) + 1j * np.random.randn(N, N)
+    operator = operator + operator.conj().T
+
+    # Benchmark Lehmann
+    start = time.perf_counter()
+    res1 = operator_spectral_function_lehmann(omega, eigenvalues, eigenvectors, operator, eta, temperature, backend="numpy")
+    end = time.perf_counter()
+    print(f"operator_spectral_function_lehmann: {end - start:.4f}s, result: {res1}")
+
+    # Benchmark Bubble
+    vertex = np.random.randn(N, N) + 1j * np.random.randn(N, N)
+    occupation = np.random.rand(N)
+    start = time.perf_counter()
+    res2 = susceptibility_bubble(omega, eigenvalues, vertex, occupation, eta, backend="numpy")
+    end = time.perf_counter()
+    print(f"susceptibility_bubble: {end - start:.4f}s, result: {res2}")
+
+if __name__ == "__main__":
+    main()

physics/spectral/quadratic/spectral_backend_quadratic.py

@@ -187,15 +187,29 @@ def greens_function_quadratic(
     # Double sum over m,n with occupation factors
     # Computes the Green's function <A| 1 / (w + i eta - (H - E0)) |B>
     # where <A| = <c0| A and |B> = B |c0>
-    for m in range(N):
-        if not occ[m]:    # needs to be occupied
-            continue
-        for n in range(N):
-            if occ[n]:    # needs to be empty
+
+    # Vectorized computation, blocked over 'm' to save memory
+    mask_n = occ == 0
+
+    if be.any(mask_n):
+        ev_n = ev[mask_n]
+        B_sub = B[mask_n, :]
+
+        for m in range(N):
+            if not occ[m]:    # needs to be occupied
                 continue
 
-            deltaE = ev[n] - ev[m]
-            G     += (A[m, n] * B[n, m]) / (z - deltaE)
+            deltaE = ev_n - ev[m]
+
+            # Extract row m of A for empty n columns
+            A_sub_m = A[m, mask_n]
+            # Extract column m of B for empty n rows
+            B_sub_nm = B_sub[:, m]
+
+            num = A_sub_m * B_sub_nm
+            denom = z - deltaE
+
+            G += be.sum(num / denom)
 
     return G
 
@@ -382,18 +396,27 @@ def greens_function_quadratic_finite_T(
     # build Fermi factor f_m
     f = be.asarray(1.0 / (1.0 + be.exp(beta * (ev.real - mu))))
 
+    # Vectorized computation, blocked over 'm' to save memory
+    f_1_minus_f = 1.0 - f
+
     for m in range(N):
-        em = ev[m]
         fm = f[m]
+        if fm == 0:
+            continue
 
-        for n in range(N):
-            fn      = f[n]
-            weight  = fm * (1.0 - fn)
-            if weight == 0:
-                continue
+        weights = fm * f_1_minus_f
+        mask = weights != 0
+
+        if not be.any(mask):
+            continue
+
+        deltaE = ev - ev[m]
+
+        num = weights[mask] * A[m, mask] * B[mask, m]
+        denom = z - deltaE[mask]
+
+        G += be.sum(num / denom)
 
-            deltaE = ev[n] - em
-            G     += weight * (A[m,n] * B[n,m]) / (z - deltaE)
     return G
 
 # -----------------------------------------

physics/spectral/spectral_backend.py

@@ -907,14 +907,20 @@ def lorentzian(delta_E):
         return (eta / np.pi) / (delta_E**2 + eta**2)
 
     # Lehmann sum over all transitions
+    # Block vectorization over 'm' to avoid N^2 memory allocations
     A = 0.0
     for m in range(N):
-        for n in range(N):
-            if be.abs(rho[m] - rho[n]) < 1e-14:
-                continue
-            delta_E = omega - (eigenvalues[n] - eigenvalues[m])
-            matrix_element_sq = be.abs(O_eigen[m, n])**2
-            A += (rho[m] - rho[n]) * matrix_element_sq * lorentzian(delta_E)
+        rho_diff = rho[m] - rho
+        mask = be.abs(rho_diff) >= 1e-14
+
+        if not be.any(mask):
+            continue
+
+        delta_E = omega - (eigenvalues - eigenvalues[m])
+        matrix_element_sq = be.abs(O_eigen[m, :])**2
+
+        A_vec = rho_diff[mask] * matrix_element_sq[mask] * lorentzian(delta_E[mask])
+        A += be.sum(A_vec)
 
     return float(be.real(A))
 
@@ -1065,14 +1071,20 @@ def susceptibility_bubble(
         occupation = be.asarray(occupation, dtype=be.float64)
 
     # Bubble: χ⁰ = \Sum _{mn} (f_m - f_n) |V_{mn}|² / (omega  + i\eta - (E_n - E_m))
+    # Block vectorization over 'm' to avoid N^2 memory allocations
     chi = 0.0 + 0.0j
     for m in range(N):
-        for n in range(N):
-            if be.abs(occupation[m] - occupation[n]) < 1e-14:
-                continue
-            denom = omega_complex + 1j * eta_complex - (eigenvalues[n] - eigenvalues[m])
-            V_mn_sq = be.abs(vertex[m, n])**2
-            chi += (occupation[m] - occupation[n]) * V_mn_sq / denom
+        occ_diff = occupation[m] - occupation
+        mask = be.abs(occ_diff) >= 1e-14
+
+        if not be.any(mask):
+            continue
+
+        denom = omega_complex + 1j * eta_complex - (eigenvalues - eigenvalues[m])
+        V_mn_sq = be.abs(vertex[m, :])**2
+
+        chi_vec = occ_diff[mask] * V_mn_sq[mask] / denom[mask]
+        chi += be.sum(chi_vec)
 
     return complex(chi)