Skip to content

Commit

Permalink
Adjust after pydata#8183
Browse files Browse the repository at this point in the history
  • Loading branch information
max-sixty committed Sep 17, 2023
1 parent 374ac99 commit 602d043
Showing 1 changed file with 9 additions and 36 deletions.
45 changes: 9 additions & 36 deletions xarray/core/rolling_exp.py
Original file line number Diff line number Diff line change
Expand Up @@ -131,42 +131,15 @@ def mean(self, keep_attrs: bool | None = None) -> T_DataWithCoords:

dim_order = self.obj.dims

from xarray.core.dataset import Dataset

if isinstance(self.obj, Dataset):
new_ds = self.obj.copy()
if not keep_attrs:
new_ds.attrs = {}
for name, var in self.obj.variables.items():
if self.dim in self.obj[name].dims:
new_ds[name] = apply_ufunc(
move_exp_nanmean,
var,
input_core_dims=[[self.dim]],
kwargs=dict(alpha=self.alpha, axis=-1),
output_core_dims=[[self.dim]],
# I'm not sure whether this line should be here (can look later,
# docs aren't _that_ clear vs. `input_core_dims`, or I'm being
# slow at understanding them.)
exclude_dims={self.dim},
keep_attrs=keep_attrs,
)
# Currently we test for attrs being removed on all arrays, even those
# without the dim, not sure if we should be doing this or not (we can
# indent a line and then only remove attrs on arrays we compute on)
if not keep_attrs:
new_ds[name].attrs = {}
return new_ds.transpose(*dim_order)
else:
return apply_ufunc(
move_exp_nanmean,
self.obj,
input_core_dims=[[self.dim]],
kwargs=dict(alpha=self.alpha, axis=-1),
output_core_dims=[[self.dim]],
exclude_dims={self.dim},
keep_attrs=keep_attrs,
).transpose(*dim_order)
return apply_ufunc(
move_exp_nanmean,
self.obj,
input_core_dims=[[self.dim]],
kwargs=dict(alpha=self.alpha, axis=-1),
output_core_dims=[[self.dim]],
exclude_dims={self.dim},
keep_attrs=keep_attrs,
).transpose(*dim_order)

def sum(self, keep_attrs: bool | None = None) -> T_DataWithCoords:
"""
Expand Down

0 comments on commit 602d043

Please sign in to comment.