A Simple Solver In Python

in Lessons
#computational chemistry

In this blog post, I would to introduce a simple Schrödinger equation solver in python. There are of course plenty of programs out there that can solve the Schrödinger equation. And many of them solve it in the context of practical theories and systems. However, often times the resulting wavefunctions of those calculations are deeply hidden within the program, which makes it tedious to experiment with. When I'm trying to do something new, sometimes I'll turn to this kind of simple program.

I'm going to write here what I think is the absolute bare minimum solver you can make. We'll only solve equations in one dimension. We'll use the eigensolver provided by SciPy. And we'll solve all the equations in real space using the finite difference method. First, let's recall the time independent Schrödinger equation:

$$\begin{equation} \big[\nabla^2 + V(x)\big] \psi = \lambda \psi. \end{equation}$$

First let's process the input. Ignore some of the import statements for now, we'll come back to these functions when we use them.

from sys import argv
from numpy import linspace, diag
from scipy.sparse.linalg import eigsh
from matplotlib import pyplot as plt

from UniversalOperators import TOperator
import SystemOperators

if __name__ == "__main__":
    # Process Input
    grid_size = int(argv[1])
    nocc = int(argv[2])
    PotentialName = argv[3]
    x_values = linspace(-6.28, 6.28, grid_size)
    grid_spacing = x_values[1] - x_values[0]

Our program takes three input parameters:

  • The number of grid points.
  • The number of eigenvectors to compute.
  • The name of the potential we want to study.

Using the linspace function in numpy, we can easily generate a list of evenly spaced grid points to represent our potential and wavefunction on.

Next we need a matrix that represents the kinetic energy:

    # Build Operator
    KineticMatrix = TOperator(grid_size, grid_spacing)

The Kinetic Energy operator is a universal operator, as opposed to the potential energy operators which are specific to the system we are trying to model. We will put the kinetic energy operator in a file UniversalOperators.py and the potential energy operators in SystemOperators.py. The kinetic energy operator can be implemented using the five point stencil approximation[1]:

from scipy.sparse import coo_matrix

def TOperator(grid_size, grid_spacing):
    '''
    Five point stencil kinetic energy operator.
    '''
    data = []
    row = []
    col = []

    for i in range(grid_size):
        if i > 1:
            row.append(i)
            col.append(i - 2)
            data.append((-0.5) * (-1.0 / (12.0 * grid_spacing**2)))
        if i > 0:
            row.append(i)
            col.append(i - 1)
            data.append((-0.5) * (16.0 / (12.0 * grid_spacing**2)))
        row.append(i)
        col.append(i)
        data.append((-0.5) * (-30.0 / (12.0 * grid_spacing**2)))
        if i + 1 < grid_size:
            row.append(i)
            col.append(i + 1)
            data.append((-0.5) * (16.0 / (12.0 * grid_spacing**2)))
        if i + 2 < grid_size:
            row.append(i)
            col.append(i + 2)
            data.append((-0.5) * (-1.0 / (12.0 * grid_spacing**2)))

    return coo_matrix((data, (row, col)), shape=(grid_size, grid_size))

In this routine, all the logic like if i > 0 is to deal with the boundary condition. In this simple example, the boundary condition is that the wavefunction goes to zero. Next we'll also have to construct a potential energy operator, and add it to our overall Hamiltonian:

    try:
        PotentialOperator = getattr(SystemOperators, PotentialName)
    except:
        exit(1)
    PotentialMatrix, potential = PotentialOperator(x_values)
    Hamiltonian = KineticMatrix + PotentialMatrix

Here we'll take advantage of Python's flexibility. The getattr function will get the function in the SystemOperators.py based on the string provided to the program. How should we implement these potential energy operators? Here's are two examples:

'''
Generate operators representing various systems.
'''
from scipy.sparse import coo_matrix
from scipy import exp
from scipy.special import lambertw

def HydrogenOperator(x_values):
    '''
    Hydrogen potential well
    '''
    grid_size = len(x_values)
    data = []
    row = []
    col = []

    for i in range(0, grid_size):
        row.append(i)
        col.append(i)
        x = x_values[i]
        data.append(1.0/abs(x))

    return coo_matrix((data, (row, col)), shape=(grid_size, grid_size)), data

def LambertOperator(x_values):
    '''
    Lambert W Function
    '''
    grid_size = len(x_values)
    data = []
    row = []
    col = []

    for i in range(0, grid_size):
        row.append(i)
        col.append(i)
        x = x_values[i]
        data_point = (1.0/(1.0 + lambertw(exp(-abs(x))))).real
        data.append(data_point)

    return coo_matrix((data, (row, col)), shape=(grid_size, grid_size)), data

The HydrogenOperator function is just a simple $\frac{1}{r}$ operator. The LambertOperator uses the Lambert W function[2] just to keep things interesting.

Now that we have the operator in matrix form, we can solve for the eigenvectors.

    # Solve The Eigenvalue Equation
    values, vectors = eigsh(Hamiltonian, k=nocc, which='SA')

And we can plot the result.

    # Plotting
    fig, ax1 = plt.subplots()
    ax2 = ax1.twinx()
    for i in range(0, nocc):
        ax1.plot(x_values, vectors.T[i], 'o--', label="K="+str(i))
    ax2.plot(x_values, PotentialMatrix.diagonal(), 'x--', label="potential")

    plt.legend()
    plt.show()

Some example output is below. We'll come back to this example later for some more interesting exercises.

Example output

[1] https://en.wikipedia.org/wiki/Five-point_stencil

[2] https://en.wikipedia.org/wiki/Lambert_W_function