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