Repository: opendatacube/odc-geo
Pull request #151 (2024-05-14 23:43:46)

Fixes for rotated GeoBox representation on xarray
- Rotated GeoBox didn't survive serialization to disk, and now does
   - Don't rely on custom `encoding["_transform"]` instead use attribute
- Improved detection of spatial dimensions with non-standard names
   - This change allows packing data arrays with different geo-registration into one `xarray.Dataset`
   - Bands with different native resolutions without upsampling is a good example of such use-case 
- Allow arbitrary names for dimensions in `xr_zeros|wrap_xr` (i.e. `dims=("T", "Y", "X", "B")`)
- Make it easy to force `y,x` names for spatial dimensions even when geobox is in 4326, `always_yx=True`
- Allow arbitrary number of non-spatial dimensions in any position relative to spatial dimensions
- Removed "LayerControl present" detection code in `.odc.explore` as it relies on internal details of `folium.Map`
   - ping: @robbibt, the idea is that if user is supplying existing map to `.explore()` then we won't try adding `LayerControl` or change map view, and just add ImageOverlay to the existing map
   - I had some issues with adding multiple images to the same map, looks like it's due to `LayerControl` and some problem/unexpected behaviour in `folium` that I have not fully debugged
- Mapping functions now expose optional `index=` parameter, used for layer ordering in folium/ipyleaflet


Patch:
diff --git a/odc/geo/_map.py b/odc/geo/_map.py
index 63544aa..181eb74 100644
--- a/odc/geo/_map.py
+++ b/odc/geo/_map.py
@@ -1,4 +1,4 @@
-from typing import Any, Optional, Tuple, Dict
+from typing import Any, Dict, Optional, Tuple
 
 import xarray as xr
 
@@ -10,17 +10,18 @@ from .geobox import GeoBox
 
 
 # pylint: disable=import-outside-toplevel, redefined-builtin, too-many-locals
-def _add_to_folium(url, bounds, map, name=None, **kw):
+def _add_to_folium(url, bounds, map, name=None, index=None, **kw):
     assert have.folium
 
     from folium.raster_layers import ImageOverlay
 
     img_overlay = ImageOverlay(url, bounds, name=name, **kw)
-    img_overlay.add_to(map)
+    if map is not None:
+        img_overlay.add_to(map, name=name, index=index)
     return img_overlay
 
 
-def _add_to_ipyleaflet(url, bounds, map, name=None, **kw):
+def _add_to_ipyleaflet(url, bounds, map, name=None, index=None, **kw):
     assert have.ipyleaflet
 
     from ipyleaflet import ImageOverlay, Map
@@ -31,7 +32,8 @@ def _add_to_ipyleaflet(url, bounds, map, name=None, **kw):
         kw.update(name=name)
 
     img_overlay = ImageOverlay(url=url, bounds=bounds, **kw)
-    map.add(img_overlay)
+    if map is not None:
+        map.add(img_overlay, index=index)
 
     return img_overlay
 
