plotStats.py

# -----------------------------------------------------------------------------
# plot_stats_mit_optimized.py – compact CLI tool to visualise MITgcm monitor output
# -----------------------------------------------------------------------------
# Usage  (from shell or notebook):
#     python plot_stats_mit_optimized.py bay/statsMIT.txt [nobay/statsMIT.txt] \
#            --prefix realistic --outdir ./figs
# -----------------------------------------------------------------------------
# Key changes vs. legacy script
#   • argparse interface (no manual path editing)
#   • single generic plotting helpers (no code repetition)
#   • pandas‑only arithmetic for bay‑nobay differences
#   • figures saved in <outdir>/<prefix>_<tag>.png
#   • all constants (column names, variable groups) at top – easy to adapt
# -----------------------------------------------------------------------------
from __future__ import annotations

import argparse
import textwrap
from pathlib import Path
from typing import Iterable, Sequence

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib as mpl

# -----------------------------------------------------------------------------
# 1.  Constants & configuration
# -----------------------------------------------------------------------------
COLS = [
    # time
    "time_tsnumber", "time_secondsf",
    # dynamic variables (min/max/mean/sd)
    "eta_max", "eta_min", "eta_mean", "eta_sd",
    "uvel_max", "uvel_min", "uvel_mean", "uvel_sd",
    "vvel_max", "vvel_min", "vvel_mean", "vvel_sd",
    "wvel_max", "wvel_min", "wvel_mean", "wvel_sd",
    "theta_max", "theta_min", "theta_mean", "theta_sd",
    "salt_max", "salt_min", "salt_mean", "salt_sd",
    # CFL
    "trAdv_CFL_u_max", "trAdv_CFL_v_max", "trAdv_CFL_w_max",
    "advcfl_u_max", "advcfl_v_max", "advcfl_w_max", "advcfl_W_hf_max",
    # energy & vorticity & expansivity
    "pe_b_mean", "ke_max", "ke_mean", "ke_vol",
    "vort_r_min", "vort_r_max",
    "vort_a_mean", "vort_a_sd", "vort_p_mean", "vort_p_sd",
    "surfExpan_theta_mean", "surfExpan_salt_mean",
]

DYN_GROUPS = {
    "eta": (r"$\eta$ (m)", "steelblue"),
    "uvel": (r"$u$ (m s$^{-1}$)", "darkorange"),
    "vvel": (r"$v$ (m s$^{-1}$)", "seagreen"),
    "wvel": (r"$w$ (m s$^{-1}$)", "purple"),
    "theta": (r"$\theta$ (°C)", "firebrick"),
    "salt": ("salinity", "gray"),
}

mpl.rcParams.update({
    "figure.autolayout": True,
    "axes.grid": True,
    "grid.color": "0.9",
})

# -----------------------------------------------------------------------------
# 2.  I/O helpers
# -----------------------------------------------------------------------------

def read_stats(fname: Path) -> pd.DataFrame:
    """Return DataFrame with shape (nsteps, nvars) from monitorMITgcm txt."""
    blocks = Path(fname).read_text().strip().split("\n\n")
    data = np.vstack([list(map(float, blk.split())) for blk in blocks]).T
    return pd.DataFrame(data, columns=COLS)


def diff_bay_nobay(bay: pd.DataFrame, nobay: pd.DataFrame) -> pd.DataFrame:
    """Return df_diff = bay − nobay (SD combined as sqrt(sum of squares))."""
    # align on time columns first
    keys = ["time_tsnumber", "time_secondsf"]
    merged = bay.merge(nobay, on=keys, suffixes=("_bay", "_nb"), how="inner")

    diff = {
        k: merged[k] if k in keys else (
            np.sqrt(merged[f"{k}_bay"] ** 2 + merged[f"{k}_nb"] ** 2)
            if k.endswith("_sd")
            else merged[f"{k}_bay"] - merged[f"{k}_nb"]
        )
        for k in COLS
    }
    return pd.DataFrame(diff)

# -----------------------------------------------------------------------------
# 3.  Plotting primitives
# -----------------------------------------------------------------------------

def _error_patch(ax, x: np.ndarray, ymean: np.ndarray, yerr: np.ndarray, **kw):
    """Thin wrapper drawing mean ± err as patch & mean as line/dots."""
    ax.fill_between(x, ymean - yerr, ymean + yerr, alpha=0.2, **kw)
    ax.plot(x, ymean, marker="o", lw=1, **kw)


