worldgen: initial public history

1 files changed, 477 insertions, 0 deletions

diff --git a/mapgen/graph.py b/mapgen/graph.py
new file mode 100644
index 0000000..126c71f
--- /dev/null
+++ b/mapgen/graph.py
@@ -0,0 +1,477 @@
+"""Operations on the H3 cell graph (CSR neighbours)."""
+from __future__ import annotations
+
+import heapq
+
+import numpy as np
+from scipy import sparse
+from scipy.sparse import csgraph
+from scipy.sparse import linalg as splinalg
+
+
+def nbr_mean(g, f):
+    return np.bincount(g.src, weights=f[g.dst], minlength=g.n) / g.counts
+
+
+def nbr_max(g, f):
+    out = np.asarray(f, dtype=np.float64).copy()
+    np.maximum.at(out, g.src, f[g.dst])
+    return out
+
+
+def nbr_min(g, f):
+    out = np.asarray(f, dtype=np.float64).copy()
+    np.minimum.at(out, g.src, f[g.dst])
+    return out
+
+
+def diffuse(g, f, iters: int, alpha: float = 0.5, mask=None):
+    f = np.asarray(f, dtype=np.float64)
+    for _ in range(int(iters)):
+        new = (1.0 - alpha) * f + alpha * nbr_mean(g, f)
+        f = new if mask is None else np.where(mask, new, f)
+    return f
+
+
+def laplacian_matrix(g):
+    """Graph Laplacian (per km²): Δf_i ≈ (4/k_i) Σ_j (f_j − f_i)/d_ij² (exact for a regular hex lattice)."""
+    w = 4.0 / (g.counts[g.src] * g.edge_km**2)
+    lap = sparse.csr_matrix((w, (g.src, g.dst)), shape=(g.n, g.n))
+    return lap - sparse.diags(np.asarray(lap.sum(axis=1)).ravel())
+
+
+def workers() -> int:
+    """Threads for independent jobs (seasons, components): WORLDGEN_THREADS, default 3. Each job computes exactly
+    what it would alone, so results never depend on it; numpy and scipy's sparse kernels release the GIL."""
+    import os
+    try:
+        return max(1, int(os.environ.get("WORLDGEN_THREADS", "3")))
+    except ValueError:
+        return 3
+
+
+def pmap(fn, items) -> list:
+    """[fn(x) for x in items] on up to workers() threads, in order."""
+    items = list(items)
+    if workers() <= 1 or len(items) <= 1:
+        return [fn(x) for x in items]
+    from concurrent.futures import ThreadPoolExecutor
+    with ThreadPoolExecutor(min(workers(), len(items))) as ex:
+        return list(ex.map(fn, items))
+
+
+def smooth_km(g, f, length_km: float, rtol: float = 1e-6, keep: bool = False):
+    """Resolution-independent smoothing: solve (I − L²Δ) s = f (screened Poisson, decay length ≈ L km).
+    keep: keep the matrix on g for the next call (loops over one grid; drop_smooth_cache(g) frees it)."""
+    f = np.asarray(f, dtype=np.float64)
+    scale = float(np.max(np.abs(f))) if f.size else 0.0
+    if length_km <= 0 or scale == 0.0:
+        return f.copy()
+    f = f / scale                                    # linear system: solve at unit scale (avoids breakdown)
+    cache = g.__dict__.get("_screened", {})
+    if length_km in cache:
+        A, inv_diag = cache[length_km]
+    else:
+        A = (sparse.identity(g.n, format="csr") - length_km**2 * laplacian_matrix(g)).tocsr()
+        inv_diag = 1.0 / A.diagonal()
+        if keep:
+            g.__dict__.setdefault("_screened", {})[length_km] = (A, inv_diag)
+    s, info = bicgstab_jacobi(A, f, f.copy(), inv_diag, rtol, 5000)
+    if info != 0:
+        raise ValueError(f"smooth_km: solver did not converge (info={info})")
+    return s * scale
+
+
+def _make_bicg_jit():
+    """bicgstab's vector updates fused into single passes (numba optional; WORLDGEN_NO_JIT=1 turns it off). Each
+    element gets the same operations in the same order as scipy's numpy statements; dot products, norms and sparse
+    products stay the very calls scipy makes, so the iterates — and the answer — are the same floats. (Threads were
+    tried and dropped: on a busy machine they wait more than they work.)"""
+    import os
+    if os.environ.get("WORLDGEN_NO_JIT"):
+        return None
+    try:
+        import numba
+    except ImportError:
+        return None
+
+    @numba.njit(cache=True, nogil=True)
+    def p_update(p, v, r, omega, beta):          # p -= omega*v; p *= beta; p += r
+        for i in range(len(p)):
+            p[i] = (p[i] - omega * v[i]) * beta + r[i]
+
+    @numba.njit(cache=True, nogil=True)
+    def scale(out, d, x):                        # out = d * x (the Jacobi preconditioner)
+        for i in range(len(x)):
+            out[i] = d[i] * x[i]
+
+    @numba.njit(cache=True, nogil=True)
+    def axpy_neg(r, a, v):                       # r -= a*v
+        for i in range(len(r)):
+            r[i] = r[i] - a * v[i]
+
+    @numba.njit(cache=True, nogil=True)
+    def x_update(x, alpha, phat, omega, shat):   # x += alpha*phat; x += omega*shat
+        for i in range(len(x)):
+            x[i] = (x[i] + alpha * phat[i]) + omega * shat[i]
+
+    @numba.njit(cache=True, nogil=True)
+    def axpy(x, a, v):                           # x += a*v
+        for i in range(len(x)):
+            x[i] = x[i] + a * v[i]
+
+    return p_update, scale, axpy_neg, x_update, axpy
+
+
+_bicg_jit = _make_bicg_jit()
+
+
+def _csr_matvec_into(A):
+    """mv(x, y): y = A @ x, by scipy's own kernel into a reused buffer (A @ x zero-fills a fresh array and calls the
+    same csr_matvec; a fresh 16 MB array costs more in page faults than the product)."""
+    try:
+        from scipy.sparse import _sparsetools
+        fn = _sparsetools.csr_matvec
+    except (ImportError, AttributeError):
+        fn = None
+    M, N = A.shape
+
+    def mv(x, y):
+        if fn is None:
+            y[:] = A @ x
+            return
+        y.fill(0.0)
+        fn(M, N, A.indptr, A.indices, A.data, x, y)
+    return mv
+
+
+JIT_MIN_N = 50_000          # smaller systems: scipy (thread start-up outweighs the gain; same answer either way)
+
+
+def bicgstab_jacobi(A, b, x0, inv_diag, rtol, maxiter):
+    """scipy.sparse.linalg.bicgstab(A, b, x0=x0, rtol=rtol, maxiter=maxiter, M=diag(inv_diag)) for a CSR matrix and
+    float64 vectors: the same iterates, statement by statement (scipy 1.12+'s pure-Python loop), with the vector
+    updates fused and everything written into preallocated buffers (memory-bound; fresh 16 MB temporaries cost
+    more in page faults than in arithmetic). Falls back to scipy without numba."""
+    b = np.asarray(b, dtype=np.float64).ravel()
+    if (_bicg_jit is None or A.shape[0] < JIT_MIN_N or not sparse.isspmatrix_csr(A)
+            and not isinstance(A, sparse.csr_array) or A.dtype != np.float64):
+        M = splinalg.LinearOperator(A.shape, matvec=lambda x: inv_diag * x)
+        return splinalg.bicgstab(A, b, x0=x0, rtol=rtol, maxiter=maxiter, M=M)
+    p_update, scale, axpy_neg, x_update, axpy = _bicg_jit
+    mv = _csr_matvec_into(A)
+    inv_diag = np.ascontiguousarray(inv_diag, dtype=np.float64)
+    x = np.array(x0, dtype=np.float64)
+    bnrm2 = np.linalg.norm(b)
+    atol = max(0.0, float(rtol) * float(bnrm2))
+    if bnrm2 == 0:
+        return b, 0
+    rhotol = np.finfo(x.dtype.char).eps ** 2
+    omegatol = rhotol
+    rho_prev, omega, alpha, p, v = None, None, None, None, None
+    r = b - A @ x if x.any() else b.copy()
+    rtilde = r.copy()
+    phat, shat, v, t = np.empty_like(r), np.empty_like(r), np.empty_like(r), np.empty_like(r)
+    for iteration in range(maxiter):
+        if np.linalg.norm(r) < atol:
+            return x, 0
+        rho = np.dot(rtilde, r)
+        if np.abs(rho) < rhotol:
+            return x, -10
+        if iteration > 0:
+            if np.abs(omega) < omegatol:
+                return x, -11
+            beta = (rho / rho_prev) * (alpha / omega)
+            p_update(p, v, r, omega, beta)
+        else:
+            p = r.copy()
+        scale(phat, inv_diag, p)
+        mv(phat, v)
+        rv = np.dot(rtilde, v)
+        if rv == 0:
+            return x, -11
+        alpha = rho / rv
+        axpy_neg(r, alpha, v)
+        s = r                                     # scipy copies r into s here and reads both unchanged until r -= omega*t
+        if np.linalg.norm(s) < atol:
+            axpy(x, alpha, phat)
+            return x, 0
+        scale(shat, inv_diag, s)
+        mv(shat, t)
+        omega = np.dot(t, s) / np.dot(t, t)
+        x_update(x, alpha, phat, omega, shat)
+        axpy_neg(r, omega, t)
+        rho_prev = rho
+    return x, maxiter
+
+
+def drop_smooth_cache(g) -> None:
+    """Free the matrices smooth_km keeps on g."""
+    g.__dict__.pop("_screened", None)
+
+
+def gradient(g, f):
+    """Tangent-plane gradient (f per km): (2/k) Σ_j (f_j − f_i)/d_ij · t_ij."""
+    w = (f[g.dst] - f[g.src]) / g.edge_km
+    s = np.stack([np.bincount(g.src, weights=w * g.edge_tangents[:, c], minlength=g.n) for c in range(3)], axis=1)
+    return s * (2.0 / g.counts)[:, None]
+
+
+def _csr(g, weights):
+    return sparse.csr_matrix((weights, (g.src, g.dst)), shape=(g.n, g.n))
+
+
+def nearest_source(g, sources, weights=None):
+    sources = np.asarray(sources, dtype=np.int64)
+    if len(sources) == 0:
+        return np.full(g.n, np.inf), np.full(g.n, -9999, dtype=np.int64)
+    w = g.edge_km if weights is None else weights
+    dist, _, src = csgraph.dijkstra(_csr(g, w), directed=True, indices=sources,
+                                    min_only=True, return_predecessors=True)
+    return dist, src.astype(np.int64)
+
+
+def distance_to(g, mask):
+    return nearest_source(g, np.flatnonzero(mask))[0]
+
+
+def priority_flood(g, z, sink_mask, eps: float = 0.01):
+    """Barnes (2014) priority-flood + ε: every non-sink cell gets a strictly descending path to a sink."""
+    sink_mask = np.asarray(sink_mask, dtype=bool)
+    if not sink_mask.any():
+        raise ValueError("priority_flood: no sink cells")
+    has_open = np.bincount(g.src, weights=(~sink_mask)[g.dst].astype(np.float64), minlength=g.n) > 0
+    seeds = np.flatnonzero(sink_mask & has_open)
+    z = np.asarray(z, dtype=np.float64)
+    if _flood_jit is not None:
+        return _flood_jit(z.copy(), sink_mask.copy(), np.asarray(g.nbr_ptr, np.int64), np.asarray(g.nbr_idx, np.int64),
+                          seeds.astype(np.int64), float(eps))
+    return _flood_py(z, sink_mask, g.nbr_ptr, g.nbr_idx, seeds, eps)
+
+
+def _flood_py(z, sink_mask, nbr_ptr, nbr_idx, seeds, eps):
+    zf = np.asarray(z, dtype=np.float64).tolist()
+    done = np.asarray(sink_mask).tolist()
+    ptr = np.asarray(nbr_ptr).tolist()
+    idx = np.asarray(nbr_idx).tolist()
+    heap = [(zf[i], i) for i in np.asarray(seeds).tolist()]
+    heapq.heapify(heap)
+    while heap:
+        zc, c = heapq.heappop(heap)
+        for k in range(ptr[c], ptr[c + 1]):
+            n = idx[k]
+            if not done[n]:
+                done[n] = True
+                if zf[n] < zc + eps:
+                    zf[n] = zc + eps
+                heapq.heappush(heap, (zf[n], n))
+    return np.array(zf)
+
+
+def _make_flood_jit():
+    """_flood_py compiled with numba when it is installed (optional: same heap order, same float steps, same result;
+    WORLDGEN_NO_JIT=1 turns it off)."""
+    import os
+    if os.environ.get("WORLDGEN_NO_JIT"):
+        return None
+    try:
+        import numba
+    except ImportError:
+        return None
+
+    @numba.njit(cache=True)
+    def flood(zf, done, ptr, idx, seeds, eps):
+        heap = [(zf[i], i) for i in seeds]
+        heapq.heapify(heap)
+        while len(heap):
+            zc, c = heapq.heappop(heap)
+            for k in range(ptr[c], ptr[c + 1]):
+                n = idx[k]
+                if not done[n]:
+                    done[n] = True
+                    if zf[n] < zc + eps:
+                        zf[n] = zc + eps
+                    heapq.heappush(heap, (zf[n], n))
+        return zf
+    return flood
+
+
+_flood_jit = _make_flood_jit()
+
+
+def _make_steep_jit():
+    """steepest_receivers' per-row maximum, compiled (numba optional, as the flood): the first edge in row order
+    with the largest slope, slopes computed as numpy does. ok=False (empty row, NaN): use the numpy path."""
+    import os
+    if os.environ.get("WORLDGEN_NO_JIT"):
+        return None
+    try:
+        import numba
+    except ImportError:
+        return None
+
+    @numba.njit(cache=True)
+    def steep(z, ptr, dst, edge):
+        n = len(ptr) - 1
+        first = np.empty(n, np.int64)
+        best = np.empty(n)
+        for i in range(n):
+            a, b = ptr[i], ptr[i + 1]
+            if a == b:
+                return first, best, False
+            bk = a
+            bs = (z[i] - z[dst[a]]) / edge[a]
+            if bs != bs:
+                return first, best, False
+            for k in range(a + 1, b):
+                sk = (z[i] - z[dst[k]]) / edge[k]
+                if sk != sk:
+                    return first, best, False
+                if sk > bs:
+                    bs, bk = sk, k
+            first[i], best[i] = bk, bs
+        return first, best, True
+    return steep
+
+
+_steep_jit = _make_steep_jit()
+
+
+def steepest_receivers(g, z):
+    """Per cell: the neighbour of steepest descent (first in neighbour order on ties), the slope and the edge length;
+    no way down → itself, 0, inf."""
+    ptr = np.asarray(g.nbr_ptr, np.int64)
+    if _steep_jit is not None and g.n:
+        first, s, ok = _steep_jit(np.asarray(z), ptr, np.asarray(g.nbr_idx, np.int64), g.edge_km)
+        if ok:
+            down = s > 0
+            recv = np.where(down, g.dst[first], np.arange(g.n))
+            return recv, np.where(down, s, 0.0), np.where(down, g.edge_km[first], np.inf)
+    slope = (z[g.src] - z[g.dst]) / g.edge_km
+    if g.n == 0 or not (np.diff(ptr) > 0).all() or np.isnan(slope).any():
+        order = np.lexsort((-slope, g.src))                  # general case (empty rows, NaN)
+        first = order[ptr[:-1]]
+    else:                                                    # same edge as the stable lexsort, without sorting
+        smax = np.maximum.reduceat(slope, ptr[:-1])
+        cand = np.flatnonzero(slope == smax[g.src])
+        rows = g.src[cand]
+        first = cand[np.concatenate([[True], rows[1:] != rows[:-1]])]
+    s = slope[first]
+    down = s > 0
+    ar = np.arange(g.n)
+    recv = np.where(down, g.dst[first], ar)
+    return recv, np.where(down, s, 0.0), np.where(down, g.edge_km[first], np.inf)
+
+
+def _gather(ptr, arr, sel):
+    counts = ptr[sel + 1] - ptr[sel]
+    tot = int(counts.sum())
+    if tot == 0:
+        return arr[:0]
+    starts = np.repeat(ptr[sel] - np.concatenate([[0], np.cumsum(counts)[:-1]]), counts)
+    return arr[starts + np.arange(tot)]
+
+
+def receiver_levels(recv):
+    recv = np.asarray(recv, dtype=np.int64)
+    n = len(recv)
+    ar = np.arange(n)
+    root = recv == ar
+    donors = ar[~root]
+    donors = donors[np.argsort(recv[donors], kind="stable")]
+    dptr = np.concatenate([[0], np.cumsum(np.bincount(recv[donors], minlength=n))])
+    levels, frontier, seen = [], ar[root], 0
+    while len(frontier):
+        levels.append(frontier)
+        seen += len(frontier)
+        frontier = _gather(dptr, donors, frontier)
+    if seen != n:
+        raise ValueError("receiver_levels: cycle in receivers")
+    return levels
+
+
+def accumulate(recv, levels, w):
+    """Sum w down the receiver tree. Per level only the receivers are touched (a full bincount per level costs
+    levels × cells); the sums are added in donor order from 0, as bincount does: the same floats."""
+    acc = np.asarray(w, dtype=np.float64).copy()
+    buf = np.zeros(len(acc))
+    if len(levels) > 1:
+        acc += 0.0                                    # as the first full-length add did: −0 becomes +0
+    for lv in reversed(levels[1:]):
+        r = recv[lv]
+        buf[r] = 0.0
+        np.add.at(buf, r, acc[lv])
+        acc[r] = acc[r] + buf[r]
+    return acc
+
+
+def _make_components_jit():
+    """Connected-component labels as scipy's connected_components numbers them — each component (every node not in
+    the mask is one by itself) by the order of its lowest node — by union-find over the edges, with no sparse matrix
+    (numba optional; WORLDGEN_NO_JIT=1 turns it off)."""
+    import os
+    if os.environ.get("WORLDGEN_NO_JIT"):
+        return None
+    try:
+        import numba
+    except ImportError:
+        return None
+
+    @numba.njit(cache=True, nogil=True)
+    def labels(n, src, dst, mask):
+        parent = np.arange(n)
+        for e in range(src.shape[0]):
+            a, b = src[e], dst[e]
+            if mask[a] and mask[b]:
+                while parent[a] != a:
+                    parent[a] = parent[parent[a]]
+                    a = parent[a]
+                while parent[b] != b:
+                    parent[b] = parent[parent[b]]
+                    b = parent[b]
+                if a != b:
+                    if a < b:
+                        parent[b] = a
+                    else:
+                        parent[a] = b
+        root_lab = np.full(n, -1, np.int32)
+        out = np.empty(n, np.int32)
+        count = 0
+        for v in range(n):
+            r = v
+            while parent[r] != r:
+                r = parent[r]
+            if root_lab[r] < 0:
+                root_lab[r] = count
+                count += 1
+            out[v] = root_lab[r] if mask[v] else -1
+        return out
+    return labels
+
+
+_components_jit = _make_components_jit()
+
+
+def components(g, mask):
+    mask = np.asarray(mask, dtype=bool)
+    if _components_jit is not None:
+        return _components_jit(g.n, g.src, g.dst, mask)
+    e = mask[g.src] & mask[g.dst]
+    m = sparse.csr_matrix((np.ones(int(e.sum())), (g.src[e], g.dst[e])), shape=(g.n, g.n))
+    _, lab = csgraph.connected_components(m, directed=False)
+    return np.where(mask, lab, -1)
+
+
+OCEAN_MIN_KM2 = 5.0e6
+
+
+def ocean_mask(g, z, min_area_km2: float = OCEAN_MIN_KM2):
+    """The connected world ocean plus any separate basin ≥ min_area_km2; smaller interior lows count as land."""
+    wet = np.asarray(z) <= 0
+    if not wet.any():
+        return wet
+    lab = components(g, wet)
+    area = np.bincount(lab[wet], weights=g.area_km2[wet])
+    keep = area >= min_area_km2
+    keep[np.argmax(area)] = True
+    return wet & keep[np.maximum(lab, 0)]