Python 石墨烯边缘磁性Hatree_Fock计算

Hatree_Fock自洽过程尽可能向量化优化 

import numpy as np
from numpy import exp, pi
from numpy.linalg import eigh
import matplotlib.pylab as plt
from numba import njit

@njit
def Fermi(E, mu, T):
    return 1/(np.exp((E - mu)/T) + 1)

@njit
def tridiag_block_matrix(H, c, u, d):
    # c, u, d are center, upper and lower blocks
    N, _ = H.shape
    n, _ = c.shape
    H[0:n, 0:n] = c
    for i in range(n, N, n):
        H[i:i+n, i:i+n] = c
        H[i-n:i, i:i+n] = d
        H[i:i+n, i-n:i] = u
    return H

@njit
def H0_Graphene(H, k, t = -1):
    A = np.array([[0, t*(1+exp(-1j* k))], [t*(1+ exp(1j*k)), 0]])
    B = np.array([[0, 0], [t, 0]])
    return tridiag_block_matrix(H, A, B, B.T)


# Hatree_Fork 自洽计算过程
def Hatree_Fock(H, ks, Nup_avg, Ndown_avg, U, T = 0.001, mu = 0,  ncc = 20):
    m, _ = Nup_avg.shape
    nk = ks.size
    for ic in range(ncc):
        Nup_avg_tmp = np.zeros((m,m), dtype='double')
        Ndown_avg_tmp = np.zeros((m,m), dtype='double')

        for i in range(nk):
            Hk0 = H0_Graphene(H, ks[i])
            Hk_up = Hk0 + U * (Ndown_avg - 0.5 * np.eye(m))
            Hk_down = Hk0 + U * (Nup_avg - 0.5 * np.eye(m))
            Ek_up, Ak_up = eigh(Hk_up)
            Ek_down, Ak_down = eigh(Hk_down)
            fermi_up = Fermi(Ek_up,mu,T).reshape(1,m)
            fermi_down = Fermi(Ek_down,mu,T).reshape(1,m)
            Nk_up = (Ak_up*Ak_up.conj()*fermi_up).sum(axis=1,dtype='double')
            Nk_down = (Ak_down*Ak_down.conj()*fermi_down).sum(axis=1,dtype='double')
            Nup_avg_tmp += np.diag(Nk_up)
            Ndown_avg_tmp += np.diag(Nk_down)

        Nup_avg = Nup_avg_tmp/nk
        Ndown_avg = Ndown_avg_tmp/nk

        if ic % 5==0:
            print(f"已经完成{ic}次迭代")
        
    return Nup_avg, Ndown_avg

def calculated_band():
    N = 128
    m = 2 * N
    H = np.zeros((m, m), dtype = np.complex64)
    nk = 128
    ks = np.linspace(0, 2*pi, nk)
    U = 1.5
    Nup_avg = 0.5 * np.eye(m)
    Ndown_avg = 0.5 * np.eye(m)
    Nup_avg[0, 0] = 0.7
    Ndown_avg[0, 0] = 0.3
    Nup_avg[-1, -1] = 0.3
    Ndown_avg[-1, -1] = 0.7
    Nup_avg, Ndown_avg = Hatree_Fock(H, ks, Nup_avg, Ndown_avg, U)
    bandUp = np.zeros((nk, m))
    bandDown = np.zeros((nk, m))
    for i in range(nk):
        Hk0 = H0_Graphene(H, ks[i])
        Hk_up = Hk0 + U * (Ndown_avg - 0.5 * np.eye(m)) 
        Hk_down = Hk0 + U * (Nup_avg - 0.5 * np.eye(m))
        Ek_up, _ = eigh(Hk_up)
        Ek_down, _ = eigh(Hk_down)
        bandUp[i, :] = Ek_up
        bandDown[i, :] = Ek_down
    plt.plot(bandUp, color = "red")
    plt.plot(bandDown, color = "blue")
    plt.show()
    plt.plot(Nup_avg - Ndown_avg)
    plt.show()
    # print(Nup_avg[0, 0], Nup_avg[-1, -1])
    # print(Ndown_avg[0, 0], Nup_avg[-1, -1])

if __name__ == '__main__':
    calculated_band()