| """Create quick field images from OneForecast prediction files.""" |
|
|
| from __future__ import annotations |
|
|
| from pathlib import Path |
| import argparse |
| import numpy as np |
| import yaml |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--config", type=Path, default=Path("conf/config.yaml")) |
| args = parser.parse_args() |
| with args.config.open("r", encoding="utf-8") as handle: |
| config = yaml.safe_load(handle) |
| root = args.config.resolve().parent.parent |
| input_dir = Path(config["visualization"]["input_dir"]) |
| output_dir = Path(config["visualization"]["output_dir"]) |
| if not input_dir.is_absolute(): |
| input_dir = root / input_dir |
| if not output_dir.is_absolute(): |
| output_dir = root / output_dir |
| output_dir.mkdir(parents=True, exist_ok=True) |
| files = sorted(input_dir.glob("prediction_*.npy")) |
| if not files: |
| raise SystemExit(f"No prediction files found in {input_dir}") |
| import matplotlib.pyplot as plt |
|
|
| channels = config["visualization"].get("channels", [0]) |
| for source in files: |
| prediction = np.load(source) |
| if prediction.shape != (1, 69, 120, 240): |
| raise ValueError(f"Expected official prediction shape [1, 69, 120, 240], got {prediction.shape}") |
| field = prediction[0] |
| for channel in channels: |
| if channel < 0 or channel >= field.shape[0]: |
| raise ValueError(f"Channel {channel} is outside prediction shape {field.shape}") |
| figure, axis = plt.subplots(figsize=(8, 3.5)) |
| image = axis.imshow(field[channel], cmap="coolwarm", aspect="auto") |
| axis.set_title(f"{source.stem}, channel {channel}") |
| axis.set_xlabel("longitude index") |
| axis.set_ylabel("latitude index") |
| figure.colorbar(image, ax=axis, shrink=0.8) |
| figure.tight_layout() |
| figure.savefig(output_dir / f"{source.stem}_ch{channel}.png", dpi=160) |
| plt.close(figure) |
| print({"input_dir": str(input_dir), "output_dir": str(output_dir), "files": len(files)}) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
