# Import the KohnSham module from ATK
from ATK.KohnSham import *

# Load molecule from VNLFile
h2o = VNLFile('h2o.vnl').readAtomicConfigurations()['h2o']

method = KohnShamMethod()
scf    = method.apply(h2o)

# calculate HOMO and LUMO eigenstates
egvs = calculateEigenstates(
    self_consistent_calculation=scf,
    quantum_numbers=( [3,4] )
    )

results = VNLFile('results.vnl')
results.addToSample(h2o, 'h2o')
results.addToSample(egvs[0], 'h2o','0')
results.addToSample(egvs[1], 'h2o','1')