Application: Harmonic Oscillator

The quantum harmonic oscillator is described by wavefunctions $$ \psi_{n}(x) = \frac{1}{\sqrt{2^n n!}} \left( \frac{m \omega}{\pi \hbar} \right)^{1/4} \exp\left(-\frac{m\omega x^2}{2\hbar}\right) H_{n}\left( \sqrt{\frac{m\omega}{\hbar}} x \right) $$ where the functions $H_{n}(z)$ are the physicists' Hermite polynomials: $$ H_{n}(z) = (-1)^n\ \mathrm{e}^{z^2} \frac{\mathrm{d}^n}{\mathrm{d}z^n} \left( \mathrm{e}^{-z^2} \right) $$ The corresponding energy levels are $$ E_{n} = \omega \left( n + \frac{1}{2} \right) $$

We now want to use SymPy to compute eigenenergies and eigenfunctions. We start by importing sympy and neccesary variables:

import sympy as sp
from sympy.abc import x, m, omega, n, z

The eigenenergies are stright-forward to compute. One can just use the definition directly and substitute $n$ by an integer, e.g. 5:

def energy(n):
    e = omega * (n + sp.Rational(1, 2))
    return e

e_5 = energy(5)

This outputs $\frac{11}{2}$. Note that atomic units are used for simplicity.

In order to evaluate the wave function, we have to at first compute the Hermite polynomial. There are multiple ways to do it. The first is to use its definition directly and compute higher order derivatives of the exponential function:

def hermite_direct(n):
    h_n = (-1)**n * sp.exp(z**2) * sp.diff(sp.exp(-z**2), (z, n))
    h_n = sp.simplify(h_n)
    return h_n

Another way is to use the recurrence relation $$ H_{n+1}(z) = 2zH_{n}(z) - H_{n}^{'}(z) $$ which can be implemented as a recursive function:

def hermite_recursive(n):
    if n == 0:
        return sp.sympify(1)
    else:
        return sp.simplify(
            2 * x * hermite_recursive(n - 1) \
            - sp.diff(hermite_recursive(n - 1), x)
        )

For n = 5, both methods lead to the polynomial $$ H_5(x) = 32 x^5 - 160 x^3 + 120 x$$

We can then use this to evaluate the wave function:

def wfn(n):
    nf = (1/sp.sqrt(2**n * sp.factorial(n))) \
         * ((m*omega)/sp.pi)**sp.Rational(1, 4)
    expf = sp.exp(-(m*omega*x**2)/2)
    hp = hermite_direct(n).subs(z, sp.sqrt(m*omega)*x)
    psi_n = sp.simplify(nf * expf * hp)

    return psi_n

psi_5 = wfn(5)

For a system with, say, $m = 1$ and $\omega = 1$, we can construct its wave function by

psi_5_param = psi_5.subs([(m, 1), (omega, 1)])

Further substituting x with any numerical value would evaluate the wave function at that point. With the evalf method, we can obtain the approximated numerical value:

psi_5_param_x_eval = psi_5_param.evalf(subs={x: sp.Rational(1, 2)})
print(psi_5_param_x_eval)
print(type(psi_5_param_x_eval))

Here, we compute the value of the wave function at $x = \frac{1}{2}$. We get 0.438575095003232 as the result. Note that the type of this object is not float, but sympy.core.numbers.Float, which is a floating-point number implementation with arbitrary precision. Actually, we can pass an additional argument n to evalf to get an accuracy of n digits: The listing

print(psi_5_param.evalf(n=50, subs={x: sp.Rational(1, 2)}))

for example, produces 0.43857509500323214478874587638315301825689754260515. This type can be easily converted to python's built-in float by calling float(). The arbitrary precision however, will be lost.

Sometimes one might just want to use SymPy to generate symbolic expressions to be used in functions. Instead of typing the results in by hand, one could use one of SymPy's built-in printers. Here, we shall use NumPyPrinter, which converts a SymPy expression to a string of python code:

from sympy.printing.pycode import NumPyPrinter
printer = NumPyPrinter()
code = printer.doprint(psi_5_param)
code = code.replace('numpy', 'np')
print(code)

This output can be copied and used to evaluate $\psi_{5}(x)$. Since we often import NumPy with the alias np, numpy from the original string is replaced.

Programming in Theoretical Chemistry

Application: Harmonic Oscillator