aboutsummaryrefslogtreecommitdiffziptar.gz
path: root/tests/test_graph.py
diff options
context:
space:
mode:
Diffstat (limited to 'tests/test_graph.py')
-rw-r--r--tests/test_graph.py362
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)
generated by cgit (git 2.34.1) at 2026-10-08 00:04:53 +0000