Computing Partial Charges in PyQuante

A user emailed asking whether PyQuante could be used to compute partial charges. Here's what I suggested. This only works with version 1.6 (currently in SVN).

#!/usr/bin/env python
"""\
Generate partial charges for the hydrogen molecule using the Mullikan
approximation. See Szabo/Ostlund eq 3.195.
"""

from PyQuante import SCF, Molecule
from PyQuante.LA2 import mkdens_spinavg
from PyQuante.NumWrap import matrixmultiply,trace
from PyQuante.cints import dist2

def closest_atom((x,y,z),atuples):
    """\
    Return the atom closest to the point (x,y,z) from a list
    of atuples.
    >>> closest_atom((0,0,0),[(1,(0,0,0)),(1,(0,0,1))])
    0
    """
    i_index = None
    r2_min = 1e10
    for i,atom in enumerate(atuples):
        atno,(ix,iy,iz) = atom
        r2 = dist2((x,y,z),(ix,iy,iz))
        if r2 < r2_min:
            i_index = i
            r2_min = r2
    assert i_index is not None, "get_atom failed to find atom"
    return i_index

def partial_charges(molecule,**opts):
    # Converge the HF wave function for this molecule:
    method = opts.get("method","HF")
    iterator = SCF(molecule,method=method)
    iterator.iterate()

    # Grab the orbitals, basis functions, and overlap matrix:
    orbs = iterator.solver.orbs
    bfs = iterator.basis_set.get()
    S = iterator.S

    # Compute the density matrix, and multiply by the overlap
    nclosed,nopen = molecule.get_closedopen()
    D = mkdens_spinavg(orbs,nclosed,nopen)
    DS = matrixmultiply(D,S)

    # Make sure the number of electrons is correct
    print "Number of electrons = ",2*trace(DS)

    # Now sum up the electrons using the Mullikan approximation.
    # Assume that the density associated with each basis function
    # goes onto the atom that owns that basis function:
    
    #qs = [0 for atom in molecule] # Use this if electron charge wanted
    qs = [atom.atno for atom in molecule] # Use this for total charge
    for ibf,bf in enumerate(bfs):
        iat = closest_atom(bf.origin(),molecule.atuples())
        qs[iat] -= 2*DS[ibf,ibf]
    return qs
    
def main():
    h2o = Molecule('h2o',[('O',(0.000, 0.000,-0.079)),
                          ('H',(0.000, 0.707, 0.628)),
                          ('H',(0.000,-0.707, 0.628))],
                   units="Angstroms")
    qs = partial_charges(h2o)
    print qs,sum(qs)
    return

if __name__ == '__main__': main()