initial support for rank and ffill

Author: 0x0L

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(s) over which to apply `{name}`.
+axis : int, optional
+    Axis over which to apply `{name}`. Either the 'dim'
+    or the 'axis' argument must be supplied.
+
+Returns
+-------
+ranked: {cls}
+    New {cls} object with `{name}` 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,17 @@ def inject_binary_ops(cls, inplace=False):
                     cls._inplace_binary_op(get_op('i' + name)))
 
 
+def inject_bottleneck_rank(cls):
+    name = 'rank'
+    f = getattr(duck_array_ops, name)
+    include_skipna = False
+    numeric_only = getattr(f, 'numeric_only', False)
+    func = cls._reduce_method(f, include_skipna, numeric_only)
+    func.__name__ = name
+    func.__doc__ = _BN_RANK_DOC.format(name=name, 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 +370,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):