File size: 2,106 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 | """ACE variable order and channel contract."""
from __future__ import annotations
import torch
PROGNOSTIC_3D = tuple(f"{name}_{level}" for name in ("T", "qT", "u", "v") for level in range(8))
PROGNOSTIC_SURFACE = ("Ts_land_or_seaice", "ps")
PROGNOSTIC_CHANNELS = PROGNOSTIC_3D + PROGNOSTIC_SURFACE
FORCING_CHANNELS = ("DSW_RF_toa", "Ts_open_ocean", "zs", "fl", "fo", "fsi")
DIAGNOSTIC_CHANNELS = (
"USW_RF_toa", "ULW_RF_toa", "USW_RF_sfc", "ULW_RF_sfc",
"DSW_RF_sfc", "DLW_RF_sfc", "P", "dTWP_adv_dt", "LHF", "SHF",
)
PROGNOSTIC_UNITS = ("K", "kg/kg", "m/s", "m/s") * 8 + ("K", "Pa")
FORCING_UNITS = ("W/m2", "K", "m", "1", "1", "1")
DIAGNOSTIC_UNITS = ("W/m2",) * 6 + ("kg/m2/s", "kg/m2/s", "W/m2", "W/m2")
INPUT_UNITS = PROGNOSTIC_UNITS + FORCING_UNITS
OUTPUT_UNITS = PROGNOSTIC_UNITS + DIAGNOSTIC_UNITS
INPUT_CHANNELS = PROGNOSTIC_CHANNELS + FORCING_CHANNELS
OUTPUT_CHANNELS = PROGNOSTIC_CHANNELS + DIAGNOSTIC_CHANNELS
def ledger() -> dict:
return {
"input_channels": list(INPUT_CHANNELS),
"output_channels": list(OUTPUT_CHANNELS),
"input_units": list(INPUT_UNITS),
"output_units": list(OUTPUT_UNITS),
"prognostic_channels": len(PROGNOSTIC_CHANNELS),
"forcing_channels": len(FORCING_CHANNELS),
"diagnostic_channels": len(DIAGNOSTIC_CHANNELS),
"input_total": len(INPUT_CHANNELS),
"output_total": len(OUTPUT_CHANNELS),
}
def validate_channels(x: torch.Tensor, expected: int, name: str = "tensor") -> None:
if x.ndim != 4:
raise ValueError(f"{name} must have shape [B,C,H,W], got {tuple(x.shape)}")
if x.shape[1] != expected:
raise ValueError(f"{name} must have {expected} channels, got {x.shape[1]}")
def split_input(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
validate_channels(x, len(INPUT_CHANNELS), "input")
n = len(PROGNOSTIC_CHANNELS)
return x[:, :n], x[:, n:]
def split_output(y: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
validate_channels(y, len(OUTPUT_CHANNELS), "output")
n = len(PROGNOSTIC_CHANNELS)
return y[:, :n], y[:, n:]
|