"""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)