Electron cloud

"""A hydrogen orbital in stipple, sampled from its probability density."""

import math
from typing import Literal

import numpy as np
from numpy.typing import NDArray

from walldye import (
    ACCENT,
    ACCENT_2,
    ACCENT_4,
    BG_ALT,
    MUTED,
    UI_ALT,
    UI_HI,
    Canvas,
    NpRng,
    P,
    Params,
    Rect,
    Vec,
    design,
    knob,
    smoothstep,
)

type Arr = NDArray[np.float64]
type Orbital = Literal["3dz2", "4f"]


class Cloud(Params):
    orbital: Literal["3dz2", "4f"] = knob(
        default="3dz2", doc="3d z-squared (two lobes and a ring) or a planar 4f (six lobes)"
    )


VARIANTS = {"4f": Cloud(orbital="4f")}

# n, where |psi|^2 peaks (Bohr radii), px per Bohr radius, half-width of the sampled cube (Bohr)
ORBITALS: dict[Orbital, tuple[int, tuple[float, float, float], float, float]] = {
    "3dz2": (3, (0.0, 0.0, 6.0), 24.0, 20.0),
    "4f": (4, (12.0, 0.0, 0.0), 14.0, 32.0),
}
TILT = math.radians(16)  # camera elevation: a ring around the z axis reads as an ellipse
DMIN = (4.6, 3.0)  # spacing between dots, px: sparse fringe to dense core
GAMMA = 0.6  # <1 flattens each lobe's interior so the silhouette carries the shape
CUT = (0.045, 0.06)  # soft |psi|^2 floor (share of peak) that draws the lobe edges
# Per tone, sparse fringe to dense core: column density floor, paint, dot diameter.
FLOORS = (0.0, 0.22, 0.5, 0.66, 0.82)
TONES = (UI_ALT, UI_HI, ACCENT_4, ACCENT_2, ACCENT)
DOTS = (2.2, 2.3, 2.5, 2.8, 3.2)
ARMS = (560, 500)  # axis half-lengths, px, unless the canvas edge comes first
TICK, REACH = 5, 480  # Bohr radii between axis ticks; px they reach from the nucleus at most


def angular(orbital: Orbital, x: Arr, y: Arr, z: Arr) -> Arr:
    """The orbital's real solid harmonic, a polynomial of degree n - 1 whose zeros are its nodes."""
    if orbital == "4f":
        return x * (x * x - 3 * z * z)  # r^3 cos 3phi in the xz plane: six lobes face the viewer
    return 3 * z * z - (x * x + y * y + z * z)


def density(orbital: Orbital, x: Arr, y: Arr, z: Arr) -> Arr:
    """|psi|^2 with a nodeless radial part, unnormalized: the harmonic squared times e^(-2r/n)."""
    n = ORBITALS[orbital][0]
    return angular(orbital, x, y, z) ** 2 * np.exp(-2 * np.sqrt(x * x + y * y + z * z) / n)


def accept(orbital: Orbital, x: Arr, y: Arr, z: Arr) -> Arr:
    """Dot probability: |psi|^2 flattened by GAMMA and cut off softly at the CUT isosurface."""
    px, py, pz = (np.array([q]) for q in ORBITALS[orbital][1])
    v = density(orbital, x, y, z) / density(orbital, px, py, pz)
    return v**GAMMA * smoothstep(CUT[0], CUT[1], v)


def sample(orbital: Orbital, rng: NpRng, count: int, extent: float) -> Arr:
    """`count` points, (count, 3), drawn from `accept` by rejection in a cube of half-width `extent`."""
    pts: list[Arr] = []
    while sum(len(p) for p in pts) < count:
        q = rng.uniform(-extent, extent, (3, 500000))
        pts.append(q[:, rng.uniform(0, 1, q.shape[1]) < accept(orbital, q[0], q[1], q[2])].T)
    return np.concatenate(pts)[:count]


def column_density(orbital: Orbital, u: Arr, v: Arr, extent: float) -> Arr:
    """Line-of-sight integral of `accept` behind each screen point (u right, v up, in Bohr radii),
    normalized to its maximum."""
    w = np.linspace(-extent, extent, 160)[:, None]
    c, s = math.cos(TILT), math.sin(TILT)
    col = accept(orbital, u + 0 * w, -v * s + w * c, v * c + w * s).sum(0)
    return col / col.max()


def thin(xy: Arr, dmin: Arr) -> list[int]:
    """Greedy thinning: the indices of the points kept, dropping point i when one already kept
    lies closer than `dmin[i]`."""
    cell = float(dmin.max())
    grid: dict[tuple[int, int], list[int]] = {}
    keep: list[int] = []
    pts: list[list[float]] = xy.tolist()
    for i, ((x, y), d) in enumerate(zip(pts, dmin.tolist(), strict=True)):
        gx, gy, d2 = int(x // cell), int(y // cell), d * d
        near = (grid.get((gx + a, gy + b), ()) for a in (-1, 0, 1) for b in (-1, 0, 1))
        if all((x - pts[j][0]) ** 2 + (y - pts[j][1]) ** 2 >= d2 for c in near for j in c):
            grid.setdefault((gx, gy), []).append(i)
            keep.append(i)
    return keep


@design(aspects="any", variants=VARIANTS)
def draw(s: Canvas[Cloud]) -> None:
    orbital = s.params.orbital
    scale, extent = ORBITALS[orbital][2:]
    # right of center on a landscape screen, a little below the middle on a portrait one
    home = (0.6354, 0.5) if s.landscape else (0.5, 0.55)
    c = s.pick(landscape=(0.6354, 0.5), portrait=(0.5, 0.55))

    x, y, z = sample(orbital, s.np_rng(7), 40000, extent).T
    u, v = x, z * math.cos(TILT) - y * math.sin(TILT)
    dens = column_density(orbital, u, v, extent)
    xy = np.stack([u * scale, -v * scale], 1)  # px from the nucleus, y down
    k = thin(xy, DMIN[0] + (DMIN[1] - DMIN[0]) * np.sqrt(dens))
    xy, dens = xy[k], dens[k]

    # the axes are fit to the screen once, from the nucleus's own place, and keep that length
    # wherever it moves
    cx, cy = home[0] * s.w, home[1] * s.h
    hx = min(ARMS[0], cx - 40, s.w - cx - 40)
    hy = min(ARMS[1], cy - 40, s.h - cy - 40)
    axis = P().M(-hx, 0).H(hx).M(0, -hy).V(hy)
    ticks = int(REACH / (TICK * scale))
    for i in range(-ticks, ticks + 1):
        if i:
            t = i * TICK * scale
            axis.M(t, -5).V(5).M(-5, t).H(5)
    # half-widths that hold the axes, their ticks and every dot with its round cap
    ex = max(hx, ticks * TICK * scale + 0.6, float(np.abs(xy[:, 0]).max()) + DOTS[-1] / 2) + 1
    ey = max(hy, ticks * TICK * scale + 0.6, float(np.abs(xy[:, 1]).max()) + DOTS[-1] / 2) + 1

    tone = np.searchsorted(FLOORS, dens, side="right") - 1
    with s.motif(c, Rect(-ex, -ey, 2 * ex, 2 * ey)):
        s.stroke(axis, BG_ALT, 1.2)
        for i, (paint, width) in enumerate(zip(TONES, DOTS, strict=True)):
            dots = P()
            for px, py in xy[tone == i].tolist():
                dots.M(px, py).H(px)
            s.stroke(dots, paint, width, cap="round")
        s.fill(P().circle(Vec(0, 0), 3.5), MUTED)