Emission nebula

"""An emission nebula with dust lanes, dithered along a Hilbert curve by Riemersma's method."""

import numpy as np
from numpy.typing import NDArray
from scipy.ndimage import map_coordinates, zoom

from walldye import (
    ACCENT,
    ACCENT_3,
    ACCENT_5,
    ACCENT_7,
    BG,
    MUTED,
    UI,
    UI_HI,
    Canvas,
    NpRng,
    P,
    Rect,
    Vec,
    design,
)
from walldye.field import cells, falloff, gauss
from walldye.geom import scatter
from walldye.pixel import Pixels, dither

type Field = NDArray[np.float64]

CELL = 3
# The nebula is computed on a fixed window of cells, the same in every aspect; the dithered
# field has faded to nothing well inside its edges, so the window never shows as a cut.
COLS, ROWS = 426, 360
CORE_AT = (660, 440)  # the core's offset inside the window
LANDSCAPE, PORTRAIT = (0.677, 0.407), (0.62, 0.36)  # the core's default place
TILT = -0.35  # the nebula's long axis, radians
LANE_AT = (-400, 460)  # a point on the dust lane, from the core
SIDE = 400  # the nebula's side of the sky starts this far left of the core
# Index 0 is empty sky; 1-4 the dithered gas, faintest first; 5-7 the stars.
PALETTE = (None, ACCENT_7, ACCENT_5, ACCENT_3, ACCENT, UI, UI_HI, MUTED)
STAR_UI, STAR_UI_HI, STAR_MUTED = 5, 6, 7
STAR_DENSITY = 32 / (860 * 1000)  # candidates per square unit of open sky
THIN_STARS = 9  # over the nebula's side of a whole window


