core.pyx

Go to the documentation of this file.

# Copyright (c) 2024 The YAC Authors
#
# SPDX-License-Identifier: BSD-3-Clause
 
import cython
import numpy as _np
from libc.stdint cimport SIZE_MAX
from cython cimport view
from libc.stdlib cimport malloc, free
from warnings import warn
from collections import namedtuple
from types import coroutine
import logging
 
logger = logging.getLogger("yac.core")
logger.addHandler(logging.NullHandler())
 
cdef extern from "Python.h":
    int Py_AtExit(void (*)())
 
cdef extern from "<mpi.h>" nogil:
    ctypedef void* MPI_Comm
    MPI_Comm MPI_Comm_f2c(int)
    MPI_Comm MPI_COMM_WORLD
    MPI_Comm MPI_COMM_NULL
 
cdef extern from *:
    """
#ifdef HAVE_CONFIG_H
#include "config.h"
#endif"""
 
cdef extern from "yac_core.h" nogil:
    char * yac_get_core_version()
 
    ctypedef int yac_int;  
    ctypedef double(*yac_coordinate_pointer)[3];
 
    cpdef enum yac_location:
        YAC_LOC_CELL
        YAC_LOC_CORNER
        YAC_LOC_EDGE
        YAC_LOC_UNDEFINED
 
    cpdef enum yac_interp_avg_weight_type:
        YAC_INTERP_AVG_ARITHMETIC
        YAC_INTERP_AVG_DIST
        YAC_INTERP_AVG_BARY
    cdef yac_interp_avg_weight_type YAC_INTERP_AVG_WEIGHT_TYPE_DEFAULT
    cdef int YAC_INTERP_AVG_PARTIAL_COVERAGE_DEFAULT
 
    cpdef enum yac_interp_ncc_weight_type:
        YAC_INTERP_NCC_AVG
        YAC_INTERP_NCC_DIST
    cdef yac_interp_ncc_weight_type YAC_INTERP_NCC_WEIGHT_TYPE_DEFAULT
    cdef int YAC_INTERP_NCC_PARTIAL_COVERAGE_DEFAULT
 
    cpdef enum yac_interp_dnn_weight_type:
        YAC_INTERP_DNN_WEIGHT_AVG
        YAC_INTERP_DNN_WEIGHT_DIST
        YAC_INTERP_DNN_WEIGHT_GAUSS
        YAC_INTERP_DNN_WEIGHT_RBF
    cpdef enum yac_interp_dnn_search_distance_type:
        YAC_INTERP_DNN_SEARCH_DISTANCE_FIXED
        YAC_INTERP_DNN_SEARCH_DISTANCE_CELL_AREA
    cdef yac_interp_dnn_weight_type YAC_INTERP_DNN_WEIGHT_DEFAULT
    cdef yac_interp_dnn_search_distance_type YAC_INTERP_DNN_SEARCH_DISTANCE_TYPE_DEFAULT
    cdef double YAC_INTERP_DNN_SEARCH_DISTANCE_DEFAULT
    cdef double YAC_INTERP_DNN_SCALE_DEFAULT
 
    cpdef enum yac_interp_method_conserv_normalisation:
        YAC_INTERP_CONSERV_DESTAREA
        YAC_INTERP_CONSERV_FRACAREA
    cdef int YAC_INTERP_CONSERV_ORDER_DEFAULT
    cdef int YAC_INTERP_CONSERV_ENFORCED_CONSERV_DEFAULT
    cdef int YAC_INTERP_CONSERV_PARTIAL_COVERAGE_DEFAULT
    cdef yac_interp_method_conserv_normalisation YAC_INTERP_CONSERV_NORMALISATION_DEFAULT
 
    cpdef enum yac_interp_spmap_weight_type:
        YAC_INTERP_SPMAP_AVG
        YAC_INTERP_SPMAP_DIST
    cpdef enum yac_interp_spmap_scale_type:
        YAC_INTERP_SPMAP_NONE
        YAC_INTERP_SPMAP_SRCAREA
        YAC_INTERP_SPMAP_INVTGTAREA
        YAC_INTERP_SPMAP_FRACAREA
    cpdef enum yac_interp_spmap_cell_area_provider:
        YAC_INTERP_SPMAP_CELL_AREA_FILE
        YAC_INTERP_SPMAP_CELL_AREA_YAC
    cdef double YAC_INTERP_SPMAP_SPREAD_DISTANCE_DEFAULT
    cdef double YAC_INTERP_SPMAP_MAX_SEARCH_DISTANCE_DEFAULT
    cdef yac_interp_spmap_weight_type YAC_INTERP_SPMAP_WEIGHTED_DEFAULT
    cdef yac_interp_spmap_scale_type YAC_INTERP_SPMAP_SCALE_TYPE_DEFAULT
    cdef double YAC_INTERP_SPMAP_SPHERE_RADIUS_DEFAULT
    cdef const char* YAC_INTERP_SPMAP_FILENAME_DEFAULT
    cdef const char* YAC_INTERP_SPMAP_VARNAME_DEFAULT
    cdef int YAC_INTERP_SPMAP_MIN_GLOBAL_ID_DEFAULT
 
    cpdef enum yac_interp_nnn_weight_type:
        YAC_INTERP_NNN_AVG #< average of n source points
        YAC_INTERP_NNN_DIST #< distance weighted average of n source points
        YAC_INTERP_NNN_GAUSS #< distance with Gauss weights of n source points
        YAC_INTERP_NNN_RBF #< radial basis functions
        YAC_INTERP_NNN_ZERO #< all weights are set to zero
    cdef yac_interp_nnn_weight_type YAC_INTERP_NNN_WEIGHTED_DEFAULT
    cdef int YAC_INTERP_NNN_N_DEFAULT
    cdef double YAC_INTERP_NNN_MAX_SEARCH_DISTANCE_DEFAULT
    cdef double YAC_INTERP_NNN_GAUSS_SCALE_DEFAULT
    cdef int YAC_INTERP_RBF_N_DEFAULT
    cdef double YAC_INTERP_RBF_MAX_SEARCH_DISTANCE_DEFAULT
    cdef double YAC_INTERP_RBF_SCALE_DEFAULT
    cdef int YAC_INTERP_RBF_KERNEL_DEFAULT
 
    cpdef enum yac_interp_file_on_missing_file:
        YAC_INTERP_FILE_MISSING_ERROR
        YAC_INTERP_FILE_MISSING_CONT
    cpdef enum yac_interp_file_on_success:
        YAC_INTERP_FILE_SUCCESS_STOP
        YAC_INTERP_FILE_SUCCESS_CONT
    cdef yac_interp_file_on_missing_file YAC_INTERP_FILE_ON_MISSING_FILE_DEFAULT
    cdef yac_interp_file_on_success YAC_INTERP_FILE_ON_SUCCESS_DEFAULT
 
    cdef const char* YAC_INTERP_CHECK_CONSTRUCTOR_KEY_DEFAULT
    cdef const char* YAC_INTERP_CHECK_DO_SEARCH_KEY_DEFAULT
 
    cdef int YAC_INTERP_CREEP_DISTANCE_DEFAULT
 
    struct yac_basic_grid_data:
        yac_int * cell_ids;
        yac_int * vertex_ids;
        yac_int * edge_ids;
        int * core_cell_mask;
        int * core_vertex_mask;
        int * core_edge_mask;
 
    struct yac_basic_grid:
        pass
    yac_basic_grid_data * yac_basic_grid_get_data(
        yac_basic_grid * grid);
    struct yac_interp_field:
        yac_location location
        size_t coordinates_idx
        size_t masks_idx
    void yac_basic_grid_to_file_parallel(
        yac_basic_grid * grid,
        const char * filename,
        MPI_Comm comm)
    const char * yac_basic_grid_get_name(yac_basic_grid * grid);
    size_t yac_basic_grid_get_data_size(yac_basic_grid * grid, yac_location location);
    void yac_basic_grid_delete(yac_basic_grid * grid);
    size_t yac_basic_grid_add_coordinates(
        yac_basic_grid * grid, yac_location location,
        yac_coordinate_pointer coordinates, size_t count);
    size_t yac_basic_grid_add_mask(
        yac_basic_grid * grid, yac_location location,
        const int * mask, size_t count, const char * mask_name);
    yac_basic_grid * yac_basic_grid_reg_2d_new(
        const char * name, size_t nbr_vertices[2], int cyclic[2],
        double *lon_vertices, double *lat_vertices)
    yac_basic_grid * yac_basic_grid_empty_new(const char* name)
    yac_basic_grid * yac_basic_grid_unstruct_new(
        const char * name, size_t nbr_vertices, size_t nbr_cells,
        int *num_vertices_per_cell, double *x_vertices, double *y_vertices,
        int *cell_to_vertex);
    yac_basic_grid * yac_basic_grid_unstruct_ll_new(
        const char * name, size_t nbr_vertices, size_t nbr_cells,
        int *num_vertices_per_cell, double *x_vertices, double *y_vertices,
        int *cell_to_vertex);
    yac_basic_grid * yac_basic_grid_unstruct_edge_new(
        const char * name, size_t nbr_vertices, size_t nbr_cells, size_t nbr_edges,
        int *num_edges_per_cell, double *x_vertices, double *y_vertices,
        int *cell_to_edge, int *edge_to_vertex);
    yac_basic_grid * yac_basic_grid_unstruct_edge_ll_new(
        const char * name, size_t nbr_vertices, size_t nbr_cells, size_t nbr_edges,
        int *num_edges_per_cell, double *x_vertices, double *y_vertices,
        int *cell_to_edge, int *edge_to_vertex);
    int yac_basic_grid_check_coordinates(
        yac_basic_grid * grid, int is_fatal);
 
    struct yac_dist_grid_pair:
        pass
    yac_dist_grid_pair * yac_dist_grid_pair_new(
        yac_basic_grid * grid_a, yac_basic_grid * grid_b,
        MPI_Comm comm);
    void yac_dist_grid_pair_delete(yac_dist_grid_pair * grid_pair);
 
    struct yac_interp_grid:
        pass
    yac_interp_grid * yac_interp_grid_new(
        yac_dist_grid_pair * grid_pair,
        const char * src_grid_name, const char * tgt_grid_name,
        size_t num_src_fields, const yac_interp_field * src_fields,
        const yac_interp_field tgt_field);
    void yac_interp_grid_delete(yac_interp_grid * interp_grid);
 
    struct interp_method:
        pass
    yac_interp_weights * yac_interp_method_do_search(
        interp_method ** method, yac_interp_grid * grid);
    void yac_interp_method_delete(interp_method ** method);
 
    struct yac_interp_stack_config:
        pass
    struct yac_interp_method:
        pass
    yac_interp_stack_config * yac_interp_stack_config_new();
    void yac_interp_stack_config_delete(
        yac_interp_stack_config * interp_stack_config);
    interp_method ** yac_interp_stack_config_generate(
        yac_interp_stack_config * interp_stack);
    
    void yac_interp_stack_config_add_average(
        yac_interp_stack_config * interp_stack_config,
        yac_interp_avg_weight_type reduction_type, int partial_coverage);
    void yac_interp_stack_config_add_ncc(
        yac_interp_stack_config * interp_stack_config,
        yac_interp_ncc_weight_type type, int partial_coverage);
    void yac_interp_stack_config_add_dnn(
        yac_interp_stack_config * interp_stack_config,
        yac_interp_dnn_weight_type type,
        yac_interp_dnn_search_distance_type search_distance_type,
        double search_distance, double scale);
    void yac_interp_stack_config_add_nnn(
        yac_interp_stack_config * interp_stack_config,
        yac_interp_nnn_weight_type type, size_t n,
        double max_search_distance, double scale);
    void yac_interp_stack_config_add_conservative(
        yac_interp_stack_config * interp_stack_config,
        int order, int enforced_conserv, int partial_coverage,
        yac_interp_method_conserv_normalisation normalisation);
    void yac_interp_stack_config_add_spmap(
        yac_interp_stack_config * interp_stack_config,
        double spread_distance, double max_search_distance,
        yac_interp_spmap_weight_type weight_type,
        yac_interp_spmap_scale_type scale_type,
        double src_sphere_radius, const char * src_filename,
        const char * src_varname, int src_min_global_id,
        double tgt_sphere_radius, const char * tgt_filename,
        const char * tgt_varname, int tgt_min_global_id);
    void yac_interp_stack_config_add_fixed(
        yac_interp_stack_config * interp_stack_config, double value);
    void yac_interp_stack_config_add_user_file(
        yac_interp_stack_config * interp_stack_config,
        const char * filename, yac_interp_file_on_missing_file on_missing_file,
        yac_interp_file_on_success on_success);
    void yac_interp_stack_config_add_check(
        yac_interp_stack_config * interp_stack_config,
        const char * constructor_key, const char * do_search_key);
    void yac_interp_stack_config_add_hcsbb(
        yac_interp_stack_config * interp_stack_config);
    void yac_interp_stack_config_add_rbf(
        yac_interp_stack_config * interp_stack_config,
        size_t n, double max_search_distance, double scale);
    void yac_interp_stack_config_add_creep(
        yac_interp_stack_config * interp_stack_config, int creep_distance);
 
    struct yac_interp_method_config:
        pass
    yac_interp_method_config * yac_interp_method_config_default_from_name_new(
        const char * method_name);
    void yac_interp_method_config_delete(
        yac_interp_method_config * config);
    void yac_interp_stack_config_add(
        yac_interp_stack_config * interp_stack_config,
        yac_interp_method_config * method_config);
    void yac_interp_method_config_set_param_str(
        yac_interp_method_config * config,
        const char * path, const char * value);
    void yac_interp_method_config_set_param_int(
        yac_interp_method_config * config,
        const char * path, int value);
    void yac_interp_method_config_set_param_dble(
        yac_interp_method_config * config,
        const char * path, double value);
    int yac_interp_method_config_get_param_int(
        const yac_interp_method_config * config,
        const char * path);
    double yac_interp_method_config_get_param_dble(
        const yac_interp_method_config * config,
        const char * path);
    char * yac_interp_method_config_get_param_str(
        const yac_interp_method_config * config,
        const char * path);
    size_t yac_interp_method_config_get_list_size(
        const yac_interp_method_config * config,
        const char * path);
    size_t yac_interp_method_config_append_list_element(
        yac_interp_method_config * config,
        const char * path);
 
    struct yac_interpolation:
        pass
    void yac_interpolation_execute_frac(
        yac_interpolation * interp, double *** src_fields,
        double *** src_frac_masks, double ** tgt_field);
    void yac_interpolation_execute_put_frac(
        yac_interpolation * interp, double *** src_fields, double *** src_frac_masks);
    int yac_interpolation_execute_put_test(yac_interpolation * interp);
    void yac_interpolation_execute_get_async(
        yac_interpolation * interp, double ** tgt_field);
    int yac_interpolation_execute_get_test(yac_interpolation * interp);
 
    struct yac_collection_selection:
        pass
    yac_collection_selection * yac_collection_selection_new(
        size_t collection_size, const size_t * selection_indices);
    void yac_collection_selection_delete(
        yac_collection_selection * collection_selection);
 
    cpdef enum yac_interp_weights_reorder_type:
        YAC_MAPPING_ON_SRC
        YAC_MAPPING_ON_TGT
    struct yac_interpolation_gen_config:
        pass
    yac_interpolation_gen_config * yac_interpolation_gen_config_new();
    void yac_interpolation_gen_config_delete(yac_interpolation_gen_config * config);
    void yac_interpolation_gen_config_set_reorder(yac_interpolation_gen_config * config, yac_interp_weights_reorder_type reorder)
    void yac_interpolation_gen_config_set_collection_size(yac_interpolation_gen_config * config, size_t collection_size);
    void yac_interpolation_gen_config_set_collection_selection(
        yac_interpolation_gen_config * config,
        const yac_collection_selection * collection_selection);
    void yac_interpolation_gen_config_set_frac_mask_fallback_value(
        yac_interpolation_gen_config * config,
        double frac_mask_fallback_value);
    void yac_interpolation_gen_config_set_scaling_factor(
        yac_interpolation_gen_config * config, double scaling_factor);
    void yac_interpolation_gen_config_set_scaling_summand(
        yac_interpolation_gen_config * config, double scaling_summand);
    void yac_interpolation_gen_config_set_yaxt_exchanger_name(
        yac_interpolation_gen_config * config, const char * name);
 
    struct yac_interp_weights:
        pass
    cpdef enum yac_weight_file_on_existing:
        YAC_WEIGHT_FILE_ERROR
        YAC_WEIGHT_FILE_KEEP
        YAC_WEIGHT_FILE_OVERWRITE
        YAC_WEIGHT_FILE_UNDEFINED
    void yac_interp_weights_write_to_file(
        yac_interp_weights * weights, const char * filename,
        const char * src_grid_name, const char * tgt_grid_name,
        size_t src_grid_size, size_t tgt_grid_size,
        yac_weight_file_on_existing on_existing);
    yac_interpolation *  yac_interp_weights_get_interpolation_ext(
        const yac_interp_weights * weights,
        const yac_interpolation_gen_config * interp_gen_config,
        int is_source, int is_target);
 
    void yac_yaxt_init(MPI_Comm comm);
    void yac_mpi_init();
    void yac_mpi_finalize();
 
    void yac_get_io_ranks(
        MPI_Comm comm, int * local_is_io, int ** io_ranks, int * num_io_ranks);
 
