diff options
Diffstat (limited to 'mapgen/render.py')
| -rw-r--r-- | mapgen/render.py | 526 |
1 files changed, 526 insertions, 0 deletions
diff --git a/mapgen/render.py b/mapgen/render.py new file mode 100644 index 0000000..2a408e6 --- /dev/null +++ b/mapgen/render.py @@ -0,0 +1,526 @@ +"""Stage `render`: equirectangular rasters, cells.npz, metadata, previews, contact sheet.""" +from __future__ import annotations + +import colorsys +import json + +import numpy as np +from PIL import Image, ImageDraw +from scipy.spatial import cKDTree + +from . import geo, projections, viewer_export +from . import plateaus as PL +from .crust import AGE_NAMES +from .environment import LEGENDS, REGIONS, ZONES +from .ice import ICE_NAMES +from .noise import fbm +from .seabed import MINERAL_NAMES, SEABED_NAMES +from .sphere import east_north, latlon_to_xyz + +CONTINUOUS = { + "elevation": ("z_surface_m", 0.5, -12000.0, "m"), + "T_mean": ("T_mean", 0.01, -100.0, "degC"), + "T_range": ("T_range", 0.01, 0.0, "degC"), + "P_ann": ("P_ann", 0.5, 0.0, "mm/yr"), + "P_jun": ("P_jun", 0.5, 0.0, "mm/yr"), + "P_dec": ("P_dec", 0.5, 0.0, "mm/yr"), + "po2": ("po2_bar", 2e-4, 0.0, "bar"), + "gravity": ("gravity_g", 1e-4, 0.0, "g"), + "pressure": ("pressure_bar", 2e-4, 0.0, "bar"), + "o2_fraction": ("o2_fraction", 2e-5, 0.0, "fraction"), + "fire": ("fire_reactivity", 1e-4, 0.0, "x"), + "vent_potential": ("vent_potential", 1e-4, 0.0, "0-1"), + "bottom_temp": ("bottom_temp_c", 0.01, -10.0, "degC"), + "sediment": ("sediment_m", 0.5, 0.0, "m"), + "plant_height": ("plant_height_x", 2e-3, 0.0, "x"), + "sst": ("sst", 0.01, -100.0, "degC"), + "productivity": ("productivity", 1e-4, 0.0, "0-1"), + "current_speed": ("current_speed", 1e-4, 0.0, "m/s"), + "upwelling": ("upwelling", 0.05, -1600.0, "m/yr"), +} +CATEGORICAL = {"plates": "plate", "age_class": "age_class", "holdridge": "holdridge", + "seasonality": "seasonality", "landform": "landform", "lithology": "lithology", + "ground": "ground", "ice": "ice", + "seabed_type": "seabed_type", "seabed_mineral": "seabed_mineral", "deposits": "deposit_main"} +DETAIL_M = np.array([60, 40, 150, 400, 80, 150, 300, 300, 300, 60, 30, 120], dtype=np.float64) # by landform +CHUNK = 128 +COAST_DETAIL_M = 250.0 + + +def _by_rows(fn, v, dtype, tail=()): + """fn applied to CHUNK-row slices of v (fn elementwise): the same values as fn(v), without fn's full-size + float64 scratch (a raster ramp builds five H×W×3 float64 temporaries: ~4 GB at 8192 px).""" + v = np.asarray(v) + if v.ndim < 2 or v.shape[0] <= CHUNK: + return fn(v) + out = np.empty(v.shape + tuple(tail), dtype) + for y0 in range(0, v.shape[0], CHUNK): + out[y0:y0 + CHUNK] = fn(v[y0:y0 + CHUNK]) + return out + + +def encode(v, scale, offset): + return _by_rows(lambda x: np.clip(np.round((np.asarray(x, dtype=np.float64) - offset) / scale), 0, 65535) + .astype(np.uint16), v, np.uint16) + + +def decode(raw, scale, offset): + return np.asarray(raw, dtype=np.float64) * scale + offset + + +def _rows_xyz(rows, W, H): + lat = 90.0 - (rows + 0.5) / H * 180.0 + lon = (np.arange(W) + 0.5) / W * 360.0 - 180.0 + LA, LO = np.meshgrid(lat, lon, indexing="ij") + return latlon_to_xyz(LA.ravel(), LO.ravel()) + + +def _make_render_jit(): + """Compiled ramp and sampling loops (numba optional; WORLDGEN_NO_JIT=1 turns it off): per pixel the same float + steps in the same order as the numpy statements they replace, so the same bytes, without the temporaries.""" + import os + if os.environ.get("WORLDGEN_NO_JIT"): + return None + try: + import numba + except ImportError: + return None + + @numba.njit(cache=True, nogil=True) + def ramp(x, lo, span, stops, out): # x: float64 (n,), out: uint8 (n, c) + m = stops.shape[0] + for j in range(x.shape[0]): + t = (x[j] - lo) / span + t = 0.0 if t < 0.0 else (1.0 if t > 1.0 else t) + t = t * (m - 1) + i = min(np.int64(t), m - 2) + f = t - i + for c in range(stops.shape[1]): + out[j, c] = np.uint8(np.int64(stops[i, c] * (1 - f) + stops[i + 1, c] * f)) + + @numba.njit(cache=True, nogil=True) + def sample(v, idx, wts, out): # out[p] = Σ_k v[idx[p,k]] * w[p,k], k left to right (as np.sum, k < 8) + for p in range(idx.shape[0]): + s = v[idx[p, 0]] * np.float64(wts[p, 0]) + for k in range(1, idx.shape[1]): + s += v[idx[p, k]] * np.float64(wts[p, k]) + out[p] = s + return ramp, sample + + +_render_jit = _make_render_jit() + + +def pixel_neighbours(g, W, H, k=3): + from .graph import workers + tree = cKDTree(g.xyz) + idx = np.empty((H, W, k), np.int32) + wts = np.empty((H, W, k), np.float32) + for y0 in range(0, H, CHUNK): + rows = np.arange(y0, min(H, y0 + CHUNK)) + d, i = tree.query(_rows_xyz(rows, W, H), k=k, workers=workers()) # per point: the same answer + w = 1.0 / np.maximum(d, 1e-9) ** 2 + w /= w.sum(axis=1, keepdims=True) + idx[rows] = i.reshape(len(rows), W, k) + wts[rows] = w.reshape(len(rows), W, k) + return idx, wts + + +def sample_cont(v, idx, wts): + out = np.empty(idx.shape[:2]) + if (_render_jit is not None and v.dtype == np.float64 and v.ndim == 1 and wts.dtype == np.float32 + and 1 <= idx.shape[-1] < 8 and idx.shape == wts.shape): + k = idx.shape[-1] + _render_jit[1](np.ascontiguousarray(v), np.ascontiguousarray(idx).reshape(-1, k), + np.ascontiguousarray(wts).reshape(-1, k), out.reshape(-1)) + return out + for y0 in range(0, idx.shape[0], CHUNK): + s = slice(y0, y0 + CHUNK) + out[s] = np.sum(v[idx[s]] * wts[s], axis=-1) + return out + + +def sample_cat(v, idx): + return v[idx[..., 0]] + + +def pixel_land(ocean_k, z): + """Pixel land mask: unanimous neighbour cells decide; at the coast (mixed) the sub-cell height does.""" + all_sea = ocean_k.all(axis=-1) + all_land = ~ocean_k.any(axis=-1) + return all_land | (~all_sea & ~all_land & (z > 0)) + + +def hillshade(z, radius_km, az=315.0, alt=45.0, exag=4.0, low_memory=False): + """low_memory: the same values, CHUNK rows at a time (one halo row each side keeps the central differences).""" + if low_memory: + H = z.shape[0] + out = np.empty(z.shape) + for y0 in range(0, H, CHUNK): + a, b = max(y0 - 1, 0), min(y0 + CHUNK + 1, H) + part = _hillshade_rows(z[a:b], a, H, radius_km, az, alt, exag) + out[y0:min(y0 + CHUNK, H)] = part[y0 - a: y0 - a + min(CHUNK, H - y0)] + return out + return _hillshade_rows(z, 0, z.shape[0], radius_km, az, alt, exag) + + +def _hillshade_rows(z, row0, H, radius_km, az, alt, exag): + """Hillshade of rows row0.. of an H-row raster (z holds those rows; edges of z use one-sided differences).""" + W = z.shape[1] + lat = 90.0 - (np.arange(row0, row0 + z.shape[0]) + 0.5) / H * 180.0 + dy = np.pi * radius_km * 1000.0 / H + dx = 2 * np.pi * radius_km * 1000.0 * np.maximum(np.cos(np.radians(lat)), 0.01) / W + gy, gx = np.gradient(z) + dzdx = gx / dx[:, None] * exag + dzdn = -gy / dy * exag + norm = np.sqrt(dzdx**2 + dzdn**2 + 1.0) + a, b = np.radians(az), np.radians(alt) + L = (np.sin(a) * np.cos(b), np.cos(a) * np.cos(b), np.sin(b)) + return np.clip((-dzdx * L[0] - dzdn * L[1] + L[2]) / norm, 0.0, 1.0) + + +def _ramp(v, lo, hi, stops): + stops = np.asarray(stops, dtype=np.float64) + + def part(x): + if (_render_jit is not None and isinstance(x, np.ndarray) and x.dtype == np.float64 and stops.ndim == 2 + and len(stops) >= 2 and not np.isnan(x).any()): + out = np.empty(x.shape + (stops.shape[1],), np.uint8) + _render_jit[0](np.ascontiguousarray(x).reshape(-1), float(lo), float(hi - lo), stops, + out.reshape(-1, stops.shape[1])) + return out + t = np.clip((x - lo) / (hi - lo), 0, 1) * (len(stops) - 1) + i = np.minimum(t.astype(np.int64), len(stops) - 2) + f = (t - i)[..., None] + return (stops[i] * (1 - f) + stops[i + 1] * f).astype(np.uint8) + return _by_rows(part, v, np.uint8, (stops.shape[-1],)) + + +def holdridge_palette(): + ramp = np.array([[216, 200, 160], [208, 196, 140], [200, 200, 120], [152, 168, 96], [106, 150, 80], + [70, 125, 68], [50, 105, 62], [37, 90, 56]], dtype=np.float64) + tint = {"polar": ((232, 236, 239), 0.9), "subpolar": ((170, 176, 160), 0.55), "boreal": ((70, 100, 80), 0.35), + "tropical": ((20, 90, 40), 0.15)} + cols = [] + for r in REGIONS: + k = len(ZONES[r]) + for j in range(k): + c = ramp[int(round((j / max(k - 1, 1)) * (len(ramp) - 1)))] + if r in tint: + c = c * (1 - tint[r][1]) + np.array(tint[r][0]) * tint[r][1] + cols.append(c) + return np.array(cols, dtype=np.uint8) + + +def category_palette(n): + return np.array([[int(255 * c) for c in colorsys.hsv_to_rgb((i * 0.618034) % 1.0, 0.55, 0.9)] + for i in range(max(n, 1))], dtype=np.uint8) + + +def _region_palette(cols): + """38 zone colours from per-region (dry, wet) colour pairs.""" + out = [] + for r in REGIONS: + dry, wet = (np.array(c, dtype=np.float64) for c in cols[r]) + k = len(ZONES[r]) + out += [dry + (wet - dry) * (j / max(k - 1, 1)) for j in range(k)] + return np.array(out, dtype=np.uint8) + + +STYLES = { # alien palettes (config [render] style); None = Earth-like + "tidal-lock": { + "zones": {"polar": ((34, 38, 46), (34, 38, 46)), "subpolar": ((70, 72, 78), (96, 104, 112)), + "boreal": ((120, 96, 64), (150, 110, 60)), "cool temperate": ((170, 120, 60), (190, 140, 70)), + "warm temperate": ((110, 70, 44), (130, 86, 50)), "subtropical": ((46, 36, 34), (60, 44, 38)), + "tropical": ((22, 20, 22), (30, 26, 28))}, + "ocean": [[4, 12, 16], [10, 34, 40], [40, 80, 84]], "lake": (60, 96, 104), "river": (70, 110, 118), + "ground": {6: (92, 86, 80)}, "sheet": (200, 222, 240), "sea_ice": (170, 200, 226), "sea_ice_seasonal": (120, 150, 170)}, + "salt-mirror": { + "zones": {"polar": ((236, 240, 244), (236, 240, 244)), "subpolar": ((228, 228, 234), (214, 214, 228)), + "boreal": ((238, 236, 228), (208, 208, 224)), "cool temperate": ((242, 238, 226), (200, 204, 222)), + "warm temperate": ((240, 230, 204), (196, 204, 220)), "subtropical": ((238, 224, 186), (192, 206, 218)), + "tropical": ((234, 214, 166), (186, 206, 216))}, + "ocean": [[70, 120, 130], [150, 195, 200], [215, 235, 235]], "lake": (150, 196, 204), "river": (150, 196, 204), + "ground": {6: (252, 250, 244)}, "sheet": (248, 250, 252), "sea_ice": (236, 244, 246), "sea_ice_seasonal": (220, 236, 238)}, +} + + +def style_palette(style): + return None if style is None else STYLES[style] + + +def relief_rgb(z, hs, zone, ground, ice, lake, land=None, vary=None, style=None, low_memory=False): + """vary: optional (brightness factor, tint) per pixel for open land (not lakes or ice): tint > 0 drier/yellower, + < 0 lusher (deeper green). style: a STYLES key (alien palette) or None. low_memory: the same pixels, made + CHUNK rows at a time (no full-size float64 scratch).""" + if low_memory: + out = np.empty(np.shape(z) + (3,), np.uint8) + for y0 in range(0, np.shape(z)[0], CHUNK): + s = slice(y0, y0 + CHUNK) + out[s] = relief_rgb(z[s], hs[s], zone[s], ground[s], ice[s], lake[s], None if land is None else land[s], + None if vary is None else tuple(np.asarray(v)[s] for v in vary), style) + return out + land = z > 0 if land is None else land + st = style_palette(style) + pal = holdridge_palette() if st is None else _region_palette(st["zones"]) + rgb = pal[np.clip(zone, 0, 37)].astype(np.float64) + ocean = _ramp(z, -6500.0 if st is None else -1500.0, 0.0, + [[11, 43, 90], [30, 90, 150], [143, 198, 224]] if st is None else st["ocean"]).astype(np.float64) + rgb = np.where(land[..., None], rgb, ocean) + grounds = ((1, (111, 143, 106)), (2, (120, 128, 100)), (5, (63, 111, 74)), (6, (239, 233, 220))) if st is None \ + else tuple(st["ground"].items()) + for code, col in grounds: + rgb = np.where((land & (ground == code))[..., None], col, rgb) + if vary is not None: + bright, tint = (np.asarray(v, dtype=np.float64)[..., None] for v in vary) + shift = np.where(tint > 0, tint * np.array([1.0, 0.6, -0.8]), -tint * np.array([-0.9, -0.2, -0.6])) + if st is not None: + shift = np.abs(tint) * np.array([0.4, 0.4, 0.4]) * np.sign(tint) + rgb = np.where(land[..., None], rgb * bright + shift, rgb) + rgb = np.where((lake & land)[..., None], (79, 143, 192) if st is None else st["lake"], rgb) + rgb = np.where(np.isin(ice, [1, 2])[..., None], (244, 248, 251) if st is None else st["sheet"], rgb) + rgb = np.where((ice == 4)[..., None], (225, 235, 242) if st is None else st["sea_ice"], rgb) + rgb = np.where((ice == 3)[..., None], 0.5 * rgb + 0.5 * np.array([220, 232, 240] if st is None else st["sea_ice_seasonal"]), rgb) + shade = np.where(land, 0.55 + 0.45 * hs, 0.85 + 0.15 * hs) + return np.clip(rgb * shade[..., None], 0, 255).astype(np.uint8) + + +RIVER_RGB = (58, 112, 176) + + +def draw_rivers(rgb, g, recv, river, strahler, min_order=2, colour=RIVER_RGB): + """Draw river segments (cell centre → receiver) of order ≥ min_order; width grows with order.""" + H, W = rgb.shape[:2] + im = Image.fromarray(rgb) + draw = ImageDraw.Draw(im) + x = (np.asarray(g.lon) + 180.0) / 360.0 * W - 0.5 + y = (90.0 - np.asarray(g.lat)) / 180.0 * H - 0.5 + cells = np.flatnonzero(river & (recv != np.arange(g.n)) & (strahler >= min_order)) + for i in cells[np.argsort(strahler[cells])]: + j = recv[i] + if abs(x[i] - x[j]) > W / 2: + continue + width = max(1, int(round(int(strahler[i]) * W / 8192))) + draw.line([(x[i], y[i]), (x[j], y[j])], fill=tuple(colour), width=width) + return np.array(im) + + +def draw_currents(rgb, g, current, ocean, per_row=60, colour=(255, 255, 255)): + """Arrows along the surface current on a lattice ≈ `per_row` across the map; length ∝ speed (1 m/s ≈ one + lattice step), sea only; currents under 2 cm/s get none.""" + H, W = rgb.shape[:2] + step = W / per_row + current = np.asarray(current, dtype=np.float64) + ocean = np.asarray(ocean, bool) + e, n = east_north(g.xyz) + tree = cKDTree(g.xyz) + ys, xs = np.meshgrid(np.arange(step / 2, H, step), np.arange(step / 2, W, step), indexing="ij") + lat, lon = 90.0 - (ys.ravel() + 0.5) / H * 180.0, (xs.ravel() + 0.5) / W * 360.0 - 180.0 + _, cell = tree.query(latlon_to_xyz(lat, lon)) + ue, un = np.sum(current[cell] * e[cell], axis=1), np.sum(current[cell] * n[cell], axis=1) + sp = np.hypot(ue, un) + im = Image.fromarray(rgb) + draw = ImageDraw.Draw(im) + for x, y, a, b, s, ok in zip(xs.ravel(), ys.ravel(), ue, un, sp, ocean[cell]): + if not ok or s < 0.02: + continue + L = min(s, 1.5) * 0.9 * step + dx, dy = a / s * L, -b / s * L + x0, y0, x1, y1 = x - dx / 2, y - dy / 2, x + dx / 2, y + dy / 2 + draw.line([(x0, y0), (x1, y1)], fill=tuple(colour), width=1) + for ang in (2.6, -2.6): # arrowhead: two barbs ±150° + c, sn = np.cos(ang), np.sin(ang) + draw.line([(x1, y1), (x1 + 0.35 * (c * dx - sn * dy), y1 + 0.35 * (sn * dx + c * dy))], + fill=tuple(colour), width=1) + return np.array(im) + + +RAMPS = { # key: (unit, lo, hi, colour stops) — shared by previews, viewer textures and legends + "elevation": ("m", -6000.0, 6000.0, [[8, 30, 70], [30, 90, 150], [140, 200, 225], [60, 120, 60], + [150, 160, 90], [140, 110, 80], [235, 235, 235]]), + "T_mean": ("degC", -40.0, 40.0, [[40, 60, 160], [240, 240, 240], [180, 30, 30]]), + "T_range": ("degC", 0.0, 60.0, [[68, 1, 84], [33, 145, 140], [253, 231, 37]]), + "P": ("mm/yr", 0.0, 4000.0, [[150, 110, 60], [230, 220, 150], [60, 150, 70], [30, 70, 170]]), + "po2": ("bar", 0.08, 0.32, [[60, 40, 90], [240, 240, 240], [200, 90, 20]]), + "gravity": ("g", 0.3, 1.8, [[20, 120, 180], [240, 240, 240], [120, 40, 40]]), + "pressure": ("bar", 0.3, 2.5, [[40, 60, 120], [240, 240, 240], [150, 60, 30]]), + "o2_fraction": ("fraction", 0.10, 0.40, [[60, 40, 90], [240, 240, 240], [200, 90, 20]]), + "fire": ("x", 0.0, 1.5, [[40, 80, 160], [240, 240, 240], [220, 120, 20]]), + "vent_potential": ("0-1", 0.0, 1.0, [[10, 20, 50], [60, 90, 160], [250, 190, 60], [255, 250, 220]]), + "bottom_temp": ("degC", -2.0, 30.0, [[30, 40, 120], [80, 160, 200], [240, 200, 120]]), + "sediment": ("m", 0.0, 3000.0, [[40, 30, 60], [140, 110, 80], [240, 220, 170]]), + "plant_height": ("x", 0.5, 3.5, [[150, 120, 60], [240, 240, 240], [30, 110, 50]]), + "sst": ("degC", -2.0, 32.0, [[30, 40, 120], [60, 150, 200], [240, 240, 200], [220, 90, 40]]), + "productivity": ("0-1", 0.0, 1.0, [[10, 20, 60], [20, 110, 120], [120, 200, 90], [240, 240, 120]]), + "current_speed": ("m/s", 0.0, 1.0, [[10, 20, 50], [40, 90, 170], [120, 200, 230], [255, 255, 255]]), + "upwelling": ("m/yr", -200.0, 200.0, [[40, 60, 160], [240, 240, 240], [30, 140, 80]]), +} +VIEWER_LAYERS = [ # id, name, source (continuous raster name | categorical name | "relief") + ("relief", "Relief", "relief"), ("biomes", "Biomes (Holdridge)", "holdridge"), + ("elevation", "Elevation", "elevation"), ("temperature", "Mean temperature", "T_mean"), + ("rainfall", "Rainfall", "P_ann"), ("seasonality", "Seasonality", "seasonality"), + ("landform", "Landform", "landform"), ("ground", "Ground", "ground"), ("ice", "Ice", "ice"), ("deposits", "Mineral deposits", "deposits"), + ("plates", "Plates", "plates"), ("o2", "O₂ partial pressure", "po2"), ("gravity", "Gravity", "gravity"), + ("pressure", "Air pressure", "pressure"), ("fire", "Fire reactivity", "fire"), + ("seabed", "Sea-floor type", "seabed_type"), ("minerals", "Sea-floor minerals", "seabed_mineral"), + ("bottom_temp", "Bottom temperature", "bottom_temp"), ("sediment", "Sediment", "sediment"), + ("vent_potential", "Vent potential", "vent_potential"), + ("currents", "Ocean currents", "current_speed"), ("sst", "Sea-surface temperature", "sst"), + ("productivity", "Sea productivity", "productivity"), +] + + +def _ramp_key(name): + return "P" if name.startswith("P_") else name + + +def _colorize(name, v): + _, lo, hi, stops = RAMPS[_ramp_key(name)] + return _ramp(v, lo, hi, stops) + + +def _preview(img: Image.Image, width: int, nearest: bool) -> Image.Image: + return img.resize((width, width // 2), Image.NEAREST if nearest else Image.LANCZOS) + + +def contact_sheet(tiles, tile_w): + cols = 3 + th = tile_w // 2 + 14 + rows = (len(tiles) + cols - 1) // cols + sheet = Image.new("RGB", (cols * tile_w, rows * th), (20, 20, 24)) + draw = ImageDraw.Draw(sheet) + for k, (name, im) in enumerate(tiles): + x, y = (k % cols) * tile_w, (k // cols) * th + sheet.paste(im.resize((tile_w, tile_w // 2)), (x, y + 14)) + draw.text((x + 4, y + 1), name, fill=(235, 235, 235)) + return sheet + + +class _Writer: + """Saves files on a background thread (PNG/zlib encoding releases the GIL) while the next layer is computed; + at most `depth` waiting, so memory stays bounded. The files are the same bytes as saved in line.""" + def __init__(self, threads: int = 2, depth: int = 4): + from concurrent.futures import ThreadPoolExecutor + self.ex, self.pending, self.depth = ThreadPoolExecutor(threads), [], depth + + def __call__(self, fn, *args, **kw): + while len(self.pending) >= self.depth: + self.pending.pop(0).result() + self.pending.append(self.ex.submit(fn, *args, **kw)) + + def close(self): + try: + for f in self.pending: + f.result() + finally: + self.ex.shutdown(wait=True, cancel_futures=True) + + +def run(ctx) -> dict: + writer = _Writer() + try: + return _run(ctx, writer) + finally: + writer.close() + + +def _run(ctx, save) -> dict: + g, cfg, d = ctx.grid, ctx.cfg, ctx.data + W = int(cfg["build"]["raster_width"]) + if ctx.res < int(cfg["build"]["res_final"]): + W = int(cfg["build"].get("dev_raster_width", W)) + H = W // 2 + pw = int(cfg["build"]["preview_width"]) + out = ctx.out_dir or ctx.root / "out" / f"r{ctx.res}" + rdir = out / "raster" + pdir = ctx.preview_dir or ctx.root / "previews" / f"r{ctx.res}" + rdir.mkdir(parents=True, exist_ok=True) + pdir.mkdir(parents=True, exist_ok=True) + idx, wts = pixel_neighbours(g, W, H) + legends = {**LEGENDS, "age_class": AGE_NAMES, "ice": ICE_NAMES, "plates": [p["id"] for p in ctx.tect["plate"]], + "seabed_type": SEABED_NAMES, "seabed_mineral": MINERAL_NAMES} + meta = {"width": W, "height": H, "projection": "equirectangular", "radius_km": g.radius_km, + "units": cfg.get("units", {}), "continuous": {}, "categorical": {}} + tiles = [] + cats = {name: sample_cat(np.asarray(d[key]), idx) for name, key in CATEGORICAL.items()} + + z = sample_cont(np.asarray(d["z_surface_m"], dtype=np.float64), idx, wts) + amp = DETAIL_M[np.clip(cats["landform"], 0, len(DETAIL_M) - 1)] + for y0 in range(0, H, CHUNK): + rows = np.arange(y0, min(H, y0 + CHUNK)) + z[rows] += amp[rows] * fbm(_rows_xyz(rows, W, H), ctx.seed + 61, 5, 64.0).reshape(len(rows), W) + lake = sample_cat(np.asarray(d["lake"]), idx) + ocean_cells = np.asarray(d["ocean"]) if "ocean" in d else np.asarray(d["z_surface_m"]) <= 0 + ocean_k = ocean_cells[idx] + coast = ocean_k.any(axis=-1) & ~ocean_k.all(axis=-1) + for y0 in range(0, H, CHUNK): # extra fine detail where the coastline runs + rows = np.arange(y0, min(H, y0 + CHUNK)) + cz = COAST_DETAIL_M * fbm(_rows_xyz(rows, W, H), ctx.seed + 67, 5, 256.0).reshape(len(rows), W) + z[rows] += np.where(coast[rows], cz, 0.0) + land = pixel_land(ocean_k, z) + if "lake_level_m" in d: # the water surface for drawing: lake cells at their level (beds stay in `elevation`) + lev = np.asarray(d["lake_level_m"], dtype=np.float64) + dz = sample_cont(np.where(np.isfinite(lev), lev - np.asarray(d["z_surface_m"], dtype=np.float64), 0.0), idx, wts) + save(Image.fromarray(encode(z + dz, *CONTINUOUS["elevation"][1:3])).save, rdir / "surface.png") + meta["continuous"]["surface"] = {"file": "surface.png", "scale": CONTINUOUS["elevation"][1], + "offset": CONTINUOUS["elevation"][2], "unit": "m"} + hs = hillshade(z, g.radius_km, low_memory=ctx.low_memory) + style = cfg.get("render", {}).get("style") + rgb = relief_rgb(z, hs, cats["holdridge"], cats["ground"], cats["ice"], lake, land, style=style, + low_memory=ctx.low_memory) + rgb = draw_rivers(rgb, g, np.asarray(d["recv"]), np.asarray(d["river"]), np.asarray(d["strahler"]), + colour=RIVER_RGB if style is None else STYLES[style]["river"]) + relief = Image.fromarray(rgb) + projections.write_all(rgb, pdir, g, ocean_cells, pw, globe_size=max(pw // 2, 64)) + save(relief.copy().save, rdir / "relief.png") + tiles.append(("relief", _preview(relief, pw, False))) + + vdir = out / "viewer" + by_src = {} # viewer layers are written as soon as their colours exist + for vid, vname, src in VIEWER_LAYERS: + by_src.setdefault(src, []).append((vid, vname)) + ventries = {} + + def emit(src, rgb_, legend): + for vid, vname in by_src.get(src, []): + ventries[vid] = viewer_export.write_layer(vdir, {"id": vid, "name": vname, "rgb": rgb_, "legend": legend}) + + emit("relief", rgb, None) + for name, (key, scale, offset, unit) in CONTINUOUS.items(): + v = z if name == "elevation" else sample_cont(np.asarray(d[key], dtype=np.float64), idx, wts) + save(Image.fromarray(encode(v, scale, offset)).save, rdir / f"{name}.png") + meta["continuous"][name] = {"file": f"{name}.png", "scale": scale, "offset": offset, "unit": unit} + crgb = _colorize(name, v) + if name == "current_speed" and "current" in d: + crgb = draw_currents(crgb, g, d["current"], ocean_cells) + if name != "elevation": + tiles.append((name, _preview(Image.fromarray(crgb), pw, False))) + if name in by_src: + unit_, lo, hi, stops = RAMPS[_ramp_key(name)] + emit(name, crgb, viewer_export.continuous_legend(unit_, lo, hi, stops)) + del crgb + for name, v in cats.items(): + save(Image.fromarray(v.astype(np.uint8), "L").save, rdir / f"{name}.png") + meta["categorical"][name] = {"file": f"{name}.png", "legend": legends[name]} + pal = holdridge_palette() if name == "holdridge" else category_palette(len(legends[name])) + crgb = pal[np.clip(v, 0, len(pal) - 1)] + tiles.append((name, _preview(Image.fromarray(crgb), pw, True))) + emit(name, crgb, viewer_export.categorical_legend(legends[name], pal)) + del crgb + (out / "fields.json").write_text(json.dumps(meta, indent=1)) + + per_cell = {k: v for k, v in d.items() if isinstance(v, np.ndarray) and v.shape[:1] == (g.n,)} + save(lambda: np.savez_compressed(out / "cells.npz", **per_cell)) + (out / "cells_meta.json").write_text(json.dumps( + {"res": ctx.res, "radius_km": g.radius_km, "n_cells": g.n, "units": cfg.get("units", {}), + "planet": {**cfg["planet"], **({"sun_lock": cfg["climate"].get("lock_at", [0.0, 0.0])} + if cfg.get("climate", {}).get("lock") else {})}, "name": cfg.get("render", {}).get("name", "World"), "style": style, + "legends": legends, "fields": sorted(per_cell), + "plateaus": [PL.features(p, ctx.seed, g.radius_km) for p in ctx.tect.get("plateau", [])]}, indent=1)) + + for name, im in tiles: + im.save(pdir / f"{name}.png") + contact_sheet(tiles, max(pw // 3, 64)).save(pdir / "contact_sheet.png") + viewer_export.write_index(vdir, [ventries[vid] for vid, _, _ in VIEWER_LAYERS]) + geo.write_all(ctx, out, land, lake) + return {} |