File size: 2,512 Bytes
7f71cfd | 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 | from pathlib import Path
import sys
import numpy as np
import torch
import yaml
ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
sys.path.insert(0, str(ROOT))
from model.stablenn_phys import StableNNPhys, rollout
CHECKPOINT_FORMAT_VERSION = "stablenn_phys_checkpoint_v1"
REQUIRED_CHECKPOINT_FIELDS = {"model", "model_config", "format_version"}
def main():
cfg = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
raw = np.load(ROOT / cfg["data"]["file"])
checkpoint = torch.load(ROOT / cfg["paths"]["checkpoint"], map_location="cpu", weights_only=False)
missing = REQUIRED_CHECKPOINT_FIELDS.difference(checkpoint)
if missing:
raise ValueError(f"Checkpoint is missing required fields: {', '.join(sorted(missing))}")
if checkpoint["format_version"] != CHECKPOINT_FORMAT_VERSION:
raise ValueError(
f"Unsupported checkpoint format_version {checkpoint['format_version']!r}; "
f"expected {CHECKPOINT_FORMAT_VERSION!r}"
)
model = StableNNPhys(checkpoint["model_config"]["hidden_size"])
model.load_state_dict(checkpoint["model"])
model.eval()
steps = cfg["runtime"]["rollout_steps"]
if steps != 64:
raise ValueError(f"runtime.rollout_steps must be 64, got {steps}")
initial = torch.from_numpy(raw["long_state"][:, 0].astype(np.float32))
surface = torch.from_numpy(raw["long_surface"][:, :steps].astype(np.float32))
advection = torch.from_numpy(raw["long_advection"][:, :steps + 1].astype(np.float32))
norm = checkpoint["normalization"]
with torch.no_grad():
prediction, tendency = rollout(model, initial, surface, advection, norm["state_mean"], norm["state_std"],
norm["tendency_mean"], norm["tendency_std"], checkpoint["dt_seconds"])
target = raw["long_state"][:, :steps + 1]
out = ROOT / cfg["paths"]["inference"]
out.parent.mkdir(parents=True, exist_ok=True)
np.savez_compressed(out, prediction=prediction.numpy(), target=target, tendency=tendency.numpy(),
surface=surface.numpy(), advection=advection.numpy(), layer_mass=raw["layer_mass"],
source=np.arange(initial.shape[0]),
lead_hours=np.arange(steps + 1) * cfg["data"]["step_hours"],
dt_seconds=np.float32(checkpoint["dt_seconds"]))
print(f"saved {out}: {steps} steps, {steps * 3 / 24:.1f} days")
if __name__ == "__main__":
main()
|