yac_plot_weights.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 interpolation weights from YAC weight files on a map."""
 
import logging
import argparse
import pathlib
 
import numpy as np
import matplotlib.pyplot as plt
import netCDF4
 
from _utils.plotting import (
    setup_figure,
    plot_grid_edges_in_ax,
    plot_weights,
    plot_fixed_values,
    add_label_in_extent,
    load_grid,
    add_keybinding_legend,
)
 
 
def load_weight_file(weightfile):
    """Load and validate YAC weight file.
 
    Args:
        weightfile: Path to weight file
 
    Returns:
        Tuple of (src_grid_name, tgt_grid_name, src_loc, tgt_loc,
                  src_address, tgt_address, weights, fixed_values)
        where fixed_values is a list of (value, dst_address_array) tuples,
        one entry per distinct fixed value (empty list when none are present).
    """
    logging.info(f"Loading weight file: {weightfile}")
    with netCDF4.Dataset(weightfile, mode="r") as wf:
        assert(wf.version == "yac weight file 1.0")
        num_weights = wf.dimensions["num_wgts"].size
        logging.debug(f"Weight file has {num_weights} weight(s)")
        if num_weights != 1:
            raise ValueError(
                f"Only files with num_weights==1 are supported, got {num_weights}"
            )
 
        src_grid_name = wf.src_grid_name
        tgt_grid_name = wf.dst_grid_name
        logging.info(f"Source grid: {src_grid_name}")
        logging.info(f"Target grid: {tgt_grid_name}")
 
        src_loc = netCDF4.chartostring(wf["src_locations"][0, :])[()].strip()
        logging.info(f"Source location: {src_loc}")
 
        tgt_loc = netCDF4.chartostring(wf["dst_location"][:])[()].strip()
        logging.info(f"Target location: {tgt_loc}")
 
        src_address = np.asarray(wf["src_address"])
        tgt_address = np.asarray(wf["dst_address"])
 
        weights = np.asarray(wf["remap_matrix"])
 
        logging.info(f"Total weight links: {len(src_address)}")
        logging.debug(f"Weight range: [{weights.min():.6f}, {weights.max():.6f}]")
 
        fixed_values = []
        if "fixed_values" in wf.variables:
            fv = np.asarray(wf["fixed_values"])
            n_per_fv = np.asarray(wf["num_dst_per_fixed_value"])
            dst_fixed = np.asarray(wf["dst_address_fixed"])
            offset = 0
            for value, count in zip(fv, n_per_fv):
                fixed_values.append((float(value), dst_fixed[offset : offset + count]))
                offset += count
            total = sum(c for _, c in [(v, len(a)) for v, a in fixed_values])
            logging.info(
                f"Fixed values: {len(fixed_values)} distinct, "
                f"{total} target points total"
            )
 
    return (
        src_grid_name,
        tgt_grid_name,
        src_loc,
        tgt_loc,
        src_address,
        tgt_address,
        weights,
        fixed_values,
    )
 
 
def filter_weights_by_index(
    src_idx,
    tgt_idx,
    src_address,
    tgt_address,
    weights,
    src_ids,
    tgt_ids,
    src_points,
    tgt_points,
):
    """Filter interpolation weights to focus on a specific source or target index.
 
    Args:
        src_idx: Global ID of the source cell/vertex to filter by (or None). This is
            the value shown as a label when pressing 's', i.e. the .gid/.vid value
            from the grid file.
        tgt_idx: Global ID of the target cell/vertex to filter by (or None). See src_idx.
        src_address: Array of source addresses from the weight file. Each value is
            global_id + 1 (1-indexed global IDs, not array positions).
        tgt_address: Array of target addresses from the weight file. Each value is
            global_id + 1 (1-indexed global IDs, not array positions).
        weights: Weight values
        src_ids: Global IDs of source grid points (.gid or .vid from grid file)
        tgt_ids: Global IDs of target grid points (.gid or .vid from grid file)
        src_points: Source coordinates, shape (2, n)
        tgt_points: Target coordinates, shape (2, n)
 
    Returns:
        Tuple of (filtered_src_address, filtered_tgt_address, filtered_weights, focus_points)
    """
    if src_idx is not None:
        logging.info(f"Filtering weights for source global ID {src_idx}")
        # Filter weight links where the source global ID matches src_idx.
        # Addresses are global_id + 1.
        weight_mask = src_address == src_idx + 1
        fixed_ids = src_ids
        fixed_global_id = src_idx
        linked_address = tgt_address[weight_mask]
        linked_ids = tgt_ids
        fixed_points = src_points
        linked_points = tgt_points
        filter_type = "source"
    else:
        logging.info(f"Filtering weights for target global ID {tgt_idx}")
        weight_mask = tgt_address == tgt_idx + 1
        fixed_ids = tgt_ids
        fixed_global_id = tgt_idx
        linked_address = src_address[weight_mask]
        linked_ids = src_ids
        fixed_points = tgt_points
        linked_points = src_points
        filter_type = "target"
 
    # linked_address values are global_id + 1; search for global_id in linked_ids
    # to find the array position of each linked point.
    sorter = np.argsort(linked_ids)
    linked_idx = sorter[np.searchsorted(linked_ids, linked_address - 1, sorter=sorter)]
 
    # Find array position of the fixed point by its global ID.
    fixed_point_idx = np.where(fixed_ids == fixed_global_id)[0]
 
    if len(fixed_point_idx) == 0:
        raise ValueError(f"Global ID {fixed_global_id} not found in {filter_type} grid")
 
    # Combine for focus points
    focus_points = np.hstack(
        [fixed_points[..., fixed_point_idx], linked_points[..., linked_idx]]
    )
 
    # Apply mask to weight data
    filtered_src_address = src_address[weight_mask]
    filtered_tgt_address = tgt_address[weight_mask]
    filtered_weights = weights[weight_mask]
 
    num_filtered = len(filtered_weights)
    logging.info(
        f"Found {num_filtered} weight links for {filter_type} global ID {fixed_global_id}"
    )
    if num_filtered > 0:
        logging.debug(
            f"Filtered weight range: [{filtered_weights.min():.6f}, {filtered_weights.max():.6f}]"
        )
        logging.debug(f"Weight sum: {filtered_weights.sum():.6f}")
 
    return filtered_src_address, filtered_tgt_address, filtered_weights, focus_points
 
 
def plot_grid(ax, corner_coords, coords, ids, core_mask, color, label_key):
    """Plot a grid on the given axis with specified coordinates, IDs, and styling.
 
    Core cells are plotted with full opacity; halo cells (core_mask==False)
    are plotted with a dimmed, dashed style to distinguish them visually.
 
    Args:
        ax: Matplotlib axis to plot on
        corner_coords: Corner coordinates of the grid
        coords: Coordinates of grid points
        ids: IDs of grid points
        core_mask: Boolean array, True for core (owned) cells, False for halo
        color: Color for core cells
        label_key: Key for labeling
    """
    logging.debug(
        f"Plotting grid with {np.sum(core_mask)} core + "
        f"{np.sum(~core_mask)} halo cells in {color}"
    )
    plot_grid_edges_in_ax(
        ax,
        corner_coords[0, core_mask],
        corner_coords[1, core_mask],
        {"zorder": 1, "color": color},
    )
    if np.any(~core_mask):
        plot_grid_edges_in_ax(
            ax,
            corner_coords[0, ~core_mask],
            corner_coords[1, ~core_mask],
            {"zorder": 1, "color": color, "alpha": 0.3, "linestyle": "dashed"},
        )
 
    add_label_in_extent(
        ax,
        coords[0, ...],
        coords[1, ...],
        ids,
        {"color": color, "visible": False},
        label_key,
    )
 
 
if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Plot YAC interpolation weights on a 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,
        help="Center of the map projection (longitude, latitude in degrees)",
        default=None,
        metavar=("LON", "LAT"),
    )
    parser.add_argument(
        "--radius",
        "-r",
        type=float,
        help="Radius in km of the region around the center to display (default: 2000)",
        default=2000,
    )
    parser.add_argument(
        "--coast-res",
        type=str,
        default="50m",
        nargs="?",
        choices=("10m", "50m", "110m"),
        help="Resolution of coastlines (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)",
    )
 
    idx_group = parser.add_mutually_exclusive_group()
    idx_group.add_argument(
        "--source-index",
        type=int,
        default=None,
        help="Focus on a source cell/vertex. (global index).",
    )
    idx_group.add_argument(
        "--target-index",
        type=int,
        default=None,
        help="Focus on a target cell/vertex. (global index)",
    )
 
    parser.add_argument(
        "weightfile",
        type=pathlib.Path,
        help="Path to YAC weight file (NetCDF format)",
    )
    parser.add_argument(
        "gridfile",
        type=pathlib.Path,
        help="Path to a grid file (NetCDF format). The file is searched for the grid name specified in the weight file.",
    )
    parser.add_argument(
        "gridfile2",
        type=pathlib.Path,
        nargs="?",
        help="Path to a second grid file to search for grid names. Use this when the source and target grids are stored in separate files.",
    )
 
    args = parser.parse_args()
 
    log_level = [logging.WARNING, logging.INFO, logging.DEBUG][min(args.verbose, 2)]
    logging.basicConfig(level=log_level, format="%(levelname)s: %(message)s")
 
    logging.info("Starting YAC weight visualization")
    logging.debug(f"Weight file: {args.weightfile}")
    logging.debug(f"Grid files: {args.gridfile}, {args.gridfile2}")
    logging.debug(f"Projection: {args.projection}")
    if args.center:
        logging.debug(f"Center: {args.center[0]:.2f}°E, {args.center[1]:.2f}°N")
    logging.debug(f"Radius: {args.radius} km")
 
    (
        src_grid_name,
        tgt_grid_name,
        src_loc,
        tgt_loc,
        src_address,
        tgt_address,
        weights,
        fixed_values,
    ) = load_weight_file(args.weightfile)
 
    logging.info("Loading source grid")
    src_corners, src_points, src_ids, src_core_mask = load_grid(
        [args.gridfile, args.gridfile2],
        src_grid_name,
        src_loc,
    )
 
    logging.info("Loading target grid")
    tgt_corners, tgt_points, tgt_ids, tgt_core_mask = load_grid(
        [args.gridfile, args.gridfile2],
        tgt_grid_name,
        tgt_loc,
    )
 
    focus_points = None
 
    if args.source_index is not None or args.target_index is not None:
        src_address, tgt_address, weights, focus_points = filter_weights_by_index(
            args.source_index,
            args.target_index,
            src_address,
            tgt_address,
            weights,
            src_ids,
            tgt_ids,
            src_points,
            tgt_points,
        )
 
    logging.info("Setting up figure and map projection")
    fig, ax = setup_figure(
        projection=args.projection,
        center=args.center,
        focus_points=focus_points,
        radius=args.radius,
        coast_res=args.coast_res,
    )
 
    logging.info("Plotting source grid (green)")
    plot_grid(ax, src_corners, src_points, src_ids, src_core_mask, "green", "S")
 
    logging.info("Plotting target grid (blue)")
    plot_grid(ax, tgt_corners, tgt_points, tgt_ids, tgt_core_mask, "blue", "T")
 
    logging.info(f"Plotting {len(weights)} weight connections")
    plot_weights(
        ax,
        src_points,
        src_ids,
        tgt_points,
        tgt_ids,
        src_address,
        tgt_address,
        weights,
    )
 
    if fixed_values:
        logging.info(f"Plotting {len(fixed_values)} fixed-value group(s)")
        plot_fixed_values(ax, tgt_points, tgt_ids, fixed_values)
 
    add_keybinding_legend(
        ax,
        [
            ("S", "toggle source labels", "green"),
            ("T", "toggle target labels", "blue"),
            ("W", "toggle weight labels"),
        ],
    )
 
    logging.info("Opening interactive plot window")
    logging.info("Press 's' to toggle source labels, 't' to toggle target labels")
    plt.show()