Skip to content

Commit

Permalink
don't drop scalar quantile coords
Browse files Browse the repository at this point in the history
  • Loading branch information
@keewis
keewis committed Nov 14, 2019
1 parent 5c441c9 commit 5f9ca06
Show file tree
Hide file tree
Showing 2 changed files with 50 additions and 28 deletions.
13 changes: 6 additions & 7 deletions xarray/core/groupby.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -584,11 +584,12 @@ def quantile(self, q, dim=None, interpolation="linear", keep_attrs=None):
Returns
-------
quantiles : Variable
If `q` is a single quantile, then the result
is a scalar. If multiple percentiles are given, first axis of
the result corresponds to the quantile and a quantile dimension
is added to the return array. The other dimensions are the
dimensions that remain after the reduction of the array.
If `q` is a single quantile, then the result is a
scalar. If multiple percentiles are given, first axis of
the result corresponds to the quantile. In either case a
quantile dimension is added to the return array. The other
dimensions are the dimensions that remain after the
reduction of the array.
See Also
--------
Expand All @@ -607,8 +608,6 @@ def quantile(self, q, dim=None, interpolation="linear", keep_attrs=None):
keep_attrs=keep_attrs,
)

if np.asarray(q, dtype=np.float64).ndim == 0:
out = out.drop_vars("quantile")
return out

def where(self, cond, other=dtypes.NA):
Expand Down
65 changes: 44 additions & 21 deletions xarray/tests/test_groupby.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -137,57 +137,73 @@ def test_da_groupby_empty():

def test_da_groupby_quantile():

array = xr.DataArray([1, 2, 3, 4, 5, 6], [("x", [1, 1, 1, 2, 2, 2])])
array = xr.DataArray(
data=[1, 2, 3, 4, 5, 6], coords={"x": [1, 1, 1, 2, 2, 2]}, dims="x"
)

# Scalar quantile
expected = xr.DataArray([2, 5], [("x", [1, 2])])
expected = xr.DataArray(
data=[2, 5], coords={"x": [1, 2], "quantile": 0.5}, dims="x"
)
actual = array.groupby("x").quantile(0.5)
assert_identical(expected, actual)

# Vector quantile
expected = xr.DataArray([[1, 3], [4, 6]], [("x", [1, 2]), ("quantile", [0, 1])])
expected = xr.DataArray(
data=[[1, 3], [4, 6]],
coords={"x": [1, 2], "quantile": [0, 1]},
dims=("x", "quantile"),
)
actual = array.groupby("x").quantile([0, 1])
assert_identical(expected, actual)

# Multiple dimensions
array = xr.DataArray(
[[1, 11, 26], [2, 12, 22], [3, 13, 23], [4, 16, 24], [5, 15, 25]],
[("x", [1, 1, 1, 2, 2]), ("y", [0, 0, 1])],
data=[[1, 11, 26], [2, 12, 22], [3, 13, 23], [4, 16, 24], [5, 15, 25]],
coords={"x": [1, 1, 1, 2, 2], "y": [0, 0, 1]},
dims=("x", "y"),
)

actual_x = array.groupby("x").quantile(0, dim=...)
expected_x = xr.DataArray([1, 4], [("x", [1, 2])])
expected_x = xr.DataArray(
data=[1, 4], coords={"x": [1, 2], "quantile": 0}, dims="x"
)
assert_identical(expected_x, actual_x)

actual_y = array.groupby("y").quantile(0, dim=...)
expected_y = xr.DataArray([1, 22], [("y", [0, 1])])
expected_y = xr.DataArray(
data=[1, 22], coords={"y": [0, 1], "quantile": 0}, dims="y"
)
assert_identical(expected_y, actual_y)

actual_xx = array.groupby("x").quantile(0)
expected_xx = xr.DataArray(
[[1, 11, 22], [4, 15, 24]], [("x", [1, 2]), ("y", [0, 0, 1])]
data=[[1, 11, 22], [4, 15, 24]],
coords={"x": [1, 2], "y": [0, 0, 1], "quantile": 0},
dims=("x", "y"),
)
assert_identical(expected_xx, actual_xx)

