from ATK.KohnSham import *

def lineOfVectorPoints(start_point,end_point,N=20):
    '''
    @return: A list containing (default: 20) k-points between two
    specified points in the Brillouin zone.
    '''
    from numpy import array
    points = []
    for i in range(N):
        new_point = (float(i)/(N-1),)*3*(array(end_point) - array(start_point))+start_point
        points.append(new_point)
    return points

# Define common symmetry points in the Brillouin zone of silicon
fcc_symmetry_points = {
    'G' : (0.000, 0.000, 0.000),
    'X' : (0.500, 0.000, 0.500), 
    'W' : (0.500, 0.250, 0.750), 
    'L' : (0.500, 0.500, 0.500), 
    'K' : (0.375, 0.375, 0.750), 
    'U' : (0.625, 0.250, 0.625) 
}

#Create line in Brillouin zone along which the
#band structure is calculated
start_point = fcc_symmetry_points['G']
end_point = fcc_symmetry_points['X']
band_k_points = lineOfVectorPoints(start_point,end_point,25)

#Restore old SCF calculation
calculation = restoreSelfConsistentCalculation('si_gga_dzp.nc')
band_structure = calculateEnergyBands(
    self_consistent_calculation = calculation,
    k_point_list = band_k_points)

#Extract calculated bands into a list
energy_bands = []
for band_index in range(band_structure.numberOfBands()):
    energy_bands.append(band_structure.band(band_index))

#print band energies in units of eV
for i in range(len(band_k_points)):
    print band_k_points[i][0],
    for band in range(band_structure.numberOfBands()):
        print energy_bands[band][i].inUnitsOf(eV),
    print