Xarray: add indexes options and better define band names

CHANGES.md

@@ -1,3 +1,73 @@
+# 7.3.0 (TBD)
+
+* add `indexes` parameter for `XarrayReader` methods. As for Rasterio, the indexes values start at `1`.
+
+    ```python
+    data = ... # DataArray of shape (2, x, y)
+
+    # before
+    with XarrayReader(data) as dst:
+        img = dst.tile(0, 0, 0)
+        assert img.count == 2
+
+    # now
+    with XarrayReader(data) as dst:
+        # Select the first `band` within the data array
+        img = dst.tile(0, 0, 0, indexes=1)
+        assert img.count == 1
+    ```
+
+* better define `band names` for `XarrayReader` objects
+
+    * band_name for `2D` dataset is extracted form the first `non-geo` coordinates value
+
+        ```python
+        data = xarray.DataArray(
+            numpy.arange(0.0, 33 * 35 * 2).reshape(2, 33, 35),
+            dims=("time", "y", "x"),
+            coords={
+                "x": numpy.arange(-170, 180, 10),
+                "y": numpy.arange(-80, 85, 5),
+                "time": [datetime(2022, 1, 1), datetime(2022, 1, 2)],
+            },
+        )
+        da = data[0]
+
+        print(da.coords["time"].data)
+        >> array('2022-01-01T00:00:00.000000000', dtype='datetime64[ns]'))
+
+        # before
+        with XarrayReader(data) as dst:
+            img = dst.info()
+            print(img.band_descriptions)[0]
+            >> ("b1", "value")
+
+        # now
+        with XarrayReader(data) as dst:
+            img = dst.info()
+            print(img.band_descriptions)[0]
+            >> ("b1", "2022-01-01T00:00:00.000000000")
+        ```
+
+    * default `band_names` is changed to DataArray's name or `array` (when no available coordinates value)
+
+        ```python
+        data = ... # DataArray of shape (x, y)
+
+        # before
+        with XarrayReader(data) as dst:
+            img = dst.info()
+            print(img.band_descriptions)[0]
+            >> ("b1", "value")
+
+        # now
+        with XarrayReader(data) as dst:
+            img = dst.info()
+            print(img.band_descriptions)[0]
+            >> ("b1", "array")
+        ```
+
+
 # 7.2.0 (2024-11-05)
 
 * Ensure compatibility between XarrayReader and other Readers by adding `**kwargs` on class methods (https://github.com/cogeotiff/rio-tiler/pull/762)

rio_tiler/io/xarray.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 import warnings
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Tuple
 
 import attr
 import numpy
@@ -28,8 +28,13 @@
 from rio_tiler.io.base import BaseReader
 from rio_tiler.models import BandStatistics, ImageData, Info, PointData
 from rio_tiler.reader import _get_width_height
-from rio_tiler.types import BBox, NoData, RIOResampling, WarpResampling
-from rio_tiler.utils import CRS_to_uri, _validate_shape_input, get_array_statistics
+from rio_tiler.types import BBox, Indexes, NoData, RIOResampling, WarpResampling
+from rio_tiler.utils import (
+    CRS_to_uri,
+    _validate_shape_input,
+    cast_to_sequence,
+    get_array_statistics,
+)
 
 try:
     import xarray
@@ -105,6 +110,7 @@ def __attrs_post_init__(self):
             for d in self.input.dims
             if d not in [self.input.rio.x_dim, self.input.rio.y_dim]
         ]
+        assert len(self._dims) in [0, 1], "Can't handle >=4D DataArray"
 
     @property
     def minzoom(self):
@@ -118,29 +124,34 @@ def maxzoom(self):
 
     @property
     def band_names(self) -> List[str]:
