I'm storing Pyomo objects in xarray arrays to allow for vectorised operations with named arrays. E.g., I might have an array of Pyomo variables indexed over the costs dimension and I want to sum them. Before v2024.6.0 this was possible without any issues. Now, if I include the min_count argument in my summation then it fails when trying to apply numpy type promotion to the Pyomo object.
What did you expect to happen?
The xarray.core.array_api_compat.result_type method should fall back to numpy.object_ when objects are present.
Minimal Complete Verifiable Example
import pyomo.kernel as pk
import xarray as xr
da = xr.DataArray([np.nan, pk.variable(lb=0, ub=10), np.nan], dims=("costs",))
try:
da.sum("costs")
print("Success without min_count")
except TypeError as e:
print(f"Failure without min_count: {e}")
try:
da.sum("costs", min_count=1)
print("Success with min_count")
except TypeError as e:
print(f"Failure with min_count: {e}")
try:
da.expand_dims(foo=["bar"]).sum("costs", min_count=1)
print("Success with extra dim and min_count")
except TypeError as e:
print(f"Failure with extra dim and min_count: {e}")
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
File ~/miniforge3/envs/calliope/lib/python3.12/site-packages/xarray/core/_aggregations.py:3175, in DataArrayAggregations.sum(self, dim, skipna, min_count, keep_attrs, **kwargs)
3087 def sum(
3088 self,
3089 dim: Dims = None,
(...)
3094 **kwargs: Any,
3095 ) -> Self:
3096 """
3097 Reduce this DataArray's data by applying ``sum`` along some dimension(s).
3098
(...)
3173 array(8.)
3174 """
-> 3175 return self.reduce(
3176 duck_array_ops.sum,
3177 dim=dim,
3178 skipna=skipna,
3179 min_count=min_count,
3180 keep_attrs=keep_attrs,
3181 **kwargs,
3182 )
File ~/miniforge3/envs/calliope/lib/python3.12/site-packages/xarray/core/dataarray.py:3839, in DataArray.reduce(self, func, dim, axis, keep_attrs, keepdims, **kwargs)
3795 def reduce(
3796 self,
3797 func: Callable[..., Any],
(...)
3803 **kwargs: Any,
3804 ) -> Self:
3805 """Reduce this array by applying `func` along some dimension(s).
3806
3807 Parameters
(...)
3836 summarized data and the indicated dimension(s) removed.
3837 """
-> 3839 var = self.variable.reduce(func, dim, axis, keep_attrs, keepdims, **kwargs)
3840 return self._replace_maybe_drop_dims(var)
File ~/miniforge3/envs/calliope/lib/python3.12/site-packages/xarray/core/variable.py:1677, in Variable.reduce(self, func, dim, axis, keep_attrs, keepdims, **kwargs)
1670 keep_attrs_ = (
1671 _get_keep_attrs(default=False) if keep_attrs is None else keep_attrs
1672 )
1674 # Note that the call order for Variable.mean is
1675 # Variable.mean -> NamedArray.mean -> Variable.reduce
1676 # -> NamedArray.reduce
-> 1677 result = super().reduce(
1678 func=func, dim=dim, axis=axis, keepdims=keepdims, **kwargs
1679 )
1681 # return Variable always to support IndexVariable
1682 return Variable(
1683 result.dims, result._data, attrs=result._attrs if keep_attrs_ else None
1684 )
File ~/miniforge3/envs/calliope/lib/python3.12/site-packages/xarray/namedarray/core.py:916, in NamedArray.reduce(self, func, dim, axis, keepdims, **kwargs)
912 if isinstance(axis, tuple) and len(axis) == 1:
913 # unpack axis for the benefit of functions
914 # like np.argmin which can't handle tuple arguments
915 axis = axis[0]
--> 916 data = func(self.data, axis=axis, **kwargs)
917 else:
918 data = func(self.data, **kwargs)
File ~/miniforge3/envs/calliope/lib/python3.12/site-packages/xarray/core/duck_array_ops.py:447, in _create_nan_agg_method.<locals>.f(values, axis, skipna, **kwargs)
445 with warnings.catch_warnings():
446 warnings.filterwarnings("ignore", "All-NaN slice encountered")
--> 447 return func(values, axis=axis, **kwargs)
448 except AttributeError:
449 if not is_duck_dask_array(values):
File ~/miniforge3/envs/calliope/lib/python3.12/site-packages/xarray/core/nanops.py:101, in nansum(a, axis, dtype, out, min_count)
99 result = sum_where(a, axis=axis, dtype=dtype, where=mask)
100 if min_count is not None:
--> 101 return _maybe_null_out(result, axis, mask, min_count)
102 else:
103 return result
File ~/miniforge3/envs/calliope/lib/python3.12/site-packages/xarray/core/nanops.py:34, in _maybe_null_out(result, axis, mask, min_count)
32 elif getattr(result, "dtype", None) not in dtypes.NAT_TYPES:
33 null_mask = mask.size - duck_array_ops.sum(mask)
---> 34 result = where(null_mask < min_count, np.nan, result)
36 return result
File ~/miniforge3/envs/calliope/lib/python3.12/site-packages/xarray/core/duck_array_ops.py:354, in where(condition, x, y)
352 """Three argument where() with better dtype promotion rules."""
353 xp = get_array_namespace(condition)
--> 354 return xp.where(condition, *as_shared_dtype([x, y], xp=xp))
File ~/miniforge3/envs/calliope/lib/python3.12/site-packages/xarray/core/duck_array_ops.py:258, in as_shared_dtype(scalars_or_arrays, xp)
252 xp = get_array_namespace(scalars_or_arrays)
254 # Pass arrays directly instead of dtypes to result_type so scalars
255 # get handled properly.
256 # Note that result_type() safely gets the dtype from dask arrays without
257 # evaluating them.
--> 258 dtype = dtypes.result_type(*scalars_or_arrays, xp=xp)
260 return [asarray(x, dtype=dtype, xp=xp) for x in scalars_or_arrays]
File ~/miniforge3/envs/calliope/lib/python3.12/site-packages/xarray/core/dtypes.py:266, in result_type(xp, *arrays_and_dtypes)
262 if xp is None:
263 xp = get_array_namespace(arrays_and_dtypes)
265 types = {
--> 266 array_api_compat.result_type(preprocess_types(t), xp=xp)
267 for t in arrays_and_dtypes
268 }
269 if any(isinstance(t, np.dtype) for t in types):
270 # only check if there's numpy dtypes – the array API does not
271 # define the types we're checking for
272 for left, right in PROMOTE_TO_OBJECT:
File ~/miniforge3/envs/calliope/lib/python3.12/site-packages/xarray/core/array_api_compat.py:42, in result_type(xp, *arrays_and_dtypes)
38 def result_type(*arrays_and_dtypes, xp) -> np.dtype:
39 if xp is np or any(
40 isinstance(getattr(t, "dtype", t), np.dtype) for t in arrays_and_dtypes
41 ):
---> 42 return xp.result_type(*arrays_and_dtypes)
43 else:
44 return _future_array_api_result_type(*arrays_and_dtypes, xp=xp)
TypeError: Cannot interpret '<pyomo.core.kernel.variable.variable object at 0x16b1dc7d0>' as a data type
What happened?
I'm storing Pyomo objects in xarray arrays to allow for vectorised operations with named arrays. E.g., I might have an array of Pyomo variables indexed over the
costsdimension and I want to sum them. Before v2024.6.0 this was possible without any issues. Now, if I include themin_countargument in my summation then it fails when trying to apply numpy type promotion to the Pyomo object.What did you expect to happen?
The
xarray.core.array_api_compat.result_typemethod should fall back tonumpy.object_when objects are present.Minimal Complete Verifiable Example
MVCE confirmation
Relevant log output
Anything else we need to know?
No response
Environment