def quick_stats_panel(ax: mpl.axes.Axes, df: pd.DataFrame, var: str, label: str, color: str):
    x = df.time_secondsf / 3600.0
    _error_patch(ax, x, df[f"{var}_mean"], df[f"{var}_sd"], color=color)
    ax.plot(x, df[f"{var}_min"], color=color, alpha=0.4, lw=0.8)
    ax.plot(x, df[f"{var}_max"], color=color, alpha=0.4, lw=0.8)
    ax.set_ylabel(label)
    ax.set_xlim(x.iloc[0], x.iloc[-1])
    ax.text(0.99, 0.9, var, transform=ax.transAxes, ha="right", va="top", fontsize=9)


# -----------------------------------------------------------------------------
# 4.  Figure factories
# -----------------------------------------------------------------------------

def fig_dynamic(df: pd.DataFrame, title: str, outfile: Path):
    fig, axes = plt.subplots(len(DYN_GROUPS), 1, figsize=(7, 9), sharex=True)
    for ax, (var, (ylab, col)) in zip(axes, DYN_GROUPS.items()):
        quick_stats_panel(ax, df, var, ylab, col)
    axes[-1].set_xlabel("time (h)")
    axes[0].set_title(title)
    fig.savefig(outfile, dpi=300)
    plt.close(fig)


def fig_cfl(df: pd.DataFrame, title: str, outfile: Path):
    time = df.time_secondsf / 3600.0
    cfl_groups = {
        "trAdv_CFL": ["u", "v", "w"],
        "advcfl": ["u", "v", "w", "W_hf"],
    }
    colors = ["coral", "olivedrab", "goldenrod", "brown"]
    fig, axes = plt.subplots(2, 1, figsize=(7, 5), sharex=True)
    for ax, (root, vars_) in zip(axes, cfl_groups.items()):
        for var, col in zip(vars_, colors):
            ax.plot(time, df[f"{root}_{var}_max"], marker="^", lw=0.8, color=col, label=var)
        ax.set_ylabel(root)
        ax.legend(ncol=len(vars_), fontsize=7)
    axes[-1].set_xlabel("time (h)")
    axes[0].set_title(title)
    fig.savefig(outfile, dpi=300)
    plt.close(fig)


# -----------------------------------------------------------------------------
# 5.  CLI driver
# -----------------------------------------------------------------------------

def parse_cli(argv: Sequence[str] | None = None) -> argparse.Namespace:
    p = argparse.ArgumentParser(
        prog="plot_stats_mit",
        formatter_class=argparse.RawTextHelpFormatter,
        description=textwrap.dedent(
            """
            Quick visualisation of MITgcm monitor stats.
            Pass one file (bay) or two files (bay nobay) to also compute Δ.
            """
        ),
    )
    p.add_argument("bay", type=Path, help="statsMIT.txt for bay run")
    p.add_argument("nobay", type=Path, nargs="?", help="statsMIT.txt for nobay run (optional)")
    p.add_argument("--prefix", default="run", help="prefix for output PNG files")
    p.add_argument("--outdir", type=Path, default=Path("./figs"))
    return p.parse_args(argv)


def main():
    args = parse_cli()
    args.outdir.mkdir(parents=True, exist_ok=True)

    df_bay = read_stats(args.bay)
    fig_dynamic(df_bay, f"Dynamics – {args.prefix} (bay)", args.outdir / f"{args.prefix}_dyn_bay.png")
    fig_cfl(df_bay, f"CFL – {args.prefix} (bay)", args.outdir / f"{args.prefix}_cfl_bay.png")

    if args.nobay is not None:
        df_nobay = read_stats(args.nobay)
        if len(df_bay) != len(df_nobay):
            raise SystemExit("❌ Bay & noBay files have different length")
        df_diff = diff_bay_nobay(df_bay, df_nobay)
        fig_dynamic(df_diff, f"Dynamics Δ(bay - nobay) – {args.prefix}", args.outdir / f"{args.prefix}_dyn_diff.png")
        fig_cfl(df_diff, f"CFL Δ(bay - nobay) – {args.prefix}", args.outdir / f"{args.prefix}_cfl_diff.png")

    print(f"✔ Figures stored in {args.outdir.resolve()}")


if __name__ == "__main__":
    main()