""" Risk raster tile engine — renders the full-Wuhan 100m risk grid as XYZ map tiles. Why this exists --------------- The model emits risk at ~140k GCN nodes per day. The product needs to display this over the full Wuhan 100m grid (~1.5M in-boundary cells) with smooth LOD. Shipping that many cells to the browser as vectors is impossible, so we rasterize server-side: 1. Build a dense per-cell risk raster R[row, col] once per (date, forecast_day): scatter each node's risk onto its 100m cell (max per cell), then nearest-fill empty cells via a Euclidean distance transform (Voronoi over nodes, quantized to the 100m grid). Cells outside the Wuhan boundary are masked out. 2. Build a max-pooled pyramid for clean LOD at low zoom. 3. Render standard 256x256 web-mercator PNG tiles by sampling the pyramid level that matches the tile's zoom. Tiles are cached; the browser just loads images. Coordinate conventions (calibrated from processed/grid_100m_index.parquet): lat = MIN_LAT + (row + 0.5) * LAT_STEP -> row 0 is SOUTH, row increases north lon = MIN_LON + (col + 0.5) * LON_STEP -> col 0 is WEST, col increases east Everything below maps lat/lon -> (row, col) the same way, so orientation is coherent end to end. Tile pixel py=0 is north (high lat -> high row); we build the RGBA array with py as the first axis so north ends up at the top of the PNG. """ from __future__ import annotations import io import json import math from functools import lru_cache from pathlib import Path import numpy as np from PIL import Image, ImageDraw from scipy.ndimage import distance_transform_edt from config import DATA_DIR, WUHAN_BOUNDS, PROJECT_ROOT # --- Grid definition (geographically-correct ~100m grid over the Wuhan bbox) --- MIN_LON = WUHAN_BOUNDS["min_lon"] MAX_LON = WUHAN_BOUNDS["max_lon"] MIN_LAT = WUHAN_BOUNDS["min_lat"] MAX_LAT = WUHAN_BOUNDS["max_lat"] NROWS = 1550 # matches processed/grid_100m_with_dem_pop.parquet row extent NCOLS = 1336 # matches its col extent LAT_STEP = (MAX_LAT - MIN_LAT) / NROWS LON_STEP = (MAX_LON - MIN_LON) / NCOLS BOUNDARY_GEOJSON = PROJECT_ROOT / "Datas" / "武汉市.geojson" # Forecast-day -> property suffix on the risk geojson features. _DAY_TO_KEY = {1: "risk_1d", 3: "risk_3d", 7: "risk_7d"} MAX_PYRAMID_LEVEL = 7 # full-res + 7 downsamples covers world zoom range TILE_PX = 256 # ---------------------------------------------------------------------------- # Affine helpers (lat/lon <-> grid row/col) # ---------------------------------------------------------------------------- def latlon_to_rowcol(lat: float, lon: float) -> tuple[int, int]: row = int((lat - MIN_LAT) / LAT_STEP) col = int((lon - MIN_LON) / LON_STEP) row = max(0, min(NROWS - 1, row)) col = max(0, min(NCOLS - 1, col)) return row, col # ---------------------------------------------------------------------------- # Boundary mask (rasterized once) # ---------------------------------------------------------------------------- @lru_cache(maxsize=1) def _boundary_mask() -> np.ndarray: """Boolean (NROWS, NCOLS) mask, True for cells inside the Wuhan boundary.""" img = Image.new("1", (NCOLS, NROWS), 0) draw = ImageDraw.Draw(img) if not BOUNDARY_GEOJSON.exists(): # No boundary file -> color the whole bbox rather than nothing. return np.ones((NROWS, NCOLS), dtype=bool) with open(BOUNDARY_GEOJSON, "r", encoding="utf-8") as f: gj = json.load(f) def _draw_ring(ring): pts = [] for lon, lat in ring: col = (lon - MIN_LON) / LON_STEP row = (lat - MIN_LAT) / LAT_STEP pts.append((col, row)) if len(pts) >= 3: draw.polygon(pts, fill=1) def _walk(geom): gtype = geom.get("type") coords = geom.get("coordinates", []) if gtype == "Polygon": for ring in coords: _draw_ring(ring) elif gtype == "MultiPolygon": for poly in coords: for ring in poly: _draw_ring(ring) if gj.get("type") == "FeatureCollection": for feat in gj.get("features", []): _walk(feat.get("geometry", {})) elif gj.get("type") == "Feature": _walk(gj.get("geometry", {})) else: _walk(gj) return np.array(img, dtype=bool) # ---------------------------------------------------------------------------- # Node loading (cached per date) # ---------------------------------------------------------------------------- @lru_cache(maxsize=8) def _load_nodes(date: str) -> tuple[np.ndarray, np.ndarray, np.ndarray]: """Return (rows, cols, risks[3]) arrays for all nodes on a date. risks is shape (n_nodes, 3) for [risk_1d, risk_3d, risk_7d]. """ filepath = DATA_DIR / f"risk_{date}.geojson" if not filepath.exists(): raise FileNotFoundError(f"No risk data for date {date}") with open(filepath, "r", encoding="utf-8") as f: gj = json.load(f) feats = gj.get("features", []) n = len(feats) rows = np.empty(n, dtype=np.int32) cols = np.empty(n, dtype=np.int32) risks = np.zeros((n, 3), dtype=np.float32) for i, feat in enumerate(feats): p = feat.get("properties", {}) lat = p.get("lat", 0.0) lon = p.get("lon", 0.0) r, c = latlon_to_rowcol(lat, lon) rows[i] = r cols[i] = c risks[i, 0] = p.get("risk_1d", 0.0) risks[i, 1] = p.get("risk_3d", 0.0) risks[i, 2] = p.get("risk_7d", 0.0) return rows, cols, risks # ---------------------------------------------------------------------------- # Risk raster + pyramid (cached per date+day) # ---------------------------------------------------------------------------- def _maxpool2(a: np.ndarray) -> np.ndarray: """Downsample by 2 taking the NaN-aware max of each 2x2 block.""" h, w = a.shape h2, w2 = (h + 1) // 2, (w + 1) // 2 out = np.full((h2, w2), np.nan, dtype=np.float32) # Pad to even dims with NaN so the reshape is clean. ph, pw = h2 * 2, w2 * 2 pad = np.full((ph, pw), np.nan, dtype=np.float32) pad[:h, :w] = a blocks = pad.reshape(h2, 2, w2, 2) # np.nanmax over the 2x2 block axes; suppress all-NaN warnings. with np.errstate(invalid="ignore"): out = np.nanmax(blocks, axis=(1, 3)) return out.astype(np.float32) @lru_cache(maxsize=6) def _risk_pyramid(date: str, day: int) -> tuple[np.ndarray, ...]: """Build the nearest-filled, boundary-masked risk raster and its LOD pyramid. Returns a tuple of arrays, level 0 = full res (NROWS, NCOLS), each subsequent level downsampled 2x. NaN marks "no data / outside boundary". """ if day not in _DAY_TO_KEY: raise ValueError(f"invalid forecast day {day}") day_idx = {1: 0, 3: 1, 7: 2}[day] rows, cols, risks = _load_nodes(date) vals = risks[:, day_idx] # Scatter nodes onto the grid, taking the max risk per cell. R = np.full((NROWS, NCOLS), -np.inf, dtype=np.float32) np.maximum.at(R, (rows, cols), vals) known = np.isfinite(R) # Nearest-fill empty cells (Voronoi over nodes, quantized to the 100m grid). if known.any(): idx = distance_transform_edt(~known, return_distances=False, return_indices=True) R = R[tuple(idx)] R = R.astype(np.float32) # Mask out everything outside the Wuhan boundary. mask = _boundary_mask() R[~mask] = np.nan pyramid = [R] for _ in range(MAX_PYRAMID_LEVEL): nxt = _maxpool2(pyramid[-1]) pyramid.append(nxt) if nxt.shape[0] <= 2 or nxt.shape[1] <= 2: break return tuple(pyramid) # ---------------------------------------------------------------------------- # Colormap (risk 0..1 -> RGBA), built once as a 256-entry LUT # ---------------------------------------------------------------------------- @lru_cache(maxsize=1) def _color_lut() -> np.ndarray: """256x4 uint8 LUT. Green -> yellow -> orange -> red, alpha grows with risk. Risk below ~0.25 is rendered transparent to keep the map readable. """ lut = np.zeros((256, 4), dtype=np.uint8) # control points: (risk, R, G, B) stops = [ (0.00, 56, 176, 0), # green (low) (0.40, 250, 204, 21), # yellow (medium) (0.60, 249, 115, 22), # orange (high) (0.80, 239, 68, 68), # red (critical) (1.00, 153, 27, 27), # dark red (extreme) ] xs = [s[0] for s in stops] for i in range(256): t = i / 255.0 # piecewise-linear RGB interpolation for k in range(len(stops) - 1): if xs[k] <= t <= xs[k + 1]: f = (t - xs[k]) / (xs[k + 1] - xs[k] + 1e-9) r = stops[k][1] + f * (stops[k + 1][1] - stops[k][1]) g = stops[k][2] + f * (stops[k + 1][2] - stops[k][2]) b = stops[k][3] + f * (stops[k + 1][3] - stops[k][3]) break else: r, g, b = stops[-1][1:] # alpha: transparent below 0.25, then ramp 90 -> 235 if t < 0.25: a = 0.0 else: a = 90 + (t - 0.25) / 0.75 * (235 - 90) lut[i] = (int(r), int(g), int(b), int(a)) return lut # ---------------------------------------------------------------------------- # Tile rendering # ---------------------------------------------------------------------------- def _tile_pixel_latlon(z: int, x: int, y: int) -> tuple[np.ndarray, np.ndarray]: """Return (lat[256,256], lon[256,256]) for each pixel center of a tile.""" n = 2.0 ** z px = (np.arange(TILE_PX) + 0.5) / TILE_PX # longitude is linear in tile-x X = (x + px) / n lon = X * 360.0 - 180.0 # shape (256,) # latitude via inverse web-mercator (nonlinear in tile-y) Y = (y + px) / n lat = np.degrees(np.arctan(np.sinh(np.pi * (1.0 - 2.0 * Y)))) # shape (256,) lon2d = np.broadcast_to(lon, (TILE_PX, TILE_PX)) # varies along axis 1 (px) lat2d = np.broadcast_to(lat[:, None], (TILE_PX, TILE_PX)) # varies along axis 0 (py) return lat2d, lon2d def _level_for_zoom(z: int) -> int: """Pick the pyramid level so ~1 source cell maps to ~1 screen pixel.""" # meters/pixel at lat ~30.6: 156543.03 * cos(lat) / 2^z ; /100m per cell cells_per_px = (156543.03 * math.cos(math.radians(30.6)) / (2.0 ** z)) / 100.0 if cells_per_px <= 1.0: return 0 return max(0, min(MAX_PYRAMID_LEVEL, int(math.floor(math.log2(cells_per_px))))) def render_tile(z: int, x: int, y: int, date: str, day: int = 1) -> bytes: """Render a single XYZ tile to PNG bytes. Fully transparent tiles return a tiny cached blank PNG. Result is cached per (z,x,y,date,day).""" return _render_tile_cached(z, x, y, date, day) @lru_cache(maxsize=1024) def _render_tile_cached(z: int, x: int, y: int, date: str, day: int) -> bytes: pyramid = _risk_pyramid(date, day) level = _level_for_zoom(z) level = min(level, len(pyramid) - 1) R = pyramid[level] factor = 2 ** level lh, lw = R.shape lat2d, lon2d = _tile_pixel_latlon(z, x, y) # lat/lon -> full-res row/col -> level row/col row = ((lat2d - MIN_LAT) / LAT_STEP).astype(np.int32) // factor col = ((lon2d - MIN_LON) / LON_STEP).astype(np.int32) // factor inside = (row >= 0) & (row < lh) & (col >= 0) & (col < lw) rc = np.clip(row, 0, lh - 1) cc = np.clip(col, 0, lw - 1) sampled = R[rc, cc] # (256,256) float32, NaN where no data valid = inside & np.isfinite(sampled) # Map risk -> LUT index (NaN cells become 0 then are zeroed-out below) lut = _color_lut() safe = np.nan_to_num(sampled, nan=0.0) idx = np.clip((safe * 255.0), 0, 255).astype(np.uint8) rgba = lut[idx] # (256,256,4) rgba[~valid] = (0, 0, 0, 0) # transparent outside data/boundary img = Image.fromarray(rgba, mode="RGBA") buf = io.BytesIO() img.save(buf, format="PNG", optimize=False) return buf.getvalue() # ---------------------------------------------------------------------------- # Point query (for click-to-inspect) # ---------------------------------------------------------------------------- @lru_cache(maxsize=16) def grid_stats(date: str, day: int = 1) -> dict: """Lightweight aggregate stats over the in-boundary 100m grid for a date/day. Computed from the cached raster, so this is cheap after the first tile build. Replaces the old heavy per-viewport LOD fetch the overlay used to do. """ R = _risk_pyramid(date, day)[0] finite = np.isfinite(R) n = int(finite.sum()) if n == 0: return {"cell_count": 0, "avg_risk": 0.0, "max_risk": 0.0, "high_risk_count": 0, "forecast_day": day, "date": date} vals = R[finite] return { "cell_count": n, "avg_risk": round(float(vals.mean()), 4), "max_risk": round(float(vals.max()), 4), "high_risk_count": int((vals >= 0.8).sum()), "forecast_day": day, "date": date, } def query_cell(lat: float, lon: float, date: str, day: int = 1) -> dict: """Return the 100m cell risk at a lat/lon for the given date. Includes all three forecast horizons (1d/3d/7d) so the click panel can show them without a separate heavy grid fetch. `risk_value` is the requested day. """ row, col = latlon_to_rowcol(lat, lon) def _sample(d: int) -> tuple[float, bool]: v = _risk_pyramid(date, d)[0][row, col] ok = bool(np.isfinite(v)) return (round(float(v), 4) if ok else 0.0), ok r1, in_b = _sample(1) r3, _ = _sample(3) r7, _ = _sample(7) current = {1: r1, 3: r3, 7: r7}[day] return { "grid_id": f"r{row}_c{col}", "row": row, "col": col, "center_lat": round(MIN_LAT + (row + 0.5) * LAT_STEP, 6), "center_lon": round(MIN_LON + (col + 0.5) * LON_STEP, 6), "risk_value": current, "risk_1d": r1, "risk_3d": r3, "risk_7d": r7, "in_boundary": in_b, "forecast_day": day, "date": date, }