-        """Return list of `band names` in DataArray."""
-        return [str(band) for d in self._dims for band in self.input[d].values] or [
-            "value"
-        ]
+        """Return list of `band descriptions` in DataArray."""
+        if not self._dims:
+            coords_name = list(self.input.coords)
+            if len(coords_name) > 3 and (coord := coords_name[2]):
+                return [str(self.input.coords[coord].data)]
+
+            return [self.input.name or "array"]
+
+        return [str(band) for d in self._dims for band in self.input[d].values]
 
     def info(self) -> Info:
         """Return xarray.DataArray info."""
-        bands = [str(band) for d in self._dims for band in self.input[d].values] or [
-            "value"
-        ]
         metadata = [band.attrs for d in self._dims for band in self.input[d]] or [{}]
 
         meta = {
             "bounds": self.bounds,
             "crs": CRS_to_uri(self.crs) or self.crs.to_wkt(),
             "band_metadata": [(f"b{ix}", v) for ix, v in enumerate(metadata, 1)],
-            "band_descriptions": [(f"b{ix}", v) for ix, v in enumerate(bands, 1)],
+            "band_descriptions": [
+                (f"b{ix}", v) for ix, v in enumerate(self.band_names, 1)
+            ],
             "dtype": str(self.input.dtype),
             "nodata_type": "Nodata" if self.input.rio.nodata is not None else "None",
             "name": self.input.name,
             "count": self.input.rio.count,
             "width": self.input.rio.width,
             "height": self.input.rio.height,
+            "dimensions": self.input.dims,
             "attrs": {
                 k: (v.tolist() if isinstance(v, (numpy.ndarray, numpy.generic)) else v)
                 for k, v in self.input.attrs.items()
@@ -149,19 +160,43 @@ def info(self) -> Info:
 
         return Info(**meta)
 
+    def _sel_indexes(
+        self, indexes: Optional[Indexes] = None
+    ) -> Tuple[xarray.DataArray, List[str]]:
+        """Select `band` indexes in DataArray."""
+        ds = self.input
+        band_names = self.band_names
+        if indexes := cast_to_sequence(indexes):
+            assert all(v > 0 for v in indexes), "Indexes value must be >= 1"
+            if ds.ndim == 2:
+                if indexes != (1,):
+                    raise ValueError(
+                        f"Invalid indexes {indexes} for array of shape {ds.shape}"
+                    )
+
+                return ds, band_names
+
+            indexes = [idx - 1 for idx in indexes]
+            ds = ds[indexes]
+            band_names = [self.band_names[idx] for idx in indexes]
+
+        return ds, band_names
+
     def statistics(
         self,
         categorical: bool = False,
         categories: Optional[List[float]] = None,
         percentiles: Optional[List[int]] = None,
         hist_options: Optional[Dict] = None,
         nodata: Optional[NoData] = None,
+        indexes: Optional[Indexes] = None,
         **kwargs: Any,
     ) -> Dict[str, BandStatistics]:
         """Return statistics from a dataset."""
         hist_options = hist_options or {}
 
-        ds = self.input
+        ds, band_names = self._sel_indexes(indexes)
+
         if nodata is not None:
             ds = ds.rio.write_nodata(nodata)
 
@@ -176,9 +211,7 @@ def statistics(
             **hist_options,
         )
 
-        return {
-            self.band_names[ix]: BandStatistics(**val) for ix, val in enumerate(stats)
-        }
+        return {band_names[ix]: BandStatistics(**val) for ix, val in enumerate(stats)}
 
     def tile(
         self,
@@ -189,6 +222,7 @@ def tile(
         reproject_method: WarpResampling = "nearest",
         auto_expand: bool = True,
         nodata: Optional[NoData] = None,
+        indexes: Optional[Indexes] = None,
         **kwargs: Any,
     ) -> ImageData:
         """Read a Web Map tile from a dataset.
@@ -211,7 +245,8 @@ def tile(
                 f"Tile(x={tile_x}, y={tile_y}, z={tile_z}) is outside bounds"
             )
 
-        ds = self.input
+        ds, band_names = self._sel_indexes(indexes)
+
         if nodata is not None:
             ds = ds.rio.write_nodata(nodata)
 
@@ -251,7 +286,7 @@ def tile(
             bounds=tile_bounds,
             crs=dst_crs,
             dataset_statistics=stats,
-            band_names=self.band_names,
+            band_names=band_names,
         )
 
     def part(
@@ -262,6 +297,7 @@ def part(
         reproject_method: WarpResampling = "nearest",
         auto_expand: bool = True,
         nodata: Optional[NoData] = None,
+        indexes: Optional[Indexes] = None,
         max_size: Optional[int] = None,
         height: Optional[int] = None,
         width: Optional[int] = None,
@@ -294,7 +330,8 @@ def part(
 
         dst_crs = dst_crs or bounds_crs
 
-        ds = self.input
+        ds, band_names = self._sel_indexes(indexes)
+
         if nodata is not None:
             ds = ds.rio.write_nodata(nodata)
 
@@ -339,7 +376,7 @@ def part(
             bounds=ds.rio.bounds(),
             crs=ds.rio.crs,
             dataset_statistics=stats,
-            band_names=self.band_names,
+            band_names=band_names,
         )
 
         output_height = height or img.height
@@ -362,6 +399,7 @@ def preview(
         height: Optional[int] = None,
         width: Optional[int] = None,
         nodata: Optional[NoData] = None,
+        indexes: Optional[Indexes] = None,
         dst_crs: Optional[CRS] = None,
         reproject_method: WarpResampling = "nearest",
         resampling_method: RIOResampling = "nearest",
@@ -388,7 +426,8 @@ def preview(
                 UserWarning,
             )
 
-        ds = self.input
+        ds, band_names = self._sel_indexes(indexes)
+
         if nodata is not None:
             ds = ds.rio.write_nodata(nodata)
 
@@ -427,7 +466,7 @@ def preview(
             bounds=ds.rio.bounds(),
             crs=ds.rio.crs,
             dataset_statistics=stats,
-            band_names=self.band_names,
+            band_names=band_names,
         )
 
         output_height = height or img.height
@@ -450,6 +489,7 @@ def point(
         lat: float,
         coord_crs: CRS = WGS84_CRS,
         nodata: Optional[NoData] = None,
+        indexes: Optional[Indexes] = None,
         **kwargs: Any,
     ) -> PointData:
         """Read a pixel value from a dataset.
@@ -472,7 +512,8 @@ def point(
         ):
             raise PointOutsideBounds("Point is outside dataset bounds")
 
-        ds = self.input
+        ds, band_names = self._sel_indexes(indexes)
+
         if nodata is not None:
             ds = ds.rio.write_nodata(nodata)
 
@@ -489,7 +530,7 @@ def point(
             arr,
             coordinates=(lon, lat),
             crs=coord_crs,
-            band_names=self.band_names,
+            band_names=band_names,
         )
 
     def feature(
@@ -500,6 +541,7 @@ def feature(
         reproject_method: WarpResampling = "nearest",
         auto_expand: bool = True,
         nodata: Optional[NoData] = None,
+        indexes: Optional[Indexes] = None,
         max_size: Optional[int] = None,
         height: Optional[int] = None,
         width: Optional[int] = None,
@@ -537,6 +579,7 @@ def feature(
             dst_crs=dst_crs,
             bounds_crs=shape_crs,
             nodata=nodata,
+            indexes=indexes,
             max_size=max_size,
             width=width,
             height=height,

tests/test_io_stac.py

@@ -1060,7 +1060,7 @@ def _get_reader(self, asset_info: AssetInfo) -> Tuple[Type[BaseReader], Dict]:
         assert info["netcdf"].crs
 
         img = stac.preview(assets=["netcdf"])
-        assert img.band_names == ["netcdf_value"]
+        assert img.band_names == ["netcdf_dataset"]
 
 
 @patch("rio_tiler.io.stac.STAC_ALTERNATE_KEY", "s3")