Enable numbagg in calculation of quantiles

xarray/core/nputils.py

@@ -206,6 +206,11 @@ def f(values, axis=None, **kwargs):
                 # to ddof=1 above.
                 if pycompat.mod_version("numbagg") < Version("0.7.0"):
                     kwargs.pop("ddof", None)
+                if name == "nanquantile":
+                    if kwargs.get("method", None) != "linear":
+                        return
+                    kwargs["quantiles"] = kwargs.pop("q")
+                    kwargs.pop("method", None)
                 return nba_func(values, axis=axis, **kwargs)
         if (
             _BOTTLENECK_AVAILABLE
@@ -285,3 +290,4 @@ def least_squares(lhs, rhs, rcond=None, skipna=False):
 nancumprod = _create_method("nancumprod")
 nanargmin = _create_method("nanargmin")
 nanargmax = _create_method("nanargmax")
+nanquantile = _create_method("nanquantile")

xarray/core/variable.py

@@ -1994,7 +1994,7 @@ def quantile(
             method = interpolation
 
         if skipna or (skipna is None and self.dtype.kind in "cfO"):
-            _quantile_func = np.nanquantile
+            _quantile_func = nputils.nanquantile
         else:
             _quantile_func = np.quantile
 

xarray/tests/conftest.py

@@ -14,9 +14,11 @@ def backend(request):
     return request.param
 
 
-@pytest.fixture(params=["numbagg", "bottleneck"])
+@pytest.fixture(params=["numbagg", "bottleneck", None])
 def compute_backend(request):
-    if request.param == "bottleneck":
+    if request.param is None:
+        options = dict(use_bottleneck=False, use_numbagg=False)
+    elif request.param == "bottleneck":
         options = dict(use_bottleneck=True, use_numbagg=False)
     elif request.param == "numbagg":
         options = dict(use_bottleneck=False, use_numbagg=True)

xarray/tests/test_dataarray.py

@@ -2889,12 +2889,13 @@
         with pytest.raises(TypeError):
             orig.mean(out=np.ones(orig.shape))
 
+    @pytest.mark.parametrize("compute_backend", ["numbagg", None], indirect=True)
     @pytest.mark.parametrize("skipna", [True, False, None])
     @pytest.mark.parametrize("q", [0.25, [0.50], [0.25, 0.75]])
     @pytest.mark.parametrize(
         "axis, dim", zip([None, 0, [0], [0, 1]], [None, "x", ["x"], ["x", "y"]])
     )
-    def test_quantile(self, q, axis, dim, skipna) -> None:
+    def test_quantile(self, q, axis, dim, skipna, compute_backend) -> None:
         va = self.va.copy(deep=True)
         va[0, 0] = np.nan
 
@@ -2912,7 +2913,7 @@
     @pytest.mark.parametrize("method", ["midpoint", "lower"])
     def test_quantile_method(self, method) -> None:
         q = [0.25, 0.5, 0.75]
         actual = DataArray(self.va).quantile(q, method=method)
 
         expected = np.nanquantile(self.dv.values, np.array(q), method=method)
 
@@ -2927,7 +2928,7 @@
             FutureWarning,
             match="`interpolation` argument to quantile was renamed to `method`",
         ):
             actual = da.quantile(q, interpolation=method)
 
         expected = da.quantile(q, method=method)
 

xarray/tests/test_dataset.py

@@ -5616,9 +5616,10 @@
         )
         assert_identical(expected, actual)
 
+    @pytest.mark.parametrize("compute_backend", ["numbagg", None], indirect=True)
     @pytest.mark.parametrize("skipna", [True, False, None])
     @pytest.mark.parametrize("q", [0.25, [0.50], [0.25, 0.75]])
-    def test_quantile(self, q, skipna) -> None:
+    def test_quantile(self, q, skipna, compute_backend) -> None:
         ds = create_test_data(seed=123)
         ds.var1.data[0, 0] = np.nan
 
@@ -5639,8 +5640,9 @@
         assert "dim3" in ds_quantile.dims
         assert all(d not in ds_quantile.dims for d in dim)
 
+    @pytest.mark.parametrize("compute_backend", ["numbagg", None], indirect=True)
     @pytest.mark.parametrize("skipna", [True, False])
-    def test_quantile_skipna(self, skipna) -> None:
+    def test_quantile_skipna(self, skipna, compute_backend) -> None:
         q = 0.1
         dim = "time"
         ds = Dataset({"a": ([dim], np.arange(0, 11))})
@@ -5655,7 +5657,7 @@
     @pytest.mark.parametrize("method", ["midpoint", "lower"])
     def test_quantile_method(self, method) -> None:
         ds = create_test_data(seed=123)
         q = [0.25, 0.5, 0.75]
 
         result = ds.quantile(q, method=method)
 
@@ -5667,7 +5669,7 @@
     def test_quantile_interpolation_deprecated(self, method) -> None:
         ds = create_test_data(seed=123)
         q = [0.25, 0.5, 0.75]
 
         with warnings.catch_warnings(record=True) as w:
             ds.quantile(q, interpolation=method)
 

xarray/tests/test_units.py

@@ -2004,6 +2004,7 @@ def test_squeeze(self, dim, dtype):
         assert_units_equal(expected, actual)
         assert_identical(expected, actual)
 
+    @pytest.mark.parametrize("compute_backend", ["numbagg", None], indirect=True)
     @pytest.mark.parametrize(
         "func",
         (
@@ -2025,7 +2026,7 @@ def test_squeeze(self, dim, dtype):
         ),
         ids=repr,
     )
-    def test_computation(self, func, dtype):
+    def test_computation(self, func, dtype, compute_backend):
         base_unit = unit_registry.m
         array = np.linspace(0, 5, 5 * 10).reshape(5, 10).astype(dtype) * base_unit
         variable = xr.Variable(("x", "y"), array)
@@ -3757,6 +3758,7 @@ def test_differentiate_integrate(self, func, variant, dtype):
         assert_units_equal(expected, actual)
         assert_identical(expected, actual)
 
+    @pytest.mark.parametrize("compute_backend", ["numbagg", None], indirect=True)
     @pytest.mark.parametrize(
         "variant",
         (
@@ -3777,7 +3779,7 @@ def test_differentiate_integrate(self, func, variant, dtype):
         ),
         ids=repr,
     )
-    def test_computation(self, func, variant, dtype):
+    def test_computation(self, func, variant, dtype, compute_backend):
         unit = unit_registry.m
 
         variants = {
@@ -3883,6 +3885,7 @@ def test_resample(self, dtype):
         assert_units_equal(expected, actual)
         assert_identical(expected, actual)
 
+    @pytest.mark.parametrize("compute_backend", ["numbagg", None], indirect=True)
     @pytest.mark.parametrize(
         "variant",
         (
@@ -3903,7 +3906,7 @@ def test_resample(self, dtype):
         ),
         ids=repr,
     )
-    def test_grouped_operations(self, func, variant, dtype):
+    def test_grouped_operations(self, func, variant, dtype, compute_backend):
         unit = unit_registry.m
 
         variants = {
@@ -5240,6 +5243,7 @@ def test_interp_reindex_like_indexing(self, func, unit, error, dtype):
         assert_units_equal(expected, actual)
         assert_equal(expected, actual)
 
+    @pytest.mark.parametrize("compute_backend", ["numbagg", None], indirect=True)
     @pytest.mark.parametrize(
         "func",
         (
@@ -5262,7 +5266,7 @@ def test_interp_reindex_like_indexing(self, func, unit, error, dtype):
             "coords",
         ),
     )
-    def test_computation(self, func, variant, dtype):
+    def test_computation(self, func, variant, dtype, compute_backend):
         variants = {
             "data": ((unit_registry.degK, unit_registry.Pa), 1, 1),
             "dims": ((1, 1), unit_registry.m, 1),
@@ -5394,6 +5398,7 @@ def test_resample(self, variant, dtype):
         assert_units_equal(expected, actual)
         assert_equal(expected, actual)
 
+    @pytest.mark.parametrize("compute_backend", ["numbagg", None], indirect=True)
     @pytest.mark.parametrize(
         "func",
         (
@@ -5415,7 +5420,7 @@ def test_resample(self, variant, dtype):
             "coords",
         ),
     )
-    def test_grouped_operations(self, func, variant, dtype):
+    def test_grouped_operations(self, func, variant, dtype, compute_backend):
         variants = {
             "data": ((unit_registry.degK, unit_registry.Pa), 1, 1),
             "dims": ((1, 1), unit_registry.m, 1),

xarray/tests/test_variable.py

@@ -1806,14 +1806,14 @@
             v = v.chunk({"x": 2})
 
         q = np.array([0.25, 0.5, 0.75])
         actual = v.quantile(q, dim="y", method=method)
 
         expected = np.nanquantile(self.d, q, axis=1, method=method)
 
         if use_dask:
             assert isinstance(actual.data, dask_array_type)
 
         np.testing.assert_allclose(actual.values, expected)
 
     @pytest.mark.parametrize("method", ["midpoint", "lower"])
     def test_quantile_interpolation_deprecation(self, method) -> None:
@@ -1842,13 +1842,15 @@
         with pytest.raises(ValueError, match=r"consists of multiple chunks"):
             v.quantile(0.5, dim="x")
 
+    @pytest.mark.parametrize("compute_backend", ["numbagg", None], indirect=True)
     @pytest.mark.parametrize("q", [-0.1, 1.1, [2], [0.25, 2]])
-    def test_quantile_out_of_bounds(self, q):
+    def test_quantile_out_of_bounds(self, q, compute_backend):
         v = Variable(["x", "y"], self.d)
 
         # escape special characters
         with pytest.raises(
-            ValueError, match=r"Quantiles must be in the range \[0, 1\]"
+            ValueError,
+            match=r"(Q|q)uantiles must be in the range \[0, 1\]",
         ):
             v.quantile(q, dim="x")