def version():
    """
    Returns the YAC core library version string.
 
    @see yac_get_core_version
    """
    return bytes.decode(yac_get_core_version())
 
cdef extern from * nogil:
    """
#ifndef YAC_UTILS_ENABLED
// we cannot link libyac_utils.a so we define the function here
#include "yac_assert.h"
void yac_read_icon_basic_grid_parallel_2(
    const char * filename, const char * gridname, MPI_Comm comm,
    struct yac_basic_grid ** basic_grid, size_t * cell_coordinate_idx,
    int ** cell_mask) {
        die("To read an ICON grid YAC must be configured with --enable-utils.");
}
struct yac_basic_grid * yac_read_scrip_basic_grid(
    const char * grid_filename, const char * mask_filename,
      const char * grid_name, int valid_mask_value, const char * name,
      int use_ll_edges, size_t * cell_coord_idx,
      size_t ** duplicated_cell_idx, yac_int ** orig_cell_global_ids,
      size_t * nbr_duplicated_cells) {
          die("To read a SCRIP grid YAC must be configured with --enable-utils.");
          return NULL;
}
#else
      #include "yac_utils.h"
#endif
    """
    void yac_read_icon_basic_grid_parallel_2(
        const char * filename, const char * gridname, MPI_Comm comm,
        yac_basic_grid ** basic_grid, size_t * cell_coordinate_idx,
        int ** cell_mask)
 
    yac_basic_grid * yac_read_scrip_basic_grid(
        const char * grid_filename, const char * mask_filename,
        const char * grid_name, int valid_mask_value, const char * name,
        int use_ll_edges, size_t * cell_coord_idx,
        size_t ** duplicated_cell_idx, yac_int ** orig_cell_global_ids,
        size_t * nbr_duplicated_cells);
 
 
