"""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 {}