pydata / pydata/xarray

ds.weighted should skip non-numeric data_vars

Open
#9,322 0 comments 0 reactions 0 assignees View on GitHub

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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.