2D bounds - simple version (xarray-contrib#370)

* 2D bounds - simple version * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix test * upd whatsnew * pre-commit black * add doc page and better docstring * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Apply suggestions from code review Co-authored-by: Deepak Cherian <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Raise on name collision - changes to doc after review * Fix typo Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Deepak Cherian <[email protected]>
aidanheerdegen · Nov 11, 2022 · 205e673 · 205e673
1 parent 614b8ef
commit 205e673
Show file tree

Hide file tree

Showing 5 changed files with 78 additions and 16 deletions.
diff --git a/cf_xarray/accessor.py b/cf_xarray/accessor.py
@@ -461,17 +461,37 @@ def wrapper(obj: DataArray | Dataset, key: str):
 }
 
 
-def _guess_bounds_dim(da, dim=None):
+def _guess_bounds_dim(da, dim=None, out_dim="bounds"):
     """
-    Guess bounds values given a 1D coordinate variable.
+    Guess bounds values given a 1D or 2D coordinate variable.
     Assumes equal spacing on either side of the coordinate label.
+    This is a coarse approximation, especially for 2D bounds on curvilinear grids.
     """
     if dim is None:
-        if da.ndim != 1:
+        if da.ndim not in [1, 2]:
             raise ValueError(
-                f"If dim is None, variable {da.name} must be 1D. Received {da.ndim}D variable instead."
+                f"If dim is None, variable {da.name} must be 1D or 2D. Received {da.ndim}D variable instead."
             )
-        (dim,) = da.dims
+        dim = da.dims
+    if not isinstance(dim, str):
+        if len(dim) > 2:
+            raise NotImplementedError(
+                "Adding bounds with more than 2 dimensions is not supported."
+            )
+        elif len(dim) == 2:
+            daX = _guess_bounds_dim(da, dim[0]).rename(bounds="Xbnds")
+            daXY = _guess_bounds_dim(daX, dim[1]).rename(bounds="Ybnds")
+            return xr.concat(
+                [
+                    daXY.isel(Xbnds=0, Ybnds=0),
+                    daXY.isel(Xbnds=0, Ybnds=1),
+                    daXY.isel(Xbnds=1, Ybnds=1),
+                    daXY.isel(Xbnds=1, Ybnds=0),
+                ],
+                out_dim,
+            )
+        else:
+            dim = dim[0]
     if dim not in da.dims:
         (dim,) = da.cf.axes[dim]
     if dim not in da.coords:
@@ -482,7 +502,7 @@ def _guess_bounds_dim(da, dim=None):
     diff = da.diff(dim)
     lower = da - diff / 2
     upper = da + diff / 2
-    bounds = xr.concat([lower, upper], dim="bounds")
+    bounds = xr.concat([lower, upper], dim=out_dim)
 
     first = (bounds.isel({dim: 0}) - diff.isel({dim: 0})).assign_coords(
         {dim: da[dim][0]}
@@ -2169,18 +2189,28 @@ 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], *, dim=None):
+    def add_bounds(
+        self,
+        keys: str | Iterable[str],
+        *,
+        dim: str | Iterable[str] | None = None,
+        output_dim: str = "bounds",
+    ):
         """
         Returns a new object with bounds variables. The bounds values are guessed assuming
-        equal spacing on either side of a coordinate label.
+        equal spacing on either side of a coordinate label. The linear estimation is only a
+        coarse approximation, especially 2D bounds on curvilinear grids. It is always better to use
+        bounds generated as part of the grid creation process. This method is purely for convenience.
 
         Parameters
         ----------
         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.
+        dim : str or Iterable[str], optional
+            Core dimension(s) along which to estimate bounds. For 2D bounds, it can
+            be a list of 2 dimension names.
+        output_dim : str
+            The name of the bounds dimension to add.
 
         Returns
         -------
@@ -2226,9 +2256,17 @@ def add_bounds(self, keys: str | Iterable[str], *, dim=None):
             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), dim=dim
+            out = _guess_bounds_dim(
+                obj[var].reset_coords(drop=True), dim=dim, out_dim=output_dim
             )
+            if output_dim in obj.dims and (new := out[output_dim].size) != (
+                old := obj[output_dim].size
+            ):
+                raise ValueError(
+                    f"The `{output_dim}` dimension already exists but has a different length than the new one "
+                    f"({old} vs {new}). Please provide another bound dimension name with `output_dim`."
+                )
+            obj.coords[bname] = out
             obj[var].attrs["bounds"] = bname
 
         return self._maybe_to_dataarray(obj)

diff --git a/cf_xarray/tests/test_accessor.py b/cf_xarray/tests/test_accessor.py
@@ -13,6 +13,7 @@
 from xarray.testing import assert_allclose, assert_identical
 
 import cf_xarray  # noqa
+from cf_xarray.helpers import vertices_to_bounds
 from cf_xarray.utils import parse_cf_standard_name_table
 
 from ..datasets import (
@@ -799,26 +800,38 @@ def test_add_bounds_multiple():
 
 
 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)},
     )
 
+    # 2D
+    expected = (
+        vertices_to_bounds(
+            np.arange(0, 13, 3).reshape(5, 1) + np.arange(-2, 2).reshape(1, 4)
+        )
+        .rename("z_bounds")
+        .assign_coords(**ds.coords)
+    )
+    actual = ds.cf.add_bounds("z").z_bounds.reset_coords(drop=True)
+    xr.testing.assert_identical(actual, expected)
+
+    # 1D
     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")
 
+    with pytest.raises(ValueError, match="The `bounds` dimension already exists"):
+        ds.cf.add_bounds("z").cf.add_bounds("x")
+
 
 def test_bounds():
     ds = airds.copy(deep=False).cf.add_bounds("lat")

diff --git a/doc/2D_bounds_error.png b/doc/2D_bounds_error.png
diff --git a/doc/bounds.md b/doc/bounds.md
@@ -10,3 +10,9 @@ See
 1. {py:func}`Dataset.cf.add_bounds`,
 1. {py:func}`cf_xarray.bounds_to_vertices`,
 1. {py:func}`cf_xarray.vertices_to_bounds`
+
+`cf_xarray` supports parsing [coordinate bounds](http://cfconventions.org/Data/cf-conventions/cf-conventions-1.10/cf-conventions.html#cell-boundaries) as encoded in the CF `bounds` attribute. A useful feature for incomplete dataset is also the automatic bounds estimation possible through `cf.add_bounds`. This method will estimate the missing bounds by finding the middle points between elements of the given coordinate, but also by extrapolating to find the outer bounds of the grid. This linear estimation works well with rectilinear grids, but it is only a coarse approximation for curvilinear and simple irregular grids.
+
+As an example, we present a "rotated pole" grid. It is defined on a rotated rectilinear grid which uses the `rlat` and  `rlon` 1D coordinates, over North America at a resolution of 0.44°. The datasets comes with 2D `lat` and `lon` coordinates. `cf_xarray` will estimate the bounds by linear interpolation (extrapolation at the edges) of the existing `lon` and `lat`, which yields good results on parts of the grid where the rotation is small. However the errors is larger in other places, as seen when visualizing the distance in degrees between the estimated bounds and the true bounds.
+
+![2d bounds error](2D_bounds_error.png)
diff --git a/doc/whats-new.rst b/doc/whats-new.rst
@@ -3,6 +3,11 @@
 What's New
 ----------
 
+v0.7.5 (unreleased)
+===================
+- ``cf.add_bounds`` can estimate 2D bounds using an approximate linear interpolation (:pr:`370`).
+  By `Pascal Bourgault`_.
+
 v0.7.4 (July 14, 2022)
 ======================