from ase import Atoms
from gpaw import GPAW, PoissonSolver, LCAO
from ase.optimize import QuasiNewton
from ase.structure import molecule
from gpaw.mpi import world

system = molecule('H2O')
system.center(vacuum=4.0)

# (Exact parameters unimportant)
calc = GPAW(txt='gpaw.txt', h=0.2, mode='lcao', basis='dzp')

system.set_calculator(calc)
E = system.get_potential_energy()

import numpy as np
W_aL = {}
for a in calc.density.D_asp:
    W_aL[a] = np.empty((calc.wfs.setups[a].lmax + 1)**2)

#calc.hamiltonian.vHt_g[:, :, :] = 1.0 # if this line is not commented,
# the integral should be exactly 1 due to normalization convention.

calc.density.ghat.integrate(calc.hamiltonian.vHt_g, W_aL)

print 'results'
for key, val in W_aL.items():
    print 'atom number', key
    print 'local integral', val[0] / np.sqrt(4.0 * np.pi)
    print