ds.weighted should skip non-numeric data_vars
Open
Nobody has claimed this yet.
bug
- Dominant language
- Python
- Stars
- 4.2k
- Forks
- 1.4k
- Avg merge
- 2d 15h
- Merged PRs (30d)
- 14
Description
What happened?
weighted reductions for non-numeric data_vars errors in contrast to e.g. ds.mean() where the variable is skipped
What did you expect to happen?
skip the variable
Minimal Complete Verifiable Example
import xarray as xr
ds = xr.Dataset(data_vars={"a": ("x", ["a"])})
w = xr.DataArray([1], dims="x")
ds.weighted(w).mean()
MVCE confirmation
- Minimal example — the example is as focused as reasonably possible to demonstrate the underlying issue in xarray.
- Complete example — the example is self-contained, including all data and the text of any traceback.
- Verifiable example — the example copy & pastes into an IPython prompt or Binder notebook, returning the result.
- New issue — a search of GitHub Issues suggests this is not a duplicate.
- Recent environment — the issue occurs with the latest version of xarray and its dependencies.
Relevant log output
---------------------------------------------------------------------------
TypeError Traceback (most recent call last)
Cell In[10], line 1
----> 1 ds.weighted(w).mean()
File ~/code/xarray/xarray/util/deprecation_helpers.py:115, in _deprecate_positional_args.<locals>._decorator.<locals>.inner(*args, **kwargs)
111 kwargs.update({name: arg for name, arg in zip_args})
113 return func(*args[:-n_extra_args], **kwargs)
--> 115 return func(*args, **kwargs)
File ~/code/xarray/xarray/core/weighted.py:498, in Weighted.mean(self, dim, skipna, keep_attrs)
490 @_deprecate_positional_args("v2023.10.0")
491 def mean(
492 self,
(...)
496 keep_attrs: bool | None = None,
497 ) -> T_Xarray:
--> 498 return self._implementation(
499 self._weighted_mean, dim=dim, skipna=skipna, keep_attrs=keep_attrs
500 )
File ~/code/xarray/xarray/core/weighted.py:559, in DatasetWeighted._implementation(self, func, dim, **kwargs)
556 def _implementation(self, func, dim, **kwargs) -> Dataset:
557 self._check_dim(dim)
--> 559 return self.obj.map(func, dim=dim, **kwargs)
File ~/code/xarray/xarray/core/dataset.py:7073, in Dataset.map(self, func, keep_attrs, args, numeric_only, **kwargs)
7070 for name, var in self.data_vars.items():
7071 if not numeric_only or np.issubdtype(var.dtype, np.number) or (var.dtype == np.bool_):
-> 7073 data_vars[name] = maybe_wrap_array(var, func(var, *args, **kwargs))
7075 attrs = self.attrs if keep_attrs else None
7077 out = type(self)(data_vars, attrs=attrs)
File ~/code/xarray/xarray/core/weighted.py:287, in Weighted._weighted_mean(self, da, dim, skipna)
279 def _weighted_mean(
280 self,
281 da: T_DataArray,
282 dim: Dims = None,
283 skipna: bool | None = None,
284 ) -> T_DataArray:
285 """Reduce a DataArray by a weighted ``mean`` along some dimension(s)."""
--> 287 weighted_sum = self._weighted_sum(da, dim=dim, skipna=skipna)
289 sum_of_weights = self._sum_of_weights(da, dim=dim)
291 return weighted_sum / sum_of_weights
File ~/code/xarray/xarray/core/weighted.py:277, in Weighted._weighted_sum(self, da, dim, skipna)
269 def _weighted_sum(
270 self,
271 da: T_DataArray,
272 dim: Dims = None,
273 skipna: bool | None = None,
274 ) -> T_DataArray:
275 """Reduce a DataArray by a weighted ``sum`` along some dimension(s)."""
--> 277 return self._reduce(da, self.weights, dim=dim, skipna=skipna)
File ~/code/xarray/xarray/core/weighted.py:232, in Weighted._reduce(da, weights, dim, skipna)
228 da = da.fillna(0.0)
230 # `dot` does not broadcast arrays, so this avoids creating a large
231 # DataArray (if `weights` has additional dimensions)
--> 232 return dot(da, weights, dim=dim)
File ~/code/xarray/xarray/util/deprecation_helpers.py:140, in deprecate_dims.<locals>.wrapper(*args, **kwargs)
132 emit_user_level_warning(
133 f"The `{old_name}` argument has been renamed to `dim`, and will be removed "
134 "in the future. This renaming is taking place throughout xarray over the "
(...)
137 PendingDeprecationWarning,
138 )
139 kwargs["dim"] = kwargs.pop(old_name)
--> 140 return func(*args, **kwargs)
File ~/code/xarray/xarray/core/computation.py:1870, in dot(dim, *arrays, **kwargs)
1867 # subscripts should be passed to np.einsum as arg, not as kwargs. We need
1868 # to construct a partial function for apply_ufunc to work.
1869 func = functools.partial(duck_array_ops.einsum, subscripts, **kwargs)
-> 1870 result = apply_ufunc(
1871 func,
1872 *arrays,
1873 input_core_dims=input_core_dims,
1874 output_core_dims=output_core_dims,
1875 join=join,
1876 dask="allowed",
1877 )
1878 return result.transpose(*all_dims, missing_dims="ignore")
File ~/code/xarray/xarray/core/computation.py:1268, in apply_ufunc(func, input_core_dims, output_core_dims, exclude_dims, vectorize, join, dataset_join, dataset_fill_value, keep_attrs, kwargs, dask, output_dtypes, output_sizes, meta, dask_gufunc_kwargs, on_missing_core_dim, *args)
1266 # feed DataArray apply_variable_ufunc through apply_dataarray_vfunc
1267 elif any(isinstance(a, DataArray) for a in args):
-> 1268 return apply_dataarray_vfunc(
1269 variables_vfunc,
1270 *args,
1271 signature=signature,
1272 join=join,
1273 exclude_dims=exclude_dims,
1274 keep_attrs=keep_attrs,
1275 )
1276 # feed Variables directly through apply_variable_ufunc
1277 elif any(isinstance(a, Variable) for a in args):
File ~/code/xarray/xarray/core/computation.py:312, in apply_dataarray_vfunc(func, signature, join, exclude_dims, keep_attrs, *args)
307 result_coords, result_indexes = build_output_coords_and_indexes(
308 args, signature, exclude_dims, combine_attrs=keep_attrs
309 )
311 data_vars = [getattr(a, "variable", a) for a in args]
--> 312 result_var = func(*data_vars)
314 out: tuple[DataArray, ...] | DataArray
315 if signature.num_outputs > 1:
File ~/code/xarray/xarray/core/computation.py:821, in apply_variable_ufunc(func, signature, exclude_dims, dask, output_dtypes, vectorize, keep_attrs, dask_gufunc_kwargs, *args)
816 if vectorize:
817 func = _vectorize(
818 func, signature, output_dtypes=output_dtypes, exclude_dims=exclude_dims
819 )
--> 821 result_data = func(*input_data)
823 if signature.num_outputs == 1:
824 result_data = (result_data,)
File ~/code/xarray/xarray/core/duck_array_ops.py:88, in einsum(*args, **kwargs)
86 return opt_einsum.contract(*args, **kwargs)
87 else:
---> 88 return np.einsum(*args, **kwargs)
File ~/.conda/envs/xarray_dev/lib/python3.10/site-packages/numpy/core/einsumfunc.py:1371, in einsum(out, optimize, *operands, **kwargs)
1369 if specified_out:
1370 kwargs['out'] = out
-> 1371 return c_einsum(*operands, **kwargs)
1373 # Check the kwargs to avoid a more cryptic error later, without having to
1374 # repeat default values here
1375 valid_einsum_kwargs = ['dtype', 'order', 'casting']
TypeError: invalid data type for einsum
Anything else we need to know?
Derived from https://github.com/pydata/xarray/issues/7027#issuecomment-2270870147, cc @JdeJong96
Compare the result of ds.mean()
Environment
main...
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with xarray/core/weighted.py, especially DatasetWeighted._implementation and Weighted._weighted_mean, then reproduce the MVCE to trace the failure through Dataset.map and dot. Done means ds.weighted(w).mean() skips the non-numeric variable, matching ds.mean(), without breaking weighted reductions for numeric variables.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- numpy, python
- Domain
- data
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Stale
- Clarity
- Clearly specified
- Newbie friendliness
- 52/100