diff --git a/uxarray/core/dataarray.py b/uxarray/core/dataarray.py old mode 100644 new mode 100755 index 693b436de..5f5833a3f --- a/uxarray/core/dataarray.py +++ b/uxarray/core/dataarray.py @@ -412,14 +412,17 @@ def nodal_average(self): # weighted_mean() # .weighted.mean() + # weights are only on grid - not supporting custom weights for now def mean(self, dim=None, *, skipna=None, keep_attrs=None, weighted=False, **kwargs): if weighted: if self._face_centered(): + grid_dim = "n_face" # use face areas as weight - weights = self.uxgrid.face_areas + weights = da.from_array(self.uxgrid.face_areas.values) elif self._edge_centered(): + grid_dim = "n_edge" # use edge magnitude as weight - weights = self.uxgrid.edge_node_distances + weights = da.from_array(self.uxgrid.edge_node_distances.values) else: # apply regular Xarray mean warnings.warn( @@ -428,13 +431,13 @@ def mean(self, dim=None, *, skipna=None, keep_attrs=None, weighted=False, **kwar "variables. Performing an unweighted mean." ) - return super().mean(dim=None, skipna=None, keep_attrs=None, **kwargs) + return super().mean(dim=dim, skipna=skipna, keep_attrs=keep_attrs, **kwargs) # compute the total weight total_weight = weights.sum() - # compute weighted mean - weighted_mean = (self * weights).sum() / total_weight + # compute weighted mean #assumption on index of dimension (last one is geometry) + weighted_mean = (self * weights).sum(dim=grid_dim) / total_weight # create a UxDataArray and return it return UxDataArray(weighted_mean, uxgrid=self.uxgrid)