don't drop scalar quantile coords

Author: keewis

xarray/core/groupby.py

@@ -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
         --------
@@ -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):

xarray/tests/test_groupby.py

@@ -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,
@@ -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)
 
 
@@ -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)
 
@@ -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)
 
@@ -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)