initial support for rank

Author: 0x0L

doc/api-hidden.rst

@@ -46,6 +46,7 @@
    Dataset.T
    Dataset.cumsum
    Dataset.cumprod
+   Dataset.rank
 
    DataArray.ndim
    DataArray.shape
@@ -91,6 +92,7 @@
    DataArray.T
    DataArray.cumsum
    DataArray.cumprod
+   DataArray.rank
 
    ufuncs.angle
    ufuncs.arccos

doc/api.rst

@@ -160,6 +160,7 @@ Computation
 :py:attr:`~Dataset.real`
 :py:attr:`~Dataset.cumsum`
 :py:attr:`~Dataset.cumprod`
+:py:attr:`~Dataset.rank`
 
 **Grouped operations**:
 :py:attr:`~core.groupby.DatasetGroupBy.assign`
@@ -312,6 +313,7 @@ Computation
 :py:attr:`~DataArray.T`
 :py:attr:`~DataArray.cumsum`
 :py:attr:`~DataArray.cumprod`
+:py:attr:`~DataArray.rank`
 
 **Grouped operations**:
 :py:attr:`~core.groupby.DataArrayGroupBy.assign_coords`

doc/whats-new.rst

@@ -208,6 +208,10 @@ Enhancements
   (:issue:`1485`).
   By `Joe Hamman <https://github.com/jhamman>`_.
 
+- New :py:meth:`~xarray.DataArray.rank` method to Datasets and DataArrays when
+  bottleneck is available (:issue:`1731`).
+  By `0x0L <https://github.com/0x0L>`_.
+
 **Performance improvements**
 
 - :py:func:`~xarray.concat` was computing variables that aren't in memory

xarray/core/duck_array_ops.py

@@ -23,6 +23,8 @@
 try:
     import bottleneck as bn
     has_bottleneck = True
+    # monkeypatch numpy with rankdata
+    np.rankdata = bn.rankdata
 except ImportError:
     # use numpy methods instead
     bn = np
@@ -240,6 +242,10 @@ def f(values, axis=None, skipna=None, **kwargs):
 cumsum = _create_nan_agg_method('cumsum', numeric_only=True, np_compat=True,
                                 no_bottleneck=True, keep_dims=True)
 
+if has_bottleneck:
+    rank = _create_nan_agg_method('rankdata', numeric_only=False, np_compat=True,
+                                  no_bottleneck=False, keep_dims=True)
+
 _fail_on_dask_array_input_skipna = partial(
     fail_on_dask_array_input,
     msg='%r with skipna=True is not yet implemented on dask arrays')

xarray/core/ops.py

@@ -51,7 +51,7 @@
                               'move_std': 'std', 'move_min': 'min',
                               'move_max': 'max', 'move_var': 'var',
                               'move_argmin': 'argmin', 'move_argmax': 'argmax',
-                              'move_median': 'median'}
+                              'move_median': 'median'}  #, 'move_rank': 'rank'}
 # TODO: wrap take, dot, sort
 
 
@@ -121,6 +121,31 @@
 """
 
 
+_BN_RANK_DOC = """\
+Ranks the data, dealing with ties and NaNs appropriately.
+
+Equal values are assigned a rank that is the average of the ranks that
+would have been otherwise assigned to all of the values within that set.
+Ranks begin at 1, not 0.
+
+NaNs in the input array are returned as NaNs.
+
+Parameters
+----------
+dim : str, optional
+    Dimension over which to apply `rank`.
+axis : int, optional
+    Axis over which to apply `rank`. Either the 'dim'
+    or the 'axis' argument must be supplied.
+
+Returns
+-------
+ranked: {cls}
+    New {cls} object with `rank` applied to its data along the
+    indicated dimension. The dtype is 'float64'.
+"""
+
+
 def fillna(data, other, join="left", dataset_join="left"):
     """Fill missing values in this object with data from the other object.
     Follows normal broadcasting and alignment rules.
@@ -305,6 +330,15 @@ def inject_binary_ops(cls, inplace=False):
                     cls._inplace_binary_op(get_op('i' + name)))
 
 
+def inject_bottleneck_rank(cls):
+    name = 'rank'
+    func = cls._reduce_method(getattr(duck_array_ops, name),
+                              include_skipna=False, numeric_only=False)
+    func.__name__ = name
+    func.__doc__ = _BN_RANK_DOC.format(cls=cls.__name__)
+    setattr(cls, name, func)
+
+
 def inject_all_ops_and_reduce_methods(cls, priority=50, array_only=True):
     # prioritize our operations over those of numpy.ndarray (priority=1)
     # and numpy.matrix (priority=10)
@@ -334,6 +368,8 @@ def inject_all_ops_and_reduce_methods(cls, priority=50, array_only=True):
 
     inject_reduce_methods(cls)
     inject_cum_methods(cls)
+    if has_bottleneck:
+        inject_bottleneck_rank(cls)
 
 
 def inject_bottleneck_rolling_methods(cls):

xarray/tests/test_dataarray.py

@@ -19,7 +19,7 @@
 from xarray.tests import (
     TestCase, ReturnItem, source_ndarray, unittest, requires_dask,
     assert_identical, assert_equal, assert_allclose, assert_array_equal,
-    raises_regex, requires_scipy)
+    raises_regex, requires_scipy, requires_bottleneck)
 
 
 class TestDataArray(TestCase):
@@ -2976,6 +2976,21 @@ def test_sortby(self):
         actual = da.sortby(['x', 'y'])
         self.assertDataArrayEqual(actual, expected)
 
+    @requires_bottleneck
+    def test_rank(self):
+        # floats
+        ar = DataArray([[3, 4, np.nan, 1]])
+        expect_0 = DataArray([[1, 1, np.nan, 1]])
+        expect_1 = DataArray([[2, 3, np.nan, 1]])
+        self.assertDataArrayEqual(ar.rank('dim_0'), expect_0)
+        self.assertDataArrayEqual(ar.rank('dim_1'), expect_1)
+        # int
+        x = DataArray([3,2,1])
+        self.assertDataArrayEqual(x.rank('dim_0'), x)
+        # str
+        y =  DataArray(['c', 'b', 'a'])
+        self.assertDataArrayEqual(y.rank('dim_0'), x)
+
 
 @pytest.fixture(params=[1])
 def da(request):