def lonlat2xyz(lon, lat):
    clat = _np.cos(lat)
    slat = _np.sin(lat)
    return clat * _np.cos(lon), clat * _np.sin(lon), slat
 
cdef _decode(const char* c_str):
    if c_str == NULL:
        return None
    return bytes.decode(c_str)
 
cdef int get_c_comm(comm, MPI_Comm* c_comm) except -1:
    """
    Convert comm into an MPI_Comm.
    Comm can be given as an fortran comm (int) or as an mpi4py.Comm
    """
    try:
        from mpi4py import MPI
        if isinstance(comm, MPI.Comm):
            comm = comm.py2f()
    except:
        pass
    if comm is None:
        c_comm[0] = MPI_COMM_WORLD
    elif isinstance(comm, int):
        c_comm[0] = MPI_Comm_f2c(comm)
    else:
        raise ValueError("Invalid communicator")
 
    return 0
 
def _to_yac_int_np_array(arr):
    cdef yac_int tmp;
    dtype = _np.asarray(<yac_int[:1]>&tmp).dtype
    assert dtype.itemsize == sizeof(yac_int), "dtype detected did not work"
    info = _np.iinfo(dtype)
    assert _np.min(arr) >= info.min and _np.max(arr) <= info.max, \
        f"Indices must be within the valid range for yac_int ({info.min:_} to {info.max:_})"
    return _np.ascontiguousarray(arr, dtype=dtype)
 
 
logger.info("Initializing MPI")
yac_mpi_init();
logger.info("Initializing yaxt")
yac_yaxt_init(MPI_COMM_WORLD);
 
Coordinates = namedtuple("Coordinates", ["basic_grid", "location", "coord_idx"])
Mask = namedtuple("Mask", ["basic_grid", "location", "mask_idx"])
 
