02 — Helmholtz Decomposition, QG Omega & PV Tendency

End-to-end demonstration on ERA5 January 2025 data (Northern Hemisphere, 1.5°, 9 levels).

Pipeline (matches core/step2_compute_tendency_terms_blocking.py):

1. Load raw ERA5 fields on full NH grid (61×240)

2. Helmholtz decomposition on total wind (u, v) on full NH (spherical FFT Laplacian, requires full longitude ring)

3. PV spatial derivatives (∂q/∂x, ∂q/∂y, ∂q/∂p) on full NH (periodic zonal boundaries)

4. Extract event-centred patch (±21° lat × ±36° lon → 29×49) at California / East Pacific

5. QG omega on full NH: SP19 (scaling) and LOG20 (SIP) for ω_adia (A+B terms)

6. Three omega components: ω_adia, ω_qg_diabatic, ω_lhr_moist (LHR via C term)

7. Diabatic / adiabatic divergent wind via independent Poisson inversions using ω_lhr_moist on patch

8. PV tendency: LHR-moist indirect contribution vs total divergent outflow on patch (total PV gradients)

9. Save NPZ of all computed fields

import numpy as np
import xarray as xr
import matplotlib.pyplot as plt
import cartopy.io.shapereader as shpreader
from shapely.geometry import box as shapely_box
import time, gc

from pvtend import (
    NHGrid, EventPatch,
    R_EARTH, OMEGA_E, R_DRY, KAPPA, G0, SP19_DRY_FRACTION,
    helmholtz_decomposition_3d, solve_qg_omega_sip, decompose_omega,
    ddx, ddy, ddp,
)
from pvtend.helmholtz import gradient, solve_poisson_spherical_fft, laplacian_spherical_fft
from pvtend.derivatives import gradient_periodic
from pvtend.moist_dry import solve_chi_from_omega
from pvtend.omega import _compute_diabatic_rhs_emanuel


# ── Coastline overlay in event-relative coordinates ──
def overlay_coastlines(ax, centre_lat, centre_lon, xlim, ylim,
                       lw=0.5, color="k", alpha=0.6, resolution="50m"):
    """Overlay coastlines in event-relative coordinates."""
    shp = shpreader.natural_earth(resolution, "physical", "coastline")
    reader = shpreader.Reader(shp)
    clip = shapely_box(centre_lon + xlim[0], centre_lat + ylim[0],
                       centre_lon + xlim[1], centre_lat + ylim[1])
    for geom in reader.geometries():
        clipped = geom.intersection(clip)
        if clipped.is_empty:
            continue
        parts = clipped.geoms if hasattr(clipped, "geoms") else [clipped]
        for part in parts:
            if hasattr(part, "coords"):
                coords = np.array(part.coords)
                ax.plot(coords[:, 0] - centre_lon,
                        coords[:, 1] - centre_lat,
                        color=color, lw=lw, alpha=alpha)


%matplotlib inline
plt.rcParams.update({"figure.dpi": 120, "font.size": 10})

1 Load ERA5 January 2025 & set up grid

All data on full NH grid (61 lat × 240 lon, 1.5°). NHGrid and EventPatch are configured for later patch extraction.

# ══════════════════════════════════════════════════════════════
# Configuration (change freely)
# ══════════════════════════════════════════════════════════════
DATA_DIR    = "era5_jan2025"
TIME_IDX    = 252               # snapshot index (≈ 2025-01-15 12Z)
PLOT_LEV    = 300               # hPa for all map plots
CENTER_LAT  = 40.0              # event centre latitude [°N]
CENTER_LON  = -120.0            # event centre longitude [°E]
QG_METHOD   = "log20"            # "sp19" (scaling) or "log20" (SIP solver)

# ══════════════════════════════════════════════════════════════
# Load all variables (full NH)
# ══════════════════════════════════════════════════════════════
print("Loading ERA5 data …")
ds_u  = xr.open_dataset(f"{DATA_DIR}/era5_u_2025_01.nc")
ds_v  = xr.open_dataset(f"{DATA_DIR}/era5_v_2025_01.nc")
ds_t  = xr.open_dataset(f"{DATA_DIR}/era5_t_2025_01.nc")
ds_z  = xr.open_dataset(f"{DATA_DIR}/era5_z_2025_01.nc")
ds_w  = xr.open_dataset(f"{DATA_DIR}/era5_w_2025_01.nc")
ds_pv = xr.open_dataset(f"{DATA_DIR}/era5_pv_2025_01.nc")

# Coordinates: flip latitude to ascending, pressure to Pa
lat_raw  = ds_u["latitude"].values          # 90 → 0 (descending)
lon      = ds_u["longitude"].values         # -180 → 178.5
levs_hpa = ds_u["pressure_level"].values    # [1000, 850, …, 100]
times    = ds_u["valid_time"].values

flip     = lat_raw[0] > lat_raw[-1]
lat      = lat_raw[::-1] if flip else lat_raw       # ascending 0 → 90
plevs_pa = (levs_hpa * 100.0).astype(np.float64)    # Pa, ascending
nlev, nlat, nlon = len(levs_hpa), len(lat), len(lon)

def _load(ds, var):
    """Load 4-D array → (nt, nlev, nlat, nlon) with ascending lat."""
    arr = ds[var].values.astype(np.float32)
    return arr[:, :, ::-1] if flip else arr

u_all  = _load(ds_u,  "u")
v_all  = _load(ds_v,  "v")
t_all  = _load(ds_t,  "t")
z_all  = _load(ds_z,  "z")        # geopotential Φ [m²/s²]
w_all  = _load(ds_w,  "w")        # omega ω [Pa/s]
pv_all = _load(ds_pv, "pv")       # potential vorticity [PVU]

for ds in [ds_u, ds_v, ds_t, ds_z, ds_w, ds_pv]:
    ds.close()

# ── Snapshot ──
ts_str  = str(times[TIME_IDX])[:16]
u_snap, v_snap = u_all[TIME_IDX], v_all[TIME_IDX]
t_snap, z_snap = t_all[TIME_IDX], z_all[TIME_IDX]
w_snap, pv_snap = w_all[TIME_IDX], pv_all[TIME_IDX]

