Wrap numpy-groupies to speed up Xarray's groupby aggregations #4473

Activity

1. shoyer commented on Sep 30, 2020

   

   shoyer

   on Sep 30, 2020

   MemberAuthor

   More actions

   A prototype implementation of the core functionality here can be found in: https://nbviewer.jupyter.org/gist/shoyer/6d6c82bbf383fb717cc8631869678737
2. max-sixty commented on Sep 30, 2020

   

   max-sixty

   on Sep 30, 2020

   Last edited by max-sixty

   Collaborator

   More actions

   This looks amazing! Thanks for finding it.

   Highly speculative, but would this also be a faster approach to stacking & unstacking? "Form 5 4" in the readme.
3. shoyer commented on Oct 1, 2020

   

   shoyer

   on Oct 1, 2020

   MemberAuthor

   More actions

   Highly speculative, but would this also be a faster approach to stacking & unstacking? "Form 5 4" in the readme.

   I'm not entirely sure, but I suspect something like the approach in #4184 might be more directly relevant for speeding up unstack (at least with NumPy arrays).
4. max-sixty commented on Oct 1, 2020

   

   max-sixty

   on Oct 1, 2020

   Collaborator

   More actions

   I'm not entirely sure, but I suspect something like the approach in #4184 might be more directly relevant for speeding up unstack (at least with NumPy arrays).

   Great. I need to think through how to do that — the approach of using MultiIndex codes to index the array directly is very elegant — I'll try applying it to stack / unstack as a project.
5. dcherian

   added

   topic-performance

   on Oct 9, 2020
6. max-sixty commented on Oct 19, 2020

   

   max-sixty

   on Oct 19, 2020

   Collaborator

   More actions

   Here's a very quick POC:

   from numpy_groupies.aggregate_numba import aggregate
   
   def npg_groupby(da: xr.DataArray, dim, func='sum'):
       group_idx, labels = pd.factorize(da.indexes[dim])
       axis = da.get_axis_num(dim)
       array = npg.aggregate(group_idx=group_idx, a=da, func=func, axis=axis)
       return array

   Run on this array:

   size_factor = 1000
   
   da = xr.DataArray(
       np.arange(1440 * size_factor).reshape(45 * size_factor, 8, 4),
       dims=("x", "y", "z"),
       coords=dict(x=list(range(45)) * size_factor, y=[1, 2, 3, 4] * 2, z=[1, 2] * 2),
   )

   It's about 2x as fast, though only generates the numpy array:

   %%timeit 
   npg_groupby(da, 'x')
   # 15 ms ± 130 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

   %%timeit
   da.groupby('x').sum()
   # 37.6 ms ± 244 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

   Any thoughts on any of:

   • What's the best way of reconstituting the coords etc, after npg produces the array?
   • Presumably we're going to have a fairly different design for this than the existing groupby operations — that design is very nested — wrapping functions and eventually calling .map to loop over each group in python.
   • Presumably we're going to need to keep the existing logic around for dask — is it reasonable for an initial version to defer to the existing logic for all dask arrays? (+ @shoyer 's thoughts above on this)
7. shoyer commented on Oct 19, 2020

   

   shoyer

   on Oct 19, 2020

   MemberAuthor

   More actions

   Something like the resample test case from #4498 might be a good example for finding 100x speed-ups. The main feature of that case is that there are a very large number of groups (only slightly fewer groups than original data points).
8. shoyer commented on Oct 19, 2020

   

   shoyer

   on Oct 19, 2020

   MemberAuthor

   More actions

   • What's the best way of reconstituting the coords etc, after npg produces the array?

   I think we can reuse the existing logic from the _combine method here:
   

   xarray/xarray/core/groupby.py

   Line 830 in 97e2625

   # assign coord when the applied function does not return that coord

   This just gives us an alternative way to calculate applied.

   • Presumably we're going to have a fairly different design for this than the existing groupby operations — that design is very nested — wrapping functions and eventually calling .map to loop over each group in python.

   Agreed. Hopefully this can live alongside in the GroupBy objects.

   • Presumably we're going to need to keep the existing logic around for dask — is it reasonable for an initial version to defer to the existing logic for all dask arrays? (+ @shoyer 's thoughts above on this)

   Yes, I agree that we should do this incrementally.
9. bmorris3 commented on Feb 21, 2022

   

   bmorris3

   on Feb 21, 2022

   More actions

   Hi @shoyer, thanks for this neat trick! What happens when bins is a sequence of bin edges, rather than a number of bins? Your example seems to break and I'm not sure how to fix it. Thanks again!
10. dcherian

    added 3 commits that reference this issue on May 13, 2022

    Enable `flox` in `GroupBy` and `resample` (pydata#5734)

    ...

    This issue will close once commit 9caa760 is merged into the 'main' branch.

    9caa760

    Enable `flox` in `GroupBy` and `resample` (pydata#5734)

    ...

    This issue will close once commit 7043bdb is merged into the 'main' branch.

    7043bdb

    Enable flox in GroupBy and resample (pydata#5734)

    ...

    This issue will close once commit e02b1c3 is merged into the 'main' branch.

    e02b1c3
11. dcherian

    closed this as completedin e02b1c3on May 15, 2022