Function test_groupby_drops_nans

xarray/tests/test_groupby.py:666–740

(shuffle: bool, chunk: Literal[False] | dict)

    ],
)
def test_groupby_drops_nans(shuffle: bool, chunk: Literal[False] | dict) -> None:
    if shuffle and chunk and not has_dask_ge_2024_08_1:
        pytest.skip()
    # GH2383
    # nan in 2D data variable (requires stacking)
    ds = xr.Dataset(
        {
            "variable": (("lat", "lon", "time"), np.arange(60.0).reshape((4, 3, 5))),
            "id": (("lat", "lon"), np.arange(12.0).reshape((4, 3))),
        },
        coords={"lat": np.arange(4), "lon": np.arange(3), "time": np.arange(5)},
    )

    ds["id"].values[0, 0] = np.nan
    ds["id"].values[3, 0] = np.nan
    ds["id"].values[-1, -1] = np.nan

    if chunk:
        ds["variable"] = ds["variable"].chunk(chunk)
    grouped = ds.groupby(ds.id)
    if shuffle:
        grouped = grouped.shuffle_to_chunks().groupby(ds.id)

    # non reduction operation
    expected1 = ds.copy()
    expected1.variable.data[0, 0, :] = np.nan
    expected1.variable.data[-1, -1, :] = np.nan
    expected1.variable.data[3, 0, :] = np.nan
    actual1 = grouped.map(lambda x: x).transpose(*ds.variable.dims)
    assert_identical(actual1, expected1)

    # reduction along grouped dimension
    actual2 = grouped.mean()
    stacked = ds.stack({"xy": ["lat", "lon"]})
    expected2 = (
        stacked.variable.where(stacked.id.notnull())
        .rename({"xy": "id"})
        .to_dataset()
        .reset_index("id", drop=True)
        .assign(id=stacked.id.values)
        .dropna("id")
        .transpose(*actual2.variable.dims)
    )
    assert_identical(actual2, expected2)

    # reduction operation along a different dimension
    actual3 = grouped.mean("time")
    expected3 = ds.mean("time").where(ds.id.notnull())
    assert_identical(actual3, expected3)

    # NaN in non-dimensional coordinate
    array = xr.DataArray([1, 2, 3], [("x", [1, 2, 3])])
    array["x1"] = ("x", [1, 1, np.nan])
    expected4 = xr.DataArray(3, [("x1", [1])])
    actual4 = array.groupby("x1").sum()
    assert_equal(expected4, actual4)

    # NaT in non-dimensional coordinate