del u_all, v_all, t_all, z_all, w_all, pv_all
gc.collect()

# ══════════════════════════════════════════════════════════════
# NH grid & EventPatch setup
# ══════════════════════════════════════════════════════════════
grid    = NHGrid(lat=lat, lon=lon)          # NHGrid works with ascending too
patcher = EventPatch(grid)                  # ±21° lat × ±36° lon

ilev = np.argmin(np.abs(levs_hpa - PLOT_LEV))

# ── Grid spacing (full NH, ascending lat) ──
lat_rad = np.deg2rad(lat)
dlat = float(np.abs(lat[1] - lat[0]))
dlon = float(np.abs(lon[1] - lon[0]))
dy     = np.deg2rad(dlat) * R_EARTH
dx_arr = np.deg2rad(dlon) * R_EARTH * np.cos(lat_rad)
dx_arr = np.maximum(dx_arr, dy * 0.1)

# ── Coriolis (full NH) ──
f_arr   = 2.0 * OMEGA_E * np.sin(lat_rad)
f_clamp = 2.0 * OMEGA_E * np.sin(np.deg2rad(5.0))
f_arr   = np.where(np.abs(f_arr) < f_clamp,
                   np.sign(f_arr + 1e-30) * f_clamp, f_arr)

print(f"Full NH grid: {nlat}×{nlon}, {nlev} levels ({levs_hpa} hPa)")
print(f"Snapshot: {ts_str}  |  PLOT_LEV = {PLOT_LEV} hPa")
print(f"Event centre: ({CENTER_LAT}°N, {CENTER_LON}°E)")
print(f"Patch shape: {patcher.patch_shape}")
print(f"|u_snap| max = {np.abs(u_snap).max():.1f} m/s  |  |w_snap| max = {np.abs(w_snap).max():.4f} Pa/s")

Loading ERA5 data …
Full NH grid: 61×240, 9 levels ([1000.  850.  700.  500.  400.  300.  250.  200.  100.] hPa)
Snapshot: 2025-01-15T12:00  |  PLOT_LEV = 300 hPa
Event centre: (40.0°N, -120.0°E)
Patch shape: (29, 49)
|u_snap| max = 101.2 m/s  |  |w_snap| max = 6.3439 Pa/s

2 Helmholtz decomposition on full NH grid

Splits total wind (u, v) at each level into rotational (ψ), divergent (χ), and harmonic components using the spherical Laplacian Poisson solver (conservative form, following xinvert/MiniUFO). Must be done on the full NH grid (full longitude ring) for spectral accuracy of the periodic zonal BCs.

%%time
# Full NH Helmholtz on TOTAL wind (spherical Laplacian): (nlev, 61, 240) → dict of same shape
helm = helmholtz_decomposition_3d(u_snap, v_snap, lat, lon)
print("Helmholtz output keys:", sorted(helm.keys()))

print(f"\nAt {PLOT_LEV} hPa (level index {ilev}):")
for comp in ["u_rot", "v_rot", "u_div", "v_div", "u_har", "v_har"]:
    print(f"  {comp:6s}  range = [{helm[comp][ilev].min():.3f}, {helm[comp][ilev].max():.3f}] m/s")
print(f"  psi    range = [{helm['psi'][ilev].min():.3e}, {helm['psi'][ilev].max():.3e}] m²/s")
print(f"  chi    range = [{helm['chi'][ilev].min():.3e}, {helm['chi'][ilev].max():.3e}] m²/s")

Helmholtz output keys: ['chi', 'divergence', 'psi', 'u_div', 'u_har', 'u_rot', 'v_div', 'v_har', 'v_rot', 'vorticity']

At 300 hPa (level index 5):
  u_rot   range = [-33.842, 95.780] m/s
  v_rot   range = [-57.907, 63.916] m/s
  u_div   range = [-14.518, 16.728] m/s
  v_div   range = [-16.977, 13.782] m/s
  u_har   range = [-22.153, 11.644] m/s
  v_har   range = [-14.617, 17.084] m/s
  psi    range = [-1.125e+08, 8.311e+07] m²/s
  chi    range = [-1.639e+07, 2.641e+07] m²/s
CPU times: user 513 ms, sys: 4.11 ms, total: 517 ms
Wall time: 516 ms

3 PV spatial derivatives on full NH grid

Compute \(\partial q/\partial x\), \(\partial q/\partial y\), \(\partial q/\partial p\) on the full NH domain using periodic zonal boundaries. Uses the total PV field (pv_snap).

# ── PV gradients on full NH (periodic lon) ──
# Total PV derivatives (bar = total field treated as background)
dpv_bar_dx = np.zeros_like(pv_snap)
dpv_bar_dy = np.zeros_like(pv_snap)
for k in range(nlev):
    dpv_bar_dx[k], dpv_bar_dy[k] = gradient_periodic(pv_snap[k], dx_arr, dy)

# ∂q_bar/∂p  (centred finite differences)
dpv_bar_dp = np.zeros_like(pv_snap)
for k in range(1, nlev - 1):
    dpv_bar_dp[k] = (pv_snap[k+1] - pv_snap[k-1]) / (plevs_pa[k+1] - plevs_pa[k-1])
dp_fd = np.diff(plevs_pa)
dpv_bar_dp[0]  = (pv_snap[1]  - pv_snap[0])  / dp_fd[0]
dpv_bar_dp[-1] = (pv_snap[-1] - pv_snap[-2]) / dp_fd[-1]

print("Full-NH PV derivatives computed:")
print(f"  |∂q/∂x| max = {np.abs(dpv_bar_dx[ilev]).max():.4e}")
print(f"  |∂q/∂y| max = {np.abs(dpv_bar_dy[ilev]).max():.4e}")
print(f"  |∂q/∂p| max = {np.abs(dpv_bar_dp[ilev]).max():.4e}")

Full-NH PV derivatives computed:
  |∂q/∂x| max = 4.0278e-11
  |∂q/∂y| max = 2.5321e-11
  |∂q/∂p| max = 5.8584e-10

4 Extract event-centred patch

Centre the ±21° lat × ±36° lon patch on California / East Pacific. Crop all full-NH fields (Helmholtz winds, PV derivatives, meteorological vars) to this patch. Compute geostrophic wind on the patch (matching core/step2 pipeline).

# ── Find patch centre index ──
i_lat, i_lon, ok = patcher.nearest_idx(CENTER_LAT, CENTER_LON)
assert ok, f"Patch does not fit at ({CENTER_LAT}, {CENTER_LON})"

# ── Relative grid ──
Y_rel, X_rel = patcher.relative_grid()
x_coords = X_rel[0, :]    # 1D relative longitude
y_coords = Y_rel[:, 0]    # 1D relative latitude

# ── Coastline overlay kwargs (reused in all plots) ──
coast_kw = dict(centre_lat=CENTER_LAT, centre_lon=CENTER_LON,
                xlim=(x_coords.min(), x_coords.max()),
                ylim=(y_coords.min(), y_coords.max()))

# ── Extract meteorological fields to patch ──
# Total (snapshot) fields
p_u_snap = patcher.extract(u_snap,  i_lat, i_lon)
p_v_snap = patcher.extract(v_snap,  i_lat, i_lon)
p_t_snap = patcher.extract(t_snap,  i_lat, i_lon)
p_z_snap = patcher.extract(z_snap,  i_lat, i_lon)
p_w_snap = patcher.extract(w_snap,  i_lat, i_lon)

# ── Extract Helmholtz components to patch ──
p_u_rot = patcher.extract(helm["u_rot"], i_lat, i_lon)
p_v_rot = patcher.extract(helm["v_rot"], i_lat, i_lon)
p_u_div = patcher.extract(helm["u_div"], i_lat, i_lon)
p_v_div = patcher.extract(helm["v_div"], i_lat, i_lon)
p_u_har = patcher.extract(helm["u_har"], i_lat, i_lon)
p_v_har = patcher.extract(helm["v_har"], i_lat, i_lon)
p_psi   = patcher.extract(helm["psi"],   i_lat, i_lon)
p_chi   = patcher.extract(helm["chi"],   i_lat, i_lon)

# ── Extract PV derivatives to patch ──
p_dpv_bar_dx = patcher.extract(dpv_bar_dx, i_lat, i_lon)
p_dpv_bar_dy = patcher.extract(dpv_bar_dy, i_lat, i_lon)
p_dpv_bar_dp = patcher.extract(dpv_bar_dp, i_lat, i_lon)

# ── Grid spacing on the patch (ascending lat) ──
patch_lat = lat[i_lat] + y_coords  # approximate absolute lats
dx_patch = np.deg2rad(grid.dlon) * R_EARTH * np.cos(np.deg2rad(patch_lat))
dx_patch = np.maximum(dx_patch, grid.dy * 0.01)
dy_patch = grid.dy

# ── Coriolis on patch ──
f_patch   = 2.0 * OMEGA_E * np.sin(np.deg2rad(patch_lat))
f_clamp_p = 2.0 * OMEGA_E * np.sin(np.deg2rad(5.0))
f_patch   = np.where(np.abs(f_patch) < f_clamp_p,
                     np.sign(f_patch + 1e-30) * f_clamp_p, f_patch)

# ── Geostrophic wind on patch (from geopotential Φ) ──
def _geostrophic_patch(phi_3d):
    """(u_g, v_g) on the patch grid."""
    ug = np.zeros_like(phi_3d)
    vg = np.zeros_like(phi_3d)
    for k in range(phi_3d.shape[0]):
        dphi_dx = ddx(phi_3d[k], dx_patch, periodic=False)
        dphi_dy = ddy(phi_3d[k], dy_patch)
        ug[k] = -(1.0 / f_patch[:, None]) * dphi_dy
        vg[k] =  (1.0 / f_patch[:, None]) * dphi_dx
    return ug, vg

p_ug_snap, p_vg_snap = _geostrophic_patch(p_z_snap)

# Patch longitude array (absolute, for solver functions)
patch_lon = CENTER_LON + x_coords

print(f"Patch shape   : {p_u_snap.shape}")

Patch shape   : (9, 29, 49)

5 Helmholtz wind decomposition on patch (2×2 panel)

All components are extracted from the full-NH Helmholtz (spherical Laplacian) to the event patch. The decomposition satisfies \(\mathbf{u} = \mathbf{u}_\psi + \mathbf{u}_\chi + \mathbf{u}_h\) (tex Eq. 1), where the Poisson equations \(\nabla^2\psi = \zeta\) and \(\nabla^2\chi = \delta\) use the full spherical Laplacian in conservative form.

2×2 panel: total wind | rotational + ψ | divergent + χ | harmonic + speed

skip = 2                          # quiver stride for patch

u_tot_p = p_u_snap[ilev];         v_tot_p = p_v_snap[ilev]
u_rot_p = p_u_rot[ilev];          v_rot_p = p_v_rot[ilev]
u_div_p = p_u_div[ilev];          v_div_p = p_v_div[ilev]
u_har_p = p_u_har[ilev];          v_har_p = p_v_har[ilev]
psi_p   = p_psi[ilev];            chi_p   = p_chi[ilev]
har_speed = np.sqrt(u_har_p**2 + v_har_p**2)

fig, axes = plt.subplots(2, 2, figsize=(14, 11))
kw = dict(cmap="RdBu_r", extend="both")

# ── (0,0) Total wind — speed contourf + quiver ──
ax = axes[0, 0]
speed = np.sqrt(u_tot_p**2 + v_tot_p**2)
cf = ax.contourf(x_coords, y_coords, speed, levels=20, cmap="YlOrRd", extend="max")
q0 = ax.quiver(x_coords[::skip], y_coords[::skip],
               u_tot_p[::skip, ::skip], v_tot_p[::skip, ::skip],
               scale=500, width=0.003, color="k")
ax.set_title("Total u, v  (speed)", fontsize=12)
plt.colorbar(cf, ax=ax, shrink=0.7, label="m/s")

# ── (0,1) Rotational + streamfunction ψ ──
ax = axes[0, 1]
vm = np.nanpercentile(np.abs(psi_p), 98)
cf = ax.contourf(x_coords, y_coords, psi_p,
                 levels=np.linspace(-vm, vm, 21), **kw)
ax.quiver(x_coords[::skip], y_coords[::skip],
          u_rot_p[::skip, ::skip], v_rot_p[::skip, ::skip],
          scale=500, width=0.003, color="k")
ax.set_title("Rotational + ψ", fontsize=12)
plt.colorbar(cf, ax=ax, shrink=0.7, label="m²/s")
# Row 0 arrow legend (total/rotational scale)
ax.quiverkey(q0, 0.90, 1.05, 10, "10 m/s", labelpos="E", fontproperties=dict(size=9))

# ── (1,0) Divergent + velocity potential χ ──
ax = axes[1, 0]
vm = np.nanpercentile(np.abs(chi_p), 98)
cf = ax.contourf(x_coords, y_coords, chi_p,
                 levels=np.linspace(-vm, vm, 21), **kw)
q1 = ax.quiver(x_coords[::skip], y_coords[::skip],
               u_div_p[::skip, ::skip], v_div_p[::skip, ::skip],
               scale=100, width=0.003, color="k")
ax.set_title("Divergent + χ", fontsize=12)
plt.colorbar(cf, ax=ax, shrink=0.7, label="m²/s")

# ── (1,1) Harmonic + speed ──
ax = axes[1, 1]
cf = ax.contourf(x_coords, y_coords, har_speed, levels=20, cmap="YlOrRd", extend="max")
ax.quiver(x_coords[::skip], y_coords[::skip],
          u_har_p[::skip, ::skip], v_har_p[::skip, ::skip],
          scale=100, width=0.003, color="k")
ax.set_title("Harmonic + |speed|", fontsize=12)
plt.colorbar(cf, ax=ax, shrink=0.7, label="m/s")
# Row 1 arrow legend (divergent/harmonic scale)
ax.quiverkey(q1, 0.90, 1.05, 2, "2 m/s", labelpos="E", fontproperties=dict(size=9))

for ax in axes.flat:
    overlay_coastlines(ax, **coast_kw)
    ax.set_xlabel("Relative longitude [°]")
    ax.set_ylabel("Relative latitude [°]")
    ax.set_aspect("equal")

fig.suptitle(f"Helmholtz decomposition (total wind) @ {PLOT_LEV} hPa — {ts_str}\n"
             f"Centre ({CENTER_LAT}°N, {CENTER_LON}°E)",
             fontsize=13, y=0.95)
fig.tight_layout()
plt.show(); plt.close("all"); gc.collect()

24

6 QG omega — two-method comparison + LHR moist decomposition

Domain choices (matching production pipeline tendency.py / _sip_on_nh()):

Method	Domain	Reason
SP19	Patch (29×49)	Trivial scaling — ω_adia = SP19_DRY_FRACTION × ω_total
LOG20 (SIP)	Full NH (61×240) → extract patch	3-D SIP solver with periodic lon BCs on full ring

Three omega components (LOG20):

Component	Definition	Physical meaning
ω_adia	SIP solve with A+B terms only	Dynamic (adiabatic) QG vertical motion
ω_qg_diabatic	ω_total − ω_adia	Diabatic residual (everything not explained by adiabatic QG)
ω_lhr_moist	ω_ABC − ω_AB	LHR-driven vertical motion (diabatic C term)

The LHR moist component uses θ̇_LHR ≈ (L_v/c_p)(θ/T) max(0, −ω ∂q_s/∂p) as the LHR approximation, fed into C_em = −(κ/p)∇²(J_em) where J_em = c_p θ̇_LHR T/θ.

2×3 panel: rows = methods, columns = [ω_total, ω_adia, ω_dia_res = ω_total − ω_adia]

%%time
# ── Full-NH geostrophic wind (for LOG20 — needs periodic longitude) ──
def _geostrophic_nh(phi_3d):
    """(u_g, v_g) on the full NH grid using periodic zonal gradients."""
    ug = np.zeros_like(phi_3d)
    vg = np.zeros_like(phi_3d)
    for k in range(phi_3d.shape[0]):
        dphi_dx, dphi_dy = gradient_periodic(phi_3d[k], dx_arr, dy)
        ug[k] = -(1.0 / f_arr[:, None]) * dphi_dy
        vg[k] =  (1.0 / f_arr[:, None]) * dphi_dx
    return ug, vg

ug_snap_nh, vg_snap_nh = _geostrophic_nh(z_snap)
print(f"Full-NH geostrophic wind: max |ug|={np.abs(ug_snap_nh).max():.1f} m/s")

# ── SP19 (total, on patch — trivial adiabatic-fraction scaling) ──
t0 = time.perf_counter()
wd_sp19 = SP19_DRY_FRACTION * p_w_snap
t_sp19 = time.perf_counter() - t0

# ── LOG20 / SIP (total, on FULL NH → extract patch) — adiabatic: A+B only ──
t0 = time.perf_counter()
wd_log20_nh, info_log20 = solve_qg_omega_sip(
    ug_snap_nh, vg_snap_nh, t_snap,
    lat, lon, plevs_pa,
    center_lat=CENTER_LAT,
    omega_b=w_snap,
)
wd_log20 = patcher.extract(wd_log20_nh, i_lat, i_lon)
t_log20 = time.perf_counter() - t0

# ── Emanuel LHR approximation → w_lhr_moist (full NH) ──
# θ = T (p₀/p)^κ
kappa_val = R_DRY / 1004.0
theta_nh = t_snap * (1e5 / plevs_pa[:, None, None]) ** kappa_val

# Saturation specific humidity:  q_s = 0.622 e_s / (p - e_s)
# Bolton:  e_s [Pa] = 611.2 exp(17.67 (T-273.15) / (T - 29.65))
T_celsius = t_snap - 273.15
e_sat = 611.2 * np.exp(17.67 * T_celsius / (t_snap - 29.65))
q_sat = 0.622 * e_sat / (plevs_pa[:, None, None] - e_sat)
q_sat = np.clip(q_sat, 0, 0.1)

# ∂q_s/∂p via centred finite differences
dqs_dp = ddp(q_sat, plevs_pa)

