File size: 2,036 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 | """Create diagnostic plots from Samudra NPZ rollout outputs."""
try:
from ._bootstrap import ROOT
except ImportError:
from _bootstrap import ROOT
import argparse
from pathlib import Path
import matplotlib.pyplot as plt
import numpy as np
DEPTHS = (2.5, 10.0, 22.5, 40.0, 65.0, 105.0, 165.0, 250.0, 375.0, 550.0, 775.0, 1050.0, 1400.0, 1850.0, 2400.0, 3100.0, 4000.0, 5000.0, 6000.0)
PROGNOSTIC_CHANNELS = tuple(
f"{name}_{level}"
for name in ("uo", "vo", "thetao", "so")
for level in range(len(DEPTHS))
) + ("zos",)
def plot_rollout(prediction_path: str, output_dir: str) -> None:
with np.load(prediction_path) as data:
predictions = np.asarray(data["predictions"])
output = Path(output_dir)
output.mkdir(parents=True, exist_ok=True)
fig, axes = plt.subplots(1, 2, figsize=(12, 4), constrained_layout=True)
axes[0].imshow(predictions[0, PROGNOSTIC_CHANNELS.index("zos")], cmap="coolwarm", aspect="auto")
axes[0].set_title("SSH forecast")
axes[1].imshow(predictions[0, PROGNOSTIC_CHANNELS.index("thetao_0")], cmap="turbo", aspect="auto")
axes[1].set_title("Surface potential temperature")
fig.savefig(output / "forecast_maps.png", dpi=150)
plt.close(fig)
indices = [PROGNOSTIC_CHANNELS.index(f"thetao_{i}") for i in range(len(DEPTHS))]
profile = predictions[:, indices].mean(axis=(0, 2, 3))
fig, ax = plt.subplots(figsize=(5, 5), constrained_layout=True)
ax.plot(profile, DEPTHS, marker="o")
ax.invert_yaxis()
ax.set(xlabel="potential temperature", ylabel="depth (m)", title="Mean temperature profile")
fig.savefig(output / "temperature_profile.png", dpi=150)
plt.close(fig)
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--prediction", default="./result/output/prediction.npz")
parser.add_argument("--output-dir", default="./result")
args = parser.parse_args()
plot_rollout(args.prediction, args.output_dir)
if __name__ == "__main__":
main()
print("save to ./result/")
|