SU(4)Hubbard平均场Python代码

标准的SU(N)共线磁序平均场计算，存在很强简并性。

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


def H0_SU4(H, k, t = 1):
    A = t * np.array([[sqrt(3),0,1,exp(-1j*k)],[0, sqrt(3),1,1],[1,1,sqrt(3),0],[exp(1j*k),1,0,sqrt(3)]])
    B = t * np.array([[0,0,0,0],[0,0,0,0],[1,0,0,0],[0,-1,0,0]])
    return tridiag_block_matrix(H, A, B.T, B)

@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

# Hatree_Fork 自洽计算过程
def Hatree_Fock(H, ks, N1avg, N2avg, N3avg, N4avg, U, T = 0.005, mu = 0,  ncc = 20):
    m, _ = N1avg.shape
    nk = ks.size
    band1 = np.zeros((nk, m))
    band2 = np.zeros((nk, m))
    band3 = np.zeros((nk, m))
    band4 = np.zeros((nk, m))
    for ic in range(ncc):
        N1avg_tmp = np.zeros((m,m), dtype='double')
        N2avg_tmp = np.zeros((m,m), dtype='double')
        N3avg_tmp = np.zeros((m,m), dtype='double')
        N4avg_tmp = np.zeros((m,m), dtype='double')
        m1avg = (N1avg - N2avg)/2
        m2avg = (N1avg+N2avg-2*N3avg)/(2*sqrt(3))
        m3avg = (N1avg+N2avg+N3avg-3*N4avg)/(2*sqrt(6))
        C = 8/5*U*(m1avg*m1avg +m2avg*m2avg + m3avg*m3avg)
        U1 = U * (m1avg/2+m2avg/(2*sqrt(3)) + m3avg/(2*sqrt(6)))
        U2 = U * (-m1avg/2+m2avg/(2*sqrt(3))+ m3avg/(2*sqrt(6)))
        U3 = U * (-2*m2avg/(2*sqrt(3)) + m3avg/(2*sqrt(6)))
        U4 = U * (-3*m3avg/(2*sqrt(6)))

        for i in range(nk):
            H0 = H0_SU4(H, ks[i])
            Hk0 = H0[1:,1:]
            Ek1, Ak1 = eigh(Hk0 - 16/5 * U1 + C)
            Ek2, Ak2 = eigh(Hk0 - 16/5 * U2 + C)
            Ek3, Ak3 = eigh(Hk0 - 16/5 * U3 + C)
            Ek4, Ak4 = eigh(Hk0 - 16/5 * U4 + C)

            if ic == ncc - 1:
                band1[i, :] = Ek1
                band2[i, :] = Ek2
                band3[i, :] = Ek3
                band4[i, :] = Ek4

            Fermi1 = Fermi(Ek1,mu,T).reshape(1,m)
            Fermi2 = Fermi(Ek2,mu,T).reshape(1,m)
            Fermi3 = Fermi(Ek3,mu,T).reshape(1,m)
            Fermi4 = Fermi(Ek4,mu,T).reshape(1,m)
            Nk1 = (np.abs(Ak1)*np.abs(Ak1)*Fermi1).sum(axis=1,dtype='double')
            Nk2 = (np.abs(Ak2)*np.abs(Ak2)*Fermi2).sum(axis=1,dtype='double')
            Nk3 = (np.abs(Ak3)*np.abs(Ak3)*Fermi3).sum(axis=1,dtype='double')
            Nk4 = (np.abs(Ak4)*np.abs(Ak4)*Fermi4).sum(axis=1,dtype='double')
            N1avg_tmp += np.diag(Nk1)
            N2avg_tmp += np.diag(Nk2)
            N3avg_tmp += np.diag(Nk3)
            N4avg_tmp += np.diag(Nk4)

        N1avg = N1avg_tmp/nk
        N2avg = N2avg_tmp/nk
        N3avg = N3avg_tmp/nk
        N4avg = N4avg_tmp/nk
    
        if ic % 5==0:
            print(f"已经完成{ic}次迭代")
        
    return band1, band2, band3, band4, N1avg, N2avg, N3avg, N4avg

def plot_band():
    N = 256
    m = 4 * N
    nk = 128
    U = 0.2
    ks = np.linspace(0, 2*pi, nk)
    H = np.zeros((m, m), dtype=np.complex64)
    N1avg = 0.25 * np.eye(m - 1)
    N2avg = 0.25 * np.eye(m - 1)
    N3avg = 0.25 * np.eye(m - 1)
    N4avg = 0.25 * np.eye(m - 1)
    N1avg[0, 0] = 0.3
    N2avg[0, 0] = 0.3
    N3avg[0, 0] = 0.3
    N4avg[0, 0] = 0.1
    band1, band2, band3, band4, N1avg, N2avg, N3avg, N4avg = Hatree_Fock(H, ks, N1avg, N2avg, N3avg, N4avg, U)
    plt.plot(band1, color = "gray")
    plt.plot(band2, color = "gray")
    plt.plot(band3, color = "gray")
    plt.plot(band4, color = "gray")
    plt.plot(band1[:, N - 1], color = "red")
    plt.plot(band2[:, N - 1], color = "red")
    plt.plot(band3[:, N - 1], color = "red")
    plt.plot(band4[:, N - 1], color = "blue")
    plt.xticks(np.arange(0, nk, nk//3), ['0', '2/3π', '4/3π', '2π'], fontsize = 12, fontweight = 'bold')
    plt.yticks(np.arange(-0.5, 0.6, 0.5), fontsize = 12, fontweight = 'bold')
    plt.ylim(-0.5, 0.5)
    plt.xlabel("k", fontsize = 13, fontweight = 'bold')
    plt.ylabel("E", fontsize = 13, fontweight = 'bold')
    plt.show()

    plt.plot(np.diag(N1avg))
    plt.plot(np.diag(N2avg))
    plt.plot(np.diag(N3avg))
    plt.plot(np.diag(N4avg))
    plt.show()

if __name__ == '__main__':
    N = 32
    m = 4 * N
    nk = 128
    U = 0.2
    ks = np.linspace(0, 2*pi, nk)
    H = np.zeros((m, m), dtype=np.complex64)
    N1avg = 0.25 * np.eye(m - 1)
    N2avg = 0.25 * np.eye(m - 1)
    N3avg = 0.25 * np.eye(m - 1)
    N4avg = 0.25 * np.eye(m - 1)
    N1avg[0, 0] = 0.3
    N2avg[0, 0] = 0.3
    N3avg[0, 0] = 0.3
    N4avg[0, 0] = 0.1
    band1, band2, band3, band4, N1avg, N2avg, N3avg, N4avg = Hatree_Fock(H, ks, N1avg, N2avg, N3avg, N4avg, U)
    np.save("./band1.npy", band1)
    np.save("./band2.npy", band2)
    np.save("./band3.npy", band3)
    np.save("./band4.npy", band4)
    np.save("./N1avg", N1avg)
    np.save("./N2avg", N2avg)
    np.save("./N3avg", N3avg)
    np.save("./N4avg", N4avg)