jpmschweitzerandClaude Opus 4.6
d364e1907d
Ruff pre-push lint caught 17 unused imports across 7 files. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
392 lines
15 KiB
Python
392 lines
15 KiB
Python
"""
|
||
Earth (GJ0d) terrain builder.
|
||
|
||
Data sources:
|
||
- Elevation: ETOPO 2022 60 arc-second (NOAA) — GeoTIFF
|
||
- Temperature: WorldClim v2.1 annual mean (10 arc-min) — GeoTIFF
|
||
- Precipitation: WorldClim v2.1 annual total (10 arc-min) — GeoTIFF
|
||
- Rivers: Natural Earth 10m rivers — GeoJSON
|
||
|
||
All sources are equirectangular with col 0 = 180°W. ETOPO and WorldClim
|
||
use col 0 = 180°W natively. Natural Earth uses -180 to 180 longitude.
|
||
"""
|
||
|
||
import json
|
||
import zipfile
|
||
import numpy as np
|
||
from pathlib import Path
|
||
|
||
from sol_data.download import ensure_cached
|
||
from sol_data.shared import (
|
||
GRID_W, GRID_H,
|
||
load_tiff_as_array,
|
||
resample_to_grid, normalize_01, compute_sea_level,
|
||
compute_hillshade, assemble_terrain,
|
||
)
|
||
|
||
# ─── Data source URLs ───────────────────────────────────────────────────────
|
||
|
||
# ETOPO 2022 60 arc-second — surface elevation (ice surface, not bedrock)
|
||
# ~130 MB GeoTIFF, 21600 x 10800, int16 metres
|
||
ETOPO_URL = "https://www.ngdc.noaa.gov/mgg/global/relief/ETOPO2022/data/60s/60s_surface_elev_gtif/ETOPO_2022_v1_60s_N90W180_surface.tif"
|
||
ETOPO_FILE = "ETOPO_2022_v1_60s_N90W180_surface.tif"
|
||
|
||
# WorldClim v2.1 — 10 arc-minute resolution (migrated to geodata.ucdavis.edu)
|
||
# Temperature: mean annual, °C × 10 (int16), in a zip
|
||
WCLIM_TEMP_URL = "https://geodata.ucdavis.edu/climate/worldclim/2_1/base/wc2.1_10m_tavg.zip"
|
||
WCLIM_TEMP_FILE = "wc2.1_10m_tavg.zip"
|
||
|
||
# Precipitation: annual total mm (int16), in a zip
|
||
WCLIM_PREC_URL = "https://geodata.ucdavis.edu/climate/worldclim/2_1/base/wc2.1_10m_prec.zip"
|
||
WCLIM_PREC_FILE = "wc2.1_10m_prec.zip"
|
||
|
||
# Natural Earth 10m rivers — GeoJSON from GitHub
|
||
RIVERS_URL = "https://raw.githubusercontent.com/nvkelso/natural-earth-vector/master/geojson/ne_10m_rivers_lake_centerlines.geojson"
|
||
RIVERS_FILE = "ne_10m_rivers_lake_centerlines.geojson"
|
||
|
||
# Earth physical constants
|
||
EARTH_OCEAN_FRACTION = 0.71
|
||
EARTH_MIN_ELEV_M = -10994.0 # Mariana Trench
|
||
EARTH_MAX_ELEV_M = 8849.0 # Everest
|
||
|
||
|
||
# ─── River filtering ────────────────────────────────────────────────────────
|
||
|
||
# Rivers to include (smart scatter: 1-2 per continent + Rhine)
|
||
INCLUDED_RIVERS = {
|
||
# Europe
|
||
"Danube", "Volga", "Rhine",
|
||
# North America
|
||
"Mississippi", "St. Lawrence",
|
||
# South America
|
||
"Amazon", "Paraná",
|
||
# Africa
|
||
"Nile", "Congo",
|
||
# West Asia
|
||
"Tigris",
|
||
# East/South Asia
|
||
"Yangtze", "Ganges", "Mekong",
|
||
# Australia
|
||
"Murray",
|
||
}
|
||
|
||
# Fuzzy matching — some NE names differ slightly
|
||
RIVER_NAME_ALIASES = {
|
||
"Parana": "Paraná",
|
||
"Chang Jiang": "Yangtze",
|
||
"Huang He": "Yellow",
|
||
"Ganga": "Ganges",
|
||
"Nil": "Nile",
|
||
"Danau": "Danube",
|
||
"Donau": "Danube",
|
||
"Rhin": "Rhine",
|
||
"Rhein": "Rhine",
|
||
"Saint Lawrence": "St. Lawrence",
|
||
"St Lawrence": "St. Lawrence",
|
||
"Río Paraná": "Paraná",
|
||
"Rio Parana": "Paraná",
|
||
}
|
||
|
||
|
||
def _match_river_name(feature_name: str) -> str:
|
||
"""Check if a Natural Earth river name matches our included set."""
|
||
if not feature_name:
|
||
return None
|
||
name = feature_name.strip()
|
||
# Direct match
|
||
if name in INCLUDED_RIVERS:
|
||
return name
|
||
# Alias match
|
||
if name in RIVER_NAME_ALIASES:
|
||
alias = RIVER_NAME_ALIASES[name]
|
||
if alias in INCLUDED_RIVERS:
|
||
return alias
|
||
# Substring match (e.g. "Mississippi River" contains "Mississippi")
|
||
for included in INCLUDED_RIVERS:
|
||
if included.lower() in name.lower() or name.lower() in included.lower():
|
||
return included
|
||
return None
|
||
|
||
|
||
# ─── Data loaders ───────────────────────────────────────────────────────────
|
||
|
||
def _load_etopo() -> np.ndarray:
|
||
"""Load ETOPO 2022 elevation data, return raw metres array."""
|
||
path = ensure_cached(ETOPO_URL, ETOPO_FILE)
|
||
print(f" loading ETOPO: {path}")
|
||
try:
|
||
arr = load_tiff_as_array(str(path))
|
||
except Exception as e:
|
||
raise RuntimeError(
|
||
f"Failed to load ETOPO GeoTIFF: {e}\n"
|
||
f"If PIL can't read this TIFF, install Pillow with TIFF support "
|
||
f"or convert to raw binary."
|
||
) from e
|
||
print(f" ETOPO shape: {arr.shape}, range: [{arr.min():.0f}, {arr.max():.0f}] m")
|
||
return arr
|
||
|
||
|
||
def _load_worldclim_temperature() -> np.ndarray:
|
||
"""
|
||
Load WorldClim v2.1 annual mean temperature.
|
||
Returns temperature in Kelvin at native resolution.
|
||
"""
|
||
zip_path = ensure_cached(WCLIM_TEMP_URL, WCLIM_TEMP_FILE)
|
||
print(f" loading WorldClim temperature: {zip_path}")
|
||
|
||
# The zip contains monthly TIFFs (tavg_01.tif to tavg_12.tif).
|
||
# Compute annual mean from all 12 months.
|
||
cache_dir = zip_path.parent
|
||
monthly_sum = None
|
||
count = 0
|
||
|
||
with zipfile.ZipFile(zip_path) as zf:
|
||
tif_names = sorted([n for n in zf.namelist() if n.endswith(".tif")])
|
||
for tif_name in tif_names:
|
||
extracted = cache_dir / Path(tif_name).name
|
||
if not extracted.exists():
|
||
zf.extract(tif_name, cache_dir)
|
||
# Handle nested paths in zip
|
||
nested = cache_dir / tif_name
|
||
if nested != extracted and nested.exists():
|
||
nested.rename(extracted)
|
||
try:
|
||
arr = load_tiff_as_array(str(extracted))
|
||
except Exception:
|
||
# Try the nested path
|
||
nested = cache_dir / tif_name
|
||
if nested.exists():
|
||
arr = load_tiff_as_array(str(nested))
|
||
else:
|
||
continue
|
||
# Replace nodata with NaN
|
||
arr[arr < -999] = np.nan
|
||
if monthly_sum is None:
|
||
monthly_sum = arr.copy()
|
||
else:
|
||
monthly_sum += arr
|
||
count += 1
|
||
|
||
if count == 0:
|
||
raise RuntimeError("No temperature TIFFs found in WorldClim archive")
|
||
|
||
# Annual mean (WorldClim tavg is °C × 10)
|
||
temp_C = (monthly_sum / count) / 10.0
|
||
# Convert to Kelvin
|
||
temp_K = temp_C + 273.15
|
||
# Replace NaN (ocean/nodata) with a reasonable ocean temperature
|
||
temp_K = np.nan_to_num(temp_K, nan=288.0)
|
||
|
||
print(f" WorldClim temp shape: {temp_K.shape}, "
|
||
f"range: [{np.nanmin(temp_K):.0f}, {np.nanmax(temp_K):.0f}] K")
|
||
return temp_K
|
||
|
||
|
||
def _load_worldclim_precipitation() -> np.ndarray:
|
||
"""
|
||
Load WorldClim v2.1 annual precipitation (sum of 12 months).
|
||
Returns precipitation in mm/year at native resolution.
|
||
"""
|
||
zip_path = ensure_cached(WCLIM_PREC_URL, WCLIM_PREC_FILE)
|
||
print(f" loading WorldClim precipitation: {zip_path}")
|
||
|
||
cache_dir = zip_path.parent
|
||
annual_sum = None
|
||
|
||
with zipfile.ZipFile(zip_path) as zf:
|
||
tif_names = sorted([n for n in zf.namelist() if n.endswith(".tif")])
|
||
for tif_name in tif_names:
|
||
extracted = cache_dir / Path(tif_name).name
|
||
if not extracted.exists():
|
||
zf.extract(tif_name, cache_dir)
|
||
nested = cache_dir / tif_name
|
||
if nested != extracted and nested.exists():
|
||
nested.rename(extracted)
|
||
try:
|
||
arr = load_tiff_as_array(str(extracted))
|
||
except Exception:
|
||
nested = cache_dir / tif_name
|
||
if nested.exists():
|
||
arr = load_tiff_as_array(str(nested))
|
||
else:
|
||
continue
|
||
arr[arr < -999] = 0.0
|
||
if annual_sum is None:
|
||
annual_sum = arr.copy()
|
||
else:
|
||
annual_sum += arr
|
||
|
||
if annual_sum is None:
|
||
raise RuntimeError("No precipitation TIFFs found in WorldClim archive")
|
||
|
||
print(f" WorldClim precip shape: {annual_sum.shape}, "
|
||
f"range: [{annual_sum.min():.0f}, {annual_sum.max():.0f}] mm/yr")
|
||
return annual_sum
|
||
|
||
|
||
def _load_rivers_geojson() -> list:
|
||
"""
|
||
Load Natural Earth rivers GeoJSON and extract polylines for included rivers.
|
||
Returns list of (name, [(row, col), ...]) in grid coordinates.
|
||
"""
|
||
path = ensure_cached(RIVERS_URL, RIVERS_FILE)
|
||
print(f" loading rivers: {path}")
|
||
|
||
with open(path) as f:
|
||
geojson = json.load(f)
|
||
|
||
rivers = []
|
||
for feature in geojson.get("features", []):
|
||
props = feature.get("properties", {})
|
||
fname = props.get("name") or props.get("name_en") or ""
|
||
matched = _match_river_name(fname)
|
||
if not matched:
|
||
continue
|
||
|
||
geom = feature.get("geometry", {})
|
||
geom_type = geom.get("type", "")
|
||
coords_list = []
|
||
|
||
if geom_type == "LineString":
|
||
coords_list = [geom["coordinates"]]
|
||
elif geom_type == "MultiLineString":
|
||
coords_list = geom["coordinates"]
|
||
else:
|
||
continue
|
||
|
||
for coords in coords_list:
|
||
path_grid = []
|
||
for lon, lat in coords:
|
||
# Convert lon/lat to grid coordinates
|
||
# Grid: row 0 = 90°N, row 255 = 90°S
|
||
# col 0 = 180°W, col 511 = 180°E
|
||
row = int((90.0 - lat) / 180.0 * GRID_H)
|
||
col = int((lon + 180.0) / 360.0 * GRID_W)
|
||
row = max(0, min(GRID_H - 1, row))
|
||
col = max(0, min(GRID_W - 1, col))
|
||
# Deduplicate: skip if same grid cell as previous point.
|
||
# Natural Earth has hundreds of lon/lat points per river,
|
||
# many of which land on the same 512x256 cell. Without
|
||
# dedup, the renderer sees len(path)=300 and draws width 6.
|
||
if path_grid and path_grid[-1] == (row, col):
|
||
continue
|
||
path_grid.append((row, col))
|
||
if len(path_grid) >= 2:
|
||
rivers.append((matched, path_grid))
|
||
|
||
# Deduplicate: keep longest segment per river name
|
||
by_name = {}
|
||
for name, path in rivers:
|
||
if name not in by_name or len(path) > len(by_name[name]):
|
||
by_name[name] = path
|
||
|
||
print(f" matched {len(by_name)} rivers: {', '.join(sorted(by_name.keys()))}")
|
||
return [(name, path) for name, path in by_name.items()]
|
||
|
||
|
||
# ─── Main builder ───────────────────────────────────────────────────────────
|
||
|
||
def build_terrain(body_def: dict) -> dict:
|
||
"""
|
||
Build Earth terrain dict from real-world data.
|
||
|
||
Returns the same dict format as planet_simulation.simulate().
|
||
"""
|
||
import sys
|
||
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
||
from planet_simulation import compute_biome
|
||
|
||
print(" Earth: loading real-world data...")
|
||
|
||
# ── 1. Elevation ────────────────────────────────────────────────────
|
||
etopo_raw = _load_etopo()
|
||
|
||
# ETOPO 2022 N90W180 is already col 0 = 180°W — no shift needed
|
||
# Resample to grid
|
||
elevation_m = resample_to_grid(etopo_raw, GRID_H, GRID_W, order=1)
|
||
|
||
# Normalise to [0, 1]
|
||
elevation = normalize_01(elevation_m, EARTH_MIN_ELEV_M, EARTH_MAX_ELEV_M)
|
||
|
||
# Sea level: Earth's ocean fraction is ~0.71
|
||
sea_level = compute_sea_level(elevation, EARTH_OCEAN_FRACTION)
|
||
surface_water = elevation < sea_level
|
||
|
||
print(f" elevation: sea_level={sea_level:.4f}, "
|
||
f"ocean={surface_water.sum()}/{GRID_H*GRID_W} cells")
|
||
|
||
# ── 2. Temperature ──────────────────────────────────────────────────
|
||
temp_raw_K = _load_worldclim_temperature()
|
||
|
||
# WorldClim uses col 0 = 180°W — no shift needed
|
||
temperature_K = resample_to_grid(temp_raw_K, GRID_H, GRID_W, order=1)
|
||
|
||
# Fill ocean areas with latitude-dependent ocean temperature
|
||
v = np.linspace(0, 1, GRID_H, dtype=np.float32)
|
||
lat_abs = np.abs(v - 0.5) * 2.0 # 0 at equator, 1 at poles
|
||
ocean_temp = 301.0 - lat_abs[:, np.newaxis] * 30.0 # ~28°C equator, ~-2°C poles
|
||
temperature_K = np.where(surface_water, ocean_temp, temperature_K)
|
||
|
||
print(f" temperature: [{temperature_K.min():.0f}, {temperature_K.max():.0f}] K")
|
||
|
||
# ── 3. Moisture ─────────────────────────────────────────────────────
|
||
precip_raw = _load_worldclim_precipitation()
|
||
|
||
# WorldClim uses col 0 = 180°W — no shift needed
|
||
precip = resample_to_grid(precip_raw, GRID_H, GRID_W, order=1)
|
||
|
||
# Normalise to [0, 1] — global max is ~10000 mm/yr (tropical rainforest)
|
||
moisture = normalize_01(precip, 0.0, 6000.0)
|
||
# Ocean moisture = high (drives adjacent land humidity)
|
||
moisture = np.where(surface_water, 0.9, moisture)
|
||
|
||
print(f" moisture: [{moisture.min():.2f}, {moisture.max():.2f}]")
|
||
|
||
# ── 4. Biome classification ─────────────────────────────────────────
|
||
# Use the existing Whittaker table with real temperature and moisture
|
||
biome = compute_biome(body_def, elevation, sea_level, surface_water,
|
||
temperature_K, moisture)
|
||
n_biomes = len(np.unique(biome))
|
||
print(f" biomes: {n_biomes} classes present")
|
||
|
||
# ── 5. Hillshade ────────────────────────────────────────────────────
|
||
hillshade = compute_hillshade(elevation)
|
||
|
||
# ── 6. Rivers ───────────────────────────────────────────────────────
|
||
named_rivers = _load_rivers_geojson()
|
||
|
||
# Clip rivers: stop each path when it hits surface water.
|
||
# Rivers like the Amazon/Nile/Rhine otherwise draw through seas.
|
||
clipped = []
|
||
for name, path in named_rivers:
|
||
clipped_path = []
|
||
for r, c in path:
|
||
if surface_water[r, c]:
|
||
break
|
||
clipped_path.append((r, c))
|
||
if len(clipped_path) >= 2:
|
||
clipped.append((name, clipped_path))
|
||
|
||
n_orig = len(named_rivers)
|
||
n_kept = len(clipped)
|
||
print(f" rivers: {n_kept}/{n_orig} kept after water clipping")
|
||
named_rivers = clipped
|
||
rivers = [path for _, path in named_rivers]
|
||
|
||
# ── 7. Assemble ─────────────────────────────────────────────────────
|
||
terrain = assemble_terrain(
|
||
elevation=elevation,
|
||
temperature_K=temperature_K,
|
||
moisture=moisture,
|
||
biome=biome,
|
||
surface_water=surface_water,
|
||
hillshade=hillshade,
|
||
rivers=rivers,
|
||
sea_level=sea_level,
|
||
)
|
||
|
||
# Store river names for the marker overlay
|
||
terrain["_river_names"] = {i: name for i, (name, _) in enumerate(named_rivers)}
|
||
|
||
return terrain
|