worldgen: initial public history

1 files changed, 362 insertions, 0 deletions

diff --git a/tests/test_graph.py b/tests/test_graph.py
new file mode 100644
index 0000000..ca324cc
--- /dev/null
+++ b/tests/test_graph.py
@@ -0,0 +1,362 @@
+import unittest
+
+import numpy as np
+
+from mapgen import graph as G
+from mapgen.sphere import latlon_to_xyz, east_north
+from tests.helpers import small_grid
+
+
+class GraphTest(unittest.TestCase):
+    def setUp(self):
+        self.g = small_grid(2)
+
+    def test_mean_max_min_diffuse(self):
+        g = self.g
+        f = np.zeros(g.n)
+        f[0] = 6.0
+        m = G.nbr_mean(g, f)
+        self.assertAlmostEqual(m[g.nbr_idx[g.nbr_ptr[0]]], 6.0 / g.counts[g.nbr_idx[g.nbr_ptr[0]]])
+        self.assertEqual(G.nbr_max(g, f)[g.nbr_idx[g.nbr_ptr[0]]], 6.0)
+        d = G.diffuse(g, f, 10)
+        self.assertLess(d.max(), 6.0)
+        self.assertGreater(np.count_nonzero(d > 1e-6), 20)
+
+    def test_gradient_of_linear_field(self):
+        g = self.g
+        f = g.xyz[:, 2] * g.radius_km  # height ∝ z → gradient points north near the equator
+        gr = G.gradient(g, f)
+        eq = np.abs(g.lat) < 20
+        _, n = east_north(g.xyz[eq])
+        cos_to_north = np.sum(gr[eq] * n, axis=1) / np.linalg.norm(gr[eq], axis=1)
+        self.assertGreater(np.median(cos_to_north), 0.99)
+        self.assertAlmostEqual(float(np.median(np.linalg.norm(gr[eq], axis=1))), 1.0, delta=0.1)
+
+    def test_distance_matches_great_circle(self):
+        g = self.g
+        i = g.cell_index(0.0, 0.0)
+        d = G.distance_to(g, np.arange(g.n) == i)
+        gc = g.radius_km * np.arccos(np.clip(g.xyz @ g.xyz[i], -1, 1))
+        far = gc > 3000
+        ratio = d[far] / gc[far]
+        self.assertTrue(np.all(ratio >= 0.999) and np.median(ratio) < 1.15)
+        self.assertTrue(np.all(np.isinf(G.distance_to(g, np.zeros(g.n, bool)))))
+
+    def test_nearest_source_labels(self):
+        g = self.g
+        a, b = g.cell_index(0.0, -90.0), g.cell_index(0.0, 90.0)
+        _, src = G.nearest_source(g, [a, b])
+        self.assertEqual(src[g.cell_index(0.0, -60.0)], a)
+        self.assertEqual(src[g.cell_index(0.0, 60.0)], b)
+
+    def test_priority_flood_fills_basin_and_drains(self):
+        g = self.g
+        ocean = g.lat < -30
+        z = np.where(ocean, -1000.0, 500.0 + 10 * g.lat)
+        pit = g.cell_index(40.0, 0.0)
+        z[pit] = -50.0  # land pit below sea level, not connected to ocean
+        zf = G.priority_flood(g, z, ocean)
+        self.assertGreater(zf[pit], z[pit])
+        recv, slope, dist = G.steepest_receivers(g, zf)
+        recv[ocean] = np.flatnonzero(ocean)
+        levels = G.receiver_levels(recv)
+        self.assertEqual(sum(len(l) for l in levels), g.n)
+        self.assertTrue(np.all(ocean[levels[0]]))
+
+    def test_priority_flood_needs_sink(self):
+        with self.assertRaisesRegex(ValueError, "no sink"):
+            G.priority_flood(self.g, np.ones(self.g.n), np.zeros(self.g.n, bool))
+
+    def test_accumulate_conserves(self):
+        g = self.g
+        ocean = g.lat < -30
+        z = np.where(ocean, -1000.0, 1000.0 + 20 * g.lat)
+        zf = G.priority_flood(g, z, ocean)
+        recv, _, _ = G.steepest_receivers(g, zf)
+        recv[ocean] = np.flatnonzero(ocean)
+        lv = G.receiver_levels(recv)
+        w = np.where(ocean, 0.0, 1.0)
+        acc = G.accumulate(recv, lv, w)
+        self.assertAlmostEqual(acc[lv[0]].sum(), w.sum())
+
+    def test_cycle_detected(self):
+        with self.assertRaisesRegex(ValueError, "cycle"):
+            G.receiver_levels(np.array([1, 0, 2]))
+
+    def test_components(self):
+        g = self.g
+        m = (np.abs(g.lat) < 10) & (np.abs(g.lon) < 20) | (np.abs(g.lat - 50) < 8) & (np.abs(g.lon) < 20)
+        lab = G.components(g, m)
+        self.assertEqual(len(np.unique(lab[m])), 2)
+        self.assertTrue(np.all(lab[~m] == -1))
+
+
+class SmoothKmTest(unittest.TestCase):
+    def test_constant_preserved_and_spike_decays(self):
+        g = small_grid(3)
+        np.testing.assert_allclose(G.smooth_km(g, np.full(g.n, 5.0), 1000.0), 5.0, rtol=1e-4)
+        i = g.cell_index(0.0, 0.0)
+        s = G.smooth_km(g, (np.arange(g.n) == i).astype(float), 1000.0)
+        d = g.radius_km * np.arccos(np.clip(g.xyz @ g.xyz[i], -1, 1))
+        near, far = s[(d > 800) & (d < 1200)].mean(), s[(d > 2800) & (d < 3200)].mean()
+        self.assertTrue(near > far > 0)
+
+    def test_resolution_independent(self):
+        vals = []
+        for res in (2, 3):
+            g = small_grid(res)
+            d = g.radius_km * np.arccos(np.clip(g.xyz @ g.xyz[g.cell_index(0.0, 0.0)], -1, 1))
+            s = G.smooth_km(g, (d < 2000).astype(float), 1500.0)
+            vals.append(s[g.cell_index(0.0, 30.0)])      # ~6700 km away
+        self.assertAlmostEqual(vals[0], vals[1], delta=0.25 * max(vals))
+        self.assertGreater(min(vals), 0.005)
+
+
+class OceanMaskTest(unittest.TestCase):
+    def test_inland_depression_is_not_ocean(self):
+        g = small_grid(3)
+        z = np.where(g.lat < 0, -3000.0, 500.0)
+        c = g.xyz[g.cell_index(40.0, 0.0)]
+        d = g.radius_km * np.arccos(np.clip(g.xyz @ c, -1, 1))
+        z[d < 600] = -50.0                                   # interior basin below sea level
+        ocean = G.ocean_mask(g, z, 1.0e6)
+        self.assertTrue(ocean[g.lat < -5].all())
+        self.assertFalse(ocean[d < 600].any())
+
+
+class SmoothKmRobustTest(unittest.TestCase):
+    def test_zero_and_tiny_fields(self):
+        g = small_grid(3)
+        np.testing.assert_array_equal(G.smooth_km(g, np.zeros(g.n), 25.0), 0.0)
+        f = np.where(g.lat > 0, 1e-9, 0.0)
+        s = G.smooth_km(g, f, 25.0)
+        self.assertTrue(np.all(np.isfinite(s)) and s.max() <= 1e-9 * (1 + 1e-6))
+
+    def test_femto_scale_field(self):
+        g = small_grid(3)
+        rng = np.random.default_rng(0)
+        f = np.where(rng.random(g.n) < 0.06, rng.random(g.n) * 8e-14, 0.0)
+        s = G.smooth_km(g, f, 25.0)
+        self.assertTrue(np.all(np.isfinite(s)))
+        self.assertAlmostEqual(float(s.sum() / f.sum()), 1.0, delta=0.05)
+
+
+def _flood_ref(g, z, sink_mask, eps=0.01):
+    """The original pure-Python priority flood (oracle for the compiled one)."""
+    import heapq
+    has_open = np.bincount(g.src, weights=(~sink_mask)[g.dst].astype(np.float64), minlength=g.n) > 0
+    zf, done = np.asarray(z, dtype=np.float64).tolist(), sink_mask.tolist()
+    ptr, idx = g.nbr_ptr.tolist(), g.nbr_idx.tolist()
+    heap = [(zf[i], i) for i in np.flatnonzero(sink_mask & has_open).tolist()]
+    heapq.heapify(heap)
+    while heap:
+        zc, c = heapq.heappop(heap)
+        for k in range(ptr[c], ptr[c + 1]):
+            n = idx[k]
+            if not done[n]:
+                done[n] = True
+                zf[n] = max(zf[n], zc + eps)
+                heapq.heappush(heap, (zf[n], n))
+    return np.array(zf)
+
+
+def _steepest_ref(g, z):
+    """The original lexsort version (oracle)."""
+    slope = (z[g.src] - z[g.dst]) / g.edge_km
+    first = np.lexsort((-slope, g.src))[g.nbr_ptr[:-1]]
+    s = slope[first]
+    down = s > 0
+    return (np.where(down, g.dst[first], np.arange(g.n)), np.where(down, s, 0.0),
+            np.where(down, g.edge_km[first], np.inf))
+
+
+def _accumulate_ref(recv, levels, w):
+    acc = np.asarray(w, dtype=np.float64).copy()
+    for lv in reversed(levels[1:]):
+        acc += np.bincount(recv[lv], weights=acc[lv], minlength=len(acc))
+    return acc
+
+
+class FastPathsTest(unittest.TestCase):
+    """The speed-ups (numba flood, sort-free receivers, per-level accumulate) give bit-identical results."""
+
+    def fields(self, g):
+        rng = np.random.default_rng(7)
+        rough = rng.normal(0, 300, g.n)
+        flat = np.round(rng.normal(0, 2, g.n))            # many exact ties: tie order matters
+        pits = np.where(rng.random(g.n) < 0.2, -50.0, rough)
+        return {"rough": rough, "flat": flat, "pits": pits}
+
+    def test_flood_matches_reference(self):
+        g = small_grid(3)
+        for name, z in self.fields(g).items():
+            sink = z < np.quantile(z, 0.1)
+            for eps in (0.01, 0.0):
+                with self.subTest(name=name, eps=eps):
+                    want = _flood_ref(g, z, sink, eps)
+                    got = G.priority_flood(g, z, sink, eps)
+                    self.assertTrue(np.array_equal(got, want))
+                    py = G._flood_py(z, sink, g.nbr_ptr, g.nbr_idx,
+                                     np.flatnonzero(sink & (np.bincount(g.src, weights=(~sink)[g.dst].astype(float),
+                                                                       minlength=g.n) > 0)), eps)
+                    self.assertTrue(np.array_equal(py, want))
+
+    def test_steepest_receivers_match_reference(self):
+        g = small_grid(3)
+        for name, z in self.fields(g).items():
+            with self.subTest(name=name):
+                for got, want in zip(G.steepest_receivers(g, z), _steepest_ref(g, z)):
+                    self.assertTrue(np.array_equal(got, want))
+
+    def test_accumulate_matches_reference(self):
+        g = small_grid(3)
+        rng = np.random.default_rng(3)
+        for name, z in self.fields(g).items():
+            with self.subTest(name=name):
+                zf = G.priority_flood(g, z, z < np.quantile(z, 0.1))
+                recv, _, _ = G.steepest_receivers(g, zf)
+                lv = G.receiver_levels(recv)
+                w = rng.random(g.n) * 1e3 * np.where(rng.random(g.n) < 0.1, -0.0, 1.0)   # with negative zeros
+                got, want = G.accumulate(recv, lv, w), _accumulate_ref(recv, lv, w)
+                self.assertTrue(np.array_equal(got, want))
+                self.assertTrue(np.array_equal(np.signbit(got), np.signbit(want)))
+
+
+class SweepTest(unittest.TestCase):
+    def test_sweep_matches_full_bincount(self):
+        from mapgen import hydrology as HY
+
+        def sweep_ref(recv, levels, water, outlets, cap):
+            acc = np.asarray(water, dtype=np.float64).copy()
+            loss = np.zeros(len(acc))
+            is_out = np.zeros(len(acc), bool)
+            is_out[outlets] = True
+            cap_cell = np.zeros(len(acc))
+            cap_cell[outlets] = cap
+            for lv in reversed(levels[1:]):
+                push = acc[lv].copy()
+                o = is_out[lv]
+                if o.any():
+                    cells = lv[o]
+                    lost = np.minimum(acc[cells], cap_cell[cells])
+                    loss[cells] = lost
+                    push[o] = acc[cells] - lost
+                acc += np.bincount(recv[lv], weights=push, minlength=len(acc))
+            return acc, loss
+
+        g = small_grid(3)
+        rng = np.random.default_rng(11)
+        z = rng.normal(0, 300, g.n)
+        zf = G.priority_flood(g, z, z < np.quantile(z, 0.1))
+        recv, _, _ = G.steepest_receivers(g, zf)
+        lv = G.receiver_levels(recv)
+        water = rng.random(g.n) * np.where(rng.random(g.n) < 0.1, -0.0, 1.0)     # with negative zeros
+        outlets = rng.choice(g.n, 200, replace=False)
+        cap = rng.random(200) * 2
+        for got, want in zip(HY._sweep(recv, lv, water, outlets, cap), sweep_ref(recv, lv, water, outlets, cap)):
+            self.assertTrue(np.array_equal(got, want))
+            self.assertTrue(np.array_equal(np.signbit(got), np.signbit(want)))
+
+
+class LeavesTest(unittest.TestCase):
+    def test_leaves_all_matches_the_walk(self):
+        from mapgen import hydrology as HY
+        g = small_grid(3)
+        rng = np.random.default_rng(4)
+        for seed in range(3):
+            z = rng.normal(0, 300, g.n)
+            ocean = z < np.quantile(z, 0.2)
+            zf = G.priority_flood(g, z, ocean)
+            lab = G.components(g, ~ocean & (zf - z > 1.0))
+            recv, _, _ = G.steepest_receivers(g, zf)
+            recv = np.where(ocean, np.arange(g.n), recv)
+            lv = G.receiver_levels(recv)
+            xs = np.flatnonzero(lab >= 0)
+            for limit in (100000, 3):
+                want = np.array([HY._leaves(recv, lab, x, limit) for x in xs])
+                got = HY._leaves_all(recv, lab, lv, limit)[xs]
+                self.assertGreater(want.sum(), 0)
+                self.assertTrue(np.array_equal(got, want), (seed, limit))
+
+
+class BicgstabJacobiTest(unittest.TestCase):
+    """The fused solver walks scipy's iterates exactly: same answers bit for bit, same exit codes."""
+    def systems(self):
+        from scipy import sparse
+        rng = np.random.default_rng(0)
+        for n in (2000, 9000):
+            i = np.repeat(np.arange(n), 6)
+            j = (i + rng.integers(-50, 50, len(i))) % n
+            L = sparse.csr_matrix((rng.random(len(i)), (i, j)), shape=(n, n))
+            S = L + L.T
+            yield (sparse.diags(np.asarray(S.sum(1)).ravel()) - S).tocsr() * 40 + sparse.identity(n, format="csr")
+            yield (sparse.identity(n, format="csr") * (1 + np.asarray(L.sum(1)).ravel().max() * 0.6) - L).tocsr()
+
+    def setUp(self):
+        self.min_n = G.JIT_MIN_N
+        G.JIT_MIN_N = 0                                      # the compiled path even on small test systems
+
+    def tearDown(self):
+        G.JIT_MIN_N = self.min_n
+
+    def test_matches_scipy_bicgstab(self):
+        from scipy.sparse import linalg as splinalg
+        rng = np.random.default_rng(1)
+        for k, A in enumerate(self.systems()):
+            for rtol, maxiter in ((1e-6, 5000), (1e-9, 5000), (1e-12, 7)):
+                b = rng.normal(size=A.shape[0])
+                inv = 1.0 / A.diagonal()
+                M = splinalg.LinearOperator(A.shape, matvec=lambda x: inv * x)
+                want = splinalg.bicgstab(A, b, x0=b * 0.5, rtol=rtol, maxiter=maxiter, M=M)
+                got = G.bicgstab_jacobi(A, b, b * 0.5, inv, rtol, maxiter)
+                with self.subTest(k=k, rtol=rtol, maxiter=maxiter):
+                    self.assertEqual(got[1], want[1])
+                    self.assertTrue(np.array_equal(got[0], want[0]))
+
+    def test_zero_right_hand_side_and_zero_start(self):
+        A = next(self.systems())
+        inv = 1.0 / A.diagonal()
+        x, info = G.bicgstab_jacobi(A, np.zeros(A.shape[0]), np.zeros(A.shape[0]), inv, 1e-6, 100)
+        self.assertEqual(info, 0)
+        self.assertFalse(x.any())
+
+
+class PmapTest(unittest.TestCase):
+    def test_order_and_same_floats_as_serial(self):
+        import os
+        from unittest import mock
+        g = small_grid(2)
+        fields = [np.random.default_rng(i).normal(size=g.n) for i in range(4)]
+        serial = [G.smooth_km(g, f, 900.0) for f in fields]
+        with mock.patch.dict(os.environ, {"WORLDGEN_THREADS": "4"}):
+            self.assertEqual(G.workers(), 4)
+            got = G.pmap(lambda f: G.smooth_km(g, f, 900.0), fields)
+        for a, b in zip(serial, got):
+            self.assertTrue(np.array_equal(a, b))
+        with mock.patch.dict(os.environ, {"WORLDGEN_THREADS": "3"}):
+            self.assertEqual(G.pmap(lambda k: k * k, range(9)), [k * k for k in range(9)])
+        with mock.patch.dict(os.environ, {"WORLDGEN_THREADS": "x"}):
+            self.assertEqual(G.workers(), 3)
+
+
+class ComponentsTest(unittest.TestCase):
+    def test_same_labels_as_scipy(self):
+        from scipy import sparse
+        from scipy.sparse import csgraph
+        g = small_grid(3)
+        rng = np.random.default_rng(11)
+
+        def scipy_labels(mask):                                # the previous implementation, as the oracle
+            e = mask[g.src] & mask[g.dst]
+            m = sparse.csr_matrix((np.ones(int(e.sum())), (g.src[e], g.dst[e])), shape=(g.n, g.n))
+            _, lab = csgraph.connected_components(m, directed=False)
+            return np.where(mask, lab, -1)
+        for p in (0.0, 0.2, 0.45, 0.6, 0.9, 1.0):
+            for _ in range(3):
+                mask = rng.random(g.n) < p
+                want, got = scipy_labels(mask), G.components(g, mask)
+                self.assertEqual(got.dtype, want.dtype)
+                self.assertTrue(np.array_equal(got, want), p)
+        smooth = G.smooth_km(g, rng.normal(size=g.n), 2000.0) > 0      # big blobs, many cells each
+        self.assertTrue(np.array_equal(G.components(g, smooth), scipy_labels(smooth)))
+        self.assertGreater(len(np.unique(G.components(g, smooth))), 2)