from __future__ import annotations import argparse import sys from pathlib import Path from typing import Any PROJECT_ROOT = Path(__file__).resolve().parent.parent MODEL_ROOT = PROJECT_ROOT / "model" SCRIPT_ROOT = PROJECT_ROOT / "scripts" for path in (PROJECT_ROOT, MODEL_ROOT, SCRIPT_ROOT): if str(path) not in sys.path: sys.path.insert(0, str(path)) import h5py import numpy as np import torch import yaml from stormer import StormCast, edm_heun_sample from data_loader import StormCastDataset def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="Run StormCast inference") parser.add_argument("--config", type=Path, default=Path("conf/config.yaml")) parser.add_argument("--regression-weights", type=Path) parser.add_argument("--diffusion-weights", type=Path) parser.add_argument("--output", type=Path) parser.add_argument("--num-steps", type=int) parser.add_argument("--diffusion-steps", type=int) parser.add_argument("--seed", type=int) return parser.parse_args() def main() -> None: args = parse_args() config_path = args.config.resolve() with config_path.open("r", encoding="utf-8") as handle: config = yaml.safe_load(handle) _resolve_paths(config, config_path.parent.parent) inference_config = config["inference"] regression_weights = args.regression_weights or Path( config["model"]["regression_weights"] ) diffusion_weights = args.diffusion_weights or Path( config["model"]["diffusion_weights"] ) output = args.output or Path(inference_config["output_dir"]) / "forecast.h5" run_inference( config=config, regression_weights=regression_weights, diffusion_weights=diffusion_weights, output=output, num_steps=args.num_steps or inference_config["num_steps"], diffusion_steps=args.diffusion_steps or inference_config["diffusion_steps"], seed=config["project"]["seed"] if args.seed is None else args.seed, ) @torch.no_grad() def run_inference( config: dict[str, Any], regression_weights: Path, diffusion_weights: Path, output: Path, num_steps: int, diffusion_steps: int, seed: int, ) -> Path: if not torch.cuda.is_available(): raise RuntimeError("StormCast inference requires a CUDA/HIP device") if num_steps < 1: raise ValueError("num_steps must be at least 1") device = torch.device("cuda", 0) torch.cuda.set_device(device) data_config = config["data"] inference_config = config["inference"] split = inference_config["split"] years_key = f"{split}_years" if years_key not in data_config: raise ValueError(f"Unknown inference split: {split}") dataset = StormCastDataset( data_root=data_config["root_dir"], years=data_config[years_key], era5_variables=data_config["era5_variables"], state_variables=data_config["state_variables"], invariant_variables=data_config["invariant_variables"], image_size=data_config["image_size"], input_steps=data_config["input_steps"], output_steps=data_config["output_steps"], normalize=data_config["normalize"], ) if num_steps > len(dataset): raise ValueError(f"Requested {num_steps} steps but split contains {len(dataset)}") regression = _load_weights(regression_weights, "regression").to(device).eval() diffusion = _load_weights(diffusion_weights, "diffusion").to(device).eval() model = StormCast(regression, diffusion) generator = torch.Generator(device=device).manual_seed(seed) output.parent.mkdir(parents=True, exist_ok=True) first = dataset[0] state = first["state"][0].unsqueeze(0).to(device, dtype=torch.float32) invariant = first["invariant"].to(device, dtype=torch.float32) state_channels = len(data_config["state_variables"]) background_channels = len(data_config["era5_variables"]) height, width = data_config["image_size"] if [height, width] != list(config["model"]["image_size"]): raise ValueError("Data and model image sizes must match") if list(data_config["era5_image_size"]) != [721, 1440]: raise ValueError("ERA5 grid must be 721 x 1440") with h5py.File(output, "w") as handle: handle.attrs["normalized"] = bool(data_config["normalize"]) handle.attrs["seed"] = seed handle.attrs["diffusion_steps"] = diffusion_steps handle.attrs["sigma_min"] = inference_config["sigma_min"] handle.attrs["sigma_max"] = inference_config["sigma_max"] handle.attrs["rho"] = inference_config["rho"] handle.attrs["regression_weights"] = str(regression_weights.resolve()) handle.attrs["diffusion_weights"] = str(diffusion_weights.resolve()) handle.attrs["state_variables"] = np.asarray( data_config["state_variables"], dtype="S" ) handle.attrs["background_variables"] = np.asarray( data_config["era5_variables"], dtype="S" ) prediction_store = handle.create_dataset( "prediction", (num_steps, state_channels, height, width), dtype="f4" ) regression_store = handle.create_dataset( "regression", (num_steps, state_channels, height, width), dtype="f4" ) target_store = handle.create_dataset( "target", (num_steps, state_channels, height, width), dtype="f4" ) background_store = handle.create_dataset( "background", (num_steps, background_channels, height, width), dtype="f4" ) time_store = handle.create_dataset("time_index", (num_steps,), dtype="i8") for index in range(num_steps): sample = dataset[index] background = sample["background"].unsqueeze(0).to( device, dtype=torch.float32 ) regression_prediction = model.predict_regression( state, background, invariant ) condition = model.diffusion_condition( state, regression_prediction, invariant ) residual = edm_heun_sample( diffusion, condition, output_channels=state_channels, num_steps=diffusion_steps, sigma_min=inference_config["sigma_min"], sigma_max=inference_config["sigma_max"], rho=inference_config["rho"], generator=generator, ) prediction = regression_prediction + residual prediction_store[index] = prediction[0].cpu().numpy() regression_store[index] = regression_prediction[0].cpu().numpy() target_store[index] = sample["state"][1].numpy() background_store[index] = sample["background"].numpy() time_store[index] = int(np.asarray(sample["time_index"]).reshape(-1)[-1]) state = prediction print(f"forecast_step={index + 1}/{num_steps}") print(f"output={output}") return output def _load_weights(path: Path, kind: str) -> torch.nn.Module: if not path.is_file(): raise FileNotFoundError( f"Missing {kind} weights: {path}. " "Train the corresponding stage first or pass an explicit weight path." ) checkpoint = torch.load(path, map_location="cpu", weights_only=False) saved_config = checkpoint.get("config") if saved_config is None: raise ValueError("Project .pt weights must include their training config") model = _build_stage_models_from_config(saved_config, kind) model.load_state_dict(checkpoint["model_state_dict"], strict=True) return model def _build_stage_models_from_config(config: dict[str, Any], kind: str) -> torch.nn.Module: from stormer import build_diffusion_model, build_regression_model data_config = config["data"] model_config = config["model"] common = { "image_size": model_config["image_size"], "state_channels": len(data_config["state_variables"]), "invariant_channels": len(data_config["invariant_variables"]), "model_channels": model_config["model_channels"], "channel_mult": model_config["channel_mult"], "num_blocks": model_config["num_blocks"], "attn_resolutions": model_config["attention_resolutions"], } if kind == "regression": return build_regression_model( **common, background_channels=len(data_config["era5_variables"]) ) return build_diffusion_model(**common) def _resolve_paths(config: dict[str, Any], project_root: Path) -> None: for section, key in (("data", "root_dir"), ("inference", "output_dir")): path = Path(config[section][key]) if not path.is_absolute(): config[section][key] = str((project_root / path).resolve()) for key in ("regression_weights", "diffusion_weights"): path = Path(config["model"][key]) if not path.is_absolute(): config["model"][key] = str((project_root / path).resolve()) if __name__ == "__main__": main()