"""eomatch.utils.crop - crop gridded matchup products to their shared valid-data extent"""
from __future__ import annotations
import logging
from typing import Iterable, List, Optional, Tuple, Union, cast
import numpy as np
import xarray as xr
from eomatch.deps import lazy_pyproj, lazy_rioxarray
__author__ = "Sam Hunt <sam.hunt@npl.co.uk>"
__all__ = ["crop_to_shared_valid_extent"]
log = logging.getLogger(__name__)
# Used to identify a dataset's measurement variables (as opposed to angles, flags, or
# other ancillary variables) when a dataset has no "meas_vars" attribute to fall back
# on. In practice this is the *only* path that ever runs: no eoio reader (nor
# eomatch.datatree.BuildMUDT.read_products) currently sets "meas_vars" on its output,
# despite eomatch.preview also documenting it as a preferred source. The underlying
# physical quantity varies by product/processing level -- reflectance, radiance,
# brightness temperature -- so this isn't restricted to any one of them. Matched as
# whole underscore-delimited tokens (see _matches_measurement_hint), not raw
# substrings, so e.g. "radiance" doesn't also match an unrelated "solar_irradiance"
# variable.
_MEASUREMENT_NAME_HINTS = ("reflectance", "radiance", "brightness_temperature")
def _matches_measurement_hint(name: str) -> bool:
"""Return ``True`` if *name* contains a :py:data:`_MEASUREMENT_NAME_HINTS` entry as
a whole underscore-delimited token, or contiguous run of tokens for multi-word hints
like ``"brightness_temperature"``.
Deliberately not a raw substring check: ``"radiance" in "solar_irradiance"`` is
``True`` even though the variable has nothing to do with radiance measurements.
:param name: variable name to check (case-insensitive).
:return: whether *name* matches a measurement hint.
"""
tokens = name.lower().split("_")
for hint in _MEASUREMENT_NAME_HINTS:
hint_tokens = hint.split("_")
n = len(hint_tokens)
if any(tokens[i : i + n] == hint_tokens for i in range(len(tokens) - n + 1)):
return True
return False
def _measurement_vars(ds: xr.Dataset) -> List[str]:
"""Identify *ds*'s measurement variables, used to determine which pixels hold real
data for :py:func:`crop_to_shared_valid_extent`.
Prefers the dataset's own ``meas_vars`` attribute (already used by
:py:mod:`eomatch.preview`) when present -- including an explicit empty list, which
is trusted as "no measurement variables" rather than triggering the fallback below;
otherwise falls back to the name-based heuristic in :py:func:`_matches_measurement_hint`.
If that heuristic finds nothing either, a warning is logged -- a dataset with no
identifiable measurement variables is silently excluded from cropping by
:py:func:`crop_to_shared_valid_extent`, so it's worth surfacing when that happens.
:param ds: dataset to inspect.
:return: names of *ds*'s measurement variables.
"""
meas_vars = ds.attrs.get("meas_vars")
if meas_vars is not None:
return [str(v) for v in meas_vars if v in ds.data_vars]
found = [str(v) for v in ds.data_vars if _matches_measurement_hint(str(v))]
if not found:
log.warning(
"No measurement variables identified for %s (checked variable names against %s) "
"-- it will be excluded from crop_to_shared_valid_extent.",
ds.attrs.get("collection"),
_MEASUREMENT_NAME_HINTS,
)
return found
def _spatial_dims(ds: xr.Dataset) -> Optional[Tuple[str, str]]:
"""Return this dataset's ``(y_dim, x_dim)`` gridded spatial dims, or ``None`` if it
has no identifiable spatial coordinate at all (e.g. point/profile/spectral products
with no imagery dimension).
Read directly off *ds*'s own spatial coordinate variables rather than guessing from
dimension *names* (fragile -- dims can be called anything: ``"line"``/``"pixel"``,
``"along_track"``/``"across_track"``, ...):
1. ``latitude``/``longitude`` coordinates' own dims, if present -- unambiguous
regardless of what those dims are actually called. A 1D lat/lon pair only
counts if each varies along a *different* dim (a rectilinear grid, e.g.
``latitude(lat)``/``longitude(lon)``); a track/profile product where both
vary along the *same* dim (e.g. ``latitude(obs)``/``longitude(obs)``, one
lat/lon pair per observation) has no 2D imagery grid at all and correctly
falls through to returning ``None`` below.
2. Otherwise, ``rioxarray``'s own ``y_dim``/``x_dim`` + CRS detection, the same
mechanism :py:mod:`eomatch.preview` already relies on. This is itself
CF-convention-aware (reads ``standard_name``/``axis`` attrs, not just common
coordinate names) -- but only for dims it can actually identify; rioxarray
has no way to guess arbitrary dim names (e.g. ``"row"``/``"col"``) that carry
none of those hints, and raises ``MissingSpatialDimensionError`` rather than
returning ``None`` in that case.
3. If step 2 finds a CRS but can't name the dims, the last two dims of a
measurement variable (see :py:func:`_measurement_vars`) are used instead --
the near-universal ``(..., y, x)`` raster dimension order, and dim-name
agnostic by construction.
:param ds: dataset to inspect.
:return: ``(y_dim, x_dim)`` names, or ``None``.
"""
if "latitude" in ds.coords and "longitude" in ds.coords:
lat_dims, lon_dims = ds["latitude"].dims, ds["longitude"].dims
if len(lat_dims) == 2 and lat_dims == lon_dims:
return str(lat_dims[0]), str(lat_dims[1])
if len(lat_dims) == 1 and len(lon_dims) == 1 and lat_dims[0] != lon_dims[0]:
return str(lat_dims[0]), str(lon_dims[0])
try:
rioxarray = lazy_rioxarray()
if ds.rio.crs is None:
return None
except Exception as exc:
# Broad on purpose: predominantly "rioxarray/pyproj isn't installed"
# (ModuleNotFoundError via lazy_rioxarray -- the "spatial" extra is optional),
# but also covers any other unexpected failure reading .rio.crs. Warned rather
# than logged quietly: unlike a clean "no CRS" result (handled above, not an
# error), reaching this branch means a product that could potentially be
# cropped is being silently skipped -- worth surfacing so "why wasn't this
# product cropped" doesn't require digging through debug logs to answer.
log.warning("_spatial_dims: rioxarray fallback did not apply to %s: %s", ds.attrs.get("collection"), exc)
return None
try:
return str(ds.rio.y_dim), str(ds.rio.x_dim)
except rioxarray.exceptions.MissingSpatialDimensionError:
for v in _measurement_vars(ds):
dims = ds[v].dims
if len(dims) >= 2:
return str(dims[-2]), str(dims[-1])
return None
def _ensure_lat_lon(ds: xr.Dataset, y_dim: str, x_dim: str) -> xr.Dataset:
"""Return *ds* with ``latitude``/``longitude`` coordinates, deriving them from a
projected grid + CRS via an exact ``pyproj`` transform if not already present.
Some products (e.g. certain Landsat reads) only carry a projected grid plus a CRS,
not lat/lon directly. The CRS is read via ``rioxarray``'s ``.rio.crs`` -- the same
mechanism :py:mod:`eomatch.preview` already relies on -- rather than assuming any
one specific attribute/variable naming convention for it.
:param ds: dataset to add coordinates to.
:param y_dim: name of *ds*'s spatial y dim, as returned by :py:func:`_spatial_dims`.
:param x_dim: name of *ds*'s spatial x dim, as returned by :py:func:`_spatial_dims`.
:return: *ds*, with ``latitude``/``longitude`` coordinates added if they were
derivable; unchanged (with a logged warning) otherwise.
"""
if "latitude" in ds.coords and "longitude" in ds.coords:
return ds
try:
lazy_rioxarray()
crs = ds.rio.crs
except Exception as exc:
log.debug("_ensure_lat_lon: could not read a CRS for %s: %s", ds.attrs.get("collection"), exc)
crs = None
if crs is None:
log.warning("Cannot derive latitude/longitude for %s: no CRS found.", ds.attrs.get("collection"))
return ds
if x_dim not in ds.coords or y_dim not in ds.coords:
# A CRS *was* found here -- this is reachable via _spatial_dims's third
# fallback (last two dims of a measurement variable), which can name spatial
# dims that were never coordinate variables to begin with. Without values for
# x_dim/y_dim there's nothing to run the pyproj transform against below.
log.warning(
"Cannot derive latitude/longitude for %s: found a CRS, but %r/%r have no coordinate values to transform.",
ds.attrs.get("collection"),
y_dim,
x_dim,
)
return ds
pyproj = lazy_pyproj()
transformer = pyproj.Transformer.from_crs(crs, "EPSG:4326", always_xy=True)
xx, yy = np.meshgrid(ds[x_dim].values, ds[y_dim].values)
lon, lat = transformer.transform(xx, yy)
return ds.assign_coords(
latitude=((y_dim, x_dim), lat.astype(np.float64)),
longitude=((y_dim, x_dim), lon.astype(np.float64)),
)
def _valid_lat_lon_bounds(ds: xr.Dataset, y_dim: str, x_dim: str) -> Optional[Tuple[float, float, float, float]]:
"""``(lat_min, lat_max, lon_min, lon_max)`` covering *ds*'s non-NaN measurement
pixels, or ``None`` if it has no measurement variables or no lat/lon.
:param ds: dataset to inspect.
:param y_dim: name of *ds*'s spatial y dim.
:param x_dim: name of *ds*'s spatial x dim.
:return: lat/lon bounds, or ``None``.
"""
if "latitude" not in ds.coords or "longitude" not in ds.coords:
return None
valid = None
for v in _measurement_vars(ds):
da = ds[v]
if y_dim not in da.dims or x_dim not in da.dims:
continue
v_valid = da.notnull()
reduce_dims = [d for d in da.dims if d not in (y_dim, x_dim)]
if reduce_dims:
v_valid = v_valid.any(reduce_dims)
valid = v_valid if valid is None else (valid | v_valid)
if valid is None:
return None
# Materialized once and reused as a plain numpy array below, rather than
# re-deriving `valid` from its (possibly lazy/dask-backed) expression graph for
# each of the emptiness check and the four bound reductions separately.
valid_mask = valid.values
if not valid_mask.any():
return None
lat = ds["latitude"].values[valid_mask]
lon = ds["longitude"].values[valid_mask]
return float(lat.min()), float(lat.max()), float(lon.min()), float(lon.max())
def _has_single_spatial_grid(ds: xr.Dataset) -> bool:
"""Return ``False`` if *ds*'s measurement variables span more than one distinct
spatial grid -- e.g. a multi-resolution product with 10m/20m/60m bands coexisting
in the same ``Dataset`` (a real, reachable shape for e.g. a Sentinel-2 read at
native resolution).
:py:func:`crop_to_shared_valid_extent` only supports a single grid per product: it
picks one ``(y_dim, x_dim)`` pair via :py:func:`_spatial_dims` and crops the whole
dataset to it. Without this check, a multi-grid product would have only the
variables on that one chosen grid cropped, while variables on every other grid
were silently left completely uncropped -- worse than not cropping the product at
all, since nothing would indicate only part of it was processed.
:param ds: dataset to inspect.
:return: ``True`` if *ds*'s measurement variables all share one spatial grid (or
there are none/only one to compare).
"""
grids = {ds[v].dims[-2:] for v in _measurement_vars(ds) if len(ds[v].dims) >= 2}
return len(grids) <= 1
def _resolve_datatree_nodes(dt: xr.DataTree) -> List[Tuple[xr.Dataset, str]]:
"""Return every one of *dt*'s populated (non-root) grid nodes, as ``(dataset,
node_path)`` pairs.
A per-product read configured with a ``"to_datatree"`` processor as its last step
(see ``eoio.processors.datatree.processor.ToDataTree``, used by e.g. the RGB-preview
``collection_read_args`` in the example scripts) returns one ``xr.DataTree`` per
product grouped by grid, rather than a flat ``xr.Dataset`` -- the root node holds
only variables with no detectable grid (typically none, for a same-resolution band
selection like an RGB triplet), so the actual measurement data lives on named child
nodes (e.g. ``"/10m"``). A same-resolution band selection has exactly one such node;
a genuinely multi-resolution read (e.g. Sentinel-2 at 10 m *and* 20 m) has more than
one -- both shapes are supported by :py:func:`crop_to_shared_valid_extent`, which
crops every populated node independently, matched against a shared bound computed
as the *union* of this product's own nodes (see its own docstring).
A node only counts as "populated" if it has at least one *measurement* variable
(matched via :py:func:`_matches_measurement_hint`, the same name-hint check
:py:func:`_measurement_vars` uses) -- not just any variable. Per-band viewing/solar
angles requested alongside a same-resolution band selection (e.g. Sentinel-2's
``vars_sel["aux"]``) commonly land on their own, coarser grid node purely because
of their own native resolution, not because the *product* is genuinely
multi-resolution -- an angles-only node must not count towards this.
:param dt: DataTree to inspect.
:return: ``(dataset, node_path)`` for every non-root node with a measurement
variable, in :py:meth:`~xr.DataTree.subtree_with_keys` order; empty if none.
"""
return [
(node.dataset, path)
for path, node in dt.subtree_with_keys
if node is not dt and any(_matches_measurement_hint(str(v)) for v in node.dataset.data_vars)
]
def _extract_datasets_for_crop(item: Union[xr.Dataset, xr.DataTree]) -> List[Tuple[xr.Dataset, Optional[str]]]:
"""Return every ``(dataset, node_path)`` to run this module's crop logic against,
for *item*.
:param item: a per-product read result -- either a plain ``Dataset``, or a
``DataTree`` produced by a per-product ``"to_datatree"`` processor (see
:py:func:`_resolve_datatree_nodes`).
:return: ``[(item, None)]`` for a plain ``Dataset``; one ``(node_dataset,
node_path)`` entry per populated grid node for a ``DataTree`` -- possibly more
than one, for a multi-resolution product; empty for a ``DataTree`` with no
populated node at all.
"""
if isinstance(item, xr.DataTree):
nodes: List[Tuple[xr.Dataset, Optional[str]]] = list(_resolve_datatree_nodes(item))
return nodes
return [(item, None)]
class _TreeWriteBack:
"""Accumulates one or more node writes into a *single* copy of a DataTree product,
across however many of its own grid nodes end up touched (lat/lon-augmented and/or
cropped) in one pass -- a naive "copy fresh from the original product, once per
write" approach would silently discard every write but the last once a product can
have more than one node.
A no-op until :py:meth:`write` is actually called -- :py:meth:`result` then returns
the original product unchanged, the same "don't pay for a copy unless something
really changed" discipline this module already applied per-node before multi-node
products were supported (e.g. when lat/lon is already present, nothing to augment).
"""
def __init__(self, item: Union[xr.Dataset, xr.DataTree]) -> None:
self._item = item
self._tree: Optional[xr.DataTree] = None
def write(self, node_path: Optional[str], ds: xr.Dataset) -> None:
"""Record *ds* as this product's new value at *node_path* (``None`` for a plain
``Dataset`` product, which can only ever have the one write)."""
if node_path is None:
self._item = ds
return
if self._tree is None:
assert isinstance(self._item, xr.DataTree)
self._tree = self._item.copy()
self._tree[node_path] = ds
def result(self) -> Union[xr.Dataset, xr.DataTree]:
"""The product with every :py:meth:`write` applied, or the original unchanged
object if :py:meth:`write` was never called."""
return self._tree if self._tree is not None else self._item
[docs]
def crop_to_shared_valid_extent(
product_ds: Iterable[Union[xr.Dataset, xr.DataTree]],
) -> List[Union[xr.Dataset, xr.DataTree]]:
"""Crop gridded products to the lat/lon bounding box where *all* of them have valid
measurement data.
Products' native grids commonly only partially overlap and/or carry NaN padding well
beyond their actual footprint (e.g. a diagonal swath inside a larger rectangular
tile) -- without this, a lot of every grid is padding that can never be compared,
which dilutes any downstream alignment or analysis. Since products aren't
necessarily on the same pixel grid at this stage, the shared region (the
image-collocated ROI) is computed from each product's own valid-data lat/lon extent,
deriving lat/lon from projected x/y + CRS first where it isn't already present.
Non-gridded products (no ``(y, x)``-like dims), a flat ``Dataset`` product spanning
more than one distinct spatial grid (e.g. an un-split multi-resolution read -- see
:py:func:`_has_single_spatial_grid`), and products with no identifiable measurement
variables or lat/lon, are passed through unchanged and excluded from the
shared-extent calculation. If fewer than two products qualify, all products are
passed through unchanged (there is nothing to collocate against). Deriving lat/lon
for a product that doesn't already have it needs the optional ``spatial`` extra
(``pyproj``/``rioxarray``); if that isn't installed, such a product is likewise
excluded rather than this function raising.
A product read with a per-collection ``"to_datatree"`` processor as the last step
of its own read pipeline (e.g. the RGB-preview ``collection_read_args`` in the
example scripts) arrives here as an ``xr.DataTree`` rather than a flat ``Dataset``
-- see :py:func:`_resolve_datatree_nodes`. Every one of its populated grid nodes is
cropped independently (e.g. a genuinely multi-resolution read, such as Sentinel-2 at
10 m *and* 20 m, has more than one), each to the *same* final shared bbox -- a
product's own contribution to that shared bbox is the union of its nodes' individual
valid extents, so e.g. Sentinel-2's 10 m and 20 m bands (covering, near enough, the
same physical scene) are matched against another product as one combined footprint,
not two independently-negotiated ones. The rest of the tree (e.g. an empty root) is
returned untouched. A ``DataTree`` with zero populated nodes is excluded, the same
as a product with no identifiable measurement variables or lat/lon.
:param product_ds: per-sensor datasets (or, per above, DataTrees), as returned by
:py:meth:`~eomatch.datatree.BuildMUDT.read_products`.
:return: new list, same order and same Dataset-vs-DataTree shape per item, cropped
where applicable.
:raises ValueError: if two or more qualifying products' valid extents don't overlap.
"""
prepared = list(product_ds)
# (index, [(y_dim, x_dim, node_path), ...one entry per this product's own grid
# node...], product_bounds) -- product_bounds already the union across those nodes.
qualifying: List[Tuple[int, List[Tuple[str, str, Optional[str]]], Tuple[float, float, float, float]]] = []
for i, item in enumerate(prepared):
nodes = _extract_datasets_for_crop(item)
if not nodes:
log.warning(
"%s is a DataTree with no populated grid node -- excluding it from the shared-extent crop.",
item.attrs.get("collection"),
)
continue
writer = _TreeWriteBack(item)
node_entries: List[Tuple[str, str, Optional[str]]] = []
node_bounds: List[Tuple[float, float, float, float]] = []
for ds, node_path in nodes:
if not _has_single_spatial_grid(ds):
log.warning(
"%s%s spans more than one spatial grid -- crop_to_shared_valid_extent "
"does not support that within a single grid node, excluding it from "
"the shared-extent crop.",
ds.attrs.get("collection"),
f" node {node_path!r}" if node_path else "",
)
continue
dims = _spatial_dims(ds)
if dims is None:
continue
y_dim, x_dim = dims
augmented = _ensure_lat_lon(ds, y_dim, x_dim)
if augmented is not ds:
# _ensure_lat_lon returns ds unchanged (same object) whenever it added
# nothing (lat/lon already present, or no CRS to derive them from) --
# skip the write in that case; _TreeWriteBack itself stays a no-op
# (no tree copy) unless at least one node in this product actually needs
# writing, here or in the crop step below.
writer.write(node_path, augmented)
bounds = _valid_lat_lon_bounds(augmented, y_dim, x_dim)
if bounds is not None:
node_entries.append((y_dim, x_dim, node_path))
node_bounds.append(bounds)
prepared[i] = writer.result()
if not node_bounds:
continue
product_bounds = (
min(b[0] for b in node_bounds),
max(b[1] for b in node_bounds),
min(b[2] for b in node_bounds),
max(b[3] for b in node_bounds),
)
qualifying.append((i, node_entries, product_bounds))
if len(qualifying) < 2:
log.info("Fewer than two gridded products with measurement data and lat/lon -- skipping shared-extent crop.")
return prepared
lat_min = max(q[2][0] for q in qualifying)
lat_max = min(q[2][1] for q in qualifying)
lon_min = max(q[2][2] for q in qualifying)
lon_max = min(q[2][3] for q in qualifying)
# Strict > only: a zero-width intersection (lat_min == lat_max) is a legitimate
# single-row/column shared extent, not "no overlap" -- the pixel-selection loop
# below already handles it correctly, and raises its own ValueError if a bbox that
# looked non-empty here somehow selects zero actual pixels.
if lat_min > lat_max or lon_min > lon_max:
raise ValueError("Products' valid-data extents do not overlap -- cannot crop to a shared extent.")
for i, node_entries, _ in qualifying:
item = prepared[i]
writer = _TreeWriteBack(item)
for y_dim, x_dim, node_path in node_entries:
ds = cast(xr.Dataset, item[node_path].dataset if node_path is not None else item)
in_bbox = (
(ds["latitude"] >= lat_min)
& (ds["latitude"] <= lat_max)
& (ds["longitude"] >= lon_min)
& (ds["longitude"] <= lon_max)
)
# Materialized once and reused for both axes below, rather than letting
# .any(x_dim)/.any(y_dim) each independently re-derive in_bbox from its
# (possibly lazy/dask-backed) expression graph.
in_bbox_mask = in_bbox.values
y_axis = in_bbox.dims.index(y_dim)
x_axis = in_bbox.dims.index(x_dim)
row_idx = np.flatnonzero(in_bbox_mask.any(axis=x_axis))
col_idx = np.flatnonzero(in_bbox_mask.any(axis=y_axis))
if row_idx.size == 0:
raise ValueError("No pixels found inside the shared lat/lon bounding box.")
sel = {y_dim: slice(row_idx[0], row_idx[-1] + 1), x_dim: slice(col_idx[0], col_idx[-1] + 1)}
writer.write(node_path, ds.isel(sel))
prepared[i] = writer.result()
return prepared
if __name__ == "__main__":
pass