cdef class BasicGrid:
    cdef yac_basic_grid* _ptr;
 
    def __init__(self, yac_basic_grid_ptr):
        raise TypeError("This class cannot be instantiated directly.")
 
    def __del__(self):
        logger.info("calling yac_basic_grid_delete")
        yac_basic_grid_delete(self._ptr)
 
    @staticmethod
    cdef BasicGrid from_ptr(yac_basic_grid* basic_grid_ptr):
        cdef BasicGrid basic_grid = BasicGrid.__new__(BasicGrid)
        basic_grid._ptr = basic_grid_ptr
        return basic_grid
 
    def to_file(self, filename, comm=None):
        filename_bytes = filename.encode()
        cdef const char* _filename = filename_bytes
        cdef MPI_Comm c_comm = MPI_COMM_NULL
        get_c_comm(comm, &c_comm)
        logger.info("calling yac_basic_grid_to_file_parallel")
        with nogil:
            yac_basic_grid_to_file_parallel(self._ptr,
                                            _filename,
                                            c_comm);
 
    @property
    def name(self):
        cdef const char * _name = yac_basic_grid_get_name(self._ptr);
        return _decode(_name)
 
    def __str__(self):
        return f"yac_basic_grid(name={self.name})"
 
    def get_data_size(self, location : yac_location):
        logger.info("calling yac_basic_grid_get_data_size")
        return yac_basic_grid_get_data_size(self._ptr, location);
 
    def set_global_index(self,
                         location : yac_location,
                         indices):
        cdef yac_basic_grid_data* data
        cdef yac_int** grid_global_ids;
        cdef size_t count = self.get_data_size(location)
        indices = _to_yac_int_np_array(indices)
        data = yac_basic_grid_get_data(self._ptr)
        assert indices.shape == (count,), "Wrong indices shape"
        if location == YAC_LOC_CELL:
            grid_global_ids = &(data.cell_ids)
        elif location == YAC_LOC_CORNER:
            grid_global_ids = &(data.vertex_ids)
        elif location == YAC_LOC_EDGE:
            grid_global_ids = &(data.edge_ids)
        else:
            raise RuntimeError("Invalid location")
        if grid_global_ids[0] == NULL:
            grid_global_ids[0] = <yac_int*>malloc(count * sizeof(yac_int));
        logger.info("setting global index")
        (<yac_int[:count]>(grid_global_ids[0]))[:] = indices
 
    def set_core_mask(self, location :yac_location, mask):
        cdef yac_basic_grid_data* data
        cdef int** grid_core_mask;
        cdef int count = self.get_data_size(location)
        data = yac_basic_grid_get_data(self._ptr)
        mask = _np.ascontiguousarray(mask, dtype=_np.intc)
        assert mask.shape == (count,), "Wrong mask shape"
        if location == YAC_LOC_CELL:
            grid_core_mask = &(data.core_cell_mask)
        elif location == YAC_LOC_CORNER:
            grid_core_mask = &(data.core_vertex_mask)
        elif location == YAC_LOC_EDGE:
            grid_core_mask = &(data.core_edge_mask)
        else:
            raise RuntimeError("Invalid location")
        if grid_core_mask[0] == NULL:
            grid_core_mask[0] = <int*>malloc(count * sizeof(int));
        logger.info("setting core mask")
        (<int[:count]>(grid_core_mask[0]))[:] = mask
 
    def add_coordinates(self,
                        location : yac_location,
                        lon = None,
                        lat = None):
        assert (lon is None) == (lat is None), "lon and lat must both be either set or unset."
        cdef double[:,::1] coords_view
        cdef size_t coords_len = 0
        cdef yac_coordinate_pointer coords_ptr = NULL
        cdef size_t idx = SIZE_MAX
        if lon is not None:
            coords_xyz = _np.stack(lonlat2xyz(lon, lat), axis=-1)
            coords_view = _np.ascontiguousarray(coords_xyz, dtype=_np.double)
            coords_ptr = <yac_coordinate_pointer>&coords_view[0,0]
            coords_len = len(coords_view)
            logger.info("calling yac_basic_grid_add_coordinates")
            idx = yac_basic_grid_add_coordinates(
                self._ptr,
                location,
                coords_ptr,
                coords_len)
        return Coordinates(basic_grid=self, location=location, coord_idx=idx)
 
    def add_mask(self,
                 location : yac_location,
                 mask,
                 mask_name : str | None = None
                 ):
        cdef int[:] mask_view = _np.ascontiguousarray(_np.asarray(mask), dtype=_np.intc)
        cdef const char* _mask_name = NULL
        if mask_name is not None:
            mask_name_bytes = mask_name.encode()
            _mask_name = mask_name_bytes
            logger.info("calling yac_basic_grid_add_mask")
        cdef size_t idx =  yac_basic_grid_add_mask(
            self._ptr,
            location,
            <int*>&mask_view[0], len(mask), _mask_name)
        return Mask(basic_grid=self, location=location, mask_idx=idx)
 
    def check_coordinates(self, is_fatal : bool = True):
        logger.info("calling yac_basic_grid_check_coordinates")
        return yac_basic_grid_check_coordinates(self._ptr, 1 if is_fatal else 0) != 0
 
    @classmethod
    def read_icon(cls,
                  filename :str,
                  gridname :str,
                  comm = None,
                  mask_valid_values=None):
 
        filename_bytes = filename.encode()
        cdef const char* _filename = filename_bytes
        gridname_bytes = gridname.encode()
        cdef const char* _gridname = gridname_bytes
        cdef MPI_Comm c_comm = MPI_COMM_NULL
        get_c_comm(comm, &c_comm)
        cdef yac_basic_grid * _yac_basic_grid
        cdef size_t cell_coordinate_idx
        cdef int * cell_mask;
 
        logger.info("calling yac_read_icon_basic_grid_parallel_2")
        yac_read_icon_basic_grid_parallel_2(
            _filename, _gridname, c_comm,
            &_yac_basic_grid,
            &cell_coordinate_idx, &cell_mask)
 
        basic_grid = BasicGrid.from_ptr(_yac_basic_grid)
        cell_coords = Coordinates(basic_grid=basic_grid,
                                  location=YAC_LOC_CELL,
                                  coord_idx=cell_coordinate_idx)
 
        _cell_mask = None
        if cell_mask != NULL and mask_valid_values is not None:
            num_cells = basic_grid.get_data_size(YAC_LOC_CELL)
            _cell_mask = basic_grid.add_mask(YAC_LOC_CELL,
                                             _np.isin(<int[:num_cells]>cell_mask, mask_valid_values))
 
        return basic_grid, InterpField(cell_coords, _cell_mask)
 
    @classmethod
    def read_scrip(cls,
                   grid_filename,
                   mask_filename,
                   grid_name,
                   name,
                   valid_mask_value = 1,
                   use_ll_edges = False,
                   ):
        cdef size_t cell_coord_idx;
        cdef size_t* duplicated_cell_idx;
        cdef yac_int* orig_cell_global_ids;
        cdef size_t nbr_duplicated_cells;
        cdef yac_basic_grid * yac_basic_grid;
        grid_filename_bytes = grid_filename.encode()
        cdef const char* _grid_filename = grid_filename_bytes
        mask_filename_bytes = mask_filename.encode()
        cdef const char* _mask_filename = mask_filename_bytes
        grid_name_bytes = grid_name.encode()
        cdef const char* _grid_name = grid_name_bytes
        name_bytes = name.encode()
        cdef const char* _name = name_bytes
        cdef int _valid_mask_value = valid_mask_value
        cdef int _use_ll_edges = use_ll_edges
        logger.info("calling yac_read_scrip_basic_grid")
        with nogil:
            yac_basic_grid = yac_read_scrip_basic_grid(
                _grid_filename,
                _mask_filename,
                _grid_name,
                _valid_mask_value,
                _name,
                _use_ll_edges,
                &cell_coord_idx,
                &duplicated_cell_idx,
                &orig_cell_global_ids,
                &nbr_duplicated_cells);
        basic_grid = BasicGrid.from_ptr(yac_basic_grid)
        return (basic_grid,
                InterpField(Coordinates(basic_grid=basic_grid,
                                         location=YAC_LOC_CELL,
                                         coord_idx=cell_coord_idx)),
                )
 
    @classmethod
    @cython.boundscheck(False)
    def reg_2d_new(cls, name, lon_vertices, lat_vertices, cyclic = [False, False]):
        cdef const double[::1] x = _np.ascontiguousarray(lon_vertices, dtype=_np.double)
        cdef const double[::1] y = _np.ascontiguousarray(lat_vertices, dtype=_np.double)
        cdef int[2] cyclic_view = cyclic
        name_bytes = name.encode()
        cdef const char* _name = name_bytes
        cdef size_t[2] _len = [len(x), len(y)]
        cdef int[2] _cyclic = cyclic
        logger.info("calling yac_basic_grid_reg_2d_new")
        with nogil:
            yac_basic_grid = yac_basic_grid_reg_2d_new(
                _name, _len, _cyclic,
                <double*>&x[0], <double*>&y[0])
        return BasicGrid.from_ptr(yac_basic_grid)
 
    @classmethod
    @cython.boundscheck(False)
    def unstruct_new(cls,
                     name : str,
                     num_vertices_per_cell,
                     x_vertices,
                     y_vertices,
                     cell_to_vertex,
                     lonlat_edges = False):
        assert len(x_vertices) == len(y_vertices)
        cdef int[:] _num_vertices_per_cell = _np.ascontiguousarray(num_vertices_per_cell, dtype=_np.intc)
        cdef double[:] _x_vertices = _np.ascontiguousarray(x_vertices, dtype=_np.double)
        cdef double[:] _y_vertices = _np.ascontiguousarray(y_vertices, dtype=_np.double)
        cdef int[:] _cell_to_vertex = _np.ascontiguousarray(cell_to_vertex, dtype=_np.intc)
        cdef int[:] _num_edges_per_cell
        cdef int[:] _edge_to_vertex
        cdef yac_basic_grid * ptr
        name_bytes = name.encode()
        cdef const char* _name = name_bytes
        cdef size_t nbr_vertices = len(x_vertices)
        cdef size_t nbr_cells = len(num_vertices_per_cell)
        if lonlat_edges:
            logger.info("calling yac_basic_grid_unstruct_ll_new")
            with nogil:
                ptr = yac_basic_grid_unstruct_ll_new(
                    _name,
                    nbr_vertices,
                    nbr_cells,
                    <int*>&_num_vertices_per_cell[0],
                    <double *>&_x_vertices[0],
                    <double *>&_y_vertices[0],
                    <int *>&_cell_to_vertex[0]);
        else:
            logger.info("calling yac_basic_grid_unstruct_new")
            with nogil:
                ptr = yac_basic_grid_unstruct_new(
                    _name,
                    nbr_vertices,
                    nbr_cells,
                    <int*>&_num_vertices_per_cell[0],
                    <double *>&_x_vertices[0],
                    <double *>&_y_vertices[0],
                    <int *>&_cell_to_vertex[0]);
        return BasicGrid.from_ptr(ptr)
 
    @classmethod
    @cython.boundscheck(False)
    def unstruct_edge_new(cls,
                          name : str,
                          num_edges_per_cell,
                          x_vertices,
                          y_vertices,
                          cell_to_edge,
                          edge_to_vertex,
                          lonlat_edges : bool = False):
        assert len(x_vertices) == len(y_vertices)
        cdef int[:] _num_edges_per_cell = _np.ascontiguousarray(num_edges_per_cell, dtype=_np.intc)
        cdef double[:] _x_vertices = _np.ascontiguousarray(x_vertices, dtype=_np.double)
        cdef double[:] _y_vertices = _np.ascontiguousarray(y_vertices, dtype=_np.double)
        cdef int[:] _cell_to_edge = _np.ascontiguousarray(cell_to_edge, dtype=_np.intc)
        cdef int[:] _edge_to_vertex = _np.ascontiguousarray(edge_to_vertex, dtype=_np.intc)
        name_bytes = name.encode()
        cdef const char* _name = name_bytes
        cdef size_t nbr_vertices = len(x_vertices)
        cdef size_t nbr_cells = len(num_edges_per_cell)
        cdef size_t nbr_edges = _np.max(_cell_to_edge)+1
        if lonlat_edges:
            logger.info("calling yac_basic_grid_unstruct_edge_ll_new")
            with nogil:
                ptr = yac_basic_grid_unstruct_edge_ll_new(
                    _name,
                    nbr_vertices,
                    nbr_cells,
                    nbr_edges,
                    <int*>&_num_edges_per_cell[0],
                    <double*>&_x_vertices[0],
                    <double*>&_y_vertices[0],
                    <int*>&_cell_to_edge[0],
                    <int*>&_edge_to_vertex[0]);
        else:
            logger.info("calling yac_basic_grid_unstruct_edge_new")
            with nogil:
                ptr = yac_basic_grid_unstruct_edge_new(
                    _name,
                    nbr_vertices,
                    nbr_cells,
                    nbr_edges,
                    <int*>&_num_edges_per_cell[0],
                    <double*>&_x_vertices[0],
                    <double*>&_y_vertices[0],
                    <int*>&_cell_to_edge[0],
                    <int*>&_edge_to_vertex[0]);
 
        return BasicGrid.from_ptr(ptr)
 
    @classmethod
    def empty(cls, name : str):
        logger.info("calling yac_basic_grid_empty_new")
        cdef yac_basic_grid * ptr = yac_basic_grid_empty_new(name.encode())
        return BasicGrid.from_ptr(ptr)
 
    @classmethod
    def from_uxgrid(cls,
                    name : str,
                    uxgrid,
                    def_edges : bool = False,
                    lonlat_edges : bool = False):
        from uxarray import INT_FILL_VALUE
        if def_edges:
            cell_to_edge = _np.asarray(uxgrid.face_edge_connectivity).reshape((-1,), order="C")
            cell_to_edge = cell_to_edge[cell_to_edge!=INT_FILL_VALUE]
            edge_to_vertex = _np.asarray(uxgrid.edge_node_connectivity).reshape((-1,), order="C")
            edge_to_vertex = edge_to_vertex[edge_to_vertex!=INT_FILL_VALUE]
            grid = cls.unstruct_edge_new(
                name,
                uxgrid.n_nodes_per_face, # n_nodes_per_face == n_edges_per_face (2D)
                _np.deg2rad(uxgrid.node_lon),
                _np.deg2rad(uxgrid.node_lat),
                cell_to_edge,
                edge_to_vertex,
                lonlat_edges = lonlat_edges)
        else:
            cell_to_vertex = _np.asarray(uxgrid.face_node_connectivity).reshape((-1,), order="C")
            cell_to_vertex = cell_to_vertex[cell_to_vertex!=INT_FILL_VALUE]
            grid = cls.unstruct_new(
                name,
                uxgrid.n_nodes_per_face, # n_nodes_per_face == n_edges_per_face (2D)
                _np.deg2rad(uxgrid.node_lon),
                _np.deg2rad(uxgrid.node_lat),
                cell_to_vertex,
                lonlat_edges = lonlat_edges)
        grid.set_global_index(YAC_LOC_CELL,
                              range(uxgrid.n_face))
        grid.set_global_index(YAC_LOC_EDGE,
                              range(uxgrid.n_edge))
        grid.set_global_index(YAC_LOC_CORNER,
                              range(uxgrid.n_node))
        return grid
 
 
cdef class InterpField:
    cdef yac_interp_field _interp_field
    cdef BasicGrid _basic_grid
 
    def __init__(self,
                 coordinates : Coordinates,
                 mask : Mask = None):
        self._basic_grid = coordinates.basic_grid
        self._interp_field.location = coordinates.location
        self._interp_field.coordinates_idx = coordinates.coord_idx
        cdef masks_idx = SIZE_MAX
        if mask is not None:
            assert mask.location == coordinates.location, \
                "coordinates and mask have different location"
            masks_idx = mask.mask_idx
        self._interp_field.masks_idx = masks_idx
 
    def __str__(self):
        return f"yac_interp_field(grid={self._basic_grid}, location={self.location}, coord_idx={self._interp_field.coordinates_idx}, mask_idx={self._interp_field.masks_idx})"
 
    @property
    def location(self):
        return self._interp_field.location
 
    @property
    def basic_grid(self):
        return self._basic_grid
 
cdef class DistGridPair:
    cdef yac_dist_grid_pair* _ptr
    _str :str
 
    def __init__(self, grid_a : BasicGrid,
                 grid_b : BasicGrid,
                 comm = None):
        cdef MPI_Comm c_comm = MPI_COMM_NULL
        get_c_comm(comm, &c_comm)
        logger.info("calling yac_dist_grid_pair_new")
        self._ptr = yac_dist_grid_pair_new(
            grid_a._ptr, grid_b._ptr,
            c_comm);
        self._str = f"yac_dist_grid_pair(grid_a={grid_a}, grid_b={grid_b}"
 
    def __str__(self):
        return self._str
 
    def __del__(self):
        logger.info("calling yac_dist_grid_pair_delete")
        yac_dist_grid_pair_delete(self._ptr);
 
cdef class InterpGrid:
    cdef yac_interp_grid* _ptr
    cdef InterpField _src_field
    cdef InterpField _tgt_field
 
    def __init__(self,
                 grid_pair : DistGridPair,
                 src_field : InterpField,  # TODO: support multiple src_fields
                 tgt_field : InterpField):
        src_grid_name = src_field._basic_grid.name
        tgt_grid_name = tgt_field._basic_grid.name
        logger.info("calling yac_interp_grid_new")
        self._ptr = yac_interp_grid_new(grid_pair._ptr,
            src_grid_name.encode(), tgt_grid_name.encode(),
            1, &src_field._interp_field,
            tgt_field._interp_field);
        self._src_field = src_field
        self._tgt_field = tgt_field
 
    def __str__(self):
        return f"yac_interp_grid(src_field={self._src_field}, tgt_field={self._tgt_field})"
 
    def __del__(self):
        logger.info("calling yac_interp_grid_delete")
        yac_interp_grid_delete(self._ptr)
 
