Orbital Lab

#!/usr/bin/env python3
"""
Generate real hydrogen-orbital isosurfaces from the Schrodinger equation.

Not an artistic approximation: this evaluates the analytic bound-state solution
    psi_nlm(r,th,ph) = R_nl(r) * Y_lm(th,ph)
on a 3D grid, forms the REAL orbitals (the physicists' p_x/p_y/p_z, d_z2, d_xy
combinations), and marching-cubes an isosurface of psi at +/- a fraction of its
peak. Sign is preserved so the two lobes of a p orbital can be coloured as the
opposite phases they actually are.
"""
import json, sys, numpy as np
from scipy.special import genlaguerre, sph_harm_y
from skimage import measure

A0 = 1.0   # Bohr radius, working in atomic units

def R_nl(n, l, r):
    """Radial part of the hydrogen wavefunction."""
    rho = 2.0 * r / (n * A0)
    from math import factorial, sqrt
    norm = sqrt((2.0/(n*A0))**3 * factorial(n-l-1) / (2*n*factorial(n+l)))
    return norm * np.exp(-rho/2) * rho**l * genlaguerre(n-l-1, 2*l+1)(rho)

def Y_real(l, m, th, ph):
    """Real spherical harmonic (the orbitals chemists actually draw)."""
    if m == 0:
        return sph_harm_y(l, 0, th, ph).real
    if m > 0:
        return (np.sqrt(2) * (-1)**m * sph_harm_y(l, m, th, ph).real)
    return (np.sqrt(2) * (-1)**m * sph_harm_y(l, -m, th, ph).imag)

def psi(n, l, m, X, Y, Z):
    r = np.sqrt(X**2 + Y**2 + Z**2) + 1e-9
    th = np.arccos(np.clip(Z/r, -1, 1))     # polar
    ph = np.arctan2(Y, X)                   # azimuth
    return R_nl(n, l, r) * Y_real(l, m, th, ph)

def isosurface(n, l, m, grid=96, extent=None, iso_frac=0.30):
    if extent is None:
        extent = 6.0 * n            # orbitals grow roughly as n^2, this is ample
    g = np.linspace(-extent, extent, grid)
    X, Y, Z = np.meshgrid(g, g, g, indexing='ij')
    P = psi(n, l, m, X, Y, Z)
    peak = np.abs(P).max()
    level = iso_frac * peak

    out = []
    for sign in (+1, -1):
        field = P * sign
        if field.max() < level:
            continue
        v, f, _, _ = measure.marching_cubes(field, level=level)
        # grid index -> world coords
        v = v / (grid - 1) * (2 * extent) - extent
        out.append((v, f, sign))
    return out, peak

def build(name, n, l, m, grid=96, iso_frac=0.30, target=1.0):
    parts, peak = isosurface(n, l, m, grid=grid, iso_frac=iso_frac)
    if not parts:
        return None
    allv = np.vstack([p[0] for p in parts])
    scale = target / np.abs(allv).max()
    positions, indices, phases = [], [], []
    voff = 0
    for v, f, sign in parts:
        v = v * scale
        positions.append(v)
        indices.append(f + voff)
        phases.append(np.full(len(v), sign, dtype=np.int8))
        voff += len(v)
    P = np.vstack(positions); F = np.vstack(indices); PH = np.concatenate(phases)
    print(f'  {name:8s} n={n} l={l} m={m:+d}  lobes={len(parts)}  verts={len(P):6d}  tris={len(F):6d}')
    return {
        'name': name, 'n': n, 'l': l, 'm': m,
        'positions': [round(float(x), 4) for x in P.ravel()],
        'indices': [int(i) for i in F.ravel()],
        'phases': [int(p) for p in PH],
    }

if __name__ == '__main__':
    grid = int(sys.argv[1]) if len(sys.argv) > 1 else 96
    print(f'Solving hydrogen orbitals on a {grid}^3 grid...')
    specs = [
        ('2pz',  2, 1,  0),   # the dumbbell: the Adom logomark's own state
        ('3dz2', 3, 2,  0),   # dumbbell wrapped in a torus
        ('3dxy', 3, 2, -2),   # four-lobe clover
        ('2px',  2, 1, +1),   # dumbbell on its side
    ]
    orbs = [o for o in (build(*s, grid=grid) for s in specs) if o]
    out = '/home/adom/project/hd-launcher-mockup/assets/orbitals.json'
    with open(out, 'w') as fh:
        json.dump({'orbitals': orbs}, fh, separators=(',', ':'))
    import os
    print(f'wrote {out}  {os.path.getsize(out):,} bytes')