1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
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))
|
