worldmap-viewer: initial public history

1 files changed, 266 insertions, 0 deletions

diff --git a/tests/test_rivers.py b/tests/test_rivers.py
new file mode 100644
index 0000000..2fb0d42
--- /dev/null
+++ b/tests/test_rivers.py
@@ -0,0 +1,266 @@
+import json
+import unittest
+
+import numpy as np
+from scipy.spatial import cKDTree
+
+import rivers as RV
+from tests.test_serve import built_world
+
+
+def net_and_arrays():
+    root = built_world()
+    d = dict(np.load(root / "out" / "r2" / "cells.npz"))
+    meta = json.loads((root / "out" / "r2" / "cells_meta.json").read_text())
+    return RV.RiverNet(d, meta["legends"], meta["radius_km"]), d
+
+
+class RiverNetTest(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.net, cls.a = net_and_arrays()
+        cls.r = np.where(cls.a["river"].astype(bool))[0]
+
+    def test_levels_never_above_ground_or_below_the_cap(self):
+        z, lev = self.a["z_surface_m"][self.r], self.net.level[self.r]
+        self.assertTrue(np.all(np.isfinite(lev)))
+        self.assertTrue(np.all(lev <= z + 1e-6))
+        self.assertTrue(np.all(z - lev <= self.net.cap[self.r] + 1e-6))
+
+    def test_levels_never_rise_downstream(self):
+        recv, riv = self.a["recv"], self.a["river"].astype(bool)
+        down = recv[self.r]
+        on = riv[down]
+        self.assertTrue(np.all(self.net.level[down[on]] <= self.net.level[self.r[on]] + 1e-6))
+
+    def test_levels_meet_their_mouth(self):
+        recv, ocean, lake = self.a["recv"], self.a["ocean"].astype(bool), self.a["lake"].astype(bool)
+        sea = self.r[ocean[recv[self.r]]]
+        self.assertTrue(len(sea) > 0)
+        k = np.isin(self.net.seg_a, sea)
+        np.testing.assert_allclose(self.net.level_b[k], 0.0)
+        lk = self.r[lake[recv[self.r]]]
+        k = np.isin(self.net.seg_a, lk)
+        np.testing.assert_allclose(self.net.level_b[k], self.a["z_filled_m"][recv[self.net.seg_a[k]]])
+
+    def test_drops_below_lakes_are_cataracts_not_cliffs(self):
+        recv, riv, lake = self.a["recv"], self.a["river"].astype(bool), self.a["lake"].astype(bool)
+        src = np.where(lake & riv[recv] & (recv != np.arange(len(recv))))[0]
+        for i in src:
+            j = recv[i]
+            dist = np.linalg.norm(self.a["g_xyz"][i] - self.a["g_xyz"][j]) * self.net.R
+            drop = self.a["z_filled_m"][i] - self.net.level[j]
+            self.assertLessEqual(drop, RV.KNICK_M_KM * dist + 1e-6 + max(0.0, self.a["z_filled_m"][i] - self.a["z_surface_m"][j]))
+
+    def test_geometry_per_segment(self):
+        n = len(self.net.seg_a)                       # river cells + lake outlets
+        self.assertGreaterEqual(n, len(self.r))
+        for k in ("a_xyz", "b_xyz", "level_a", "level_b", "half_w", "floor", "wall", "meander", "reach", "seg_len"):
+            self.assertEqual(len(getattr(self.net, k)), n, k)
+        self.assertTrue(np.all(self.net.half_w >= 0.03) and np.all(self.net.wall > 0))
+
+
+class StyleTest(unittest.TestCase):
+    def test_hard_dry_rock_makes_steeper_narrower_valleys(self):
+        g = RV.valley_style("plateau", "—", "granite/gneiss", 400.0)
+        s = RV.valley_style("plateau", "—", "sandstone/shale", 1800.0)
+        self.assertGreater(g[1], 3 * s[1])       # walls
+        self.assertLess(g[0], s[0])              # floor
+        self.assertGreater(g[1], 2500)           # the granite gorge: sheer
+
+    def test_floodplains_are_wide_and_meander(self):
+        f = RV.valley_style("plain", "floodplain", "sandstone/shale", 900.0)
+        m = RV.valley_style("mountains", "—", "granite/gneiss", 900.0)
+        self.assertGreater(f[0], 10 * m[0])
+        self.assertGreater(f[2], 0)
+        self.assertEqual(m[2], 0)
+
+    def test_channel_width(self):
+        self.assertAlmostEqual(RV.half_width_km(2.0), 0.004 * np.sqrt(2.0 * 31.7), places=9)
+
+
+class ValleyTest(unittest.TestCase):
+    def test_profile(self):
+        d = np.array([0.0, 0.05, 0.2, 1.0, 2.0])
+        v = RV.valley_surface(d, 100.0, 0.1, 0.3, 2000.0)
+        self.assertEqual(v[0], 100.0)                        # channel: the water level
+        self.assertEqual(v[1], 100.0)
+        self.assertTrue(101.0 <= v[2] <= 103.0)              # floor
+        self.assertAlmostEqual(v[4] - v[3], 2000.0, places=6)   # wall: 2,000 m per km
+
+    def test_points_are_a_global_lattice(self):
+        net, _ = net_and_arrays()
+        s = int(np.argmax(net.seg_len))
+        c = net.a_xyz[s]
+        p1, _ = net.segment_points(s, c, 50.0, 5.0)
+        p2, _ = net.segment_points(s, c, 400.0, 5.0)
+        self.assertGreater(len(p2), len(p1))
+        a = {tuple(np.round(v, 12)) for v in p1}
+        b = {tuple(np.round(v, 12)) for v in p2}
+        self.assertTrue(a <= b, "a wider query returns a superset of the same points")
+
+    def test_antimeridian_segment(self):
+        net, _ = net_and_arrays()
+        a = np.array([np.cos(np.radians(10)) * np.cos(np.radians(179.9)), np.cos(np.radians(10)) * np.sin(np.radians(179.9)), np.sin(np.radians(10))])
+        b = np.array([np.cos(np.radians(10)) * np.cos(np.radians(-179.9)), np.cos(np.radians(10)) * np.sin(np.radians(-179.9)), np.sin(np.radians(10))])
+        pts = RV.densify(a, b, 0.0, 1.0, 1.0 / net.R)   # 1 km steps (a chord in planet radii)
+        self.assertTrue(np.all(np.abs(np.linalg.norm(pts, axis=1) - 1) < 1e-12))
+        self.assertLess(np.max(np.linalg.norm(np.diff(pts, axis=0), axis=1)) * net.R, 1.01)
+
+
+class RoughWallTest(unittest.TestCase):
+    def test_walls_keep_the_terrain_roughness_but_the_channel_stays_flat(self):
+        d = np.array([0.0, 0.3, 1.0, 2.0])
+        smooth = RV.valley_surface(d, 100.0, 0.1, 0.1, 500.0)
+        rough = RV.valley_surface(d, 100.0, 0.1, 0.1, 500.0, rough=np.array([80.0, 80.0, 80.0, -40.0]))
+        self.assertEqual(rough[0], smooth[0])                   # the water level
+        self.assertAlmostEqual(rough[2] - smooth[2], 80.0)      # a wall well away from the floor: full roughness
+        self.assertAlmostEqual(rough[3] - smooth[3], -40.0)
+        self.assertTrue(0 < rough[1] - smooth[1] < 80.0)        # fading in above the floor
+
+
+LEG = {"landform": ["ocean", "plain", "hills"], "ground": ["—", "floodplain"], "lithology": ["granite/gneiss", "sandstone/shale"],
+       "age_class": ["oceanic", "craton (3g era)"]}
+
+
+def synthetic(cells):
+    """A tiny world. cells: (lat, lon, z, discharge, landform, ground, receiver index, kind: river|ocean|lake|land)."""
+    ll = np.radians([(c[0], c[1]) for c in cells])
+    a = {"g_xyz": np.stack([np.cos(ll[:, 0]) * np.cos(ll[:, 1]), np.cos(ll[:, 0]) * np.sin(ll[:, 1]), np.sin(ll[:, 0])], 1),
+         "z_surface_m": np.array([c[2] for c in cells], float), "z_filled_m": np.array([c[2] for c in cells], float),
+         "discharge_km3_yr": np.array([c[3] for c in cells], float),
+         "landform": np.array([LEG["landform"].index(c[4]) for c in cells]), "ground": np.array([LEG["ground"].index(c[5]) for c in cells]),
+         "recv": np.array([c[6] for c in cells]), "river": np.array([c[7] == "river" for c in cells]),
+         "ocean": np.array([c[7] == "ocean" for c in cells]), "lake": np.array([c[7] == "lake" for c in cells]),
+         "lithology": np.ones(len(cells), int), "age_class": np.ones(len(cells), int), "P_ann": np.full(len(cells), 900.0)}
+    return RV.RiverNet(a, LEG, 12742.0)
+
+
+def xyz(lat, lon):
+    la, lo = np.radians(lat), np.radians(lon)
+    return np.array([[np.cos(la) * np.cos(lo), np.cos(la) * np.sin(lo), np.sin(la)]])
+
+
+class PerValleyTest(unittest.TestCase):
+    # a big floodplain river along the equator (east to the sea) and a small hill stream 22 km north joining it
+    net = synthetic([(0, 0.0, 100, 5000, "plain", "floodplain", 1, "river"), (0, 0.3, 90, 5000, "plain", "floodplain", 2, "river"),
+                     (0, 0.6, 80, 5000, "plain", "floodplain", 3, "river"), (0, 0.9, 70, 5000, "plain", "floodplain", 4, "river"),
+                     (0, 1.2, -50, 0, "ocean", "—", 4, "ocean"),
+                     (0.1, 0.3, 400, 2, "hills", "—", 6, "river"), (0.1, 0.6, 380, 2, "hills", "—", 2, "river")])
+
+    def test_a_big_valley_is_not_crowded_out_by_a_nearby_stream(self):
+        q = xyz(0.05, 0.45)                                           # 11 km from both centrelines
+        V, ch, dr = self.net.valleys(q, 0.5)
+        big = float(np.interp(0.45, [0.3, 0.6], [self.net.level[1], self.net.level[2]]))
+        self.assertLessEqual(V[0], big + 3.0 + 1e-6, "on the big river's floodplain floor")
+
+    def test_the_answer_does_not_depend_on_the_other_query_points(self):
+        q = xyz(0.05, 0.45)
+        many = np.vstack([q, xyz(3.0, 20.0), xyz(-0.2, 0.1), xyz(0.08, 0.5)])
+        a, b = self.net.valleys(q, 0.5)[0], self.net.valleys(many, 0.5)[0]
+        self.assertEqual(a[0], b[0])
+
+    def test_floors_stay_inside_the_reach(self):
+        self.assertTrue(np.all(self.net.floor <= RV.FLOOR_MAX_KM + 1e-9))
+        self.assertTrue(np.all(self.net.reach >= self.net.half_w + self.net.floor))
+
+    def test_rough_walls_never_dip_below_the_water(self):
+        v = RV.valley_surface(np.array([0.6, 1.0]), 100.0, 0.1, 0.3, 500.0, rough=np.array([-500.0, -500.0]))
+        self.assertTrue(np.all(v >= 101.0))
+
+
+class BankTest(unittest.TestCase):
+    def test_a_river_doubling_back_below_itself_keeps_its_banks_above_its_water(self):
+        # a canyon river runs east, then back west 5.5 km north and 400 m lower: its lower floor must not
+        # undercut the upper reach (dry pits below the water beside it)
+        net = synthetic([(0, 0.0, 900, 300, "hills", "—", 1, "river"), (0, 0.3, 700, 300, "hills", "—", 2, "river"),
+                         (0.05, 0.0, 300, 300, "hills", "—", 3, "river"), (0.05, -0.3, -50, 0, "ocean", "—", 3, "ocean")])
+        la, lo = np.meshgrid(np.linspace(-0.03, 0.08, 40), np.linspace(-0.05, 0.35, 80), indexing="ij")
+        q = np.vstack([xyz(a, b) for a, b in zip(la.ravel(), lo.ravel())])
+        V, ch, _ = net.valleys(q, 0.1)
+        best, lev, bank = np.full(len(q), np.inf), np.zeros(len(q)), np.zeros(len(q), bool)
+        c = q.mean(0) / np.linalg.norm(q.mean(0))
+        for s in range(len(net.seg_a)):                               # the nearest reach, its water, its floor zone
+            p, t = net.segment_points(s, c, 60.0, 0.05)
+            d, k = cKDTree(p).query(q)
+            d *= net.R
+            b = d < best
+            best[b], bank[b] = d[b], (d <= net.half_w[s] + net.floor[s])[b]
+            lev[b] = (net.level_a[s] + (net.level_b[s] - net.level_a[s]) * t[k])[b]
+        dry = ~ch & bank
+        self.assertGreater(dry.sum(), 100)
+        self.assertEqual(int((V[dry] < lev[dry]).sum()), 0, "floor below the water of the reach beside it")
+
+
+class LakeOutletTest(unittest.TestCase):
+    def test_a_lake_outlet_joins_its_river_starting_at_the_lake_surface(self):
+        net = synthetic([(0, 0.0, 1100, 0, "hills", "—", 1, "lake"), (0, 0.3, 1100, 300, "plain", "—", 2, "river"),
+                         (0, 0.6, 600, 300, "plain", "—", 3, "river"), (0, 0.9, -50, 0, "ocean", "—", 3, "ocean")])
+        k = np.where(net.seg_a == 0)[0]
+        self.assertEqual(len(k), 1, "the lake outlet has a segment")
+        self.assertEqual(net.level_a[k[0]], 1100.0)
+        self.assertEqual(net.seg_b[k[0]], 1)
+
+
+class GroundTest(unittest.TestCase):
+    def test_given_levels_reach_their_shoulders_and_sub_pixel_streams_are_not_drawn(self):
+        cells = [(0, 0.0, 800, 5000, "plateau", "—", 1, "river"), (0, 0.3, 700, 5000, "plateau", "—", 2, "river"),
+                 (0, 0.6, -50, 0, "ocean", "—", 2, "ocean")]
+        levels = np.array([400.0, 300.0, 0.0])
+        ground = np.array([800.0, 700.0, 0.0])
+        LEG2 = dict(LEG, landform=LEG["landform"] + ["plateau"])
+        ll = np.radians([(c[0], c[1]) for c in cells])
+        arr = {"g_xyz": np.stack([np.cos(ll[:, 0]) * np.cos(ll[:, 1]), np.cos(ll[:, 0]) * np.sin(ll[:, 1]), np.sin(ll[:, 0])], 1),
+               "z_surface_m": ground.copy(), "z_filled_m": ground.copy(), "discharge_km3_yr": np.array([5000.0, 5000.0, 0.0]),
+               "landform": np.array([3, 3, 0]), "ground": np.zeros(3, int), "recv": np.array([1, 2, 2]),
+               "river": np.array([True, True, False]), "ocean": np.array([False, False, True]), "lake": np.zeros(3, bool),
+               "lithology": np.zeros(3, int), "age_class": np.ones(3, int), "P_ann": np.full(3, 300.0)}
+        net = RV.RiverNet(arr, LEG2, 12742.0, levels=levels, ground=ground)
+        k = 0
+        self.assertGreaterEqual(net.reach[k], net.half_w[k] + net.floor[k] + (ground[0] - levels[0]) / net.wall[k] - 1e-9)
+        small = synthetic([(0, 0.0, 100, 0.2, "plain", "—", 1, "river"), (0, 0.3, 90, 0.2, "plain", "—", 2, "river"),
+                           (0, 0.6, -50, 0, "ocean", "—", 2, "ocean")])
+        q = xyz(0.0, 0.15)
+        self.assertFalse(small.channel(q, 5.0)[1][0], "a 0.2 km³/yr stream is not drawn at 5 km per pixel")
+        self.assertTrue(small.channel(q, 0.01)[1][0], "but is at street level")
+
+
+class ChunkTest(unittest.TestCase):
+    def test_a_coarse_tile_splits_into_compact_chunks(self):
+        net = synthetic([(0, 0.0, 100, 5, "plain", "—", 1, "river"), (0, 0.3, 90, 5, "plain", "—", 2, "river"),
+                         (0, 0.6, -50, 0, "ocean", "—", 2, "ocean")])
+        la, lo = np.meshgrid(np.linspace(-52.0, -46.4, 256), np.linspace(-45.0, -33.75, 256), indexing="ij")   # a z5 tile
+        a, o = np.radians(la.ravel()), np.radians(lo.ravel())
+        chunks = net._chunks(np.stack([np.cos(a) * np.cos(o), np.cos(a) * np.sin(o), np.sin(a)], 1))
+        self.assertLess(len(chunks), 300, "not thousands of one-row strips")
+        self.assertEqual(sorted(np.concatenate([c[0] for c in chunks]).tolist()), list(range(256 * 256)))
+
+
+class RiverNetStateTest(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        import refine as RF
+        from tests.test_serve import built_world
+        w = RF.WorldCells(built_world() / "out" / "r2")
+        cls.a = {k: np.asarray(w.a[k]) for k in RF.NET_KEYS}
+        cls.legends, cls.R = w.meta["legends"], float(w.meta["radius_km"])
+
+    def test_state_round_trip_answers_the_same(self):
+        import rivers as RV
+        net = RV.RiverNet(self.a, self.legends, self.R)
+        arrays, meta = net.state()
+        self.assertEqual(set(arrays), set(RV.STATE))
+        back = RV.RiverNet.from_state(arrays, meta)
+        for k in RV.STATE:
+            np.testing.assert_array_equal(getattr(back, k), getattr(net, k))
+        self.assertEqual(back.max_extent, net.max_extent)
+        p = np.asarray(self.a["g_xyz"][:200], dtype=np.float64)
+        np.testing.assert_array_equal(back.valleys(p, 5.0)[1], net.valleys(p, 5.0)[1])
+        np.testing.assert_array_equal(back.channel(p, 5.0)[1], net.channel(p, 5.0)[1])
+
+    def test_tree_is_built_on_first_use(self):
+        import rivers as RV
+        net = RV.RiverNet.from_state(*RV.RiverNet(self.a, self.legends, self.R).state())
+        self.assertIsNone(net._tree)
+        self.assertIsNotNone(net.tree)
+        self.assertIs(net.tree, net.tree)