根据wannier90_hr.dat构造哈密顿量，并根据考虑的近邻数画出能带

作者：吕叶竹

首先准备POSCAR文件，wannier90_hr.dat，以及下面的代码，放在一个文件夹中，直接运行即可，其中15行设置费米能级，第102和第103行用来设置考虑的近邻数，176行用来设置y轴范围，代码如下：

######################################################################################################
import numpy as np
import matplotlib.pyplot as plt
######################################################################################################
import numpy as np
f=open("POSCAR","r")
ln=f.readlines()
f.close()
       
deg_ws=[]
lat=[]
a=[]

k=0
fermi=-2.2586 

for j in range(2,5):
    sp=ln[j].split()
            # get reduced lattice vector components
    ham_R1=float(sp[0])
    ham_R2=float(sp[1])
    ham_R3=float(sp[2])
           
    ham_key=(ham_R1,ham_R2,ham_R3)
    lat.append(ham_key)

lat=np.array(lat,dtype=float)

lat0=lat[0]
lat1=lat[1]
lat2=lat[2]

lat_per=np.copy(lat)
lat_per=lat_per[list(range(3))] 
#######################################################################################################################################
path=[[0.00000,  0.00000,   0.00000],[0.50000,  0.50000,   0.00000],[0.00000,  0.50000,   0.00000], [0.00000,  0.00000,   0.00000], [0.50000,  0.00000,   0.00000]]#高对称点坐标
k_list=np.array(path)
n_nodes=k_list.shape[0]
nk=101
k_metric = np.linalg.inv(np.dot(lat_per,lat_per.T))
#print(k_metric)
k_node=np.zeros(n_nodes,dtype=float)
#print(k_node)
for n in range(1,n_nodes):
            dk = k_list[n]-k_list[n-1]
            dklen = np.sqrt(np.dot(dk,np.dot(k_metric,dk)))
            k_node[n]=k_node[n-1]+dklen
    
        # Find indices of nodes in interpolated list
node_index=[0]
for n in range(1,n_nodes-1):
    frac=k_node[n]/k_node[-1]
    node_index.append(int(round(frac*(nk-1))))
node_index.append(nk-1)
k_dist=np.zeros(nk,dtype=float)
k_vec=np.zeros((nk,3),dtype=float)
#print(k_vec)
k_vec[0]=k_list[0]
for n in range(1,n_nodes):
    n_i=node_index[n-1]
    n_f=node_index[n]
    kd_i=k_node[n-1]
    kd_f=k_node[n]
    k_i=k_list[n-1]
    k_f=k_list[n]
    for j in range(n_i,n_f+1):
        frac=float(j-n_i)/float(n_f-n_i)
        k_dist[j]=kd_i+frac*(kd_f-kd_i)
        k_vec[j]=k_i+frac*(k_f-k_i)


#提取hr.dat的数据到矩阵中  求H矩阵（未除以简并）
###########################################################################################################
import numpy as np
f=open("wannier90_hr.dat","r")
ln=f.readlines()
f.close()
        #
        # get number of wannier functions
num_wan=int(ln[1])
        # get number of Wigner-Seitz points
num_ws=int(ln[2])
        # get degenereacies of Wigner-Seitz points
deg_ws=[]
key=[]
a=[]

k=0
for j in range(3,len(ln)):
    sp=ln[j].split()
    for s in sp:
        deg_ws.append(int(s))
    if len(deg_ws)==num_ws:
        last_j=j
        break
    if len(deg_ws)>num_ws:
        raise Exception("Too many degeneracies for WS points!")
deg_ws=np.array(deg_ws,dtype=int)
ham_r={} # format is ham_r[(R1,R2,R3)]["h"][i,j] for < i | H | j+R >
ind_R=0 # which R vector in line is this?
near1=5
near2=5  #Set the nearest neighbor considered
for j in range(last_j+1 ,len(ln)):
    sp=ln[j].split()
            # get reduced lattice vector components
    ham_R1=int(sp[0])
    ham_R2=int(sp[1])
    ham_R3=int(sp[2])
            # get Wannier indices
    if abs(ham_R1)<=near1 and abs(ham_R2)<=near2:
        ham_i=int(sp[3])-1
        ham_j=int(sp[4])-1
                # get matrix element
        ham_val=float(sp[5])+1.0j*float(sp[6])
        ham_key=(ham_R1,ham_R2,ham_R3)
        key.append(ham_key)
        if (ham_key in ham_r)==False:
            ham_r[ham_key]={
                "h":np.zeros((num_wan,num_wan),dtype=complex),
                "deg":deg_ws[ind_R]
                }
            ind_R+=1
        ham_r[ham_key]["h"][ham_i,ham_j]=ham_val
############################################################################################################
print(ind_R)
#提取R（约化坐标下），组成数组   求Rij
###########################################################################################################
for l in range(ind_R):  #range(hr.dat中数据行/轨道数的平方=R的数目)
    a.append(key[k])
    k=k+num_wan*num_wan #k=k+轨道数的平方
Rc=np.asarray(a)
###########################################################################################################

#构造H(k),矩阵中每一项是各个R对应的tij*e^(ik*Rij)求和。 求本征值
###########################################################################################################
ret_eval=np.zeros((num_wan,nk),dtype=float)
ham=np.zeros([num_wan,num_wan],dtype = complex)
Ham=np.zeros([num_wan,num_wan],dtype = complex)
H=np.zeros([num_wan,num_wan],dtype = complex)

for m,k in enumerate(k_vec):
    for i in range(num_wan):
        for j in range(num_wan):
            for R in a:
                cc=np.asarray(R)
                R_ij=cc
            #print(R_ij)
                phase=np.exp((2.0j)*np.pi*np.dot(k,R_ij))
            #print(phase)
                ham[i,j]=((ham_r[R]["h"][i,j]/float(ham_r[R]["deg"])))*phase
                Ham[i,j]=Ham[i,j]+ham[i,j]
    H=Ham
    Ham=np.zeros([num_wan,num_wan],dtype = complex)
    eval=np.linalg.eigvalsh(H)
    eval=np.array(eval.real-fermi,dtype=float)
    args=eval.argsort()
    eval=eval[args]
    ret_eval[:,m]=eval[:]

##########################################################################################################
#画能带图
##########################################################################################################
# k_label=(r'$\Gamma$',r'$M$',r'$K$', r'$\Gamma$')
# k_label=(r'$M$',r'$K$',r'$\Gamma$', r'$K$', r'$M$')
k_label=(r'$G$',r'$R$',r'$Y$', r'$G$', r'$X$')
fig, ax = plt.subplots()
for i in range(ret_eval.shape[0]):
    ax.plot(k_dist,ret_eval[i],"k-")
for n in range(len(k_node)):
    ax.axvline(x=k_node[n],linewidth=0.5, color='k')
ax.set_xlabel("Path in k-space")
ax.set_ylabel("Band energy (eV)")
ax.set_xlim(k_dist[0],k_dist[-1])
ax.set_xticks(k_node)
ax.set_xticklabels(k_label)
plt.ylim(-0.5,0.5)
fig.tight_layout()
fig.savefig("band.pdf")