cdef class InterpolationStack:
    cdef yac_interp_stack_config* _ptr
    interpolation_register = {}
 
    def __init__(self):
        logger.info("calling yac_interp_stack_config_new")
        self._ptr = yac_interp_stack_config_new()
 
    @classmethod
    def from_list(cls, interp_list):
        interp_stack = InterpolationStack()
        interp_list = [(t, {}) if type(t) is str else t for t in interp_list]
        for name, kwargs in interp_list:
            interp_stack.add(name, **kwargs)
        return interp_stack
 
    def __del__(self):
        logger.info("calling yac_interp_stack_config_delete")
        yac_interp_stack_config_delete(self._ptr);
 
    def register_interpolation(name):
        def _register(method):
            InterpolationStack.interpolation_register[name] = method
            return method
        return _register
 
    def add(self, name, *args, **kwargs):
        self.interpolation_register[name](self, *args, **kwargs)
 
    @register_interpolation("average")
    def add_average(self,
                    weight_type : yac_interp_avg_weight_type | str = YAC_INTERP_AVG_WEIGHT_TYPE_DEFAULT,
                    partial_coverage : bool = YAC_INTERP_AVG_PARTIAL_COVERAGE_DEFAULT):
        if isinstance(weight_type, str):
            weight_type = getattr(yac_interp_avg_weight_type, weight_type)
        logger.info("calling yac_interp_stack_config_add_average")
        yac_interp_stack_config_add_average(
            self._ptr, weight_type, partial_coverage);
 
    @register_interpolation("ncc")
    def add_ncc(self,
                weight_type : yac_interp_ncc_weight_type | str = YAC_INTERP_NCC_WEIGHT_TYPE_DEFAULT,
                partial_coverage : bool = YAC_INTERP_NCC_PARTIAL_COVERAGE_DEFAULT):
        if isinstance(weight_type, str):
            weight_type = getattr(yac_interp_ncc_weight_type, weight_type)
        logger.info("calling yac_interp_stack_config_add_ncc")
        yac_interp_stack_config_add_ncc(self._ptr, weight_type, partial_coverage);
 
    @register_interpolation("dnn")
    def add_dnn(self,
                weight_type : yac_interp_dnn_weight_type | str = YAC_INTERP_DNN_WEIGHT_DEFAULT,
                search_distance_type : yac_interp_dnn_search_distance_type | str = YAC_INTERP_DNN_SEARCH_DISTANCE_TYPE_DEFAULT,
                search_distance : float = YAC_INTERP_DNN_SEARCH_DISTANCE_DEFAULT,
                scale : float = YAC_INTERP_DNN_SCALE_DEFAULT):
        if isinstance(weight_type, str):
            weight_type = getattr(yac_interp_dnn_weight_type, weight_type)
        if isinstance(search_distance_type, str):
            search_distance_type = getattr(yac_interp_dnn_search_distance_type, search_distance_type)
        logger.info("calling yac_interp_stack_config_add_dnn")
        yac_interp_stack_config_add_dnn(
            self._ptr, weight_type, search_distance_type, search_distance, scale);
 
    @register_interpolation("nnn")
    def add_nnn(self,
                nnn_type : yac_interp_nnn_weight_type | str = YAC_INTERP_NNN_WEIGHTED_DEFAULT,
                n : int = YAC_INTERP_NNN_N_DEFAULT,
                max_search_distance: float = YAC_INTERP_NNN_MAX_SEARCH_DISTANCE_DEFAULT,
                scale : float = YAC_INTERP_NNN_GAUSS_SCALE_DEFAULT,
                ):
        if isinstance(nnn_type, str):
            nnn_type = getattr(yac_interp_nnn_weight_type, nnn_type)
        logger.info("calling yac_interp_stack_config_add_nnn")
        yac_interp_stack_config_add_nnn(
            self._ptr, nnn_type, n, max_search_distance, scale);
 
    @register_interpolation("conservative")
    def add_conservative(self,
                         order : int = YAC_INTERP_CONSERV_ORDER_DEFAULT,
                         enforced_conserv : bool = YAC_INTERP_CONSERV_ENFORCED_CONSERV_DEFAULT,
                         partial_coverage : bool = YAC_INTERP_CONSERV_PARTIAL_COVERAGE_DEFAULT,
                         normalisation : yac_interp_method_conserv_normalisation | str = YAC_INTERP_CONSERV_NORMALISATION_DEFAULT
                         ):
        assert order in {1,2}, "Invalid order"
        if isinstance(normalisation, str):
            normalisation = getattr(yac_interp_method_conserv_normalisation, normalisation)
        logger.info("calling yac_interp_stack_config_add_conservative")
        yac_interp_stack_config_add_conservative(
            self._ptr, order, enforced_conserv, partial_coverage, normalisation);
 
    @register_interpolation("source_to_target_map")
    def add_source_to_target_map(self,
                                 spread_distance : float = YAC_INTERP_SPMAP_SPREAD_DISTANCE_DEFAULT,
                                 max_search_distance : float = YAC_INTERP_SPMAP_MAX_SEARCH_DISTANCE_DEFAULT,
                                 weight_type : yac_interp_spmap_weight_type | str = YAC_INTERP_SPMAP_WEIGHTED_DEFAULT,
                                 scale_type : yac_interp_spmap_scale_type | str = YAC_INTERP_SPMAP_SCALE_TYPE_DEFAULT,
                                 src_sphere_radius : float = 0.0,
                                 src_filename : str | None = None,
                                 src_varname : str | None = None,
                                 src_min_global_id : int = YAC_INTERP_SPMAP_MIN_GLOBAL_ID_DEFAULT,
                                 tgt_sphere_radius : float = 0.0,
                                 tgt_filename : str | None = None,
                                 tgt_varname : str | None = None,
                                 tgt_min_global_id : int = YAC_INTERP_SPMAP_MIN_GLOBAL_ID_DEFAULT,
                                 ):
        cdef const char* _src_filename = NULL
        cdef const char* _src_varname = NULL
        cdef const char* _tgt_filename = NULL
        cdef const char* _tgt_varname = NULL
        if src_filename is not None:
            src_filename_bytes = bytes(src_filename)
            _src_filename = src_filename_bytes
        if src_varname is not None:
            src_varname_bytes = bytes(src_varname)
            _src_varname = src_varname_bytes
        if tgt_filename is not None:
            tgt_filename_bytes = bytes(tgt_filename)
            _tgt_filename = tgt_filename_bytes
        if src_varname is not None:
            src_varname_bytes = bytes(src_varname)
            _src_varname = src_varname_bytes
        if isinstance(weight_type, str):
            weight_type = getattr(yac_interp_spmap_weight_type, weight_type)
        if isinstance(scale_type, str):
            scale_type = getattr(yac_interp_spmap_scale_type, scale_type)
        logger.info("calling yac_interp_stack_config_add_spmap")
        yac_interp_stack_config_add_spmap(
            self._ptr, spread_distance, max_search_distance, weight_type,
            scale_type, src_sphere_radius, _src_filename, _src_varname,
            src_min_global_id, tgt_sphere_radius, _tgt_filename, _tgt_varname,
            tgt_min_global_id);
 
    @register_interpolation("fixed")
    def add_fixed(self, value : float):
        logger.info("calling yac_interp_stack_config_add_fixed")
        yac_interp_stack_config_add_fixed(
            self._ptr, value);
 
    @register_interpolation("user_file")
    def add_user_file(self,
                      filename : str | None = None,
                      on_missing_file : yac_interp_file_on_missing_file | str = YAC_INTERP_FILE_ON_MISSING_FILE_DEFAULT,
                      on_success : yac_interp_file_on_success | str = YAC_INTERP_FILE_ON_SUCCESS_DEFAULT):
        cdef const char* _filename = NULL
        if filename is not None:
            filename_bytes = filename.encode()
            _filename = filename_bytes
        if isinstance(on_missing_file, str):
            on_missing_file = getattr(yac_interp_file_on_missing_file, on_missing_file)
        if isinstance(on_success, str):
            on_success = getattr(yac_interp_file_on_success, on_success)
        logger.info("calling yac_interp_stack_config_add_user_file")
        yac_interp_stack_config_add_user_file(
            self._ptr, _filename, on_missing_file, on_success);
 
    @register_interpolation("check")
    def add_check(self,
                  constructor_key : str = _decode(YAC_INTERP_CHECK_CONSTRUCTOR_KEY_DEFAULT),
                  do_search_key : str = _decode(YAC_INTERP_CHECK_DO_SEARCH_KEY_DEFAULT)):
        logger.info("calling yac_interp_stack_config_add_check")
        yac_interp_stack_config_add_check(self._ptr,
            constructor_key.encode(), do_search_key.encode());
 
    @register_interpolation("bernstein_bezier")
    def add_bernstein_bezier(self):
        logger.info("calling yac_interp_stack_config_add_hcsbb")
        yac_interp_stack_config_add_hcsbb(self._ptr);
 
    @register_interpolation("rbf")
    def add_rbf(self,
                n : int = YAC_INTERP_RBF_N_DEFAULT,
                max_search_distance : float = YAC_INTERP_RBF_MAX_SEARCH_DISTANCE_DEFAULT,
                scale : float = YAC_INTERP_RBF_SCALE_DEFAULT,
                rbf_kernel : int = YAC_INTERP_RBF_KERNEL_DEFAULT
                ):
        logger.info("calling yac_interp_stack_config_add_rbf")
        yac_interp_stack_config_add_rbf(self._ptr, n, max_search_distance, scale);
 
    @register_interpolation("creep")
    def add_creep(self,
                  creep_distance : int = YAC_INTERP_CREEP_DISTANCE_DEFAULT):
        logger.info("calling yac_interp_stack_config_add_creep")
        yac_interp_stack_config_add_creep(self._ptr, creep_distance);
 
    def add_from_config(self, config : InterpMethodConfig):
        """Add a method config object to this interpolation stack.
 
        The stack takes ownership of the config; it must not be used
        (or deleted) after this call.
        """
        logger.info("calling yac_interp_stack_config_add")
        yac_interp_stack_config_add(self._ptr, config._ptr)
        config._ptr = NULL  # prevent double-free
 
    # This method should eventually replace `add`
    def add_method(self, method, **kwargs):
        cfg = InterpMethodConfig(method, **kwargs)
        self.add_from_config(cfg)
 
