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,
    yac_weight_file_on_existing
)
 
import uxarray as ux
import netCDF4
import numpy as np
from argparse import ArgumentParser
import json
from pathlib import Path
import logging
import re
 
 
def load_grid(filename, gridname=None):
    match filename:
        case f if m := re.match(r"^healpix(\d{1,2})$", f):
            return ux.Grid.from_healpix(zoom=int(m[1]))
        case f if m := re.match(r"^reg((?P<lon_from>-?\d+\.?\d*):(?P<lon_to>-?\d+\.?\d*):)?(?P<lon_n>\d+),((?P<lat_from>-?\d+\.?\d*):(?P<lat_to>-?\d+\.?\d*):)?(?P<lat_n>\d+)$", f):
            return ux.Grid.from_structured(lon=np.linspace(start=float(m["lon_from"] or 0),
                                                           stop=float(m["lon_to"] or 360),
                                                           num=int(m["lon_n"])),
                                           lat=np.linspace(start=float(m["lat_from"] or -90),
                                                           stop=float(m["lat_to"] or 90),
                                                           num=int(m["lat_n"])),
                                           )
        case _:
            if gridname is not None:
                return ux.open_multigrid(filename, gridnames=[gridname])[gridname]
            return ux.open_grid(filename)
 
 
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)
 
 
def parse_interpolation(s):
    try:
        return json.loads(s)
    except Exception as e:
        logging.debug(f"Cannot parse interpolation as json: {e}")
        return s, {}
 
 
yac_location.__str__ = lambda loc: loc.name
 
 
parser = ArgumentParser(prog="yac_generate_weights.py", description="Generate YAC weights files")
parser.add_argument("src_grid_path", type=str,
    help=(
        "Path to the source grid file, or a shorthand: "
        "'healpix<N>' for a HEALPix grid at zoom level N (e.g. 'healpix5'), "
        "or 'reg[lon_from:lon_to:]<lon_n>,[lat_from:lat_to:]<lat_n>' for a "
        "regular structured grid (e.g. 'reg360,180' or 'reg-180:180:360,-90:90:181')."
    )
)
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=None)
 
parser.add_argument("tgt_grid_path", type=str,
    help=(
        "Path to the target grid file, or a shorthand: "
        "'healpix<N>' for a HEALPix grid at zoom level N (e.g. 'healpix5'), "
        "or 'reg[lon_from:lon_to:]<lon_n>,[lat_from:lat_to:]<lat_n>' for a "
        "regular structured grid (e.g. 'reg360,180' or 'reg-180:180:360,-90:90:181')."
    )
)
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", "--tgt_grid_name", type=str, default=None)
 
parser.add_argument(
    "interpolation_stack",
    type=parse_interpolation,
    nargs="+",
    help="format: method_name[,JSON mapping]",
)
parser.add_argument("--output-file", type=Path, default="weights.nc")
parser.add_argument("--output-grids-file",
                    type=Path,
                    default="grids.nc",
                    help="Write the source and target grids to this file for debugging purposes.")
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 = load_grid(args.src_grid_path, args.src_grid_name)
src_grid = BasicGrid.from_uxgrid(args.src_grid_name or "source_grid", 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 = load_grid(args.tgt_grid_path, args.tgt_grid_name)
tgt_grid = BasicGrid.from_uxgrid(args.tgt_grid_name or "target_grid", 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 netCDF4.Dataset(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(f"Writing grid file {args.output_grids_file}")
if args.output_grids_file.exists():
    logging.info(f"Overwriting existing file")
    args.output_grids_file.unlink()
src_grid.to_file(str(args.output_grids_file))
tgt_grid.to_file(str(args.output_grids_file))
 
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),
    src_grid_name=args.src_grid_name or "source_grid",
    tgt_grid_name=args.tgt_grid_name or "target_grid",
    on_existing=yac_weight_file_on_existing.YAC_WEIGHT_FILE_OVERWRITE,
)