# θ̇_LHR ≈ (L_v / c_p)(θ / T) max(0, -ω ∂q_s/∂p)
L_v = 2.501e6  # latent heat of vaporisation [J/kg]
c_p = 1004.0
condensation_rate = np.maximum(0.0, -w_snap * dqs_dp)  # >0 for ascending
theta_dot_lhr_nh = (L_v / c_p) * (theta_nh / t_snap) * condensation_rate
print(f"θ̇_LHR max = {theta_dot_lhr_nh.max():.4e} K/s, "
      f"nonzero fraction = {(theta_dot_lhr_nh > 0).mean():.1%}")

# C_em = -(κ/p) ∇²_spherical(J_em),  J_em = c_p θ̇_LHR T/θ
t0 = time.perf_counter()
C_em_nh = _compute_diabatic_rhs_emanuel(
    theta_dot_lhr_nh, t_snap, theta_nh, plevs_pa, lat, lon,
)

# Solve A+B+C_em → ω_ABC
w_abc_nh, info_abc = solve_qg_omega_sip(
    ug_snap_nh, vg_snap_nh, t_snap,
    lat, lon, plevs_pa,
    center_lat=CENTER_LAT,
    omega_b=w_snap,
    rhs_c=C_em_nh,
)
w_abc = patcher.extract(w_abc_nh, i_lat, i_lon)
t_em = time.perf_counter() - t0

# ── Three omega components ──
# w_qg_diabatic = total - adia (diabatic residual)
wm_sp19     = p_w_snap - wd_sp19
wm_log20    = p_w_snap - wd_log20     # = w_qg_diabatic for LOG20
w_qg_diabatic  = wm_log20                # alias
# w_lhr_moist = ω_ABC - ω_AB (LHR diabatic contribution, by linearity)
w_em_moist  = w_abc - wd_log20

# Extract LHR diagnostics to patch
theta_dot_lhr_patch = patcher.extract(theta_dot_lhr_nh, i_lat, i_lon)
C_em_patch = patcher.extract(C_em_nh, i_lat, i_lon)

print(f"\nSP19:  {t_sp19:.4f}s  |  LOG20 (SIP, full NH): {t_log20:.2f}s  |  Emanuel: {t_em:.2f}s")
print(f"LOG20 SIP: {info_log20['iters']} iters, "
      f"residual={info_log20['final_residual']:.2e}, "
      f"numba={info_log20['numba']}, terms={info_log20['terms']}")
print(f"ABC   SIP: {info_abc['iters']} iters, "
      f"residual={info_abc['final_residual']:.2e}, terms={info_abc['terms']}")
W_PLOT_LEV = 500
ilev_w = np.argmin(np.abs(levs_hpa - W_PLOT_LEV))
print(f"|wd_log20| max = {np.abs(wd_log20[ilev_w]).max():.4f} Pa/s")
print(f"|w_qg_diabatic| max = {np.abs(w_qg_diabatic[ilev_w]).max():.4f} Pa/s")
print(f"|w_em_moist| max = {np.abs(w_em_moist[ilev_w]).max():.4f} Pa/s")

# ── 2×3 plot: rows = [SP19, LOG20], cols = [w_total, w_adia, w_dia_res] ──
methods = [("SP19 (patch, scaling)",  p_w_snap, wd_sp19,  wm_sp19),
           ("LOG20 / SIP (full NH)",  p_w_snap, wd_log20, wm_log20)]
col_labels = ["ω_total", "ω_adia", "ω_dia_res (diabatic residual)"]

fig, axes = plt.subplots(2, 3, figsize=(18, 10))

for row, (mname, wt, wd, wm) in enumerate(methods):
    row_fields = [wt[ilev_w], wd[ilev_w], wm[ilev_w]]
    vm_row = max(np.nanpercentile(np.abs(f), 98) for f in row_fields)
    if vm_row < 1e-15:
        vm_row = 1.0
    for col, (clabel, field) in enumerate(zip(col_labels, [wt, wd, wm])):
        ax = axes[row, col]
        f2d = field[ilev_w]
        cf = ax.contourf(x_coords, y_coords, f2d,
                         levels=np.linspace(-vm_row, vm_row, 21),
                         cmap="RdBu_r", extend="both")
        overlay_coastlines(ax, **coast_kw)
        ax.set_aspect("equal")
        ax.set_title(f"{mname}: {clabel}", fontsize=11)
        plt.colorbar(cf, ax=ax, shrink=0.7, label="Pa/s")

for ax in axes.flat:
    ax.set_xlabel("Rel. lon [°]"); ax.set_ylabel("Rel. lat [°]")

fig.suptitle(f"Total omega: 2-method comparison @ {W_PLOT_LEV} hPa — {ts_str}\n"
             f"LOG20/SIP on full NH (61×240), extracted to patch",
             fontsize=14, y=0.93)
fig.tight_layout()
plt.show(); plt.close("all"); gc.collect()

Full-NH geostrophic wind: max |ug|=235.3 m/s
θ̇_LHR max = 4.2244e-03 K/s, nonzero fraction = 46.7%

SP19:  0.0000s  |  LOG20 (SIP, full NH): 1.57s  |  Emanuel: 0.35s
LOG20 SIP: 300 iters, residual=6.88e-12, numba=True, terms=AB
ABC   SIP: 300 iters, residual=4.57e-12, terms=ABC
|wd_log20| max = 0.7093 Pa/s
|w_qg_moist| max = 2.0354 Pa/s
|w_em_moist| max = 0.8396 Pa/s

CPU times: user 2.26 s, sys: 101 ms, total: 2.36 s
Wall time: 3.43 s

21

6b LHR moist decomposition: w_qg_diabatic vs w_lhr_moist

Row 1 (map at PLOT_LEV): w_total | w_qg_diabatic (= w_total - w_adia) | w_lhr_moist (= w_ABC - w_AB, own cbar range)

Row 2 (x-p cross-section at y_rel=+15°, log-p vertical axis): LHR proxy theta_dot_LHR | C_em = -(kappa/p) nabla^2 J_em (diabatic Q-term) | w_lhr_moist

The LHR approximation uses saturation-specific-humidity lapse:

