Today I’ll discuss a numerical method that is widely applicable to problems in quantum mechanics and quantum chemistry, although it is rarely found in textbooks. Suppose we want to solve the stationary Schrödinger equation in its most general form

where 𝐻 is the Hamiltonian operator, 𝜓_𝑛 is an eigenvector (wave functions) of the Hamiltonian and 𝐸_𝑛 the corresponding (energy) eigenvalue. We want to calculate all the different eigenvectors with their energies. In particular, we are interested in the ground state, i.e. the eigenvector with the lowest energy.

The method we discuss here is a variant of gradient descent. But it’s gradient descent in disguise. Consider what happens if we apply 𝐻 to some arbitrary wave function 𝜓_init. We can expand 𝜓 in terms of the (yet unknown) eigenvectors:

where 𝜓_0 shall denote the eigenvector with the lowest energy and 𝜓_𝑁 the one with the highest energy. Then we have

because

the term with the largest eigenvalue gets bigger and bigger in comparison with the other terms. So in the limit, only that term survives and the resulting vector lies completely in the direction of 𝜓_𝑁.

Now, we don’t want the vector with the largest eigenvalue, but the one with the lowest. How can we achieve that? Easy! We just map 𝐻 linearly to some other operator, which has the same eigenvectors, but its eigenvalues are reversed. The simplest choice is

where 𝜖 is some constant to be chosen later. When we apply 𝐾 to some eigenvalue of 𝐻, we get

In other words, the eigenvectors of 𝐻 are also eigenvectors or 𝐾, but with different eigenvalues. For the H-eigenvector 𝜓_0 with lowest H-eigenvalue 𝐸_0 we get

so the corresponding K-eigenvalue is less than 1.

For the H-eigenvector 𝜓_𝑁 with highest H-eigenvalue 𝐸_𝑁, we get

and the corresponding value is 1.

So we see that the spectrum has been reversed. Now we use the idea above to iterate to the K-eigenvector with highest K-eigenvalue, and this will be the H-eigenvector with lowest H-eigenvalue, the ground state.

The iteration scheme in detail is like that, where the upper index in braces counts iterations:

2. Iterate

where

and N is the normalizing operation

We normalize in each step so that we don’t run out of the range of valid floating point numbers.

Finally, we must choose the step size 𝜖. In general, it must be

but since we don’t know the energies beforehand, that doesn’t help very much. There are sophisticated methods available to choose an optimal step size, but some experimentation shall suffice for us here.

But why did I call it gradient descent in disguise? Because 𝐻𝜓 is directly related to the functional derivative of the energy with respect to the wave function:

where 𝛿 denotes a functional derivative. So this recipe is really a gradient descent.

Now let’s write a general routine for our gradient descent:

import numpy as np

def gradient_descent(H, n_steps, eps, psi_init, integral):
    
    def normalize(psi):
        return psi / np.sqrt(integral(psi**2))
    
    psi = normalize(psi_init)
        
    for j in range(n_steps):        
        E = integral(psi * H(psi))
        psi = normalize(psi - eps * (H(psi) - E*psi))
    
    E = integral(psi * H(psi))
    
    return psi, E

We kept this function completely general, so it can be applied to all kind of problems. Let’s apply it to the 1D harmonic oscillator, first. We need to implement the Hamiltonian and the integral operator:

from findiff import FinDiff

x = np.linspace(-5, 5, 100)
dx = x[1] - x[0]
laplace = FinDiff(0, dx, 2)

V = 0.5 * x**2

def H(psi):
    return -0.5 * laplace(psi) + V * psi 


def integral(f):
    return np.trapz(f) * dx

Now we can compute the ground state of the harmonic oscillator:

psi = np.exp(-np.abs(x))
psi, E = gradient_descent(H, 100, 0.01, psi, integral)
E

Out: 0.4997480777288626

The exact value should be 0.5.

import matplotlib.pyplot as plt

plt.plot(x, psi)

Or the hydrogen atom:

x = y = z = np.linspace(-10, 10, 80)
dx = dy = dz = x[1] - x[0]
laplace = FinDiff(0, dx, 2) + FinDiff(1, dy, 2) + FinDiff(2, dz, 2)

X, Y, Z = np.meshgrid(x, y, z, indexing='ij')
R = np.sqrt(X**2 + Y**2 + Z**2)
V = -2 / R

def H(psi):
    return -laplace(psi) + V * psi 


def integral(f):
    return np.trapz(np.trapz(np.trapz(f))) * dx * dy * dz

psi = np.exp(-R**2)

psi, E = gradient_descent(H, 600, 0.01, psi, integral)
E

Out: -0.9848171001223128

The exact value should be -1, so the result is quite crude. This is because we didn’t take care of the singularity at the origin. But that’s another story.

There is much more that we could discuss, like taking into account constraints to calculate excited states. But I will leave that for a separate post. So stay tuned. Thanks for reading!