python画石墨烯的能带图和态密度

学习python已经1个多月了，最近需要画一下石墨烯的能带图和态密度，接下来利用python对角化石墨烯中C原子的Hamilton量，取本征值，画一下能带图吧。

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from numpy import *
from  math import *
def Hamiltonian3(pos,a1,a2,k,t):
    H=mat(zeros((2,2)))
    eps=0.01
    for i in range(0,2):
        for j in range(0,2):
            for m in range(-1,2):
                for n in range(-1,2):
                    d=pos[i]+n*a1+m*a2-pos[j]
                    # 最近邻判断条件
                    if abs(linalg.norm(d)-1)<eps:
                        # 满足最近邻条件的原子间跳跃积分求和
                        H[i,j]=H[i,j]+t * e ** (1j*dot(k,d))
    return H
# 元胞内的两个C原子坐标
pos1=array([0,0])
pos2=array([0,1])
pos=[pos1,pos2]
# 取晶格基矢
a1=array([sqrt(3),0])
a2=array([sqrt(3)/2,3/2])

#新建一个名叫 fig的画图窗口
fig = plt.figure()
# 准备绘制三维图形
ax = Axes3D(fig)
# k点的取值范围
Kx = arange(-pi, pi, 0.01)
Ky = arange(-pi, pi, 0.01)
n=len(Kx)
# 初始化能量E为n×n的0矩阵
E1=mat(zeros((n,n)))
t=-1
for i in range(n):
   for j in range(n):
       # 取k点
       k=array([Kx[i],Ky[j]])
       # 计算Hamilton
       H=Hamiltonian3(pos,a1,a2,k,t)
       # 取本征值
       Es=linalg.eig(H)[0]
       # 存储本征值
       E1[i,j]=Es[0]
E2=-E1
#画图
ax.plot_surface(Kx,Ky,E1, rstride=1, cstride=1, cmap='rainbow')
ax.plot_surface(Kx,Ky,E2, rstride=1, cstride=1, cmap='rainbow')
plt.show()

画出来的如如下。

image.png