1 files changed, 54 insertions, 0 deletions
diff --git a/mapgen/erosion.py b/mapgen/erosion.py
new file mode 100644
index 0000000..5384988
--- /dev/null
+++ b/mapgen/erosion.py
@@ -0,0 +1,54 @@
+"""Stage `erosion`: implicit stream-power incision (Braun & Willett 2013) + km-scale hillslope smoothing."""
+from __future__ import annotations
+
+import numpy as np
+
+from .config import params
+from .crust import CRATON, LIP, OCEANIC, POST_OROGEN, PRE_OROGEN, RIFT, SCAR, VOLCANO
+from .elevation import solve_sea_level_connected
+from .graph import (accumulate, drop_smooth_cache, ocean_mask, priority_flood, receiver_levels, smooth_km,
+                    steepest_receivers)
+
+DEFAULTS = {"min_sea_km2": 5.0e6, "steps": 12, "dt_myr": 5.0, "k": 0.02, "m": 0.5, "hillslope_km": 25.0, "sediment_fill": 0.15}
+K_MULT = {OCEANIC: 0.0, CRATON: 1.5, PRE_OROGEN: 1.2, POST_OROGEN: 1.0, RIFT: 1.0, LIP: 0.6, SCAR: 0.9, VOLCANO: 0.8}
+
+
+def erode(g, z, kmult, P):
+    z = np.asarray(z, dtype=np.float64).copy()
+    ar = np.arange(g.n)
+    for _ in range(int(P["steps"])):
+        ocean = ocean_mask(g, z, P["min_sea_km2"])
+        zb = np.where(ocean, 0.0, z)                     # base level = sea level, not the seafloor
+        zf = priority_flood(g, zb, ocean)
+        zb = np.where(ocean, zb, zb + P["sediment_fill"] * (zf - zb))   # sediment infills closed basins
+        recv, _, dist = steepest_receivers(g, zf)
+        recv = np.where(ocean, ar, recv)
+        levels = receiver_levels(recv)
+        area = accumulate(recv, levels, g.area_km2)
+        F = P["k"] * kmult * P["dt_myr"] * area ** P["m"] / np.where(np.isfinite(dist), dist, 1.0)
+        F[ocean] = 0.0
+        znew = zb.copy()
+        for lv in levels[1:]:
+            r = recv[lv]
+            znew[lv] = np.minimum(zb[lv], (zb[lv] + F[lv] * znew[r]) / (1.0 + F[lv]))
+        znew = smooth_km(g, znew, P["hillslope_km"], keep=True)   # hillslope/sub-grid smoothing, fixed length in km
+        z = np.where(ocean, z, znew)
+    drop_smooth_cache(g)
+    return z
+
+
+def run(ctx) -> dict:
+    g = ctx.grid
+    P = params(ctx.cfg, "erosion", DEFAULTS)
+    z0, age = ctx.need("elevation_m", "age_class")
+    kmult = np.vectorize(K_MULT.get)(age).astype(np.float64)
+    z0 = np.asarray(z0, dtype=np.float64)
+    added = np.asarray(ctx.data.get("land_added", np.zeros(g.n, bool)), bool)   # land the sea-level solve didn't
+    target = ctx.cfg["build"]["land_fraction"] + g.area_km2[added & (z0 > 0)].sum() / g.area_km2.sum()   # see
+    z = erode(g, z0, kmult, P)
+    z = solve_sea_level_connected(g, z, target, P["min_sea_km2"])
+    z = np.maximum(z, -11000.0)
+    # the sea is the connected world ocean; a separate basin ≥ min_sea_km2 shaped the coasts above as sea and stays
+    # open water for the climate (open_water), but its water is a lake (hydrology), not sea
+    return {"elevation_eroded_m": z.astype(np.float32), "ocean": ocean_mask(g, z, np.inf),
+            "open_water": ocean_mask(g, z, P["min_sea_km2"])}