Source code for pycif.plugins.transforms.basic.regrid.utils.conservative.conservative

from __future__ import annotations

import numpy as np

from .......utils.classes.domains import Domain
from .......utils.geometry.utils import get_cell_area, get_vertices, is_unstructured
from .weights import compute_weights_unstructured




def conservative(
    domain_in: Domain,
    domain_out: Domain,
    chunk_size: int | None = None,
    processes: int | None = None,
) -> dict[str, np.ndarray]:
    """Compute indices and weigths for conservative interpolation between
    unstructured or regular grids.

    Args:
        domain_in (Domain): Input domain plugin.
        domain_out (Domain): Output domain plugin.
        chunk_size (int, optional): Chunk size in the latitude bins. Defaults to None.
        processes (int, optional): Number of processes. Defaults to None.

    Returns:
        dict with keys "i", "j" and "wgt"
    """
    lon_vertices_out, lat_vertices_out = get_vertices(domain_out)
    lon_vertices_in, lat_vertices_in = get_vertices(domain_in)

    cell_area_out = get_cell_area(domain_out)
    cell_area_in = get_cell_area(domain_in)

    indices, weights = compute_weights_unstructured(
        lon_vertices_out,
        lat_vertices_out,
        cell_area_out,
        lon_vertices_in,
        lat_vertices_in,
        cell_area_in,
        chunk_size,
        processes,
    )

    if is_unstructured(domain_in):
        ind_j = indices
        ind_i = np.zeros_like(indices)

    else:
        shape = (domain_in.nlat, domain_in.nlon)  # type: ignore
        # pylint: disable=unbalanced-tuple-unpacking
        ind_i, ind_j = np.unravel_index(indices, shape, order="F")

    output = {"i": ind_i, "j": ind_j, "wgt": weights}

    return output