from ATK.KohnSham import *

# Set the basis set to be SZ
basis_set = basisSetParameters(type=SingleZeta)

# Define the density grid cut-off
electron_density_params = electronDensityParameters(
    mesh_cutoff=100*Rydberg
    )

# Define the k-point sampling
bz_integration = brillouinZoneIntegrationParameters( (1,1,100) )

# Define the DFT method
dft_method = KohnShamMethod(
    basis_set_parameters=basis_set,
    brillouin_zone_integration_parameters=bz_integration,
    electron_density_parameters=electron_density_params
    )

# Open file for storing the results
f = open('energy.dat', 'w')

# Loop over lattice constants
import numpy
for a in numpy.arange(2.9,3.1,0.01):

    # Create the unit cell for 1D chain
    unit_cell = [[9.0, 0.0, 0.0],
                [ 0.0, 9.0, 0.0],
                [ 0.0, 0.0,  a ]]*Angstrom

    # Create the atomic configuration
    li_chain = PeriodicAtomConfiguration(unit_cell,[Lithium], 
                                         [ (0.0, 0.0, 0.0)*Angstrom ] )
    
    # Execute the self-consistent calculation
    dft_calculation = dft_method.apply(li_chain)
    
    # Calculate the system energy
    energy = calculateTotalEnergy(dft_calculation)
    
    # Write lattice constant and energy to file and screen 
    s = '%.4f %.10f \n' % (a, energy.inUnitsOf(Rydberg))
    f.write(s)
    print s,

# Close file and exit
f.close()