Plotting of grid data

We will use the Mayavi package to plot grid data. This can be installed with

conda install -c conda-forge mayavi

or

mamba install mayavi

If you want to use Mayavi with Jupyter notebooks, you will need to install the ipyevents package with

conda install -c conda-forge ipyevents

or

mamba install ipyevents

If you want to have interactive plots in your notebooks, you also need to install the x3dom Javascript extension with

jupyter nbextension install --py mayavi --user

We start by importing the necessary packages.

import numpy as np
from mayavi import mlab

If you are using Jupyter notebooks, you will need to run the following cell to initialize Mayavi:

mlab.init_notebook()

Molecular orbitals are one-electron wavefunctions. In our HF calculations, we represent them in the basis of atomic orbitals, which are themselves one-electron wavefunctions. To visualize them in 3D-space, we need to evaluate them on a spatial grid. We shall at first define a function to evaluate a contracted Gaussian basis function in space:

def evaluate_gaussian_basis(g, x, y, z):
    i, j, k = g.ijk
    x0, y0, z0 = g.A
    g_val = 0.0
    for n, d, alpha in zip(g.norm_const, g.coefs, g.exps):
        g_val += n * d * (x - x0)**i * (y - y0)**j * (z - z0)**k \
            * np.exp(-alpha * ((x - x0)**2 + (y - y0)**2 + (z - z0)**2))
    return g_val

Of course, we have to build a spatial grid first. This can be done with the following function:

def build_grid(xlim, ylim, zlim, nx, ny, nz):
    x_ = np.linspace(*xlim, nx)
    y_ = np.linspace(*ylim, ny)
    z_ = np.linspace(*zlim, nz)
    x, y, z = np.meshgrid(x_, y_, z_, indexing='ij')
    return x, y, z

Afterwards, we define a function to transform the MOs from the AO basis to the position basis:

def evaluate_mo_grid(mol, grid):
    x, y, z = grid
    
    # Build MO coefficient matrix in spin-orbit basis
    n_spatial_orb = len(mol.mo_energy)
    c_spin = np.zeros((n_spatial_orb * 2, n_spatial_orb * 2))
    c_spin[:n_spatial_orb:, ::2] = mol.mo_coeff
    c_spin[n_spatial_orb:, 1::2] = mol.mo_coeff
    
    # Evaluate AOs on the grid
    ao_grid = np.zeros((n_spatial_orb * 2, 
                        x.shape[0], x.shape[1], x.shape[2]))
    
    for i, g in enumerate(mol.basisfunctions):
        g_grid = evaluate_gaussian_basis(g, x, y, z)
        ao_grid[i] = g_grid
        ao_grid[i + n_spatial_orb] = g_grid
    
    # Transform to MOs
    mo_grid = np.einsum('ij,ixyz->jxyz', c_spin, ao_grid)
    
    return mo_grid

In the end, we can define a function to plot the MOs, along with the atoms and bonds for reference:

R_VDW = {
    'H': 2.0787,
    'C': 3.2125,
    'N': 2.9291,
    'O': 2.8724,
}

R_COV = {
    'H': 0.6047,
    'C': 1.4173,
    'N': 1.3417,
    'O': 1.1905,
}

ATOM_COLOR = {
    'H': (0.8, 0.8, 0.8),
    'C': (0.0, 1.0, 0.0),
    'N': (0.0, 0.0, 1.0),
    'O': (1.0, 0.0, 0.0),
}

def visualize_cube(mol, grid, density, isovalues, colors, figure, **kwargs):
    # Draw atoms
    for a in mol.atomlist:
        p = mlab.points3d(
            *a.coord, R_VDW[a.symbol], color=ATOM_COLOR[a.symbol], 
            scale_factor=0.5, figure=figure,
        )
    
    # Draw bonds
    for i, a in enumerate(mol.atomlist):
        ca = a.coord
        sa = a.symbol
        ra = R_COV[sa]
        for b in mol.atomlist[:i]:
            cb = b.coord
            sb = b.symbol
            rb = R_COV[sb]
            
            vab = cb - ca
            rab = np.linalg.norm(vab)
            if rab < (ra + rb) * 1.2:
                mid = ca + vab * (ra / (ra + rb))
                p = mlab.plot3d(*np.vstack((ca, mid)).T, tube_radius=0.2, 
                                color=ATOM_COLOR[sa], figure=figure)
                p = mlab.plot3d(*np.vstack((cb, mid)).T, tube_radius=0.2, 
                                color=ATOM_COLOR[sb], figure=figure)

    # Draw isosurface
    for ival, c in zip(isovalues, colors):
        p = mlab.contour3d(*grid, density, color=c, contours=[ival], 
                           figure=figure, **kwargs)
    
    return p

Now, we can take a HF calculation of water

# Coordinates are in the unit of Angstrom.
o1 = Atom('O', [ 0.000,  0.000,  0.000], unit='A')
h1 = Atom('H', [ 0.758,  0.587,  0.000], unit='A')
h2 = Atom('H', [-0.758,  0.587,  0.000], unit='A')

mol = HartreeFock()
mol.set_atomlist([o1, h1, h2])
mol.get_basis('sto-3g')

mol.initialize()
e_scf = mol.run_hf(verbose=0)

construct a grid

XLIM, YLIM, ZLIM = (-5.0, 5.0), (-5.0, 5.0), (-5.0, 5.0)
NX, NY, NZ = 80, 80, 80

grid = build_grid(XLIM, YLIM, ZLIM, NX, NY, NZ)
mo_grid = evaluate_mo_grid(mol, grid)

and plot HOMO and LUMO of water:

ISOSURFACE_COLORS = [(1, 0, 0), (0, 0, 1)]
ORBITAL_ISOVALUES = [0.05, -0.05]

ihomo = (mol.nocc - 1) * 2
ilumo = mol.nocc * 2

fig_homo = mlab.figure()
visualize_cube(mol, grid, mo_grid[ihomo], ORBITAL_ISOVALUES, 
               ISOSURFACE_COLORS, fig_homo)

fig_lumo = mlab.figure()
visualize_cube(mol, grid, mo_grid[ilumo], ORBITAL_ISOVALUES, 
               ISOSURFACE_COLORS, fig_lumo)

# Remove this line if you are using Jupyter notebook
mlab.show()

While the MOs are just mathematical constructions, the electron density is a measurable quantity.

The spin-free one-electron density is defined as $$ \rho(r) = \int \Psi^*(x, x_2, \cdots, x_N) \Psi(x, x_2, \cdots, x_N) \ \mathrm{d}\sigma \mathrm{d}x_2 \cdots \mathrm{d}x_N $$ where \(r\) stands for the spatial coordinate, \(\sigma\) the spin coordinate and \(x\) a combination of both of them. For a wavefunction expressed as a Slater determinant, this expression simplifies to $$ \rho(r) = \sum_{i=1}^{N} \int \phi_i(x) \phi_i(x)\ \mathrm{d} \sigma $$ where \(\phi_i\)'s are occupied spin-orbitals.

We can calculate and visualize the density as follows:

DENSITY_COLORS = [(0.3, 0.3, 0.3)]
DENSITY_ISOVALUES = [0.05]

density_grid = np.zeros((NX, NY, NZ))
for d in mo_grid[:mol.nocc * 2]:
    density_grid += np.abs(d)**2

fig_density = mlab.figure()
visualize_cube(mol, grid, density_grid, DENSITY_ISOVALUES,
               DENSITY_COLORS, fig_density, opacity=0.5)

# Remove this line if you are using Jupyter notebook
mlab.show()