|
| 1 | +#!/usr/bin/env python3 |
| 2 | +""" |
| 3 | +plot_diagnostics_binary.py |
| 4 | +========================== |
| 5 | +Sidecar plotter for MITgcm binary diagnostics output. |
| 6 | +
|
| 7 | +Watches an experiment directory (e.g. repeat-year-50/) for new state3D/state2D |
| 8 | +binary files and renders surface field plots (SST, SSS, SSH, KE) into each |
| 9 | +run's plots/ subdirectory. |
| 10 | +
|
| 11 | +Designed to run on the login node alongside SLURM jobs. |
| 12 | +
|
| 13 | +Usage: |
| 14 | + nohup uv run python spectre_utils/plot_diagnostics_binary.py \ |
| 15 | + simulations/glorysv12-curvilinear repeat-year-50 \ |
| 16 | + --poll 120 & |
| 17 | +
|
| 18 | + # Or watch a single run: |
| 19 | + uv run python spectre_utils/plot_diagnostics_binary.py \ |
| 20 | + simulations/glorysv12-curvilinear repeat-year-50 --run 001 --poll 0 |
| 21 | +""" |
| 22 | + |
| 23 | +import argparse |
| 24 | +import os |
| 25 | +import re |
| 26 | +import sys |
| 27 | +import time |
| 28 | +from datetime import datetime, timedelta |
| 29 | +from pathlib import Path |
| 30 | + |
| 31 | +import numpy as np |
| 32 | + |
| 33 | +_plt = None |
| 34 | + |
| 35 | + |
| 36 | +def _import_mpl(): |
| 37 | + global _plt |
| 38 | + if _plt is None: |
| 39 | + import matplotlib |
| 40 | + matplotlib.use("Agg") |
| 41 | + import matplotlib.pyplot as plt |
| 42 | + _plt = plt |
| 43 | + |
| 44 | + |
| 45 | +# --------------------------------------------------------------------------- |
| 46 | +# Grid reader |
| 47 | +# --------------------------------------------------------------------------- |
| 48 | + |
| 49 | +def read_model_grid(horizgridfile, Nx, Ny): |
| 50 | + """Read curvilinear grid coordinates from horizgridfile.bin.""" |
| 51 | + arr = np.fromfile(horizgridfile, dtype=">f8") |
| 52 | + fields = arr.reshape(16, Ny + 1, Nx + 1) |
| 53 | + return fields[0, :Ny, :Nx], fields[1, :Ny, :Nx] # xC, yC |
| 54 | + |
| 55 | + |
| 56 | +# --------------------------------------------------------------------------- |
| 57 | +# Binary diagnostics reader |
| 58 | +# --------------------------------------------------------------------------- |
| 59 | + |
| 60 | +def parse_diag_meta(meta_path): |
| 61 | + """Parse a MITgcm diagnostics .meta file.""" |
| 62 | + text = Path(meta_path).read_text() |
| 63 | + |
| 64 | + dim_match = re.search(r"dimList\s*=\s*\[\s*([\d\s,]+)\]", text) |
| 65 | + dims = [int(x) for x in dim_match.group(1).replace(",", " ").split()] |
| 66 | + ndims = len(dims) // 3 |
| 67 | + nx, ny = dims[0], dims[3] |
| 68 | + nz = dims[6] if ndims >= 3 else 1 |
| 69 | + |
| 70 | + prec_match = re.search(r"dataprec\s*=\s*\[\s*'(\w+)'\s*\]", text) |
| 71 | + dtype = np.dtype(">f8") if prec_match and "64" in prec_match.group(1) else np.dtype(">f4") |
| 72 | + |
| 73 | + nrec_match = re.search(r"nrecords\s*=\s*\[\s*(\d+)\s*\]", text) |
| 74 | + nrecords = int(nrec_match.group(1)) |
| 75 | + |
| 76 | + fld_match = re.search(r"fldList\s*=\s*\{([^}]+)\}", text) |
| 77 | + fields = re.findall(r"'(\w+)\s*'", fld_match.group(1)) if fld_match else [] |
| 78 | + |
| 79 | + return { |
| 80 | + "nx": nx, "ny": ny, "nz": nz, "ndims": ndims, |
| 81 | + "dtype": dtype, "nrecords": nrecords, "fields": fields, |
| 82 | + } |
| 83 | + |
| 84 | + |
| 85 | +def read_surface_field(data_path, meta, field_name): |
| 86 | + """Read the surface (k=0) slice of a named field from binary diagnostics.""" |
| 87 | + if field_name not in meta["fields"]: |
| 88 | + return None |
| 89 | + |
| 90 | + nx, ny, nz = meta["nx"], meta["ny"], meta["nz"] |
| 91 | + dtype = meta["dtype"] |
| 92 | + rec_size = nx * ny |
| 93 | + field_idx = meta["fields"].index(field_name) |
| 94 | + |
| 95 | + if meta["ndims"] == 3: |
| 96 | + # 3D file: each field has nz levels |
| 97 | + offset = field_idx * nz * rec_size |
| 98 | + else: |
| 99 | + # 2D file: each field is one record |
| 100 | + offset = field_idx * rec_size |
| 101 | + |
| 102 | + data = np.fromfile(data_path, dtype=dtype, count=rec_size, offset=offset * np.dtype(dtype).itemsize) |
| 103 | + return data.reshape(ny, nx) |
| 104 | + |
| 105 | + |
| 106 | +# --------------------------------------------------------------------------- |
| 107 | +# Run discovery |
| 108 | +# --------------------------------------------------------------------------- |
| 109 | + |
| 110 | +def discover_experiment_runs(simulation_dir, experiment): |
| 111 | + """Find run subdirectories under experiment/ that have diagnostic output.""" |
| 112 | + exp_dir = os.path.join(simulation_dir, experiment) |
| 113 | + if not os.path.isdir(exp_dir): |
| 114 | + return [] |
| 115 | + runs = [] |
| 116 | + for d in sorted(os.listdir(exp_dir)): |
| 117 | + full = os.path.join(exp_dir, d) |
| 118 | + if os.path.isdir(full): |
| 119 | + runs.append(d) |
| 120 | + return runs |
| 121 | + |
| 122 | + |
| 123 | +def find_diag_timesteps(run_dir, prefix="state3D"): |
| 124 | + """Find timestep suffixes for which both .data and .meta exist.""" |
| 125 | + timesteps = set() |
| 126 | + for f in Path(run_dir).glob(f"{prefix}.*.data"): |
| 127 | + m = re.search(rf"{prefix}\.(\d{{10}})\.data", f.name) |
| 128 | + if m: |
| 129 | + ts = m.group(1) |
| 130 | + meta = f.with_suffix(".meta") |
| 131 | + if meta.exists(): |
| 132 | + timesteps.add(ts) |
| 133 | + return sorted(timesteps) |
| 134 | + |
| 135 | + |
| 136 | +# --------------------------------------------------------------------------- |
| 137 | +# Plotting |
| 138 | +# --------------------------------------------------------------------------- |
| 139 | + |
| 140 | +FIELD_CONFIG = { |
| 141 | + "SST": {"label": "Sea Surface Temperature", "unit": "\u00b0C", |
| 142 | + "cmap": "RdYlBu_r", "vmin": 2, "vmax": 30}, |
| 143 | + "SSS": {"label": "Sea Surface Salinity", "unit": "PSU", |
| 144 | + "cmap": "viridis", "vmin": 33, "vmax": 37}, |
| 145 | + "SSH": {"label": "Sea Surface Height", "unit": "m", |
| 146 | + "cmap": "RdBu_r", "vmin": -1.5, "vmax": 1.5}, |
| 147 | + "KE": {"label": "Surface Kinetic Energy", "unit": "m\u00b2/s\u00b2", |
| 148 | + "cmap": "hot_r", "vmin": 0, "vmax": 0.5}, |
| 149 | +} |
| 150 | + |
| 151 | + |
| 152 | +def plot_field(field_data, field_name, xC, yC, title_extra, output_path): |
| 153 | + _import_mpl() |
| 154 | + cfg = FIELD_CONFIG[field_name] |
| 155 | + fig, ax = _plt.subplots(1, 1, figsize=(12, 6)) |
| 156 | + masked = np.ma.masked_where( |
| 157 | + (field_data == 0) | np.isnan(field_data) | (field_data <= -999), field_data |
| 158 | + ) |
| 159 | + im = ax.pcolormesh(xC, yC, masked, cmap=cfg["cmap"], |
| 160 | + vmin=cfg["vmin"], vmax=cfg["vmax"], shading="auto") |
| 161 | + cb = fig.colorbar(im, ax=ax, shrink=0.8, pad=0.02) |
| 162 | + cb.set_label(cfg["unit"], fontsize=10) |
| 163 | + ax.set_title(f'{cfg["label"]} \u2014 {title_extra}', fontsize=13) |
| 164 | + ax.set_xlabel("Longitude") |
| 165 | + ax.set_ylabel("Latitude") |
| 166 | + ax.set_aspect("equal") |
| 167 | + fig.tight_layout() |
| 168 | + fig.savefig(output_path, dpi=120, bbox_inches="tight") |
| 169 | + _plt.close(fig) |
| 170 | + |
| 171 | + |
| 172 | +# --------------------------------------------------------------------------- |
| 173 | +# Processing |
| 174 | +# --------------------------------------------------------------------------- |
| 175 | + |
| 176 | +def process_run(run_dir, run_name, xC, yC, plotted, deltaT, start_date, nx, ny): |
| 177 | + """Process one run directory. Returns number of new plots created.""" |
| 178 | + t0 = datetime.strptime(start_date, "%Y-%m-%d") |
| 179 | + plots_dir = os.path.join(run_dir, "plots") |
| 180 | + new_count = 0 |
| 181 | + |
| 182 | + # Find available timesteps |
| 183 | + ts_3d = find_diag_timesteps(run_dir, "state3D") |
| 184 | + ts_2d = find_diag_timesteps(run_dir, "state2D") |
| 185 | + |
| 186 | + for ts in ts_3d: |
| 187 | + if ts in plotted: |
| 188 | + continue |
| 189 | + |
| 190 | + iter_num = int(ts) |
| 191 | + model_date = (t0 + timedelta(seconds=iter_num * deltaT)).strftime("%Y-%m-%d") |
| 192 | + title = f"{run_name} \u2014 {model_date}" |
| 193 | + |
| 194 | + data_path = os.path.join(run_dir, f"state3D.{ts}.data") |
| 195 | + meta_path = os.path.join(run_dir, f"state3D.{ts}.meta") |
| 196 | + meta = parse_diag_meta(meta_path) |
| 197 | + |
| 198 | + os.makedirs(plots_dir, exist_ok=True) |
| 199 | + |
| 200 | + try: |
| 201 | + # SST |
| 202 | + sst = read_surface_field(data_path, meta, "THETA") |
| 203 | + if sst is not None: |
| 204 | + out = os.path.join(plots_dir, f"SST_{ts}.png") |
| 205 | + if not os.path.exists(out): |
| 206 | + plot_field(sst, "SST", xC, yC, title, out) |
| 207 | + new_count += 1 |
| 208 | + |
| 209 | + # SSS |
| 210 | + sss = read_surface_field(data_path, meta, "SALT") |
| 211 | + if sss is not None: |
| 212 | + out = os.path.join(plots_dir, f"SSS_{ts}.png") |
| 213 | + if not os.path.exists(out): |
| 214 | + plot_field(sss, "SSS", xC, yC, title, out) |
| 215 | + new_count += 1 |
| 216 | + |
| 217 | + # KE |
| 218 | + u = read_surface_field(data_path, meta, "UVEL") |
| 219 | + v = read_surface_field(data_path, meta, "VVEL") |
| 220 | + if u is not None and v is not None: |
| 221 | + ke = 0.5 * (u ** 2 + v ** 2) |
| 222 | + out = os.path.join(plots_dir, f"KE_{ts}.png") |
| 223 | + if not os.path.exists(out): |
| 224 | + plot_field(ke, "KE", xC, yC, title, out) |
| 225 | + new_count += 1 |
| 226 | + |
| 227 | + except Exception as e: |
| 228 | + print(f" [{run_name}] Error processing state3D {ts}: {e}") |
| 229 | + continue |
| 230 | + |
| 231 | + # SSH from state2D |
| 232 | + if ts in ts_2d: |
| 233 | + try: |
| 234 | + data2d_path = os.path.join(run_dir, f"state2D.{ts}.data") |
| 235 | + meta2d = parse_diag_meta(os.path.join(run_dir, f"state2D.{ts}.meta")) |
| 236 | + etan = read_surface_field(data2d_path, meta2d, "ETAN") |
| 237 | + if etan is not None: |
| 238 | + out = os.path.join(plots_dir, f"SSH_{ts}.png") |
| 239 | + if not os.path.exists(out): |
| 240 | + plot_field(etan, "SSH", xC, yC, title, out) |
| 241 | + new_count += 1 |
| 242 | + except Exception as e: |
| 243 | + print(f" [{run_name}] Error processing state2D {ts}: {e}") |
| 244 | + |
| 245 | + plotted.add(ts) |
| 246 | + |
| 247 | + return new_count |
| 248 | + |
| 249 | + |
| 250 | +# --------------------------------------------------------------------------- |
| 251 | +# Main |
| 252 | +# --------------------------------------------------------------------------- |
| 253 | + |
| 254 | +def main(): |
| 255 | + parser = argparse.ArgumentParser( |
| 256 | + description="Sidecar plotter for MITgcm binary diagnostics" |
| 257 | + ) |
| 258 | + parser.add_argument("simulation_dir", |
| 259 | + help="Path to simulation directory (e.g. simulations/glorysv12-curvilinear)") |
| 260 | + parser.add_argument("experiment", |
| 261 | + help="Experiment subdirectory (e.g. repeat-year-50)") |
| 262 | + parser.add_argument("--run", default=None, |
| 263 | + help="Watch a single run (e.g. 001) instead of all runs") |
| 264 | + parser.add_argument("--poll", type=int, default=120, |
| 265 | + help="Seconds between polls (0 = single pass, no loop)") |
| 266 | + parser.add_argument("--start-date", default="2002-07-01") |
| 267 | + parser.add_argument("--dt", type=float, default=360.0, help="Model timestep in seconds") |
| 268 | + args = parser.parse_args() |
| 269 | + |
| 270 | + Nx, Ny = 768, 424 |
| 271 | + simulation_dir = os.path.abspath(args.simulation_dir) |
| 272 | + exp_dir = os.path.join(simulation_dir, args.experiment) |
| 273 | + |
| 274 | + horizgridfile = os.path.join(simulation_dir, "input", "horizgridfile.bin") |
| 275 | + xC, yC = read_model_grid(horizgridfile, Nx, Ny) |
| 276 | + print(f"Grid loaded: {Ny}x{Nx}, lon [{xC.min():.1f},{xC.max():.1f}], lat [{yC.min():.1f},{yC.max():.1f}]") |
| 277 | + print(f"Watching: {exp_dir}") |
| 278 | + if args.run: |
| 279 | + print(f"Single run: {args.run}") |
| 280 | + |
| 281 | + plotted_cache = {} # run_name → set of plotted timesteps |
| 282 | + |
| 283 | + while True: |
| 284 | + if args.run: |
| 285 | + runs = [args.run] |
| 286 | + else: |
| 287 | + runs = discover_experiment_runs(simulation_dir, args.experiment) |
| 288 | + |
| 289 | + for run_name in runs: |
| 290 | + run_dir = os.path.join(exp_dir, run_name) |
| 291 | + if not os.path.isdir(run_dir): |
| 292 | + continue |
| 293 | + |
| 294 | + if run_name not in plotted_cache: |
| 295 | + plotted_cache[run_name] = set() |
| 296 | + |
| 297 | + n = process_run(run_dir, f"{args.experiment}/{run_name}", |
| 298 | + xC, yC, plotted_cache[run_name], |
| 299 | + args.dt, args.start_date, Nx, Ny) |
| 300 | + if n > 0: |
| 301 | + print(f"[{args.experiment}/{run_name}] Plotted {n} new images") |
| 302 | + |
| 303 | + if args.poll <= 0: |
| 304 | + break |
| 305 | + time.sleep(args.poll) |
| 306 | + |
| 307 | + |
| 308 | +if __name__ == "__main__": |
| 309 | + main() |
0 commit comments