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

@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_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 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)
    
# @njit
def calculated_band(H, ks, t = 1):
    nk = ks.size
    m, _ = H.shape
    band = np.zeros((nk, m - 1))
    Hk = np.zeros((m, m), dtype=np.complex64)
    for i in range(nk):
        Hk = H0_SU4(Hk, ks[i])
        E, _ = eigh(Hk[1:,1:])
        band[i, :] = E
    return band

def plot_SU4_Band():
    nk = 128
    N = 256
    ks = np.linspace(0, 2*pi, nk)
    Hk = np.zeros((4*N, 4*N), dtype=np.complex64)
    band = calculated_band(Hk, ks)
    plt.plot(band, color = "gray")
    plt.plot(band[:, N - 1], color = "red")
    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()

if __name__ == '__main__':
    plot_SU4_Band()

需要理解numpy数组的深浅拷贝，传指针引用等思想。