diff --git a/cf_xarray/helpers.py b/cf_xarray/helpers.py index 33304523..df2e8f7e 100644 --- a/cf_xarray/helpers.py +++ b/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( diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 4d28675e..27dae4c6 100644 --- a/doc/whats-new.rst +++ b/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) ======================