cdef class InterpMethodConfig:
    """User-oriented wrapper around a YAC interpolation method configuration.
 
    Typical usage::
 
        cfg = InterpMethodConfig("nnn")
        cfg.set_param_int("n", 8)
        cfg.set_param_str("weighted", "gauss_weighted")
        stack = InterpolationStack()
        stack.add_from_config(cfg)
    """
    cdef yac_interp_method_config* _ptr
 
    def __init__(self, method_name : str, **kwargs):
        method_name_bytes = method_name.encode()
        logger.info("calling yac_interp_method_config_default_from_name_new")
        self._ptr = yac_interp_method_config_default_from_name_new(
            method_name_bytes)
        for k,v in kwargs.items():
            self.set_param(k,v)
 
    def __del__(self):
        if self._ptr != NULL:
            logger.info("calling yac_interp_method_config_delete")
            yac_interp_method_config_delete(self._ptr)
            self._ptr = NULL
 
    def set_param_str(self, path : str, value : str):
        """Set a parameter by dot-separated path from a string value."""
        logger.info("calling yac_interp_method_config_set_param_str")
        yac_interp_method_config_set_param_str(
            self._ptr, path.encode(), value.encode())
 
    def set_param_int(self, path : str, value : int):
        """Set a scalar integer parameter by dot-separated path."""
        logger.info("calling yac_interp_method_config_set_param_int")
        yac_interp_method_config_set_param_int(
            self._ptr, path.encode(), value)
 
    def set_param_dble(self, path : str, value : float):
        """Set a scalar double parameter by dot-separated path."""
        logger.info("calling yac_interp_method_config_set_param_dble")
        yac_interp_method_config_set_param_dble(
            self._ptr, path.encode(), value)
 
    def set_param(self, path : str, value : str | int | float | dict | list):
        """Set a parameter by dot-separated path, dispatching on type."""
        if isinstance(value, str):
            self.set_param_str(path, value)
        elif isinstance(value, int):
            self.set_param_int(path, value)
        elif isinstance(value, float):
            self.set_param_dble(path, value)
        elif isinstance(value, dict):
            for k,v in value.items():
                self.set_param(path+"."+k, v)
        elif isinstance(value, list):
            for v in value:
                i = self.append_list_element(path)
                self.set_param(path+"."+str(i), v)
        else:
            raise TypeError(f"Unsupported type {type(value)} for set_param")
 
    def get_param_int(self, path : str) -> int:
        """Get the integer value of a parameter by dot-separated path."""
        logger.info("calling yac_interp_method_config_get_param_int")
        return yac_interp_method_config_get_param_int(
            self._ptr, path.encode())
 
    def get_param_dble(self, path : str) -> float:
        """Get the double value of a parameter by dot-separated path."""
        logger.info("calling yac_interp_method_config_get_param_dble")
        return yac_interp_method_config_get_param_dble(
            self._ptr, path.encode())
 
    def get_param_str(self, path : str) -> str:
        """Get the string value of a parameter by dot-separated path."""
        cdef char * raw
        logger.info("calling yac_interp_method_config_get_param_str")
        raw = yac_interp_method_config_get_param_str(
            self._ptr, path.encode())
        result = raw.decode()
        free(raw)
        return result
 
    def get_list_size(self, path : str) -> int:
        """Return the number of elements in the list parameter at *path*."""
        logger.info("calling yac_interp_method_config_get_list_size")
        return yac_interp_method_config_get_list_size(
            self._ptr, path.encode())
 
    def append_list_element(self, path : str) -> int:
        """Append a default element to the list parameter at *path*.
 
        Returns the zero-based index of the newly appended element so that its
        sub-parameters can be set with paths like ``"listname.<index>.subparam"``.
        """
        logger.info("calling yac_interp_method_config_append_list_element")
        return yac_interp_method_config_append_list_element(
            self._ptr, path.encode())
 
cdef class InterpWeights:
    cdef yac_interp_weights* _ptr;
    cdef _num_src_points, _num_tgt_points;
 
    def __init__(self):
        raise TypeError("This class cannot be instantiated directly.")
 
    @staticmethod
    cdef InterpWeights from_ptr(yac_interp_weights* interp_weights_ptr,
                                num_src_points : int,
                                num_tgt_points : int):
        cdef InterpWeights interp_weights = InterpWeights.__new__(InterpWeights)
        interp_weights._ptr = interp_weights_ptr
        interp_weights._num_src_points = num_src_points
        interp_weights._num_tgt_points = num_tgt_points
        return interp_weights
 
    def write_to_file(self,
                      filename : str,
                      src_grid_name : str = "src_grid",
                      tgt_grid_name : str = "tgt_grid",
                      num_global_src_points : int = 0,  
                      num_global_tgt_points : int = 0,  
                      on_existing = YAC_WEIGHT_FILE_ERROR,
                      ):
        filename_bytes = filename.encode()
        src_grid_name_bytes = src_grid_name.encode()
        tgt_grid_name_bytes = tgt_grid_name.encode()
        cdef const char* filename_cstr = filename_bytes
        cdef const char* src_grid_name_cstr = src_grid_name_bytes
        cdef const char* tgt_grid_name_cstr = tgt_grid_name_bytes
        cdef yac_interp_weights* _ptr = self._ptr;
        cdef int num_global_src_points_c = num_global_src_points
        cdef int num_global_tgt_points_c = num_global_tgt_points
        cdef yac_weight_file_on_existing on_existing_c = on_existing
        logger.info("calling yac_interp_weights_write_to_file")
        with nogil:
            yac_interp_weights_write_to_file(
                _ptr,
                filename_cstr,
                src_grid_name_cstr,
                tgt_grid_name_cstr,
                num_global_src_points_c,
                num_global_tgt_points_c,
                on_existing_c);
 
    def get_interpolation(self,
                          reorder : yac_interp_weights_reorder_type = YAC_MAPPING_ON_SRC,
                          collection_size : int | None = None,
                          frac_mask_fallback_value : float | None = None,
                          scaling_factor : float = 1.0,
                          scaling_summand : float = 0.0,
                          yaxt_exchanger_name : str | None = None,
                          is_source : bool = True,
                          is_target : bool = True,
                          collection_selection : List[int] | None = None):
        assert (collection_size is not None) ^ (collection_selection is not None), \
          "Either collection_size or collection_selection must be provided"
        if collection_size is not None:
            collection_selection = list(range(collection_size))
        collection_size = len(collection_selection)
 
        cdef yac_collection_selection* coll_sel;
        cdef yac_interpolation_gen_config * config = yac_interpolation_gen_config_new();
        cdef size_t[:] coll_sel_np = _np.ascontiguousarray(collection_selection, dtype=_np.uintp)
        yac_interpolation_gen_config_set_reorder(config, reorder)
 
        coll_sel = yac_collection_selection_new(collection_size, <size_t*>&coll_sel_np[0]);
        yac_interpolation_gen_config_set_collection_selection(config,coll_sel);
        yac_collection_selection_delete(coll_sel);
 
        if frac_mask_fallback_value is not None:
            yac_interpolation_gen_config_set_frac_mask_fallback_value(config, frac_mask_fallback_value);
        yac_interpolation_gen_config_set_scaling_factor(config, scaling_factor);
        yac_interpolation_gen_config_set_scaling_summand(config, scaling_summand);
        if yaxt_exchanger_name is not None:
            yac_interpolation_gen_config_set_yaxt_exchanger_name(config, yaxt_exchanger_name.encode());
 
        logger.info("calling yac_interp_weights_get_interpolation_ext")
        cdef yac_interpolation * interp_ptr = yac_interp_weights_get_interpolation_ext(
            self._ptr, config, is_source, is_target);
 
        yac_interpolation_gen_config_delete(config);
        return Interpolation.from_ptr(interp_ptr,
                                      self._num_src_points,
                                      self._num_tgt_points,
                                      collection_size,
                                      )
 
