From ae4c3f726ac039a8c3678a8ba3adf9a603d7904a Mon Sep 17 00:00:00 2001 From: nawaaaaaAaar <198336470+nawaaaaaAaar@users.noreply.github.com> Date: Thu, 22 Oct 2026 14:31:05 +0000 Subject: [PATCH 1/2] Broadcast vertices-to-bounds conversion over extra dimensions --- cf_xarray/helpers.py | 65 +++++++++++++++++++++++--------- cf_xarray/tests/test_helpers.py | 66 +++++++++++++++++++++++++++++++++ 2 files changed, 114 insertions(+), 17 deletions(-) diff --git a/cf_xarray/helpers.py b/cf_xarray/helpers.py index 76febed9..7d446bb2 100644 --- a/cf_xarray/helpers.py +++ b/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, + 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, + ) diff --git a/cf_xarray/tests/test_helpers.py b/cf_xarray/tests/test_helpers.py index e822063b..5ac515ab 100644 --- a/cf_xarray/tests/test_helpers.py +++ b/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 From 807839c0eb41ef20d2d7b104972a234783da6f00 Mon Sep 17 00:00:00 2001 From: nawaaaaaAaar <198336470+nawaaaaaAaar@users.noreply.github.com> Date: Thu, 22 Oct 2026 14:40:09 +0000 Subject: [PATCH 2/2] Describe existing ndarray input support in type annotation --- cf_xarray/helpers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cf_xarray/helpers.py b/cf_xarray/helpers.py index 7d446bb2..9196fb1a 100644 --- a/cf_xarray/helpers.py +++ b/cf_xarray/helpers.py @@ -482,7 +482,7 @@ def _sort_vertices(vertices: np.ndarray, order: str) -> np.ndarray: def vertices_to_bounds( - vertices: DataArray, + vertices: DataArray | np.ndarray, out_dims: Sequence[str] = ("bounds", "x", "y"), core_dims: Sequence[Hashable] | None = None, ) -> DataArray: