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