actual_yy = array.groupby("y").quantile(0)
expected_yy = xr.DataArray(
[[1, 26], [2, 22], [3, 23], [4, 24], [5, 25]],
[("x", [1, 1, 1, 2, 2]), ("y", [0, 1])],
data=[[1, 26], [2, 22], [3, 23], [4, 24], [5, 25]],
coords={"x": [1, 1, 1, 2, 2], "y": [0, 1], "quantile": 0},
dims=("x", "y"),
)
assert_identical(expected_yy, actual_yy)

times = pd.date_range("2000-01-01", periods=365)
x = [0, 1]
foo = xr.DataArray(
np.reshape(np.arange(365 * 2), (365, 2)),
coords=dict(time=times, x=x),
coords={"time": times, "x": x},
dims=("time", "x"),
)
g = foo.groupby(foo.time.dt.month)

actual = g.quantile(0, dim=...)
expected = xr.DataArray(
[
data=[
0.0,
62.0,
120.0,
Expand All @@ -201,12 +217,17 @@ def test_da_groupby_quantile():
610.0,
670.0,
],
[("month", np.arange(1, 13))],
coords={"month": np.arange(1, 13), "quantile": 0},
dims="month",
)
assert_identical(expected, actual)

actual = g.quantile(0, dim="time")[:2]
expected = xr.DataArray([[0.0, 1], [62.0, 63]], [("month", [1, 2]), ("x", [0, 1])])
expected = xr.DataArray(
data=[[0.0, 1], [62.0, 63]],
coords={"month": [1, 2], "x": [0, 1], "quantile": 0},
dims=("month", "x"),
)
assert_identical(expected, actual)


Expand All @@ -216,7 +237,9 @@ def test_ds_groupby_quantile():
)

# Scalar quantile
expected = xr.Dataset({"a": ("x", [2, 5])}, coords={"x": [1, 2]})
expected = xr.Dataset(
data_vars={"a": ("x", [2, 5])}, coords={"quantile": 0.5, "x": [1, 2]}
)
actual = ds.groupby("x").quantile(0.5)
assert_identical(expected, actual)

Expand All @@ -240,24 +263,24 @@ def test_ds_groupby_quantile():
)

actual_x = ds.groupby("x").quantile(0, dim=...)
expected_x = xr.Dataset({"a": ("x", [1, 4])}, coords={"x": [1, 2]})
expected_x = xr.Dataset({"a": ("x", [1, 4])}, coords={"x": [1, 2], "quantile": 0})
assert_identical(expected_x, actual_x)

actual_y = ds.groupby("y").quantile(0, dim=...)
expected_y = xr.Dataset({"a": ("y", [1, 22])}, coords={"y": [0, 1]})
expected_y = xr.Dataset({"a": ("y", [1, 22])}, coords={"y": [0, 1], "quantile": 0})
assert_identical(expected_y, actual_y)

actual_xx = ds.groupby("x").quantile(0)
expected_xx = xr.Dataset(
{"a": (("x", "y"), [[1, 11, 22], [4, 15, 24]])},
coords={"x": [1, 2], "y": [0, 0, 1]},
coords={"x": [1, 2], "y": [0, 0, 1], "quantile": 0},
)
assert_identical(expected_xx, actual_xx)

actual_yy = ds.groupby("y").quantile(0)
expected_yy = xr.Dataset(
{"a": (("x", "y"), [[1, 26], [2, 22], [3, 23], [4, 24], [5, 25]])},
coords={"x": [1, 1, 1, 2, 2], "y": [0, 1]},
coords={"x": [1, 1, 1, 2, 2], "y": [0, 1], "quantile": 0},
).transpose()
assert_identical(expected_yy, actual_yy)

Expand Down Expand Up @@ -290,14 +313,14 @@ def test_ds_groupby_quantile():
],
)
},
coords={"month": np.arange(1, 13)},
coords={"month": np.arange(1, 13), "quantile": 0},
)
assert_identical(expected, actual)

actual = g.quantile(0, dim="time").isel(month=slice(None, 2))
expected = xr.Dataset(
data_vars={"a": (("month", "x"), [[0.0, 1], [62.0, 63]])},
coords={"month": [1, 2], "x": [0, 1]},
coords={"month": [1, 2], "x": [0, 1], "quantile": 0},
)
assert_identical(expected, actual)

Expand Down

0 comments on commit 5f9ca06

Please sign in to comment.