scripts/inference.py

"""OneForecast inference entry point with the shared ERA5 adapter."""

from __future__ import annotations

import argparse
from pathlib import Path
import sys

import numpy as np
import torch
import yaml

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from model.era5_adapter import OFFICIAL_VARIABLES, OneForecastERA5Adapter
from model.oneforecast import build_model, check_checkpoint_compatibility, read_official_checkpoint


def _resolve_path(value: str | Path, config_path: Path) -> Path:
    path = Path(value).expanduser()
    return path if path.is_absolute() else (config_path.parent.parent / path).resolve()


def _load_config(path: Path) -> dict:
    with path.open("r", encoding="utf-8") as handle:
        config = yaml.safe_load(handle)
    config["datapipe"]["dataset_dir"] = str(_resolve_path(config["datapipe"]["dataset_dir"], path))
    config["model"]["official_checkpoint_path"] = str(
        _resolve_path(config["model"]["official_checkpoint_path"], path)
    )
    config["model"]["checkpoint_path"] = config["model"]["official_checkpoint_path"]
    config["inference"]["trained_model_path"] = str(
        _resolve_path(config["inference"]["trained_model_path"], path)
    )
    config["inference"]["official_checkpoint_path"] = str(
        _resolve_path(config["inference"]["official_checkpoint_path"], path)
    )
    config["inference"]["output_dir"] = str(_resolve_path(config["inference"]["output_dir"], path))
    return config


def _resolve_device(name: str) -> torch.device:
    """Map the logical DCU name to the backend exposed by this PyTorch build."""
    requested = str(name).lower()
    if requested == "dcu":
        if torch.cuda.is_available():
            return torch.device("cuda")
        privateuse = torch._C._get_privateuse1_backend_name()
        if privateuse != "privateuseone":
            return torch.device(privateuse)
        raise RuntimeError("runtime.device=dcu, but this PyTorch build exposes no usable accelerator")
    if requested == "auto":
        return torch.device("cuda" if torch.cuda.is_available() else "cpu")
    device = torch.device(requested)
    if device.type == "cuda" and not torch.cuda.is_available():
        raise RuntimeError("runtime.device=cuda, but torch.cuda.is_available() is False")
    return device


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", type=Path, default=Path("conf/config.yaml"))
    parser.add_argument("--check-data", action="store_true")
    parser.add_argument("--check-model", action="store_true")
    parser.add_argument("--check-checkpoint", action="store_true")
    parser.add_argument("--model-source", choices=("trained", "official"), default=None)
    args = parser.parse_args()
    config = _load_config(args.config.resolve())
    if tuple(config["datapipe"]["variables"]) != OFFICIAL_VARIABLES:
        raise ValueError("datapipe.variables must exactly match the official 69-channel order")
    if args.model_source is not None:
        config["inference"]["model_source"] = args.model_source
    if args.check_data:
        settings = config["datapipe"]
        adapter = OneForecastERA5Adapter(
            settings["dataset_dir"], settings["test_years"], batch_size=1,
            input_steps=settings["input_steps"], output_steps=settings["output_steps"],
            normalize=settings["normalize"], num_workers=settings["num_workers"],
        )
        print(adapter.inspect())
        return
    if args.check_model:
        configured_init = config["model"].get("weight_init", "scratch")
        config["model"]["weight_init"] = "scratch"
        with __import__("torch").device("meta"):
            model = build_model(config, build_graph=False)
        print({"model": type(model).__name__, "parameters": sum(p.numel() for p in model.parameters()),
               "configured_weight_init": configured_init})
        return
    if args.check_checkpoint:
        with __import__("torch").device("meta"):
            model = build_model(config, build_graph=False)
        report = check_checkpoint_compatibility(
            model, config["model"]["official_checkpoint_path"]
        )
        print(report)
        if not report.compatible:
            raise SystemExit(1)
        return
    settings = config["datapipe"]
    if settings["input_steps"] != 1 or settings["output_steps"] != 1:
        raise SystemExit("OneForecast inference currently requires input_steps=1 and output_steps=1")
    device = _resolve_device(config["runtime"].get("device", "cpu"))
    config["model"]["weight_init"] = "scratch"
    model = build_model(config).to(device)
    source = config["inference"].get("model_source", "trained")
    checkpoint_path = config["inference"][
        "trained_model_path" if source == "trained" else "official_checkpoint_path"
    ]
    state, _ = read_official_checkpoint(checkpoint_path)
    model.load_state_dict(state)
    model.eval()
    adapter = OneForecastERA5Adapter(
        _resolve_path(settings["dataset_dir"], args.config), settings["test_years"],
        batch_size=1, input_steps=1, output_steps=1,
        normalize=settings["normalize"], num_workers=settings["num_workers"],
    )
    loader, _ = adapter.get_dataloader("test")
    output_dir = Path(config["inference"]["output_dir"])
    output_dir.mkdir(parents=True, exist_ok=True)
    max_batches = int(config["inference"].get("max_batches", -1))
    processed = 0
    with torch.no_grad():
        for index, batch in enumerate(loader):
            inputs, targets = batch[0].float().to(device), batch[1].float().to(device)
            if inputs.ndim == 5 or targets.ndim == 5:
                raise ValueError("OneForecast currently supports input_steps=1 and output_steps=1 only")
            if inputs.ndim != 4:
                raise ValueError(f"Expected batched input with four dimensions, got {inputs.shape}")
            if inputs.shape[-2] == 121:
                inputs = inputs[..., :120, :]
            if targets.shape[-2] == 121:
                targets = targets[..., :120, :]
            if inputs.shape[-2:] != (120, 240) or targets.shape[-2:] != (120, 240):
                raise ValueError(f"Expected official model grid 120x240, got {inputs.shape} and {targets.shape}")
            prediction = model(torch.nan_to_num(inputs))
            if settings["normalize"]:
                means, stds = adapter.selected_statistics()
                prediction = prediction.cpu() * torch.from_numpy(stds).float() + torch.from_numpy(means).float()
            np.save(output_dir / f"prediction_{index:05d}.npy", prediction.cpu().numpy())
            processed += 1
            if max_batches >= 0 and index + 1 >= max_batches:
                break
    print({"output_dir": str(output_dir), "batches": processed})


if __name__ == "__main__":
    main()