def do_search(interpolation_stack : InterpolationStack, interp_grid : InterpGrid):
    cdef interp_method** interpstack = yac_interp_stack_config_generate(interpolation_stack._ptr)
    cdef yac_interp_weights * w
    with nogil:
        w = yac_interp_method_do_search(
            interpstack, interp_grid._ptr);
    logger.info("calling yac_interp_method_delete")
    yac_interp_method_delete(interpstack);
    free(interpstack);
 
    src_loc = interp_grid._src_field.location
    num_src_points = interp_grid._src_field._basic_grid.get_data_size(src_loc)
 
    tgt_loc = interp_grid._tgt_field.location
    num_tgt_points = interp_grid._tgt_field._basic_grid.get_data_size(tgt_loc)
 
    return InterpWeights.from_ptr(w, num_src_points, num_tgt_points)
 
def compute_weights(interpolation_stack : InterpolationStack,
                    src_field : InterpField | None,
                    tgt_field : InterpField | None,
                    comm = None):
    dgp = DistGridPair(src_field.basic_grid, tgt_field.basic_grid)
    ig = InterpGrid(dgp,
                    src_field,
                    tgt_field)
    return do_search(interpolation_stack,
                     ig)
 
cdef void* fill_pointer_array(array):
    cdef double[::1] view1d
    if len(array.shape) == 1:
        view1d = array
        return &view1d[0]
    cdef void** ptr = <void**> malloc(array.shape[0] * sizeof(ptr[0]))
    cdef int i
    for i in range(array.shape[0]):
        slic = array[i, ...]
        assert slic.base is not None
        ptr[i] = fill_pointer_array(slic)
    return ptr
 
cdef void free_pointer_array(void* ptr, shape):
    cdef int i
    if len(shape) > 1:
        for i in range(shape[0]):
            free_pointer_array((<void**>ptr)[i], shape[1:])
        free(ptr)
 
cdef class Interpolation:
    cdef yac_interpolation* _ptr;
    cdef int _collection_size;
    cdef _num_src_points, _num_tgt_points
 
    def __init__(self):
        raise TypeError("This class cannot be instantiated directly.")
 
    @staticmethod
    cdef Interpolation from_ptr(yac_interpolation* interpolation_ptr,
                                num_src_points, num_tgt_points,
                                collection_size):
        cdef Interpolation interp = Interpolation.__new__(Interpolation)
        interp._ptr = interpolation_ptr
        interp._num_src_points = num_src_points
        interp._num_tgt_points = num_tgt_points
        interp._collection_size = collection_size
        return interp
 
    def __call__(self, src = None, *, tgt = None, frac_mask = None):
        cdef double*** _src = NULL
        cdef double*** _frac_masks = NULL
        cdef double** _tgt = NULL
        if src is not None:
            src = src.reshape((-1, 1, self._num_src_points))
            _src = <double***>fill_pointer_array(src)
        if frac_mask is not None:
            frac_mask = frac_mask.reshape((-1, 1, self._num_src_points))
            _frac_masks = <double***>fill_pointer_array(frac_mask)
 
        if tgt is None and self._num_tgt_points > 0:
            tgt = _np.empty((self._collection_size, self._num_tgt_points), dtype=_np.double)
 
        if tgt is not None:
            tgt = tgt.reshape(self._collection_size, self._num_tgt_points)
        _tgt = <double**>fill_pointer_array(tgt)
 
        logger.info("calling yac_interpolation_execute_frac")
        with nogil:
            yac_interpolation_execute_frac(
                self._ptr, _src, _frac_masks, _tgt);
        if src is not None:
            free_pointer_array(_src, src.shape)
        if frac_mask is not None:
            free_pointer_array(_frac_masks, frac_mask.shape)
        if tgt is not None:
            free_pointer_array(_tgt, tgt.shape)
        return tgt
 
    @coroutine
    def execute_async(self, src = None, *, tgt = None, frac_mask = None):
        cdef double*** _src = NULL
        cdef double*** _frac_masks = NULL
        cdef double** _tgt = NULL
        if src is not None:
            src = src.reshape((-1, 1, self._num_src_points))
            _src = <double***>fill_pointer_array(src)
        if frac_mask is not None:
            frac_mask = frac_mask.reshape((-1, 1, self._num_src_points))
            _frac_masks = <double***>fill_pointer_array(frac_mask)
 
        if tgt is None and self._num_tgt_points > 0:
            tgt = _np.empty((self._collection_size, self._num_tgt_points), dtype=_np.double)
        if tgt is not None:
            tgt = tgt.reshape(self._collection_size, self._num_tgt_points)
        _tgt = <double**>fill_pointer_array(tgt)
 
        # wait until put buffer is free
        while yac_interpolation_execute_put_test(self._ptr) == 0:
            yield
 
        with nogil:
            yac_interpolation_execute_put_frac(self._ptr, _src, _frac_masks);
            yac_interpolation_execute_get_async(self._ptr, _tgt);
 
        # wait for the data
        while yac_interpolation_execute_get_test(self._ptr) == 0:
            yield
 
        if src is not None:
            free_pointer_array(_src, src.shape)
        if frac_mask is not None:
            free_pointer_array(_frac_masks, frac_mask.shape)
        if tgt is not None:
            free_pointer_array(_tgt, tgt.shape)
        return tgt
 
def get_io_ranks(comm=None):
    """
    Returns a list of ranks designated for I/O operations in the given
    communicator.
    """
    cdef MPI_Comm c_comm = MPI_COMM_NULL
    get_c_comm(comm, &c_comm)
    cdef int local_is_io
    cdef int *io_ranks_ptr
    cdef int num_io_ranks
    logger.info("calling yac_get_io_ranks")
    with nogil:
        yac_get_io_ranks(c_comm, &local_is_io, &io_ranks_ptr, &num_io_ranks)
    io_ranks = [io_ranks_ptr[i] for i in range(num_io_ranks)]
    free(io_ranks_ptr)
    return bool(local_is_io), io_ranks
 
if Py_AtExit(yac_mpi_finalize) < 0:
    warn("could not register yac_mpi_finalize with Py_AtExit()")