The Hydrogen Atom Solved With Python

Diogoleite asked me recently, after reading my article “ Two Lines of Python to Solve The Schrödinger Equation ”, how one could use the so-called shooting method to solve the Schrödinger equation for the hydrogen atom. The method I showed in that article is not well suited for the case of hydrogen, so let’s discuss that special case here.

First of all, what is the problem with the method of my older article when applying it to the hydrogen atom? The Schrödinger equation is a partial differential eigenvalue equation. That is, it’s not only a differential equation but it also contains an undetermined parameter (energy). Only for certain values of this parameter does the Schrödinger equation yield physical results. These valid parameter values are the energy states of the physical system. To solve this problem, I discretized space, which turned the partial differential equation into an algebraic equation. So we were left with a normal matrix eigenvalue problem, which is easy to solve with the standard scipy package. In principle, this does also work with the hydrogen atom. But the potential of the hydrogen atom has a singularity at the origin, and that’s a problem for numerics when discretizing space. There are also methods to handle this, but extra steps are needed.

So, coming back to Diogoleite ‘s original question, how do we solve the hydrogen atom with the shooting method? The Schrödinger equation for the hydrogen atom reads

You may notice that there is a factor 2 in the potential. This is coming from my choice of units. I’m using an atomic unit system where ℏ is 1 and energies are in units of Rydberg. But that’s just details that don’t matter.

The usual approach is to go to spherical coordinates and make a separation ansatz

into a radial part and an angular part. The angular part just yields the spherical harmonics and we don’t have to worry about that anymore. The only problem left to be solved is the radial part. The radial part of the Schrödinger equation then reads

where 𝑙 is a non-negative integer corresponding to the index 𝑙 in the spherical harmonics. For a given value of 𝐸, we can solve this equation numerically using standard techniques. So we bring it into standard form by defining

So that we get the first-order system

Let’s define a function for 𝑦′:

import numpy as np

def y_prime(y, r, l, E):
    u, u_prime = y
    return np.array(
        [
            u_prime, 
            (l*(l+1)/r**2 - 2/r - E) * u
        ]
    )

We have the boundary condition 𝑢(∞)=0 and 𝑢(0)=0. In the method of shooting, we use one condition as the initial condition and then naively integrate up to the point where the other boundary condition shall be applied. Usually, that second boundary condition will not automatically be fulfilled. So we need to adjust our free parameter 𝐸 such that the result for 𝑢u also matches the second boundary condition. That’s shooting.

When we solve the radial equation for a given 𝐸, we have two choices:

or

It turns out that in our case, the first option is numerically unstable, while the second option is stable. So we integrate from outside to inside. Of course, we don’t integrate from infinity, but just some large 𝑟. Here is the code for solving the radial equation:

def solve_radial(E, l, rr):
    """Solves the radial differential equation for given E and l."""
    
    # For stability reasons, we integrate
    # from outside (large r) to inside (small r),
    # so reverse the radial values:
    rr_rev = rr[::-1]
    u_prime_oo = -1.E-6 # value for the initial slope is quite arbitrary

    # Set the "initial" conditions
    y_oo = [0, u_prime_oo]
    
    # Solve the radial equation
    y_rev = integrate.odeint(
        y_prime, y_oo, rr_rev, args=(l, E))
    
    # Reverse order again so that have u ordered normally:
    u = y_rev[:, 0][::-1]
    
    # Normalize state:
    norm = integrate.simps(u**2, x=rr)
    u /= np.sqrt(norm)
    
    return u

For some given 𝐸, the function gives us a solution. But we have to take care that the solution is physical. It is only physical if it satisfies the boundary condition 𝑢(0)=0. However, we can only integrate to small but nonzero 𝑟r because the ODE has a singularity at 𝑟=0. So we extrapolate our computed solution to 𝑟=0 to estimate its value 𝑢(0):

def shoot(E, l, r):
    u = solve_radial(E, l, r) / r**l
    
    # Extrapolate u to the origin r=0.
    return u[0] - r[0] * (u[1] - u[0])/(r[1] - r[0]), u

How does that look like for various 𝐸? Let’s check:

rr = np.logspace(-6, 5, 500)
EE = np.linspace(-1.5, 0.1, 100)
u0s = [
    shoot(E, 0, rr)[0] for E in EE
]

plt.plot(EE, u0s)
plt.grid()
plt.xlabel("E")
plt.ylabel("u(0)")

The physical values for the energy are those that satisfy 𝑢(0)=0. So we can graphically evaluate the energy eigenvalues from this plot. But of course, we want the computer to do this for us. So we start from some low enough energy and increase the energy until 𝑢(0) changes sign. Then we know an 𝐸-interval where 𝑢(0)=0 is contained. We can then use a root-finding routine from scipy to find the exact true value for 𝐸:

from scipy.optimize import root_scalar

E_bounds = [] # saves the physical energy values
u_bounds = [] # saves the physical wave functions
n_states = 7 # max. number of states we are looking for

dE = 0.01 # scan resolution to look for sign changes
E = -1.5 # starting energy
l = 0  # the angular momentum

u0_last, _ = shoot(E, l, rr)

# Scan energies from below and watch for sign changes of u(0):

while len(E_bounds) < n_states and E < 0:
    E += dE
    u0, u = shoot(E, l, rr)
    
    if u0 * u0_last < 0:  
        # sign switch between last and current E  
        
        E_bound = root_scalar(
            lambda E_: shoot(E_, l, rr)[0], x0=E-dE, x1=E).root
        
        E_bounds.append(E_bound)
        # get the wave function at this energy:
        _, u_bound = shoot(E_bound, l, rr)

        u_bounds.append(u_bound)
    
    u0_last = u0

This gives us the energy eigenvalues E_bounds:

[-1.0000000140820482,
 -0.24999999677082366,
 -0.11111111288956713,
 -0.06250000354217922,
 -0.03999762571325783,
 -0.02769212300497268,
 -0.01899299852275164]

And the corresponding wave functions:

for u in u_bounds:
    plt.plot(rr, u)

plt.xlim(0, 30)
plt.xlabel("r")
plt.ylabel("u(r)")

That’s it for the shooting method applied to the hydrogen atom. Thanks for reading!