Function test_weighted_quantile_3D

xarray/tests/test_weighted.py:627–651

(dim, q, add_nans, skipna)

@pytest.mark.parametrize("add_nans", (True, False))
@pytest.mark.parametrize("skipna", (None, True, False))
def test_weighted_quantile_3D(dim, q, add_nans, skipna):
    dims = ("a", "b", "c")
    coords = dict(a=[0, 1, 2], b=[0, 1, 2, 3], c=[0, 1, 2, 3, 4])

    data = np.arange(60).reshape(3, 4, 5).astype(float)

    # add approximately 25 % NaNs (https://stackoverflow.com/a/32182680/3010700)
    if add_nans:
        c = int(data.size * 0.25)
        data.ravel()[np.random.choice(data.size, c, replace=False)] = np.nan

    da = DataArray(data, dims=dims, coords=coords)

    # Weights are all ones, because we will compare against DataArray.quantile (non-weighted)
    weights = xr.ones_like(da)

    result = da.weighted(weights).quantile(q, dim=dim, skipna=skipna)
    expected = da.quantile(q, dim=dim, skipna=skipna)

    assert_allclose(expected, result)

    ds = da.to_dataset(name="data")
    result2 = ds.weighted(weights).quantile(q, dim=dim, skipna=skipna)

    assert_allclose(expected, result2.data)


@pytest.mark.parametrize(