| """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/") |
|
|
