diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 87834d2e676..dd1db41fe23 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -71,6 +71,10 @@ New Features By `Richard Kleijn `_ . - Add a ``combine_attrs`` parameter to :py:func:`open_mfdataset` (:pull:`4971`). By `Justus Magin `_. +- Enable passing arrays with a subset of dimensions to + :py:meth:`DataArray.clip` & :py:meth:`Dataset.clip`; these methods now use + :py:func:`xarray.apply_ufunc`; (:pull:`5184`). + By `Maximilian Roos `_. - Disable the `cfgrib` backend if the `eccodes` library is not installed (:pull:`5083`). By `Baudouin Raoult `_. - Added :py:meth:`DataArray.curvefit` and :py:meth:`Dataset.curvefit` for general curve fitting applications. (:issue:`4300`, :pull:`4849`) By `Sam Levang `_. diff --git a/xarray/core/_typed_ops.py b/xarray/core/_typed_ops.py index a5a9b99e275..d1e68a6fc0d 100644 --- a/xarray/core/_typed_ops.py +++ b/xarray/core/_typed_ops.py @@ -144,9 +144,6 @@ def round(self, *args, **kwargs): def argsort(self, *args, **kwargs): return self._unary_op(ops.argsort, *args, **kwargs) - def clip(self, *args, **kwargs): - return self._unary_op(ops.clip, *args, **kwargs) - def conj(self, *args, **kwargs): return self._unary_op(ops.conj, *args, **kwargs) @@ -195,7 +192,6 @@ def conjugate(self, *args, **kwargs): __invert__.__doc__ = operator.invert.__doc__ round.__doc__ = ops.round_.__doc__ argsort.__doc__ = ops.argsort.__doc__ - clip.__doc__ = ops.clip.__doc__ conj.__doc__ = ops.conj.__doc__ conjugate.__doc__ = ops.conjugate.__doc__ @@ -338,9 +334,6 @@ def round(self, *args, **kwargs): def argsort(self, *args, **kwargs): return self._unary_op(ops.argsort, *args, **kwargs) - def clip(self, *args, **kwargs): - return self._unary_op(ops.clip, *args, **kwargs) - def conj(self, *args, **kwargs): return self._unary_op(ops.conj, *args, **kwargs) @@ -389,7 +382,6 @@ def conjugate(self, *args, **kwargs): __invert__.__doc__ = operator.invert.__doc__ round.__doc__ = ops.round_.__doc__ argsort.__doc__ = ops.argsort.__doc__ - clip.__doc__ = ops.clip.__doc__ conj.__doc__ = ops.conj.__doc__ conjugate.__doc__ = ops.conjugate.__doc__ @@ -532,9 +524,6 @@ def round(self, *args, **kwargs): def argsort(self, *args, **kwargs): return self._unary_op(ops.argsort, *args, **kwargs) - def clip(self, *args, **kwargs): - return self._unary_op(ops.clip, *args, **kwargs) - def conj(self, *args, **kwargs): return self._unary_op(ops.conj, *args, **kwargs) @@ -583,7 +572,6 @@ def conjugate(self, *args, **kwargs): __invert__.__doc__ = operator.invert.__doc__ round.__doc__ = ops.round_.__doc__ argsort.__doc__ = ops.argsort.__doc__ - clip.__doc__ = ops.clip.__doc__ conj.__doc__ = ops.conj.__doc__ conjugate.__doc__ = ops.conjugate.__doc__ diff --git a/xarray/core/_typed_ops.pyi b/xarray/core/_typed_ops.pyi index 0eaced16d9f..4a6c2dc7b4e 100644 --- a/xarray/core/_typed_ops.pyi +++ b/xarray/core/_typed_ops.pyi @@ -64,7 +64,6 @@ class DatasetOpsMixin: def __invert__(self: T_Dataset) -> T_Dataset: ... def round(self: T_Dataset, *args, **kwargs) -> T_Dataset: ... def argsort(self: T_Dataset, *args, **kwargs) -> T_Dataset: ... - def clip(self: T_Dataset, *args, **kwargs) -> T_Dataset: ... def conj(self: T_Dataset, *args, **kwargs) -> T_Dataset: ... def conjugate(self: T_Dataset, *args, **kwargs) -> T_Dataset: ... @@ -235,7 +234,6 @@ class DataArrayOpsMixin: def __invert__(self: T_DataArray) -> T_DataArray: ... def round(self: T_DataArray, *args, **kwargs) -> T_DataArray: ... def argsort(self: T_DataArray, *args, **kwargs) -> T_DataArray: ... - def clip(self: T_DataArray, *args, **kwargs) -> T_DataArray: ... def conj(self: T_DataArray, *args, **kwargs) -> T_DataArray: ... def conjugate(self: T_DataArray, *args, **kwargs) -> T_DataArray: ... @@ -406,7 +404,6 @@ class VariableOpsMixin: def __invert__(self: T_Variable) -> T_Variable: ... def round(self: T_Variable, *args, **kwargs) -> T_Variable: ... def argsort(self: T_Variable, *args, **kwargs) -> T_Variable: ... - def clip(self: T_Variable, *args, **kwargs) -> T_Variable: ... def conj(self: T_Variable, *args, **kwargs) -> T_Variable: ... def conjugate(self: T_Variable, *args, **kwargs) -> T_Variable: ... diff --git a/xarray/core/common.py b/xarray/core/common.py index 2b0eac89146..087bc7e4dfe 100644 --- a/xarray/core/common.py +++ b/xarray/core/common.py @@ -380,6 +380,28 @@ def squeeze( dims = get_squeeze_dims(self, dim, axis) return self.isel(drop=drop, **{d: 0 for d in dims}) + def clip(self, min=None, max=None, *, keep_attrs: bool = None): + """ + Return an array whose values are limited to ``[min, max]``. + At least one of max or min must be given. + + Refer to `numpy.clip` for full documentation. + + See Also + -------- + numpy.clip : equivalent function + """ + from .computation import apply_ufunc + + if keep_attrs is None: + # When this was a unary func, the default was True, so retaining the + # default. + keep_attrs = _get_keep_attrs(default=True) + + return apply_ufunc( + np.clip, self, min, max, keep_attrs=keep_attrs, dask="allowed" + ) + def get_index(self, key: Hashable) -> pd.Index: """Get an index for a dimension, with fall-back to a default RangeIndex""" if key not in self.dims: diff --git a/xarray/core/ops.py b/xarray/core/ops.py index 50495b84246..27740d53d45 100644 --- a/xarray/core/ops.py +++ b/xarray/core/ops.py @@ -291,7 +291,6 @@ def inplace_to_noninplace_op(f): # _typed_ops.py uses the following wrapped functions as a kind of unary operator argsort = _method_wrapper("argsort") -clip = _method_wrapper("clip") conj = _method_wrapper("conj") conjugate = _method_wrapper("conjugate") round_ = _func_slash_method_wrapper(duck_array_ops.around, name="round") diff --git a/xarray/core/variable.py b/xarray/core/variable.py index 79850c0db36..469174f35e2 100644 --- a/xarray/core/variable.py +++ b/xarray/core/variable.py @@ -1665,6 +1665,21 @@ def fillna(self, value): def where(self, cond, other=dtypes.NA): return ops.where_method(self, cond, other) + def clip(self, min=None, max=None): + """ + Return an array whose values are limited to ``[min, max]``. + At least one of max or min must be given. + + Refer to `numpy.clip` for full documentation. + + See Also + -------- + numpy.clip : equivalent function + """ + from .computation import apply_ufunc + + return apply_ufunc(np.clip, self, min, max, dask="allowed") + def reduce( self, func, diff --git a/xarray/tests/test_dataarray.py b/xarray/tests/test_dataarray.py index c18f9acfca1..accb601fc1b 100644 --- a/xarray/tests/test_dataarray.py +++ b/xarray/tests/test_dataarray.py @@ -6434,28 +6434,43 @@ def test_idxminmax_dask(self, op, ndim): assert_equal(getattr(ar0_dsk, op)(dim="x"), getattr(ar0_raw, op)(dim="x")) +@pytest.fixture(params=["numpy", pytest.param("dask", marks=requires_dask)]) +def backend(request): + return request.param + + @pytest.fixture(params=[1]) -def da(request): +def da(request, backend): if request.param == 1: times = pd.date_range("2000-01-01", freq="1D", periods=21) - values = np.random.random((3, 21, 4)) - da = DataArray(values, dims=("a", "time", "x")) - da["time"] = times - return da + da = DataArray( + np.random.random((3, 21, 4)), + dims=("a", "time", "x"), + coords=dict(time=times), + ) if request.param == 2: - return DataArray([0, np.nan, 1, 2, np.nan, 3, 4, 5, np.nan, 6, 7], dims="time") + da = DataArray([0, np.nan, 1, 2, np.nan, 3, 4, 5, np.nan, 6, 7], dims="time") if request.param == "repeating_ints": - return DataArray( + da = DataArray( np.tile(np.arange(12), 5).reshape(5, 4, 3), coords={"x": list("abc"), "y": list("defg")}, dims=list("zyx"), ) + if backend == "dask": + return da.chunk() + elif backend == "numpy": + return da + else: + raise ValueError + @pytest.fixture def da_dask(seed=123): + # TODO: if possible, use the `da` fixture parameterized with backends rather than + # this. pytest.importorskip("dask.array") rs = np.random.RandomState(seed) times = pd.date_range("2000-01-01", freq="1D", periods=21) @@ -6596,6 +6611,7 @@ def test_rolling_properties(da): @pytest.mark.parametrize("name", ("sum", "mean", "std", "min", "max", "median")) @pytest.mark.parametrize("center", (True, False, None)) @pytest.mark.parametrize("min_periods", (1, None)) +@pytest.mark.parametrize("backend", ["numpy"], indirect=True) def test_rolling_wrapped_bottleneck(da, name, center, min_periods): bn = pytest.importorskip("bottleneck", minversion="1.1") @@ -6898,12 +6914,7 @@ def test_rolling_keep_attrs_deprecated(): data = np.linspace(10, 15, 100) coords = np.linspace(1, 10, 100) - da = DataArray( - data, - dims=("coord"), - coords={"coord": coords}, - attrs=attrs_da, - ) + da = DataArray(data, dims=("coord"), coords={"coord": coords}, attrs=attrs_da) # deprecated option with pytest.warns( @@ -7206,6 +7217,7 @@ def test_fallback_to_iris_AuxCoord(self, coord_values): @pytest.mark.parametrize( "window_type, window", [["span", 5], ["alpha", 0.5], ["com", 0.5], ["halflife", 5]] ) +@pytest.mark.parametrize("backend", ["numpy"], indirect=True) def test_rolling_exp(da, dim, window_type, window): da = da.isel(a=0) da = da.where(da > 0.2) @@ -7225,6 +7237,7 @@ def test_rolling_exp(da, dim, window_type, window): @requires_numbagg +@pytest.mark.parametrize("backend", ["numpy"], indirect=True) def test_rolling_exp_keep_attrs(da): attrs = {"attrs": "da"} da.attrs = attrs @@ -7313,3 +7326,35 @@ def test_deepcopy_obj_array(): x0 = DataArray(np.array([object()])) x1 = deepcopy(x0) assert x0.values[0] is not x1.values[0] + + +def test_clip(da): + with raise_if_dask_computes(): + result = da.clip(min=0.5) + assert result.min(...) >= 0.5 + + result = da.clip(max=0.5) + assert result.max(...) <= 0.5 + + result = da.clip(min=0.25, max=0.75) + assert result.min(...) >= 0.25 + assert result.max(...) <= 0.75 + + with raise_if_dask_computes(): + result = da.clip(min=da.mean("x"), max=da.mean("a")) + assert result.dims == da.dims + assert_array_equal( + result.data, + np.clip(da.data, da.mean("x").data[:, :, np.newaxis], da.mean("a").data), + ) + + with_nans = da.isel(time=[0, 1]).reindex_like(da) + with raise_if_dask_computes(): + result = da.clip(min=da.mean("x"), max=da.mean("a")) + result = da.clip(with_nans) + # The values should be the same where there were NaNs. + assert_array_equal(result.isel(time=[0, 1]), with_nans.isel(time=[0, 1])) + + # Unclear whether we want this work, OK to adjust the test when we have decided. + with pytest.raises(ValueError, match="arguments without labels along dimension"): + result = da.clip(min=da.mean("x"), max=da.mean("a").isel(x=[0, 1])) diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py index fb15f15df1a..96570a8415b 100644 --- a/xarray/tests/test_dataset.py +++ b/xarray/tests/test_dataset.py @@ -6069,16 +6069,16 @@ def test_dir_unicode(data_set): def ds(request): if request.param == 1: return Dataset( - { - "z1": (["y", "x"], np.random.randn(2, 8)), - "z2": (["time", "y"], np.random.randn(10, 2)), - }, - { - "x": ("x", np.linspace(0, 1.0, 8)), - "time": ("time", np.linspace(0, 1.0, 10)), - "c": ("y", ["a", "b"]), - "y": range(2), - }, + dict( + z1=(["y", "x"], np.random.randn(2, 8)), + z2=(["time", "y"], np.random.randn(10, 2)), + ), + dict( + x=("x", np.linspace(0, 1.0, 8)), + time=("time", np.linspace(0, 1.0, 10)), + c=("y", ["a", "b"]), + y=range(2), + ), ) if request.param == 2: @@ -6845,3 +6845,18 @@ def test_deepcopy_obj_array(): x0 = Dataset(dict(foo=DataArray(np.array([object()])))) x1 = deepcopy(x0) assert x0["foo"].values[0] is not x1["foo"].values[0] + + +def test_clip(ds): + result = ds.clip(min=0.5) + assert result.min(...) >= 0.5 + + result = ds.clip(max=0.5) + assert result.max(...) <= 0.5 + + result = ds.clip(min=0.25, max=0.75) + assert result.min(...) >= 0.25 + assert result.max(...) <= 0.75 + + result = ds.clip(min=ds.mean("y"), max=ds.mean("y")) + assert result.dims == ds.dims diff --git a/xarray/tests/test_variable.py b/xarray/tests/test_variable.py index c951e7d3118..5e8e02a82f8 100644 --- a/xarray/tests/test_variable.py +++ b/xarray/tests/test_variable.py @@ -48,6 +48,11 @@ ] +@pytest.fixture +def var(): + return Variable(dims=list("xyz"), data=np.random.rand(3, 4, 5)) + + class VariableSubclassobjects: def test_properties(self): data = 0.5 * np.arange(10) @@ -2510,3 +2515,27 @@ def test_DaskIndexingAdapter(self): v = Variable(dims=("x", "y"), data=CopyOnWriteArray(DaskIndexingAdapter(da))) self.check_orthogonal_indexing(v) self.check_vectorized_indexing(v) + + +def test_clip(var): + # Copied from test_dataarray (would there be a way to combine the tests?) + result = var.clip(min=0.5) + assert result.min(...) >= 0.5 + + result = var.clip(max=0.5) + assert result.max(...) <= 0.5 + + result = var.clip(min=0.25, max=0.75) + assert result.min(...) >= 0.25 + assert result.max(...) <= 0.75 + + result = var.clip(min=var.mean("x"), max=var.mean("z")) + assert result.dims == var.dims + assert_array_equal( + result.data, + np.clip( + var.data, + var.mean("x").data[np.newaxis, :, :], + var.mean("z").data[:, :, np.newaxis], + ), + ) diff --git a/xarray/util/generate_ops.py b/xarray/util/generate_ops.py index 9cd5dd1b8f5..b6b7f8cbac7 100644 --- a/xarray/util/generate_ops.py +++ b/xarray/util/generate_ops.py @@ -66,7 +66,6 @@ OTHER_UNARY_METHODS = ( ("round", "ops.round_"), ("argsort", "ops.argsort"), - ("clip", "ops.clip"), ("conj", "ops.conj"), ("conjugate", "ops.conjugate"), )