| 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() |
|
|
