(use_flox: bool, shuffle: bool)
| 2987 | @pytest.mark.parametrize("shuffle", [True, False]) |
| 2988 | @pytest.mark.parametrize("use_flox", [True, False]) |
| 2989 | def test_multiple_groupers(use_flox: bool, shuffle: bool) -> None: |
| 2990 | da = DataArray( |
| 2991 | np.array([1, 2, 3, 0, 2, np.nan]), |
| 2992 | dims="d", |
| 2993 | coords=dict( |
| 2994 | labels1=("d", np.array(["a", "b", "c", "c", "b", "a"])), |
| 2995 | labels2=("d", np.array(["x", "y", "z", "z", "y", "x"])), |
| 2996 | ), |
| 2997 | name="foo", |
| 2998 | ) |
| 2999 | |
| 3000 | groupers: dict[str, Grouper] |
| 3001 | groupers = dict(labels1=UniqueGrouper(), labels2=UniqueGrouper()) |
| 3002 | gb = da.groupby(groupers) |
| 3003 | if shuffle: |
| 3004 | gb = gb.shuffle_to_chunks().groupby(groupers) |
| 3005 | repr(gb) |
| 3006 | |
| 3007 | expected = DataArray( |
| 3008 | np.array([[1.0, np.nan, np.nan], [np.nan, 2.0, np.nan], [np.nan, np.nan, 1.5]]), |
| 3009 | dims=("labels1", "labels2"), |
| 3010 | coords={ |
| 3011 | "labels1": np.array(["a", "b", "c"], dtype=object), |
| 3012 | "labels2": np.array(["x", "y", "z"], dtype=object), |
| 3013 | }, |
| 3014 | name="foo", |
| 3015 | ) |
| 3016 | with xr.set_options(use_flox=use_flox): |
| 3017 | actual = gb.mean() |
| 3018 | assert_identical(actual, expected) |
| 3019 | |
| 3020 | # ------- |
| 3021 | coords = {"a": ("x", [0, 0, 1, 1]), "b": ("y", [0, 0, 1, 1])} |
| 3022 | square = DataArray(np.arange(16).reshape(4, 4), coords=coords, dims=["x", "y"]) |
| 3023 | groupers = dict(a=UniqueGrouper(), b=UniqueGrouper()) |
| 3024 | gb = square.groupby(groupers) |
| 3025 | if shuffle: |
| 3026 | gb = gb.shuffle_to_chunks().groupby(groupers) |
| 3027 | repr(gb) |
| 3028 | with xr.set_options(use_flox=use_flox): |
| 3029 | actual = gb.mean() |
| 3030 | expected = DataArray( |
| 3031 | np.array([[2.5, 4.5], [10.5, 12.5]]), |
| 3032 | dims=("a", "b"), |
| 3033 | coords={"a": [0, 1], "b": [0, 1]}, |
| 3034 | ) |
| 3035 | assert_identical(actual, expected) |
| 3036 | |
| 3037 | expected = square.astype(np.float64) |
| 3038 | expected["a"], expected["b"] = broadcast(square.a, square.b) |
| 3039 | with xr.set_options(use_flox=use_flox): |
| 3040 | assert_identical( |
| 3041 | square.groupby(x=UniqueGrouper(), y=UniqueGrouper()).mean(), expected |
| 3042 | ) |
| 3043 | |
| 3044 | b = xr.DataArray( |
| 3045 | np.random.default_rng(0).random((2, 3, 4)), |
| 3046 | coords={"xy": (("x", "y"), [["a", "b", "c"], ["b", "c", "c"]], {"foo": "bar"})}, |
nothing calls this directly
no test coverage detected
searching dependent graphs…