yac_plot_grids.py

Go to the documentation of this file.

#!/usr/bin/env python3
#
# Copyright (c) 2026 The YAC Authors
#
# SPDX-License-Identifier: BSD-3-Clause
#
# /// script
# requires-python = ">=3.9"
# dependencies = ["numpy", "matplotlib", "netCDF4", "cartopy"]
# ///
 
"""Plot YAC debug grid files on an interactive map."""
 
import logging
import argparse
from itertools import cycle
 
import numpy as np
import matplotlib.pyplot as plt
import netCDF4
 
from _utils.plotting import (
    setup_figure,
    add_label_in_extent,
    plot_grid_edges_in_ax,
    load_grid,
    add_keybinding_legend,
)
 
# Cycle through these colours when more than one grid is plotted.
_GRID_COLORS = ["green", "blue", "red", "orange", "purple", "brown", "pink", "gray"]
 
 
def find_grids_in_file(filename):
    ds = netCDF4.Dataset(filename, mode="r")
    yac_grid_names = [d[3:] for d in ds.dimensions
                       if d.startswith("nv_")]
    if len(yac_grid_names) == 0:
        logging.warning(f"No yac grids found in {filename!r}")
    return yac_grid_names
 
if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Plot YAC debug grid files on an interactive map"
    )
    parser.add_argument(
        "-v",
        "--verbose",
        action="count",
        default=0,
        help="Increase verbosity (-v for INFO, -vv for DEBUG)",
    )
    parser.add_argument(
        "--center",
        "-c",
        type=float,
        nargs=2,
        metavar=("LON", "LAT"),
        default=(0, 0),
        help="Center of the map projection (longitude, latitude in degrees)",
    )
    parser.add_argument(
        "--radius",
        "-r",
        type=float,
        default=2000,
        help="Radius in km of the region to display (default: 2000). Set to 0 to auto-scale the view to fit all plotted grid data.",
    )
    parser.add_argument(
        "--coast-res",
        type=str,
        default="50m",
        nargs="?",
        choices=("10m", "50m", "110m"),
        help=(
            "Coastline resolution (default: 50m). "
            "Use --coast-res without value to disable."
        ),
    )
    parser.add_argument(
        "--projection",
        type=str,
        default="orthographic",
        choices=("orthographic", "stereographic", "platecarree"),
        help="Map projection type (default: orthographic)",
    )
    parser.add_argument(
        "gridfile",
        type=str,
        help="Path to YAC debug grid file (NetCDF format)",
    )
    parser.add_argument(
        "--only",
        type=str,
        nargs="*",
        default=None,
        help="Only plot grids with these names",
    )
    parser.add_argument(
        "--corner",
        action="store_true",
        help="Load indices on corners",
    )
    parser.add_argument(
        "--ranks",
        type=int,
        nargs="*",
        default=None,
        help="Plot only the part of the grid that was stored on the given MPI ranks")
 
    args = parser.parse_args()
 
    log_level = [logging.WARNING, logging.INFO, logging.DEBUG][min(args.verbose, 2)]
    logging.basicConfig(format="%(levelname)s: %(message)s")
    logging.getLogger("_utils").setLevel(level=log_level)
    logger = logging.getLogger("plotting")
    logger.setLevel(level=log_level)
 
 
    logger.info("Starting YAC grid visualisation")
    logger.debug(f"CLI arguments: {args}")
 
    logger.info("Setting up figure and map projection")
    fig, ax = setup_figure(
        projection=args.projection,
        center=args.center,
        focus_points=None,
        radius=args.radius,
        coast_res=args.coast_res,
    )
 
    legend = []
    for gridname, color, key in zip(args.only or find_grids_in_file(args.gridfile),
                                    _GRID_COLORS, "ABCDEFGHIJKLMNO"):
        logger.info(f"Plotting {gridname}")
        (clo, cla), (lo, la), ids, mask = load_grid([args.gridfile], gridname, "CORNER" if args.corner else "CELL", args.ranks)
        plot_grid_edges_in_ax(ax, clo[mask], cla[mask], plot_kwargs={"color": color})
        plot_grid_edges_in_ax(ax, clo[~mask], cla[~mask], plot_kwargs={"color": color, "alpha": 0.3, "linestyle": "dashed"})
 
        if args.corner:
            ids, select = np.unique(ids.ravel(), return_index=True)
            lo = lo.ravel()[select]
            la = la.ravel()[select]
 
        add_label_in_extent(
            ax,
            lo,
            la,
            ids,
            {"color": "black", "visible": False},
            key,
        )
        legend.append((key, f"toggle {gridname} labels", color))
 
    add_keybinding_legend(ax, legend)
    if args.radius==0:
        ax.autoscale(True)
 
    logger.info("Opening interactive plot window")
    plt.show()