Skip to content

Commit

Permalink
Duck array ops for all and any
Browse files Browse the repository at this point in the history
  • Loading branch information
tomwhite committed Dec 13, 2024
1 parent 49502fc commit b3cd8eb
Show file tree
Hide file tree
Showing 2 changed files with 14 additions and 6 deletions.
18 changes: 13 additions & 5 deletions xarray/core/duck_array_ops.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,8 +16,6 @@

import numpy as np
import pandas as pd
from numpy import all as array_all # noqa: F401
from numpy import any as array_any # noqa: F401
from numpy import ( # noqa: F401
isclose,
isnat,
Expand Down Expand Up @@ -319,7 +317,7 @@ def allclose_or_equiv(arr1, arr2, rtol=1e-5, atol=1e-8):
if lazy_equiv is None:
with warnings.catch_warnings():
warnings.filterwarnings("ignore", r"All-NaN (slice|axis) encountered")
return bool(isclose(arr1, arr2, rtol=rtol, atol=atol, equal_nan=True).all())
return bool(array_all(isclose(arr1, arr2, rtol=rtol, atol=atol, equal_nan=True)))
else:
return lazy_equiv

Expand All @@ -333,7 +331,7 @@ def array_equiv(arr1, arr2):
with warnings.catch_warnings():
warnings.filterwarnings("ignore", "In the future, 'NAT == x'")
flag_array = (arr1 == arr2) | (isnull(arr1) & isnull(arr2))
return bool(flag_array.all())
return bool(array_all(flag_array))
else:
return lazy_equiv

Expand All @@ -349,7 +347,7 @@ def array_notnull_equiv(arr1, arr2):
with warnings.catch_warnings():
warnings.filterwarnings("ignore", "In the future, 'NAT == x'")
flag_array = (arr1 == arr2) | isnull(arr1) | isnull(arr2)
return bool(flag_array.all())
return bool(array_all(flag_array))
else:
return lazy_equiv

Expand Down Expand Up @@ -536,6 +534,16 @@ def f(values, axis=None, skipna=None, **kwargs):
cumsum_1d.numeric_only = True


def array_all(array, axis=None, keepdims=False):
xp = get_array_namespace(array)
return xp.all(array, axis=axis, keepdims=keepdims)


def array_any(array, axis=None, keepdims=False):
xp = get_array_namespace(array)
return xp.any(array, axis=axis, keepdims=keepdims)


_mean = _create_nan_agg_method("mean", invariant_0d=True)


Expand Down
2 changes: 1 addition & 1 deletion xarray/core/weighted.py
Original file line number Diff line number Diff line change
Expand Up @@ -171,7 +171,7 @@ def __init__(self, obj: T_Xarray, weights: T_DataArray) -> None:

def _weight_check(w):
# Ref https://github.com/pydata/xarray/pull/4559/files#r515968670
if duck_array_ops.isnull(w).any():
if duck_array_ops.array_any(duck_array_ops.isnull(w)):
raise ValueError(
"`weights` cannot contain missing values. "
"Missing values can be replaced by `weights.fillna(0)`."
Expand Down

0 comments on commit b3cd8eb

Please sign in to comment.