worldmap-viewer: initial public history

1 files changed, 320 insertions, 0 deletions

diff --git a/rivers.py b/rivers.py
new file mode 100644
index 0000000..9a1688f
--- /dev/null
+++ b/rivers.py
@@ -0,0 +1,320 @@
+"""River valleys for the deep zoom: a graded water level per river cell, a valley style from
+landform, ground, rock and climate, and the valley cut into procedural heights. Procedural (`idea`), deterministic,
+continuous across tiles. Lives outside mapgen/ so changing it never invalidates the build cache."""
+from __future__ import annotations
+
+import math
+import threading
+
+import numpy as np
+from scipy.spatial import cKDTree
+
+from mapgen.noise import value_noise
+
+THETA = 0.45            # graded slope ∝ Q^-θ (discharge stands in for drainage area)
+KS = 2.0                # m/km at Q = 1 km³/yr on a craton in average rock
+KNICK_M_KM = 15.0       # steepest drop below lakes and steep reaches: a cataract, not a cliff
+MAX_REACH_KM = 40.0
+FLOOR_MAX_KM = 15.0     # the widest valley floor (half-width beyond the channel)
+FADE_M = 2000.0         # a valley narrower than 2 px is lifted by up to this much: it fades in, never pops
+CHUNK_KM = 100.0        # spread-out queries (long profiles) are answered in compact chunks
+AGE_KS = {"craton (3g era)": 1.0, "pre-Lightening orogen": 1.5, "post-Lightening orogen": 3.0, "rift": 1.5,
+          "Lightening basalt province": 2.0, "collapse scar": 1.5, "overshoot volcano": 3.0, "oceanic": 1.0}
+ROCK = {"granite/gneiss": 1.6, "metamorphic": 1.6, "basalt": 1.8, "andesite": 1.4, "limestone": 1.3,
+        "sandstone/shale": 0.7, "oceanic basalt": 1.5}                  # hardness: graded steepness and caps
+CAP_M = {"plain": 150, "hills": 400, "mountains": 1200, "plateau": 1500, "rift valley": 600, "escarpment": 800,
+         "volcanic arc": 900, "volcanic massif": 900, "basalt plateau": 1200, "dunes": 60, "badlands": 300, "ocean": 0}
+STYLE = {"plain": (15, 40, 1.0), "hills": (5, 150, 0.5), "mountains": (1.5, 600, 0.0), "plateau": (1.2, 1500, 0.0),
+         "rift valley": (5, 300, 0.3), "escarpment": (1.5, 800, 0.0), "volcanic arc": (1.5, 600, 0.0),
+         "volcanic massif": (1.5, 600, 0.0), "basalt plateau": (1.2, 1500, 0.0), "dunes": (8, 60, 0.5),
+         "badlands": (3, 400, 0.2), "ocean": (15, 40, 1.0)}          # landform → (floor × half-width, wall m/km, meander)
+GROUND = {"floodplain": (30, 20, 1.5), "delta": (30, 20, 1.5), "wetland": (20, 20, 1.2), "bog": (20, 20, 1.2)}
+DETAIL_M = {"ocean": 300, "plain": 50, "hills": 240, "mountains": 900, "plateau": 120, "rift valley": 300,
+            "escarpment": 500, "volcanic arc": 700, "volcanic massif": 700, "basalt plateau": 120, "dunes": 60,
+            "badlands": 240}   # ≈ 2 × the procedural detail's amplitude (tiles.DEEP): how far above the cell it reaches
+WALL_ROCK = {"granite/gneiss": 1.5, "metamorphic": 1.5, "basalt": 1.5, "oceanic basalt": 1.5, "limestone": 1.4,
+             "andesite": 1.2, "sandstone/shale": 0.6}
+STATE = ("cap", "level", "seg_a", "seg_b", "a_xyz", "b_xyz", "level_a", "level_b", "seg_len", "half_w", "floor",
+         "wall", "meander", "amp", "width", "reach")      # what a RiverNet computes (serve cache: saved, not recomputed)
+
+
+def half_width_km(q_km3_yr):
+    """Channel half-width: 4·√Q(m³/s) metres (hydraulic geometry)."""
+    return 0.004 * np.sqrt(np.asarray(q_km3_yr, dtype=np.float64) * 31.7)
+
+
+DRAW_MIN_PX = 0.02                        # streams narrower than this share of a pixel are not drawn (refined
+ALWAYS_HW = float(half_width_km(2.0))     # 0.2 km³/yr streams appear from ≈ 0.5 km/px); world rivers (≥ 2 km³/yr)
+                                          # are drawn at every zoom
+
+
+def valley_style(landform: str, ground: str, rock: str, rain_mm: float):
+    """(floor × half-width, wall m/km, meander factor) for a river cell."""
+    floor, wall, meander = GROUND.get(ground) or STYLE.get(landform, (5, 150, 0.5))
+    k = WALL_ROCK.get(rock, 1.0)
+    dry = 1.4 if rain_mm < 500 else 0.8 if rain_mm > 1500 else 1.0
+    return floor / k, wall * k * dry, meander
+
+
+def valley_surface(d_km, level, half_w, floor, wall, rough=0.0):
+    """Height (m) of a valley at distance d from its centreline: the water level in the channel, a floor 1–3 m above
+    it, then walls rising at `wall` m/km. `rough` (the terrain's own procedural detail, m) roughens the walls,
+    fading in over the first 0.5 km above the floor, so they look eroded rather than planar."""
+    d = np.asarray(d_km, dtype=np.float64)
+    fl = np.clip((d - half_w) / np.maximum(floor, 1e-9), 0.0, 1.0)
+    up = np.maximum(d - half_w - floor, 0.0)
+    wall_m = np.maximum(0.0, wall * up + np.clip(up / 0.5, 0.0, 1.0) * rough)   # rough walls never dip below the floor
+    return np.where(d <= half_w, level, level + 1.0 + 2.0 * fl + wall_m)
+
+
+def _rownorm(p):
+    """np.linalg.norm(p, axis=1) for (n, 3), without its per-call overhead: the same squares summed in the same order."""
+    return np.sqrt(p[:, 0] * p[:, 0] + p[:, 1] * p[:, 1] + p[:, 2] * p[:, 2])
+
+
+def _vnorm(v):
+    """np.linalg.norm of one 3-vector (it is sqrt(v·v))."""
+    return math.sqrt(v.dot(v))
+
+
+def densify(a, b, t0, t1, step):
+    """Unit vectors along a→b (a short great-circle segment) at the lattice t = k·step/|b−a|, k integer, t in [t0, t1]."""
+    d = b - a
+    L = _vnorm(d)
+    if L == 0:
+        return a[None, :]
+    k = np.arange(np.ceil(t0 * L / step), np.floor(t1 * L / step) + 1)
+    t = np.append(k * step / L, [] if t1 < 1 else [1.0])   # the segment's end joins the next one
+    if len(t) == 0:
+        return np.zeros((0, 3))
+    p = a[None, :] * (1 - t[:, None]) + b[None, :] * t[:, None]
+    return p / _rownorm(p)[:, None]
+
+
+class RiverNet:
+    def __init__(self, a: dict, legends: dict, radius_km: float, levels=None, ground=None):
+        """levels: given water levels (a refined area); ground: the heights the valleys are cut from (default the
+        cells' surface; a refined area passes its valley-shoulder heights, so walls reach up to them)."""
+        self.R = float(radius_km)
+        riv = np.asarray(a["river"]).astype(bool)
+        recv = np.asarray(a["recv"]).astype(np.int64)
+        ocean, lake = np.asarray(a["ocean"]).astype(bool), np.asarray(a["lake"]).astype(bool)
+        z = np.asarray(a["z_surface_m"], dtype=np.float64)
+        zf = np.asarray(a["z_filled_m"], dtype=np.float64)
+        q = np.asarray(a["discharge_km3_yr"], dtype=np.float64)
+        xyz = np.asarray(a["g_xyz"], dtype=np.float64)
+        pick = lambda leg, key, table, default: np.array([table.get(n, default) for n in legends[leg]])[np.asarray(a[key])]
+        rock = pick("lithology", "lithology", ROCK, 1.0)
+        ks = KS * pick("age_class", "age_class", AGE_KS, 1.0) * rock
+        self.cap = pick("landform", "landform", CAP_M, 400) * rock / 1.6
+        n = len(riv)
+        r = np.where(riv)[0]
+        dist = np.zeros(n)
+        dist[r] = np.linalg.norm(xyz[r] - xyz[recv[r]], axis=1) * self.R
+        if levels is None:
+            lev = np.full(n, np.nan)
+            lev[lake] = zf[lake]
+            lev[ocean] = 0.0
+            for i in r[np.argsort(zf[r], kind="stable")]:           # upstream: graded on the receiver's level, capped
+                j = recv[i]
+                b = lev[j] if np.isfinite(lev[j]) else z[j]
+                lev[i] = min(z[i], max(b + ks[i] * q[i] ** -THETA * dist[i], z[i] - self.cap[i]))
+            src = np.where((riv | lake) & ~ocean & (recv != np.arange(n)))[0]
+            d_src = np.linalg.norm(xyz[src] - xyz[recv[src]], axis=1) * self.R
+            for i, di in zip(src[np.argsort(-zf[src], kind="stable")], d_src[np.argsort(-zf[src], kind="stable")]):
+                j = recv[i]                                         # downstream: drops limited to a cataract
+                if riv[j]:
+                    lev[j] = max(lev[j], min(z[j], lev[i] - KNICK_M_KM * di))
+            for i in r[np.argsort(-zf[r], kind="stable")]:          # never rising downstream
+                j = recv[i]
+                if riv[j] and lev[j] > lev[i]:
+                    lev[j] = lev[i]
+        else:   # given water levels (a refined area: erosion has already cut the valleys; receivers keep theirs)
+            lev = np.asarray(levels, dtype=np.float64).copy()
+            lev[lake] = zf[lake]
+            lev[ocean] = 0.0
+        self.level = lev
+        out = np.where(lake & riv[recv] & (recv != np.arange(n)))[0]     # lake outlets flow on to their river
+        dist[out] = np.linalg.norm(xyz[out] - xyz[recv[out]], axis=1) * self.R
+        a_cell = np.concatenate([r, out])
+        b_cell = recv[a_cell]
+        style_cell = np.concatenate([r, recv[out]])                      # an outlet looks like the river it feeds
+        self.seg_a, self.seg_b = a_cell, b_cell
+        self.a_xyz, self.b_xyz = xyz[a_cell], xyz[b_cell]
+        self.level_a = np.where(lake[a_cell], zf[a_cell], lev[a_cell])
+        self.level_b = np.where(ocean[b_cell], 0.0, np.where(lake[b_cell], zf[b_cell], lev[b_cell]))
+        self.seg_len = dist[a_cell]
+        self.half_w = half_width_km(q[style_cell])
+        names = {k: np.asarray(legends[k], dtype=object) for k in ("landform", "ground", "lithology")}
+        st = np.array([valley_style(names["landform"][a["landform"][i]], names["ground"][a["ground"][i]],
+                                    names["lithology"][a["lithology"][i]], float(a["P_ann"][i])) for i in style_cell]).reshape(-1, 3)
+        self.floor = np.minimum(st[:, 0] * self.half_w, FLOOR_MAX_KM)
+        self.wall, self.meander = st[:, 1], st[:, 2]
+        self.amp = 2.5 * 2 * self.half_w * self.meander                     # meander swing (km)
+        zg = z if ground is None else np.asarray(ground, dtype=np.float64)
+        depth = np.maximum(np.maximum(zg[a_cell], self.level_a) - self.level_a, 0.0)
+        detail = pick("landform", "landform", DETAIL_M, 300)[style_cell]
+        self.width = 2 * (self.half_w + self.floor + depth / self.wall)            # as seen at the cell's mean ground
+        self.reach = self.half_w + self.floor + np.minimum(MAX_REACH_KM, (depth + detail) / self.wall)   # + detail relief
+        self._tree, self._tree_lock = None, threading.Lock()   # built on first use (see tree)
+        self.max_extent = float(np.max(self.seg_len / 2 + self.reach + self.amp)) if len(a_cell) else 0.0
+
+    @property
+    def tree(self):
+        """KD-tree over segment midpoints (unit vectors), built on first use; None without segments."""
+        if self._tree is None and len(self.seg_a):
+            with self._tree_lock:
+                if self._tree is None:
+                    mid = np.asarray(self.a_xyz, dtype=np.float64) + np.asarray(self.b_xyz, dtype=np.float64)
+                    self._tree = cKDTree(mid / np.linalg.norm(mid, axis=1, keepdims=True))
+        return self._tree
+
+    def state(self):
+        """(arrays, meta) that from_state turns back into the same net without recomputing it."""
+        return {k: getattr(self, k) for k in STATE}, {"R": self.R, "max_extent": self.max_extent}
+
+    @classmethod
+    def from_state(cls, arrays: dict, meta: dict) -> "RiverNet":
+        net = cls.__new__(cls)
+        for k in STATE:
+            setattr(net, k, np.asarray(arrays[k]))      # memory maps as plain arrays: same data, cheap indexing
+        net.R, net.max_extent = float(meta["R"]), float(meta["max_extent"])
+        net._tree, net._tree_lock = None, threading.Lock()
+        return net
+
+    # --- geometry ---------------------------------------------------------------------------------------------
+    def candidates(self, center, radius_km):
+        """Segments whose valley can reach any point within radius_km of a unit vector."""
+        if self.tree is None:
+            return []
+        idx = np.asarray(self.tree.query_ball_point(center, 2 * np.sin(min(np.pi, (radius_km + self.max_extent) / self.R) / 2)),
+                         dtype=np.int64)
+        if not len(idx):
+            return []
+        a, ab = self.a_xyz[idx], self.b_xyz[idx] - self.a_xyz[idx]
+        tc = np.clip(np.einsum("ij,ij->i", center - a, ab) / np.maximum(np.einsum("ij,ij->i", ab, ab), 1e-30), 0.0, 1.0)
+        d = np.linalg.norm(center - (a + tc[:, None] * ab), axis=1) * self.R
+        return idx[d <= radius_km + self.reach[idx] + self.amp[idx] + 1.0].tolist()
+
+    def segment_points(self, s, center, within_km, spacing_km):
+        """Centreline of segment s near a unit vector: its own global lattice (so any query gets the same points),
+        meandered. Returns (unit vectors, t along the segment)."""
+        a, b = self.a_xyz[s], self.b_xyz[s]
+        ab = b - a
+        L = max(_vnorm(ab), 1e-12)
+        tc = float(np.dot(center - a, ab) / max(np.dot(ab, ab), 1e-30))
+        span = (within_km + self.amp[s]) / self.R / L
+        t0, t1 = max(0.0, tc - span), min(1.0, tc + span)
+        if t0 > t1:
+            return np.zeros((0, 3)), np.zeros(0)
+        p = densify(a, b, t0, t1, max(spacing_km, 0.0005) / self.R)
+        if len(p) == 0:
+            return p, np.zeros(0)
+        t = np.clip(np.dot(p - a, ab) / max(np.dot(ab, ab), 1e-30), 0.0, 1.0)
+        if self.amp[s] > 0 and self.seg_len[s] > 0:
+            side = np.empty_like(p)                    # np.cross(ab, p), the same products and differences
+            side[:, 0] = ab[1] * p[:, 2] - ab[2] * p[:, 1]
+            side[:, 1] = ab[2] * p[:, 0] - ab[0] * p[:, 2]
+            side[:, 2] = ab[0] * p[:, 1] - ab[1] * p[:, 0]
+            side /= np.maximum(_rownorm(side), 1e-30)[:, None]
+            lam = 11 * 2 * self.half_w[s]
+            off = self.amp[s] * value_noise(p * (self.R / lam), 4242) * np.sin(np.pi * t)
+            p = p + side * (off / self.R)[:, None]
+        return p / _rownorm(p)[:, None], t
+
+    def _chunks(self, xyz):
+        """Split query points into compact groups (halving along their widest axis): (indices, centre, radius km)."""
+        out, stack = [], [np.arange(len(xyz))]
+        while stack:
+            ix = stack.pop()
+            c = xyz[ix].mean(axis=0)
+            c = c / max(np.linalg.norm(c), 1e-12) if np.linalg.norm(c) > 1e-9 else xyz[ix[0]]
+            r = float(np.max(np.linalg.norm(xyz[ix] - c, axis=1))) * self.R
+            if r <= CHUNK_KM or len(ix) <= 16:
+                out.append((ix, c, 2 * self.R * np.arcsin(min(1.0, r / (2 * self.R)))))   # chord → arc
+            else:
+                p = xyz[ix]
+                o = np.argsort(p[:, int(np.argmax(p.max(axis=0) - p.min(axis=0)))], kind="stable")
+                stack += [ix[o[: len(ix) // 2]], ix[o[len(ix) // 2:]]]
+        return out
+
+    def parts(self, xyz, px_km, spacing_km=None):
+        """Per query point, the lowest valley surface over every valley that reaches it, in parts:
+        (base = level + floor rise + fade lift, wall rise, roughness share). base is +inf where no valley reaches.
+        No lower valley undercuts a reach's banks: each reach's floor is a lower bound, falling away at its wall slope
+        beyond it (a reach doubling back below itself leaves a terrace, not dry pits below the water beside it)."""
+        n = len(xyz)
+        base, wall_up, rough_f = np.full(n, np.inf), np.zeros(n), np.zeros(n)
+        best, bank = np.full(n, np.inf), np.full(n, -np.inf)
+        spacing = spacing_km or px_km
+        for ix, c, r in self._chunks(xyz):
+            q = xyz[ix]
+            for s in self.candidates(c, r):
+                alpha = float(np.clip(self.width[s] / px_km - 1.0, 0.0, 1.0))
+                if alpha <= 0:
+                    continue                                   # narrower than a pixel: no valley to see
+                p, t = self.segment_points(s, c, r + self.reach[s], spacing)
+                if len(p) == 0:
+                    continue
+                d, k = cKDTree(p).query(q, distance_upper_bound=self.reach[s] / self.R)
+                hit = k < len(p)
+                if not hit.any():
+                    continue
+                dk = d[hit] * self.R
+                lv = self.level_a[s] + (self.level_b[s] - self.level_a[s]) * t[k[hit]]
+                fl = np.clip((dk - self.half_w[s]) / max(self.floor[s], 1e-9), 0.0, 1.0)
+                up = np.maximum(dk - self.half_w[s] - self.floor[s], 0.0)
+                b = lv + 1.0 + 2.0 * fl + (1.0 - alpha) * FADE_M
+                w = self.wall[s] * up
+                j = ix[hit]
+                better = b + w < best[j]
+                jb = j[better]
+                best[jb], base[jb], wall_up[jb] = (b + w)[better], b[better], w[better]
+                rough_f[jb] = np.clip(up / 0.5, 0.0, 1.0)[better]
+                np.maximum.at(bank, j, b - w)
+        low = bank > best
+        base[low], wall_up[low], rough_f[low] = bank[low], 0.0, 0.0
+        return base, wall_up, rough_f
+
+    def channel(self, xyz, px_km):
+        """Channel water per query point (true width), channel as drawn (≥ 0.6 px), and the water level there."""
+        n = len(xyz)
+        ch, dr, lev = np.zeros(n, bool), np.zeros(n, bool), np.full(n, np.nan)
+        for ix, c, r in self._chunks(xyz):
+            pts, lv, hw = [], [], []
+            for s in self.candidates(c, r):
+                if self.half_w[s] < min(DRAW_MIN_PX * px_km, ALWAYS_HW):
+                    continue                                   # a stream far below a pixel wide: not drawn yet
+                reach = max(self.half_w[s], 0.6 * px_km)
+                a, ab = self.a_xyz[s], self.b_xyz[s] - self.a_xyz[s]
+                tc = np.clip(np.dot(c - a, ab) / max(np.dot(ab, ab), 1e-30), 0.0, 1.0)
+                if _vnorm(c - (a + tc * ab)) * self.R > r + reach + self.amp[s] + 1.0:
+                    continue
+                p, t = self.segment_points(s, c, r + reach, px_km / 2)
+                keep = _rownorm(p - c) * self.R <= r + reach + px_km   # only points that reach the chunk
+                p, t = p[keep], t[keep]
+                if len(p):
+                    pts.append(p)
+                    lv.append(self.level_a[s] + (self.level_b[s] - self.level_a[s]) * t)
+                    hw.append(np.full(len(p), self.half_w[s]))
+            if not pts:
+                continue
+            P, LV, HW = np.concatenate(pts), np.concatenate(lv), np.concatenate(hw)
+            bound = max(float(HW.max()), 0.6 * px_km) / self.R
+            d, k = cKDTree(P).query(xyz[ix], distance_upper_bound=bound)
+            hit = k < len(P)
+            kk = np.where(hit, k, 0)
+            dk = np.where(hit, d * self.R, np.inf)
+            ch[ix] = hit & (dk <= HW[kk])
+            dr[ix] = hit & (dk <= np.maximum(HW[kk], 0.6 * px_km))
+            lev[ix] = np.where(hit, LV[kk], np.nan)
+        return ch, dr, lev
+
+    def valleys(self, xyz, px_km, rough=None, spacing_km=None, with_base=False):
+        """Valley surface V (m; +inf where none), channel and drawn channel at query points (+ base, see parts)."""
+        base, wall_up, f = self.parts(xyz, px_km, spacing_km)
+        rg = 0.0 if rough is None else np.asarray(rough, dtype=np.float64)
+        V = base + np.maximum(0.0, wall_up + f * rg)
+        ch, dr, lev = self.channel(xyz, px_km)
+        V = np.where(ch, lev, V)
+        return (V, ch, dr, base) if with_base else (V, ch, dr)