Convert nD bounds arrays to (n-1)D vertex arrays

cf_xarray/helpers.py

@@ -6,7 +6,10 @@
 
 
 def bounds_to_vertices(
-    bounds: DataArray, bounds_dim: str, order: Optional[str] = "counterclockwise"
+    bounds: DataArray,
+    bounds_dim: str,
+    core_dims=None,
+    order: Optional[str] = "counterclockwise",
 ) -> DataArray:
     """
     Convert bounds variable to vertices. There 2 covered cases:
@@ -18,15 +21,18 @@ def bounds_to_vertices(
     Parameters
     ----------
     bounds : DataArray
-        The bounds to convert. Must be of shape (N, 2) or (N, M, 4).
+        The bounds to convert.
     bounds_dim : str
         The name of the bounds dimension of `bounds` (the one of length 2 or 4).
     order : {'counterclockwise', 'clockwise', None}
-        Valid for 2D coordinates only (bounds of shape (N, M, 4), ignored otherwise.
+        Valid for 2D coordinates only (i.e. bounds of shape (..., N, M, 4), ignored otherwise.
         Order the bounds are given in, assuming that ax0-ax1-upward is a right handed
         coordinate system, where ax0 and ax1 are the two first dimensions of `bounds`.
         If None, the counterclockwise version is computed and then verified. If the
         check fails the clockwise version is returned. See Notes for more details.
+    core_dims : list, optional
+        List of core dimensions for apply_ufunc. This must not include bounds_dims.
+        The shape of (*core_dims, bounds_dim) must be (N, 2) or (N, M, 4).
 
     Returns
     -------
@@ -44,41 +50,60 @@ def bounds_to_vertices(
 
     Please refer to the CF conventions document : http://cfconventions.org/Data/cf-conventions/cf-conventions-1.8/cf-conventions.html#cell-boundaries.
     """
-    # Get old and new dimension names and retranspose array to have bounds dim at axis 0.
-    bnd_dim = (
-        bounds_dim if isinstance(bounds_dim, str) else bounds.get_axis_num(bounds_dim)
+
+    if core_dims is None:
+        core_dims = [dim for dim in bounds.dims if dim != bounds_dim]
+
+    output_sizes = {f"{dim}_vertices": bounds.sizes[dim] + 1 for dim in core_dims}
+    output_core_dims = list(output_sizes.keys())
+
+    n_core_dims = len(core_dims)
+    nbounds = bounds[bounds_dim].size
+
+    if not (n_core_dims == 2 and nbounds == 4) and not (
+        n_core_dims == 1 and nbounds == 2
+    ):
+        raise ValueError(
+            f"Bounds format not understood. Got {bounds.dims} with shape {bounds.shape}."
+        )
+
+    return xr.apply_ufunc(
+        _bounds_helper,
+        bounds,
+        input_core_dims=[core_dims + [bounds_dim]],
+        dask="parallelized",
+        kwargs={"n_core_dims": n_core_dims, "nbounds": nbounds, "order": order},
+        output_core_dims=[output_core_dims],
+        dask_gufunc_kwargs=dict(output_sizes=output_sizes),
+        output_dtypes=[bounds.dtype],
     )
-    old_dims = [dim for dim in bounds.dims if dim != bnd_dim]
-    new_dims = [f"{dim}_vertices" for dim in old_dims]
-    values = bounds.transpose(bnd_dim, *old_dims).data
-    if len(old_dims) == 2 and bounds.ndim == 3 and bounds[bnd_dim].size == 4:
+
+
+def _bounds_helper(values, n_core_dims, nbounds, order):
+    if n_core_dims == 2 and nbounds == 4:
         # Vertices case (2D lat/lon)
         if order in ["counterclockwise", None]:
             # Names assume we are drawing axis 1 upward et axis 2 rightward.
-            bot_left = values[0, :, :]
-            bot_right = values[1, :, -1:]
-            top_right = values[2, -1:, -1:]
-            top_left = values[3, -1:, :]
+            bot_left = values[..., :, :, 0]
+            bot_right = values[..., :, -1:, 1]
+            top_right = values[..., -1:, -1:, 2]
+            top_left = values[..., -1:, :, 3]
             vertex_vals = np.block([[bot_left, bot_right], [top_left, top_right]])
         if order is None:  # We verify if the ccw version works.
             calc_bnds = vertices_to_bounds(vertex_vals).values
             order = "counterclockwise" if np.all(calc_bnds == values) else "clockwise"
         if order == "clockwise":
-            bot_left = values[0, :, :]
-            top_left = values[1, -1:, :]
-            top_right = values[2, -1:, -1:]
-            bot_right = values[3, :, -1:]
+            bot_left = values[..., :, :, 0]
+            top_left = values[..., -1:, :, 1]
+            top_right = values[..., -1:, -1:, 2]
+            bot_right = values[..., :, -1:, 3]
             # Our asumption was wrong, axis 1 is rightward and axis 2 is upward
             vertex_vals = np.block([[bot_left, bot_right], [top_left, top_right]])
-    elif len(old_dims) == 1 and bounds.ndim == 2 and bounds[bnd_dim].size == 2:
+    elif n_core_dims == 1 and nbounds == 2:
         # Middle points case (1D lat/lon)
-        vertex_vals = np.concatenate((values[0, :], values[1, -1:]))
-    else:
-        raise ValueError(
-            f"Bounds format not understood. Got {bounds.dims} with shape {bounds.shape}."
-        )
+        vertex_vals = np.concatenate((values[..., :, 0], values[..., -1:, 1]), axis=-1)
 
-    return xr.DataArray(vertex_vals, dims=new_dims)
+    return vertex_vals
 
 
 def vertices_to_bounds(

doc/whats-new.rst

@@ -6,6 +6,7 @@ What's New
 v0.6.1 (unreleased)
 ===================
 - Support detecting pint-backed Variables with units-based criteria. By `Deepak Cherian`_.
+- Support reshaping nD bounds arrays to (n-1)D vertex arrays. By `Deepak Cherian`_.
 
 v0.6.0 (June 29, 2021)
 ======================