Functions for creating images of hydrogen orbitals

hydrogen-orbitals.ipynb

hydrogen_orbitals.py

import numpy as np
import matplotlib.pyplot as plt
from scipy.special import sph_harm, genlaguerre
from ai import cs # coordinate trasnformations


A0_STAR = 5.2946541e-11 # The (reduced) Bohr radius A0_STAR is a constant of the universe


def big_psi(qnum_n, qnum_l, radius):
    """
    The radial component of the elctron probability denisty.

    qnum_n: integer, 0 <= qnum_n
    qnum_l: integer, 0 <= qnum_l <= qnum_n-1
    radius: real, 0 < raius

    big_psi(qnum_n, qnum_l, radius): real
    """

    rho = (2*radius) / (qnum_n*A0_STAR)
    L = genlaguerre(qnum_n - qnum_l - 1, 2*qnum_l + 1)(rho)

    factorial = lambda x: np.prod(np.arange(1, qnum_n + 1))   # good enough factorial for small numbers
    sqrt_term = np.sqrt(((2/qnum_n*A0_STAR)**3) * factorial(qnum_n - qnum_l - 1) / (2*qnum_n*factorial(qnum_n + qnum_l)))

    result = sqrt_term * np.exp(-rho/2) * (rho**qnum_l) * L

    return result


def psi_normsq(qnum_n, qnum_l, qnum_m, radius, theta, phi):
    """
    Multiply the radial component big_psi by the appropriate spherical harmonic and take the norm-square.

    qnum_n: integer, 0 <= qnum_n
    qnum_l: integer, 0 <= qnum_l <= qnum_n-1
    qnum_m: inetger, -qnum_l <= qnum_m <= qnum_l
    radius: real, 0 < radius
    theta:  real, angle
    phi:    real, angle

    psi_normsq(qnum_n, qnum_l, qnum_m, radius, theta, phi): real, > 0
    """

    err_msg = "n, l and m must be integers with n = 0,1,2,...; l = 0,1,...,n-1; m = -l,...,l"
    if not all([int(x) == x for x in [qnum_n, qnum_l, qnum_m]]):
        raise ValueError(err_msg)
    if not (qnum_n >= 0 and (0 <= qnum_l <= qnum_n - 1) and -qnum_l <= qnum_m <= qnum_l):
        raise ValueError(err_msg)

    Y = sph_harm(qnum_m, qnum_l, theta, phi)   # note l, m swap order due to implementation of sph_harm

    result = np.abs(big_psi(qnum_n, qnum_l, radius) * Y)**2

    return result


def psi_xy(qnum_n, qnum_l, qnum_m, x, y):
    """
    value of psi_normsq in the x-y plane

    qnum_n: integer, 0 <= qnum_n
    qnum_l: integer, 0 <= qnum_l <= qnum_n-1
    qnum_m: integer, -qnum_l <= qnum_m <= qnum_l
    x: real
    y: real

    psi_xy(qnum_n, qnum_l, qnum_m, x, y): real, > 0
    """

    radius, theta, phi = cs.cart2sp(x, y, 0)
    result = psi_normsq(qnum_n, qnum_l, qnum_m, radius, theta, phi)

    return result


def create_image(data, cmap='viridis', filename=None):
    """
    Creates (and, optionally, saves) image from numpy data

    data:     numpy real array, assumed square and non-negative
    cmap:     matplotlib colormap
    filename: default None (no save in this case), otherwise expects string saves image as filename
    """
    sizes = np.shape(data)
    fig = plt.figure(figsize=(5, 5))
    ax = plt.Axes(fig, [0., 0., 1., 1.])
    ax.set_axis_off()
    fig.add_axes(ax)
    ax.imshow(data, cmap)
    if filename is not None:
        plt.savefig(filename, dpi=sizes[0])
        plt.close()


def hydrogen_orbital_image(qnum_n, qnum_l, qnum_m, scale, npix=2048, cmap='viridis', save=False):
    """
    creates (and, optionally, saves) image of hydrogen orbital with given quantun numbers at a given scale

    qnum_n: integer, 0 <= qnum_n
    qnum_l: integer, 0 <= qnum_l <= qnum_n-1
    qnum_m: integer, -qnum_l <= qnum_m <= qnum_l
    scale: real, > 0, 9e-10 is a good starting point for low quantum numbers
    npix:  integer, width/height of image (in pixels)
    cmap:  matplotlib colormap
    save:  boolean, if True saves image with filname in format hydrogen_atom_NLM_scale.png
    """

    # check quantum numbers are in correct range
    err_msg = "Quantum numbers n, l and m must be integers with n = 0,1,2,...; l = 0,1,...,n-1; m = -l,...,l"
    if not all([int(x) == x for x in [qnum_n, qnum_l, qnum_m]]):
        raise ValueError(err_msg)
    if not (qnum_n >= 0 and (0 <= qnum_l <= qnum_n - 1) and -qnum_l <= qnum_m <= qnum_l):
        raise ValueError(err_msg)

    x_axis = np.linspace(-scale, scale, npix)
    y_axis = np.linspace(-scale, scale, npix)

    xx, yy = np.meshgrid(x_axis, y_axis)
    img_data = psi_xy(qnum_n, qnum_l, qnum_m, xx, yy)

    if save:
        filename = 'hydrogen_atom_{}{}{}_{}.png'.format(qnum_n, qnum_l, qnum_m, scale)
    else:
        filename = None
    create_image(img_data, cmap, filename)

# TODO: it would be good to implenent an autoscale option to select an appropriate scale for the orbital image without input from the user.
# I can think of a hacky way to do this. Perhaps better to use analytic properties of big_psi though?