File size: 2,111 Bytes
38f9214 | 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 | """Input/output contract for the isolated fake MetNet-3 package."""
from dataclasses import dataclass, field
from typing import Dict, Mapping, Tuple
import torch
@dataclass
class InputSchema:
high_channels: int = 4
low_channels: int = 3
omo_channels: int = 2
hrrr_channels: int = 8 # ASSUMPTION: proxy, not the paper's 617 channels.
goes_channels: int = 4 # ASSUMPTION: proxy, not the paper's 16 channels.
static_channels: int = 2 # coordinates/topography grouping; elevation is explicit below.
high_frames: int = 3 # ADAPTATION: paper describes 11 MRMS frames.
omo_frames: int = 3 # ADAPTATION: paper describes 9 OMO frames.
ground_targets: int = 6
precipitation_bins: int = 16 # ADAPTATION: paper uses 512 bins.
ground_bins: int = 8 # ADAPTATION: paper uses 256/180 bins.
@property
def total_static_channels(self) -> int:
return self.static_channels + 2 # elevation/topography + coordinates
def _require(batch: Mapping[str, torch.Tensor], key: str) -> torch.Tensor:
if key not in batch:
raise KeyError(f"missing MetNet-3 input: {key}")
return batch[key]
def validate_batch(batch: Mapping[str, torch.Tensor], schema: InputSchema) -> Tuple[int, int, int]:
"""Validate fake/adapter tensors and return (batch, height, width)."""
high = _require(batch, "mrms_high")
if high.ndim != 5:
raise ValueError("mrms_high must be [B,T,C,H,W]")
b, _, _, h, w = high.shape
for key in ("mrms_low", "omo", "hrrr", "goes"):
value = _require(batch, key)
if value.shape[0] != b or value.ndim != 5:
raise ValueError(f"{key} must be batched [B,T,C,H,W]")
for key in ("elevation", "coordinates", "topography_embedding", "omo_input_mask", "current_time", "lead_time"):
value = _require(batch, key)
if key in ("current_time", "lead_time"):
if value.shape != (b, 1):
raise ValueError(f"{key} must have shape [B, 1]")
elif value.shape[-2:] != (h, w):
raise ValueError(f"{key} spatial shape mismatch")
return b, h, w
|