Simple groupby operations in xarray

[JulienBrn]

I have been using xarray for a while and all my troubles seem to be around how groupby operations work, especially with a multiindex.
Let us take a simple common example : a dataset with a multiindex and one wishes to groupby some variable to compute the mean. Then one wishes to substract the mean of each group to the variable.

[Edit] Added code that works and updated issues status

Example code

Note: the example has been changed to reflect a real use case where we want to compute the luminosity of each frame normalized by the average luminosity of its video. To justify the use of a multiindex, the videos have different number of frames.

import pandas as pd, xarray as xr, numpy as np

xr.set_options(use_flox=False)

d = xr.Dataset()
d["video"] = xr.DataArray(["vid_1"]*100 + ["vid_2"]*10 + ["vid_3"] * 50, dims="vidframe")
d["frame"] = xr.DataArray(list(range(100)) + list(range(10)) + list(range(50)), dims="vidframe")
d["luminosity"] = xr.DataArray(np.random.random(d.sizes["vidframe"]), dims="vidframe")

d = d.set_index(vidframe = ["video", "frame"])
print(d)

m = d.groupby("video").mean()
print(m)

res = d - m
print(res)

Encountered issues

There are several issues with this code:

1. [Unsolved] If I do not specify use_flox=False, then d.groupby("video").mean() fails.
2. [Solved] res = d - m fails because m is an array of size 10 because of the grouping and it seems that xarray does not automatically call pandas merge on the multiindex beforehand... [Solution] one is expected to do d.groupby("video") - m (see code that works section and dcherian's reply)
3. [Bug?] With d.set_index(vidframe = ["video", "frame"]) , d.groupby("video") - m fails. Solution is to not use a multiindex
4. [Not addressed] Usually I like to put many related things in a single dataset (see code that works section). However, it is extremely tedious to spend time renaming dimensions... Any other way to do this nicely?
5. [Not addressed] I have no clue why the sparse solution fails (see end of post)

Code that works

import pandas as pd, xarray as xr, numpy as np
xr.set_options(use_flox=False)

d = xr.Dataset()
d["video"] = xr.DataArray(["vid_1"]*100 + ["vid_2"]*10 + ["vid_3"] * 50, dims="vidframe")
d["frame"] = xr.DataArray(list(range(100)) + list(range(10)) + list(range(50)), dims="vidframe")
d["luminosity"] = xr.DataArray(np.random.random(d.sizes["vidframe"]), dims="vidframe")
d = d.set_coords(["video", "frame"])
print(d)

d["video_luminosity"] = d["luminosity"].groupby("video").mean().rename(dict(video="video_dim"))

print(d)
d["relative_luminosty"] = d["luminosity"].groupby("video") - d["video_luminosity"].rename(dict(video_dim="video"))
print(d)

Sparse attempted solution

Another solution is simply promote each level of a multiindex to a dimension and use a sparse array backend to avoid memory blowup. This has its own issues, but it using it provides the most readable code.

[EDIT] I actually tried with the following code and I get errors... Is xarray fully integrated with sparse or am I doing a mistake?
[EDIT] Works with the correct code (i.e d["luminosity"].mean("frame")), however, fails with d["luminosity"].mean("video"), but I do not see why it should fail.

import pandas as pd, numpy as np, xarray as xr

df = pd.DataFrame()
df["video"] = ["vid_1"]*100 + ["vid_2"]*10 + ["vid_3"] * 50
df["frame"] = list(range(100)) + list(range(10)) + list(range(50))
df["luminosity"] = np.random.random(len(df.index))
df = df.set_index(["video", "frame"])
d = xr.Dataset.from_dataframe(df, sparse=True)

print(d)
d["video_luminosity"] = d["luminosity"].mean("frame")
print(d)
d["relative_luminosity"] = d["luminosity"] - d["video_luminosity"]
print(d)

[dcherian]

I think (1) is fixed if you upgrade flox.

Re (2): First to do the arithmetic you'd need (https://docs.xarray.dev/en/stable/user-guide/groupby.html#grouped-arithmetic)

res = d.groupby('x') - m

But that's still a bug (#8952)

How am I supposed to make this work in xarray ?

Do you really need the MultiIndex? It works just fine if I comment out the set_index line.

[JulienBrn]

1. Just upgraded flox to version flox-0.9.6 and I still get this error

 xarray/core/groupby.py", line 1000, in _flox_reduce
     raise ValueError(f"cannot reduce over dimensions {grouper.name!r}")

2. Maybe having a toy example is not a good idea, let us consider the case where we have several videos with a variable number of frames. The index is a multiindex of (video, frame). Let us say the value variable (here v) is the luminosity of each frame. Let us say we want to compute the luminosity of each frame minus the luminosity of the average frame in that video. As I understand this is a good use case of MultiIndex.

[EDIT] Example changed to reflect this use case. The new videodimension represents dimension x
[EDIT] Link clearly suggests this is how I should do things, yet it does not work