"""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" # /areas//: 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/eras//).""" 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 /areas//, so an area an era left as it was is built once; each world's set folder /-m/ 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.' (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.' (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