Return stackstac.stack() as xr.Dataset

[FlorisCalkoen]

Hi, is there a specific reason for stackstac.stack() to return the stack as xr.DataArray and not as xr.Dataset? In the type of analyses that I do it is often easier to hold the data in a xr.Dataset (see below). Would it be possible to add an optional argument to stackstac.stack() to return the stack as xr.Dataset?

import numpy as np
import planetary_computer
import pystac_client
import rasterio
import stackstac
import xarray as xr

catalog = pystac_client.Client.open(
    "https://planetarycomputer.microsoft.com/api/stac/v1",
    modifier=planetary_computer.sign_inplace,
)

roi = {
    "type": "Polygon",
    "coordinates": [
        [
            [151.2945334307, -33.7448472377],
            [151.3229999588, -33.7448472377],
            [151.3229999588, -33.6917695125],
            [151.2945334307, -33.6917695125],
            [151.2945334307, -33.7448472377],
        ]
    ],
}

bbox = [151.2945334307, -33.7448472377, 151.3229999588, -33.6917695125]

search = catalog.search(
    collections=["sentinel-2-l2a"],
    intersects=roi,
    datetime="2022-07-01/2023-01-01",
    query={"eo:cloud_cover": {"lt": 50}},
)

items = search.item_collection()

stack = stackstac.stack(
    items[:3],
    assets=["B02", "B03", "B04", "B08", "B11", "SCL"],
    bounds_latlon=bbox,
    snap_bounds=False,
    resampling=rasterio.enums.Resampling.bilinear,
    rescale=False,
    xy_coords="center",
    sortby_date="asc",
    band_coords=True,
)

# rename to common names when common names are available
common_names = stack.common_name.where(~stack.common_name.isnull(), stack.band)
stack = stack.where(lambda x: x > 0, other=np.nan).assign_coords(
    band=lambda x: common_names.rename("band"),  # use common names
)

da = stack.compute()

Compute spectral indices as xr.DataArray

For me it would be it would be convenient to have the spectral indices in the same object as the surface reflectance data. However, as the data comes in xr.DataArray, with several coordinates associated to the different dimensions, it is awkward to concat a new xr.DataArray along a dimension:

# extract arrays that are used to compute spectral index
green = da.sel(band="green")
nir = da.sel(band="nir")

# and compute the index.
ndwi = (green - nir) / (green + nir)

# concat at this point fails because of the "band" coordinates
# xr.concat([da, ndwi], "band")  

# so first all metadata has to be added..
ndwi = ndwi.expand_dims(dim={"band": np.array(["NDWI"])}, axis=1)
ndwi = ndwi.assign_coords(
    title=("band", np.array(["Normalized-Difference Water Index (NDWI)"]))
)
ndwi = ndwi.assign_coords(gsd=("band", np.array([10.0])))
ndwi = ndwi.assign_coords(common_name=("band", np.array(["NDWI"])))
ndwi = ndwi.assign_coords(center_wavelength=("band", np.array([None])))
ndwi = ndwi.assign_coords(full_width_half_max=("band", np.array([None])))

# and now it can be merged
xr.concat([da, ndwi], "band")

Compute spectral indices from xr.Dataset

When the data would be hold in a xr.Dataset such operation would be much more obvious:

# This function copies coordinates to attributes because coordinates associated to
# a dimension that is dropped will also be dropped.
def stackstac_to_dataset(stack):
    stack = stack.copy()
    ds = stack.to_dataset("band")

    # loop over coordinates that have band dimension and save those as attrs in variables
    for coord, da in ds.band.coords.items():
        if "band" in da.dims:
            for i, band in enumerate(stack.band.values):
                ds[band].attrs[coord] = da.values[i]

    # dropping band dimension will also drop coords with ref to band dimension
    ds = ds.drop_dims("band")
    return ds


ds = stackstac_to_dataset(stack)

# compute spectral index
ds["NDWI"] = (ds.green - ds.nir) / (ds.green + ds.nir)
ds

Would it be possible to add a as_dataset argument to stackstac.stack()?