File size: 4,686 Bytes
929e312 | 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 | """Autoregressive inference for the paper-version Samudra model."""
try:
from ._bootstrap import ROOT
except ImportError:
from _bootstrap import ROOT
import argparse
from pathlib import Path
import numpy as np
import torch
import yaml
from model.samudra import build_model
STATE_CHANNELS = 77
BOUNDARY_CHANNELS = 4
def load_data(path: str | Path) -> tuple[np.ndarray, np.ndarray]:
"""Load and validate native Samudra time-major arrays."""
with np.load(path) as data:
prognostic = np.asarray(data["prognostic"], dtype=np.float32)
boundary = np.asarray(data["boundary"], dtype=np.float32)
if prognostic.ndim != 4 or prognostic.shape[1] != STATE_CHANNELS:
raise ValueError("prognostic must have shape [time, 77, lat, lon]")
if boundary.ndim != 4 or boundary.shape[1] != BOUNDARY_CHANNELS:
raise ValueError("boundary must have shape [time, 4, lat, lon]")
if prognostic.shape[0] != boundary.shape[0] or prognostic.shape[2:] != boundary.shape[2:]:
raise ValueError("prognostic and boundary time/grid dimensions must match")
return prognostic, boundary
def load_checkpoint(model: torch.nn.Module, path: str, device: torch.device) -> None:
try:
checkpoint = torch.load(path, map_location=device, weights_only=True)
except TypeError:
checkpoint = torch.load(path, map_location=device)
state = checkpoint.get("model", checkpoint.get("state_dict", checkpoint))
model.load_state_dict(state)
@torch.no_grad()
def rollout(model, prognostic, boundary, steps, device):
# Training uses one consecutive boundary frame for each recurrent pass.
available_steps = boundary.shape[0] - 1
if prognostic.shape[0] < 2 or available_steps < 1:
raise ValueError(
"rollout requires at least two states and two boundary time samples"
)
if steps > available_steps:
raise ValueError(
f"requested {steps} model steps, but input data provides only {available_steps}; "
f"use --steps <= {available_steps} or provide a longer data file"
)
previous = torch.from_numpy(prognostic[0]).to(device)
current = torch.from_numpy(prognostic[1]).to(device)
predictions = []
for step in range(steps):
forcing = torch.from_numpy(boundary[step + 1]).to(device)
inputs = torch.cat((previous, current, forcing), dim=0).unsqueeze(0)
prediction = model(inputs).squeeze(0)
previous, current = prediction[:STATE_CHANNELS], prediction[STATE_CHANNELS:]
predictions.extend((previous.cpu().numpy(), current.cpu().numpy()))
return np.stack(predictions)
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--config", default="./conf/config.yaml")
parser.add_argument("--data", default="./data/test.npz")
parser.add_argument("--checkpoint", default=None)
parser.add_argument("--output", default=None)
parser.add_argument("--steps", type=int, default=None)
parser.add_argument("--device", default=None)
args = parser.parse_args()
if args.steps is not None and args.steps < 1:
raise ValueError("--steps must be positive")
with open(args.config, encoding="utf-8") as handle:
config = yaml.safe_load(handle)
device = torch.device(args.device or ("cuda" if torch.cuda.is_available() else "cpu"))
model = build_model(config).to(device)
checkpoint = args.checkpoint or config["inference"].get("checkpoint") or "./data/checkpoints/model_bak.pth"
output_path = args.output or config["inference"].get("output_dir", "./result/output")
if Path(output_path).suffix != ".npz":
output_path = str(Path(output_path) / "prediction.npz")
load_checkpoint(model, checkpoint, device)
model.eval()
prognostic, boundary = load_data(args.data)
configured_steps = int(config["inference"].get("rollout_steps", 1))
if configured_steps < 1:
raise ValueError("inference.rollout_steps must be positive")
available_steps = boundary.shape[0] - 1
steps = args.steps if args.steps is not None else min(configured_steps, available_steps)
if args.steps is None and steps < configured_steps:
print(
f"input data supports {available_steps} model steps; "
f"using rollout_steps={steps}"
)
predictions = rollout(model, prognostic, boundary, steps, device)
output = Path(output_path)
output.parent.mkdir(parents=True, exist_ok=True)
np.savez_compressed(output, predictions=predictions, forcing=boundary[1 : steps + 1])
print(f"saved predictions: {output} shape={predictions.shape}")
if __name__ == "__main__":
main()
|