File size: 1,358 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 | from __future__ import annotations
import argparse
from pathlib import Path
import matplotlib.pyplot as plt
import xarray as xr
from common import load_config, resolve_path
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--config", default="conf/config.yaml")
parser.add_argument("--input", default=None)
parser.add_argument("--channel", default=None)
args = parser.parse_args()
cfg = load_config(args.config)
viz = cfg["visualization"]
source = Path(args.input) if args.input else resolve_path(viz["input_file"], cfg)
channel = args.channel or viz["channel"]
data = xr.open_dataarray(source).sel(channel=channel)
if "time" in data.dims:
data = data.isel(time=-1)
output = resolve_path(viz["output_file"], cfg)
output.parent.mkdir(parents=True, exist_ok=True)
fig, ax = plt.subplots(figsize=(12, 5), constrained_layout=True)
image = ax.pcolormesh(data.lon, data.lat, data, shading="auto", cmap=viz["cmap"])
valid_time = str(data.time.values) if "time" in data.coords else ""
ax.set(title=f"FuXi 2.1 {channel} | {valid_time}", xlabel="Longitude", ylabel="Latitude")
fig.colorbar(image, ax=ax, label=channel)
fig.savefig(output, dpi=160)
plt.close(fig)
print(f"Saved visualization to {output}")
if __name__ == "__main__":
main()
|