@@ -56,6 +58,7 @@ def add_to(
     map: Any,
     *,
     name: Optional[str] = None,
+    index: Optional[int] = None,
     fmt: str = "png",
     max_size: int = 4096,
     resampling: str = "nearest",
@@ -167,12 +170,13 @@ def add_to(
     if _add_to is None:
         return url, bounds
 
-    return _add_to(url, bounds, map, name=name, **kw)
+    return _add_to(url, bounds, map, name=name, index=index, **kw)
 
 
 def explore(
     xx: Any,
     map: Optional[Any] = None,
+    *,
     bands: Optional[Tuple[str, str, str]] = None,
     vmin: Optional[float] = None,
     vmax: Optional[float] = None,
@@ -245,12 +249,10 @@ def explore(
             "Please install it before using `.explore()`."
         )
 
-    from folium import Map, LayerControl
+    from folium import LayerControl, Map
 
-    # Update any supplied kwargs with custom params
     map_kwds = {} if map_kwds is None else map_kwds
-    kwargs.update(cmap=cmap, vmin=vmin, vmax=vmax, robust=robust, resampling=resampling)
-    map_kwds.update(tiles=tiles, attr=attr)
+    new_map = map is None
 
     # If input is a dataset, convert to an RGBA array
     if isinstance(xx, xr.Dataset):
@@ -258,11 +260,19 @@ def explore(
 
     # Create folium Map if required
     if map is None:
-        map = Map(**map_kwds)
+        map = Map(tiles=tiles, attr=attr, **map_kwds)
 
     # Add to map and raise a friendly error if data has unsuitable dims
     try:
-        xx.odc.add_to(map, **kwargs)
+        xx.odc.add_to(
+            map,
+            cmap=cmap,
+            vmin=vmin,
+            vmax=vmax,
+            robust=robust,
+            resampling=resampling,
+            **kwargs,
+        )
     except ValueError as e:
         raise ValueError(
             "Only 2D single-band (x, y) or 3D multi-band (x, y, band) "
@@ -271,14 +281,12 @@ def explore(
             "or `.sel()`: `da.isel(time=0).odc.explore()`."
         ) from e
 
-    # Zoom map to extent of data
-    map.fit_bounds(xx.odc.map_bounds())
+    if new_map:
+        # Zoom map to extent of data
+        map.fit_bounds(xx.odc.map_bounds())
 
-    # Add a layer control if requested and not already added
-    layer_control_added = any(
-        isinstance(child, LayerControl) for child in map._children.values()
-    )
-    if layer_control and not layer_control_added:
-        LayerControl().add_to(map)
+        # Add a layer control if requested
+        if layer_control:
+            LayerControl().add_to(map)
 
     return map
diff --git a/odc/geo/_version.py b/odc/geo/_version.py
index 0f4d995..fdca860 100644
--- a/odc/geo/_version.py
+++ b/odc/geo/_version.py
@@ -1,3 +1,3 @@
 """version information only."""
 
-__version__ = "0.4.4"
+__version__ = "0.4.5"
diff --git a/odc/geo/_xr_interop.py b/odc/geo/_xr_interop.py
index d63e394..b149eca 100644
--- a/odc/geo/_xr_interop.py
+++ b/odc/geo/_xr_interop.py
@@ -38,6 +38,7 @@ from .geom import Geometry
 from .math import (
     affine_from_axis,
     approx_equal_affine,
+    is_affine_st,
     maybe_int,
     resolution_from_affine,
 )
@@ -63,6 +64,14 @@ _DEFAULT_CRS_COORD_NAME = "spatial_ref"
 # these attributes are pruned during reproject
 SPATIAL_ATTRIBUTES = ("crs", "crs_wkt", "grid_mapping", "gcps", "epsg")
 
+# dimensions with these names are considered spatial
+STANDARD_SPATIAL_DIMS = [
+    ("y", "x"),
+    ("yc", "xc"),
+    ("latitude", "longitude"),
+    ("lat", "lon"),
+]
+
 
 @dataclass
 class GeoState:
@@ -151,14 +160,30 @@ def spatial_dims(
     :returns: ``None`` if no dimensions with expected names are found
     :returns: ``('y', 'x') | ('latitude', 'longitude') | ('lat', 'lon')``
     """
-    guesses = [("y", "x"), ("latitude", "longitude"), ("lat", "lon")]
+
+    def skip_dim(dim: str) -> bool:
+        if dim in ("time", "band", "bands", "wavelength", "wavelengths"):
+            return True
+
+        # skip dimensions without coord of the same name
+        if dim not in xx.coords:
+            return True
+
+        coord = xx.coords[dim]
+        # Primary coordinate for spatial dimension must have floating point type
+        if coord.dtype.kind != "f":
+            return True
+
+        return False
 
     _dims = [str(dim) for dim in xx.dims]
     dims = set(_dims)
-    for guess in guesses:
+    for guess in STANDARD_SPATIAL_DIMS:
         if dims.issuperset(guess):
             return guess
 
+    _dims = [dim for dim in _dims if not skip_dim(str(dim))]
+
     if relaxed and len(_dims) >= 2:
         return _dims[-2], _dims[-1]
 
@@ -335,11 +360,21 @@ def crop(
 
     .. seealso:: :py:meth:`odc.geo.xr.mask`
     """
+    meta: ODCExtension = xx.odc
+    sdims = meta.spatial_dims
+    gbox = meta.geobox
+
+    if sdims is None or gbox is None:
+        raise ValueError("Can't locate spatial dimensions")
+
+    if not isinstance(gbox, GeoBox):
+        raise ValueError("Can't crop GCPGeoBox")
+
     # Create new geobox with pixel grid of `xx` but enclosing `poly`.
-    poly_geobox = xx.odc.geobox.enclosing(poly)
+    poly_geobox = gbox.enclosing(poly)
 
     # Calculate ROI slices into `xx` for intersection between both geoboxes.
-    roi = xx.odc.geobox.overlap_roi(poly_geobox)
+    roi = gbox.overlap_roi(poly_geobox)
 
     # Verify that `poly` overlaps with `xx` by checking if the returned
     # ROI is empty
@@ -349,9 +384,6 @@ def crop(
         )
 
     # Crop spatial dims of `xx` using ROI
-    sdims = spatial_dims(xx)
-    if sdims is None:
-        raise ValueError("Can't locate spatial dimensions")
     xx_cropped = xx.isel({sdims[0]: roi[0], sdims[1]: roi[1]})
 
     # Optionally mask data outside rasterized `poly`
@@ -362,20 +394,42 @@ def crop(
 
 
 def xr_coords(
-    gbox: SomeGeoBox, crs_coord_name: Optional[str] = _DEFAULT_CRS_COORD_NAME
+    gbox: SomeGeoBox,
+    crs_coord_name: Optional[str] = _DEFAULT_CRS_COORD_NAME,
+    always_yx: bool = False,
+    dims: Optional[Tuple[str, str]] = None,
 ) -> Dict[Hashable, xarray.DataArray]:
     """
     Dictionary of Coordinates in xarray format.
 
+    :param gbox:
+      :py:class:`~odc.geo.geobox.GeoBox` or :py:class:`~odc.geo.gcp.GCPGeoBox`
+
     :param crs_coord_name:
-       Use custom name for CRS coordinate, default is "spatial_ref". Set to ``None`` to not generate
-       CRS coordinate at all.
+       Use custom name for CRS coordinate, default is "spatial_ref". Set to
+       ``None`` to not generate CRS coordinate at all.
+
+    :param always_yx:
+       If True, always use names ``y,x`` for spatial coordinates even for
+       geographic geoboxes.
+
+    :param dims:
+       Use custom names for spatial dimensions, default is to use ``y,x`` or
+       ``latitude, longitude`` based on projection used. Dimensions are supplied
+       in "array" order, i.e. ``('y', 'x')``.
 
     :returns:
-       Dictionary ``name:str -> xr.DataArray``. Where names are either ``y,x`` for projected or
+       Dictionary ``name:str -> xr.DataArray``. Where names are either as
+       supplied by ``dims=`` or otherwise ``y,x`` for projected or
        ``latitude, longitude`` for geographic.
 
     """
+    if dims is None:
+        if always_yx:
+            dims = ("y", "x")
+        else:
+            dims = gbox.dimensions
+
     attrs = {}
     crs = gbox.crs
     if crs is not None:
@@ -386,8 +440,7 @@ def xr_coords(
 
     if isinstance(gbox, GCPGeoBox):
         coords: Dict[Hashable, xarray.DataArray] = {
-            name: _mk_pixel_coord(name, sz, None)
-            for name, sz in zip(gbox.dimensions, gbox.shape)
+            name: _mk_pixel_coord(name, sz) for name, sz in zip(dims, gbox.shape)
         }
         gcps = gbox.gcps()
     else:
@@ -395,12 +448,11 @@ def xr_coords(
         if gbox.axis_aligned:
             coords = {
                 name: _coord_to_xr(name, coord, **attrs)
-                for name, coord in gbox.coordinates.items()
+                for name, coord in zip(dims, gbox.coordinates.values())
             }
         else:
             coords = {
-                name: _mk_pixel_coord(name, sz, transform)
-                for name, sz in zip(gbox.dimensions, gbox.shape)
+                name: _mk_pixel_coord(name, sz) for name, sz in zip(dims, gbox.shape)
             }
 
     if crs_coord_name is not None and crs is not None:
@@ -414,14 +466,11 @@ def xr_coords(
 def _mk_pixel_coord(
     name: str,
     sz: int,
-    transform: Optional[Affine],
 ) -> xarray.DataArray:
     data = numpy.arange(0.5, sz, dtype="float32")
     xx = xarray.DataArray(
         data, coords={name: data}, dims=(name,), attrs={"units": "pixel"}
     )
-    if transform is not None:
-        xx.encoding["_transform"] = transform[:6]
     return xx
 
 
@@ -531,21 +580,24 @@ def _extract_transform(
         except ValueError:
             return None
 
-    if original_transform is not None and approx_equal_affine(
-        transform, original_transform
-    ):
-        transform = original_transform
+    if original_transform is not None:
+        if not is_affine_st(original_transform):
+            # non-axis aligned geobox detected
+            # adjust transform
+            #  world <- pix' <- pix
+            transform = original_transform * transform
+
+        if approx_equal_affine(transform, original_transform):
+            transform = original_transform
 
-    if not gcp and (_pix2world := _xx.encoding.get("_transform", None)) is not None:
-        # non-axis aligned geobox detected
-        # adjust transform
-        #  world <- pix' <- pix
-        transform = Affine(*_pix2world) * transform
     return transform
 
 
 def _locate_geo_info(src: XarrayObject) -> GeoState:
     # pylint: disable=too-many-locals
+    if len(src.dims) < 2:
+        return GeoState()
+
     sdims = spatial_dims(src, relaxed=True)
     if sdims is None:
         return GeoState()
@@ -1011,6 +1063,8 @@ def wrap_xr(
     time=None,
     nodata=None,
     crs_coord_name: Optional[str] = _DEFAULT_CRS_COORD_NAME,
+    always_yx: bool = False,
+    dims: Optional[Tuple[str, ...]] = None,
     axis: Optional[int] = None,
     **attrs,
 ) -> xarray.DataArray:
@@ -1021,37 +1075,74 @@ def wrap_xr(
     :param gbox: Geobox, must same shape as last two axis of ``im``
     :param time: optional time axis value(s), defaults to None
     :param nodata: optional `nodata` value, defaults to None
+    :param crs_coord_name: allows to change name of the crs coordinate variable
+    :param always_yx: If True, always use names ``y,x`` for spatial coordinates
+    :param dims: Custom names for spatial dimensions
+    :param axis: Which axis of the input array corresponds to Y,X
     :param attrs: Any other attributes to set on the result
     :return: xarray DataArray
     """
+    # pylint: disable=too-many-locals,too-many-arguments
+    assert dims is None or len(dims) == im.ndim
+
     if axis is None:
         axis = 1 if time is not None else 0
+    elif axis < 0:  # handle numpy style negative axis
+        axis = int(im.ndim) + axis
 
     if im.ndim == 2 and axis == 1:
         im = im[numpy.newaxis, ...]
 
-    assert axis in (0, 1)  # upto 1 extra dimension on the left only
-    assert im.ndim - axis - 2 in (0, 1)  # upto 1 extra dimension on the right only
+    assert axis >= 0
+    assert im.ndim - axis - 2 >= 0
     assert im.shape[axis : axis + 2] == gbox.shape
 
-    prefix_dims: Tuple[str, ...] = ("time",) if axis == 1 else ()
-    postfix_dims: Tuple[str, ...] = ("band",) if im.ndim - axis > 2 else ()
+    def _prefix_dims(n):
+        if n == 0:
+            return ()
+        if n == 1:
+            return ("time",)
+        return ("time", *[f"dim_{i}" for i in range(n - 1)])
+
+    def _postfix_dims(n):
+        if n == 0:
+            return ()
+        if n == 1:
+            return ("band",)
+        return (f"b_{i}" for i in range(n))
+
+    sdims: Optional[Tuple[str, str]] = None
+    if dims is None:
+        sdims = ("y", "x") if always_yx else gbox.dimensions
+        dims = (*_prefix_dims(axis), *sdims, *_postfix_dims(im.ndim - axis - 2))
+    else:
+        sdims = dims[axis], dims[axis + 1]
 
-    dims = (*prefix_dims, *gbox.dimensions, *postfix_dims)
-    coords = xr_coords(gbox, crs_coord_name=crs_coord_name)
+    prefix_dims = dims[:axis]
+    postfix_dims = dims[axis + 2 :]
+
+    coords = xr_coords(
+        gbox,
+        crs_coord_name=crs_coord_name,
+        always_yx=always_yx,
+        dims=sdims,
+    )
 
     if time is not None:
         if not isinstance(time, xarray.DataArray):
             if len(prefix_dims) > 0 and isinstance(time, (str, datetime)):
                 time = [time]
 
-            time = xarray.DataArray(time, dims=prefix_dims).astype("datetime64[ns]")
+            time = xarray.DataArray(time, dims=prefix_dims[:1]).astype("datetime64[ns]")
 
         coords["time"] = time
+
     if postfix_dims:
-        coords["band"] = xarray.DataArray(
-            [f"b{i}" for i in range(im.shape[-1])], dims=postfix_dims
-        )
+        for a, dim in enumerate(postfix_dims):
+            nb = im.shape[axis + 2 + a]
+            coords[dim] = xarray.DataArray(
+                [f"b{i}" for i in range(nb)], dims=(dim,), name=dim
+            )
 
     if nodata is not None:
         attrs = {"nodata": nodata, **attrs}
@@ -1081,6 +1172,9 @@ def xr_zeros(
     :param crs_coord_name: allows to change name of the crs coordinate variable
 
     :return: :py:class:`xarray.DataArray` filled with zeros (numpy or dask)
+
+    .. seealso:: :py:meth:`odc.geo.xr.wrap_xr`
+
     """
     if time is not None:
         _shape: Tuple[int, ...] = (len(time), *geobox.shape.yx)


Test patch:
diff --git a/odc/geo/testutils.py b/odc/geo/testutils.py
index bb27e95..ba402bc 100644
--- a/odc/geo/testutils.py
+++ b/odc/geo/testutils.py
@@ -201,7 +201,6 @@ def purge_crs_info(xx: xr.DataArray) -> xr.DataArray:
         for attr in attributes_to_clear:
             _x.attrs.pop(attr, None)
         _x.encoding.pop("grid_mapping", None)
-        _x.encoding.pop("_transform", None)
         return _x
 
     # remove non-dimensional coordinate, which is CRS in our case
diff --git a/tests/test_mpu_fs.py b/tests/test_mpu_fs.py
index cead7f6..4036b0f 100644
--- a/tests/test_mpu_fs.py
+++ b/tests/test_mpu_fs.py
@@ -6,6 +6,8 @@ from uuid import uuid4
 
 import pytest
 
+_ = pytest.importorskip("odc.geo.cog")
+
 from odc.geo.cog._mpu_fs import MPUFileSink
 
 # pylint: disable=protected-access
@@ -16,7 +18,7 @@ def slurp(path: Path) -> bytes:
         return f.read()
 
 
-@pytest.mark.parametrize("parts_base", [f"parts-{uuid4().hex[:8]}", None])
+@pytest.mark.parametrize("parts_base", [f"parts-3782781", None])
 @pytest.mark.parametrize("num_parts", [1, 2, 3, 10])
 def test_filesink(tmp_path: Path, parts_base: str | None, num_parts: int):
     dst = Path(tmp_path / f"{uuid4().hex}.bin")
diff --git a/tests/test_xr_interop.py b/tests/test_xr_interop.py
index 7447155..10e74ca 100644
--- a/tests/test_xr_interop.py
+++ b/tests/test_xr_interop.py
@@ -18,6 +18,7 @@ from odc.geo.xr import (
     ODCExtensionDs,
     rasterize,
     register_geobox,
+    spatial_dims,
     wrap_xr,
     xr_coords,
     xr_reproject,
@@ -25,6 +26,11 @@ from odc.geo.xr import (
 )
 
 # pylint: disable=redefined-outer-name,import-outside-toplevel,protected-access
+TEST_GEOBOXES_SMALL_AXIS_ALIGNED = [
+    GeoBox.from_bbox((-10, -2, 5, 4), "epsg:4326", tight=True, resolution=0.2),
+    GeoBox.from_bbox((-10, -2, 5, 4), "epsg:3857", tight=True, resolution=1),
+    GeoBox.from_bbox((-10, -2, 5, 4), "epsg:3857", tight=True, resolution=resxy_(1, 2)),
+]
 
 
 @pytest.fixture
@@ -87,8 +93,7 @@ def test_geobox_xr_coords():
     assert cc["spatial_ref"].attrs["spatial_ref"] == gbox.crs.wkt
     assert isinstance(cc["spatial_ref"].attrs["grid_mapping_name"], str)
     assert isinstance(cc["spatial_ref"].attrs["GeoTransform"], str)
-    assert cc["x"].encoding["_transform"] == _gbox.affine[:6]
-    assert cc["y"].encoding["_transform"] == _gbox.affine[:6]
+    assert xr_zeros(_gbox, dtype="uint16").odc.geobox == _gbox
 
     # geographic CRS
     A = mkA(0, scale=(0.1, -0.1), translation=(10, 30))
@@ -176,6 +181,7 @@ def test_odc_extension(xx_epsg4326: xr.DataArray, geobox_epsg4326: GeoBox):
     assert xx.odc.output_geobox("epsg:3857").crs == "epsg:3857"
     assert xx.odc.map_bounds() == gbox.map_bounds()
     assert xx.odc.output_geobox("utm").crs.epsg is not None
+    assert xx.odc.aspect == gbox.aspect
 
     # this drops encoding/attributes, but crs/geobox should remain the same
     _xx = xx * 10.0
@@ -352,6 +358,9 @@ def test_wrap_xr():
     assert xx.dims == gbox.dims
     assert xx.attrs == {}
 
+    assert wrap_xr(data, gbox, always_yx=True).dims == ("y", "x")
+    assert wrap_xr(data, gbox, dims=("Y", "X")).dims == ("Y", "X")
+
     xx = wrap_xr(data, gbox, nodata=None)
     assert xx.attrs == {}
 
@@ -377,6 +386,31 @@ def test_wrap_xr():
     assert xx.band.data.tolist() == ["b0"]
 
 
+@pytest.mark.parametrize("gbox", TEST_GEOBOXES_SMALL_AXIS_ALIGNED)
+@pytest.mark.parametrize("nprefix", [0, 1, 2])
+@pytest.mark.parametrize("npostfix", [0, 1, 2])
+def test_wrap_xr_nd(gbox: GeoBox, nprefix: int, npostfix: int):
+    shape = (1,) * nprefix + gbox.shape + (3,) * npostfix
+    data = np.zeros(shape, dtype="uint16")
+    xx = wrap_xr(data, gbox, axis=nprefix)
+    assert xx.odc.geobox == gbox
+    assert xx.odc.ydim == nprefix
+    assert xx.dims[:nprefix] == ("time", "dim_0", "dim_1", "dim_2", "dim_3")[:nprefix]
+
+    if npostfix == 1:
+        assert xx.dims[-1] == "band"
+    if npostfix > 1:
+        assert xx.dims[nprefix + 2 :] == tuple(f"b_{i}" for i in range(npostfix))
+
+    _dims = tuple(f"custom_{dim}" for dim in xx.dims)
+    _dims = _dims[:nprefix] + ("y", "x") + _dims[nprefix + 2 :]
+
+    yy = wrap_xr(data, gbox, axis=nprefix, dims=_dims)
+    assert yy.dims == _dims
+    assert yy.odc.geobox == gbox
+    assert yy.odc.spatial_dims == _dims[nprefix : nprefix + 2]
+
+
 @pytest.mark.parametrize("xx_time", [None, ["2020-01-30"]])
 @pytest.mark.parametrize("xx_chunks", [None, (-1, -1), (4, 4)])
 def test_xr_reproject(xx_epsg4326: xr.DataArray):
@@ -481,13 +515,6 @@ def test_is_dask_collection():
     assert is_dask_collection is dask.is_dask_collection
 
 
-TEST_GEOBOXES_SMALL_AXIS_ALIGNED = [
-    GeoBox.from_bbox((-10, -2, 5, 4), "epsg:4326", tight=True, resolution=0.2),
-    GeoBox.from_bbox((-10, -2, 5, 4), "epsg:3857", tight=True, resolution=1),
-    GeoBox.from_bbox((-10, -2, 5, 4), "epsg:3857", tight=True, resolution=resxy_(1, 2)),
-]
-
-
 @pytest.mark.parametrize("geobox", TEST_GEOBOXES_SMALL_AXIS_ALIGNED)
 @pytest.mark.parametrize(
     "bad_geo_transform",
@@ -703,3 +730,33 @@ def test_mask(xx_epsg4326, poly):
     xx_ds_masked = xx_ds.odc.mask(poly=poly)
     assert xx_ds_masked.odc.geobox == xx_ds.odc.geobox
     assert xx_ds_masked.test.isnull().any()
+
+
+def test_spatial_dims():
+    gbox = GeoBox.from_bbox([0, -10, 100, 20], "epsg:4326", tight=True, shape=(13, 29))
+    xx0 = xr_zeros(gbox, "int16", always_yx=True, time=["2020-01-01", "2020-01-02"])
+    assert spatial_dims(xx0) == ("y", "x")
+
+    assert wrap_xr(
+        xx0.data,
+        gbox,
+        dims=("T", "Y", "X"),
+        axis=1,
+    ).odc.spatial_dims == ("Y", "X")
+    assert wrap_xr(
+        xx0.data[..., np.newaxis],
+        gbox,
+        dims=("time", "Y", "X", "wavelength"),
+        axis=1,
+    ).odc.spatial_dims == ("Y", "X")
+
+    xx = wrap_xr(
+        np.zeros((2, 3, *gbox.shape.yx, 2, 3), dtype="int16"),
+        gbox,
+        axis=2,
+        dims=("A", "B", "Y", "X", "C", "D"),
+    )
+    assert "X" in xx.coords and "Y" in xx.coords
+    assert xx.odc.geobox == gbox
+    assert xx.odc.ydim == 2
+    assert xx.odc.spatial_dims == xx.dims[2:4]