\[\dot{\theta}_{\mathrm{LHR}} \approx \frac{L_v}{c_p}\frac{\theta}{T}\max\!\bigl(0,\;-\omega\,\partial q_s/\partial p\bigr)\]

# ── 2×3 plot ──
# Row 1: w_total / w_qg_diabatic / w_lhr_moist  (map at PLOT_LEV)
#         w_lhr_moist uses its own cbar range
# Row 2: x–p cross-section at y_rel=+15  (LHR θ̇, C_em, w_lhr_moist), log-p axis

from matplotlib.ticker import ScalarFormatter

XP_YREL = 15.0  # cross-section latitude offset [°]
j_xp = np.argmin(np.abs(y_coords - XP_YREL))  # index for y_rel ≈ +15

fig, axes = plt.subplots(2, 3, figsize=(18, 11))

# ── Row 1: map view at PLOT_LEV ──
map_fields = [
    ("ω_total",     p_w_snap[ilev]),
    ("ω_qg_diabatic",  w_qg_diabatic[ilev]),
    ("ω_lhr_moist",  w_em_moist[ilev]),
]

# Shared range for w_total and w_qg_diabatic only
vm_shared = max(np.nanpercentile(np.abs(f), 98) for lbl, f in map_fields[:2])
if vm_shared < 1e-15:
    vm_shared = 1.0
# Separate range for w_em_moist
vm_em = np.nanpercentile(np.abs(w_em_moist[ilev]), 98)
if vm_em < 1e-15:
    vm_em = 1.0

for col, (label, fld) in enumerate(map_fields):
    ax = axes[0, col]
    vm_use = vm_em if col == 2 else vm_shared
    cf = ax.contourf(x_coords, y_coords, fld,
                     levels=np.linspace(-vm_use, vm_use, 21),
                     cmap="RdBu_r", extend="both")
    overlay_coastlines(ax, **coast_kw)
    ax.set_aspect("equal")
    ax.set_title(f"{label} @ {PLOT_LEV} hPa", fontsize=11)
    ax.set_xlabel("Rel. lon [°]"); ax.set_ylabel("Rel. lat [°]")
    plt.colorbar(cf, ax=ax, shrink=0.7, label="Pa/s")

# ── Row 2: x–p cross-sections at y_rel=+15, log-p axis ──
xp_fields = [
    ("θ̇_LHR [K/s]",           theta_dot_lhr_patch[:, j_xp, :], "Oranges"),
    ("C_em (Q-term) [s⁻³]",   C_em_patch[:, j_xp, :],          "RdBu_r"),
    ("ω_lhr_moist [Pa/s]",      w_em_moist[:, j_xp, :],          "RdBu_r"),
]

for col, (label, fld, cmap) in enumerate(xp_fields):
    ax = axes[1, col]
    vm_xp = np.nanpercentile(np.abs(fld), 98)
    if vm_xp < 1e-20:
        vm_xp = 1.0
    if cmap == "Oranges":
        levs_xp = np.linspace(0, vm_xp, 21)
    else:
        levs_xp = np.linspace(-vm_xp, vm_xp, 21)
    cf = ax.contourf(x_coords, levs_hpa, fld,
                     levels=levs_xp, cmap=cmap, extend="both")
    ax.set_yscale("log")
    ax.set_ylim(levs_hpa.max(), levs_hpa.min())  # high pressure at bottom
    ax.yaxis.set_major_formatter(ScalarFormatter())
    ax.yaxis.set_minor_formatter(ScalarFormatter())
    ax.set_yticks(levs_hpa)
    ax.set_title(f"{label}  (y_rel=+{XP_YREL:.0f}° section)", fontsize=10)
    ax.set_xlabel("Rel. lon [°]"); ax.set_ylabel("Pressure [hPa]")
    plt.colorbar(cf, ax=ax, shrink=0.7, label=label.split("[")[-1].rstrip("]") if "[" in label else "")

fig.suptitle(f"LHR moist decomposition @ {PLOT_LEV} hPa — {ts_str}\n"
             f"Row 1: maps  |  Row 2: x–p at y_rel=+{XP_YREL:.0f}°  |  "
             f"Centre ({CENTER_LAT}°N, {CENTER_LON}°E)",
             fontsize=13, y=0.97)
fig.tight_layout()
plt.show(); plt.close("all"); gc.collect()

70002

7 Diabatic / adiabatic divergent wind on patch

Uses decompose_omega() with ω_lhr_moist (LHR-driven) as the moist component, then recovers divergent wind for each component via independent Poisson inversions (tex Eq. 6, 15):

\[\nabla^2 \chi_{\rm lhr\_moist} = -\partial\omega_{\rm lhr\_moist}/\partial p, \quad \nabla^2 \chi_{\rm adia} = -\partial\omega_{\rm adia}/\partial p\]

Both χ are solved independently using the spherical Laplacian, rather than computing one as a residual.

1×3 panel: total divergent (from Helmholtz on total wind) | adiabatic divergent | lhr-moist divergent

%%time
# ── Divergent wind from w_lhr_moist and w_adia via independent Poisson ──
# ω_adia = wd_log20 (LOG20 A+B solve)
# ω_lhr_moist = w_em_moist (LHR contribution = ω_ABC − ω_AB)

omega_dry_total = wd_log20
omega_em_moist  = w_em_moist

# Independent Poisson inversions for each component
chi_em, u_div_em, v_div_em = solve_chi_from_omega(
    omega_em_moist, patch_lat, patch_lon, plevs_pa,
)
chi_d, u_div_adiabatic, v_div_adiabatic = solve_chi_from_omega(
    omega_dry_total, patch_lat, patch_lon, plevs_pa,
)

# QG-diabatic divergent wind (from full diabatic residual ω_qg_diabatic = ω_total − ω_adia)
chi_qg, u_div_qg, v_div_qg = solve_chi_from_omega(
    w_qg_diabatic, patch_lat, patch_lon, plevs_pa,
)

# For backward compat aliases used in §8
u_div_diabatic = u_div_em
v_div_diabatic = v_div_em
omega_diabatic_total = omega_em_moist

print(f"|u_div_lhr_moist| max = {np.abs(u_div_em[ilev]).max():.3f} m/s")
print(f"|u_div_qg_diabatic| max = {np.abs(u_div_qg[ilev]).max():.3f} m/s")
print(f"|u_div_adia|      max = {np.abs(u_div_adiabatic[ilev]).max():.3f} m/s")

# ── 1×3 panel: divergent wind total / dry / em-moist @ PLOT_LEV ──
fig, axes = plt.subplots(1, 3, figsize=(18, 5))
div_panels = [
    ("Div total",       p_u_div[ilev],  p_v_div[ilev]),
    ("Div adiabatic",  u_div_adiabatic[ilev], v_div_adiabatic[ilev]),
    ("Div lhr-moist",   u_div_em[ilev],  v_div_em[ilev]),
]

for ax, (title, uk, vk) in zip(axes, div_panels):
    speed = np.sqrt(uk**2 + vk**2)
    # Scale quiver arrows proportional to max speed in each panel
    ref_speed = max(float(np.nanpercentile(speed, 98)), 0.01)
    ref_nice = float(np.round(ref_speed, -int(np.floor(np.log10(ref_speed)))))  # round to 1 sig fig
    if ref_nice < 1e-10:
        ref_nice = ref_speed
    qscale = ref_nice * 20  # 20 arrows worth across the domain

    cf = ax.contourf(x_coords, y_coords, speed, levels=20, cmap="YlOrRd", extend="max")
    q = ax.quiver(x_coords[::skip], y_coords[::skip],
                  uk[::skip, ::skip], vk[::skip, ::skip],
                  scale=qscale, width=0.003, color="k")
    ax.quiverkey(q, 0.88, 1.06, ref_nice, f"{ref_nice:.1f} m/s",
                 labelpos="E", fontproperties=dict(size=8))
    overlay_coastlines(ax, **coast_kw)
    ax.set_aspect("equal")
    ax.set_xlabel("Rel. lon [°]"); ax.set_ylabel("Rel. lat [°]")
    ax.set_title(title, fontsize=12)
    plt.colorbar(cf, ax=ax, shrink=0.7, label="m/s")

fig.suptitle(f"Divergent wind: total / adia / lhr-moist @ {PLOT_LEV} hPa (LOG20+LHR) — {ts_str}",
             fontsize=14, y=1.05)
fig.tight_layout()
plt.show(); plt.close("all"); gc.collect()

|u_div_lhr_moist| max = 2.068 m/s
|u_div_qg_diabatic| max = 6.410 m/s
|u_div_adia|      max = 6.008 m/s

CPU times: user 823 ms, sys: 3.98 ms, total: 827 ms
Wall time: 827 ms

58875

8 PV tendency on patch: three estimates

All terms use total PV gradients (\(\partial q/\partial x\), \(\partial q/\partial y\), \(\partial q/\partial p\)).

1. Indirect LHR-moist (LHR divergent wind advecting total PV):

\[\mathcal{T}_{\rm lhr\_moist} = -u_{\chi,\rm lhr}\frac{\partial q}{\partial x} - v_{\chi,\rm lhr}\frac{\partial q}{\partial y} - \omega_{\rm lhr}\frac{\partial q}{\partial p}\]

2. QG-diabatic (full diabatic residual divergent wind from \(\omega_{\rm qg\_diabatic} = \omega - \omega_{\rm adia}\)):

\[\mathcal{T}_{\rm qg\_diabatic} = -u_{\chi,\rm qg}\frac{\partial q}{\partial x} - v_{\chi,\rm qg}\frac{\partial q}{\partial y} - \omega_{\rm qg\_diabatic}\frac{\partial q}{\partial p}\]

3. Total divergent outflow (Helmholtz divergent wind + diabatic residual ω):

\[\mathcal{T}_{\rm div} = -u_{\rm div}\frac{\partial q}{\partial x} - v_{\rm div}\frac{\partial q}{\partial y} - \omega_{\rm dia\_res}\frac{\partial q}{\partial p}\]

where \(\omega_{\rm dia\_res} = \omega - \omega_{\rm adia}\).

# ── PV tendency terms on the patch ──
# Uses TOTAL PV gradient (∂q/∂x, ∂q/∂y, ∂q/∂p)
# and LHR-moist divergent wind from w_em_moist Poisson solve.

# ω_dia_res = ω_total − ω_adia  (already computed in §6 as wm_log20)
w_diabatic = wm_log20

# Indirect LHR-moist PV tendency
T_diabatic = -(u_div_diabatic * p_dpv_bar_dx
            + v_div_diabatic * p_dpv_bar_dy
            + omega_diabatic_total * p_dpv_bar_dp)

# QG-diabatic PV tendency (full diabatic divergent wind from ω_qg_diabatic Poisson)
T_qg_diabatic = -(u_div_qg * p_dpv_bar_dx
               + v_div_qg * p_dpv_bar_dy
               + w_qg_diabatic * p_dpv_bar_dp)

# Total divergent outflow PV tendency (Helmholtz div wind + ω_dia_res)
T_div = -(p_u_div * p_dpv_bar_dx
          + p_v_div * p_dpv_bar_dy
          + w_diabatic * p_dpv_bar_dp)

# ── Stats at PLOT_LEV ──
tm_lev = T_diabatic[ilev]
tq_lev = T_qg_diabatic[ilev]
td_lev = T_div[ilev]
mask = (np.abs(td_lev) > 1e-20) & np.isfinite(tm_lev) & np.isfinite(td_lev)
corr_em  = np.corrcoef(tm_lev[mask], td_lev[mask])[0, 1] if mask.sum() > 10 else np.nan
mask_qg = (np.abs(td_lev) > 1e-20) & np.isfinite(tq_lev) & np.isfinite(td_lev)
corr_qg  = np.corrcoef(tq_lev[mask_qg], td_lev[mask_qg])[0, 1] if mask_qg.sum() > 10 else np.nan
ratio_em = np.nanmean(np.abs(tm_lev)) / (np.nanmean(np.abs(td_lev)) + 1e-30)
ratio_qg = np.nanmean(np.abs(tq_lev)) / (np.nanmean(np.abs(td_lev)) + 1e-30)
print(f"At {PLOT_LEV} hPa:")
print(f"  corr(T_lhr_moist, T_div) = {corr_em:.4f}  |  |T_lhr|/|T_div| = {ratio_em:.3f}")
print(f"  corr(T_qg_diabatic, T_div) = {corr_qg:.4f}  |  |T_qd|/|T_div| = {ratio_qg:.3f}")

# ── 1×3 contourf: T_lhr_moist / T_qg_diabatic / T_div ──
fig, axes = plt.subplots(1, 3, figsize=(20, 5))

# Shared colour scale across all three
vm_pv = max(np.nanpercentile(np.abs(tm_lev), 98),
            np.nanpercentile(np.abs(tq_lev), 98),
            np.nanpercentile(np.abs(td_lev), 98))
if vm_pv < 1e-30:
    vm_pv = 1.0
levs_cf = np.linspace(-vm_pv, vm_pv, 21)

for ax, field, title in zip(axes, [tm_lev, tq_lev, td_lev],
                             [r"$\mathcal{T}_{\rm lhr\_moist}$",
                              r"$\mathcal{T}_{\rm qg\_diabatic}$",
                              r"$\mathcal{T}_{\rm div}$"]):
    cf = ax.contourf(x_coords, y_coords, field, levels=levs_cf,
                     cmap="RdBu_r", extend="both")
    overlay_coastlines(ax, **coast_kw)
    ax.set_aspect("equal")
    ax.set_xlabel("Rel. lon [°]"); ax.set_ylabel("Rel. lat [°]")
    ax.set_title(title, fontsize=13)
    plt.colorbar(cf, ax=ax, shrink=0.7, label="PVU/s")

fig.suptitle(f"PV tendency (total ∂q) @ {PLOT_LEV} hPa — {ts_str}\n"
             f"corr(lhr,div)={corr_em:.3f}  corr(qg_dia,div)={corr_qg:.3f}",
             fontsize=13, y=1.02)
fig.tight_layout()
plt.show(); plt.close("all"); gc.collect()

At 300 hPa:
  corr(T_lhr_moist, T_div) = 0.1984  |  |T_lhr|/|T_div| = 0.100
  corr(T_qg_diabatic, T_div) = 0.7492  |  |T_qd|/|T_div| = 0.869

2446

9 Save all computed fields to NPZ

Save all patch-level fields to a compressed NPZ archive for downstream analysis. Format matches the core step2_compute_tendency_terms_blocking.py output.

import os

out_dir = DATA_DIR
out_path = os.path.join(out_dir, "era5_step2_jan2025.npz")

# ── Vertical-average weighting (pressure thickness weights) ──
dp_weights = np.zeros(nlev)
for k in range(nlev):
    if k == 0:
        dp_weights[k] = (plevs_pa[1] - plevs_pa[0]) / 2.0
    elif k == nlev - 1:
        dp_weights[k] = (plevs_pa[-1] - plevs_pa[-2]) / 2.0
    else:
        dp_weights[k] = (plevs_pa[k+1] - plevs_pa[k-1]) / 2.0
dp_weights /= dp_weights.sum()

def wavg(f3d):
    """Pressure-weighted vertical average → 2D."""
    return np.nansum(f3d * dp_weights[:, None, None], axis=0)

np.savez_compressed(
    out_path,
    # ── Metadata ──
    Y_rel=Y_rel,
    X_rel=X_rel,
    levels=levs_hpa,
    plevs_pa=plevs_pa,
    center_lat=CENTER_LAT,
    center_lon=CENTER_LON,
    timestamp=ts_str,
    qg_method="log20",
    # ── Patch meteorological fields (3D: nlev, nlat_p, nlon_p) ──
    t_snap=p_t_snap,
    z_snap=p_z_snap,
    w_snap=p_w_snap,
    u_snap=p_u_snap,
    v_snap=p_v_snap,
    # ── Helmholtz wind (3D, spherical FFT Laplacian) ──
    u_rot=p_u_rot,
    v_rot=p_v_rot,
    u_div=p_u_div,
    v_div=p_v_div,
    u_har=p_u_har,
    v_har=p_v_har,
    psi=p_psi,
    chi=p_chi,
    # ── Geostrophic wind (3D) ──
    ug_snap=p_ug_snap,
    vg_snap=p_vg_snap,
    # ── PV derivatives (3D) ──
    dpv_bar_dx=p_dpv_bar_dx,
    dpv_bar_dy=p_dpv_bar_dy,
    dpv_bar_dp=p_dpv_bar_dp,
    # ── QG omega: adiabatic (A+B), full (A+B+C_em) ──
    omega_dry_sp19=wd_sp19,
    omega_dry_log20=wd_log20,
    omega_moist_sp19=wm_sp19,
    omega_moist_log20=wm_log20,
    # ── Three components: w_adia, w_qg_diabatic, w_lhr_moist ──
    omega_dry_total=wd_log20,
    w_qg_diabatic=w_qg_diabatic,
    w_em_moist=w_em_moist,
    # ── LHR diagnostics (patch) ──
    theta_dot_lhr=theta_dot_lhr_patch,
    C_em=C_em_patch,
    # ── LHR/adiabatic divergent wind from w_lhr_moist Poisson (3D) ──
    u_div_diabatic=u_div_diabatic,
    v_div_diabatic=v_div_diabatic,
    u_div_adiabatic=u_div_adiabatic,
    v_div_adiabatic=v_div_adiabatic,
    omega_diabatic_total=omega_diabatic_total,
    # ── PV tendency (3D) ──
    T_diabatic=T_diabatic,
    T_div=T_div,
    # ── Vertical averages (2D) ──
    T_diabatic_wavg=wavg(T_diabatic),
    T_div_wavg=wavg(T_div),
    u_rot_wavg=wavg(p_u_rot),
    v_rot_wavg=wavg(p_v_rot),
    u_div_wavg=wavg(p_u_div),
    v_div_wavg=wavg(p_v_div),
    u_div_diabatic_wavg=wavg(u_div_diabatic),
    v_div_diabatic_wavg=wavg(v_div_diabatic),
)

print(f"Saved → {out_path}  ({os.path.getsize(out_path)/1e6:.1f} MB)")

Saved → era5_jan2025/era5_step2_jan2025.npz  (2.1 MB)