diff options
Diffstat (limited to 'tests/test_hydrology.py')
| -rw-r--r-- | tests/test_hydrology.py | 133 |
1 files changed, 133 insertions, 0 deletions
diff --git a/tests/test_hydrology.py b/tests/test_hydrology.py new file mode 100644 index 0000000..3e9314a --- /dev/null +++ b/tests/test_hydrology.py @@ -0,0 +1,133 @@ +import unittest + +import numpy as np + +from mapgen import graph as G +from mapgen import hydrology as HY +from tests.helpers import make_ctx + + +def basin_ctx(p_basin, pet_basin): + ctx = make_ctx(3) + g = ctx.grid + ocean = g.lat < -20 + z = np.where(ocean, -3000.0, 200.0 + 30.0 * (g.lat + 20)) # land rises northward + c = g.xyz[g.cell_index(40.0, 0.0)] + d = g.radius_km * np.arccos(np.clip(g.xyz @ c, -1, 1)) + z = np.where(d < 800, z - 1500 * (1 - d / 800), z) # closed basin + ctx.data.update({"elevation_eroded_m": z.astype(np.float32), "P_ann": np.full(g.n, p_basin), + "PET": np.full(g.n, pet_basin)}) + return ctx, d + + +class HydrologyTest(unittest.TestCase): + def test_strahler(self): + recv = np.array([2, 2, 6, 5, 5, 6, 6]) + levels = G.receiver_levels(recv) + o = HY.strahler(recv, levels, np.ones(7, bool)) + np.testing.assert_array_equal(o, [1, 1, 2, 1, 1, 2, 3]) + + def test_arid_basin_is_endorheic_with_salt(self): + ctx, d = basin_ctx(100.0, 1500.0) + out = HY.run(ctx) + inside = d < 300 + self.assertTrue(out["endorheic"][inside].all()) + self.assertTrue(out["salt_flat"][inside].any() or out["lake"][inside].any()) + + def test_humid_basin_overflows(self): + ctx, d = basin_ctx(2500.0, 400.0) + out = HY.run(ctx) + inside = d < 300 + self.assertTrue(out["lake"][inside].all()) + self.assertFalse(out["endorheic"][inside].any()) + + def test_mass_balance_and_rivers(self): + ctx, _ = basin_ctx(2500.0, 400.0) + out = HY.run(ctx) + g = ctx.grid + roots = out["recv"] == np.arange(g.n) + total = np.sum(out["runoff_mm"] * g.area_km2) * 1e-6 + self.assertAlmostEqual(out["discharge_km3_yr"][roots].sum(), total, delta=total * 1e-6) + self.assertTrue(out["river"].any()) + self.assertTrue(np.all(out["strahler"][out["river"]] >= 1)) + land = ctx.data["elevation_eroded_m"] > 0 + nonendo = land & ~out["endorheic"] & ~roots + self.assertTrue(np.all(out["z_filled_m"][out["recv"][nonendo]] < out["z_filled_m"][nonendo])) + + +class InlandDepressionTest(unittest.TestCase): + def test_humid_depression_below_sea_level_becomes_lake(self): + ctx, d = basin_ctx(2500.0, 400.0) + g = ctx.grid + z = ctx.data["elevation_eroded_m"].astype(np.float64) + z[d < 300] = -40.0 + ctx.data["elevation_eroded_m"] = z.astype(np.float32) + ctx.data["ocean"] = G.ocean_mask(g, z, 1.0e6) + out = HY.run(ctx) + self.assertTrue(out["lake"][d < 300].all()) + + +def two_pit_ctx(p, pet): + ctx = make_ctx(3) + g = ctx.grid + ocean = g.lat < -20 + z = np.where(ocean, -3000.0, 200.0 + 30.0 * (g.lat + 20)) + for lon in (-3.5, 3.5): + c = g.xyz[g.cell_index(40.0, lon)] + d = g.radius_km * np.arccos(np.clip(g.xyz @ c, -1, 1)) + z = np.where(d < 900, z - (1500 + 200 * (lon > 0)) * (1 - d / 900), z) # two pits, one depression + ctx.data.update({"elevation_eroded_m": z.astype(np.float32), "P_ann": np.full(g.n, p), "PET": np.full(g.n, pet)}) + return ctx, z + + +class TerminalTest(unittest.TestCase): + def test_every_land_root_is_lake_or_salt(self): + ctx, z = two_pit_ctx(700.0, 1400.0) + out = HY.run(ctx) + roots = (out["recv"] == np.arange(len(z))) & (z > 0) | (out["recv"] == np.arange(len(z))) & out["endorheic"] + self.assertTrue(out["endorheic"].any()) + self.assertTrue(np.all(out["lake"][roots] | out["salt_flat"][roots])) + self.assertEqual(int((roots & out["endorheic"]).sum()), 1) # one terminal for the depression + + def test_mass_balance_with_lake_evaporation(self): + ctx, z = two_pit_ctx(1000.0, 1200.0) + out = HY.run(ctx) + g = ctx.grid + roots = out["recv"] == np.arange(g.n) + total = np.sum(out["runoff_mm"] * g.area_km2) * 1e-6 + self.assertGreater(out["lake_loss_km3_yr"].sum(), 0) + self.assertAlmostEqual(out["discharge_km3_yr"][roots].sum() + out["lake_loss_km3_yr"].sum(), total, + delta=total * 1e-6) + + +class LakeBedTest(unittest.TestCase): + def test_lakes_carved_below_level_and_rerun_is_stable(self): + ctx, d = basin_ctx(2500.0, 400.0) + z0 = np.asarray(ctx.data["elevation_eroded_m"], dtype=np.float64) + out = HY.run(ctx) + lake, lev, z = out["lake"], out["lake_level_m"], out["elevation_eroded_m"].astype(np.float64) + self.assertTrue(lake.any()) + self.assertTrue(np.all(np.isnan(lev[~lake]))) + self.assertTrue(np.all(z[lake] <= lev[lake] - HY.DEFAULTS["lake_min_m"] + 1e-3)) + np.testing.assert_array_equal(z[~lake], z0[~lake].astype(np.float32)) # only lake beds change + ctx.data.update({"elevation_eroded_m": out["elevation_eroded_m"]}) # again (eras rerun hydrology) + again = HY.run(ctx) + np.testing.assert_array_equal(again["lake"], lake) + np.testing.assert_allclose(again["lake_level_m"][lake], lev[lake]) + np.testing.assert_array_equal(again["elevation_eroded_m"], out["elevation_eroded_m"]) + + def test_bigger_lakes_deeper(self): + from tests.helpers import make_ctx as mk + g = mk(3).grid + P = {**HY.DEFAULTS, "lake_depth_exp": 0.3} + z = np.zeros(g.n) + lake = np.zeros(g.n, bool) + small, big = g.cell_index(10.0, 0.0), np.flatnonzero(g.xyz @ g.xyz[g.cell_index(-30.0, 90.0)] > 0.97) + lake[small] = lake[big] = True + lid = np.where(lake, 0, -1) + lid[big] = 1 + bed = HY.lake_beds(g, z, np.where(lake, 0.0, np.nan), lake, lid, np.zeros(g.n, bool), np.zeros(g.n, bool), + np.full(g.n, 1000.0), 1, P) + self.assertLess(bed[big].min(), -P["lake_min_m"]) + self.assertLessEqual(bed[small], -P["lake_min_m"]) + self.assertTrue(np.all(bed[~lake] == 0)) |