#!/usr/bin/env python
from scipy import *
from pylab import *

def eps(k):
    " Energy eigenvalues for graphene tigh binding with tp'=0"
    return sqrt(3+2*(cos(k[0])+cos(k[1])+cos(k[0]-k[1])))


def GivePath(arr, Np):
    """ given array of k-points in the path, generates a
    dense list of momentum points in the choosen bath
    (Np points between each pair of points in the array)
    """
    kmsh = arange(0, 1, 1./Np)
    tot_path = []
    for i in range(len(arr)-1):
        path = map(lambda x: 2*pi*(arr[i] + (arr[i+1]-arr[i])*x), kmsh)
        tot_path.extend(path)
    return tot_path


# The three points in terms of
# reciprocal vectors (b1,b2)
G = array([0,0])
K = array([1/3.,-1/3.])
M = array([0,-1/2.])
# Number of points between to momentum points to plot
Np = 100

# Path we want to plot
spath = ['M', 'G', 'K', 'M']
rpath = [M, G, K, M]

# Computes array of points in 1BZ
arrp = GivePath(rpath, Np)

# computes band energy for these points
Y1 = map(eps, arrp)
# and energy for the second band
Y2 = map(lambda x: -eps(x), arrp)

# produces title for the plot containing path in 1BZ
strpath=""
for s in spath:
    strpath += s
    

# actual plotting
title('Band structure of graphene in the path '+strpath)
xticks(Np*arange(len(rpath)), spath)
ylabel("Energy in units of hopping t")
rc('lines', lw=3)

plot(Y1)
plot(Y2)

show()