Deltares / xugrid

Xarray and unstructured grids
https://deltares.github.io/xugrid/
MIT License
58 stars 8 forks source link

xugrid types are not preserved consistently #34

Open Huite opened 1 year ago

Huite commented 1 year ago

Xarray isn't quite commutative in operations, but largely so:

da + 1  # -> DataArray
1 + da  # -> DataArray

This isn't true for broadcasting, the first argument's dimensions go first.

However, for UgridDataArrays, it works rather poorly:

uda + da  # -> UgridDataArray
da + uda  # -> DataArray

This also came up, even worse(?):

xr.DataArray.notnull(uda)  # -> np.array

If I'm not mistaken, xarray does the right thing in general, e.g. calling a numpy operation will return a DataArray (IIRC). We should try the same approach to preserve the xugrid type.

JoerivanEngelen commented 1 year ago

Just for completeness and to avoid discouraging potential users:

Whereas this returns a numpy array:

xr.DataArray.notnull(uda)  # -> np.array

This does the right thing:

uda.notnull()  # ->  UgridDataArray
Huite commented 1 year ago
  1. xr.DataArray.notnull(uda) seems quite obscure, generally only uda.notnull() will be called.

  2. The xarray top level functions should be explicitly wrapped, I've made a start with ones_like, zeros_like, etc.

  3. Unfortunately, there doesn't seem to be a good way to address this: da + uda # -> DataArray

While numpy has a procedure to determine operator precedence (and e.g. defer to xarray), xarray doesn't implement this procedure (AFAICT). So the line above just calls the __add__ method of da. Once we're in xarray, there's no way to re-attach a UGRID topology.

I guess we'll have to live with this: arithmetic is already non-commutative due to xarray's broadcasting rules. For broadcasting a grid to temporal data, it is a little cumbersome:

out = UgridDataArray(time * uda, grid=uda.ugrid.grid)

# or:

out = (uda * time).transpose("time", face_dimension)

Alternatively, we could allow promotion with a None grid, then modify the binary_ops to look at the grid of other. But it complicates things with maker things much more ergonomic.

Best to add a clear warning in the documentation instead.