settled-reach/tooling/planet-gen/sol_data/shared.py

"""
Shared utilities for loading and processing real-world planetary data.

All loaders produce arrays compatible with the planet_simulation terrain dict:
  - Grid size: GRID_H x GRID_W (256 x 512)
  - Elevation: float32 [0, 1] normalised
  - Temperature: float32 in absolute Kelvin (normalised to [0,1] later)
  - Moisture: float32 [0, 1]
  - Sea level: float elevation threshold
"""

import math
import numpy as np
from scipy.ndimage import zoom
from PIL import Image

# Planetary DEMs can exceed PIL's default decompression bomb limit
Image.MAX_IMAGE_PIXELS = None

# Match planet_simulation grid
GRID_W = 512
GRID_H = 256


# ─── Loading ────────────────────────────────────────────────────────────────

def load_tiff_as_array(path: str) -> np.ndarray:
    """
    Load a GeoTIFF/TIFF as a numpy array via PIL.

    PIL handles uncompressed and LZW-compressed TIFFs with 8/16/32-bit
    integer or float samples. For multi-band, returns (H, W, bands).
    For single-band, returns (H, W).
    """
    img = Image.open(path)
    arr = np.array(img, dtype=np.float32)
    return arr


def load_raw_binary(path: str, width: int, height: int,
                    dtype: str = ">i2", offset: int = 0) -> np.ndarray:
    """
    Load a raw binary raster (PDS IMG, .bin, etc).

    Parameters
    ----------
    path   : file path
    width  : number of columns
    height : number of rows
    dtype  : numpy dtype string (e.g. ">i2" for big-endian int16)
    offset : byte offset to skip (header size)
    """
    dt = np.dtype(dtype)
    expected_bytes = width * height * dt.itemsize
    with open(path, "rb") as f:
        f.seek(offset)
        raw = f.read(expected_bytes)
    if len(raw) < expected_bytes:
        raise ValueError(
            f"Expected {expected_bytes} bytes, got {len(raw)}. "
            f"Check width/height/dtype/offset."
        )
    arr = np.frombuffer(raw, dtype=dt).reshape(height, width).astype(np.float32)
    return arr


def load_image_as_elevation(path: str, invert: bool = False) -> np.ndarray:
    """
    Load a greyscale or RGB image and convert to float32 elevation.
    For RGB, uses luminance. For greyscale, uses the single channel.
    """
    img = Image.open(path).convert("L")
    arr = np.array(img, dtype=np.float32) / 255.0
    if invert:
        arr = 1.0 - arr
    return arr


# ─── Resampling ─────────────────────────────────────────────────────────────

def resample_to_grid(arr: np.ndarray, target_h: int = GRID_H,
                     target_w: int = GRID_W,
                     order: int = 1) -> np.ndarray:
    """
    Resample a 2D array to target grid size.

    order: 0=nearest, 1=bilinear, 3=cubic
    """
    if arr.shape == (target_h, target_w):
        return arr.astype(np.float32)
    zoom_y = target_h / arr.shape[0]
    zoom_x = target_w / arr.shape[1]
    return zoom(arr, (zoom_y, zoom_x), order=order).astype(np.float32)


# ─── Normalisation ──────────────────────────────────────────────────────────

def normalize_01(arr: np.ndarray, lo: float = None, hi: float = None) -> np.ndarray:
    """Normalise array to [0, 1]."""
    if lo is None:
        lo = float(arr.min())
    if hi is None:
        hi = float(arr.max())
    if hi - lo < 1e-9:
        return np.zeros_like(arr, dtype=np.float32)
    return np.clip((arr - lo) / (hi - lo), 0.0, 1.0).astype(np.float32)


def compute_sea_level(elevation: np.ndarray, ocean_fraction: float) -> float:
    """
    Compute sea_level threshold such that ocean_fraction of cells are below it.
    """
    if ocean_fraction <= 0.0:
        return 0.0
    if ocean_fraction >= 1.0:
        return 1.0
    return float(np.percentile(elevation, ocean_fraction * 100.0))


# ─── Longitude shift ────────────────────────────────────────────────────────

def shift_longitude(arr: np.ndarray, shift_cols: int) -> np.ndarray:
    """
    Roll array along the longitude (column) axis.

    The pipeline uses col 0 = 180°W. If source data uses col 0 = 0° (Greenwich),
    shift by half the width to align.
    """
    return np.roll(arr, shift_cols, axis=1)


def greenwich_to_dateline(arr: np.ndarray) -> np.ndarray:
    """
    Shift from col 0 = 0° (Greenwich) to col 0 = 180°W (dateline).
    Standard for most NASA/NOAA global datasets → pipeline convention.
    """
    return shift_longitude(arr, arr.shape[1] // 2)


# ─── Hillshade ──────────────────────────────────────────────────────────────

def compute_hillshade(elevation: np.ndarray,
                      sun_azimuth_deg: float = 315.0,
                      sun_altitude_deg: float = 45.0) -> np.ndarray:
    """
    Compute hillshade from elevation grid. Matches planet_simulation.compute_hillshade().
    """
    scale = elevation.shape[1] / 8.0
    gy, gx = np.gradient(elevation * scale)
    mag = np.sqrt(gx**2 + gy**2 + 1.0)
    nx = -gx / mag
    ny = -gy / mag
    nz = 1.0 / mag

    az = math.radians(sun_azimuth_deg)
    alt = math.radians(sun_altitude_deg)
    lx = math.cos(alt) * math.sin(az)
    ly = -math.cos(alt) * math.cos(az)
    lz = math.sin(alt)

    shade = np.clip(nx * lx + ny * ly + nz * lz, 0.0, 1.0)
    return shade.astype(np.float32)


# ─── Analytical temperature models ─────────────────────────────────────────

def temperature_equilibrium_K(luminosity_solar: float, distance_au: float,
                              albedo: float = 0.3) -> float:
    """
    Stefan-Boltzmann equilibrium temperature in Kelvin.
    """
    L_sun = 3.828e26  # watts
    sigma = 5.670e-8
    d_m = distance_au * 1.496e11
    T_eq = ((luminosity_solar * L_sun * (1 - albedo)) /
            (16 * math.pi * sigma * d_m**2)) ** 0.25
    return T_eq


def temperature_grid_analytical(
    base_T_K: float,
    grid_h: int = GRID_H,
    grid_w: int = GRID_W,
    elevation: np.ndarray = None,
    lapse_rate_K_per_unit: float = 40.0,
    lat_gradient_K: float = 60.0,
) -> np.ndarray:
    """
    Analytical temperature grid: equator-to-pole gradient + elevation lapse.

    Parameters
    ----------
    base_T_K             : equatorial temperature in Kelvin
    elevation            : normalised [0,1] elevation grid (optional)
    lapse_rate_K_per_unit: temperature drop per unit elevation
    lat_gradient_K       : total temperature drop from equator to pole
    """
    v = np.linspace(0, 1, grid_h, dtype=np.float32)
    lat_frac = np.abs(v - 0.5) * 2.0  # 0 at equator, 1 at poles
    lat_temp = lat_frac[:, np.newaxis] * lat_gradient_K  # broadcast to (H, W)
    temp = np.full((grid_h, grid_w), base_T_K, dtype=np.float32)
    temp -= lat_temp
    if elevation is not None:
        temp -= elevation * lapse_rate_K_per_unit
    return temp


# ─── Terrain dict assembly ──────────────────────────────────────────────────

def assemble_terrain(
    elevation: np.ndarray,
    temperature_K: np.ndarray,
    moisture: np.ndarray,
    biome: np.ndarray,
    surface_water: np.ndarray,
    hillshade: np.ndarray,
    rivers: list,
    sea_level: float,
) -> dict:
    """
    Assemble the terrain dict in the format expected by render_heightmap
    and render_globe. Temperature is normalised to [0,1] for the output
    (matching planet_simulation.simulate() lines 891-893).
    """
    H, W = elevation.shape
    river_grid = np.zeros((H, W), dtype=bool)
    for path in rivers:
        for r, c in path:
            if 0 <= r < H and 0 <= c < W:
                river_grid[r, c] = True

    # Normalise temperature to [0,1] for renderer display
    t_min, t_max = temperature_K.min(), temperature_K.max()
    temp_norm = ((temperature_K - t_min) / (t_max - t_min + 1e-9)).astype(np.float32)

    return {
        "elevation":     elevation.astype(np.float32),
        "temperature":   temp_norm,
        "moisture":      moisture.astype(np.float32),
        "biome":         biome.astype(np.int8),
        "surface_water": surface_water.astype(bool),
        "hillshade":     hillshade.astype(np.float32),
        "river_grid":    river_grid,
        "rivers":        rivers,
        "sea_level":     float(sea_level),
        "_grid_w":       W,
        "_grid_h":       H,
    }