xarray groupby multiple variables

[miguelcarcamov]

Hi there,

I have the following xarray dataset:

xarr_ds = xr.Dataset(
                data_vars=dict(
                    U=(["row"], uvw_grid[:,0]),
                    V=(["row"], uvw_grid[:,1]),
                    W=(["row"], uvw_grid[:,2]),
                    DATA=(["row", "chan", "corr"], data),
                    WEIGHT=(["row", "corr"], weight),
                    FLAG=(["row", "chan", "corr"], flag)
                ),
                coords=dict(
                     ROWID=(["row"], row_id)
                )
            )

I would like to do something like in xarr_ds.groupby(["U", "V"]) - like in dask. However, this is not possible in dask because the cells cannot be a matrix. Also, this is not possible in xarray as well since one can only groupby by only one variable. Is there a workaround for this? Has anyone done a multiple variable groupby on a xarray dataset?

This is what I am trying to do:

import xarray as xr
import dask.array as da
import numpy as np
nrows = 1000
nchan = 16
ncorr = 2
N = 256

_id = da.arange(0, nrows)
uvw_grid = da.random.uniform(size=(nrows, 3)).astype(np.int32) + N//2
data = da.random.uniform(size=(nrows, nchan, ncorr)).astype(np.float32) + 1j*da.random.uniform(size=(nrows, nchan, ncorr)).astype(np.float32)
weight = da.random.uniform(size=(nrows, ncorr))


xarr_ds = xr.Dataset(
                data_vars=dict(
                    U=(["row"], uvw_grid[:,0]),
                    V=(["row"], uvw_grid[:,1]),
                    W=(["row"], uvw_grid[:,2]),
                    DATA=(["row", "chan", "corr"], data),
                    WEIGHT=(["row","corr"], weight),
                ),
                coords=dict(
                     ROWID=(["row"], _id)
                )
            )

xarr_ds


xarr_ds.groupby(["U", "V"])

Cheers

[max-sixty]

This is something we'd really like to improve, and some efforts (e.g. numpy_groupies) will move us closer.

There are some approaches that will allow for multiple-dimension groupbys now (edit: much of the time xarray handles this already, and doesn't require manually stacking-groupby-unstacking) — and to respond quickly rather than completely — the approach I normally take is to stack the relevant dimensions, run the groupby operation, and then unstack.

@miguelcarcamov could you make your example copy-pastable per #5404?

[miguelcarcamov]

I would need to swap the Z dim with row dim right? I think is not possible. If I do that I get :ValueError: replacement dimension 'Z' is not a 1D variable along the old dimension 'row'

[max-sixty]

I need to think more about whether that's the correct behavior, but for the moment, setting the dims to be one of the indexes (i.e. adding .swap_dims(row='U') makes it work, I think:

In [25]: ds.swap_dims(row='U').set_index(Z=['U','V']).groupby('Z').mean()
Out[25]:
<xarray.Dataset>
Dimensions:    (Z: 1, chan: 16, corr: 2)
Coordinates:
  * Z          (Z) MultiIndex
  - Z_level_0  (Z) int64 128
  - Z_level_1  (Z) int64 128
Dimensions without coordinates: chan, corr
Data variables:
    W          (Z) float64 128.0
    DATA       (Z, chan, corr) complex64 (0.4959001+0.5078582j) ... (0.498601...
    WEIGHT     (Z, corr) float64 0.5128 0.497

Whereas excluding that retains the row dim, which isn't what we want:

In [26]: ds.set_index(Z=['U','V']).groupby('Z').mean()
Out[26]:
<xarray.Dataset>
Dimensions:    (Z: 1, chan: 16, corr: 2, row: 1000)
Coordinates:
    ROWID      (row) int64 0 1 2 3 4 5 6 7 8 ... 992 993 994 995 996 997 998 999
  * Z          (Z) MultiIndex
  - Z_level_0  (Z) int64 128
  - Z_level_1  (Z) int64 128
Dimensions without coordinates: chan, corr, row
Data variables:
    W          (Z, row) int32 128 128 128 128 128 128 ... 128 128 128 128 128
    DATA       (Z, row, chan, corr) complex64 (0.33360422+0.23278244j) ... (0...
    WEIGHT     (Z, row, corr) float64 0.5932 0.8688 0.02971 ... 0.5938 0.9018

[miguelcarcamov]

Right! The first code works beautifully, although I think that to avoid the swap_dims it would be better to create the dataset using the coordinate "U" from the beginning. That's what I did and it works! Thank you!

[miguelcarcamov]

I have another question kind of related with this, but also not related. When creating a dataset with xarray arrays with dask arrays inside then when doing calculation there seems to be no dask array in the variables. However, if I force the dataset to use dask arrays using .data the variables are indeed dask arrays. I was wondering if xarray chooses to use dask array or numpy arrays depending on the memory of the machine or if this is done automatically.

Cheers

[max-sixty]

@miguelcarcamov could you open a new discussion with a repro?

[miguelcarcamov]

Yep @max-sixty, I have been checking that stack is probably the solution, although I don't get how to do it on my case. I have added the code to my question. Another question is:

It will work if I do UV=(["row", "uv], uvw_grid[:,0:2]), instead of separating each coordinate of uvw_grid?

[miguelcarcamov]

UPDATE:
I have tried to move U V to coordinates with strings 'u','v'. Then stacking:
stacked = xarr_ds.stack(z=["u", "v"])

However, this takes a lot of time and uses a lot of RAM (more than 8GB) since the line cannot end because it fills the whole memory.

[dcherian]

Now possible with flox: https://flox.readthedocs.io/en/latest/intro.html#grouping-by-multiple-variables