File size: 3,149 Bytes
5c365c5 | 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 | """Climate metrics and conservation diagnostics from the paper equations."""
from __future__ import annotations
import torch
def _weights(area_weights: torch.Tensor, field: torch.Tensor) -> torch.Tensor:
weights = area_weights.to(device=field.device, dtype=field.dtype)
if weights.ndim == 1:
weights = weights[:, None]
if weights.shape != field.shape[-2:]:
raise ValueError(f"area weights must match spatial shape {field.shape[-2:]}, got {tuple(weights.shape)}")
return weights / weights.sum().clamp_min(torch.finfo(field.dtype).eps)
def area_weighted_global_mean(field: torch.Tensor, area_weights: torch.Tensor) -> torch.Tensor:
weights = _weights(area_weights, field)
return (field * weights).sum(dim=(-2, -1))
def time_mean_pattern_rmse(pred: torch.Tensor, truth: torch.Tensor, area_weights: torch.Tensor) -> torch.Tensor:
if pred.shape != truth.shape:
raise ValueError("pred and truth must have equal shapes")
error = (pred - truth).mean(dim=0) if pred.ndim == 4 else (pred - truth).mean(dim=1)
weights = _weights(area_weights, error)
return torch.sqrt((error.square() * weights).sum(dim=(-2, -1)).clamp_min(0.0))
def global_time_mean_bias(pred: torch.Tensor, truth: torch.Tensor, area_weights: torch.Tensor) -> torch.Tensor:
if pred.shape != truth.shape:
raise ValueError("pred and truth must have equal shapes")
pred_gm = area_weighted_global_mean(pred, area_weights)
truth_gm = area_weighted_global_mean(truth, area_weights)
return (pred_gm - truth_gm).mean(dim=0 if pred.ndim == 4 else 1)
def total_water_path(qt: torch.Tensor, dp: torch.Tensor, gravity: float = 9.80665) -> torch.Tensor:
"""Compute TWP=(1/g) sum_k qT_k dp_k; qt/dp are [...,8,H,W]."""
if qt.shape != dp.shape or qt.shape[-3] != 8:
raise ValueError("qt and dp must have shape [...,8,H,W]")
return (qt * dp).sum(dim=-3) / gravity
def moisture_budget_violation(twp_t: torch.Tensor, twp_next: torch.Tensor, evaporation: torch.Tensor, precipitation: torch.Tensor, advective_tendency: torch.Tensor) -> torch.Tensor:
"""Eq.(1) residual using a one-step tendency."""
return (twp_next - twp_t) - (evaporation - precipitation + advective_tendency)
def dry_air_surface_pressure(surface_pressure: torch.Tensor, twp: torch.Tensor, gravity: float = 9.80665) -> torch.Tensor:
return surface_pressure - gravity * twp
def forecast_metrics(pred: torch.Tensor, truth: torch.Tensor, area_weights: torch.Tensor) -> dict[str, float]:
"""Return aggregate Eq.(5-7) metrics for [T,C,H,W] or [B,T,C,H,W]."""
if pred.shape != truth.shape:
raise ValueError("pred and truth must have equal shapes")
if pred.ndim == 5:
pred = pred[0]
truth = truth[0]
rmse = time_mean_pattern_rmse(pred, truth, area_weights)
bias = global_time_mean_bias(pred, truth, area_weights)
return {
"mean_time_mean_rmse": float(rmse.mean().item()),
"mean_global_time_mean_bias": float(bias.mean().item()),
"T7_time_mean_rmse": float(rmse[7].item()),
"precipitation_time_mean_rmse": float(rmse[40].item()),
}
|