yac_generate_weights.py

Go to the documentation of this file.

#!/usr/bin/env python3
 
# Copyright (c) 2025 The YAC Authors
#
# SPDX-License-Identifier: BSD-3-Clause
 
 
 
try:
    from mpi4py import MPI  # calls MPI_Init with some special treatments
except:
    pass
 
from yac.core import (BasicGrid,
                      yac_location,
                      InterpField,
                      InterpolationStack,
                      compute_weights,
                      )
 
import uxarray as ux
import numpy as np
from argparse import ArgumentParser
import json
from pathlib import Path
import logging
 
 
def parse_interp_method(s):
    if "," in s:
        s, args = s.split(",")
        args = json.loads(args)
    else:
        args = {}
    assert type(args) is dict, "Invalid argument spec for interpolation"
    return s, args
 
 
def get_coords(loc, uxgrid, yac_grid):
    if loc == yac_location.YAC_LOC_CELL:
        lon, lat = np.deg2rad([uxgrid.face_lon, uxgrid.face_lat])
    elif loc == yac_location.YAC_LOC_CORNER:
        lon = lat = None
    elif loc == yac_location.YAC_LOC_EDGE:
        lon, lat = np.deg2rad([uxgrid.edge_lon, uxgrid.edge_lat])
    return yac_grid.add_coordinates(loc, lon, lat)
 
 
yac_location.__str__ = lambda loc: loc.name
 
 
parser = ArgumentParser(prog="gen_weights.py",
                        description="Generate YAC weights files")
parser.add_argument("src_grid_path", type=str)
parser.add_argument("--src-location",
                    type=lambda loc: yac_location[loc],
                    choices=[yac_location.YAC_LOC_CELL, yac_location.YAC_LOC_CORNER, yac_location.YAC_LOC_EDGE],
                    default=yac_location.YAC_LOC_CELL
                    )
parser.add_argument("--src-mask-var", type=str,
                    help="Variable name of the mask")
parser.add_argument("--src-mask-filename", type=str,
                    help="Filename to read the mask from.")
parser.add_argument("--src-mask-valid-values", type=int, nargs="*", default=[1],
                    help="Valid values for the mask")
parser.add_argument("--src-grid-name", type=str, default="src_grid")
 
parser.add_argument("tgt_grid_path", type=str)
parser.add_argument("--tgt-location",
                    type=lambda loc: yac_location[loc],
                    choices=[yac_location.YAC_LOC_CELL, yac_location.YAC_LOC_CORNER, yac_location.YAC_LOC_EDGE],
                    default=yac_location.YAC_LOC_CELL
                    )
parser.add_argument("--tgt-mask-var", type=str,
                    help="Variable name of the mask")
parser.add_argument("--tgt-mask-filename", type=str,
                    help="Filename to read the mask from.")
parser.add_argument("--tgt-mask-valid-values", type=int, nargs="*", default=[1],
                    help="Valid values for the mask")
parser.add_argument("--tgt-grid-name", type=str, default="tgt_grid")
 
parser.add_argument("interpolation_stack",
                    type=parse_interp_method,
                    nargs="+",
                    help="format: method_name[,JSON mapping]")
parser.add_argument("--output-file", type=Path,
                    default="weights.nc")
parser.add_argument("--log-level",
                    default=logging.INFO,
                    type=lambda x: getattr(logging, x),
                    help="Configure the logging level.",
                    )
 
args = parser.parse_args()
 
logging.basicConfig(level=args.log_level)
 
logging.info("Reading src grid")
src_uxgrid = ux.open_grid(args.src_grid_path)
src_grid = BasicGrid.from_uxgrid(args.src_grid_name, src_uxgrid)
src_coords = get_coords(args.src_location, src_uxgrid, src_grid)
if args.src_mask_var is not None:
    src_mask_file = args.src_mask_filename or args.src_grid_path
    with ux.open_dataset(args.src_grid_path, src_mask_file) as ds:
        m = np.isin(ds[args.src_mask_var], args.src_mask_valid_values)
        src_mask = src_grid.add_mask(args.src_location, m)
else:
    src_mask = None
src_field = InterpField(src_coords, src_mask)
 
logging.info("Reading tgt grid")
tgt_uxgrid = ux.open_grid(args.tgt_grid_path)
tgt_grid = BasicGrid.from_uxgrid(args.tgt_grid_name, tgt_uxgrid)
tgt_coords = get_coords(args.tgt_location, tgt_uxgrid, tgt_grid)
if args.tgt_mask_var is not None:
    tgt_mask_file = args.tgt_mask_filename or args.tgt_grid_path
    with ux.open_dataset(args.tgt_grid_path, tgt_mask_file) as ds:
        m = np.isin(ds[args.tgt_mask_var], args.tgt_mask_valid_values)
        tgt_mask = tgt_grid.add_mask(args.tgt_location, m)
else:
    tgt_mask = None
tgt_field = InterpField(tgt_coords, tgt_mask)
 
interp_stack = InterpolationStack.from_list(args.interpolation_stack)
 
logging.info("Computing weights")
weights = compute_weights(interp_stack,
                          src_field,
                          tgt_field)
 
logging.info("Write weights to file")
weights.write_to_file(str(args.output_file))