diff options
Diffstat (limited to 'tests/test_refine.py')
| -rw-r--r-- | tests/test_refine.py | 805 |
1 files changed, 805 insertions, 0 deletions
diff --git a/tests/test_refine.py b/tests/test_refine.py new file mode 100644 index 0000000..3284dc5 --- /dev/null +++ b/tests/test_refine.py @@ -0,0 +1,805 @@ +import json +import tempfile +import unittest +from pathlib import Path + +import h3.api.basic_int as h3 +import numpy as np + +import refine as RF +from mapgen import config as C + + +def square(lat, lon, d, rid="r1", name="test"): + return {"id": rid, "name": name, "outline": [[lat - d, lon - d], [lat - d, lon + d], [lat + d, lon + d], [lat + d, lon - d]]} + + +class AreasTest(unittest.TestCase): + def test_regions_file(self): + self.assertEqual(RF.load_regions(Path(tempfile.mkdtemp()) / "none.json"), []) + p = Path(tempfile.mkdtemp()) / "regions.json" + p.write_text(json.dumps({"regions": [square(10, 20, 1)]})) + self.assertEqual(RF.load_regions(p)[0]["name"], "test") + + def test_cells_inside_the_outline(self): + (cells,) = RF.area_cells([square(10, 20, 1)], 4) + lat = np.array([h3.cell_to_latlng(int(c))[0] for c in cells]) + lon = np.array([h3.cell_to_latlng(int(c))[1] for c in cells]) + self.assertTrue(len(cells) > 10) + self.assertTrue(np.all(np.abs(lat - 10) <= 1.1) and np.all(np.abs(lon - 20) <= 1.1)) + self.assertTrue(np.all(np.diff(cells.astype(np.int64)) > 0), "sorted, unique") + + def test_touching_outlines_merge_separate_ones_dont(self): + a, b = square(10, 20, 1, "a"), square(10, 22, 1, "b") # share the lon 21 edge + c = square(-30, 100, 1, "c") + areas = RF.area_cells([a, b, c], 4) + self.assertEqual(len(areas), 2) + self.assertEqual(sorted(len(x) for x in areas)[1], len(RF.area_cells([a], 4)[0]) + len(RF.area_cells([b], 4)[0]) - len(np.intersect1d(RF.area_cells([a], 4)[0], RF.area_cells([b], 4)[0]))) + + def test_hash_follows_the_cells_only(self): + (x,) = RF.area_cells([square(10, 20, 1, "a", "one")], 4) + (y,) = RF.area_cells([square(10, 20, 1, "b", "renamed")], 4) + (w,) = RF.area_cells([square(10, 20, 1.5)], 4) + self.assertEqual(RF.area_hash(x), RF.area_hash(y)) + self.assertNotEqual(RF.area_hash(x), RF.area_hash(w)) + + def test_sea_only_where_it_reaches_the_open_ocean(self): + """World ocean cut off by a refined land strip is no sea (a lagoon: the lakes take it); as the world's own rule, + a separate basin counts as sea only when it is that big.""" + (cells,) = RF.area_cells([square(10, 20, 3)], 4) + g, halo = RF.subgrid(cells, 12742.0) + d = np.hypot(g.lat - 10, g.lon - 20) + z = np.where((d > 0.9) & (d < 1.6), 5.0, -10.0) # a ring of land around a wet middle + ocean_world = np.ones(g.n, bool) + inner, outer = d <= 0.9, d >= 1.6 + self.assertTrue(inner.sum() > 3 and (outer & halo).any()) + sea = RF._sea(g, z, ocean_world, halo) + self.assertTrue(sea[outer].all(), "open water out to the area's edge: sea") + self.assertFalse(sea[inner].any(), "behind the strip: not sea") + big = RF._sea(g, z, ocean_world, halo, min_km2=float(g.area_km2[inner].sum()) * 0.9) + self.assertTrue(big[inner].all(), "a separate basin as big as a world sea stays sea") + self.assertFalse(RF._sea(g, z, np.zeros(g.n, bool), halo).any(), "never where the world has no ocean") + + def test_world_lake_keeps_its_world_outflow(self): + """A world lake that drains to the world sea keeps that way out at its world level: the fine relief must not + dam it higher (r5 2026-10-04: lakes at 5 m stood at 29 m and drowned their shores). Only its world drainage + path is cut; a lake that already drains low enough is left alone.""" + (cells,) = RF.area_cells([square(10, 20, 3)], 4) + g, halo = RF.subgrid(cells, 12742.0) + parent = np.where(g.lon < 19, 0, np.where(g.lon < 21, 1, 2)) # world: lake → land → ocean (west to east) + parent = np.where((g.lat > 11.5) & (parent == 1), 3, parent) # and land off the drainage path + world = {"lake": np.array([1, 0, 0, 0], bool), "ocean": np.array([0, 0, 1, 0], bool), + "lake_id": np.array([7, -1, -1, -1]), "lake_level_m": np.array([5.0, np.nan, np.nan, np.nan]), + "recv": np.array([1, 2, 2, 3])} + z = np.select([parent == 0, parent == 2], [-20.0, -900.0], 30.0) # a 30 m sill between lake and a deep sea + sea = RF._sea(g, z, world["ocean"][parent], halo) + lake = (parent == 0) & ~halo + + def level(zz): + return float(RF.priority_flood(g, zz, sea)[lake].max()) + self.assertGreater(level(z), 25.0, "the fine sill dams the lake") + cut = RF.world_outflows(g, z, halo, sea, parent, world) + self.assertLessEqual(level(cut), 5.0, "the lake drains at its world level") + self.assertTrue((cut <= z).all(), "only lowers") + self.assertTrue((cut[parent == 3] == z[parent == 3]).all(), "off the world path: untouched") + changed = cut < z + self.assertTrue(0 < changed.sum() < (parent == 1).sum() / 3, "a channel, not the whole cell") + self.assertTrue((cut[changed] >= 0.0).all(), "never cut below sea level (the sea's surface, not its bed)") + low = np.where(parent == 1, 2.0, z) + self.assertTrue((RF.world_outflows(g, low, halo, sea, parent, world) == low).all(), "drains already: untouched") + + def test_subgrid_with_halo(self): + (cells,) = RF.area_cells([square(10, 20, 1)], 4) + g, halo = RF.subgrid(cells, 12742.0) + self.assertEqual(int((~halo).sum()), len(cells)) + self.assertTrue(halo.sum() > 0) + pairs = set(zip(g.src.tolist(), g.dst.tolist())) + self.assertTrue(all((j, i) in pairs for i, j in pairs), "symmetric neighbours") + inner = np.where(~halo)[0] + self.assertTrue(np.all(g.counts[inner] >= 5), "every area cell keeps all its neighbours (pentagons: 5)") + self.assertAlmostEqual(float(g.area_km2[inner].mean()), h3.average_hexagon_area(4, unit="km^2") * 4, delta=0.3 * h3.average_hexagon_area(4, unit="km^2") * 4) + + +from tests.test_serve import built_world +from tests.helpers import big_river +import serve +import tiles as T + + +class FieldsTest(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.root = built_world() + cls.world = RF.WorldCells(cls.root / "out" / "r2") + cls.src = T.TileSource(serve.World(cls.root, 2), seed=7, cache_dir=Path(tempfile.mkdtemp())) + i = big_river(cls.world.a) + lat, lon = float(cls.world.a["g_lat"][i]), float(cls.world.a["g_lon"][i]) + (cls.cells,) = RF.area_cells([square(lat, lon, 4)], 4) + cls.g, cls.halo = RF.subgrid(cls.cells, 12742.0) + cls.f = RF.initial_fields(cls.g, cls.halo, cls.world, cls.src) + cls.cfg = C.load(cls.root)[0] + cls.r = RF.refine_area(cls.g, cls.halo, cls.world, cls.src, cls.cfg) + + def test_parents_are_the_world_cells(self): + p = self.f["parent"] + np.testing.assert_array_equal(self.world.a["g_ids"][p], [h3.cell_to_parent(int(c), 2) for c in self.g.ids]) + + def test_elevation_is_the_world_surface_plus_coarse_relief(self): + z = self.src.relief_at(self.g.lat, self.g.lon, RF.RELIEF_MIN_KM) + np.testing.assert_allclose(self.f["elevation_m"], z) + + def test_temperature_follows_the_lapse_rate(self): + p = self.f["parent"] + dz = (self.f["elevation_m"] - self.world.a["z_surface_m"][p]) / 1000.0 + np.testing.assert_allclose(self.f["T_mean"], self.world.a["T_mean"][p] - RF.LAPSE_C_PER_KM * dz, atol=1e-4) + np.testing.assert_array_equal(self.f["lithology"], self.world.a["lithology"][p]) + + def test_crossings_find_the_big_river(self): + inflow, exits, lev = RF.crossings(self.g, self.halo, self.world) + self.assertTrue(inflow.sum() > 0 or len(exits) > 0, "the widest world river crosses the square") + self.assertTrue(np.all(self.halo[exits])) + self.assertTrue(np.all(inflow[self.halo] == 0), "inflow enters inside the area") + + def test_crossings_match_scanning_every_world_river(self): + """The near-area prefilter drops only segments that cannot cross (oracle: the unfiltered scan).""" + from unittest import mock + got = RF.crossings(self.g, self.halo, self.world) + with mock.patch.object(RF, "world_near", lambda w, cells, rings=3: np.arange(len(w.ids))): + want = RF.crossings(self.g, self.halo, self.world) + for a, b in zip(got, want): + np.testing.assert_array_equal(a, b) + + def test_erosion_keeps_the_halo_and_the_sea_and_cuts_valleys(self): + f = self.f + sea = f["ocean_world"] & (f["elevation_m"] <= 0) + inflow, _, _ = RF.crossings(self.g, self.halo, self.world) + water = RF.runoff_km3(self.g, f) + inflow + z = RF.erode_area(self.g, f["elevation_m"], self.halo, sea, water, np.ones(self.g.n)) + np.testing.assert_array_equal(z[self.halo], f["elevation_m"][self.halo]) + np.testing.assert_array_equal(z[sea], f["elevation_m"][sea]) + self.assertTrue((f["elevation_m"] - z).max() > 5.0, "valleys cut") + self.assertTrue(np.all(z <= f["elevation_m"] + 1e-6), "erosion only lowers") + + def test_schema_and_halo(self): + for k in ("z_surface_m", "river", "recv", "discharge_km3_yr", "lake", "holdridge", "ground", "ice", "po2_bar"): + self.assertEqual(len(self.r[k]), self.g.n, k) + np.testing.assert_array_equal(self.r["halo"], self.halo) + self.assertFalse(self.r["river"][self.halo].any()) + + def test_water_balance(self): + s = RF.stats(self.r) + self.assertLess(s["water_balance_error"], 0.01) + + def test_exit_where_the_world_exits(self): + _, exits, _ = RF.crossings(self.g, self.halo, self.world) + into = self.r["recv"][~self.halo] + big = np.where(self.r["river"] & ~self.halo & self.halo[self.r["recv"]])[0] # rivers leaving the area + if len(exits): + leave = set(self.r["recv"][big].tolist()) + self.assertTrue(set(exits.tolist()) & leave, "the main river leaves through a world exit cell") + + def test_no_river_dead_ends(self): + r, recv = self.r, self.r["recv"] + riv = np.where(r["river"])[0] + ok = r["river"][recv[riv]] | r["lake"][recv[riv]] | r["ocean"][recv[riv]] | r["halo"][recv[riv]] | r["endorheic"][recv[riv]] + self.assertTrue(np.all(ok)) + + def test_sea_stays_sea(self): + sea = self.f["ocean_world"] & (self.f["elevation_m"] <= 0) + self.assertTrue(np.all(self.r["ocean"][sea])) + + def test_deterministic(self): + again = RF.refine_area(self.g, self.halo, self.world, self.src, self.cfg) + for k in ("z_surface_m", "discharge_km3_yr", "river"): + np.testing.assert_array_equal(again[k], self.r[k]) + + def test_streams_from_a_fifth_of_a_km3(self): + q = self.r["discharge_km3_yr"] + land = ~self.r["ocean"] & ~self.r["lake"] & ~self.halo + np.testing.assert_array_equal(self.r["river"][land], q[land] >= RF.RIVER_MIN_KM3_YR) + + def test_lakes_only_in_real_hollows(self): + deep = self.r["depression_depth_m"] > RF.LAKE_MIN_DEPTH_M + self.assertFalse((self.r["lake"] & ~deep & ~self.r["endorheic"] & ~self.halo).any(), "no lakes in shallow dips") + + def test_shoulder_heights_bridge_the_rivers(self): + env = RF.shoulders(self.r) + riv = self.r["river"] & ~self.halo + self.assertTrue(np.all(env[riv] >= self.r["z_surface_m"][riv] - 1e-6)) + self.assertTrue(np.all(env[~riv] == self.r["z_surface_m"][~riv])) + + def test_an_inflowing_river_carries_its_water_into_the_area(self): + inflow, _, _ = RF.crossings(self.g, self.halo, self.world) + if inflow.max() > 0: + q = self.r["discharge_km3_yr"][~self.halo] + self.assertGreaterEqual(q.max(), 0.9 * inflow.max(), "not bounced straight back out through the halo") + outlets = RF.outlets(self.g, self.halo, self.world, self.f["parent"]) + self.assertTrue(outlets.any() and np.all(self.halo[outlets])) + self.assertLess(outlets.sum(), self.halo.sum(), "only some halo cells are ways out") + + def test_no_block_steps_between_world_cells(self): + z, p, g = self.r["elevation_eroded_m"], self.f["parent"], self.g + inner = ~self.halo[g.src] & ~self.halo[g.dst] & ~self.r["ocean"][g.src] & ~self.r["ocean"][g.dst] + d = np.abs(z[g.src] - z[g.dst])[inner] + across = (p[g.src] != p[g.dst])[inner] + if across.sum() > 10 and (~across).sum() > 10: + self.assertLess(np.percentile(d[across], 95), 1.3 * np.percentile(d[~across], 95) + 5.0) + + def test_world_cell_means_of_bedrock_are_kept_smoothly(self): + z, p = self.r["elevation_eroded_m"], self.f["parent"] + for w in np.unique(p[~self.halo]): + m = (p == w) & ~self.halo + floor = (self.r["canyon_m"][m] > 0).any() or (self.r["fan_m"][m] > 0).any() # the sea-floor pass: sea only + if m.sum() == 49 and not self.world.a["ocean"][w] and not floor: + self.assertAlmostEqual(float(z[m].mean()), float(self.world.a["elevation_eroded_m"][w]), delta=25.0) + + def test_no_fake_lakes_against_walls(self): + g, r = self.g, self.r + wall = self.halo & ~r["outlet"] + touching = set(r["lake_id"][g.src[wall[g.dst] & r["lake"][g.src]]].tolist()) + for lid in touching: # a wall lower than the lake would have let it spill out + m = r["lake"] & (r["lake_id"] == lid) + walls = g.dst[m[g.src] & wall[g.dst]] + level = float(r["z_filled_m"][m].max()) + self.assertTrue(np.all(r["elevation_eroded_m"][walls] >= level - 1e-6), f"lake {lid} dammed by a wall") + + def test_temperatures_follow_the_final_heights(self): + p = self.f["parent"] + dz = (self.r["elevation_eroded_m"] - self.world.a["elevation_eroded_m"][p]) / 1000.0 + np.testing.assert_allclose(self.r["T_mean"], self.world.a["T_mean"][p] - RF.LAPSE_C_PER_KM * dz, atol=1e-3) + + def test_border_report(self): + s = RF.stats(self.r) + self.assertIn("exit_level_error_m", s) + if s["exit_level_error_m"] is not None: + self.assertLess(abs(s["exit_level_error_m"]), 50.0, "rivers leave at the world river's level") + + def test_zone_fields_come_from_the_world_cells(self): + from mapgen import fields as FL + p, a, r = self.f["parent"], self.world.a, self.r + for k in ("gravity_g", "o2_fraction", "fire_reactivity", "plant_height_x"): + np.testing.assert_array_equal(r[k], a[k][p], err_msg=k) + H = float(self.cfg["planet"]["scale_height_km"]) * 1000.0 + sea = a["pressure_bar"][p] / FL.pressure(1.0, a["z_surface_m"][p], H) + np.testing.assert_allclose(r["pressure_bar"], sea * FL.pressure(1.0, r["z_surface_m"], H), rtol=1e-5) + np.testing.assert_allclose(r["po2_bar"], r["o2_fraction"] * r["pressure_bar"], rtol=1e-5) + + +def clear_regions(root): # built_world() is shared across tests: leave no regions behind + import shutil + shutil.rmtree(root / "out" / "r2" / "regions", ignore_errors=True) + (root / "places" / "regions.json").unlink(missing_ok=True) + + +class BuildTest(unittest.TestCase): + def tearDown(self): + clear_regions(built_world()) + + def test_build_caches_per_world_and_skips_when_done(self): + root = built_world() + w = RF.WorldCells(root / "out" / "r2") + i = big_river(w.a) + reg = root / "places" / "regions.json" + reg.parent.mkdir(exist_ok=True) + reg.write_text(json.dumps({"regions": [square(float(w.a["g_lat"][i]), float(w.a["g_lon"][i]), 2.5)]})) # its own area (built_world is shared) + m1 = RF.build_areas(root, 2, reg, log=lambda s: None) + self.assertEqual(len(m1), 1) + self.assertTrue(m1[0]["built"]) + d = Path(m1[0]["dir"]) + self.assertTrue((d / "cells.npz").exists() and (d / "meta.json").exists()) + m2 = RF.build_areas(root, 2, reg, log=lambda s: None) + self.assertFalse(m2[0]["built"], "cached: not rebuilt") + stale = d.parent.parent / "0000000000-m1" + stale.mkdir() + RF.build_areas(root, 2, reg, log=lambda s: None) + self.assertFalse(stale.exists(), "results of other world builds are removed") + + def test_redrawn_and_deleted_regions_leave_no_old_area(self): + root = built_world() + w = RF.WorldCells(root / "out" / "r2") + i = big_river(w.a) + la, lo = float(w.a["g_lat"][i]), float(w.a["g_lon"][i]) + reg = root / "places" / "regions.json" + reg.parent.mkdir(exist_ok=True) + reg.write_text(json.dumps({"regions": [square(la, lo, 2)]})) + (m1,) = RF.build_areas(root, 2, reg, log=lambda s: None) + reg.write_text(json.dumps({"regions": [square(la, lo, 3)]})) # redrawn larger + (m2,) = RF.build_areas(root, 2, reg, log=lambda s: None) + self.assertFalse(Path(m1["dir"]).exists(), "the old area is gone") + (Path(m2["dir"]) / "meta.json").write_text("{broken") # a crash mid-save + (m3,) = RF.build_areas(root, 2, reg, log=lambda s: None) + self.assertTrue(m3["built"], "a broken result is rebuilt, not fatal") + reg.write_text(json.dumps({"regions": []})) + self.assertEqual(RF.build_areas(root, 2, reg, log=lambda s: None), []) + self.assertEqual(RF.built_areas(root / "out" / "r2" / "regions", Path(m2["dir"]).parent.name.rsplit("-m", 1)[0]), []) + + + def test_a_failed_build_keeps_the_previous_result(self): + from unittest import mock + root = built_world() + w = RF.WorldCells(root / "out" / "r2") + i = big_river(w.a) + la, lo = float(w.a["g_lat"][i]), float(w.a["g_lon"][i]) + tmp = Path(tempfile.mkdtemp()) + reg = tmp / "regions.json" + reg.write_text(json.dumps({"regions": [square(la, lo, 2)]})) + (m1,) = RF.build_areas(root, 2, reg, log=lambda s: None, regions_root=tmp / "res") + reg.write_text(json.dumps({"regions": [square(la, lo, 3)]})) # extended … + with mock.patch.object(RF, "refine_area", side_effect=RuntimeError("boom")), self.assertRaises(RuntimeError): + RF.build_areas(root, 2, reg, log=lambda s: None, regions_root=tmp / "res") # … and the build fails + self.assertTrue((Path(m1["dir"]) / "cells.npz").exists(), "the old area stays until a build succeeds") + + +class LandOutlineTest(unittest.TestCase): + def test_outline_of_a_landmass_covers_it_simply(self): + root = built_world() + a = RF.WorldCells(root / "out" / "r2").a + land = ~a["ocean"].astype(bool) + i = int(np.argmax(land * a["z_surface_m"])) # a point on the highest landmass + outline = RF.land_outline(a, float(a["g_lat"][i]), float(a["g_lon"][i]), step_deg=2.0) + self.assertGreaterEqual(len(outline), 4) + self.assertLessEqual(len(outline), 500) + import serve + self.assertFalse(serve.RegionStore._crosses([(lo, la) for la, lo in outline]), "a simple polygon") + cells = set(h3.polygon_to_cells(h3.LatLngPoly([tuple(p) for p in outline]), 2)) + comp = RF.landmass(a, float(a["g_lat"][i]), float(a["g_lon"][i])) + inside = np.isin(a["g_ids"][comp], np.array(sorted(cells), dtype=np.uint64)) + self.assertGreater(inside.mean(), 0.97, "the landmass lies inside its outline") + + +class RecordTest(unittest.TestCase): + def test_the_record_keeps_what_was_built_not_later_edits(self): + from unittest import mock + root = built_world() + w = RF.WorldCells(root / "out" / "r2") + i = big_river(w.a) + la, lo = float(w.a["g_lat"][i]), float(w.a["g_lon"][i]) + tmp = Path(tempfile.mkdtemp()) + reg = tmp / "regions.json" + a, b = square(la, lo, 2, "a"), square(-la, lo + 90, 2, "b") + reg.write_text(json.dumps({"regions": [a]})) + real = RF.refine_area + + def edit_meanwhile(*args, **kw): + reg.write_text(json.dumps({"regions": [a, b]})) # region b drawn during the build + return real(*args, **kw) + with mock.patch.object(RF, "refine_area", side_effect=edit_meanwhile): + RF.build_areas(root, 2, reg, log=lambda s: None, regions_root=tmp / "res") + key = next((tmp / "res").glob("*-m*")).name.rsplit("-m", 1)[0] + rec = RF.read_record(tmp / "res", key)["regions"] + self.assertEqual(set(rec), {"a"}, "b was not built: it must not look built") + + +def _subgrid_reference(cells, radius_km): # the first implementation (sets and lists): the oracle + inner = set(int(c) for c in cells) + ring = {n for c in inner for n in h3.grid_ring(c, 1)} - inner + ids = np.array(sorted(inner | ring), dtype=np.uint64) + rings = [[n for n in h3.grid_ring(int(c), 1)] for c in ids] + pos = {int(c): k for k, c in enumerate(ids)} + nb = [[pos[n] for n in r if n in pos] for r in rings] + ptr = np.concatenate([[0], np.cumsum([len(x) for x in nb])]).astype(np.int64) + idx = np.array([k for x in nb for k in x], dtype=np.int64) + return ids, ptr, idx, ~np.isin(ids, cells) + + +def _area_cells_reference(regions, fine_res): # the first implementation (flood fill over every cell) + cells = set() + for r in regions: + cells.update(h3.polygon_to_cells(h3.LatLngPoly([tuple(p) for p in r["outline"]]), fine_res)) + left, out = set(cells), [] + while left: + seed = left.pop() + comp, stack = [seed], [seed] + while stack: + for n in h3.grid_ring(stack.pop(), 1): + if n in left: + left.discard(n) + comp.append(n) + stack.append(n) + out.append(np.array(sorted(comp), dtype=np.uint64)) + return sorted(out, key=lambda a: int(a[0])) + + +class LeanBuildTest(unittest.TestCase): + def test_subgrid_is_the_same_grid(self): + (cells,) = RF.area_cells([square(10, 20, 4)], 4) + g, halo = RF.subgrid(cells, 12742.0) + ids, ptr, idx, h = _subgrid_reference(cells, 12742.0) + np.testing.assert_array_equal(g.ids, ids) + np.testing.assert_array_equal(g.nbr_ptr, ptr) + np.testing.assert_array_equal(g.nbr_idx, idx) + np.testing.assert_array_equal(halo, h) + np.testing.assert_allclose(g.lat, [h3.cell_to_latlng(int(c))[0] for c in ids]) + + def test_areas_are_the_same_areas(self): + cases = [[square(10, 20, 3, "a"), square(10, 25.9, 3, "b")], # touching: one area + [square(10, 20, 3, "a"), square(10, 40, 3, "b")], # apart: two + [square(10, 20, 6, "a"), square(10, 20, 1, "b")], # nested + [square(-30, 170, 3, "a"), square(-30, -176, 3, "b"), square(40, 0, 2, "c")]] # across ±180 + for regions in cases: + got, want = RF.area_cells(regions, 4), _area_cells_reference(regions, 4) + self.assertEqual(len(got), len(want)) + for a, b in zip(got, want): + np.testing.assert_array_equal(a, b) + + def test_a_build_does_not_load_the_refined_regions(self): + from unittest import mock + root = built_world() + tmp = Path(tempfile.mkdtemp()) + reg = tmp / "r.json" + reg.write_text(json.dumps({"regions": [square(10, 20, 3)]})) + RF.build_areas(root, 2, reg, log=lambda s: None, regions_root=tmp / "res") + reg.write_text(json.dumps({"regions": [square(10, 20, 3), square(-20, 60, 3, "b")]})) + real = RF.RegionSet + + def only_empty(root_, *a, **k): + rs = real(root_, *a, **k) + if not rs.empty: + raise AssertionError("a build loaded the refined regions (memory for nothing)") + return rs + with mock.patch.object(RF, "RegionSet", side_effect=only_empty): + RF.build_areas(root, 2, reg, log=lambda s: None, regions_root=tmp / "res") + + +class FingerprintTest(unittest.TestCase): + def test_region_fingerprint_changes_when_a_result_is_rebuilt(self): + d = Path(tempfile.mkdtemp()) + area = RF.results_dir(d, "abc") / "a1" + area.mkdir(parents=True) + np.savez(area / "cells.npz", g_ids=np.array([1], np.uint64)) + (area / "meta.json").write_text("{}") + f1 = RF.regions_fingerprint(d, "abc") + import os, time + st = (area / "cells.npz").stat() + os.utime(area / "cells.npz", ns=(st.st_atime_ns, st.st_mtime_ns + 10**9)) # rebuilt (new model, same cells) + self.assertNotEqual(RF.regions_fingerprint(d, "abc"), f1) + + +class HaloHeightsTest(unittest.TestCase): + def test_halo_keeps_real_heights_in_the_output(self): + root = built_world() + w = RF.WorldCells(root / "out" / "r2") + src = T.TileSource(serve.World(root, 2), seed=7, cache_dir=Path(tempfile.mkdtemp())) + i = big_river(w.a) + (cells,) = RF.area_cells([square(float(w.a["g_lat"][i]), float(w.a["g_lon"][i]), 4)], 4) + g, halo = RF.subgrid(cells, 12742.0) + r = RF.refine_area(g, halo, w, src, C.load(root)[0]) + self.assertLess(float(np.max(r["z_surface_m"])), 20000.0, "no routing walls leak into heights") + + +class CoastTest(unittest.TestCase): + def test_coastal_area_keeps_land_and_rivers_above_the_sea(self): + root = built_world() + w = RF.WorldCells(root / "out" / "r2") + a = w.a + mouths = np.where(a["river"].astype(bool) & a["ocean"][a["recv"]].astype(bool))[0] + i = int(mouths[np.argmax(a["discharge_km3_yr"][mouths])]) + (cells,) = RF.area_cells([square(float(a["g_lat"][i]), float(a["g_lon"][i]), 4)], 4) + g, halo = RF.subgrid(cells, 12742.0) + src = T.TileSource(serve.World(root, 2), seed=7, cache_dir=Path(tempfile.mkdtemp())) + r = RF.refine_area(g, halo, w, src, C.load(root)[0]) + inner = ~halo + self.assertTrue((r["ocean"] & inner).any(), "the test area reaches the sea") + self.assertFalse((r["river"] & inner & (r["elevation_eroded_m"] < 0)).any(), "no river below sea level") + lake_low = r["lake"] & inner & ~r["endorheic"] & (r["z_filled_m"] < 0) + self.assertFalse(lake_low.any(), "no open lake below sea level") + coast_land = inner & ~r["ocean"] & np.isin(np.arange(g.n), g.src[r["ocean"][g.dst]]) + self.assertTrue(np.all(r["elevation_eroded_m"][coast_land] >= 0), "land by the sea is above it") + + +ERA_CUT = """ +[[event]] +name = "cut" +kind = "disintegrate" +center = [{lat}, {lon}] +radius_km = 1000.0 +depth_m = 3000.0 +[[event]] +name = "haze" +kind = "zone" +shape = "circle" +center = [{lat}, {lon}] +radius_km = 1500.0 +edge_km = 5.0 +fields = {{ gravity_g = 0.5, pressure_bar = 2.0, fire_reactivity = 0.4 }} +[eras] +order = ["before", "after"] +default = "after" +[eras.before] +label = "Before" +events = [] +[eras.after] +label = "After" +events = ["cut", "haze"] +""" + +_ERA_WORLD = [] + + +def era_world(): + """(root, far point, cut point, km between them): a small world with an era 'after' that cuts into the land + farthest from the biggest river and holds a zone there (built once per test run).""" + if not _ERA_WORLD: + import atexit + import shutil + from mapgen import eras as ER, pipeline as P + from mapgen.sphere import gc_dist_km + from mapgen.testing import small_world + root = Path(tempfile.mkdtemp(prefix="worldmap-viewer-era-")) + atexit.register(shutil.rmtree, root, True) + small_world(root) + ctx = P.build(root, 2, log=lambda m: None) + g, a = ctx.grid, ctx.data + i = big_river(a) + d = gc_dist_km(g.xyz, g.xyz[i], g.radius_km) + j = int(np.argmax(np.where(~np.asarray(a["ocean"]), d, -1.0))) + (root / "config" / "eras.toml").write_text(ERA_CUT.format(lat=round(float(g.lat[j]), 3), + lon=round(float(g.lon[j]), 3))) + ER.build_all(root, 2, log=lambda m: None, base=ctx) + _ERA_WORLD.append((root, (float(g.lat[i]), float(g.lon[i])), (float(g.lat[j]), float(g.lon[j])), float(d[j]))) + return _ERA_WORLD[0] + + +def write_regions(path: Path, *squares) -> Path: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps({"regions": list(squares)})) + return path + + +class EraAreasTest(unittest.TestCase): + def test_an_area_the_era_left_alone_is_built_once(self): + root, far, cut, dist = era_world() + self.assertGreater(dist, 5500.0, "the fixture needs the river far from the cut") + tmp = Path(tempfile.mkdtemp()) + reg = write_regions(tmp / "regions.json", square(*far, 2, "far"), square(*cut, 2, "cut")) + metas = RF.build_areas(root, 2, reg, log=lambda s: None, regions_root=tmp / "res") + self.assertEqual(sorted({m["era"] for m in metas}), ["after", "before"]) + by = {(m["era"], m["hash"]): m for m in metas} + hf = RF.area_hash(RF.area_cells([square(*far, 2)], 4)[0]) + hc = RF.area_hash(RF.area_cells([square(*cut, 2)], 4)[0]) + self.assertEqual(by[("before", hf)]["key"], by[("after", hf)]["key"]) + self.assertNotEqual(by[("before", hc)]["key"], by[("after", hc)]["key"]) + self.assertEqual(sorted(p.name for p in RF.areas_dir(tmp / "res").iterdir() if p.is_dir()), + sorted({m["key"] for m in metas})) + self.assertEqual(len({m["key"] for m in metas}), 3) + sets = [p for p in (tmp / "res").iterdir() if p.name != RF.AREAS] + self.assertEqual(len(sets), 2) + for s in sets: + links = RF.built_areas(tmp / "res", s.name.rsplit("-m", 1)[0]) + self.assertEqual(len(links), 2) + self.assertTrue(all(p.is_symlink() for p in links)) + self.assertIn(json.loads((s / "era.json").read_text())["era"], ("before", "after")) + again = RF.build_areas(root, 2, reg, log=lambda s: None, regions_root=tmp / "res") + self.assertFalse(any(m["built"] for m in again), "cached for both worlds") + + def test_key_follows_the_values_the_area_reads(self): + import copy + root, far, _, _ = era_world() + cfg, _ = C.load(root) + (cells,) = RF.area_cells([square(*far, 2)], 4) + w = RF.WorldCells(root / "out" / "r2") + k0 = RF.area_key(w, cells, cfg, []) + self.assertEqual(RF.area_key(RF.WorldCells(root / "out" / "r2"), cells, cfg, []), k0, "deterministic") + idx = RF.world_near(w, cells) + par = np.unique([h3.cell_to_parent(int(c), 2) for c in cells]) + self.assertTrue(np.isin(np.searchsorted(w.ids, par.astype(np.uint64)), idx).all(), "the cells under it") + near = RF.WorldCells(root / "out" / "r2") + near.a["T_mean"] = near.a["T_mean"].copy() + near.a["T_mean"][idx[0]] += 1.0 + self.assertNotEqual(RF.area_key(near, cells, cfg, []), k0) + away = RF.WorldCells(root / "out" / "r2") + away.a["T_mean"] = away.a["T_mean"].copy() + away.a["T_mean"][np.setdiff1d(np.arange(len(away.ids)), idx)[0]] += 1.0 + self.assertEqual(RF.area_key(away, cells, cfg, []), k0, "values far away don't matter") + cfg2 = copy.deepcopy(cfg) + cfg2["build"]["seed"] = int(cfg2["build"]["seed"]) + 1 + self.assertNotEqual(RF.area_key(w, cells, cfg2, []), k0) + + def test_an_unbuilt_era_waits_with_a_message(self): + root = built_world() + p = root / "config" / "eras.toml" + p.write_text(ERA_CUT.format(lat=0.0, lon=0.0)) + try: + tmp = Path(tempfile.mkdtemp()) + w = RF.WorldCells(root / "out" / "r2") + i = big_river(w.a) + reg = write_regions(tmp / "regions.json", square(float(w.a["g_lat"][i]), float(w.a["g_lon"][i]), 2)) + logs = [] + metas = RF.build_areas(root, 2, reg, log=logs.append, regions_root=tmp / "res") + self.assertEqual({m["era"] for m in metas}, {"before"}) + self.assertTrue(any("era after is not built" in s for s in logs), logs) + finally: + p.unlink() + + def test_refined_cells_keep_the_eras_air(self): + from mapgen import environment as EN, eras as ER, fields as FL + root, _, cut, _ = era_world() + tmp = Path(tempfile.mkdtemp()) + reg = write_regions(tmp / "regions.json", square(*cut, 2, "cut")) + metas = RF.build_areas(root, 2, reg, log=lambda s: None, regions_root=tmp / "res") + (m,) = [m for m in metas if m["era"] == "after"] + w = RF.WorldCells(ER.era_dir(root, 2, "after")) + with np.load(Path(m["dir"]) / "cells.npz") as z: + r = {k: z[k] for k in z.files} + p = r["parent"] + np.testing.assert_array_equal(r["gravity_g"], w.a["gravity_g"][p]) + self.assertTrue(np.any(np.isclose(r["fire_reactivity"], 0.4)), "the zone reaches the refined cells") + cfg = C.load(root)[0] + H = float(cfg["planet"]["scale_height_km"]) * 1000.0 + sea_p = r["pressure_bar"] / FL.pressure(1.0, r["z_surface_m"], H) + wet = np.sqrt(sea_p / float(cfg["planet"]["sea_level_pressure_bar"])) + zone, _ = EN.holdridge(r["biotemp"], r["P_ann"] * wet, r["T_min"]) + inner = ~r["halo"] + np.testing.assert_array_equal(r["holdridge"][inner], zone[inner]) + + +class LayoutTest(unittest.TestCase): + def setUp(self): + self.root = built_world() + w = RF.WorldCells(self.root / "out" / "r2") + i = big_river(w.a) + self.at = (float(w.a["g_lat"][i]), float(w.a["g_lon"][i])) + self.sq = square(*self.at, 2) + self.tmp = Path(tempfile.mkdtemp()) + self.reg = write_regions(self.tmp / "regions.json", self.sq) + self.res = self.tmp / "res" + + def fingerprint(self): + return T.TileSource(serve.World(self.root, 2), seed=int(C.load(self.root)[0]["build"]["seed"]), + regions_dir=self.res, with_regions=False).fingerprint + + def test_old_layout_is_replaced(self): + h = RF.area_hash(RF.area_cells([self.sq], 4)[0]) + old = RF.results_dir(self.res, self.fingerprint()) / h + old.mkdir(parents=True) + (old / "cells.npz").write_bytes(b"old") + (self.res / "0123456789-m5" / "x").mkdir(parents=True) + (m,) = RF.build_areas(self.root, 2, self.reg, log=lambda s: None, regions_root=self.res) + link = Path(m["dir"]) + self.assertTrue(link.is_symlink()) + with np.load(link / "cells.npz") as z: + self.assertIn("g_ids", z.files) + self.assertFalse((self.res / "0123456789-m5").exists()) + + def test_a_dangling_link_is_rebuilt(self): + import shutil + (m,) = RF.build_areas(self.root, 2, self.reg, log=lambda s: None, regions_root=self.res) + shutil.rmtree(Path(m["dir"]).resolve()) + (m2,) = RF.build_areas(self.root, 2, self.reg, log=lambda s: None, regions_root=self.res) + self.assertTrue(m2["built"]) + self.assertTrue((Path(m2["dir"]) / "cells.npz").exists()) + + def test_a_new_model_rebuilds_the_area(self): + from unittest import mock + (m,) = RF.build_areas(self.root, 2, self.reg, log=lambda s: None, regions_root=self.res) + with mock.patch.object(RF, "MODEL", RF.MODEL + 1): + (m2,) = RF.build_areas(self.root, 2, self.reg, log=lambda s: None, regions_root=self.res) + self.assertNotEqual(m2["key"], m["key"]) + self.assertTrue(m2["built"]) + self.assertEqual(m2["model"], RF.MODEL + 1) + + def test_unused_areas_and_links_go(self): + (m,) = RF.build_areas(self.root, 2, self.reg, log=lambda s: None, regions_root=self.res) + write_regions(self.reg, square(self.at[0] + 5.0, self.at[1] + 5.0, 2)) # redrawn elsewhere + (m2,) = RF.build_areas(self.root, 2, self.reg, log=lambda s: None, regions_root=self.res) + self.assertFalse(Path(m["dir"]).is_symlink() or Path(m["dir"]).exists()) + self.assertEqual([p.name for p in RF.areas_dir(self.res).iterdir() if p.is_dir()], [m2["key"]]) + + +class SeaFloorRefineTest(unittest.TestCase): + @classmethod + def setUpClass(cls): + from unittest import mock + import seafloor as SF + cls.root = built_world() + w = RF.WorldCells(cls.root / "out" / "r2") + a = w.a + river, ocean = a["river"].astype(bool), a["ocean"].astype(bool) + mouth = river & ocean[a["recv"]] & ~ocean + i = int(np.argmax(np.where(mouth, a["discharge_km3_yr"], -1.0))) + cls.at = (float(a["g_lat"][i]), float(a["g_lon"][i])) + (cells,) = RF.area_cells([square(*cls.at, 3)], 4) + cls.g, cls.halo = RF.subgrid(cells, 12742.0) + src = T.TileSource(serve.World(cls.root, 2), seed=7, cache_dir=Path(tempfile.mkdtemp())) + cfg = C.load(cls.root)[0] + cls.r = RF.refine_area(cls.g, cls.halo, w, src, cfg) + real = SF.run + + def untouched(g, z, *args): # the same pass, heights left as they were + _, f, v = real(g, z, *args) + return np.asarray(z, dtype=np.float64), f, v + with mock.patch.object(SF, "run", side_effect=untouched): + cls.r0 = RF.refine_area(cls.g, cls.halo, w, src, cfg) + + def test_the_pass_changes_only_the_sea(self): + r, r0 = self.r, self.r0 + np.testing.assert_array_equal(r["ocean"], r0["ocean"]) + land = ~r0["ocean"] + np.testing.assert_array_equal(r["elevation_eroded_m"][land], r0["elevation_eroded_m"][land]) + np.testing.assert_array_equal(r["river"], r0["river"]) + np.testing.assert_array_equal(r["elevation_eroded_m"][self.halo], r0["elevation_eroded_m"][self.halo]) + sea = r["ocean"] & ~self.halo + self.assertTrue(sea.any(), "the square reaches the sea") + self.assertTrue(np.all(r["elevation_eroded_m"][sea] <= np.maximum(r0["elevation_eroded_m"][sea], -5.0))) + + def test_vents_and_sea_fields_only_in_the_sea(self): + r = self.r + cells = np.searchsorted(self.g.ids, r["vents"]["cell"]) + self.assertTrue(np.all(r["ocean"][cells])) + land = ~r["ocean"] + self.assertTrue(np.all(r["seabed_type"][land] == 0) and np.all(r["vent"][land] == 0)) + for k in ("seabed_type", "seabed_mineral", "bottom_temp_c", "sediment_m", "vent", "canyon_m", "fan_m", + "plateau_id"): + self.assertEqual(len(r[k]), self.g.n, k) + + def test_build_saves_the_vents(self): + tmp = Path(tempfile.mkdtemp()) + reg = write_regions(tmp / "regions.json", square(*self.at, 3, "coast")) + (m,) = RF.build_areas(self.root, 2, reg, log=lambda s: None, regions_root=tmp / "res") + self.assertEqual(m["model"], RF.MODEL) + self.assertIn("vents", m) + self.assertIn("canyon_max_m", m) + with np.load(Path(m["dir"]) / "vents.npz") as z: + self.assertEqual(set(z.files), {"cell", "lat", "lon", "type", "temp_c", "flow", "mineral"}) + self.assertEqual(len(z["cell"]), m["vents"]) + + def test_preview_of_a_refined_region(self): + import seafloor as SF + tmp = Path(tempfile.mkdtemp()) + reg = write_regions(tmp / "regions.json", square(*self.at, 3, "coast")) + RF.build_areas(self.root, 2, reg, log=lambda s: None, regions_root=tmp / "res") + d = SF.preview_region(self.root, 2, "coast", regions_path=reg, regions_root=tmp / "res", out_dir=tmp / "prev") + self.assertGreater((d / "coast-base.jpg").stat().st_size, 1000) + with self.assertRaises(SystemExit): + SF.preview_region(self.root, 2, "nowhere", regions_path=reg, regions_root=tmp / "res", out_dir=tmp / "prev") + + +class PlateauOutlineTest(unittest.TestCase): + def test_outline_holds_the_plateau_and_its_margin(self): + from mapgen import plateaus as PL + from mapgen.sphere import gc_dist_km, latlon_to_xyz + from tests.helpers import small_grid + p = {"name": "east-flank", "center": [-39.9, -18.0], "area_km2": 3.0e6, "elongation": 1.9, + "azimuth_deg": 14.6, "top_m": [1500.0, 3000.0]} + o = RF.plateau_outline(p, 7, 12742.0) + self.assertEqual(len(o), 96) + self.assertFalse(serve.RegionStore._crosses([(lo, la) for la, lo in o]), "a simple polygon") + cells = np.array(sorted(h3.polygon_to_cells(h3.LatLngPoly([tuple(q) for q in o]), 3)), dtype=np.uint64) + g = small_grid(3) + pid = PL.cell_ids(g.xyz, [p], 7, 12742.0) + self.assertGreater(float(np.isin(g.ids[pid == 0], cells).mean()), 0.99) + q = np.array(o) + d = gc_dist_km(latlon_to_xyz(q[:, 0], q[:, 1]), latlon_to_xyz(*p["center"]), 12742.0) + self.assertTrue(np.all(d >= 300.0)) + + +class FastPathsTest(unittest.TestCase): + """Speed-ups with the same results as the per-cell h3 calls and loops they replaced.""" + + def test_cell_parents_match_h3(self): + rng = np.random.default_rng(1) + base = [c for r0 in h3.get_res0_cells() for c in h3.cell_to_children(r0, 1)] + for res_child, res_parent in ((7, 5), (4, 2), (5, 5), (3, 0)): + cells = np.array([h3.cell_to_center_child(int(c), res_child) for c in rng.choice(base, 50)] + + [h3.cell_to_children(int(c), res_child)[-1] for c in rng.choice(base, 50)], np.uint64) + want = [h3.cell_to_parent(int(c), res_parent) for c in cells] + np.testing.assert_array_equal(RF.cell_parents(cells, res_parent), np.array(want, np.uint64)) + with self.assertRaises(ValueError): + RF.cell_parents(np.array([h3.cell_to_center_child(int(base[0]), 3)], np.uint64), 4) + + def test_k_mult_matches_the_lookup(self): + age = np.array([0, 1, 2, 3, 4, 5, 6, 7, 8, -1, 3.9, 9], dtype=np.float64) + want = np.vectorize(lambda c: RF.K_MULT.get(int(c), 1.0))(age).astype(np.float64) + np.testing.assert_array_equal(RF.k_mult(age), want) + np.testing.assert_array_equal(RF.k_mult(age.astype(np.int8)[:9]), want[:9]) + + def test_subgrid_in_parallel_is_the_serial_grid(self): + from unittest import mock + import h3par + (cells,) = RF.area_cells([square(10, 20, 4)], 4) + g1, h1 = RF.subgrid(cells, 12742.0) + with mock.patch.object(h3par, "MIN_PARALLEL", 1), mock.patch.dict("os.environ", {"MAPVIEW_H3_WORKERS": "3"}): + g2, h2 = RF.subgrid(cells[::-1].copy(), 12742.0) # unsorted input too + for k in ("ids", "lat", "lon", "xyz", "area_km2", "nbr_ptr", "nbr_idx"): + np.testing.assert_array_equal(getattr(g1, k), getattr(g2, k)) + np.testing.assert_array_equal(h1, h2) |