diff options
Diffstat (limited to 'refine.py')
| -rw-r--r-- | refine.py | 1273 |
1 files changed, 1273 insertions, 0 deletions
diff --git a/refine.py b/refine.py new file mode 100644 index 0000000..4c0da9b --- /dev/null +++ b/refine.py @@ -0,0 +1,1273 @@ +"""Regional refinement (spec docs/superpowers/specs/2026-09-25-regional-refinement-design.md): inside drawn regions, +re-run erosion, rivers, lakes, ground and biomes on H3 cells two resolutions finer than the world, with the world +build as the boundary. Generated terrain (`idea`). Outside mapgen/ so the world build cache is unaffected.""" +from __future__ import annotations + +import copy +import hashlib +import json +import math +import os +import shutil +import tempfile +import threading +import time +from pathlib import Path +from types import SimpleNamespace + +import h3.api.basic_int as h3 +import numpy as np + +import h3par + +from mapgen import fields as FL, hydrology as HY, ice as IC +from mapgen.config import params +from mapgen.environment import DEFAULTS as ENV_DEFAULTS, ground, holdridge +from mapgen.graph import (OCEAN_MIN_KM2, accumulate, components, drop_smooth_cache, priority_flood, receiver_levels, + smooth_km, steepest_receivers) +from scipy.spatial import cKDTree + +from mapgen.grid import Grid +from mapgen.sphere import latlon_to_xyz + +import seafloor as SF +import servecache + +MODEL = 9 # bump when the refinement model changes (part of the cache key); 2: lake depth, shoulders; + # 3: halo outlets and walls; 4: 1 km hillslope (canyons stay narrow); 5: smooth bedrock + # means, connected sea, lake-aware incision, border report; 6: sea-floor pass (canyons, + # fans, ponds, volcanic fields, vents), zone fields from the world cells; 7: world lakes + # (deep, carved beds) fill to their surface and are ways out; no lake re-carving; + # 8: sea only where it reaches the open ocean (walled-off world ocean: lagoon lakes); + # 9: world lakes keep their world way out (no fine sill dams them above their level) +FINER = 2 # fine resolution = world resolution + FINER (49 children per world cell) +LAPSE_C_PER_KM = 6.5 +RELIEF_MIN_KM = 5.0 # procedural relief kept at wavelengths ≥ ≈ the fine cell spacing +COPY = ("continental", "age_class", "plate", "lithology", "landform", "seasonality", "P_ann", "P_jun", "P_dec", "PET", + "dist_ocean_km", "m_o2_zones", "m_gravity_zones", "deposits", "deposit_main") +LAPSE = ("T_mean", "T_min", "T_jun", "T_dec", "biotemp") + + + +def load_regions(path: Path) -> list[dict]: + path = Path(path) + return json.loads(path.read_text()).get("regions", []) if path.exists() else [] + + +def area_cells(regions: list[dict], fine_res: int) -> list[np.ndarray]: + """Cells of the union of all outlines, split into connected areas: outlines that overlap or touch merge (each + outline's own cells count as one piece). Arrays, not sets: a continent has millions of cells.""" + pieces, caps = [], [] + for r in regions: + c = np.unique(np.array(h3.polygon_to_cells(h3.LatLngPoly([tuple(p) for p in r["outline"]]), fine_res), + dtype=np.uint64)) + if len(c): + pieces.append(c) + o = np.asarray(r["outline"], dtype=np.float64) + v = latlon_to_xyz(o[:, 0], o[:, 1]).reshape(-1, 3) + m = v.mean(axis=0) + m = m / max(np.linalg.norm(m), 1e-12) + caps.append((m, float(np.max(np.arccos(np.clip(v @ m, -1, 1)))))) # a cap holding the outline + parent = list(range(len(pieces))) + + def find(k): + while parent[k] != k: + parent[k] = parent[parent[k]] + k = parent[k] + return k + order = sorted(range(len(pieces)), key=lambda k: len(pieces[k])) + for a_ in range(len(order)): + for b_ in range(a_ + 1, len(order)): + i_, j_ = order[a_], order[b_] + (ci, ri), (cj, rj) = caps[i_], caps[j_] + apart = math.acos(float(np.clip(ci @ cj, -1, 1))) > ri + rj + 0.01 # caps apart (+ ≈ 2 cells): no touch + if not apart and find(i_) != find(j_) and _touch(pieces[i_], pieces[j_]): + parent[find(i_)] = find(j_) + groups = {} + for k in range(len(pieces)): + groups.setdefault(find(k), []).append(pieces[k]) + out = [np.unique(np.concatenate(g)) if len(g) > 1 else g[0] for g in groups.values()] + return sorted(out, key=lambda a: int(a[0])) + + +def _touch(small: np.ndarray, big: np.ndarray) -> bool: + """Whether two cell sets overlap or are neighbours (the smaller one's cells and their rings against the other).""" + if np.isin(small, big, assume_unique=True).any(): + return True + for c in small.tolist(): # only the edge of the smaller one can touch + ring = np.array(h3.grid_ring(c, 1), dtype=np.uint64) + k = np.searchsorted(big, ring) + k[k >= len(big)] = 0 + if (big[k] == ring).any(): + return True + return False + + +def landmass(a, lat: float, lon: float) -> np.ndarray: + """Indices of the world's land cells connected to the land cell nearest (lat, lon) (a continent or island).""" + ids = np.asarray(a["g_ids"]) + land = ~np.asarray(a["ocean"]).astype(bool) + L = np.where(land)[0] + start = L[cKDTree(np.asarray(a["g_xyz"])[L]).query(latlon_to_xyz(lat, lon).reshape(3))[1]] + seen, stack = {int(start)}, [int(start)] + while stack: + for n in h3.grid_ring(int(ids[stack.pop()]), 1): + j = int(np.searchsorted(ids, np.uint64(n))) + if j < len(ids) and ids[j] == n and land[j] and j not in seen: + seen.add(j) + stack.append(j) + return np.array(sorted(seen), dtype=np.int64) + + +def land_outline(a, lat: float, lon: float, step_deg: float = 0.5, margin: int = 1, max_points: int = 500) -> list: + """A simple outline polygon [[lat, lon], …] around the landmass at (lat, lon): its cells on a step_deg grid, + grown by `margin` grid cells (the coast and a strip of sea), holes filled, traced along the grid lines.""" + from scipy.ndimage import binary_dilation, binary_fill_holes, label + comp = landmass(a, lat, lon) + la = np.asarray(a["g_lat"])[comp] + lo = np.asarray(a["g_lon"])[comp] + c0 = float(np.degrees(np.arctan2(np.sin(np.radians(lo)).mean(), np.cos(np.radians(lo)).mean()))) + rel = (lo - c0 + 180.0) % 360.0 - 180.0 # longitudes around the landmass's middle + while True: + r0, q0 = np.floor(la.min() / step_deg) - margin - 1, np.floor(rel.min() / step_deg) - margin - 1 + R = int(np.ceil(la.max() / step_deg) - r0 + margin + 2) + Q = int(np.ceil(rel.max() / step_deg) - q0 + margin + 2) + m = np.zeros((R, Q), bool) + m[(np.floor(la / step_deg) - r0).astype(int), (np.floor(rel / step_deg) - q0).astype(int)] = True + m = binary_fill_holes(binary_dilation(m, np.ones((3, 3), bool), iterations=margin) if margin else m) + lab, _ = label(m) # the piece holding the landmass + m = lab == lab[int(np.floor(la[0] / step_deg) - r0), int(np.floor(rel[0] / step_deg) - q0)] + while True: # no cells touching at a corner only + d1 = m[:-1, :-1] & m[1:, 1:] & ~m[:-1, 1:] & ~m[1:, :-1] + d2 = m[:-1, 1:] & m[1:, :-1] & ~m[:-1, :-1] & ~m[1:, 1:] + if not (d1.any() or d2.any()): + break + m[:-1, 1:] |= d1 + m[:-1, :-1] |= d2 + m = binary_fill_holes(m) + nxt = {} # boundary edges, inside on the left + for r, q in zip(*np.nonzero(m)): + if r == 0 or not m[r - 1, q]: + nxt[(r, q + 1)] = (r, q) + if r == R - 1 or not m[r + 1, q]: + nxt[(r + 1, q)] = (r + 1, q + 1) + if q == 0 or not m[r, q - 1]: + nxt[(r, q)] = (r + 1, q) + if q == Q - 1 or not m[r, q + 1]: + nxt[(r + 1, q + 1)] = (r, q + 1) + start = min(nxt) + loop, v = [start], nxt[start] + while v != start: + loop.append(v) + v = nxt[v] + pts = [p for k, p in enumerate(loop) # corners only + if (loop[k - 1][0] - p[0], loop[k - 1][1] - p[1]) != (p[0] - loop[(k + 1) % len(loop)][0], + p[1] - loop[(k + 1) % len(loop)][1])] + if len(pts) <= max_points: + return [[round(float((r + r0) * step_deg), 6), round(float((q + q0) * step_deg + c0 + 180.0) % 360.0 - 180.0, 6)] + for r, q in pts] + step_deg *= 1.25 + + +def plateau_outline(p: dict, seed: int, radius_km: float, margin_km: float = 300.0, n: int = 96) -> list: + """An outline [[lat, lon], …] around a sunken plateau (mapgen.plateaus' noise-warped ellipse) grown by margin_km: + per bearing, the farthest point still inside the plateau plus the margin (star-shaped: a simple polygon).""" + from mapgen import plateaus as PL + from mapgen.sphere import east_north, great_circle_point, xyz_to_latlon + c = latlon_to_xyz(*p["center"]) + e, nn = east_north(c[None]) + a, _ = PL.semi_axes(p) + steps = np.linspace(0.0, a * (1.0 + PL.EDGE_WARP) / (1.0 - PL.EDGE_WARP), 400) + out = [] + for az in np.linspace(0.0, 2.0 * np.pi, n, endpoint=False): + t = np.cos(az) * nn[0] + np.sin(az) * e[0] + pts = np.cos(steps / radius_km)[:, None] * c + np.sin(steps / radius_km)[:, None] * t + inside = np.flatnonzero(PL.rho(pts, p, seed, radius_km) < 1.0) + r = float(steps[inside[-1]]) if len(inside) else 0.0 + la, lo = xyz_to_latlon(great_circle_point(c, t, r + margin_km, radius_km)) + out.append([round(float(la), 4), round(float(lo), 4)]) + return out + + +def region_key(region: dict) -> str: + """A region's key in the build record: its id, or its outline for hand-written regions without one.""" + return str(region.get("id") or "outline-" + outline_hash(region)) + + +def outline_hash(region: dict) -> str: + return hashlib.sha1(json.dumps(region["outline"]).encode()).hexdigest()[:12] + + +def region_areas(regions: list[dict], areas: list[np.ndarray], fine_res: int) -> dict: + """{region id: {"outline": outline hash, "area": hash of the area holding it (None: too small for a cell)}}.""" + out = {} + for r in regions: + cells = h3.polygon_to_cells(h3.LatLngPoly([tuple(p) for p in r["outline"]]), fine_res) + c = np.uint64(cells[0]) if cells else None + hit = None + for a in areas: + k = int(np.searchsorted(a, c)) if c is not None else len(a) + if k < len(a) and a[k] == c: + hit = area_hash(a) + break + out[region_key(r)] = {"outline": outline_hash(r), "area": hit} + return out + + +def read_record(regions_root: Path, key: str): + """What the last successful build made of which outline (build_areas writes it), or None.""" + try: + rec = json.loads((results_dir(regions_root, key) / "outlines.json").read_text()) + except (OSError, ValueError): + return None + ok = isinstance(rec, dict) and isinstance(rec.get("regions"), dict) and all( + isinstance(v, dict) and "outline" in v and "area" in v for v in rec["regions"].values()) + return rec if ok else None # a damaged record counts as none (a rebuild) + + +def area_hash(cells: np.ndarray) -> str: + return hashlib.sha1(np.asarray(cells, dtype=np.uint64).tobytes()).hexdigest()[:12] + + +AREAS = "areas" # <regions_root>/areas/<key>/: results by content, shared between worlds +KEY_SKIP = ("era_mask",) # era bookkeeping (and zone_* weights), not terrain: equal ground, equal key +KEY_RINGS = 3 # world cells around an area whose values its refinement reads (relief raster, + # mean correction, river crossings) + + +def areas_dir(regions_root: Path) -> Path: + return Path(regions_root) / AREAS + + +def world_dirs(root: Path, res: int, log=print) -> list: + """(era, world folder) of every world the viewer can show: the base build first (named after the first era + without events, else "base"), then each era with events that has been built (out/r<res>/eras/<era>/).""" + from mapgen import config as C, eras as ER + root = Path(root) + _, tect = C.load(root) + order = (tect.get("eras") or {}).get("order", []) + plain = [n for n in order if not C.era_events(tect, n)] + out = [(plain[0] if plain else "base", root / "out" / f"r{res}")] + for n in order: + if not C.era_events(tect, n): + continue + d = ER.era_dir(root, res, n) + if (d / "cells.npz").exists(): + out.append((n, d)) + else: + log(f"refine: era {n} is not built (run mapgen.py build): its areas wait") + return out + + +def cell_parents(cells, res: int) -> np.ndarray: + """h3.cell_to_parent for an array of cells (all of resolution ≥ res), by the index bits: the resolution field + set to res, the digits below it set to 7 (unused).""" + cells = np.asarray(cells, dtype=np.uint64) + if len(cells) and int(((cells >> np.uint64(52)) & np.uint64(15)).min()) < res: + raise ValueError(f"cell_parents: a cell is coarser than resolution {res}") + unused = 0 + for r in range(res + 1, 16): + unused |= 7 << ((15 - r) * 3) + out = (cells & np.uint64(~(15 << 52) & (2**64 - 1))) | np.uint64(res << 52) + return out | np.uint64(unused) + + +def world_near(world: "WorldCells", cells, rings: int = KEY_RINGS) -> np.ndarray: + """Indices of the world cells under the area and `rings` rings around them.""" + par = np.unique(cell_parents(cells, world.res)) + near = np.unique(np.fromiter((n for c in par.tolist() for n in h3.grid_disk(c, rings)), dtype=np.uint64)) + k = np.minimum(np.searchsorted(world.ids, near), len(world.ids) - 1) + return k[world.ids[k] == near].astype(np.int64) + + +def area_key(world: "WorldCells", cells, cfg: dict, plateaus: list) -> str: + """Content key of an area's refinement: the model, its cells, the config and every world value it reads (the + world cells under and around it, their river levels, the plateaus there). An era that left those values as they + were gives the same key: the area is built once for both.""" + import tiles + idx = world_near(world, cells) + h = hashlib.sha256(f"m{MODEL}:t{tiles.VERSION}:f{FINER}".encode()) + h.update(np.asarray(cells, dtype=np.uint64).tobytes()) + h.update(json.dumps(cfg, sort_keys=True, default=str).encode()) + for k in sorted(world.a.keys()): + if k in KEY_SKIP or k.startswith("zone_"): + continue + v = world.a.peek(k) + if v.shape[:1] == world.ids.shape: + h.update(k.encode()) + h.update(np.ascontiguousarray(v[idx]).tobytes()) + h.update(np.ascontiguousarray(world.river_level[idx]).tobytes()) + pid = np.asarray(world.a.peek("plateau_id"))[idx] if "plateau_id" in world.a else np.zeros(0, np.int64) + near = [int(i) for i in np.unique(pid) if 0 <= i < len(plateaus)] + h.update(json.dumps([plateaus[i] for i in near], sort_keys=True, default=str).encode()) + return h.hexdigest()[:16] + + +def _write_json(path: Path, obj) -> None: + fd, tmp = tempfile.mkstemp(dir=path.parent, suffix=".json") + with os.fdopen(fd, "w") as fh: + json.dump(obj, fh) + os.replace(tmp, path) + + +def _remove(p: Path) -> None: + if p.is_symlink(): + p.unlink(missing_ok=True) + elif p.is_dir(): + shutil.rmtree(p, ignore_errors=True) + + +def _link(set_dir: Path, name: str, target: Path) -> Path: + """set_dir/name → target as a relative symlink, replaced atomically (an old real folder goes first).""" + link = set_dir / name + rel = os.path.relpath(target, set_dir) + if link.is_symlink() and os.readlink(link) == rel: + return link + tmp = set_dir / f".{name}.link-{os.getpid()}" + tmp.unlink(missing_ok=True) + os.symlink(rel, tmp) + if link.exists() and not link.is_symlink(): + shutil.rmtree(link) + os.replace(tmp, link) + return link + + +def _rings(ids) -> tuple[np.ndarray, np.ndarray]: + """Neighbours of each cell (h3.grid_ring order), flat, with per-cell counts (pentagons have five).""" + return h3par.run(h3par.rings, ids) + + +def subgrid(cells: np.ndarray, radius_km: float): + """The area cells plus a one-cell halo ring, as a Grid whose neighbours stay inside that set (numpy: a continent's + millions of cells must not become Python sets and lists).""" + cells = np.asarray(cells, dtype=np.uint64) + if len(cells) > 1 and not (cells[1:] > cells[:-1]).all(): + cells = np.unique(cells) + nb_a, cnt_a = _rings(cells) + u = np.unique(nb_a) + p = np.minimum(np.searchsorted(cells, u), len(cells) - 1) + ring = u[cells[p] != u] # the halo: neighbours outside the area + del u, p + ids = np.sort(np.concatenate([cells, ring]), kind="stable") # sorted: the area and its ring + nb_r, cnt_r = _rings(ring) # the area's own rings are known: only the ring's + pc, pr = np.searchsorted(ids, cells), np.searchsorted(ids, ring) + order = np.empty(len(ids), np.int64) # each id's ring, in ids order + order[pc] = np.arange(len(cells)) + order[pr] = len(cells) + np.arange(len(ring)) + cnt = np.concatenate([cnt_a, cnt_r]) + start = np.concatenate([[0], np.cumsum(cnt)[:-1]]) + counts = cnt[order] + pool = np.concatenate([nb_a, nb_r]) + del nb_a, nb_r + first = np.repeat(start[order] - np.concatenate([[0], np.cumsum(counts)[:-1]]), counts) + flat = pool[first + np.arange(int(counts.sum()))] + del pool, first + k = np.searchsorted(ids, flat) + k[k >= len(ids)] = 0 + keep = ids[k] == flat # neighbours inside the set (in ring order) + per = np.add.reduceat(keep.astype(np.int64), np.concatenate([[0], np.cumsum(counts)[:-1]])) if len(ids) else counts + ptr = np.concatenate([[0], np.cumsum(per)]).astype(np.int64) + idx = k[keep].astype(np.int64) + del flat, k, keep + res = h3.get_resolution(int(ids[0])) + ll, area = h3par.run(h3par.centres_areas, ids) + g = Grid(res, radius_km, ids, ll[:, 0], ll[:, 1], latlon_to_xyz(ll[:, 0], ll[:, 1]), area * radius_km ** 2, ptr, idx) + halo = np.zeros(len(ids), bool) + halo[pr] = True + return g, halo + + +class LazyArrays(dict): + """An npz file's arrays, each read on first use (a refinement needs a few of the world's fields, not all).""" + def __init__(self, path: Path): + super().__init__() + self._npz = np.load(path) + self._keys = set(self._npz.files) + + def __missing__(self, k): + if k not in self._keys: + raise KeyError(k) + v = self[k] = self._npz[k] + return v + + def __contains__(self, k): + return k in self._keys + + def keys(self): + return self._keys + + def __iter__(self): + return iter(self._keys) + + def __len__(self): + return len(self._keys) + + def peek(self, k): + """The array without keeping it (a key hashes every field once).""" + return dict.__getitem__(self, k) if dict.__contains__(self, k) else self._npz[k] + + +class WorldCells: + def __init__(self, out_dir: Path): + self.a = LazyArrays(Path(out_dir) / "cells.npz") + self.meta = json.loads((Path(out_dir) / "cells_meta.json").read_text()) + self.ids = self.a["g_ids"] + self.res = h3.get_resolution(int(self.ids[0])) + self._levels = None + self._tree = None + + @property + def tree(self) -> cKDTree: + if self._tree is None: + self._tree = cKDTree(np.asarray(self.a["g_xyz"], dtype=np.float64)) + return self._tree + + @property + def river_level(self) -> np.ndarray: + """The world rivers' graded water levels (rivers.py, spec §10): how deep a river may cut where it leaves.""" + if self._levels is None: + import rivers as RV + self._levels = RV.RiverNet(self.a, self.meta["legends"], float(self.meta["radius_km"])).level + return self._levels + + def index(self, fine_ids) -> np.ndarray: + return np.searchsorted(self.ids, cell_parents(fine_ids, self.res)) + + +def initial_fields(g, halo, world: WorldCells, src) -> dict: + p = world.index(g.ids) + a = world.a + z = src.relief_at(g.lat, g.lon, RELIEF_MIN_KM) + dz_km = (z - a["z_surface_m"][p]) / 1000.0 + out = {"parent": p, "elevation_m": z, "ocean_world": a["ocean"][p].astype(bool)} + for k in COPY: + if k in a: + out[k] = a[k][p] + for k in LAPSE: + out[k] = a[k][p] - LAPSE_C_PER_KM * dz_km + out["T_range"] = a["T_range"][p] + out["biotemp"] = np.clip(out["biotemp"], 0.0, 30.0) + return out + + +EROSION = {"steps": 8, "dt_myr": 1.0, "k": 0.004, "m": 0.8, "hillslope_km": 1.0, "sediment_fill": 0.15, + "area_per_km3": 2000.0} # water flux → equivalent drainage area (0.5 m/yr runoff) +K_MULT = {0: 0.0, 1: 1.5, 2: 1.2, 3: 1.0, 4: 1.0, 5: 0.6, 6: 0.9, 7: 0.8} # age class → erodibility (as the world) + + +def k_mult(age_class) -> np.ndarray: + """K_MULT per cell (1.0 for classes it lacks).""" + c = np.asarray(age_class) + c = np.trunc(c).astype(np.int64) if c.dtype.kind == "f" else c.astype(np.int64) + table = np.array([K_MULT.get(i, 1.0) for i in range(max(K_MULT) + 1)], dtype=np.float64) + ok = (c >= 0) & (c < len(table)) + return np.where(ok, table[np.where(ok, c, 0)], 1.0) + + +def runoff_km3(g, f) -> np.ndarray: + p, pet = f["P_ann"], np.maximum(f["PET"], 1e-6) + aet = p / np.sqrt(1.0 + (p / pet) ** 2) # Pike (1964), as the world + return np.maximum(p - aet, 0.0) * g.area_km2 * 1e-6 + + +def crossings(g, halo, world): + """World river segments crossing the area outline: inflow (km³/yr) at the first inside cell, and the halo cells + where they leave, with the world river's level there.""" + a = world.a + riv = np.where(a["river"].astype(bool))[0] + near = np.zeros(len(world.ids), bool) # only segments near the area can cross its outline + near[world_near(world, g.ids)] = True + riv = riv[near[riv] | near[np.asarray(a["recv"])[riv]]] + res = h3.get_resolution(int(g.ids[0])) + pos = {int(c): k for k, c in enumerate(g.ids)} + inflow, exits, levels = np.zeros(g.n), [], [] + for i in riv: + j = int(a["recv"][i]) + if j == i: + continue + pa, pb = a["g_xyz"][i], a["g_xyz"][j] + n = max(2, int(np.linalg.norm(pb - pa) * g.radius_km / 1.0)) # 1 km steps + t = np.linspace(0, 1, n)[:, None] + pts = pa * (1 - t) + pb * t + pts /= np.linalg.norm(pts, axis=1, keepdims=True) + lat, lon = np.degrees(np.arcsin(pts[:, 2])), np.degrees(np.arctan2(pts[:, 1], pts[:, 0])) + k = [pos.get(h3.latlng_to_cell(float(la), float(lo), res), -1) for la, lo in zip(lat, lon)] + inside = [kk >= 0 and not halo[kk] for kk in k] + for s in range(1, len(k)): + if not inside[s - 1] and inside[s]: + inflow[k[s]] += float(a["discharge_km3_yr"][i]) + if inside[s - 1] and not inside[s] and k[s] >= 0: + exits.append(k[s]) + lv = world.river_level + li = lv[i] if np.isfinite(lv[i]) else a["z_surface_m"][i] + lj = lv[j] if np.isfinite(lv[j]) else (0.0 if a["ocean"][j] else a["z_surface_m"][j]) + levels.append(float(min(li, lj))) + return inflow, np.array(exits, dtype=np.int64), np.array(levels) + + +WALL_M = 1.0e6 # halo cells that are no way out: walls for routing + + +def outlets(g, halo, world, parent, z=None) -> np.ndarray: + """Halo cells where water may leave the area: where the world's own flow leaves it (a world cell with children + inside drains into this halo cell's world cell), or the world sea. Elsewhere (e.g. upstream of an inflowing + river) the halo is a wall, so inflows cross the area instead of turning back out.""" + a = world.a + inside = np.unique(parent[~halo]) + targets = np.setdiff1d(np.asarray(a["recv"])[inside], inside) + out = halo & (np.isin(parent, targets) | np.asarray(a["ocean"])[parent].astype(bool)) + if z is not None: # and wherever the ground outside lies below the area cells next to it (water runs out there) + inner_nb = ~halo[g.dst] + low_in = np.full(g.n, np.inf) + np.minimum.at(low_in, g.src[inner_nb], np.asarray(z)[g.dst[inner_nb]]) + out |= halo & (np.asarray(z) < low_in) + if not out.any(): # a closed basin in the world: its lowest border cell drains + k = np.where(halo)[0] + out[k[np.argmin(np.asarray(a["z_surface_m"])[parent[k]])]] = True + return out + + +def erode_area(g, z, halo, sea, water, kmult, P=EROSION, outlet=None, progress=None): + """Stream-power incision (implicit, as the world build) driven by water flux, with the halo and the sea fixed; + water leaves through the sea and the outlet halo cells only (default: the whole halo).""" + z = np.asarray(z, dtype=np.float64).copy() + fixed = halo | sea + sinks = sea | (halo if outlet is None else outlet) + wall = halo & ~sinks + ar = np.arange(g.n) + for step in range(int(P["steps"])): + if progress: + progress("erosion", 0.1 + 0.5 * step / int(P["steps"])) + base = np.where(sea, 0.0, z) # rivers cut to sea level, never toward the seabed + zf = priority_flood(g, np.where(wall, WALL_M, base), sinks) + zb = np.where(fixed, base, z + P["sediment_fill"] * (zf - z)) + recv, _, dist = steepest_receivers(g, zf) + recv = np.where(fixed, ar, recv) + levels = receiver_levels(recv) + q = accumulate(recv, levels, water) + F = P["k"] * kmult * P["dt_myr"] * (q * P["area_per_km3"]) ** P["m"] / np.where(np.isfinite(dist), dist, 1.0) + F[fixed] = 0.0 + floor = zb.copy() # the level of the way out each cell drains to + for lv in levels[1:]: + floor[lv] = floor[recv[lv]] + zn = zb.copy() + for lv in levels[1:]: + r = recv[lv] + cut = np.minimum(zb[lv], (zb[lv] + F[lv] * zn[r]) / (1.0 + F[lv])) + zn[lv] = np.maximum(cut, np.minimum(floor[lv], zb[lv])) # never below it (hollows are no base level) + zn = smooth_km(g, zn, P["hillslope_km"], keep=True) + z = np.where(fixed, z, np.minimum(z, zn)) + drop_smooth_cache(g) + return z + + +def mean_correction(xyz, z, parent, halo, target, world, iters=12, hold=None, memo=None) -> np.ndarray: + """A smooth height correction (the compact 4-of-5 kernel over nearby world cell centres, evaluated at each fine + cell) that makes each world cell's area children average the target (the world's bedrock): no block steps + between world cells. Tapers to 0 over the two cell rows next to the halo (held at world values), so exits and + the edge meet the world without a step. Cells in `hold` (e.g. the fixed sea) count in the means but are not moved. + memo: a dict shared by calls on the same cells (the kernel and the taper are computed once).""" + z = np.asarray(z, dtype=np.float64) + inner = ~halo + wx = np.asarray(world.a["g_xyz"], dtype=np.float64) + if memo: + near, idx, w, has, taper = memo["geometry"] + else: + pids = np.unique(parent) + ring = world.tree.query(wx[pids], k=7)[1].ravel() # the parents and their world neighbours + near = np.unique(np.concatenate([pids, ring])) + nw = max(1, h3par.workers()) # KD-tree queries on threads: same answers + d, idx = cKDTree(wx[near]).query(np.asarray(xyz, dtype=np.float64), k=5, workers=nw) + w = np.clip(1.0 - d[:, :4] / np.maximum(d[:, 4:5], 1e-12), 0.0, None) ** 2 + w = w / np.maximum(w.sum(1, keepdims=True), 1e-300) + idx = idx[:, :4] + has = np.bincount(parent[inner], minlength=len(wx)) > 0 + xyz = np.asarray(xyz, dtype=np.float64) + step = np.median(cKDTree(xyz).query(xyz, k=2, workers=nw)[0][:, 1]) # the fine cell spacing + taper = (np.clip((cKDTree(xyz[halo]).query(xyz, workers=nw)[0] / step - 1.0) / 2.0, 0.0, 1.0) if halo.any() + else 1.0) + if memo is not None: + memo["geometry"] = (near, idx.astype(np.int32), w, has, taper) + corr = np.zeros(len(z)) + for _ in range(iters): + zc = z + corr + sums = np.bincount(parent[inner], weights=zc[inner], minlength=len(wx)) + cnt = np.bincount(parent[inner], minlength=len(wx)) + resid = np.where(has, np.asarray(target, dtype=np.float64) - sums / np.maximum(cnt, 1), 0.0) + corr += taper * (resid[near][idx] * w).sum(1) + corr[halo] = 0.0 + if hold is not None: + corr[hold] = 0.0 + return corr + + +RIVER_MIN_KM3_YR = 0.2 +LAKE_MIN_DEPTH_M = 15.0 # only hollows deeper than this hold lakes (5 km cells: shallow dips are noise) + + +def _p_for_runoff(target_mm, pet): + """Precipitation (mm) whose Pike runoff is target_mm (bisection): lets the world hydrology code carry inflows.""" + lo, hi = np.zeros_like(target_mm), target_mm + pet * 4 + 10 + for _ in range(60): + mid = (lo + hi) / 2 + run = mid - mid / np.sqrt(1 + (mid / pet) ** 2) + lo, hi = np.where(run < target_mm, mid, lo), np.where(run < target_mm, hi, mid) + return (lo + hi) / 2 + + +def _sea(g, z, ocean_world, halo, min_km2: float = OCEAN_MIN_KM2) -> np.ndarray: + """Sea: cells at or below 0 m under the world's ocean that reach it beyond the area (through the halo), or a + separate basin as big as the world counts as sea (its own rule). World ocean that the fine heights wall off is + no sea: a lagoon, which the lakes take (author 2026-10-03). Inland dry basins stay land.""" + wet = np.asarray(z) <= 0 + lab = components(g, wet) + ow = wet & np.asarray(ocean_world) + seas = np.unique(lab[ow & np.asarray(halo)]) + area = np.bincount(lab[ow], weights=np.asarray(g.area_km2)[ow], minlength=int(lab.max()) + 1 if lab.size else 0) + seas = np.union1d(seas[seas >= 0], np.flatnonzero(area >= min_km2)) + return wet & np.isin(lab, seas) + + +OUTFLOW_TOL_M = 1.0 # a world lake may stand this much above its world level before its way out is cut +OUTFLOW_DROP_M = 0.5 # the cut way out starts this far below the world lake's level + + +def world_outflows(g, z, halo, sinks, parent, a) -> np.ndarray: + """z with each world lake's way out kept: a lake that the fine relief dams above its world level (a sill the world + cells' means hide) gets a channel cut along its world drainage path, down to where that path already drains lower + (the sea, a lower lake, a way out of the area), never below sea level. The channel follows the least digging (Dijkstra on the height above + the lake's level) and falls evenly from just below the lake's level. Lakes that drain low enough stay as they are.""" + import heapq + z = np.asarray(z, dtype=np.float64).copy() + sinks = np.asarray(sinks, bool) + lake_w, ocean_w = np.asarray(a["lake"]).astype(bool), np.asarray(a["ocean"]).astype(bool) + lid, lev_w, recv = np.asarray(a["lake_id"]), np.asarray(a["lake_level_m"], np.float64), np.asarray(a["recv"]) + inner_lake = ~halo & lake_w[parent] + ids = np.unique(lid[parent[inner_lake]]) + ids = ids[ids >= 0] + if not ids.size: + return z + inside_w = np.zeros(len(lake_w), bool) + inside_w[parent[~halo]] = True + levels = {int(L): float(np.nanmax(lev_w[lid == L])) for L in ids} + zf = None + for L in sorted(levels, key=levels.get): + lvl = levels[L] + if not np.isfinite(lvl): + continue + own = lid == L + water = inner_lake & own[parent] & (z < lvl) + if not water.any(): + continue + if zf is None: + zf = priority_flood(g, z, sinks) + if zf[water].min() <= lvl + OUTFLOW_TOL_M: + continue + exits = np.flatnonzero(own & ~own[recv]) # the world path: lake exit → sea or lower + path = set(exits.tolist()) + for c in exits: + c = int(recv[c]) + while c not in path and inside_w[c]: + path.add(c) + if ocean_w[c] or (lake_w[c] and lid[c] != L) or recv[c] == c: + break + c = int(recv[c]) + corridor = np.isin(parent, list(path)) & (~halo | sinks) + seeds = np.flatnonzero(corridor & (zf <= lvl - OUTFLOW_DROP_M) & ~(own[parent] & (z < lvl))) + goal = corridor & own[parent] & (z < lvl) + if not seeds.size or not goal.any(): + continue + cost = np.full(g.n, np.inf) + prev = np.full(g.n, -1, np.int64) + cost[seeds] = 0.0 + heap = [(0.0, int(s)) for s in seeds] + heapq.heapify(heap) + hit = -1 + while heap: + c0, i = heapq.heappop(heap) + if c0 > cost[i]: + continue + if goal[i]: + hit = i + break + for j in g.nbr_idx[g.nbr_ptr[i]:g.nbr_ptr[i + 1]]: + if not corridor[j]: + continue + cj = c0 + max(z[j] - lvl, 0.0) + 1e-3 + if cj < cost[j]: + cost[j], prev[j] = cj, i + heapq.heappush(heap, (cj, int(j))) + if hit < 0: + continue + chain = [hit] # lake → … → a cell that drains lower + while prev[chain[-1]] >= 0: + chain.append(int(prev[chain[-1]])) + chain = np.array(chain[1:], dtype=np.int64) + if not chain.size: + continue + top = lvl - OUTFLOW_DROP_M + end = min(max(float(zf[chain[-1]]), 0.0), top) # the sea's surface, not its bed + z[chain] = np.minimum(z[chain], top + (end - top) * np.arange(1, len(chain) + 1) / len(chain)) + zf = None + return z + + +def _lake_water(g, z, parent, world, halo) -> tuple[np.ndarray, np.ndarray]: + """(water, surface): cells below a world lake's surface connected to its cells (as _sea for the world ocean), and + that lake's surface (m; NaN elsewhere). The world's sea is never lake water.""" + a = world.a + sea = _sea(g, z, np.asarray(a["ocean"]).astype(bool)[parent], halo) + lake = np.asarray(a["lake"]).astype(bool)[parent] + lev = np.asarray(a["lake_level_m"] if "lake_level_m" in a else a["z_filled_m"], dtype=np.float64)[parent] + water, surf = np.zeros(g.n, bool), np.full(g.n, np.nan) + for L in np.unique(lev[lake]): + wet = (np.asarray(z) < L) & ~sea + lab = components(g, wet) + own = np.unique(lab[wet & lake & (lev == L)]) + m = wet & np.isin(lab, own[own >= 0]) & ~water + water |= m + surf[m] = L + return water, surf + + +ZONE_KEYS = ("gravity_g", "o2_fraction", "fire_reactivity", "plant_height_x") + + +def zone_fields(world, parent, z_surface, sea_p, scale_height_m) -> dict: + """Air, gravity and fire of the fine cells: their world cell's values (config zones and era zone events alike); + the pressure falls from the world cell's sea-level pressure with the fine cell's own height.""" + out = {k: np.asarray(world.a[k])[parent] for k in ZONE_KEYS} + out["pressure_bar"] = sea_p * FL.pressure(1.0, z_surface, scale_height_m) + out["po2_bar"] = np.asarray(out["o2_fraction"], dtype=np.float64) * out["pressure_bar"] + return out + + +def refine_area(g, halo, world, src, cfg, progress=None, plateaus=None) -> dict: + tell = progress or (lambda stage, f: None) + tell("fields", 0.05) + f = initial_fields(g, halo, world, src) + p = f["parent"] + bedrock = np.asarray(world.a["elevation_eroded_m"], dtype=np.float64) # ice comes back on top later + mc = None if servecache.low_memory() else {} # the correction's kernel, built once + z0 = f["elevation_m"] + mean_correction(g.xyz, f["elevation_m"], p, halo, bedrock, world, memo=mc) + inflow, exits, exit_lev = crossings(g, halo, world) + lw, surf = _lake_water(g, z0, p, world, halo) # halo under a world lake stands at its surface + on_lake = halo & lw # (water fills it, not drains via its bed) and + z0 = np.where(on_lake, surf, z0) # is a way out (into the world lake) + np.minimum.at(z0, exits, exit_lev) # world exits are the lowest ways out + sea = _sea(g, z0, f["ocean_world"], halo) + water = runoff_km3(g, f) + inflow + kmult = k_mult(f["age_class"]) + out_cells = outlets(g, halo, world, p, z0) + out_cells[exits] = True + out_cells |= on_lake + z0 = world_outflows(g, z0, halo, sea | out_cells, p, world.a) # no fine sill dams a world lake (MODEL 9) + z = erode_area(g, z0, halo, sea, water, kmult, outlet=out_cells, progress=tell) + z = z + mean_correction(g.xyz, z, p, halo, bedrock, world, hold=sea, memo=mc) # means back, smoothly + del mc + ocean = _sea(g, z, f["ocean_world"], halo) # land that sank to the sea joins it + tell("sea floor", 0.62) # canyons, fans, ponds, volcanoes, vents + z, sea_floor, vents = SF.run(g, z, ocean, halo, p, world.a, plateaus or [], int(cfg["build"]["seed"])) + out_cells |= outlets(g, halo, world, p, z) # the raised area may now spill over more halo + z = world_outflows(g, z, halo, ocean | out_cells, p, world.a) # again on the final heights (means put back) + # temperatures follow the final heights (lapse rate from the world's bedrock height, as the world climate) + dz_km = (z - bedrock[p]) / 1000.0 + for k in LAPSE: + f[k] = np.asarray(world.a[k])[p] - LAPSE_C_PER_KM * dz_km + f["biotemp"] = np.clip(f["biotemp"], 0.0, 30.0) + tell("rivers", 0.65) + # the world hydrology code, on the sub-grid: outlets = the ways out; inflow carried as extra runoff + pet = np.maximum(f["PET"], 1e-6) + run_mm = np.where(ocean, 0.0, np.maximum(f["P_ann"] - f["P_ann"] / np.sqrt(1 + (f["P_ann"] / pet) ** 2), 0.0)) + extra_mm = inflow / np.maximum(g.area_km2 * 1e-6, 1e-12) + p_hy = np.where(extra_mm > 0, _p_for_runoff(run_mm + extra_mm, pet), f["P_ann"]) + p_hy = np.where(halo, 0.0, p_hy) # no rain from outside the area + hcfg = copy.deepcopy(cfg) + hcfg.setdefault("hydrology", {})["river_min_km3_yr"] = RIVER_MIN_KM3_YR + hcfg["hydrology"]["min_depth_m"] = LAKE_MIN_DEPTH_M + ctx = SimpleNamespace(grid=g, cfg=hcfg, seed=int(cfg["build"]["seed"]), data={ + "elevation_eroded_m": np.where(halo & ~out_cells, WALL_M, z), "P_ann": p_hy, "PET": pet, "ocean": ocean | out_cells, + "lake_carve": np.zeros(g.n, bool)}) # lake beds come carved from the world + ctx.need = lambda *keys: [ctx.data[k] for k in keys] + hy = HY.run(ctx) + hy["runoff_mm"] = run_mm + river = hy["river"] & ~halo + tell("ground", 0.85) + # environment: life zones and ground from the refined fields (lithology, landform from the world cells) + P = params(cfg, "environment", ENV_DEFAULTS) + H = float(cfg["planet"]["scale_height_km"]) * 1000.0 + sea_p = (np.asarray(world.a["pressure_bar"], dtype=np.float64)[p] + / FL.pressure(1.0, np.asarray(world.a["z_surface_m"])[p], H)) # the world cell's sea-level air + wet = np.sqrt(sea_p / float(cfg["planet"]["sea_level_pressure_bar"])) # dense air: rain acts wetter (eras) + zone, _ = holdridge(f["biotemp"], f["P_ann"] * wet, f["T_min"]) + gr = ground(g, z, f["lithology"], f["T_mean"], f["T_min"], f["P_ann"], f["dist_ocean_km"], river, hy["strahler"], + hy["discharge_km3_yr"], hy["lake"] & ~halo, hy["salt_flat"] & ~halo, P, ocean) + ctx.data.update({"T_mean": f["T_mean"], "T_jun": f["T_jun"], "T_dec": f["T_dec"], "P_ann": f["P_ann"], "ocean": ocean, + "elevation_eroded_m": z}) # real heights again (no routing walls) + ic = IC.run(ctx) + fl = zone_fields(world, p, ic["z_surface_m"], sea_p, H) + exit_level = np.full(g.n, np.nan) + np.fmin.at(exit_level, exits, exit_lev) + out = {"g_ids": g.ids, "g_lat": g.lat, "g_lon": g.lon, "g_xyz": g.xyz, "g_area_km2": g.area_km2, + "parent": p, "halo": halo, "elevation_eroded_m": z, "ocean": ocean, "holdridge": zone, "ground": gr, + **{k: v for k, v in hy.items() if k not in ("river", "elevation_eroded_m", "lake_cut_m")}, "river": river, **ic, **fl, + **{k: f[k] for k in (*COPY, *LAPSE, "T_range")}} + out["lake"] = hy["lake"] & ~halo + out["inflow_km3_yr"] = inflow + out["outlet"] = out_cells + out["exit_level_m"] = exit_level + out["z_filled_m"] = np.where(halo & ~out_cells, z, out["z_filled_m"]) # no routing walls in the saved result + out.update(sea_floor) + out["vents"] = vents # per vent (build_areas: vents.npz) + return out + + +def stats(a) -> dict: + """Build report: size, rivers, largest lakes and the water balance (sources vs outflow + lake evaporation).""" + inner = ~a["halo"] + q, recv = a["discharge_km3_yr"], a["recv"] + sink = a["ocean"] | a["outlet"] + out_q = q[~sink & sink[recv]].sum() # flow into the sea or out of the area + ar = np.arange(len(q)) + terminal = q[~sink & (recv == ar)].sum() # endorheic basins keep (evaporate) theirs + lost = a.get("lake_loss_km3_yr", np.zeros(len(q)))[~sink].sum() + src = (a["runoff_mm"][inner & ~a["ocean"]] * a["g_area_km2"][inner & ~a["ocean"]] * 1e-6).sum() + a["inflow_km3_yr"].sum() + lakes = np.bincount(np.maximum(a["lake_id"][a["lake"]], 0), weights=a["g_area_km2"][a["lake"]]) if a["lake"].any() else np.zeros(1) + # border: where world rivers leave, the refined water level just inside vs the world river's level there + lev = a.get("exit_level_m", np.full(len(q), np.nan)) + errs = [float(a["z_filled_m"][(recv == e) & inner].min() - lev[e]) for e in np.where(np.isfinite(lev))[0] + if ((recv == e) & inner).any()] + return {"cells": int(inner.sum()), "rivers": int(a["river"].sum()), + "largest_lakes_km2": sorted(lakes.tolist(), reverse=True)[:3], + "water_balance_error": float(abs(out_q + terminal + lost - src) / max(src, 1e-9)), + "exits": int(np.isfinite(lev).sum()), "exits_reached": len(errs), + "exit_level_error_m": max(errs, key=abs) if errs else None, + "canyon_max_m": float(a["canyon_m"][inner].max()) if "canyon_m" in a and inner.any() else 0.0, + "fan_max_m": float(a["fan_m"][inner].max()) if "fan_m" in a and inner.any() else 0.0} + + +def results_dir(regions_root: Path, key: str) -> Path: + return Path(regions_root) / f"{key}-m{MODEL}" + + +def built_areas(regions_root: Path, key: str) -> list[Path]: + d = results_dir(regions_root, key) + return sorted(p for p in d.glob("*") if (p / "cells.npz").exists()) if d.exists() else [] + + +def regions_fingerprint(regions_root: Path, key: str, dirs: list | None = None) -> str: + """Changes whenever any built result changes (model, cells, rebuild): tile URLs must follow. dirs: a listing + already made (built_areas).""" + h = hashlib.sha1(f"m{MODEL}".encode()) + for p in (built_areas(regions_root, key) if dirs is None else dirs): + st = (p / "cells.npz").stat() + h.update(f"{p.name}:{st.st_size}:{st.st_mtime_ns}".encode()) + return h.hexdigest()[:6] + + +def _save(d: Path, arrays: dict, meta: dict, vents: dict | None = None) -> None: + d.mkdir(parents=True, exist_ok=True) + (d / "cells.npz").unlink(missing_ok=True) # cells.npz last: its presence means "complete" + fd, tmp = tempfile.mkstemp(dir=d, suffix=".json") + with os.fdopen(fd, "w") as fh: + fh.write(json.dumps(meta, indent=1) + "\n") + os.replace(tmp, d / "meta.json") + if vents is not None: + fd, tmp = tempfile.mkstemp(dir=d, suffix=".npz") + with os.fdopen(fd, "wb") as fh: + np.savez_compressed(fh, **vents) + os.replace(tmp, d / "vents.npz") + fd, tmp = tempfile.mkstemp(dir=d, suffix=".npz") + with os.fdopen(fd, "wb") as fh: + np.savez_compressed(fh, **arrays) + os.replace(tmp, d / "cells.npz") + + +def _meta(d: Path): + try: + return json.loads((d / "meta.json").read_text()) + except (OSError, ValueError): + return None + + +def build_areas(root: Path, res: int, regions_path: Path, log=print, regions_root: Path | None = None, + progress=None) -> list[dict]: + """Refine every connected area of regions_path's outlines for every world the viewer can show (the base build + and each built era). Results are keyed by content (area_key) under <regions_root>/areas/<key>/, so an area an era + left as it was is built once; each world's set folder <regions_root>/<world fingerprint>-m<MODEL>/ links its areas + by cell hash (what RegionSet reads) and records its outlines (outlines.json) and era (era.json). Other worlds' + sets and unused areas go only after every area built (a failed build keeps the map as it was). + progress(stage, fraction) reports a running build (fraction over all areas still to build).""" + import gc + + import serve + import tiles + from mapgen import config as C + root = Path(root) + cfg, tect = C.load(root) + plateaus = tect.get("plateau", []) + regions_root = Path(regions_root or (root / "out" / f"r{res}" / "regions")) + store = areas_dir(regions_root) + store.mkdir(parents=True, exist_ok=True) + regions = load_regions(regions_path) # once: the records must describe what was built + areas = area_cells(regions, res + FINER) + try: + memo = json.loads((store / "keys.json").read_text()) + except (OSError, ValueError): + memo = {} + if progress: + progress("keys", 0.0) + sets = [] + for era, out_dir in world_dirs(root, res, log): + world = serve.World(root, res, out=out_dir) + src = tiles.TileSource(world, int(cfg["build"]["seed"]), regions_dir=regions_root, with_regions=False) + wc = WorldCells(out_dir) + keys = {} + for cells in areas: + m = f"{src.fingerprint}:m{MODEL}:{area_hash(cells)}" # a new model: new keys + if m not in memo: + memo[m] = area_key(wc, cells, cfg, plateaus) + keys[area_hash(cells)] = memo[m] + sets.append((era, world, src, wc, keys)) + todo, queued = [], set() + for era, world, src, wc, keys in sets: + for cells in areas: + k = keys[area_hash(cells)] + done = (store / k / "cells.npz").exists() and _meta(store / k) is not None + if not done and k not in queued: + queued.add(k) + todo.append((cells, k, era, world, src, wc)) + for n, (cells, k, era, world, src, wc) in enumerate(todo): + tell = (lambda stage, f, n=n: progress(stage, (n + f) / len(todo))) if progress else None + R = float(world.cells_meta["radius_km"]) + km2 = len(cells) * float(np.mean([h3.cell_area(int(c), unit="rads^2") for c in cells[:200]])) * R ** 2 + if km2 > 2.0e6: + log(f"refine: area {area_hash(cells)} is {km2:,.0f} km² (over ≈ 2,000,000): this will take a while") + t0 = time.time() + g, halo = subgrid(cells, R) + a = refine_area(g, halo, wc, src, cfg, progress=tell, plateaus=plateaus) + vents = a.pop("vents", None) + meta = {"hash": area_hash(cells), "key": k, "model": MODEL, "era": era, "world": src.fingerprint, + "area_km2": km2, "vents": 0 if vents is None else int(len(vents["cell"])), "seconds": round(time.time() - t0, 1), **stats(a)} + if tell: + tell("saving", 0.95) + _save(store / k, a, meta, vents) + log(f"refine: area {meta['hash']} ({era}): {meta['cells']} cells, {meta['rivers']} river cells, " + f"{meta['seconds']} s") + del a, g, halo # the next area must not share memory with this one + gc.collect() + trim_memory() + metas, built_now, keep_sets = [], {t[1] for t in todo}, set() + for era, world, src, wc, keys in sets: # the records: links, outlines, era + sd = results_dir(regions_root, src.fingerprint) + sd.mkdir(parents=True, exist_ok=True) + keep_sets.add(sd.name) + for cells in areas: + h, k = area_hash(cells), keys[area_hash(cells)] + link = _link(sd, h, store / k) + metas.append({**_meta(store / k), "dir": str(link), "era": era, "built": k in built_now}) + _write_json(sd / "outlines.json", {"regions": region_areas(regions, areas, res + FINER)}) + _write_json(sd / "era.json", {"era": era}) + fps = {s[2].fingerprint for s in sets} + _write_json(store / "keys.json", {m: v for m, v in memo.items() + if m.split(":", 1)[0] in fps and m.split(":")[1] == f"m{MODEL}"}) + # only after every area built: other worlds' sets, links of redrawn or deleted regions, unused areas go + names = {area_hash(c) for c in areas} + for d in regions_root.iterdir(): + if d.name != AREAS and d.name not in keep_sets and (d.is_dir() or d.is_symlink()): + _remove(d) + used = set() + for s in keep_sets: + for p in (regions_root / s).iterdir(): + if p.name.startswith(servecache.PREFIX) or p.name.startswith(".") or not (p.is_dir() or p.is_symlink()): + continue # serve caches: the server prunes them + if p.name in names and p.is_symlink(): + used.add(os.path.basename(os.readlink(p))) + else: + _remove(p) + for p in store.iterdir(): + if p.is_dir() and p.name not in used: + _remove(p) + return metas + + +NET_KEYS = ("g_xyz", "g_ids", "river", "recv", "ocean", "lake", "z_surface_m", "z_filled_m", "discharge_km3_yr", + "lithology", "age_class", "landform", "ground", "P_ann") # what rivers.RiverNet reads +BLEND_KM = 10.0 +LAKE_FLOOR_M = 30.0 # lake water reaches this far below a lake's deepest cell (procedural detail) + + +def shoulders(a, radius_km: float = 12742.0) -> np.ndarray: + """Heights with each river cell raised to its valley shoulders (the mean of its dry neighbours), so smooth + interpolation keeps the plateau and the valley shape can be cut in at sub-cell scale.""" + z = np.asarray(a["z_surface_m"], dtype=np.float64).copy() + riv = np.asarray(a["river"]).astype(bool) & ~np.asarray(a["halo"]).astype(bool) + dry = (~np.asarray(a["river"]).astype(bool) & ~np.asarray(a["lake"]).astype(bool) & ~np.asarray(a["ocean"]).astype(bool) + & ~np.asarray(a["halo"]).astype(bool)) + if not riv.any() or not dry.any(): + return z + tree = cKDTree(a["g_xyz"][dry]) + spacing = np.sqrt(np.mean(a["g_area_km2"])) * 1.07 / radius_km # neighbour distance on the unit sphere + d, k = tree.query(a["g_xyz"][riv], k=min(6, int(dry.sum())), distance_upper_bound=1.5 * spacing) + d, k = np.atleast_2d(d), np.atleast_2d(k) + ok = np.isfinite(d) + zd = np.where(ok, z[dry][np.minimum(k, dry.sum() - 1)], 0.0) + mean = np.where(ok.any(1), zd.sum(1) / np.maximum(ok.sum(1), 1), -np.inf) + z[riv] = np.maximum(z[riv], mean) + return z + + +def touches_trees(trees, spacing_km: float, radius_km: float, z, x, y, margin_km: float = 0.0) -> bool: + """Whether tile z/x/y comes within (2 cell spacings + margin) of the cells in any of the KD-trees.""" + if not trees: + return False + span = 180.0 / 2 ** z + f = np.array([0.0, 0.5, 1.0]) + LA, LO = np.meshgrid(90.0 - (y + f) * span, -180.0 + (x + f) * span, indexing="ij") + p = latlon_to_xyz(LA, LO).reshape(-1, 3) + r = float(np.max(np.linalg.norm(p - p[4], axis=1))) + (2 * spacing_km + margin_km) / radius_km + return any(t.query_ball_point(p[4], r, return_length=True) > 0 for t in trees) + + +def trim_memory() -> None: + """Give freed memory back to the system (glibc keeps it otherwise).""" + try: + import ctypes + ctypes.CDLL("libc.so.6").malloc_trim(0) + except (OSError, AttributeError): + pass + + +def _compute_region_set(dirs, legends: dict, radius_km: float, spill=None): + """Everything a RegionSet serves, computed from its areas' results (as the old RegionSet.__init__ did): + (arrays, meta) for the serve cache. Keys: 'a.<field>' (the concatenated inner cells, sorted by H3 id, recv + remapped), 'halo_xyz', 'area' (index into meta['area_names'] per fine cell), 'z_env', 'lake_level', + 'lake_floor', 'net.<state>' (rivers.STATE) when a river net exists. spill: a folder to keep the arrays in + (memory maps; same values) instead of RAM.""" + R = float(radius_km) + keep = servecache.spiller(spill) + files = [np.load(p / "cells.npz") for p in dirs] # read key by key: each array once, no copies kept + try: + halo = [f["halo"].astype(bool) for f in files] + inner = [~h for h in halo] + ids = [f["g_ids"] for f in files] + recv_p = [f["recv"].astype(np.int64) for f in files] + river_p = [f["river"].astype(bool) for f in files] + xyz_p = [f["g_xyz"] for f in files] + halo_xyz = np.concatenate([x[h] for x, h in zip(xyz_p, halo)]) + order = np.argsort(np.concatenate([i_[m] for i_, m in zip(ids, inner)])) + area_of = np.concatenate([np.full(int(m.sum()), p, np.int32) for p, m in enumerate(inner)])[order] + # halo cells rivers leave through (for the river net): per part, their indices + ends = [] + for m, h, r, rv in zip(inner, halo, recv_p, river_p): + out = np.where(m & rv & h[r])[0] + ends.append((out, *np.unique(r[out], return_inverse=True)) if len(out) else None) + keys = set.intersection(*(set(f.files) for f in files)) + a, end_vals = {}, {} + for k in sorted(keys): + vs = [f[k] for f in files] + if any(len(v) != len(i_) for v, i_ in zip(vs, ids)): + continue + a[k] = keep(f"a.{k}", np.concatenate([v[m] for v, m in zip(vs, inner)])[order]) + if k in NET_KEYS: + end_vals[k] = [v[e[1]] if e is not None else v[:0] for v, e in zip(vs, ends)] + del vs + finally: + for f in files: + f.close() + # receivers index into each part's own arrays: remap to the sorted inner cells (halo/outside → self) + inv = np.empty_like(order) + inv[order] = np.arange(len(order)) + recv_all, off = [], 0 + for r, m in zip(recv_p, inner): + new = -np.ones(len(m), np.int64) + new[np.where(m)[0]] = np.arange(off, off + m.sum()) + rr = new[r[m]] + recv_all.append(np.where(rr >= 0, rr, np.arange(off, off + m.sum()))) + off += m.sum() + a["recv"] = keep("a.recv", inv[np.concatenate(recv_all)][order]) + ids_all = a["g_ids"] + z_env = shoulders(a, R) # plateau across river cells + lake = np.asarray(a["lake"]).astype(bool) + lake_level = np.where(lake, a["z_filled_m"], np.nan) + # each lake cell's bed (+ LAKE_FLOOR_M for the fine detail): ground far below it lies past a dam or a cliff, + # not under water (no water walls hanging over a drop) + lake_floor = np.where(lake, lake_level - np.asarray(a["depression_depth_m"], dtype=np.float64) - LAKE_FLOOR_M, + np.nan) + # the river net runs on into the halo cells where rivers leave, so they meet the world's rivers at the edge + # (built from the few arrays it needs: a copy of every field would double the memory) + n = len(ids_all) + net_recv = a["recv"].copy() + extra, base = [], n + for part, e in enumerate(ends): + if e is None: + continue + out, h, k = e + net_recv[np.searchsorted(ids_all, ids[part][out])] = base + k + extra.append(part) + base += len(h) + nkeys = [k for k in NET_KEYS if k in a] + net_a = {k: np.concatenate([a[k], *[end_vals[k][p] for p in extra]]) for k in nkeys if k != "recv"} + net_a["recv"] = np.concatenate([net_recv, np.arange(n, base)]) + net_a["river"] = np.concatenate([a["river"].astype(bool), np.zeros(base - n, bool)]) + net_a["lake"] = np.concatenate([a["lake"].astype(bool), np.zeros(base - n, bool)]) + ground = np.concatenate([z_env, net_a["z_surface_m"][n:]]) + arrays = {f"a.{k}": v for k, v in a.items()} + arrays.update(halo_xyz=keep("halo_xyz", halo_xyz), area=keep("area", area_of), z_env=keep("z_env", z_env), + lake_level=keep("lake_level", lake_level), lake_floor=keep("lake_floor", lake_floor)) + meta = {"area_names": [p.name for p in dirs], "res": h3.get_resolution(int(ids_all[0])), + "spacing_km": float(np.sqrt(np.mean(a["g_area_km2"])) * 1.07), "net": None} + try: + import rivers as RV + net = RV.RiverNet(net_a, legends, R, levels=net_a["z_surface_m"], ground=ground) + st, meta["net"] = net.state() + arrays.update({f"net.{k}": v for k, v in st.items()}) + except Exception: + meta["net"] = None + return arrays, meta + + +CODE_FILES = (Path(__file__), Path(__file__).with_name("rivers.py")) # what computes a region set's served arrays + + +def code_key() -> str: + """Hash of the code that computes a region set's arrays: new code, new serve cache (stored river nets, shoulders + and lake floors are values of this code, not only of the areas' results).""" + h = hashlib.sha1() + for p in CODE_FILES: + h.update(Path(p).read_bytes()) + return h.hexdigest()[:8] + + +class _AreaTrees: + """Per-area KD-trees of a RegionSet, built on first use (one area: the set's main tree).""" + def __init__(self, rs): + self.rs, self._built = rs, {} + + def __contains__(self, name) -> bool: + return name in self.rs.area_names + + def __getitem__(self, name): + if name not in self._built: + rs = self.rs + if len(rs.area_names) == 1: + self._built[name] = rs.tree + else: + k = rs.area_names.index(name) + self._built[name] = cKDTree(np.asarray(rs.arrays["g_xyz"], dtype=np.float64)[np.asarray(rs.area) == k]) + return self._built[name] + + +class RegionSet: + """All refined areas built for the current world: fine cells for tiles, 3D heights and the inspector. The arrays + come from the serve cache (memory-mapped; computed and written on first load); search trees are built on first + use (warm() builds them before render workers fork, so they share them).""" + @classmethod + def none(cls) -> "RegionSet": + """No refined areas (e.g. for a build, which needs only the world's relief).""" + rs = cls.__new__(cls) + rs.fingerprint, rs.area_names, rs.empty, rs.R, rs.cache = "", [], True, 0.0, None + rs.area_keys = {} + rs._tree = rs._halo_tree = None + rs._lock = threading.Lock() + rs.area_trees = _AreaTrees(rs) + return rs + + def __init__(self, regions_root: Path, key: str, legends: dict, radius_km: float, use_cache: bool = True): + self.R = float(radius_km) + dirs = built_areas(regions_root, key) + self.fingerprint = regions_fingerprint(regions_root, key, dirs) + self.area_names = [p.name for p in dirs] # area hashes + self.area_keys = {p.name: (Path(os.readlink(p)).name if p.is_symlink() else p.name) for p in dirs} + self.empty = not dirs + self.cache = None + self._tree = self._halo_tree = None + self._lock = threading.Lock() + self.area_trees = _AreaTrees(self) + if self.empty: + return + spill = None + if servecache.low_memory() and use_cache: # fields on disk while the cache is written + spill = Path(tempfile.mkdtemp(prefix=".spill-", dir=results_dir(regions_root, key))) + build = lambda: _compute_region_set(dirs, legends, self.R, spill=spill) + try: + if use_cache: + self.cache = servecache.cache_dir(results_dir(regions_root, key), f"{self.fingerprint}-{code_key()}") + arrays, meta = servecache.load_or_build(self.cache, build) + else: + arrays, meta = build() + self._from_cache(arrays, meta) + finally: + if spill is not None: + shutil.rmtree(spill, ignore_errors=True) + trim_memory() # the loading's scratch memory + + def _from_cache(self, arrays: dict, meta: dict) -> None: + import rivers as RV + self.arrays = {k[2:]: v for k, v in arrays.items() if k.startswith("a.")} + self.ids = self.arrays["g_ids"] + self.res = int(meta["res"]) + self.spacing_km = float(meta["spacing_km"]) + self.area = arrays["area"] + self._halo_xyz = arrays["halo_xyz"] + self.z_env, self.lake_level, self.lake_floor = arrays["z_env"], arrays["lake_level"], arrays["lake_floor"] + st = {k[4:]: v for k, v in arrays.items() if k.startswith("net.")} + self.net = RV.RiverNet.from_state(st, meta["net"]) if meta.get("net") else None + + @property + def tree(self): + if self._tree is None and not self.empty: + with self._lock: + if self._tree is None: + self._tree = cKDTree(np.asarray(self.arrays["g_xyz"], dtype=np.float64)) + return self._tree + + @property + def halo_tree(self): + if self._halo_tree is None and not self.empty: + with self._lock: + if self._halo_tree is None: + self._halo_tree = cKDTree(np.asarray(self._halo_xyz, dtype=np.float64)) + return self._halo_tree + + def warm(self) -> None: + """Build the search trees now (before forking render workers, so they share them).""" + if self.empty: + return + self.tree + self.halo_tree + if self.net is not None: + self.net.tree + + def nearest(self, xyz, k=1): + d, i = self.tree.query(xyz, k=k) + return i, d * self.R + + def edge_km(self, xyz): + return self.halo_tree.query(xyz)[0] * self.R + + def weight(self, xyz): + """0 outside an area, rising to 1 over BLEND_KM inside its edge.""" + if self.empty: + return np.zeros(len(xyz)) + d_in = self.tree.query(xyz, distance_upper_bound=1.5 * self.spacing_km / self.R)[0] * self.R # inf: far outside + e = self.halo_tree.query(xyz, distance_upper_bound=(2 * BLEND_KM + 2 * self.spacing_km) / self.R)[0] * self.R + depth = np.full(len(d_in), -np.inf) + m = np.isfinite(d_in) + depth[m] = (e[m] - d_in[m]) / 2 # ≈ distance inside the area edge; 0 on it + t = np.clip(depth / BLEND_KM, 0.0, 1.0) # (bounded queries: deep inside → inf → 1) + return t * t * (3 - 2 * t) + + def touches(self, z, x, y, areas=None, margin_km: float = 0.0) -> bool: + """Whether tile z/x/y (with a margin for borders) can show any refined area (or one of `areas`).""" + if self.empty or (areas is not None and not areas): + return False + trees = [self.tree] if areas is None else [self.area_trees[h] for h in areas if h in self.area_trees] + return touches_trees(trees, self.spacing_km, self.R, z, x, y, margin_km) + + def index_of(self, lat, lon): + if self.empty: + return None + c = np.uint64(h3.latlng_to_cell(float(lat), float(lon), self.res)) + i = int(np.searchsorted(self.ids, c)) + return i if i < len(self.ids) and self.ids[i] == c else None |