1 files changed, 111 insertions, 0 deletions
diff --git a/tests/test_erosion.py b/tests/test_erosion.py
new file mode 100644
index 0000000..2adefc0
--- /dev/null
+++ b/tests/test_erosion.py
@@ -0,0 +1,111 @@
+import unittest
+
+import numpy as np
+
+from mapgen import erosion as ER
+from mapgen.crust import CRATON
+from tests.helpers import make_ctx
+
+
+def mountain_ctx(k=None):
+    cfg = {"build": {"land_fraction": 0.05}}
+    if k is not None:
+        cfg["erosion"] = {"k": k}
+    ctx = make_ctx(3, cfg=cfg)
+    g = ctx.grid
+    d = g.radius_km * np.arccos(np.clip(g.xyz @ g.xyz[g.cell_index(0.0, 0.0)], -1, 1))
+    z = np.where(d < 6000, 6000 * np.clip(1 - d / 6000, 0, None) ** 1.2 + 50, -4000.0)
+    ctx.data.update({"elevation_m": z.astype(np.float32), "age_class": np.full(g.n, CRATON, np.int8)})
+    return ctx, z
+
+
+class ErosionTest(unittest.TestCase):
+    def test_erode_lowers_land(self):
+        ctx, z0 = mountain_ctx()
+        kmult = np.full(ctx.grid.n, 1.5)
+        z1 = ER.erode(ctx.grid, z0, kmult, ER.DEFAULTS)
+        land = z0 > 0
+        self.assertLess(z1[land].mean(), z0[land].mean())
+        self.assertLessEqual(z1.max(), z0.max() + 1e-6)
+
+    def test_stage_keeps_land_fraction(self):
+        ctx, _ = mountain_ctx()
+        z = ER.run(ctx)["elevation_eroded_m"]
+        g = ctx.grid
+        self.assertTrue(np.all(np.isfinite(z)))
+        self.assertAlmostEqual(g.area_km2[z > 0].sum() / g.area_km2.sum(), 0.05, delta=0.01)
+
+    def test_stable_with_huge_k(self):
+        ctx, _ = mountain_ctx(k=1e3)
+        z = ER.run(ctx)["elevation_eroded_m"]
+        self.assertTrue(np.all(np.isfinite(z)))
+        self.assertGreaterEqual(z.min(), -11000)
+
+    def test_stage_outputs_connected_ocean(self):
+        ctx, z0 = mountain_ctx()
+        g = ctx.grid
+        c = g.xyz[g.cell_index(0.0, 0.0)]
+        d = g.radius_km * np.arccos(np.clip(g.xyz @ c, -1, 1))
+        z0 = np.where(d < 700, -500.0, z0)                     # deep interior pit in the mountain block
+        ctx.data["elevation_m"] = z0.astype(np.float32)
+        out = ER.run(ctx)
+        self.assertIn("ocean", out)
+        self.assertFalse(out["ocean"][d < 300].any())
+        self.assertTrue(out["ocean"][d > 7000].all())
+
+    def test_stage_adds_the_elevation_stages_added_land(self):
+        ctx, z0 = mountain_ctx()
+        ctx.cfg["build"]["land_fraction"] = 0.02
+        g = ctx.grid
+        mtn = z0 > 0
+        ctx.data.update({"continental": mtn, "continental_base": mtn})   # lifted shelf: no new crust at all
+        added = mtn & (g.lon >= 0)
+        ctx.data["land_added"] = added
+        extra = g.area_km2[added].sum() / g.area_km2.sum()
+        z = ER.run(ctx)["elevation_eroded_m"]
+        self.assertAlmostEqual(g.area_km2[z > 0].sum() / g.area_km2.sum(), 0.02 + extra, delta=0.005)
+
+class BaseLevelTest(unittest.TestCase):
+    def test_erosion_keeps_coastline(self):
+        ctx = make_ctx(3)
+        g = ctx.grid
+        block = (np.abs(g.lat) < 30) & (np.abs(g.lon) < 50)
+        z0 = np.where(block, 800.0 + 2500.0 * np.exp(-((g.lon + 45) / 6.0) ** 2), -6000.0)   # coastal range by a deep sea
+        z1 = ER.erode(g, z0, np.full(g.n, 1.5), ER.DEFAULTS)
+        np.testing.assert_array_equal(z1 > 0, z0 > 0)
+        self.assertTrue(np.all(z1[~block] == z0[~block]))
+
+
+class HillslopeResolutionTest(unittest.TestCase):
+    def test_hillslope_smoothing_is_resolution_independent(self):
+        from tests.helpers import small_grid
+        drops = []
+        for res in (3, 4):
+            g = small_grid(res)
+            c = g.xyz[g.cell_index(20.0, 0.0)]
+            d = g.radius_km * np.arccos(np.clip(g.xyz @ c, -1, 1))
+            z0 = np.where(g.lat > -40, 500.0 + 3000.0 * np.exp(-(d / 600.0) ** 2), -4000.0)
+            P = dict(ER.DEFAULTS, k=0.0)                       # isolate hillslope smoothing
+            z1 = ER.erode(g, z0, np.ones(g.n), P)
+            drops.append(z0.max() - z1[d < 150].max())
+        self.assertAlmostEqual(drops[0] / max(drops[1], 1e-9), 1.0, delta=0.35)
+
+
+class SedimentFillTest(unittest.TestCase):
+    def test_shallow_basins_fill_deep_ones_remain(self):
+        g = make_ctx(3).grid
+        z0 = np.where(g.lat > -20, 200.0 + 2.0 * (g.lat + 20), -3000.0)
+        depth = {}
+        for name, (lat, dz) in {"shallow": (20.0, 60.0), "deep": (50.0, 2500.0)}.items():
+            c = g.xyz[g.cell_index(lat, 0.0)]
+            d = g.radius_km * np.arccos(np.clip(g.xyz @ c, -1, 1))
+            z0 = np.where(d < 700, z0 - dz * (1 - d / 700), z0)
+            depth[name] = d < 250
+        from mapgen.graph import priority_flood
+        def fill_after(frac):
+            z1 = ER.erode(g, z0, np.full(g.n, 1.0), dict(ER.DEFAULTS, k=0.0, sediment_fill=frac))
+            return priority_flood(g, z1, z1 <= -1000) - z1
+        self.assertGreater(fill_after(0.0)[depth["shallow"]].max(), 20.0)     # without sediment it stays a basin
+        fill = fill_after(ER.DEFAULTS["sediment_fill"])
+        self.assertLess(fill[depth["shallow"]].max(), 0.2 * 60.0)             # noise-scale basins mostly filled
+        self.assertGreater(fill[depth["deep"]].max(), 20.0)