scripts/inference.py

from __future__ import annotations

import argparse
from pathlib import Path
import sys
from typing import Any

import numpy as np
import torch
from torch.utils.data import DataLoader


PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))
SCRIPTS_DIR = PROJECT_ROOT / "scripts"
if str(SCRIPTS_DIR) not in sys.path:
    sys.path.insert(0, str(SCRIPTS_DIR))

from model.checkpoint import load_checkpoint
from era5_adapter import WMAEERA5Dataset
from train import build_model, device_summary, load_config, resolve_device, resolve_project_path, validate_config


def build_inference_dataset(config: dict[str, Any], config_path: Path) -> WMAEERA5Dataset:
    data = config["data"]
    if int(data["input_steps"]) != 1 or int(data["output_steps"]) != 1:
        raise ValueError("W-MAE inference currently requires input_steps=1 and output_steps=1.")
    if not data["variables"]:
        raise ValueError(
            "data.variables is empty. Supply the verified, explicitly ordered 20-channel list before inference."
        )
    return WMAEERA5Dataset(
        dataset_dir=resolve_project_path(data["dataset_dir"], config_path),
        years=data["test_years"],
        variables=data["variables"],
        task="pretrain",
        input_steps=1,
        output_steps=1,
        normalize=bool(data["normalize"]),
    )


def sample_time_index(time_index: Any, batch_index: int, batch_size: int) -> Any:
    """Extract one sample from DataLoader's sequence-major default collation."""
    if isinstance(time_index, (list, tuple)):
        if time_index and isinstance(time_index[0], (list, tuple)):
            return time_index[0][batch_index]
        if len(time_index) == batch_size:
            return time_index[batch_index]
    return time_index


def run_inference(
    model: torch.nn.Module,
    loader: DataLoader,
    output_dir: Path,
    mask_ratio: float,
    device: torch.device,
    max_samples: int | None = None,
    log_interval: int = 1,
) -> int:
    output_dir.mkdir(parents=True, exist_ok=True)
    model.eval()
    written = 0
    with torch.no_grad():
        total_batches = len(loader)
        for batch_index, (inputs, _, _, step_idx, time_index) in enumerate(loader, start=1):
            if inputs.ndim != 4:
                raise ValueError(f"W-MAE inference expects a 4D batch, got {tuple(inputs.shape)}.")
            inputs = inputs.to(device)
            result = model(inputs, mask_ratio=mask_ratio)
            reconstruction = model.unpatchify(result.prediction)
            error = reconstruction - inputs
            mask = result.mask.reshape(result.mask.shape[0], *model.patch_embed.grid_size)
            for sample_index in range(inputs.shape[0]):
                if max_samples is not None and written >= max_samples:
                    print(
                        f"inference batch={batch_index}/{total_batches} samples_written={written}",
                        flush=True,
                    )
                    return written
                sample_time = sample_time_index(time_index, sample_index, inputs.shape[0])
                if isinstance(sample_time, (list, tuple)):
                    sample_time = sample_time[0]
                np.savez_compressed(
                    output_dir / f"sample_{written:06d}.npz",
                    input=inputs[sample_index].cpu().numpy(),
                    reconstruction=reconstruction[sample_index].cpu().numpy(),
                    error=error[sample_index].cpu().numpy(),
                    mask=mask[sample_index].cpu().numpy(),
                    step_idx=np.asarray(step_idx[sample_index].item()),
                    time_index=np.asarray(sample_time),
                )
                written += 1
            if batch_index == 1 or batch_index % log_interval == 0 or batch_index == total_batches:
                print(
                    f"inference batch={batch_index}/{total_batches} samples_written={written}",
                    flush=True,
                )
    return written


def main() -> None:
    parser = argparse.ArgumentParser(description="Run W-MAE reconstruction inference on ERA5 samples.")
    parser.add_argument("--config", type=Path, default=PROJECT_ROOT / "conf" / "config.yaml")
    parser.add_argument("--checkpoint", type=Path, default="./data/checkpoint/model_bak.pth")
    parser.add_argument("--checkpoint-source-root", type=Path, default=None)
    parser.add_argument("--non-strict-checkpoint", action="store_true")
    parser.add_argument("--output-dir", type=Path, default=None)
    parser.add_argument(
        "--device",
        default="auto",
        help="Inference device: auto (default), cuda (AMD DCU through HIP), or cpu.",
    )
    parser.add_argument("--mask-ratio", type=float, default=None)
    parser.add_argument("--max-samples", type=int, default=None)
    parser.add_argument("--log-interval", type=int, default=1)
    args = parser.parse_args()
    if args.max_samples is not None and args.max_samples <= 0:
        raise ValueError("--max-samples must be positive when provided.")
    if args.log_interval <= 0:
        raise ValueError("--log-interval must be positive.")

    config_path = args.config.resolve()
    print(
        f"inference started: config={config_path} checkpoint={args.checkpoint} device={args.device}",
        flush=True,
    )
    config = load_config(config_path)
    validate_config(config)
    if args.checkpoint is None:
        raise ValueError("Inference requires an explicit --checkpoint path.")

    device = resolve_device(args.device)
    print(f"runtime: {device_summary(device)}", flush=True)
    print("building model", flush=True)
    model = build_model(config).to(device)
    print(f"model ready: parameters_device={next(model.parameters()).device}", flush=True)
    source_root = (
        resolve_project_path(args.checkpoint_source_root, config_path)
        if args.checkpoint_source_root
        else None
    )
    report = load_checkpoint(
        model,
        resolve_project_path(args.checkpoint, config_path),
        strict=not args.non_strict_checkpoint,
        map_location=device,
        source_root=source_root,
    )
    print(f"loaded checkpoint: {report}", flush=True)

    data_config = config["data"]
    print("building inference dataset", flush=True)
    dataset = build_inference_dataset(config, config_path)
    loader = DataLoader(
        dataset,
        batch_size=int(data_config["batch_size"]),
        shuffle=False,
        num_workers=int(data_config["num_workers"]),
    )
    print(f"dataset ready: samples={len(dataset)} batches={len(loader)}", flush=True)
    output_dir = resolve_project_path(args.output_dir, config_path) if args.output_dir else (
        resolve_project_path(config["project"]["output_dir"], config_path) / "inference"
    )
    mask_ratio = float(config["model"]["mask_ratio"] if args.mask_ratio is None else args.mask_ratio)
    if mask_ratio not in {0.0, 0.75}:
        raise ValueError("W-MAE supports mask_ratio 0.0 or 0.75 only.")
    print(f"inference running: mask_ratio={mask_ratio} output_dir={output_dir}", flush=True)
    count = run_inference(model, loader, output_dir, mask_ratio, device, args.max_samples, args.log_interval)
    print(f"wrote {count} reconstruction samples to {output_dir}", flush=True)
    print("inference completed successfully", flush=True)


if __name__ == "__main__":
    main()