Fixes #134: add xr.Dataset as input type for appropriate modules

Add transparent Dataset support via two decorators in a new
dataset_support module. Single-input functions (slope, aspect,
curvature, hillshade, focal.mean, 10 classify functions, 3 proximity
functions) iterate over data_vars and return a Dataset. Multi-input
functions (10 multispectral indices) accept a Dataset with band-name
kwargs. zonal.stats merges per-variable DataFrames with prefixed
columns. Includes 18 new tests.

Author: brendancol

xrspatial/aspect.py

@@ -19,6 +19,7 @@
 from xrspatial.utils import _extract_latlon_coords
 from xrspatial.utils import cuda_args
 from xrspatial.utils import ngjit
+from xrspatial.dataset_support import supports_dataset
 
 
 def _geodesic_cuda_dims(shape):
@@ -270,6 +271,7 @@ def _run_dask_cupy_geodesic(data, lat_2d, lon_2d, a2, b2, z_factor):
 # Public API
 # =====================================================================
 
+@supports_dataset
 def aspect(agg: xr.DataArray,
            name: Optional[str] = 'aspect',
            method: str = 'planar',

xrspatial/classify.py

@@ -25,6 +25,7 @@ class cupy(object):
 import numpy as np
 
 from xrspatial.utils import ArrayTypeFunctionMapping, cuda_args, ngjit, not_implemented_func
+from xrspatial.dataset_support import supports_dataset
 
 
 @ngjit
@@ -83,6 +84,7 @@ def _run_dask_cupy_binary(data, values_cupy):
     return out
 
 
+@supports_dataset
 def binary(agg, values, name='binary'):
     """
     Binarize a data array based on a set of values. Data that equals to a value in the set will be
@@ -266,6 +268,7 @@ def _bin(agg, bins, new_values):
     return out
 
 
+@supports_dataset
 def reclassify(agg: xr.DataArray,
                bins: List[int],
                new_values: List[int],
@@ -416,6 +419,7 @@ def _quantile(agg, k):
     return out
 
 
+@supports_dataset
 def quantile(agg: xr.DataArray,
              k: int = 4,
              name: Optional[str] = 'quantile') -> xr.DataArray:
@@ -723,6 +727,7 @@ def _run_dask_cupy_natural_break(agg, num_sample, k):
     return out
 
 
+@supports_dataset
 def natural_breaks(agg: xr.DataArray,
                    num_sample: Optional[int] = 20000,
                    name: Optional[str] = 'natural_breaks',
@@ -854,6 +859,7 @@ def _run_equal_interval(agg, k, module):
     return out
 
 
+@supports_dataset
 def equal_interval(agg: xr.DataArray,
                    k: int = 5,
                    name: Optional[str] = 'equal_interval') -> xr.DataArray:
@@ -952,6 +958,7 @@ def _run_std_mean(agg, module):
     return out
 
 
+@supports_dataset
 def std_mean(agg: xr.DataArray,
              name: Optional[str] = 'std_mean') -> xr.DataArray:
     """
@@ -1044,6 +1051,7 @@ def _run_dask_head_tail_breaks(agg):
     return out
 
 
+@supports_dataset
 def head_tail_breaks(agg: xr.DataArray,
                      name: Optional[str] = 'head_tail_breaks') -> xr.DataArray:
     """
@@ -1096,6 +1104,7 @@ def _run_dask_cupy_percentiles(data, pct):
     return _run_percentiles(data_cpu, pct, da)
 
 
+@supports_dataset
 def percentiles(agg: xr.DataArray,
                 pct: Optional[List] = None,
                 name: Optional[str] = 'percentiles') -> xr.DataArray:
@@ -1212,6 +1221,7 @@ def _run_dask_cupy_maximum_breaks(agg, k):
     return out
 
 
+@supports_dataset
 def maximum_breaks(agg: xr.DataArray,
                    k: int = 5,
                    name: Optional[str] = 'maximum_breaks') -> xr.DataArray:
@@ -1312,6 +1322,7 @@ def _run_dask_cupy_box_plot(agg, hinge):
     return out
 
 
+@supports_dataset
 def box_plot(agg: xr.DataArray,
              hinge: float = 1.5,
              name: Optional[str] = 'box_plot') -> xr.DataArray:

xrspatial/curvature.py

@@ -25,6 +25,7 @@ class cupy(object):
 from xrspatial.utils import cuda_args
 from xrspatial.utils import get_dataarray_resolution
 from xrspatial.utils import ngjit
+from xrspatial.dataset_support import supports_dataset
 
 
 @ngjit
@@ -107,6 +108,7 @@ def _run_dask_cupy(data: da.Array,
     return out
 
 
+@supports_dataset
 def curvature(agg: xr.DataArray,
               name: Optional[str] = 'curvature') -> xr.DataArray:
     """

xrspatial/dataset_support.py

@@ -0,0 +1,80 @@
+"""Decorators for transparent xr.Dataset support on xr.DataArray functions."""
+
+from __future__ import annotations
+
+import functools
+import inspect
+
+import xarray as xr
+
+
+def supports_dataset(func):
+    """Decorator that lets single-input DataArray functions accept a Dataset.
+
+    When a Dataset is passed as the first argument, the wrapped function
+    is called on each data variable and the results are collected into
+    a new Dataset.
+    """
+    sig = inspect.signature(func)
+    has_name_param = 'name' in sig.parameters
+
+    @functools.wraps(func)
+    def wrapper(agg, *args, **kwargs):
+        if isinstance(agg, xr.Dataset):
+            results = {}
+            for var_name in agg.data_vars:
+                kw = dict(kwargs)
+                if has_name_param:
+                    kw['name'] = var_name
+                results[var_name] = func(agg[var_name], *args, **kw)
+            return xr.Dataset(results, attrs=agg.attrs)
+        return func(agg, *args, **kwargs)
+
+    return wrapper
+
+
+def supports_dataset_bands(**band_param_map):
+    """Decorator for multi-input functions that take separate band DataArrays.
+
+    Enables passing a single Dataset with keyword arguments that map
+    band aliases to Dataset variable names.
+
+    Example::
+
+        @supports_dataset_bands(nir='nir_agg', red='red_agg')
+        def ndvi(nir_agg, red_agg, name='ndvi'): ...
+
+        # Enables:
+        ndvi(ds, nir='band_8', red='band_4')
+    """
+
+    def decorator(func):
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            if args and isinstance(args[0], xr.Dataset):
+                ds = args[0]
+                func_kwargs = {}
+                used = set()
+                for alias, param in band_param_map.items():
+                    if alias not in kwargs:
+                        raise TypeError(
+                            f"'{alias}' keyword required when passing a Dataset"
+                        )
+                    var_name = kwargs[alias]
+                    if var_name not in ds.data_vars:
+                        raise ValueError(
+                            f"'{var_name}' not in Dataset. "
+                            f"Available: {list(ds.data_vars)}"
+                        )
+                    func_kwargs[param] = ds[var_name]
+                    used.add(alias)
+                # Pass through remaining kwargs (name, soil_factor, etc.)
+                for k, v in kwargs.items():
+                    if k not in used:
+                        func_kwargs[k] = v
+                return func(**func_kwargs)
+            return func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator

xrspatial/focal.py

@@ -29,6 +29,7 @@ class cupy(object):
 
 from xrspatial.convolution import convolve_2d, custom_kernel
 from xrspatial.utils import ArrayTypeFunctionMapping, cuda_args, ngjit, not_implemented_func
+from xrspatial.dataset_support import supports_dataset
 
 # TODO: Make convolution more generic with numba first-class functions.
 
@@ -158,6 +159,7 @@ def _mean(data, excludes):
     return out
 
 
+@supports_dataset
 def mean(agg, passes=1, excludes=[np.nan], name='mean'):
     """
     Returns Mean filtered array using a 3x3 window.

xrspatial/hillshade.py

@@ -14,6 +14,7 @@
 
 from .gpu_rtx import has_rtx
 from .utils import calc_cuda_dims, has_cuda_and_cupy, is_cupy_array, is_cupy_backed
+from .dataset_support import supports_dataset
 
 
 def _run_numpy(data, azimuth=225, angle_altitude=25):
@@ -99,6 +100,7 @@ def _run_cupy(d_data, azimuth, angle_altitude):
     return output
 
 
+@supports_dataset
 def hillshade(agg: xr.DataArray,
               azimuth: int = 225,
               angle_altitude: int = 25,

xrspatial/multispectral.py

@@ -11,6 +11,7 @@
 
 from xrspatial.utils import (ArrayTypeFunctionMapping, cuda_args, ngjit, not_implemented_func,
                              validate_arrays)
+from xrspatial.dataset_support import supports_dataset_bands
 
 # 3rd-party
 try:
@@ -75,6 +76,7 @@ def _arvi_dask_cupy(nir_data, red_data, blue_data):
     return out
 
 
+@supports_dataset_bands(nir='nir_agg', red='red_agg', blue='blue_agg')
 def arvi(nir_agg: xr.DataArray,
          red_agg: xr.DataArray,
          blue_agg: xr.DataArray,
@@ -215,6 +217,7 @@ def _evi_dask_cupy(nir_data, red_data, blue_data, c1, c2, soil_factor, gain):
     return out
 
 
+@supports_dataset_bands(nir='nir_agg', red='red_agg', blue='blue_agg')
 def evi(nir_agg: xr.DataArray,
         red_agg: xr.DataArray,
         blue_agg: xr.DataArray,
@@ -374,6 +377,7 @@ def _gci_dask_cupy(nir_data, green_data):
     return out
 
 
+@supports_dataset_bands(nir='nir_agg', green='green_agg')
 def gci(nir_agg: xr.DataArray,
         green_agg: xr.DataArray,
         name='gci'):
@@ -451,6 +455,7 @@ def gci(nir_agg: xr.DataArray,
 
 
 # NBR ----------
+@supports_dataset_bands(nir='nir_agg', swir2='swir2_agg')
 def nbr(nir_agg: xr.DataArray,
         swir2_agg: xr.DataArray,
         name='nbr'):
@@ -529,6 +534,7 @@ def nbr(nir_agg: xr.DataArray,
                      attrs=nir_agg.attrs)
 
 
+@supports_dataset_bands(swir1='swir1_agg', swir2='swir2_agg')
 def nbr2(swir1_agg: xr.DataArray,
          swir2_agg: xr.DataArray,
          name='nbr2'):
@@ -614,6 +620,7 @@ def nbr2(swir1_agg: xr.DataArray,
 
 
 # NDVI ----------
+@supports_dataset_bands(nir='nir_agg', red='red_agg')
 def ndvi(nir_agg: xr.DataArray,
          red_agg: xr.DataArray,
          name='ndvi'):
@@ -691,6 +698,7 @@ def ndvi(nir_agg: xr.DataArray,
 
 
 # NDMI ----------
+@supports_dataset_bands(nir='nir_agg', swir1='swir1_agg')
 def ndmi(nir_agg: xr.DataArray,
          swir1_agg: xr.DataArray,
          name='ndmi'):
@@ -874,6 +882,7 @@ def _savi_dask_cupy(nir_data, red_data, soil_factor):
 
 
 # SAVI ----------
+@supports_dataset_bands(nir='nir_agg', red='red_agg')
 def savi(nir_agg: xr.DataArray,
          red_agg: xr.DataArray,
          soil_factor: float = 1.0,
@@ -1006,6 +1015,7 @@ def _sipi_dask_cupy(nir_data, red_data, blue_data):
     return out
 
 
+@supports_dataset_bands(nir='nir_agg', red='red_agg', blue='blue_agg')
 def sipi(nir_agg: xr.DataArray,
          red_agg: xr.DataArray,
          blue_agg: xr.DataArray,
@@ -1142,6 +1152,7 @@ def _ebbi_dask_cupy(red_data, swir_data, tir_data):
     return out
 
 
+@supports_dataset_bands(red='red_agg', swir='swir_agg', tir='tir_agg')
 def ebbi(red_agg: xr.DataArray,
          swir_agg: xr.DataArray,
          tir_agg: xr.DataArray,

xrspatial/proximity.py

@@ -10,6 +10,7 @@
 from numba import prange
 
 from xrspatial.utils import get_dataarray_resolution, ngjit
+from xrspatial.dataset_support import supports_dataset
 
 EUCLIDEAN = 0
 GREAT_CIRCLE = 1
@@ -648,6 +649,7 @@ def _process_dask(raster, xs, ys):
 
 # ported from
 # https://github.com/OSGeo/gdal/blob/master/gdal/alg/gdalproximity.cpp
+@supports_dataset
 def proximity(
     raster: xr.DataArray,
     x: str = "x",
@@ -783,6 +785,7 @@ def proximity(
     return result
 
 
+@supports_dataset
 def allocation(
     raster: xr.DataArray,
     x: str = "x",
@@ -915,6 +918,7 @@ def allocation(
     return result
 
 
+@supports_dataset
 def direction(
     raster: xr.DataArray,
     x: str = "x",

xrspatial/slope.py

@@ -28,6 +28,7 @@ class cupy(object):
 from xrspatial.utils import cuda_args
 from xrspatial.utils import get_dataarray_resolution
 from xrspatial.utils import ngjit
+from xrspatial.dataset_support import supports_dataset
 
 
 def _geodesic_cuda_dims(shape):
@@ -267,6 +268,7 @@ def _run_dask_cupy_geodesic(data, lat_2d, lon_2d, a2, b2, z_factor):
 # Public API
 # =====================================================================
 
+@supports_dataset
 def slope(agg: xr.DataArray,
           name: str = 'slope',
           method: str = 'planar',