worldgen: initial public history

1 files changed, 121 insertions, 0 deletions

diff --git a/mapgen/grid.py b/mapgen/grid.py
new file mode 100644
index 0000000..8138b5f
--- /dev/null
+++ b/mapgen/grid.py
@@ -0,0 +1,121 @@
+"""H3 hexagonal grid as a CSR cell graph. Stage `grid`."""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from functools import cached_property
+
+import h3.api.basic_int as h3
+import numpy as np
+
+from .sphere import latlon_to_xyz, tangent_dir
+
+
+EDGE_BLOCK = 1 << 20
+
+
+def edge_blocks(n: int, block: int | None = None):
+    """Slices over n edges, EDGE_BLOCK at a time: per-edge (row-wise) formulas give the same values block by block,
+    with (block, 3) float64 temporaries instead of (edges, 3) ones (290 MB each at r5, 2 GB at r6)."""
+    block = block or EDGE_BLOCK
+    return [slice(a, min(a + block, n)) for a in range(0, n, block)]
+
+
+def rowdot_at(A, idx, t) -> np.ndarray:
+    """np.sum(A[idx] * t, axis=1) — per edge, A's row at an end dotted with the edge's vector — in edge blocks."""
+    out = np.empty(len(idx))
+    for s in edge_blocks(len(idx)):
+        out[s] = np.sum(A[idx[s]] * t[s], axis=1)
+    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). Same ids as h3, without a Python call per cell."""
+    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)
+
+
+@dataclass
+class Grid:
+    res: int
+    radius_km: float
+    ids: np.ndarray
+    lat: np.ndarray
+    lon: np.ndarray
+    xyz: np.ndarray
+    area_km2: np.ndarray
+    nbr_ptr: np.ndarray
+    nbr_idx: np.ndarray
+
+    @property
+    def n(self) -> int:
+        return len(self.ids)
+
+    @cached_property
+    def counts(self):
+        return np.diff(self.nbr_ptr)
+
+    @cached_property
+    def src(self):
+        return np.repeat(np.arange(self.n), self.counts)
+
+    @property
+    def dst(self):
+        return self.nbr_idx
+
+    @cached_property
+    def edge_km(self):
+        out = np.empty(len(self.dst))
+        for s in edge_blocks(len(self.dst)):
+            d = np.sum(self.xyz[self.src[s]] * self.xyz[self.dst[s]], axis=1)
+            out[s] = self.radius_km * np.arccos(np.clip(d, -1.0, 1.0))
+        return out
+
+    @cached_property
+    def edge_tangents(self):
+        out = np.empty((len(self.dst), 3))
+        for s in edge_blocks(len(self.dst)):
+            out[s] = tangent_dir(self.xyz[self.src[s]], self.xyz[self.dst[s]])
+        return out
+
+    @cached_property
+    def spacing_km(self) -> float:
+        return float(self.edge_km.mean())
+
+    def cell_index(self, lat: float, lon: float) -> int:
+        c = np.uint64(h3.latlng_to_cell(float(lat), float(lon), self.res))
+        return int(np.searchsorted(self.ids, c))
+
+    def to_arrays(self) -> dict:
+        return {"g_ids": self.ids, "g_lat": self.lat, "g_lon": self.lon, "g_xyz": self.xyz,
+                "g_area_km2": self.area_km2, "g_nbr_ptr": self.nbr_ptr, "g_nbr_idx": self.nbr_idx}
+
+    @classmethod
+    def from_arrays(cls, d: dict, res: int, radius_km: float) -> "Grid":
+        return cls(res, radius_km, d["g_ids"], d["g_lat"], d["g_lon"], d["g_xyz"], d["g_area_km2"],
+                   d["g_nbr_ptr"], d["g_nbr_idx"])
+
+
+def build_grid(res: int, radius_km: float) -> Grid:
+    cells: list[int] = []
+    for r0 in h3.get_res0_cells():
+        cells.extend(h3.cell_to_children(r0, res))
+    ids = np.array(sorted(cells), dtype=np.uint64)
+    ll = np.array([h3.cell_to_latlng(int(c)) for c in ids], dtype=np.float64)
+    area = np.array([h3.cell_area(int(c), unit="rads^2") for c in ids]) * radius_km**2
+    rings = [h3.grid_ring(int(c), 1) for c in ids]
+    counts = np.array([len(r) for r in rings], dtype=np.int64)
+    ptr = np.concatenate([[0], np.cumsum(counts)]).astype(np.int64)
+    flat = np.array([c for r in rings for c in r], dtype=np.uint64)
+    idx = np.searchsorted(ids, flat).astype(np.int64)
+    return Grid(res, radius_km, ids, ll[:, 0], ll[:, 1], latlon_to_xyz(ll[:, 0], ll[:, 1]), area, ptr, idx)
+
+
+def run(ctx) -> dict:
+    return build_grid(ctx.res, float(ctx.cfg["planet"]["radius_km"])).to_arrays()