UgridDataArray.weighted does not preserve grid

[veenstrajelmer]

When using UgridDataArray.weighted, the grid property is lost. It has to be manually re-added. Same for UgridDataset.weighted.

import xugrid as xu
import xarray as xr
import numpy as np
import dfm_tools as dfmt

file_nc = dfmt.data.fm_curvedbend_map(return_filepath=True)
uds = xu.open_dataset(file_nc, chunks={'time':1}) 

weights = xr.DataArray(np.diff(uds['mesh2d_interface_sigma']), dims='mesh2d_nLayers')

uda = uds['mesh2d_sa1']

# # Select / integrate depth
uda = uda.weighted(weights).mean(dim='mesh2d_nLayers',keep_attrs=True)  # Depth integration
# uda.grid #fails, weignted converted it back to a xr.Dataset
uda_restored = xu.UgridDataArray(uda, grid=uds.grid) #UgridDataset expects `grids=[uds.grid]`
uda_restored.grid

[Huite]

This means the wrapping code is failing: https://github.com/Deltares/xugrid/blob/main/xugrid/core/wrap.py#L59

To explain, any access of a method or attribute on a UgridDataArray or UgridDataset is called within a wrapped context, which checks the result type and reattaches the topology if needed (depends on the resulting dims).

It's failing here because the return type of weighted is not a DataArray or Dataset, but rather a DataArrayWeighted object: https://docs.xarray.dev/en/stable/generated/xarray.DataArray.weighted.html

Now that you report this: I'm guessing the same problem occurs with .groupby, since that should return a DataArrayGroupby object. I haven't checked carefully how many of these objects there are, and which complexity they have. A quick search in the xarray source seems to suggest:

• DataArrayGroupBy
• DataArrayRolling
• DataArrayCoarsen
• DataArrayResample
• DataArrayWeighted

(and probably simlarly so for Dataset.)

These are all quite similar in setup. Since they all have the reduction methods defined on them, it may be best to write more wrapping code for them as well, including the grid topologies along the way in a composite object (just like UgridDataArray).
This means expanding the set of type checks in the wrapping code. We can make a dict of mappings, I suppose.
Then the wrapped classes (e.g. UgridDataArrayGroupby) need their own wrapped methods, ensuring they return a UgridDataArray or UgridDataset as needed.

This would also be fixed by xarray's explicit indexes, since if we can store the UGRID state in an ordinary index, it will be propagated to the GroupBy and Weighted objects as well.

[Huite]

Pragmatically: it might be easiest to just overload the methods for now.

[hcwinsemius]

I can confirm that .groupby also looses the grid attribute. This may be more complex as each group should receive only a part of the grid instead of the full.