Allow where to receive a callable (#3827)

* allow where to receive a callable

* Update xarray/core/common.py

Co-Authored-By: keewis <[email protected]>

* docstring

* whatsnew

Co-authored-by: keewis <[email protected]>

Author: max-sixty

doc/whats-new.rst

@@ -43,7 +43,9 @@ New Features
   in 0.14.1) is now on by default. To disable, use
   ``xarray.set_options(display_style="text")``.
   By `Julia Signell <https://github.com/jsignell>`_.
-
+- :py:meth:`Dataset.where` and :py:meth:`DataArray.where` accept a lambda as a
+  first argument, which is then called on the input; replicating pandas' behavior.
+  By `Maximilian Roos <https://github.com/max-sixty>`_
 
 Bug fixes
 ~~~~~~~~~

xarray/core/common.py

@@ -1119,6 +1119,15 @@ def where(self, cond, other=dtypes.NA, drop: bool = False):
 
         >>> import numpy as np
         >>> a = xr.DataArray(np.arange(25).reshape(5, 5), dims=('x', 'y'))
+        >>> a
+        <xarray.DataArray (x: 5, y: 5)>
+        array([[ 0,  1,  2,  3,  4],
+            [ 5,  6,  7,  8,  9],
+            [10, 11, 12, 13, 14],
+            [15, 16, 17, 18, 19],
+            [20, 21, 22, 23, 24]])
+        Dimensions without coordinates: x, y
+
         >>> a.where(a.x + a.y < 4)
         <xarray.DataArray (x: 5, y: 5)>
         array([[  0.,   1.,   2.,   3.,  nan],
@@ -1127,6 +1136,7 @@ def where(self, cond, other=dtypes.NA, drop: bool = False):
                [ 15.,  nan,  nan,  nan,  nan],
                [ nan,  nan,  nan,  nan,  nan]])
         Dimensions without coordinates: x, y
+
         >>> a.where(a.x + a.y < 5, -1)
         <xarray.DataArray (x: 5, y: 5)>
         array([[ 0,  1,  2,  3,  4],
@@ -1135,6 +1145,7 @@ def where(self, cond, other=dtypes.NA, drop: bool = False):
                [15, 16, -1, -1, -1],
                [20, -1, -1, -1, -1]])
         Dimensions without coordinates: x, y
+
         >>> a.where(a.x + a.y < 4, drop=True)
         <xarray.DataArray (x: 4, y: 4)>
         array([[  0.,   1.,   2.,   3.],
@@ -1143,6 +1154,14 @@ def where(self, cond, other=dtypes.NA, drop: bool = False):
                [ 15.,  nan,  nan,  nan]])
         Dimensions without coordinates: x, y
 
+        >>> a.where(lambda x: x.x + x.y < 4, drop=True)
+        <xarray.DataArray (x: 4, y: 4)>
+        array([[  0.,   1.,   2.,   3.],
+               [  5.,   6.,   7.,  nan],
+               [ 10.,  11.,  nan,  nan],
+               [ 15.,  nan,  nan,  nan]])
+        Dimensions without coordinates: x, y
+
         See also
         --------
         numpy.where : corresponding numpy function
@@ -1152,6 +1171,9 @@ def where(self, cond, other=dtypes.NA, drop: bool = False):
         from .dataarray import DataArray
         from .dataset import Dataset
 
+        if callable(cond):
+            cond = cond(self)
+
         if drop:
             if other is not dtypes.NA:
                 raise ValueError("cannot set `other` if drop=True")

xarray/tests/test_dataarray.py

@@ -2215,6 +2215,12 @@ def test_where(self):
         actual = arr.where(arr.x < 2, drop=True)
         assert_identical(actual, expected)
 
+    def test_where_lambda(self):
+        arr = DataArray(np.arange(4), dims="y")
+        expected = arr.sel(y=slice(2))
+        actual = arr.where(lambda x: x.y < 2, drop=True)
+        assert_identical(actual, expected)
+
     def test_where_string(self):
         array = DataArray(["a", "b"])
         expected = DataArray(np.array(["a", np.nan], dtype=object))

xarray/tests/test_dataset.py

@@ -4349,13 +4349,22 @@ def test_where(self):
         assert actual.a.name == "a"
         assert actual.a.attrs == ds.a.attrs
 
+        # lambda
+        ds = Dataset({"a": ("x", range(5))})
+        expected = Dataset({"a": ("x", [np.nan, np.nan, 2, 3, 4])})
+        actual = ds.where(lambda x: x > 1)
+        assert_identical(expected, actual)
+
     def test_where_other(self):
         ds = Dataset({"a": ("x", range(5))}, {"x": range(5)})
         expected = Dataset({"a": ("x", [-1, -1, 2, 3, 4])}, {"x": range(5)})
         actual = ds.where(ds > 1, -1)
         assert_equal(expected, actual)
         assert actual.a.dtype == int
 
+        actual = ds.where(lambda x: x > 1, -1)
+        assert_equal(expected, actual)
+
         with raises_regex(ValueError, "cannot set"):
             ds.where(ds > 1, other=0, drop=True)