Convert vertices to bounds for nD arrays

[jbusecke]

Thank you for developing this amazing package, I am only slowly appreciating its full amazingness.

I am currently trying to convert an array of layer depths that varies in time and space, which represent cell vertices in the vertical into a bound style array.

This is my original array lev_vertices:

For a minimal example you can create a similar array with:

lev_vertices = xr.DataArray(np.arange(9), dims='test').expand_dims(x=2,y=2, time=2)
lev_vertices

I would like to convert the dimension test into bounds. This works as expected when operating on a single profile:

cf_xarray.vertices_to_bounds(
    lev_bounds.isel(time=0,x=0, y=0).load(), out_dims=["bnds", "test"]
)

But I cant figure out how to make this work on the full array...

cf_xarray.vertices_to_bounds(
    lev_vertices.load(), out_dims=["bnds", "test"]
)

gives

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-76-4487305bbcdc> in <module>
----> 1 cf_xarray.vertices_to_bounds(
      2     lev_vertices.load(), out_dims=["bnds", "test"]
      3 )

/scratch/gpfs/jbusecke/conda_tigercpu/envs/aguadv_omz_busecke_2021/lib/python3.8/site-packages/cf_xarray/helpers.py in vertices_to_bounds(vertices, out_dims)
    116         )
    117     else:
--> 118         raise ValueError(
    119             f"vertices format not understood. Got {vertices.dims} with shape {vertices.shape}."
    120         )

ValueError: vertices format not understood. Got ('x', 'y', 'time', 'test') with shape (2, 2, 2, 9).

I am wondering if I am missing an option to ignore dimensions here? I was thinking of using xr.map_blocks, but that would impose restrictions on my chunking (which might affect performance). So I first wanted to check in here, and see if there is an easy fix.

[jbusecke]

This little snippet actually does exactly what I want:

lev_bounds = xr.concat([lev_vertices.isel(test=slice(0,-1)), lev_vertices.isel(test=slice(1,None))], dim='bnds')
lev_bounds

But I think it would be awesome if this could be replicated within cf-xarray.

EDIT: Thats not 100% correct. The actual coordinate values need to be changed to a common value...

[aulemahal]

I think the reason your use case is not implemented in cf_xarray is that it doesn't respect the cf conventions? Or at least, our implementation is limited to static horizontal grids.

The code could be adapted to a more general use case like yours, we'll have to see if it fits the scope of the package.

But overall you are right, in fact the current code does the same thing as you suggested, but assuming vertices is an array only defined over the dimensions to transform. In your case, you have x, y and time as additional dimensions.

bnd_vals = np.stack((vertices[:-1], vertices[1:]), axis=0)
return xr.DataArray(bnd_vals, dims=('bounds', 'test')])

[dcherian]

I think we can get this to work with apply_ufunc and using some numpy indexing tricks (i.e. using ...). We know that the dimension name "test" so that's the core dimension and will be the last axis when apply_ufunc passes it down. So bnd_vals = np.stack((vertices[..., :-1], vertices[..., 1:]), axis=0) should work?

PR welcome!