diff --git a/benchmark_quadratic.py b/benchmark_quadratic.py new file mode 100644 index 0000000..ab2fba5 --- /dev/null +++ b/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() diff --git a/benchmark_spectral.py b/benchmark_spectral.py new file mode 100644 index 0000000..184f36d --- /dev/null +++ b/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() diff --git a/physics/spectral/quadratic/spectral_backend_quadratic.py b/physics/spectral/quadratic/spectral_backend_quadratic.py index 8850d99..a50fce5 100644 --- a/physics/spectral/quadratic/spectral_backend_quadratic.py +++ b/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 # where = 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 # ----------------------------------------- diff --git a/physics/spectral/spectral_backend.py b/physics/spectral/spectral_backend.py index 858ba6f..2dedf07 100644 --- a/physics/spectral/spectral_backend.py +++ b/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)