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)