from __future__ import annotations
import warnings
from functools import wraps
import numpy as np
import pandas as pd
from numpy.typing import NDArray
from scipy.sparse import coo_array
from .......utils.geometry.equal_area_proj import WorldEqualArea, optimize_proj_bins
from .......utils.geometry.polygons import vertices_to_coords
from .......utils.geometry.utils import standardize_lon
from .......utils.check.errclass import CifValueError
try:
imported_njit = True # pylint: disable=invalid-name
from numba import njit
except ImportError:
imported_njit = False # pylint: disable=invalid-name
def njit(func):
return wraps(func)
@njit
def _get_sparse_coords_kernel(
n_out: int,
ind_out: NDArray[np.signedinteger],
) -> NDArray[np.signedinteger]:
n = ind_out.shape[0]
coords = np.zeros((2, n), dtype=np.int64)
counts = np.zeros(n_out, dtype=np.int64)
for i in range(n):
coords[0, i] = ind_out[i]
coords[1, i] = counts[ind_out[i]]
counts[ind_out[i]] += 1
return coords
def _extract_indices_weights_numba(
n_cells: int,
intersection: pd.DataFrame,
) -> tuple[NDArray[np.signedinteger], NDArray[np.floating]]:
n_out = n_cells
n_in = intersection.i_out.value_counts().max()
i_out = intersection.i_out.to_numpy()
i_in = intersection.i_in.to_numpy()
weight = intersection.weight.to_numpy()
coords = _get_sparse_coords_kernel(n_out, i_out)
indices = coo_array((i_in, coords), shape=(n_out, n_in), dtype=np.int64)
weights = coo_array((weight, coords), shape=(n_out, n_in), dtype=np.float64)
return indices.toarray(), weights.toarray() # type: ignore
def _extract_indices_weights_python(
n_cells: int,
intersection: pd.DataFrame,
) -> tuple[NDArray[np.signedinteger], NDArray[np.floating]]:
n_out = n_cells
n_in = intersection.i_out.value_counts().max()
indices = np.full((n_out, n_in), -1, dtype=int)
weights = np.full((n_out, n_in), np.nan, dtype=float)
def fill_data(df: pd.DataFrame) -> None:
i_out = df.i_out.iat[0]
n_in = df.shape[0]
indices[i_out, :n_in] = df.i_in.to_numpy()
weights[i_out, :n_in] = df.weight.to_numpy()
intersection.groupby("i_out")[["i_out", "i_in", "weight"]].apply(
fill_data) # type: ignore
return indices, weights
def _extract_indices_weights(
n_cells: int,
intersection: pd.DataFrame,
use_numba: bool = True,
) -> tuple[NDArray[np.signedinteger], NDArray[np.floating]]:
if use_numba and imported_njit:
return _extract_indices_weights_numba(n_cells, intersection)
else:
warnings.warn(
"Numba is not installed, using slow Python implementation",
UserWarning,
)
return _extract_indices_weights_python(n_cells, intersection)
[docs]
def compute_weights_unstructured(
lon_vertices_out: NDArray[np.floating],
lat_vertices_out: NDArray[np.floating],
cell_area_out: NDArray[np.floating] | None,
lon_vertices_in: NDArray[np.floating],
lat_vertices_in: NDArray[np.floating],
cell_area_in: NDArray[np.floating] | None,
chunk_size: int | None = None,
processes: int | None = None,
use_numba: bool = True,
) -> tuple[NDArray[np.signedinteger], NDArray[np.floating]]:
"""Compute indices and weigths for conservative interpolation between
unstructured grids.
Args:
lon_vertices_out (2D array (M1, N1)): Output grid longitudes of the vertices of each cell.
lat_vertices_out (2D array (M1, N1)): Output grid latitudes of the vertices of each cell.
cell_area_out (1D array (M1) or None): Precomputed areas of each cell of the output grid.
lon_vertices_in (2D array (M2, N2)): Input grid longitudes of the vertices of each cell.
lat_vertices_in (2D array (M2, N2)): Input grid latitudes of the vertices of each cell.
cell_area_in (1D array (M2) or None): Precomputed areas of each cell of the input grid.
chunk_size (int, optional): Chunk size in the latitude bins. Defaults to None.
processes (int, optional): Number of processes. Defaults to None.
use_numba (bool, optional): Use Numba if available, only used for unit tests. Defaults to True.
Returns:
2D array of int (M1, K), 2D array of float (M1, K): Indices and weights
for interpolation.
"""
if lon_vertices_in.shape != lat_vertices_in.shape:
raise CifValueError(
"input longitude and latitude vertices have different shapes")
if lon_vertices_out.shape != lat_vertices_out.shape:
raise CifValueError(
"output longitude and latitude vertices have different shapes")
n_out, _ = lon_vertices_out.shape
# Get projection bins
mid_split, poles_split = optimize_proj_bins(lat_vertices_out, lat_vertices_in)
# Get output projections
lon_vertices_out = standardize_lon(lon_vertices_out)
geom_out = WorldEqualArea(
lon_vertices_out,
lat_vertices_out,
cell_area=cell_area_out,
index_name="i_out",
mid_split=mid_split,
poles_split=poles_split,
chunk_size=chunk_size,
processes=processes
)
# Get input projections
lon_vertices_in = standardize_lon(lon_vertices_in)
coords_in = vertices_to_coords(lon_vertices_in, lat_vertices_in)
unique, unique_indices = np.unique(coords_in, return_index=True, axis=0)
lon_vertices_in = unique[:, :, 0]
lat_vertices_in = unique[:, :, 1]
geom_in = WorldEqualArea(
lon_vertices_in,
lat_vertices_in,
indices=unique_indices,
cell_area=cell_area_in.flatten() if cell_area_in is not None else None,
index_name="i_in",
mid_split=mid_split,
poles_split=poles_split,
chunk_size=chunk_size,
processes=processes
)
# Compute intersection area
intersection = geom_out.intersection_area(geom_in)
if intersection.empty:
raise CifValueError("Input and output domains are not overlapping")
# Intersection area to weights
ind_out = intersection[geom_out.index_name]
intersection["weight"] = (
intersection["area"] / geom_out.area.loc[ind_out].to_numpy()
)
# Format indices and weights
indices, weights = _extract_indices_weights(n_out, intersection, use_numba)
return indices, weights
