Infer 1D bounds for nD variables. (#328)

Author: dcherian

cf_xarray/accessor.py

@@ -456,20 +456,32 @@ def wrapper(obj: DataArray | Dataset, key: str):
 }
 
 
-def _guess_bounds_dim(da):
+def _guess_bounds_dim(da, dim=None):
     """
     Guess bounds values given a 1D coordinate variable.
     Assumes equal spacing on either side of the coordinate label.
     """
-    assert da.ndim == 1
+    if dim is None:
+        if da.ndim != 1:
+            raise ValueError(
+                f"If dim is None, variable {da.name} must be 1D. Received {da.ndim}D variable instead."
+            )
+        (dim,) = da.dims
+    if dim not in da.dims:
+        (dim,) = da.cf.axes[dim]
+    if dim not in da.coords:
+        raise NotImplementedError(
+            "Adding bounds for unindexed dimensions is not supported currently."
+        )
 
-    dim = da.dims[0]
     diff = da.diff(dim)
     lower = da - diff / 2
     upper = da + diff / 2
     bounds = xr.concat([lower, upper], dim="bounds")
 
-    first = (bounds.isel({dim: 0}) - diff[0]).assign_coords({dim: da[dim][0]})
+    first = (bounds.isel({dim: 0}) - diff.isel({dim: 0})).assign_coords(
+        {dim: da[dim][0]}
+    )
     result = xr.concat([first, bounds], dim=dim)
 
     return result
@@ -2144,7 +2156,7 @@ def get_bounds_dim_name(self, key: str) -> str:
         assert self._obj.sizes[bounds_dim] in [2, 4]
         return bounds_dim
 
-    def add_bounds(self, keys: str | Iterable[str]):
+    def add_bounds(self, keys: str | Iterable[str], *, dim=None):
         """
         Returns a new object with bounds variables. The bounds values are guessed assuming
         equal spacing on either side of a coordinate label.
@@ -2153,6 +2165,9 @@ def add_bounds(self, keys: str | Iterable[str]):
         ----------
         keys : str or Iterable[str]
             Either a single variable name or a list of variable names.
+        dim : str, optional
+            Core dimension along whch to estimate bounds. If None, ``keys``
+            must refer to 1D variables only.
 
         Returns
         -------
@@ -2198,7 +2213,9 @@ def add_bounds(self, keys: str | Iterable[str]):
             bname = f"{var}_bounds"
             if bname in obj.variables:
                 raise ValueError(f"Bounds variable name {bname!r} will conflict!")
-            obj.coords[bname] = _guess_bounds_dim(obj[var].reset_coords(drop=True))
+            obj.coords[bname] = _guess_bounds_dim(
+                obj[var].reset_coords(drop=True), dim=dim
+            )
             obj[var].attrs["bounds"] = bname
 
         return self._maybe_to_dataarray(obj)

cf_xarray/tests/test_accessor.py

@@ -779,6 +779,28 @@ def test_add_bounds_multiple():
     assert {"x1_bounds", "x2_bounds"} <= set(multiple.cf.add_bounds("X").variables)
 
 
+def test_add_bounds_nd_variable():
+
+    ds = xr.Dataset(
+        {"z": (("x", "y"), np.arange(12).reshape(4, 3))},
+        coords={"x": np.arange(4), "y": np.arange(3)},
+    )
+
+    expected = (
+        xr.concat([ds.z - 1.5, ds.z + 1.5], dim="bounds")
+        .rename("z_bounds")
+        .transpose("bounds", "y", "x")
+    )
+    with pytest.raises(ValueError):
+        ds.cf.add_bounds("z")
+
+    actual = ds.cf.add_bounds("z", dim="x").z_bounds.reset_coords(drop=True)
+    xr.testing.assert_identical(expected, actual)
+
+    with pytest.raises(NotImplementedError):
+        ds.drop_vars("x").cf.add_bounds("z", dim="x")
+
+
 def test_bounds():
     ds = airds.copy(deep=True).cf.add_bounds("lat")