File size: 2,925 Bytes
a13f4b9 | 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 | """Run validated, batched autoregressive solar forecasting inference."""
import importlib.util
from pathlib import Path
import numpy as np
import torch
import yaml
ROOT = Path(__file__).resolve().parents[1]
def main():
import argparse
parser = argparse.ArgumentParser(); parser.add_argument("--config", type=Path, default=ROOT / "conf/config.yaml"); parser.add_argument("--data", type=Path); parser.add_argument("--checkpoint", type=Path); parser.add_argument("--output-dir", type=Path); parser.add_argument("--batch-size", type=int, default=4); parser.add_argument("--device", choices=("auto", "cpu", "cuda"), default=None); args = parser.parse_args()
cfg = yaml.safe_load(args.config.read_text())
spec = importlib.util.spec_from_file_location("surya_model", ROOT / "model/surya.py")
module = importlib.util.module_from_spec(spec); spec.loader.exec_module(module)
model = module.Surya(**cfg["model"])
checkpoint = args.checkpoint or ROOT / cfg["paths"]["checkpoint"]
if not checkpoint.exists(): raise FileNotFoundError("Run training before inference")
device_name = args.device or cfg["runtime"]["device"]
if device_name == "cuda" and not torch.cuda.is_available(): raise RuntimeError("CUDA requested but unavailable")
device = torch.device("cuda" if torch.cuda.is_available() and device_name != "cpu" else "cpu")
model.load_state_dict(torch.load(checkpoint, map_location="cpu", weights_only=False)["model"]); model.to(device).eval()
data = np.load(args.data or ROOT / cfg["data"]["root"] / "test.npz"); raw_inputs, raw_targets = data["inputs"], data["targets"]
if raw_inputs.ndim != 5 or raw_inputs.shape[1:] != (2, 13, cfg["data"]["image_size"], cfg["data"]["image_size"]): raise ValueError("test.npz violates the BTCHW protocol")
mean = np.asarray(cfg["data"]["channel_mean"], dtype=np.float32)[None, None, :, None, None]; std = np.asarray(cfg["data"]["channel_std"], dtype=np.float32)[None, None, :, None, None]
inputs = (np.sign(raw_inputs) * np.log1p(np.abs(raw_inputs)) - mean) / std
predictions = []
with torch.no_grad():
for start in range(0, len(inputs), args.batch_size): predictions.append(model(torch.from_numpy(inputs[start:start+args.batch_size]).to(device), steps=cfg["data"]["forecast_steps"]).cpu().numpy())
predictions = np.concatenate(predictions, axis=0)
# Evaluation consumes the original physical-value representation.
predictions = np.sign(predictions * std[:, :1] + mean[:, :1]) * (np.expm1(np.abs(predictions * std[:, :1] + mean[:, :1])))
output = args.output_dir or ROOT / cfg["paths"]["inference_dir"]; output.mkdir(parents=True, exist_ok=True)
np.savez_compressed(output / "forecast.npz", inputs=raw_inputs, targets=raw_targets,
predictions=predictions, activity=data["activity"])
print("inference=", output / "forecast.npz")
if __name__ == "__main__": main()
|