| 22 | @pytest.mark.parametrize("keepdims", [True, False]) |
| 23 | @pytest.mark.parametrize("nan", [True, False]) |
| 24 | def test_numel(dtype, keepdims, nan): |
| 25 | x = np.ones((2, 3, 4)) |
| 26 | if nan: |
| 27 | y = np.random.default_rng().uniform(-1, 1, size=(2, 3, 4)) |
| 28 | x[y < 0] = np.nan |
| 29 | numel = da.reductions.nannumel |
| 30 | |
| 31 | def _sum(arr, **kwargs): |
| 32 | n = np.sum(np.ma.masked_where(np.isnan(arr), arr), **kwargs) |
| 33 | return n.filled(0) if isinstance(n, np.ma.MaskedArray) else n |
| 34 | |
| 35 | else: |
| 36 | numel = da.reductions.numel |
| 37 | _sum = np.sum |
| 38 | |
| 39 | assert_eq( |
| 40 | numel(x, axis=(), keepdims=keepdims, dtype=dtype), |
| 41 | _sum(x, axis=(), keepdims=keepdims, dtype=dtype), |
| 42 | ) |
| 43 | assert_eq( |
| 44 | numel(x, axis=0, keepdims=keepdims, dtype=dtype), |
| 45 | _sum(x, axis=0, keepdims=keepdims, dtype=dtype), |
| 46 | ) |
| 47 | |
| 48 | for length in range(x.ndim): |
| 49 | for sub in itertools.combinations([d for d in range(x.ndim)], length): |
| 50 | assert_eq( |
| 51 | numel(x, axis=sub, keepdims=keepdims, dtype=dtype), |
| 52 | _sum(x, axis=sub, keepdims=keepdims, dtype=dtype), |
| 53 | ) |
| 54 | |
| 55 | for length in range(x.ndim): |
| 56 | for sub in itertools.combinations([d for d in range(x.ndim)], length): |
| 57 | ssub = np.random.default_rng().shuffle(list(sub)) |
| 58 | assert_eq( |
| 59 | numel(x, axis=ssub, keepdims=keepdims, dtype=dtype), |
| 60 | _sum(x, axis=ssub, keepdims=keepdims, dtype=dtype), |
| 61 | ) |
| 62 | |
| 63 | |
| 64 | def reduction_0d_test(da_func, darr, np_func, narr): |