计算物理投影对角化性能优化——numpy + numba

一般来说，对角化哈密顿量总是依赖布里渊区的一系列q，然后将对角化得到的能带和波函数存下来，这个算法的通用写法是：

@jit(nopython=True)
def solve_Heff(nk, V):
    L, nk = V.shape
    Ek = np.zeros((nk, nk))
    Sk = np.zeros((nk, nk))
    for i in range(nk):
        Heff = np.diag(Mterm2(V, nk, i)) - Mterm3(V, nk, i)
        E, Psi = LA.eigh(Heff)
        n = np.power(np.absolute(Psi), 2)
        Sk_tmp = -n * np.log(n + 1e-5)
        Ek[:,i] = E
        Sk[:,i] = np.sum(Sk_tmp, axis = 0)
    return Ek, Sk

哈密顿量一般情况下总是按照一定规则生成，这里用常规的二重循环生成哈密顿量矩阵。

@jit(nopython=True)
def  Mterm2(V, nk, q):
    M2 = np.zeros(nk, dtype=np.complex64)
    Ap = np.power(np.absolute(V),2).sum(axis=1).copy()  # 预处理一个常数项
    for i in range(nk):
        M2[i] += np.sum(Ap*np.power(np.absolute(V[:,(i - q + nk) % nk]), 2))
    return M2 / nk

@jit(nopython=True)
def Mterm3(V, nk, q):
    M3 = np.zeros((nk, nk), dtype=np.complex64)
    for i in range(nk):
        for j in range(nk):
            M3[i, j] += np.sum(np.conj(V[:, (i - q + nk) % nk]) * V[:, i] * V[:,(j - q + nk) % nk] * np.conj(V[:,j]))
    return M3 / nk

这段程序用numpy+numba部分向量化优化达到性能最大，任务管理器观察发现，此程序将CPU跑满。