def value_noise(rng: NpRng, cols: int, rows: int, feature: int) -> Field:
    """Smooth cubic value noise in about [-1, 1], (rows, cols), with features about `feature`
    cells wide."""
    g = rng.uniform(-1, 1, (rows // feature + 4, cols // feature + 4))
    return zoom(g, feature, order=3)[:rows, :cols]


def fbm(s: Canvas, key: int, feature: int, octaves: int = 4) -> Field:
    """Unit-variance fBm over the window; octave o comes from stream key + o."""
    f = np.zeros((ROWS, COLS))
    for o in range(octaves):
        f += 0.55**o * value_noise(s.np_rng(key + o), COLS, ROWS, max(2, feature >> o))
    return f / f.std()


def field(s: Canvas, dx: Field, dy: Field) -> Field:
    """Gas density in [0, 1] at offsets (dx, dy) from the core, one value per window cell."""
    wx, wy = fbm(s, 1, 60), fbm(s, 2, 60)
    jj, ii = np.mgrid[0:ROWS, 0:COLS].astype(float)

    def warp(f: Field) -> Field:
        return map_coordinates(f, [jj + wy * 22, ii + wx * 22], mode="nearest")

    gas = warp(fbm(s, 3, 110, 3))
    dust = np.abs(warp(fbm(s, 4, 150, 3)))  # ridged: dark lanes along zero crossings
    grain = fbm(s, 5, 7, 2)  # keeps every region off a quantization plateau
    ca, sa = np.cos(TILT), np.sin(TILT)

    def radius(ax: float, ay: float, wobble: float) -> Field:
        u, v = dx / ax, dy / ay
        return np.hypot(u * ca - v * sa, u * sa + v * ca) + wobble * gas

    veil = falloff(radius(540, 380, 0.3), 1, 1.3)  # wide thin gas: mostly the faintest step
    body = falloff(radius(300, 220, 0.25), 1, 1.5)  # the dense part, kept small
    hot = gauss(np.hypot(dx, dy), 70 * np.sqrt(2))
    v = veil * (0.42 + 0.1 * gas) + 0.28 * body + 0.24 * hot
    # one straight dust lane, tilted and passing well below-right of the core
    lane_dy = LANE_AT[1] - (dx - LANE_AT[0]) * 0.62 + wy * 45
    lane = gauss((dy - lane_dy) / (55 + 20 * gas), 1)
    # the core itself stays whole: no dust ridges splitting it into two pockets
    dust_k = 1 - 0.6 * hot
    v *= (1 - dust_k * 0.4 * (1 - np.clip(dust / 0.5, 0, 1))) * (1 - 0.8 * lane)
    v += 0.08 * grain * np.clip(v * 6, 0, 1) * (1 - 0.8 * hot)
    return np.clip(v, 0, 1)


def overlap(a: Rect, b: Rect) -> Rect | None:
    """The intersection of two rects, or None when they do not overlap."""
    x0, y0 = max(a.x, b.x), max(a.y, b.y)
    x1, y1 = min(a.x1, b.x1), min(a.y1, b.y1)
    return Rect(x0, y0, x1 - x0, y1 - y0) if x1 > x0 and y1 > y0 else None


@design(aspects="any")
def draw(s: Canvas) -> None:
    def corner(x: float, y: float) -> tuple[int, int]:
        """The window's top-left for a core at (x, y), snapped to whole cells."""
        return round((x - CORE_AT[0]) / CELL) * CELL, round((y - CORE_AT[1]) / CELL) * CELL

    # The core on the right third of a landscape screen, the left kept as open sky; high on a
    # portrait one, bleeding off the right edge.
    at = s.pick(landscape=LANDSCAPE, portrait=PORTRAIT)
    core = Vec(*corner(at.x, at.y)) + CORE_AT
    fx, fy = LANDSCAPE if s.landscape else PORTRAIT
    hx, hy = corner(fx * s.w, fy * s.h)  # the window where the core sits by default
    win = Rect(-CORE_AT[0], -CORE_AT[1], COLS * CELL, ROWS * CELL)
    xs, ys = cells(win, CELL)
    v = field(s, xs, ys)
    gas = dither(v, 5, method="riemersma")

    def plot(px: Pixels, p: NDArray[np.float64], index: int) -> None:
        px.grid[int(p[1] // CELL), int(p[0] // CELL)] = index

    # stars in the open sky: everywhere but the nebula's side of its window
    sky = Pixels(s.w // CELL, s.h // CELL, PALETTE)
    r = s.rng(9)
    open_sky = s.inset(40)
    side = Rect(hx + CORE_AT[0] - SIDE, hy, win.x1 + SIDE, win.h)
    n = round(STAR_DENSITY * open_sky.w * open_sky.h)
    for p in scatter(n, open_sky, r, accept=lambda q: not side.contains(q), tries=1):
        plot(sky, p, STAR_MUTED if r.random() < 0.2 else STAR_UI_HI)
    # where the nebula sits by default, its gas gives way to the stars it would cover
    px = Pixels(COLS, ROWS, PALETTE)
    px.grid[:, :] = gas
    i0, j0 = hx // CELL, hy // CELL
    a0, b0 = max(i0, 0), max(j0, 0)
    a1, b1 = min(i0 + COLS, sky.cols), min(j0 + ROWS, sky.rows)
    px.grid[b0 - j0 : b1 - j0, a0 - i0 : a1 - i0][sky.grid[b0:b1, a0:a1] > 0] = 0
    # More stars behind the nebula's side, seen once it moves away. The side is trimmed to whole
    # cells, so where the nebula sits by default it hides no part of the stars around it.
    quiet = Rect(-(SIDE // CELL) * CELL, win.y, win.x1 + (SIDE // CELL) * CELL, win.h)
    behind = overlap(Rect(hx + CORE_AT[0] + quiet.x, hy, quiet.w, quiet.h), open_sky)
    if behind is not None:
        r2 = s.rng(10)
        for p in scatter(round(STAR_DENSITY * behind.w * behind.h), behind, r2, tries=1):
            plot(sky, p, STAR_MUTED if r2.random() < 0.2 else STAR_UI_HI)
    sky.draw(s, CELL)

    # fainter stars over the nebula's side, only where the gas is thin (the faintest step at most)
    def thin(q: Vec) -> bool:
        i, j = int((q.x - hx) // CELL), int((q.y - hy) // CELL)
        return bool(gas[j, i] <= 1 and v[j, i] < 0.3)

    thin_rect = Rect(hx + 60, hy + 30, win.w - 88, win.h - 60)
    near = overlap(thin_rect, s.inset(30))
    if near is not None:
        n = round(THIN_STARS * near.w * near.h / (thin_rect.w * thin_rect.h))
        for p in scatter(n, near, r, accept=thin):
            plot(px, p - (hx, hy), STAR_UI_HI if r.random() < 0.3 else STAR_UI)

    with s.motif(core, win, handle=Rect(-300, -230, 600, 460), cell=CELL):
        s.fill(P().rect(*quiet), BG)  # the nebula's side keeps the open sky's stars out
        px.draw(s, CELL, (win.x, win.y))