Broadcast vertices-to-bounds conversion over extra dimensions

cf_xarray/helpers.py

@@ -482,7 +482,9 @@ def _sort_vertices(vertices: np.ndarray, order: str) -> np.ndarray:
 
 
 def vertices_to_bounds(
-    vertices: DataArray, out_dims: Sequence[str] = ("bounds", "x", "y")
+    vertices: DataArray | np.ndarray,
+    out_dims: Sequence[str] = ("bounds", "x", "y"),
+    core_dims: Sequence[Hashable] | None = None,
 ) -> DataArray:
     """
     Convert vertices to CF-compliant bounds.
@@ -496,10 +498,14 @@ def vertices_to_bounds(
     Parameters
     ----------
     vertices : DataArray
-        The vertices to convert. Must be of shape (N + 1) or (N + 1, M + 1).
+        The vertices to convert. Core dimensions must have shape (N + 1)
+        or (N + 1, M + 1). Other dimensions are broadcast over.
     out_dims : Sequence[str],
         The name of the dimension in the output. The first is the 'bounds'
         dimension and the following are the coordinate dimensions.
+    core_dims : Sequence[Hashable], optional
+        The one or two vertex dimensions to convert. Defaults to all dimensions.
+        Dask-backed inputs must have a single chunk along each core dimension.
 
     Returns
     -------
@@ -525,22 +531,47 @@ def vertices_to_bounds(
     ----------
     Please refer to the CF conventions document : http://cfconventions.org/Data/cf-conventions/cf-conventions-1.8/cf-conventions.html#cell-boundaries.
     """
-    if vertices.ndim == 1:
-        bnd_vals = np.stack((vertices[:-1], vertices[1:]), axis=0)
-    elif vertices.ndim == 2:
-        bnd_vals = np.stack(
-            (
-                vertices[:-1, :-1],
-                vertices[:-1, 1:],
-                vertices[1:, 1:],
-                vertices[1:, :-1],
-            ),
-            axis=0,
-        )
-    else:
+    if not isinstance(vertices, DataArray):
+        vertices = DataArray(vertices)
+    if core_dims is None:
+        core_dims = list(vertices.dims)
+    n_core_dims = len(core_dims)
+    if n_core_dims not in (1, 2):
         raise ValueError(
             f"vertices format not understood. Got {vertices.dims} with shape {vertices.shape}."
         )
-    return xr.DataArray(bnd_vals, dims=out_dims[: vertices.ndim + 1]).transpose(
-        ..., out_dims[0]
+    if len(out_dims) < n_core_dims + 1:
+        raise ValueError("out_dims must name the bounds and each core dimension.")
+
+    output_dims = list(out_dims[1 : n_core_dims + 1]) + [out_dims[0]]
+    output_sizes = {
+        out_dim: vertices.sizes[dim] - 1
+        for dim, out_dim in zip(core_dims, output_dims, strict=False)
+    }
+    output_sizes[out_dims[0]] = 2**n_core_dims
+    return xr.apply_ufunc(
+        _vertices_helper,
+        vertices,
+        input_core_dims=[list(core_dims)],
+        output_core_dims=[output_dims],
+        exclude_dims=set(core_dims),
+        kwargs={"n_core_dims": n_core_dims},
+        dask="parallelized",
+        dask_gufunc_kwargs={"output_sizes": output_sizes},
+        output_dtypes=[vertices.dtype],
     )
+
+
+def _vertices_helper(vertices, n_core_dims):
+    if n_core_dims == 1:
+        return np.stack((vertices[..., :-1], vertices[..., 1:]), axis=-1)
+    else:
+        return np.stack(
+            (
+                vertices[..., :-1, :-1],
+                vertices[..., :-1, 1:],
+                vertices[..., 1:, 1:],
+                vertices[..., 1:, :-1],
+            ),
+            axis=-1,
+        )

cf_xarray/tests/test_helpers.py

@@ -1,4 +1,5 @@
 import numpy as np
+import pytest
 import xarray as xr
 from numpy.testing import assert_array_equal
 from xarray.testing import assert_equal
@@ -125,6 +126,71 @@ def test_vertices_to_bounds() -> None:
     assert_array_equal(mollwds.lon_bounds, lon_b)
 
 
+@pytest.mark.parametrize("ndim", [1, 2])
+@pytest.mark.parametrize("chunked", [False, True])
+def test_vertices_to_bounds_broadcast_dimensions(ndim, chunked):
+    core_dims = ["a", "b"][:ndim]
+    shape = (2,) + (4,) * ndim
+    vertices = xr.DataArray(
+        np.arange(np.prod(shape)).reshape(shape),
+        dims=["time", *core_dims],
+        coords={"time": [10, 20]},
+    ).transpose(*core_dims, "time")
+    if chunked:
+        pytest.importorskip("dask.array")
+        vertices = vertices.chunk({"time": 1})
+    out_dims = ["bounds", "x", "y"][: ndim + 1]
+
+    result = cfxr.vertices_to_bounds(vertices, out_dims=out_dims, core_dims=core_dims)
+
+    expected = xr.concat(
+        [
+            cfxr.vertices_to_bounds(vertices.isel(time=i), out_dims=out_dims)
+            for i in range(2)
+        ],
+        dim=vertices.time,
+    )
+    assert_equal(result, expected)
+    assert result.dtype == vertices.dtype
+    assert result.dims == ("time", *out_dims[1:], "bounds")
+    if chunked:
+        assert isinstance(result.data, DaskArray)
+        assert result.chunks[0] == (1, 1)
+
+
+def test_vertices_to_bounds_same_core_dimension_name():
+    vertices = xr.DataArray(np.arange(12).reshape(3, 4), dims=("time", "depth"))
+
+    result = cfxr.vertices_to_bounds(
+        vertices, out_dims=("bounds", "depth"), core_dims=["depth"]
+    )
+
+    assert result.sizes == {"time": 3, "depth": 3, "bounds": 2}
+    assert_array_equal(result[..., 0], vertices[:, :-1])
+    assert_array_equal(result[..., 1], vertices[:, 1:])
+
+
+def test_vertices_to_bounds_datetime_broadcast():
+    dates = np.arange("2000-01-01", "2000-01-05", dtype="datetime64[D]")
+    vertices = xr.DataArray(dates, dims="edge").expand_dims(time=[1, 2])
+
+    result = cfxr.vertices_to_bounds(
+        vertices, out_dims=("bounds", "depth"), core_dims=["edge"]
+    )
+
+    assert result.dtype == vertices.dtype
+    assert_array_equal(result[0, :, 0], dates[:-1])
+    assert_array_equal(result[0, :, 1], dates[1:])
+
+
+def test_vertices_to_bounds_invalid_dimensions():
+    vertices = xr.DataArray(np.ones((2, 3, 4)), dims=("a", "b", "c"))
+    with pytest.raises(ValueError, match="vertices format not understood"):
+        cfxr.vertices_to_bounds(vertices)
+    with pytest.raises(ValueError, match="out_dims"):
+        cfxr.vertices_to_bounds(vertices, out_dims=["bounds"], core_dims=["c"])
+
+
 @requires_cftime
 def test_bounds_to_vertices_cftime() -> None:
     import cftime