worldgen: initial public history

1 files changed, 231 insertions, 0 deletions

diff --git a/mapgen/hydrology.py b/mapgen/hydrology.py
new file mode 100644
index 0000000..d7adde1
--- /dev/null
+++ b/mapgen/hydrology.py
@@ -0,0 +1,231 @@
+"""Stage `hydrology`: depression filling, lakes vs endorheic basins, discharge, rivers, Strahler order."""
+from __future__ import annotations
+
+import numpy as np
+from scipy import sparse
+from scipy.sparse import csgraph
+
+from .config import params
+from .crust import RIFT
+from .graph import accumulate, components, distance_to, priority_flood, receiver_levels, steepest_receivers
+from .noise import fbm
+from .pipeline import StageError
+
+DEFAULTS = {"fill_eps_m": 0.01, "min_depth_m": 1.0, "river_min_km3_yr": 2.0,
+            "salt_flat_max_p_mm": 300.0, "salt_flat_fraction": 0.2, "dry_lake_fraction": 0.05,
+            # lake beds: deepest point = k × area^exp m (Earth-like: ≈95 m at
+            # 1,000 km², ≈230 m at 20,000 km², ≈610 m at 500,000 km²) × 0.5–2 (noise at the lake), × rift_x in
+            # rifts, × arid_x for dry terminal lakes; never shallower than lake_min_m (refinement keeps ≥ 15 m)
+            "lake_depth_k": 12.0, "lake_depth_exp": 0.3, "lake_min_m": 25.0, "lake_max_m": 1800.0,
+            "lake_rift_x": 2.5, "lake_arid_x": 0.4, "lake_arid_p_mm": 400.0, "lake_shore": 0.25}
+
+
+def strahler(recv, levels, river):
+    n = len(recv)
+    order = np.zeros(n, np.int8)
+    mx = np.zeros(n, np.int8)
+    cnt = np.zeros(n, np.int16)
+    for lv in reversed(levels):
+        cells = lv[river[lv]]
+        if len(cells) == 0:
+            continue
+        order[cells] = np.where(cnt[cells] >= 2, mx[cells] + 1, np.maximum(mx[cells], 1))
+        nonroot = cells[recv[cells] != cells]
+        r, o = recv[nonroot], order[nonroot]
+        np.maximum.at(mx, r, o)
+        np.add.at(cnt, r, (o == mx[r]).astype(np.int16))
+    return order
+
+
+def _groups(lab):
+    idx = np.argsort(lab, kind="stable")
+    ls = lab[idx]
+    start = np.searchsorted(ls, 0)
+    idx, ls = idx[start:], ls[start:]
+    cuts = np.flatnonzero(np.diff(ls)) + 1
+    return np.split(idx, cuts) if len(idx) else []
+
+
+def _leaves(recv, lab, x, limit=100000):
+    """Does the flow from depression cell x leave its depression for good (not back in over shallow ground)?"""
+    own, y = lab[x], recv[x]
+    for _ in range(limit):
+        if lab[y] == own:
+            return False
+        if lab[y] >= 0 or recv[y] == y:
+            return True
+        y = recv[y]
+    return True
+
+
+def _leaves_all(recv, lab, levels, limit=100000):
+    """_leaves for every cell at once: walking down from recv[x], the first cell that is in a depression or a root
+    (stop) and how many steps away it is; x leaves unless that cell is in x's own depression (or the walk is longer
+    than limit steps)."""
+    n = len(recv)
+    stop, d = np.arange(n), np.zeros(n, np.int64)
+    for lv in levels[1:]:                                    # roots first: a cell's receiver is done before it
+        free = lv[lab[lv] < 0]
+        stop[free] = stop[recv[free]]
+        d[free] = d[recv[free]] + 1
+    y = recv
+    return (d[y] >= limit) | (lab[stop[y]] != lab)
+
+
+def _route_inside(g, dep, lab, recv, targets):
+    """Within each depression, point every cell along the shortest intra-depression path to its target cell."""
+    e = dep[g.src] & dep[g.dst] & (lab[g.src] == lab[g.dst])
+    m = sparse.csr_matrix((g.edge_km[e], (g.src[e], g.dst[e])), shape=(g.n, g.n))
+    _, pred, _ = csgraph.dijkstra(m, directed=False, indices=targets, min_only=True, return_predecessors=True)
+    out = recv.copy()
+    inside = dep & (pred >= 0)
+    out[inside] = pred[inside]
+    return out
+
+
+def _sweep(recv, levels, water, outlets, cap):
+    """Accumulate flow upstream→downstream; at each spilling-lake outlet remove up to `cap` (lake evaporation)."""
+    acc = np.asarray(water, dtype=np.float64).copy()
+    loss = np.zeros(len(acc))
+    is_out = np.zeros(len(acc), bool)
+    is_out[outlets] = True
+    cap_cell = np.zeros(len(acc))
+    cap_cell[outlets] = cap
+    buf = np.zeros(len(acc))                    # per level only the receivers change (graph.accumulate): same floats
+    if len(levels) > 1:
+        acc += 0.0
+    for lv in reversed(levels[1:]):
+        push = acc[lv].copy()
+        o = is_out[lv]
+        if o.any():
+            cells = lv[o]
+            lost = np.minimum(acc[cells], cap_cell[cells])
+            loss[cells] = lost
+            push[o] = acc[cells] - lost
+        r = recv[lv]
+        buf[r] = 0.0
+        np.add.at(buf, r, push)
+        acc[r] = acc[r] + buf[r]
+    return acc, loss
+
+
+def lake_levels(z, zf, lab, lake, lake_id, endo):
+    """Each lake cell's water level (NaN elsewhere): a spilling lake stands at its spill height; a terminal lake at the
+    lowest ground of its basin it does not cover (the next cell to flood). Carving the beds leaves both unchanged."""
+    level = np.full(len(z), np.nan)
+    m = int(lab.max()) + 1 if len(lab) and lab.max() >= 0 else 0
+    dry = np.full(m, np.inf)                                    # per depression: its lowest uncovered ground
+    sel = (lab >= 0) & ~lake
+    np.minimum.at(dry, lab[sel], z[sel])
+    for c in _groups(np.where(lake, lake_id, -1)):
+        k = lab[c[0]]
+        level[c] = dry[k] if endo[c[0]] and k >= 0 and np.isfinite(dry[k]) else zf[c].min()
+    return level
+
+
+def lake_beds(g, z, level, lake, lake_id, endo, rift, p_ann, seed, P, only=None):
+    """Lake beds carved below their level (heights unchanged elsewhere): each lake's deepest point from its area
+    (P lake_*), noise keyed by where the lake lies (not its id), the depth rising from lake_shore × that at the shore
+    to all of it at the cell farthest from shore; never above the ground (min), never shallower than lake_min_m.
+    only: carve just the lakes touching these cells (eras: where events changed the ground)."""
+    z = np.asarray(z, dtype=np.float64).copy()
+    if not lake.any():
+        return z
+    shore = distance_to(g, ~lake) if (~lake).any() else np.full(g.n, g.spacing_km)
+    for c in _groups(np.where(lake, lake_id, -1)):
+        if only is not None and not only[c].any():
+            continue
+        area = g.area_km2[c].sum()
+        ctr = g.xyz[c].mean(0)
+        ctr = ctr / max(np.linalg.norm(ctr), 1e-12)
+        d = P["lake_depth_k"] * area ** P["lake_depth_exp"] * 2.0 ** float(fbm(ctr[None], seed + 4421, 3, 6.0)[0] * 1.4)
+        if rift[c].mean() > 0.3:
+            d *= P["lake_rift_x"]
+        if endo[c[0]] and p_ann[c].mean() < P["lake_arid_p_mm"]:
+            d *= P["lake_arid_x"]
+        d = float(np.clip(d, P["lake_min_m"], P["lake_max_m"]))
+        t = shore[c] / max(shore[c].max(), 1e-9)
+        prof = np.maximum(d * (P["lake_shore"] + (1 - P["lake_shore"]) * t ** 0.6), P["lake_min_m"])
+        z[c] = np.minimum(z[c], level[c] - prof)
+    return z
+
+
+def run(ctx) -> dict:
+    g = ctx.grid
+    P = params(ctx.cfg, "hydrology", DEFAULTS)
+    z, p_ann, pet = ctx.need("elevation_eroded_m", "P_ann", "PET")
+    z = z.astype(np.float64)
+    ar = np.arange(g.n)
+    ocean = np.asarray(ctx.data["ocean"]) if "ocean" in ctx.data else z <= 0
+    aet = p_ann / np.sqrt(1.0 + (p_ann / np.maximum(pet, 1e-6)) ** 2)      # Pike (1964)
+    runoff = np.where(ocean, 0.0, np.maximum(p_ann - aet, 0.0))
+    water = runoff * g.area_km2 * 1e-6                                       # km³/yr per cell
+    zf = priority_flood(g, z, ocean, P["fill_eps_m"])
+    depth = zf - z
+    dep = ~ocean & (depth > P["min_depth_m"])
+    lab = components(g, dep)
+    recv1, _, _ = steepest_receivers(g, zf)
+    recv1 = np.where(ocean, ar, recv1)
+    levels1 = receiver_levels(recv1)
+    q1 = accumulate(recv1, levels1, water)
+    evap_net = np.maximum(pet - p_ann, 0.0) * g.area_km2 * 1e-6              # full-lake evaporation, km³/yr
+
+    groups = _groups(lab)
+    leaves = _leaves_all(recv1, lab, levels1) if len(groups) else None
+    outlets = np.zeros(len(groups), np.int64)
+    terminals = np.zeros(len(groups), np.int64)
+    cap = np.zeros(len(groups))
+    for k, c in enumerate(groups):
+        ext = c[lab[recv1[c]] != lab[c[0]]]
+        ext = ext[leaves[ext]]
+        outlets[k] = ext[np.argmax(q1[ext])] if len(ext) else c[np.argmax(q1[c])]
+        terminals[k] = c[np.argmin(z[c])]
+        cap[k] = evap_net[c].sum()
+
+    # 1) every depression drains to its spill outlet; one upstream-first sweep decides spill vs endorheic
+    recv_a = _route_inside(g, dep, lab, recv1, outlets) if len(groups) else recv1
+    acc_a, _ = _sweep(recv_a, receiver_levels(recv_a), water, outlets, cap)
+    inflow = acc_a[outlets]
+    spill = inflow >= cap
+    # 2) endorheic depressions drain to their single lowest cell instead
+    targets = np.where(spill, outlets, terminals)
+    recv = _route_inside(g, dep, lab, recv1, targets) if len(groups) else recv1
+    recv[terminals[~spill]] = terminals[~spill]
+    try:
+        levels = receiver_levels(recv)
+    except ValueError as e:
+        raise StageError(f"hydrology: {e}") from e
+    q, loss = _sweep(recv, levels, water, outlets[spill], cap[spill])
+
+    lake = np.zeros(g.n, bool)
+    endo = np.zeros(g.n, bool)
+    salt = np.zeros(g.n, bool)
+    lake_id = np.full(g.n, -1, np.int32)
+    for k, c in enumerate(groups):
+        if spill[k]:
+            lake[c] = True
+            lake_id[c] = k
+            continue
+        endo[c] = True
+        cs = c[np.argsort(z[c])]
+        nl = int(np.searchsorted(np.cumsum(evap_net[cs]), inflow[k]))
+        if inflow[k] > 0:
+            nl = max(nl, 1)                         # the terminal always holds some water
+        lake[cs[:nl]] = True
+        lake_id[cs[:nl]] = k
+        if nl == 0 or (nl < P["dry_lake_fraction"] * len(cs) and p_ann[c].mean() < P["salt_flat_max_p_mm"]):
+            a = np.cumsum(g.area_km2[cs])
+            ns = max(nl + 1, int(np.searchsorted(a, P["salt_flat_fraction"] * a[-1])))
+            salt[cs[nl:ns]] = True
+    river = ~ocean & ~lake & (q >= P["river_min_km3_yr"])
+    level = lake_levels(z, zf, lab, lake, lake_id, endo)
+    rift = np.asarray(ctx.data["age_class"]) == RIFT if "age_class" in ctx.data else np.zeros(g.n, bool)
+    zb = lake_beds(g, z, level, lake, lake_id, endo, rift, p_ann, ctx.seed, P, ctx.data.get("lake_carve"))
+    cut = np.asarray(ctx.data.get("lake_cut_m", 0.0), dtype=np.float64) + (z - zb)   # sea masks see the uncut ground
+    z = zb
+    depth = zf - z
+    dtype = np.asarray(ctx.data["elevation_eroded_m"]).dtype
+    return {"elevation_eroded_m": z.astype(dtype), "lake_level_m": level.astype(np.float32),
+            "lake_cut_m": np.broadcast_to(cut, z.shape).astype(np.float32), "z_filled_m": zf, "recv": recv, "discharge_km3_yr": q, "runoff_mm": runoff, "aet_mm": aet,
+            "lake": lake, "lake_id": lake_id, "endorheic": endo, "salt_flat": salt, "river": river,
+            "strahler": strahler(recv, levels, river), "depression_depth_m": depth, "lake_loss_km3_yr": loss}