Skip to content

Implement weighted groupby #5480

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

Conversation

dcherian
Copy link
Contributor

@dcherian dcherian commented Jun 17, 2021

Initial proof-of-concept. Suggestions to improve this are very welcome.

Here's some convenient testing code

import xarray as xr

ds = xr.tutorial.open_dataset('rasm').load()
month_length = ds.time.dt.days_in_month
weights = month_length.groupby('time.season') / month_length.groupby('time.season').sum()

actual = ds.weighted(month_length).groupby("time.season").mean()
expected = (ds * weights).groupby('time.season').sum(skipna=False)
xr.testing.assert_allclose(actual, expected)

I've added info to the repr

ds.weighted(month_length).groupby("time.season")
WeightedDatasetGroupBy, grouped over 'season'
4 groups with labels 'DJF', 'JJA', 'MAM', 'SON'.
weighted along dimensions: time by 'days_in_month'

@dcherian dcherian force-pushed the weighted-groupby branch 2 times, most recently from f6c41ae to 820cab9 Compare June 17, 2021 03:03
return combined


class DatasetGroupBy(DatasetGroupByCombine, DatasetGroupbyArithmetic):
Copy link
Contributor Author

@dcherian dcherian Jun 25, 2021

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pulled out combine because the inheritance order is DatasetGroupBy → DatasetGroupByArithmetic → ImplementsDatasetReduce → IncludeReduceMethods . That last one overwrites the reductions on WeighedGroupBy if it inherits from DatasetGroupBy

@dcherian dcherian changed the title Implement groupby.weighted Implement weighted groupby Jun 25, 2021
@dcherian
Copy link
Contributor Author

This attempt subclasses GroupBy and works as .weighted().groupby()

@shoyer let me know what you think

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants