from __future__ import annotations
from logging import debug, warning
import numpy as np
from numpy.typing import NDArray
from .......utils.check.errclass import CifValueError
from .......utils.classes.domains import Domain
from .conservative import conservative
def _unwrap_corners(zlonc: NDArray[np.floating]) -> NDArray[np.floating]:
# Bring all corners of a footprint in the same revolution as its first
# corner to deal with footprints crossing the dateline
return zlonc[0] + (zlonc - zlonc[0] + 180) % 360 - 180
def _order_corners(
zlonc: NDArray[np.floating],
zlatc: NDArray[np.floating],
) -> tuple[NDArray[np.floating], NDArray[np.floating]]:
# Sort corners by angle around the footprint center to get valid
# (non self-intersecting) polygons whatever the order in the product;
# this is exact for convex footprints
xlon = _unwrap_corners(zlonc)
angles = np.arctan2(zlatc - zlatc.mean(axis=0), xlon - xlon.mean(axis=0))
order = np.argsort(angles, axis=0)
return (
np.take_along_axis(zlonc, order, axis=0),
np.take_along_axis(zlatc, order, axis=0),
)
def _footprint_area(
zlonc: NDArray[np.floating],
zlatc: NDArray[np.floating],
) -> NDArray[np.floating]:
# Shoelace formula in degrees; only used to detect degenerate footprints
xlon = _unwrap_corners(zlonc)
return 0.5 * np.abs(np.sum(
xlon * np.roll(zlatc, -1, axis=0) - np.roll(xlon, -1, axis=0) * zlatc,
axis=0,
))
def _empty_weights(nobs: int) -> dict[str, NDArray]:
return {
"i": np.zeros((0, 1), dtype=int),
"j": np.zeros((0, 1), dtype=int),
"wgt": np.zeros((0, 1)),
"filtered": np.zeros(0, dtype=int),
"non_filtered": np.arange(nobs),
}
[docs]
def conservative_sparse(
domain_in: Domain,
footprints: Domain,
chunk_size: int | None = None,
processes: int | None = None,
min_coverage: float = 0.5,
) -> dict[str, NDArray]:
"""Compute conservative weights from a grid to observation footprints.
Each observation is the area-weighted average of the grid cells
overlapping its footprint. Observations with invalid footprints or whose
footprint is covered by the grid by less than ``min_coverage`` are
considered outside the domain.
Args:
domain_in (Domain): Gridded domain.
footprints (Domain): Unstructured domain with the corners of the
observation footprints.
chunk_size (int, optional): Chunk size in the latitude bins.
processes (int, optional): Number of processes.
min_coverage (float): Minimum fraction of a footprint to be covered
by the grid.
Returns:
dict with keys "i", "j", "wgt", "filtered" and "non_filtered"
"""
if not getattr(footprints, "unstructured_domain", False):
raise CifValueError(
"The regrid method 'fast-conservative' needs the corners of the "
"observation footprints with sparse data. Please activate "
"'use_corners' and specify 'ncorners' for your observations."
)
zlonc = np.asarray(footprints.zlonc, dtype=float)
zlatc = np.asarray(footprints.zlatc, dtype=float)
nobs = zlonc.shape[1]
# Discard footprints with missing corners or no area
valid = np.all(np.isfinite(zlonc) & np.isfinite(zlatc), axis=0)
zlonc, zlatc = _order_corners(zlonc[:, valid], zlatc[:, valid])
is_degenerate = _footprint_area(zlonc, zlatc) <= 0
valid[np.flatnonzero(valid)[is_degenerate]] = False
zlonc, zlatc = zlonc[:, ~is_degenerate], zlatc[:, ~is_degenerate]
if not np.all(valid):
warning(
f"{np.sum(~valid)} observations out of {nobs} have missing or "
"degenerate footprint corners. They are considered outside the "
"domain."
)
if not np.any(valid):
return _empty_weights(nobs)
# Footprints outside the latitude range of the grid are outside the
# domain; this also avoids computing intersections with no common
# latitude band, which are not supported by the conservative routine
lat_grid = np.asarray(domain_in.zlatc, dtype=float)
in_lat_range = (zlatc.max(axis=0) >= np.nanmin(lat_grid)) \
& (zlatc.min(axis=0) <= np.nanmax(lat_grid))
valid[np.flatnonzero(valid)[~in_lat_range]] = False
zlonc, zlatc = zlonc[:, in_lat_range], zlatc[:, in_lat_range]
if not np.any(valid):
debug("No observation footprint in the latitude range of the grid")
return _empty_weights(nobs)
domain_valid = Domain()
domain_valid.unstructured_domain = True
domain_valid.zlonc = zlonc
domain_valid.zlatc = zlatc
domain_valid.nlon = domain_valid.nlat = zlonc.shape[1]
try:
weights = conservative(domain_in, domain_valid, chunk_size, processes)
except CifValueError as e:
# No footprint overlaps the grid, e.g., all observations of the
# period are outside a regional domain
if "not overlapping" not in str(e):
raise
debug("No observation footprint overlaps the grid")
return _empty_weights(nobs)
# Padding may be filled with NaNs and negative indices
i, j = weights["i"], weights["j"]
wgt = np.nan_to_num(weights["wgt"], nan=0.0)
# Keep observations sufficiently covered by the grid and renormalize
# their weights so that partially covered footprints are not biased
coverage = wgt.sum(axis=1)
inside = coverage >= min_coverage
i, j, wgt = i[inside], j[inside], wgt[inside] / coverage[inside, np.newaxis]
# Remove columns with padding only and make padding point to a valid
# cell with a zero weight
columns = np.any(wgt > 0, axis=0)
i, j, wgt = i[:, columns], j[:, columns], wgt[:, columns]
i = np.where(wgt > 0, i, 0)
j = np.where(wgt > 0, j, 0)
filtered = np.flatnonzero(valid)[inside]
non_filtered = np.setdiff1d(np.arange(nobs), filtered)
if filtered.size == 0:
return _empty_weights(nobs)
return {
"i": i,
"j": j,
"wgt": wgt,
"filtered": filtered,
"non_filtered": non_filtered,
}
