File size: 2,125 Bytes
04fe5fd | 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 | """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()
|