# Import the KohnSham module from ATK
from ATK.KohnSham import *
   
# Load molecule from VNLFile
h2o = VNLFile("h2o.vnl").readAtomicConfigurations()["h2o"]

# List of basis set and their names 
BasisSetList = [['SZ'  , SingleZeta],              
                ['SZP' , SingleZetaPolarized],
                ['DZ'  , DoubleZeta],              
                ['DZP' , DoubleZetaPolarized],     
                ['DZDP', DoubleZetaDoublePolarized]]

# Variable for holding the previous energy
E0 = None

print '    Basis set    TotalEnergy (eV)    Difference    '
print '---------------------------------------------------'

for basis in BasisSetList:
    name,basis_type = basis
    bsp = basisSetParameters(type=basis_type)
    method = KohnShamMethod(basis_set_parameters=bsp)
    E = calculateTotalEnergy(method.apply(h2o)).inUnitsOf(eV)

    if E0:
        print "     %-9s %15.6f %15.6f" % (name, E, (E-E0))
    else:
        print "     %-9s %15.6f"        % (name, E)
    E0 = E