Home

Demo

from k1lib.imports import *

Purpose is to plot electron cloud orbitals around a simple hydrogen atom using spherical harmonic equations. I've longed to produce these plots by myself. Now it's finally here, and it's glorious.

Eqns taken from https://en.wikipedia.org/wiki/Table_of_spherical_harmonics

Quick recap is that assuming we can divide our wave function like so: $\psi_{n,l,m}(r,\theta,\phi)=R_{n,l}(r)y^m_l(\theta,\phi)$, where $n, l, m$ are quantum numbers. Vary these numbers to get different differential equations, and solve them to get the full wave function $\psi$.

from math import pi; from torch import sin, cos
r2 = lambda x: x**0.5; e = lambda x: math.e**x
pi2 = 2*pi; pi4 = 4*pi; ep0 = 8.85418782e-12
h = 6.62607015e-34; hbar = h/pi2
me = 9.1093837e-31; qe = 1.60217663e-19
a0 = (pi4*ep0*hbar**2)/(me*qe**2)
tensor = torch.tensor

def y(m, l):
    if l == 0: return lambda theta, phi: 0.5/r2(pi)
    elif l == 1:
        if  m == -1: return lambda theta, phi: 0.5*r2(3/pi2)*e(-1j*phi)*sin(theta)
        elif m == 0: return lambda theta, phi: 0.5*r2(3/pi)*cos(theta)
        elif m == 1: return lambda theta, phi: -0.5*r2(3/pi2)*e(1j*phi)*sin(theta)
    elif l == 2:
        if   m == -2: return lambda theta, phi: 0.25*r2(15/pi2)*e(-2j*phi)*sin(theta)**2
        elif m == -1: return lambda theta, phi: 0.5 *r2(15/pi2)*e(-1j*phi)*sin(theta)*cos(theta)
        elif m == 0:  return lambda theta, phi: 0.25*r2(5/pi)*(3*cos(theta)**2-1)
        elif m == 1:  return lambda theta, phi: -0.5 *r2(15/pi2)*e(1j*phi)*sin(theta)*cos(theta)
        elif m == 2:  return lambda theta, phi: 0.25*r2(15/pi2)*e(2j*phi)*sin(theta)**2
    elif l == 3:
        if   m == -3: return lambda theta, phi: 0.125*r2(35/pi)*e(-3j*phi)*sin(theta)**3
        elif m == -2: return lambda theta, phi: 0.25*r2(105/pi2)*e(-2j*phi)*sin(theta)**2*cos(theta)
        elif m == -1: return lambda theta, phi: 0.125*r2(21/pi)*e(-1j*phi)*sin(theta)*(5*cos(theta)**2-1)
        elif m == 0:  return lambda theta, phi: 0.25*r2(7/pi)*(5*cos(theta)**3-3*cos(theta))
        elif m == 1:  return lambda theta, phi: -0.125*r2(21/pi)*e(1j*phi)*sin(theta)*(5*cos(theta)**2-1)
        elif m == 2:  return lambda theta, phi: 0.25*r2(105/pi2)*e(2j*phi)*sin(theta)**2*cos(theta)
        elif m == 3:  return lambda theta, phi: -0.125*r2(35/pi)*e(3j*phi)*sin(theta)**3

def R(n, l, Z=1, a0=a0):
    if n == 1 and l == 0:
        return lambda r: 2*(Z/a0)**1.5*e(-Z*r/a0)
    elif n == 2 and l == 0:
        return lambda r: 1/r2(8)*(Z/a0)**1.5*(2-Z*r/a0)*e(-Z*r/(2*a0))
    elif n == 2 and l == 1:
        return lambda r: 1/r2(24)*(Z/a0)**1.5*(2*Z*r/a0)*e(-Z*r/(2*a0))

def psi(n, l, m, Z=1, a0=a0):
    _y = y(m, l); _R = R(n, l)
    return lambda r, theta, phi: _y(theta, phi)*_R(r)
p = psi(2, 1, 0)
p(tensor(1e-9), tensor(1), tensor(2))

tensor(4.1689e+11)

def psii(n, l, m, Z=1, a0=a0):
    _y = y(m, l); return lambda r, theta, phi: _y(theta, phi)/(r**n)

def gen(bounds=4e-10, n=50, p=psi(2, 1, 0), scales=[1, 1, 1]):
    sx, sy, sz = scales
    x = torch.linspace(-bounds*sx, bounds*sx, n)[:, None, None].expand(n, n, n)
    y = torch.linspace(-bounds*sy, bounds*sy, n)[None, :, None].expand(n, n, n)
    z = torch.linspace(-bounds*sz, bounds*sz, n)[None, None, :].expand(n, n, n)
    xyProj = r2(x**2+y**2); radius = r2(x**2+y**2+z**2)
    theta = torch.arcsin(xyProj/radius); phi = torch.arccos(x/xyProj)
    return abs(p(radius, theta, phi))**2

def plot(p=psi(2, 1, 0), offset=0, bounds=4e-10, n=50, scales=[1, 1, 1]):
    heatmap = gen(bounds, n, p, scales); mean = heatmap.mean()
    std = torch.tensor(heatmap.reshape([n**3]) | op().item().all() | sort(None) | filt(op() != inf) | toList()).std()
    plt.k3d(8).march(heatmap.numpy(), mean + std*offset, edgecolor=[1, 0, 0, 0.05]);

gen(n=300).sum()*(2*4e-10/50)**3

tensor(806.4459)

Lowest electron energy state

plot(psi(1, 0, 0));

This looks like a pi orbital:

plot(psi(2, 1, 0));

These look like donuts:

plot(psi(2, 1, 1));

This appears to be a duplicate of the one above, because they're degenerate states!

plot(psi(2, 1, -1));

And this is the simple spherical case, but just higher in energy

plot(psi(2, 0, 0));

Ignoring R

Because R's formulas are kinda hard to find, and I saw some sites only publish graphs with just Y only, so here goes. Actually, now I'm dividing everything by r, because without it, there's no way to control scales to your will so that it displays nicely.

s orbital:

plot(psii(1, 0, 0))

P-z orbital:

plot(psii(1.2, 1, 0))

P-x and P-y orbitals:

plot(psii(1, 1, 1))

d orbitals:

plot(psii(1.5, 2, 0))

plot(psii(1, 2, 1))

plot(psii(1.1, 2, 2))

plot(psii(1.5, 3, 0))

# thumbnail
plot(psii(1.1, 3, 1))

plot(psii(1, 3, 2))

plot(psii(1.2, 3, 3))

Damn, these are quite beautiful. And they do match with other sources.