OneScience-Group/OneForecast

1"""OneForecast inference entry point with the shared ERA5 adapter."""2 3from __future__ import annotations4 5import argparse6from pathlib import Path7import sys8 9import numpy as np10import torch11import yaml12 13sys.path.insert(0, str(Path(__file__).resolve().parents[1]))14 15from model.era5_adapter import OFFICIAL_VARIABLES, OneForecastERA5Adapter16from model.oneforecast import build_model, check_checkpoint_compatibility, read_official_checkpoint17 18 19def _resolve_path(value: str | Path, config_path: Path) -> Path:20    path = Path(value).expanduser()21    return path if path.is_absolute() else (config_path.parent.parent / path).resolve()22 23 24def _load_config(path: Path) -> dict:25    with path.open("r", encoding="utf-8") as handle:26        config = yaml.safe_load(handle)27    config["datapipe"]["dataset_dir"] = str(_resolve_path(config["datapipe"]["dataset_dir"], path))28    config["model"]["official_checkpoint_path"] = str(29        _resolve_path(config["model"]["official_checkpoint_path"], path)30    )31    config["model"]["checkpoint_path"] = config["model"]["official_checkpoint_path"]32    config["inference"]["trained_model_path"] = str(33        _resolve_path(config["inference"]["trained_model_path"], path)34    )35    config["inference"]["official_checkpoint_path"] = str(36        _resolve_path(config["inference"]["official_checkpoint_path"], path)37    )38    config["inference"]["output_dir"] = str(_resolve_path(config["inference"]["output_dir"], path))39    return config40 41 42def _resolve_device(name: str) -> torch.device:43    """Map the logical DCU name to the backend exposed by this PyTorch build."""44    requested = str(name).lower()45    if requested == "dcu":46        if torch.cuda.is_available():47            return torch.device("cuda")48        privateuse = torch._C._get_privateuse1_backend_name()49        if privateuse != "privateuseone":50            return torch.device(privateuse)51        raise RuntimeError("runtime.device=dcu, but this PyTorch build exposes no usable accelerator")52    if requested == "auto":53        return torch.device("cuda" if torch.cuda.is_available() else "cpu")54    device = torch.device(requested)55    if device.type == "cuda" and not torch.cuda.is_available():56        raise RuntimeError("runtime.device=cuda, but torch.cuda.is_available() is False")57    return device58 59 60def main() -> None:61    parser = argparse.ArgumentParser()62    parser.add_argument("--config", type=Path, default=Path("conf/config.yaml"))63    parser.add_argument("--check-data", action="store_true")64    parser.add_argument("--check-model", action="store_true")65    parser.add_argument("--check-checkpoint", action="store_true")66    parser.add_argument("--model-source", choices=("trained", "official"), default=None)67    args = parser.parse_args()68    config = _load_config(args.config.resolve())69    if tuple(config["datapipe"]["variables"]) != OFFICIAL_VARIABLES:70        raise ValueError("datapipe.variables must exactly match the official 69-channel order")71    if args.model_source is not None:72        config["inference"]["model_source"] = args.model_source73    if args.check_data:74        settings = config["datapipe"]75        adapter = OneForecastERA5Adapter(76            settings["dataset_dir"], settings["test_years"], batch_size=1,77            input_steps=settings["input_steps"], output_steps=settings["output_steps"],78            normalize=settings["normalize"], num_workers=settings["num_workers"],79        )80        print(adapter.inspect())81        return82    if args.check_model:83        configured_init = config["model"].get("weight_init", "scratch")84        config["model"]["weight_init"] = "scratch"85        with __import__("torch").device("meta"):86            model = build_model(config, build_graph=False)87        print({"model": type(model).__name__, "parameters": sum(p.numel() for p in model.parameters()),88               "configured_weight_init": configured_init})89        return90    if args.check_checkpoint:91        with __import__("torch").device("meta"):92            model = build_model(config, build_graph=False)93        report = check_checkpoint_compatibility(94            model, config["model"]["official_checkpoint_path"]95        )96        print(report)97        if not report.compatible:98            raise SystemExit(1)99        return100    settings = config["datapipe"]101    if settings["input_steps"] != 1 or settings["output_steps"] != 1:102        raise SystemExit("OneForecast inference currently requires input_steps=1 and output_steps=1")103    device = _resolve_device(config["runtime"].get("device", "cpu"))104    config["model"]["weight_init"] = "scratch"105    model = build_model(config).to(device)106    source = config["inference"].get("model_source", "trained")107    checkpoint_path = config["inference"][108        "trained_model_path" if source == "trained" else "official_checkpoint_path"109    ]110    state, _ = read_official_checkpoint(checkpoint_path)111    model.load_state_dict(state)112    model.eval()113    adapter = OneForecastERA5Adapter(114        _resolve_path(settings["dataset_dir"], args.config), settings["test_years"],115        batch_size=1, input_steps=1, output_steps=1,116        normalize=settings["normalize"], num_workers=settings["num_workers"],117    )118    loader, _ = adapter.get_dataloader("test")119    output_dir = Path(config["inference"]["output_dir"])120    output_dir.mkdir(parents=True, exist_ok=True)121    max_batches = int(config["inference"].get("max_batches", -1))122    processed = 0123    with torch.no_grad():124        for index, batch in enumerate(loader):125            inputs, targets = batch[0].float().to(device), batch[1].float().to(device)126            if inputs.ndim == 5 or targets.ndim == 5:127                raise ValueError("OneForecast currently supports input_steps=1 and output_steps=1 only")128            if inputs.ndim != 4:129                raise ValueError(f"Expected batched input with four dimensions, got {inputs.shape}")130            if inputs.shape[-2] == 121:131                inputs = inputs[..., :120, :]132            if targets.shape[-2] == 121:133                targets = targets[..., :120, :]134            if inputs.shape[-2:] != (120, 240) or targets.shape[-2:] != (120, 240):135                raise ValueError(f"Expected official model grid 120x240, got {inputs.shape} and {targets.shape}")136            prediction = model(torch.nan_to_num(inputs))137            if settings["normalize"]:138                means, stds = adapter.selected_statistics()139                prediction = prediction.cpu() * torch.from_numpy(stds).float() + torch.from_numpy(means).float()140            np.save(output_dir / f"prediction_{index:05d}.npy", prediction.cpu().numpy())141            processed += 1142            if max_batches >= 0 and index + 1 >= max_batches:143                break144    print({"output_dir": str(output_dir), "batches": processed})145 146 147if __name__ == "__main__":148    main()149