1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
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)
|
