Hydrogenic Orbitals v1

Made with Chatgpt 5, Python 3, Google Colab notebook, and DajaVu Sans.

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Orbitals up to n = 8 (204 states), square tiles + fixed 24px DejaVu Sans labels with AA,
with TOP_MARGIN and ROW_SPACING controls, using all 24-bit RGB colors exactly once.
"""

import math, time, numpy as np
from PIL import Image, ImageDraw, ImageFont

# ----------------------------
# Layout controls
# ----------------------------
WIDTH, HEIGHT = 4096, 4096
GRID_COLS, GRID_ROWS = 17, 12
TOP_MARGIN   = 40   # px – adds empty space at the very top
ROW_SPACING  = 20   # px – uniform gap between rows (affects whitespace areas)

# ----------------------------
# Physics / render controls
# ----------------------------
N_MAX = 8
Z_SAMPLES = 17
GAMMA = 0.70
BLOCK_SIZE = 262144
SEED = 1337

# Typography (fixed)
FONT_SIZE = 24
FONT_PATHS = [
    "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
    "/usr/local/share/fonts/DejaVuSans.ttf",
    "/usr/share/fonts/truetype/dejavu/DejaVuSans-Book.ttf",
    "DejaVuSans.ttf",
]
LABEL_VERT_PAD   = 3    # top padding inside the label band
LABEL_SIDE_PAD   = 8    # horizontal padding (still centered)
LABEL_BOX_HEIGHT = 36   # fixed label band height (fits 24px + pads)
LABEL_GAP        = 6    # gap between tile bottom and label band

np.random.seed(SEED)

# ----------------------------
# Math (no SciPy)
# ----------------------------
def double_factorial(n: int) -> int:
    if n <= 0: return 1
    res = 1
    for k in range(n, 0, -2): res *= k
    return res

def assoc_legendre(l: int, m: int, x: np.ndarray) -> np.ndarray:
    m = abs(m)
    if l < m: return np.zeros_like(x, np.float32)
    x = x.astype(np.float32)
    Pmm = ((-1)**m) * double_factorial(2*m - 1) * (1.0 - x*x)**(0.5*m)
    if l == m: return Pmm
    Pm1m = x * (2*m + 1) * Pmm
    if l == m+1: return Pm1m
    p2, p1 = Pmm, Pm1m
    for L in range(m+2, l+1):
        p = ((2*L-1)*x*p1 - (L+m-1)*p2) / (L-m)
        p2, p1 = p1, p
    return p1

def real_sph_harm_sq(l: int, m: int, cos_theta: np.ndarray, phi: np.ndarray) -> np.ndarray:
    mabs = abs(m)
    num = math.factorial(l - mabs); den = math.factorial(l + mabs)
    Nlm = math.sqrt((2*l + 1)/(4*math.pi) * (num/den))
    Plm = assoc_legendre(l, mabs, cos_theta)
    if m == 0:
        Y = Nlm * Plm
    else:
        fac = math.sqrt(2.0) * Nlm * Plm
        Y = fac * (np.cos(mabs*phi) if m>0 else np.sin(mabs*phi))
    return (Y*Y).astype(np.float32)

def radial_envelope(n: int, l: int, r: np.ndarray) -> np.ndarray:
    rr = np.clip(r, 0.0, 1.0).astype(np.float32)
    nr = max(0, n - l - 1)
    cent = np.power(rr + 1e-6, 2*l)
    k = 3.0 / max(1, n)
    expo = np.exp(-2.0 * k * rr)
    nodes = np.cos(np.pi * nr * rr)**2 if nr>0 else 1.0
    return (cent * expo * nodes).astype(np.float32)

def tile_orbital_intensity(n: int, l: int, m: int, side: int, z_samples: int) -> np.ndarray:
    xs = np.linspace(-1.0, 1.0, side, dtype=np.float32)
    ys = np.linspace(-1.0, 1.0, side, dtype=np.float32)
    X, Y = np.meshgrid(xs, ys)
    intensity = np.zeros((side, side), dtype=np.float32)
    zs = np.linspace(-1.0, 1.0, z_samples, dtype=np.float32)
    XY2 = X*X + Y*Y
    phi = (np.arctan2(Y, X) + 2*np.pi) % (2*np.pi)
    for z in zs:
        r2 = XY2 + z*z
        inside = r2 <= 1.0
        r = np.sqrt(np.maximum(r2, 1e-9))
        cos_theta = np.zeros_like(r, np.float32); cos_theta[inside] = z / r[inside]
        A = np.zeros_like(r, np.float32); A[inside] = real_sph_harm_sq(l, m, cos_theta[inside], phi[inside])
        R = np.zeros_like(r, np.float32); R[inside] = radial_envelope(n, l, r[inside])
        intensity += A * R
    if intensity.max() > 0: intensity /= intensity.max()
    return intensity.astype(np.float32)

# ----------------------------
# Font + AA helpers
# ----------------------------
def load_dejavu(size: int) -> ImageFont.FreeTypeFont:
    last_err = None
    for p in FONT_PATHS:
        try:
            f = ImageFont.truetype(p, size)
            if hasattr(f, "getmetrics"):
                return f
        except Exception as e:
            last_err = e
            continue
    raise RuntimeError(
        "DejaVu Sans not found. Install DejaVuSans.ttf and update FONT_PATHS. "
        f"Last error: {last_err}"
    )

def render_text_alpha(text: str, max_w: int, max_h: int, side_pad: int, top_pad: int) -> np.ndarray:
    """Fixed 24px rendering; top-aligned inside the label band. Returns uint8 alpha array."""
    font = load_dejavu(FONT_SIZE)
    # measure
    scratch = Image.new("L", (max_w, max_h), 0)
    drw = ImageDraw.Draw(scratch)
    bbox = drw.textbbox((0, 0), text, font=font, anchor="lt")
    w = bbox[2] - bbox[0]
    # horizontal center; vertical top align
    x = max(0, (max_w - w)//2)
    y = max(0, top_pad)
    img = Image.new("L", (max_w, max_h), 0)
    drw = ImageDraw.Draw(img)
    drw.text((x, y), text, fill=255, font=font)
    return np.array(img, dtype=np.uint8)

# ----------------------------
# Layout helpers
# ----------------------------
def split_evenly(total: int, parts: int):
    base = total // parts; rem = total % parts
    return [base + 1 if i < rem else base for i in range(parts)]

def enumerate_orbitals(n_max: int):
    for n in range(1, n_max+1):
        for l in range(0, n):
            for m in range(-l, l+1):
                yield (n, l, m)

# ----------------------------
# Color space + matching
# ----------------------------
def rgb_cube_arrays():
    r = np.repeat(np.arange(256, dtype=np.uint16), 256*256)
    g = np.tile(np.repeat(np.arange(256, dtype=np.uint16), 256), 256)
    b = np.tile(np.arange(256, dtype=np.uint16), 256*256)
    Y = (0.2126*r + 0.7152*g + 0.0722*b).astype(np.float32)
    rf = (r/255.0).astype(np.float32); gf = (g/255.0).astype(np.float32); bf = (b/255.0).astype(np.float32)
    maxc = np.maximum(np.maximum(rf, gf), bf); minc = np.minimum(np.minimum(rf, gf), bf)
    v = maxc; delta = maxc - minc
    s = np.zeros_like(v); nz = v > 0; s[nz] = delta[nz] / v[nz]
    h = np.zeros_like(v); mask = delta > 1e-20
    mr = mask & (maxc == rf); mg = mask & (maxc == gf); mb = mask & (maxc == bf)
    h[mr] = ((gf[mr] - bf[mr]) / delta[mr]) % 6.0
    h[mg] = ((bf[mg] - rf[mg]) / delta[mg]) + 2.0
    h[mb] = ((rf[mb] - gf[mb]) / delta[mb]) + 4.0
    h = (h / 6.0) % 1.0
    return r.astype(np.uint8), g.astype(np.uint8), b.astype(np.uint8), h, s, v, Y

def circ_dist(a: np.ndarray, b: float) -> np.ndarray:
    d = np.abs(a - b); return np.minimum(d, 1.0 - d)

def pref_metric_for_group(h: np.ndarray, s: np.ndarray, L: int) -> np.ndarray:
    if L == 0: return (s + 0.02 * circ_dist(h, 0.5)).astype(np.float32)
    centers = {1:[0.00,0.08,0.17], 2:[0.33,0.45], 3:[0.66,0.83]}.get(L, [0.25,0.58,0.92])
    d = circ_dist(h, centers[0])
    for c in centers[1:]: d = np.minimum(d, circ_dist(h, c))
    return (d - 0.05*s).astype(np.float32)

# ----------------------------
# Build intensities + masks + label alpha (with TOP_MARGIN & ROW_SPACING)
# ----------------------------
def build_canvas_intensity_masks_labels():
    # Column widths stay evenly split
    col_widths = split_evenly(WIDTH, GRID_COLS)
    x_off = np.cumsum([0] + col_widths[:-1])

    # Row heights with explicit top margin and fixed spacing
    BOTTOM_MARGIN = ROW_SPACING  # guarantee at least this much
    usable_height = HEIGHT - TOP_MARGIN - BOTTOM_MARGIN - ROW_SPACING * (GRID_ROWS - 1)
    row_heights = split_evenly(usable_height, GRID_ROWS)

    # y offsets: start at TOP_MARGIN, then add row height + spacing per row
    y_off = [TOP_MARGIN]
    for rh in row_heights[:-1]:
      y_off.append(y_off[-1] + rh + ROW_SPACING)
    y_off = np.array(y_off)


    # Square tile side must fit width and (row height - label band - gap)
    side_by_row = [max(1, rh - LABEL_BOX_HEIGHT - LABEL_GAP) for rh in row_heights]
    side = min(min(side_by_row), min(col_widths))

    I = np.zeros((HEIGHT, WIDTH), np.float32)
    Lgroup = np.full((HEIGHT, WIDTH), 255, np.uint8)      # ℓ for tiles
    tile_mask = np.zeros((HEIGHT, WIDTH), bool)
    label_alpha = np.zeros((HEIGHT, WIDTH), np.uint8)     # AA for labels

    orbitals = list(enumerate_orbitals(N_MAX))
    assert len(orbitals) == GRID_COLS * GRID_ROWS

    idx = 0
    for r in range(GRID_ROWS):
        for c in range(GRID_COLS):
            n, l, m = orbitals[idx]
            rh = row_heights[r]; cw = col_widths[c]
            y0, y1 = y_off[r], y_off[r] + rh
            x0, x1 = x_off[c], x_off[c] + cw

            # Tile packed at the TOP of the row area
            ys = y0
            ye = ys + side
            xs = x0 + (cw - side)//2
            xe = xs + side

            # Orbital intensity
            tile = tile_orbital_intensity(n, l, m, side, Z_SAMPLES)
            I[ys:ye, xs:xe] = tile
            Lgroup[ys:ye, xs:xe] = l
            tile_mask[ys:ye, xs:xe] = True

            # Label band directly under the tile
            lbl_y0 = ye + LABEL_GAP
            lbl_y1 = lbl_y0 + LABEL_BOX_HEIGHT
            text = f"n={n} ℓ={l} m={m if m>=0 else '−'+str(abs(m))}"
            alpha = render_text_alpha(
                text=text,
                max_w=cw,
                max_h=LABEL_BOX_HEIGHT,
                side_pad=LABEL_SIDE_PAD,
                top_pad=LABEL_VERT_PAD
            )
            la = label_alpha[lbl_y0:lbl_y1, x0:x1]
            label_alpha[lbl_y0:lbl_y1, x0:x1] = np.maximum(la, alpha)

            idx += 1

    # Normalize/gamma over tile pixels only
    tile_vals = I[tile_mask]; vmax = tile_vals.max() if tile_vals.size else 0.0
    if vmax > 0: I[tile_mask] = (tile_vals / vmax) ** GAMMA
    return I, Lgroup, tile_mask, label_alpha, side

# ----------------------------
# Color assignment (labels → orbitals → whitespace)
# ----------------------------
def rgb_match_and_render(I: np.ndarray, Lgroup: np.ndarray, tile_mask: np.ndarray, label_alpha: np.ndarray):
    H, W = I.shape; N = H*W
    I_flat = I.reshape(-1); L_flat = Lgroup.reshape(-1)
    tmask = tile_mask.reshape(-1); la = label_alpha.reshape(-1)

    lbl_idx = np.nonzero(la > 0)[0]
    orb_idx = np.nonzero(tmask)[0]
    wht_idx = np.nonzero(~(tmask | (la > 0)))[0]
    N_lbl, N_orb, N_wht = len(lbl_idx), len(orb_idx), len(wht_idx)
    print(f"Orbital pixels: {N_orb:,} | Label pixels: {N_lbl:,} | Whitespace: {N_wht:,}")

    rC, gC, bC, hC, sC, vC, YC = rgb_cube_arrays()
    assert rC.size == N

    col_order = np.argsort(YC, kind="stable")

    # Labels: darkest → highest alpha (AA)
    col_lbl = col_order[:N_lbl]
    rest = col_order[N_lbl:]
    sort_lbl = np.argsort(-la[lbl_idx], kind="stable")
    lbl_sorted = lbl_idx[sort_lbl]
    outR = np.empty(N, np.uint8); outG = np.empty(N, np.uint8); outB = np.empty(N, np.uint8)
    outR[lbl_sorted] = rC[col_lbl]; outG[lbl_sorted] = gC[col_lbl]; outB[lbl_sorted] = bC[col_lbl]

    # Orbitals: global intensity → luminance with ℓ-hue nudges
    col_orb = rest[:N_orb]; col_wht = rest[N_orb:]
    orb_sorted = orb_idx[np.argsort(I_flat[orb_idx], kind="stable")]

    blocks = (N_orb + BLOCK_SIZE - 1) // BLOCK_SIZE
    for bi in range(blocks):
        s = bi*BLOCK_SIZE; e = min((bi+1)*BLOCK_SIZE, N_orb); Lb = e - s
        pix_block = orb_sorted[s:e]; cols_block = col_orb[s:e]
        groups = L_flat[pix_block]; h_block = hC[cols_block]; s_block = sC[cols_block]
        used = np.zeros(Lb, bool); chosen = {}
        uniq, counts = np.unique(groups, return_counts=True)
        for L in uniq[np.argsort(-counts)]:
            need = counts[uniq==L][0]; avail = np.where(~used)[0]
            P = pref_metric_for_group(h_block[avail], s_block[avail], int(L))
            chosen_rel = avail[np.argpartition(P, min(need-1, len(avail)-1))[:need]]
            used[chosen_rel] = True; chosen[int(L)] = chosen_rel
        assert used.all()
        for L in uniq:
            rel = chosen[int(L)]; cols_g = cols_block[rel]; pix_g = pix_block[groups==L]
            if pix_g.size>1: pix_g = pix_g[np.argsort(I_flat[pix_g], kind="stable")]
            if cols_g.size>1: cols_g = cols_g[np.argsort(YC[cols_g], kind="stable")]
            outR[pix_g] = rC[cols_g]; outG[pix_g] = gC[cols_g]; outB[pix_g] = bC[cols_g]

        if (bi+1)%8==0 or bi==blocks-1:
            print(f"  Blocks {bi+1}/{blocks} ({100*(bi+1)/blocks:.1f}%)")

    # Whitespace: leftover colors, shuffled (noise)
    cols_wht_shuf = col_wht[np.random.permutation(len(col_wht))]
    outR[wht_idx] = rC[cols_wht_shuf]; outG[wht_idx] = gC[cols_wht_shuf]; outB[wht_idx] = bC[cols_wht_shuf]

    img = np.stack([outR, outG, outB], axis=1).reshape(HEIGHT, WIDTH, 3)
    return img

# ----------------------------
# Main
# ----------------------------
def main():
    print("Building intensities, masks, and label AA with margins/spacings...")
    I, Lgroup, tile_mask, label_alpha, side = build_canvas_intensity_masks_labels()
    print(f"Tile side = {side}px | Label alpha pixels = {int((label_alpha>0).sum()):,}")

    print("Assigning colors (labels → orbitals → whitespace)...")
    img = rgb_match_and_render(I, Lgroup, tile_mask, label_alpha)

    expected_sum = 65536 * (255 * 256 // 2)  # 2,137,014,272
    chsum = img.reshape(-1, 3).sum(axis=0, dtype=np.uint64)
    print(f"Channel sums R,G,B = {tuple(int(x) for x in chsum)}; expected {expected_sum}")

    ts = time.strftime("%Y%m%d_%H%M%S")
    outfile = f"orbitals_n8_204_allRGB_labels24_margins_{ts}.png"
    print(f"Saving {outfile} ...")
    Image.fromarray(img, "RGB").save(outfile, compress_level=0)
    print("Done.")

if __name__ == "__main__":
    main()