Skip to content

Instantly share code, notes, and snippets.

@tkmharris
Last active December 24, 2019 16:37
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
Star You must be signed in to star a gist
Save tkmharris/239d634d2bc32ae3dc8628cf9821260c to your computer and use it in GitHub Desktop.
Functions for creating images of hydrogen orbitals
Display the source blob
Display the rendered blob
Raw
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
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?
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment