| 951 | |
| 952 | |
| 953 | def test_weighted_reduction(): |
| 954 | # Weighted reduction |
| 955 | def w_sum(x, weights=None, dtype=None, computing_meta=False, **kwargs): |
| 956 | """`chunk` callable for (weighted) sum""" |
| 957 | if computing_meta: |
| 958 | return x |
| 959 | if weights is not None: |
| 960 | x = x * weights |
| 961 | return np.sum(x, dtype=dtype, **kwargs) |
| 962 | |
| 963 | # Arrays |
| 964 | a = 1 + np.ma.arange(60).reshape(6, 10) |
| 965 | a[2, 2] = np.ma.masked |
| 966 | dx = da.from_array(a, chunks=(4, 5)) |
| 967 | # Weights |
| 968 | w = np.linspace(1, 2, 6).reshape(6, 1) |
| 969 | |
| 970 | # No weights (i.e. normal sum) |
| 971 | x = da.reduction(dx, w_sum, np.sum, dtype=dx.dtype) |
| 972 | assert_eq(x, np.sum(a), check_shape=True) |
| 973 | |
| 974 | # Weighted sum |
| 975 | x = da.reduction(dx, w_sum, np.sum, dtype="f8", weights=w) |
| 976 | assert_eq(x, np.sum(a * w), check_shape=True) |
| 977 | |
| 978 | # Non-broadcastable weights (short axis) |
| 979 | with pytest.raises(ValueError): |
| 980 | da.reduction(dx, w_sum, np.sum, weights=[1, 2, 3]) |
| 981 | |
| 982 | # Non-broadcastable weights (too many dims) |
| 983 | with pytest.raises(ValueError): |
| 984 | da.reduction(dx, w_sum, np.sum, weights=[[[2]]]) |
| 985 | |
| 986 | |
| 987 | def test_cumreduction_no_rechunk_on_1d_array(): |