From ca993090d5f0e8aab9199b22eca4038a039bddd6 Mon Sep 17 00:00:00 2001 From: Maximilian Roos Date: Wed, 4 Oct 2023 12:03:43 -0700 Subject: [PATCH 1/6] Allow a function in `.sortby` method --- doc/whats-new.rst | 5 ++++- xarray/core/common.py | 3 ++- xarray/core/dataarray.py | 28 ++++++++++++++++++++-------- xarray/core/dataset.py | 23 +++++++++++++++++------ 4 files changed, 43 insertions(+), 16 deletions(-) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index e73a1a7fa62..c7f2d1aca26 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -23,9 +23,12 @@ New Features ~~~~~~~~~~~~ - :py:meth:`DataArray.where` & :py:meth:`Dataset.where` accept a callable for - the ``other`` parameter, passing the object as the first argument. Previously, + the ``other`` parameter, passing the object as the only argument. Previously, this was only valid for the ``cond`` parameter. (:issue:`8255`) By `Maximilian Roos `_. +- :py:meth:`DataArray.sortby` & :py:meth:`Dataset.sortby` accept a callable for + the ``variables`` parameter, passing the object as the only argument. + By `Maximilian Roos `_. Breaking changes ~~~~~~~~~~~~~~~~ diff --git a/xarray/core/common.py b/xarray/core/common.py index 2a4c4c200d4..f571576850c 100644 --- a/xarray/core/common.py +++ b/xarray/core/common.py @@ -1073,7 +1073,8 @@ def where(self, cond: Any, other: Any = dtypes.NA, drop: bool = False) -> Self: ---------- cond : DataArray, Dataset, or callable Locations at which to preserve this object's values. dtype must be `bool`. - If a callable, it must expect this object as its only parameter. + If a callable, the callable is passed this object, and the result is used as + the value for cond. other : scalar, DataArray, Dataset, or callable, optional Value to use for locations in this object where ``cond`` is False. By default, these locations are filled with NA. If a callable, it must diff --git a/xarray/core/dataarray.py b/xarray/core/dataarray.py index ef4389f3c6c..617eebce33b 100644 --- a/xarray/core/dataarray.py +++ b/xarray/core/dataarray.py @@ -4921,7 +4921,7 @@ def dot( def sortby( self, - variables: Hashable | DataArray | Sequence[Hashable | DataArray], + variables: Hashable | DataArray | Sequence[Hashable | DataArray] | Callable, ascending: bool = True, ) -> Self: """Sort object by labels or values (along an axis). @@ -4942,9 +4942,10 @@ def sortby( Parameters ---------- - variables : Hashable, DataArray, or sequence of Hashable or DataArray - 1D DataArray objects or name(s) of 1D variable(s) in - coords whose values are used to sort this array. + variables : Hashable, DataArray, sequence of Hashable or DataArray, or Callable + 1D DataArray objects or name(s) of 1D variable(s) in coords whose values are + used to sort this array. If a callable, the callable is passed this object, + and the result is used as the value for cond. ascending : bool, default: True Whether to sort by ascending or descending order. @@ -4964,22 +4965,33 @@ def sortby( Examples -------- >>> da = xr.DataArray( - ... np.random.rand(5), + ... np.arange(5,0,-1), ... coords=[pd.date_range("1/1/2000", periods=5)], ... dims="time", ... ) >>> da - array([0.5488135 , 0.71518937, 0.60276338, 0.54488318, 0.4236548 ]) + array([5, 4, 3, 2, 1]) Coordinates: * time (time) datetime64[ns] 2000-01-01 2000-01-02 ... 2000-01-05 >>> da.sortby(da) - array([0.4236548 , 0.54488318, 0.5488135 , 0.60276338, 0.71518937]) + array([1, 2, 3, 4, 5]) Coordinates: - * time (time) datetime64[ns] 2000-01-05 2000-01-04 ... 2000-01-02 + * time (time) datetime64[ns] 2000-01-05 2000-01-04 ... 2000-01-01 + + >>> da.sortby(lambda x: x) + + array([1, 2, 3, 4, 5]) + Coordinates: + * time (time) datetime64[ns] 2000-01-05 2000-01-04 ... 2000-01-01 """ + # We need to convert the callable here rather than pass it through to the + # dataset method, since otherwise the dataset method would try to call the + # callable with the dataset as the object + if callable(variables): + variables = variables(self) ds = self._to_temp_dataset().sortby(variables, ascending=ascending) return self._from_temp_dataset(ds) diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py index 459e2f3fce7..f47855fd1fc 100644 --- a/xarray/core/dataset.py +++ b/xarray/core/dataset.py @@ -7824,7 +7824,7 @@ def roll( def sortby( self, - variables: Hashable | DataArray | list[Hashable | DataArray], + variables: Hashable | DataArray | list[Hashable | DataArray] | Callable, ascending: bool = True, ) -> Self: """ @@ -7846,9 +7846,10 @@ def sortby( Parameters ---------- - variables : Hashable, DataArray, or list of hashable or DataArray - 1D DataArray objects or name(s) of 1D variable(s) in - coords/data_vars whose values are used to sort the dataset. + kariables : Hashable, DataArray, sequence of Hashable or DataArray, or Callable + 1D DataArray objects or name(s) of 1D variable(s) in coords whose values are + used to sort this array. If a callable, the callable is passed this object, + and the result is used as the value for cond. ascending : bool, default: True Whether to sort by ascending or descending order. @@ -7874,8 +7875,7 @@ def sortby( ... }, ... coords={"x": ["b", "a"], "y": [1, 0]}, ... ) - >>> ds = ds.sortby("x") - >>> ds + >>> ds.sortby("x") Dimensions: (x: 2, y: 2) Coordinates: @@ -7884,9 +7884,20 @@ def sortby( Data variables: A (x, y) int64 3 4 1 2 B (x, y) int64 7 8 5 6 + >>> ds.sortby(lambda x: -x['y']) + + Dimensions: (x: 2, y: 2) + Coordinates: + * x (x) Date: Wed, 4 Oct 2023 19:08:30 +0000 Subject: [PATCH 2/6] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- xarray/core/dataarray.py | 2 +- xarray/core/dataset.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/xarray/core/dataarray.py b/xarray/core/dataarray.py index 617eebce33b..2d64559479c 100644 --- a/xarray/core/dataarray.py +++ b/xarray/core/dataarray.py @@ -4965,7 +4965,7 @@ def sortby( Examples -------- >>> da = xr.DataArray( - ... np.arange(5,0,-1), + ... np.arange(5, 0, -1), ... coords=[pd.date_range("1/1/2000", periods=5)], ... dims="time", ... ) diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py index f47855fd1fc..58140cb7940 100644 --- a/xarray/core/dataset.py +++ b/xarray/core/dataset.py @@ -7884,7 +7884,7 @@ def sortby( Data variables: A (x, y) int64 3 4 1 2 B (x, y) int64 7 8 5 6 - >>> ds.sortby(lambda x: -x['y']) + >>> ds.sortby(lambda x: -x["y"]) Dimensions: (x: 2, y: 2) Coordinates: From b06884424c695f990ffee47a7c8fc4d9552d06fb Mon Sep 17 00:00:00 2001 From: Maximilian Roos <5635139+max-sixty@users.noreply.github.com> Date: Thu, 5 Oct 2023 12:12:32 -0700 Subject: [PATCH 3/6] Update dataarray.py Co-authored-by: Michael Niklas --- xarray/core/dataarray.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xarray/core/dataarray.py b/xarray/core/dataarray.py index 2d64559479c..9606ce02a73 100644 --- a/xarray/core/dataarray.py +++ b/xarray/core/dataarray.py @@ -4921,7 +4921,7 @@ def dot( def sortby( self, - variables: Hashable | DataArray | Sequence[Hashable | DataArray] | Callable, + variables: Hashable | DataArray | Sequence[Hashable | DataArray] | Callable[[Self], Hashable | DataArray | Sequence[Hashable | DataArray]], ascending: bool = True, ) -> Self: """Sort object by labels or values (along an axis). From 5525d41d36e2fded5225b5fa2adb76db00211113 Mon Sep 17 00:00:00 2001 From: Maximilian Roos <5635139+max-sixty@users.noreply.github.com> Date: Thu, 5 Oct 2023 12:12:40 -0700 Subject: [PATCH 4/6] Update dataset.py Co-authored-by: Michael Niklas --- xarray/core/dataset.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py index 58140cb7940..b8f16863c38 100644 --- a/xarray/core/dataset.py +++ b/xarray/core/dataset.py @@ -7824,7 +7824,7 @@ def roll( def sortby( self, - variables: Hashable | DataArray | list[Hashable | DataArray] | Callable, + variables: Hashable | DataArray | Sequence[Hashable | DataArray] | Callable[[Self], Hashable | DataArray | list[Hashable | DataArray]], ascending: bool = True, ) -> Self: """ From 78feaa5e2dd6723b8a31fc3e279d4d97e0d8ff63 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 5 Oct 2023 19:13:18 +0000 Subject: [PATCH 5/6] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- xarray/core/dataarray.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/xarray/core/dataarray.py b/xarray/core/dataarray.py index 9606ce02a73..8b3e999b78c 100644 --- a/xarray/core/dataarray.py +++ b/xarray/core/dataarray.py @@ -4921,7 +4921,10 @@ def dot( def sortby( self, - variables: Hashable | DataArray | Sequence[Hashable | DataArray] | Callable[[Self], Hashable | DataArray | Sequence[Hashable | DataArray]], + variables: Hashable + | DataArray + | Sequence[Hashable | DataArray] + | Callable[[Self], Hashable | DataArray | Sequence[Hashable | DataArray]], ascending: bool = True, ) -> Self: """Sort object by labels or values (along an axis). From d2c6bfd9eada82025ae432b8ba7039ab3f0ec698 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 5 Oct 2023 19:13:57 +0000 Subject: [PATCH 6/6] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- xarray/core/dataset.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py index b8f16863c38..533eed2c848 100644 --- a/xarray/core/dataset.py +++ b/xarray/core/dataset.py @@ -7824,7 +7824,10 @@ def roll( def sortby( self, - variables: Hashable | DataArray | Sequence[Hashable | DataArray] | Callable[[Self], Hashable | DataArray | list[Hashable | DataArray]], + variables: Hashable + | DataArray + | Sequence[Hashable | DataArray] + | Callable[[Self], Hashable | DataArray | list[Hashable | DataArray]], ascending: bool = True, ) -> Self: """