plotting real space density

density.py

import psi4
import numpy as np

# from qcdb
def INT_NCART(am):
    """Gives the number of cartesian functions for an angular momentum.
    #define INT_NCART(am) ((am>=0) ? ((((am)+2)*((am)+1))>>1) : 0)

    """
    return (((abs(am) + 2) * (abs(am) + 1)) >> 1)

# taken and modified from qcdb
def compute_phi_ao(shell, xyz_ao, x, y, z):
    """Returns the values of a basis function sitting at xyz_ao at point x,y,z"""
    phi = 0
    am = shell.am()
    nprim = shell.nprimitive()
    a = shell.exps()
    c = shell.coefs()

    dx = x - xyz_ao[0]
    dy = y - xyz_ao[1]
    dz = z - xyz_ao[2]
    rr = dx * dx + dy * dy + dz * dz

    cexpr = 0
    for iprim in range(nprim):
        cexpr += c[iprim] * np.exp(-a[iprim] * rr)

    for l in range(INT_NCART(am)):
        components = basis.exp_ao[am][l]
        phi += pow(dx, components[0]) * \
                        pow(dy, components[1]) * \
                        pow(dz, components[2]) * \
                        cexpr
    return phi


def calc_D(cgto, Da, z):
    x = y = 0
    density = 0
    for i, phi_i, xyz_aoi in cgto:
        for j, phi_j, xyz_aoj in cgto:
            # should be using gaussian product rule here..
            density += Da[i][j] * compute_phi_ao(phi_i, xyz_aoi, x, y, z) \
                                * compute_phi_ao(phi_j, xyz_aoj, x, y, z)
    return density

# Determinant of a 2x2 matrix leads to same loop as above, with normalization??
def calc_amplitude(cgto,nalpha, Ca, z):
    x = y = 0
    amp = 0
    for i, phi_i, xyz_aoi in cgto:
        for j, phi_j, xyz_aoj in cgto:
            amp += Ca[i][j] * compute_phi_ao(phi_i, xyz_aoi, x, y, z)\
                            * compute_phi_ao(phi_j, xyz_aoj, x, y, z)
    return amp/np.sqrt(nalpha*2)


psi4.core.set_output_file('output.dat', False)

mol = psi4.geometry("""
0 1
H 0.0 0.0 0
H 0.0 0.0 1.4
units bohr
no_reorient
symmetry c1
# no_com
""")

psi4.set_options({'basis': 'sto-3g', 'scf_type': 'pk', 'e_convergence': 1e-8})
psi4.set_memory('1 GiB')

e, wfn = psi4.energy("hf", return_wfn=True)

Da = wfn.Da_subset("AO").np
Da *= 2.0  # !!
Ca = wfn.Ca_subset("AO", "ALL").np
# print(Ca)
nalpha = wfn.nalpha()
nel = nalpha * 2
# print(nalpha)
# print(Da)

molecule = psi4.qcdb.Molecule(mol.to_dict())
basis = psi4.qcdb.BasisSet.pyconstruct(molecule, "BASIS", "sto-3g")

nao = basis.nao()

# testing
# for i in range(nao):
#   for j in range(nao):
#     x=0
#     for k in range(nalpha/2):
#       x+=2.0*Ca[i][k]*Ca[j][k]
#     Da[i][j]=x
# print(Da)

cgto = []
# grab all cGTOs and their aufpunkte
c = 0
for icenter in range(0, molecule.natom()):
    xyz_ao = molecule.xyz(icenter)
    for ishell in range(0, basis.nshell_on_center(icenter)):
        cgto.append((c, basis.shell(icenter, ishell), xyz_ao))
        c += 1

# print(nao,c)
# scanning params in Angstrom
scan_line = np.linspace(-3.0, 3.0, num=1000)

for zval in scan_line:
    z = zval / 0.5291 # xyz_ao is in bohr
    val=calc_D(cgto,Da,z)
    # val = calc_amplitude(cgto,nalpha, Ca, z)

    print(zval, val)