Molecular orbital values on grid points in PySCF?

Question

I am looking for a way to easily evaluate individual molecular orbitals on a grid(assuming a single determinant method for now) in PySCF.

I am aware on how to generate efficient grids for real-space integration (see here for example), and I've also managed to find out how to evaluate the electron density on these grid points, with the following general code for example

def gen_matrix_element(function_to_eval, molecule, myhf, args):
    matelements = np.asarray([])

    for position in grid.coords:
        real_space_value = function_to_eval(r, theta, phi, args) 
        matelements = np.append(matelements, real_space_value)

    dm1 = myhf.make_rdm1(ao_repr=True)[0] + myhf.make_rdm1(ao_repr=True)[1]
    ao_value =  dft.numint.eval_ao(molecule, grid.coords, deriv=0)
    rho = dft.numint.eval_rho(molecule, ao_value, dm1) # Density on the same grid

    combined = grid.weights * matelements
    integral = np.dot(rho.T, combined.T)
    return integral

I am not quite certain, however, what is the procedure if I would want to return the integral value for each molecular orbital individually. Is there a similar function to dft.numint.eval_rho for molecular orbitals?

I also assume that if I sum over all the occupied molecular orbitals, I get the same result as I would get having used the full electron density - is this correct?

---

Edit: I've tried the implementation suggested above, but I am a bit unsure on what the numbers mean.

My current code is the following

#!/usr/bin/env python

import numpy as np
import pyscf
from pyscf import gto, scf, tools, dft, ao2mo
from scipy.special import sph_harm
import math as m

###
# Initialize the molecule and the calculation
###
def setup_system():
    mol = gto.M( # Change atom here
        atom = [["He", (0.0, 0.0, 0.0)], ["He", (1.0, 0.0, 0.0)]],
        basis = {'He': 'aug-cc-pvtz'},
        spin=0,
        charge=0,
        verbose=0
    )
    mol.build()

    # Unperturbed calculators
    mf0 = dft.UKS(mol)
    mf0.xc = 'PBE0'
    e0 = mf0.kernel()

    # Grid for spatial integration
    grid = pyscf.dft.gen_grid.Grids(mol)
    grid.level = 7
    grid.build()

    return mol, mf0, grid, e0

def print_mo_on_grid(molecule, calculator, grid):
    ao = molecule.eval_gto('GTOval', grid.coords)
    mo = np.einsum('gi,xij->xgj', ao, calculator.mo_coeff)
    return ao, calculator.mo_coeff, mo


molecule, calculator, grid, e0 = setup_system() # setting up the calculation

ao_vals, mo_coeff, mo_vals = print_mo_on_grid(molecule, calculator, grid)

print("ao vals shape: ",ao_vals.shape)
print("mo vals shape: ",mo_vals.shape)
print("mo coeff shape: ",mo_coeff.shape)

This code results in the following printout:

ao vals shape:  (79368, 46)
mo vals shape:  (2, 79368, 46)
mo coeff shape:  (2, 46, 46)

I understand the shape of the ao_vals variable: for 79368 grid points, I get the value of each of the 46 basis functions. I already don't fully get the shape of the MO coeff tensor - I assume the first number is 2 because of the UKS calculation I'm doing, but why is that even needed? Can't I have a (46, 46) matrix? And I guess this weird shape messes up the mo_vals variable too - I'd have expected a vector of 79368 elements, but I don't know how to interpret the current shape at all.

Answer 1

You can evaluate molecular orbitals on a grid by first evaluating the AOs on the grid with

 ao = mol.eval_gto('GTOval', coords)

and then contracting with the orbital coefficients $$ \psi_{i\sigma} = \sum_\alpha C^\sigma_{i\alpha} \chi_\alpha $$

I don't remember off-hand which way the indices are in PySCF, but this should be easy to figure out. The code should look something like

mo = lib.einsum('gi,xij->xgj', ao, mf.mo_coeff)

addendum regarding edit above:

ao vals shape:  (79368, 46)
mo vals shape:  (2, 79368, 46)
mo coeff shape:  (2, 46, 46)

This means you have 79368 grid points, 46 basis functions, and 46 spatial molecular orbitals (in general you may have fewer molecular orbitals than atomic orbitals as the atomic orbital basis set may have linear dependencies). For unrestricted wave functions, PySCF uses arrays with additional indices for the spin; I believe [0] is spin-up and [1] is spin-down.