1 files changed, 166 insertions, 0 deletions
diff --git a/mapgen/ocean.py b/mapgen/ocean.py
new file mode 100644
index 0000000..67dfaf7
--- /dev/null
+++ b/mapgen/ocean.py
@@ -0,0 +1,166 @@
+"""Ocean circulation for stage `climate`: wind-driven surface currents, sea-surface temperature, Ekman upwelling
+and a productivity index.
+
+Currents: the Stommel model on the sphere, r∇²ψ + βψ_x = k·curl(τ/ρH), solved for the stream function ψ over
+every cell; land is the same fluid with `land_friction`× the friction (Brinkman penalisation), so coasts block the
+flow and islands need no special treatment. u = k × ∇ψ. Western boundary currents come out ≈ r/β wide.
+"""
+from __future__ import annotations
+
+import numpy as np
+from scipy import sparse
+from scipy.sparse import linalg as splinalg
+
+from .graph import bicgstab_jacobi, distance_to, gradient, pmap, smooth_km
+from .grid import rowdot_at
+from .sphere import east_north
+
+DEFAULTS = {
+    "enabled": True, "friction_days": 5.0, "layer_m": 150.0, "stress_k": 1.0, "land_friction": 1000.0,
+    "rho_air": 1.2, "drag": 1.3e-3, "direct_max_cells": 500000,
+    "relax_days": 300.0, "kappa_m2s": 1000.0,
+    "ekman_min_lat": 3.0, "upwell_coast_km": 100.0,
+    "prod_upwell": 0.6, "prod_shelf": 0.4, "prod_mix": 0.3, "upwell_ref_m_yr": 100.0,
+    "prod_front": 0.4, "front_min": 0.5, "front_ref": 1.0,
+}
+RHO_W = 1025.0
+YEAR_S = 3.15576e7
+
+
+def omega(day_hours):
+    return 2.0 * np.pi / (float(day_hours) * 3600.0)
+
+
+def divergence(g, V):
+    """Graph divergence of a tangent field (V per km): (2/k_i) Σ_j V_j·t_ij / d_ij (exact for linear V on a hex
+    lattice; the V_i terms cancel)."""
+    w = rowdot_at(V, g.dst, g.edge_tangents) / g.edge_km
+    return np.bincount(g.src, weights=w, minlength=g.n) * (2.0 / g.counts)
+
+
+def _edge_operator(g, w):
+    """Sparse operator f ↦ Σ_j w_ij (f_j − f_i)."""
+    A = sparse.csr_matrix((w, (g.src, g.dst)), shape=(g.n, g.n))
+    return A - sparse.diags(np.asarray(A.sum(axis=1)).ravel())
+
+
+def wind_stress(wind, P):
+    """Bulk formula τ = k·ρ_air·C_d·|w|·w (N/m²)."""
+    wind = np.asarray(wind, dtype=np.float64)
+    return P["stress_k"] * P["rho_air"] * P["drag"] * np.linalg.norm(wind, axis=1, keepdims=True) * wind
+
+
+def streamfunction(g, ocean, wind, P, day_hours):
+    """ψ (m²/s) of the wind-driven surface flow; one direct sparse solve over every cell."""
+    Om = omega(day_hours)
+    R = g.radius_km * 1e3
+    beta = 2.0 * Om * np.cos(np.radians(g.lat)) / R                    # 1/(m·s)
+    beta30 = 2.0 * Om * np.cos(np.radians(30.0)) / R
+    r = max(1.0 / (P["friction_days"] * 86400.0), beta30 * g.spacing_km * 1e3)   # boundary layer ≥ one cell
+    F = np.where(ocean[:, None], wind_stress(wind, P), 0.0) / (RHO_W * P["layer_m"])
+    curl = divergence(g, np.cross(F, g.xyz)) * 1e-3                    # k·curl F = ∇·(F×k), 1/s²
+    fr = np.where(ocean, 1.0, P["land_friction"])
+    fe = 2.0 / (1.0 / fr[g.src] + 1.0 / fr[g.dst])                    # harmonic mean: coasts count as land
+    L = _edge_operator(g, 4.0 / (g.counts[g.src] * g.edge_km ** 2) * fe)          # ∇·(fr∇), per km²
+    e, _ = east_north(g.xyz)
+    Dx = _edge_operator(g, (2.0 / g.counts[g.src]) * rowdot_at(e, g.src, g.edge_tangents) / g.edge_km)
+    A = (L + sparse.diags(beta / r * 1e3) @ Dx).tolil()                # β/r·1e3: per km
+    b = 1e6 * curl / r
+    # gauge: ψ is defined up to a constant; pin it deep inland, where land friction soaks up the solve's
+    # compatibility residual (pinned at sea it would leave a point vortex there)
+    k = int(np.argmax(distance_to(g, ocean))) if (~ocean).any() else 0
+    A[k, :] = 0.0
+    A[k, k] = 1.0
+    b[k] = 0.0
+    return splinalg.spsolve(A.tocsc(), b)
+
+
+def velocity(g, psi):
+    """u = k × ∇ψ (m/s), tangent 3-vectors."""
+    return np.cross(g.xyz, gradient(g, psi) * 1e-3)
+
+
+def _coarse_currents(g, ocean, wind, P, day_hours):
+    """Grids too big for a direct solve: solve on the parent H3 resolution, carry u down, smooth."""
+    from .grid import build_grid, cell_parents
+    cg = build_grid(g.res - 1, g.radius_km)
+    parent = np.searchsorted(cg.ids, cell_parents(g.ids, g.res - 1))
+    cnt = np.maximum(np.bincount(parent, minlength=cg.n), 1)
+    c_ocean = np.bincount(parent, weights=ocean.astype(float), minlength=cg.n) / cnt > 0.5
+    c_wind = np.stack([np.bincount(parent, weights=wind[:, k], minlength=cg.n) / cnt for k in range(3)], axis=1)
+    cu = currents(cg, c_ocean, c_wind, P, day_hours)
+    u = np.stack(pmap(lambda k: smooth_km(g, cu[parent, k], cg.spacing_km), range(3)), axis=1)
+    return u - np.sum(u * g.xyz, axis=1, keepdims=True) * g.xyz       # back into the tangent plane
+
+
+def currents(g, ocean, wind, P, day_hours):
+    """Surface current (n,3) m/s; 0 on land."""
+    ocean = np.asarray(ocean, bool)
+    wind = np.asarray(wind, dtype=np.float64)
+    if not ocean.any():
+        return np.zeros((g.n, 3))
+    if g.n > P["direct_max_cells"]:
+        u = _coarse_currents(g, ocean, wind, P, day_hours)
+    else:
+        u = velocity(g, streamfunction(g, ocean, wind, P, day_hours))
+    return np.where(ocean[:, None], u, 0.0)
+
+
+def _masked_edges(g, mask, w):
+    return _edge_operator(g, np.where(mask[g.src] & mask[g.dst], w, 0.0))
+
+
+def sst(g, ocean, u, T_eq, P):
+    """Annual-mean sea-surface temperature (°C): steady u·∇T − κ∇²T = λ(T_eq − T) over the ocean (upwind
+    advection, no flux into land); land keeps T_eq."""
+    ocean = np.asarray(ocean, bool)
+    T_eq = np.asarray(T_eq, dtype=np.float64)
+    lam = 1.0 / (P["relax_days"] * 86400.0)
+    up = (4.0 / g.counts[g.src]) * np.maximum(-rowdot_at(u, g.src, g.edge_tangents), 0.0) / (g.edge_km * 1e3)
+    adv = -_masked_edges(g, ocean, up)                                 # Σ_upstream c_ij (T_i − T_j), 1/s
+    dif = _masked_edges(g, ocean, 4.0 / (g.counts[g.src] * g.edge_km ** 2)) * (P["kappa_m2s"] * 1e-6)
+    A = (adv - dif) / lam + sparse.identity(g.n)
+    A = (sparse.diags(ocean.astype(float)) @ A + sparse.diags((~ocean).astype(float))).tocsr()
+    T, info = bicgstab_jacobi(A, T_eq, T_eq.copy(), 1.0 / A.diagonal(), 1e-9, 5000)
+    if info != 0:
+        raise ValueError(f"sst: solver did not converge (info={info})")
+    return T
+
+
+def upwelling(g, ocean, wind, P, day_hours):
+    """Ekman pumping (m/yr, + up): w = ∇·M, M = τ×k/(ρf), |f| floored at `ekman_min_lat`. Transport pointing off a
+    coast leaves the coast cell (land carries none), so coastal upwelling needs no separate rule. Smoothed over
+    `upwell_coast_km`; 0 on land."""
+    ocean = np.asarray(ocean, bool)
+    Om = omega(day_hours)
+    fmin = 2.0 * Om * np.sin(np.radians(P["ekman_min_lat"]))
+    f = np.where(g.lat >= 0, 1.0, -1.0) * np.maximum(np.abs(2.0 * Om * np.sin(np.radians(g.lat))), fmin)
+    tau = np.where(ocean[:, None], wind_stress(wind, P), 0.0)
+    M = np.cross(tau, g.xyz) / (RHO_W * f[:, None])                    # m²/s
+    w = np.where(ocean, divergence(g, M) * 1e-3 * YEAR_S, 0.0)
+    return np.where(ocean, smooth_km(g, w, P["upwell_coast_km"]), 0.0)
+
+
+def sst_front(g, ocean, sst_c):
+    """|∇SST| over the sea (°C per 100 km); edges to land count as flat, so coasts are no front."""
+    sst_c = np.asarray(sst_c, dtype=np.float64)
+    both = ocean[g.src] & ocean[g.dst]
+    w = np.where(both, sst_c[g.dst] - sst_c[g.src], 0.0) / g.edge_km
+    grad = np.stack([np.bincount(g.src, weights=w * g.edge_tangents[:, c], minlength=g.n) for c in range(3)], axis=1)
+    return np.where(ocean, np.linalg.norm(grad, axis=1) * (2.0 / g.counts) * 100.0, 0.0)
+
+
+def productivity(g, ocean, w, z, t_range, sst_c, P):
+    """0–1 sea productivity: upwelling (saturating), shallow shelf, winter mixing and SST fronts (where warm and cold
+    currents meet, e.g. a Brazil–Malvinas confluence; gradients under `front_min` °C/100 km add nothing), dimmed
+    toward the poles."""
+    ocean = np.asarray(ocean, bool)
+    depth = np.maximum(-np.asarray(z, dtype=np.float64), 0.0)
+    x = np.maximum(w, 0.0) / P["upwell_ref_m_yr"]
+    shelf = np.clip((1000.0 - depth) / 800.0, 0.0, 1.0)
+    mix = np.clip(np.asarray(t_range) / 20.0, 0.0, 1.0) * np.clip((20.0 - np.asarray(sst_c)) / 20.0, 0.0, 1.0)
+    light = 0.3 + 0.7 * np.cos(np.radians(g.lat))
+    xf = np.maximum(sst_front(g, ocean, sst_c) - P["front_min"], 0.0) / P["front_ref"]
+    p = (P["prod_upwell"] * x / (1.0 + x) + P["prod_shelf"] * shelf + P["prod_mix"] * mix
+         + P["prod_front"] * xf / (1.0 + xf)) * light
+    return np.where(ocean, np.clip(p, 0.0, 1.0), 0.0)