File size: 5,977 Bytes
9191802 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 | """Autoregressive inference using a project-produced FuXi 2.1 checkpoint."""
from __future__ import annotations
import argparse
from datetime import datetime, timedelta
import numpy as np
import torch
import xarray as xr
from onescience.datapipes.climate.era5 import ERA5Dataset
from common import load_config, resolve_path
from model.FuXi21 import FuXi21
from variables import c85_from_config
CHECKPOINT_FORMAT = "fuxi21_reconstructed_checkpoint_v1"
def select_device(requested: str) -> torch.device:
if requested not in {"auto", "cpu", "cuda"}:
raise ValueError("inference.device must be auto, cpu, or cuda")
if requested == "cuda" or (requested == "auto" and torch.cuda.is_available()):
if not torch.cuda.is_available():
raise RuntimeError("inference.device=cuda, but no CUDA/HIP device is available")
return torch.device("cuda")
return torch.device("cpu")
def load_array(path_value: str | None, cfg: dict, shape: tuple[int, ...], name: str) -> torch.Tensor:
if path_value is None:
raise ValueError(f"model.{name}_file is required outside the smoke profile")
path = resolve_path(path_value, cfg)
value = torch.from_numpy(np.load(path)).float()
if tuple(value.shape) != shape:
raise ValueError(f"{name} must have shape {shape}, got {tuple(value.shape)}")
return value
def build_model(cfg: dict) -> FuXi21:
model_cfg = cfg["model"]
profile_name = model_cfg["profile"]
profile = model_cfg["profiles"][profile_name]
height, width = profile["grid_size"]
if profile_name == "smoke":
static_fields = torch.zeros(6, height, width)
channel_mask = torch.ones(85, height, width)
else:
static_fields = load_array(model_cfg["static_fields_file"], cfg, (6, height, width), "static_fields")
channel_mask = load_array(model_cfg["channel_mask_file"], cfg, (85, height, width), "channel_mask")
return FuXi21(
static_fields,
channel_mask,
activation_checkpointing=False,
**profile,
)
def temporal_features(valid_time: datetime, step: int, device: torch.device) -> tuple[torch.Tensor, ...]:
return (
torch.tensor([step], device=device, dtype=torch.float32),
torch.tensor([(valid_time.hour * 60 + valid_time.minute) / 1440], device=device),
torch.tensor([min(365, valid_time.timetuple().tm_yday) / 365], device=device),
)
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--config", default="conf/config.yaml")
parser.add_argument("--device", choices=("auto", "cpu", "cuda"), default=None)
parser.add_argument("--preflight-only", action="store_true")
args = parser.parse_args()
cfg = load_config(args.config)
if cfg.get("protocol") != "non_official_protocol":
raise ValueError("Inference config must declare protocol: non_official_protocol")
infer_cfg = cfg["inference"]
channels, diagnostics = c85_from_config(cfg)
checkpoint_path = resolve_path(infer_cfg["checkpoint"], cfg)
if not checkpoint_path.is_file():
raise FileNotFoundError(f"Project checkpoint not found: {checkpoint_path}")
device = select_device(args.device or infer_cfg["device"])
checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=True)
if checkpoint.get("format") != CHECKPOINT_FORMAT:
raise ValueError(f"Checkpoint must use format {CHECKPOINT_FORMAT}")
if checkpoint.get("protocol") != "non_official_protocol":
raise ValueError("Checkpoint protocol must be non_official_protocol")
if checkpoint.get("model_profile") != cfg["model"]["profile"]:
raise ValueError("Checkpoint model profile does not match the configured model profile")
if args.preflight_only:
print(f"checkpoint={checkpoint_path}, profile={checkpoint['model_profile']}, device={device}")
return
model = build_model(cfg).to(device)
model.load_state_dict(checkpoint["model"])
model.eval()
split = infer_cfg["split"]
split_cfg = cfg["data"]["splits"][split]
dataset = ERA5Dataset(
dataset_dir=str(resolve_path(cfg["paths"]["data_root"], cfg)),
used_years=split_cfg["years"],
used_variables=channels,
input_steps=cfg["data"]["input_steps"],
output_steps=cfg["data"]["output_steps"],
normalize=True,
)
state, _, _, _, time_index = dataset[0]
crop_size = cfg["data"]["crop_size"]
if crop_size is not None:
state = state[..., : crop_size[0], : crop_size[1]]
state = state.unsqueeze(0).to(device)
valid_time = datetime.strptime(time_index[-1], "%Y%m%d%H")
interval = timedelta(hours=cfg["data"]["time_step_hours"])
diagnostic_indices = [channels.index(name) for name in diagnostics]
forecasts = []
valid_times = []
for step in range(infer_cfg["steps"]):
with torch.inference_mode():
state = model(state, *temporal_features(valid_time, step, device))
forecasts.append(state[:, -1].float().cpu().numpy()[0])
valid_times.append(np.datetime64(valid_time))
if infer_cfg["zero_diagnostic_feedback"]:
state[:, -1, diagnostic_indices] = 0
valid_time += interval
height, width = forecasts[0].shape[-2:]
output_path = resolve_path(infer_cfg["output_file"], cfg)
output_path.parent.mkdir(parents=True, exist_ok=True)
xr.DataArray(
np.stack(forecasts),
dims=("time", "channel", "lat", "lon"),
coords={
"time": valid_times,
"channel": channels,
"lat": np.linspace(90, -90, height),
"lon": np.arange(width) * (360 / width),
},
attrs={"checkpoint_format": CHECKPOINT_FORMAT, "protocol": cfg["protocol"]},
name="forecast",
).to_netcdf(output_path)
print(f"Saved {len(forecasts)} forecast step(s) to {output_path}")
if __name__ == "__main__":
main()
|