bands.py

import matplotlib.pyplot as plt
import numpy as np

from momentum import Momentum
from tSnS import tSnS


def Bandstructure(model, k_points, bandset=2, **kwargs):

	"""
	#Calculate & Plot Bandstructure along irreducible 
	#path with given number of k_points
	"""

	#################################
	#Calculate Bandstructure
	#################################
	#Create path along irreducible path in 1. BZ
	#High symmetry points
	S     = 0.5*(model.G1+model.G2)
	X     = 0.5*model.G1
	Y     = 0.5*model.G2
	Gamma = 0.0*model.G1

	irr_path = [Gamma, Y, S, Gamma, X, S]
	k_path, high_sym_index = BZ.k_path(k_points, irr_path = irr_path)

	band_colors   = kwargs.get('band_colors', ['k']*model.N)
	high_sym_marker  = [r'$\Gamma$', r'$Y$', r'$S$', r'$\Gamma$', r'$X$', r'$S$']

	e = []
	for kpoint in k_path: 

		H0 = model.BH_Hamiltonian(kpoint, bandset=bandset)
		e_temp = np.linalg.eigvalsh(H0)
		e+= [e_temp]
	e = np.stack(e)


	#################################
	#Plot Bandstructure
	#################################

	fig = plt.figure(figsize = (3,2))
	ax =fig.add_subplot(1,1,1)
	ax.set_ylabel(r'$\epsilon_b (eV)$')
	ax.set_ylim(e.min(), e.max())
	
	#Plot bands
	x = np.arange(0,k_points,1)
	
	for i in range (len(e[0])):
		ax.plot(x, e[:,i], color = band_colors[i], linewidth = 1.7)
		
	#Indicate high symmetry points
	for i in range(len(high_sym_index)):
		ax.axvline(high_sym_index[i], color='k', alpha=0.9, linewidth = 0.5)

	ax.set_xticks(high_sym_index)
	ax.set_xticklabels(high_sym_marker)
	ax.set_xlim(high_sym_index[0], high_sym_index[-1])

	plt.tight_layout()
	plt.show()



if __name__ == '__main__': 

	model = tSnS()
	BZ    = Momentum(model.G1, model.G2)
	Bandstructure(model, 1000, bandset=2)