How to groupby_bins using another variable?

[josuemtzmo]

Hello,

I'm trying to make more efficient the computation of water mass transform using the following code on a dataset with 12x6560x6540 (time,y,x). The extract below creates a mask to then multiply the mask (Heaviside function) to a given field (salinity component of the buoyancy flux) and sum it over time.

for i in range(len(isopycnal_bins)-1):
    bin_mask = rho.where(np.bitwise_and(rho <= isopycnal_bins[i+1], rho > isopycnal_bins[i])) * 0 + 1
    masked_transform = (buoy_salt_f * bin_mask).sum(dim = 'time')

According to me, it should be possible to do this same computation using groupby_bins with the code below:

buoy_salt_f.groupby_bins(rho,isopycnal_bins,labels = isopycnal_center).sum('time')

However I don't manage to make it work and I get the following error:

File ~/datawork/conda-env/jupyter-lab/lib/python3.9/site-packages/xarray/core/variable.py:870, in Variable.__getitem__(self, key)
    857 """Return a new Variable object whose contents are consistent with
    858 getting the provided key from the underlying data.
    859 
   (...)
    867 array `x.values` directly.
    868 """
    869 dims, indexer, new_order = self._broadcast_indexes(key)
--> 870 data = as_indexable(self._data)[indexer]
    871 if new_order:
    872     data = np.moveaxis(data, range(len(new_order)), new_order)

File ~/datawork/conda-env/jupyter-lab/lib/python3.9/site-packages/xarray/core/indexing.py:1374, in DaskIndexingAdapter.__getitem__(self, key)
   1372 key = key.tuple
   1373 try:
-> 1374     return self.array[key]
   1375 except NotImplementedError:
   1376     # manual orthogonal indexing.
   1377     # TODO: port this upstream into dask in a saner way.
   1378     value = self.array

File ~/datawork/conda-env/jupyter-lab/lib/python3.9/site-packages/dask/array/core.py:1978, in Array.__getitem__(self, index)
   1970     index = (index,)
   1972 from dask.array.slicing import (
   1973     normalize_index,
   1974     slice_with_bool_dask_array,
   1975     slice_with_int_dask_array,
   1976 )
-> 1978 index2 = normalize_index(index, self.shape)
   1979 dependencies = {self.name}
   1980 for i in index2:

File ~/datawork/conda-env/jupyter-lab/lib/python3.9/site-packages/dask/array/slicing.py:917, in normalize_index(idx, shape)
    915 for axis, (i, d) in enumerate(zip(idx, none_shape)):
    916     if d is not None:
--> 917         check_index(axis, i, d)
    918 idx = tuple(map(sanitize_index, idx))
    919 idx = tuple(map(normalize_slice, idx, none_shape))

File ~/datawork/conda-env/jupyter-lab/lib/python3.9/site-packages/dask/array/slicing.py:980, in check_index(axis, ind, dimension)
    975             raise IndexError(
    976                 f"Boolean array with size {ind.size} is not long enough "
    977                 f"for axis {axis} with size {dimension}"
    978             )
    979     elif (ind >= dimension).any() or (ind < -dimension).any():
--> 980         raise IndexError(
    981             f"Index is out of bounds for axis {axis} with size {dimension}"
    982         )
    983 elif isinstance(ind, slice):
    984     return

IndexError: Index is out of bounds for axis 0 with size 0

Without the sum over the time axis, the code seems to work as expected, returning a groupby_bins object:

group_data = buoy_salt_f.groupby_bins(rho,isopycnal_bins,labels = isopycnal_center)

returns:

DataArrayGroupBy, grouped over 'rhop_sig0_bins'
156 groups with labels 1.025e+03, ..., 1.031e+03.

however if I check the dims by group_data.dims, I get the same error as above.

Is this currently supported in xarray? or is there a more efficient way to do this computation? or am I doing something wrong?

[dcherian]

We'd need a simple example to debug. Going through groupby_bins right now will compute that rho array, so I don't think that's what you want. Use flox instead

[josuemtzmo]

Thank you so much, using flox solved the issue! Here it is the solution to the problem described above:

salt, density = xr.align( buoy_salt_f, rho)
binned_buoy_salt = xarray_reduce(
    salt,
    density,
    func="sum",
    expected_groups=isopycnal_bins,
    dim="time_counter"
)

[dcherian]

Great! Tell your friends!