diff options
Diffstat (limited to 'mapgen/hydrology.py')
| -rw-r--r-- | mapgen/hydrology.py | 231 |
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} |