forked from xarray-contrib/xarray-array-testing
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathreduction.py
More file actions
100 lines (80 loc) · 3.87 KB
/
reduction.py
File metadata and controls
100 lines (80 loc) · 3.87 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
from contextlib import nullcontext
import hypothesis.strategies as st
import numpy as np
import pytest
import xarray.testing.strategies as xrst
from hypothesis import given
from xarray_array_testing.base import DuckArrayTestMixin
class ReductionTests(DuckArrayTestMixin):
@staticmethod
def expected_errors(op, **parameters):
return nullcontext()
@pytest.mark.parametrize("op", ["mean", "sum", "prod", "std", "var"])
@given(st.data())
def test_variable_numerical_reduce(self, op, data):
variable = data.draw(xrst.variables(array_strategy_fn=self.array_strategy_fn))
with self.expected_errors(op, variable=variable):
# compute using xr.Variable.<OP>()
actual = getattr(variable, op)().data
# compute using xp.<OP>(array)
expected = getattr(self.xp, op)(variable.data)
assert isinstance(actual, self.array_type(op)), f"wrong type: {type(actual)}"
self.assert_equal(actual, expected)
@pytest.mark.parametrize("op", ["all", "any"])
@given(st.data())
def test_variable_boolean_reduce(self, op, data):
variable = data.draw(xrst.variables(array_strategy_fn=self.array_strategy_fn))
with self.expected_errors(op, variable=variable):
# compute using xr.Variable.<OP>()
actual = getattr(variable, op)().data
# compute using xp.<OP>(array)
expected = getattr(self.xp, op)(variable.data)
assert isinstance(actual, self.array_type(op)), f"wrong type: {type(actual)}"
self.assert_equal(actual, expected)
@pytest.mark.parametrize("op", ["max", "min"])
@given(st.data())
def test_variable_order_reduce(self, op, data):
variable = data.draw(xrst.variables(array_strategy_fn=self.array_strategy_fn))
with self.expected_errors(op, variable=variable):
# compute using xr.Variable.<OP>()
actual = getattr(variable, op)().data
# compute using xp.<OP>(array)
expected = getattr(self.xp, op)(variable.data)
assert isinstance(actual, self.array_type(op)), f"wrong type: {type(actual)}"
self.assert_equal(actual, expected)
@pytest.mark.parametrize("op", ["argmax", "argmin"])
@given(st.data())
def test_variable_order_reduce_index(self, op, data):
variable = data.draw(xrst.variables(array_strategy_fn=self.array_strategy_fn))
with self.expected_errors(op, variable=variable):
# compute using xr.Variable.<OP>()
actual = {k: v.item() for k, v in getattr(variable, op)(dim=...).items()}
# compute using xp.<OP>(array)
index = getattr(self.xp, op)(variable.data)
unraveled = np.unravel_index(index, variable.shape)
expected = dict(zip(variable.dims, unraveled))
self.assert_equal(actual, expected)
@pytest.mark.parametrize(
"op",
[
"cumsum",
pytest.param(
"cumprod",
marks=pytest.mark.skip(reason="not yet included in the array api"),
),
],
)
@given(st.data())
def test_variable_cumulative_reduce(self, op, data):
array_api_names = {"cumsum": "cumulative_sum", "cumprod": "cumulative_prod"}
variable = data.draw(xrst.variables(array_strategy_fn=self.array_strategy_fn))
with self.expected_errors(op, variable=variable):
# compute using xr.Variable.<OP>()
actual = getattr(variable, op)().data
# compute using xp.<OP>(array)
# Variable implements n-d cumulative ops by iterating over dims
expected = variable.data
for axis in range(variable.ndim):
expected = getattr(self.xp, array_api_names[op])(expected, axis=axis)
assert isinstance(actual, self.array_type(op)), f"wrong type: {type(actual)}"
self.assert_equal(actual, expected)