投影严格对角化Python实现

下面Python实现尽可能利用了numpy向量化的技术。当k点比较大时，计算量还是比较大。

from math import sqrt, pi
import numpy as np
from numpy.lib.twodim_base import mask_indices
import numpy.linalg as LA
import matplotlib.pylab as plt
from scipy.linalg import block_diag

# 生成块三对角矩阵函数
def tridiag(c, u, d, N): 
    # c, u, d are center, upper and lower blocks, repeat N times
    cc = block_diag(*([c]*N)) 
    shift = c.shape[1]
    uu = block_diag(*([u]*N)) 
    uu = np.hstack((np.zeros((uu.shape[0], shift)), uu[:,:-shift]))
    dd = block_diag(*([d]*N)) 
    dd = np.hstack((dd[:,shift:],np.zeros((uu.shape[0], shift))))
    return cc+uu+dd

# 生成Hamiltonian动能部分
def Hamiltonian_H0_SU4(k,N,t=-1):
    """
    t 是最近邻hopping系数
    k 是Bloch波矢量
    N 是宽度
    """
    A = np.matrix([[sqrt(3),0,1,np.exp(-1j*k)],[0, sqrt(3),1,1],[1,1,sqrt(3),0],[np.exp(1j*k),1,0,sqrt(3)]])
    B = np.matrix([[0,0,0,0],[0,0,0,0],[1,0,0,0],[