Skip to content

Commit

Permalink
initial support for rank
Browse files Browse the repository at this point in the history
  • Loading branch information
0x0L authored and 0x0L committed Nov 24, 2017
1 parent 4caae2e commit 8cc142b
Show file tree
Hide file tree
Showing 5 changed files with 58 additions and 1 deletion.
1 change: 1 addition & 0 deletions doc/api-hidden.rst
Original file line number Diff line number Diff line change
Expand Up @@ -91,6 +91,7 @@
DataArray.T
DataArray.cumsum
DataArray.cumprod
DataArray.rank

ufuncs.angle
ufuncs.arccos
Expand Down
1 change: 1 addition & 0 deletions doc/api.rst
Original file line number Diff line number Diff line change
Expand Up @@ -312,6 +312,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`
Expand Down
3 changes: 3 additions & 0 deletions doc/whats-new.rst
Original file line number Diff line number Diff line change
Expand Up @@ -208,6 +208,9 @@ Enhancements
(:issue:`1485`).
By `Joe Hamman <https://github.com/jhamman>`_.

- New :py:meth:`~xarray.DataArray.rank`. Requires bottleneck (:issue:`1731`).
By `0x0L <https://github.com/0x0L>`_.

**Performance improvements**

- :py:func:`~xarray.concat` was computing variables that aren't in memory
Expand Down
37 changes: 37 additions & 0 deletions xarray/core/dataarray.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
from .accessors import DatetimeAccessor
from .alignment import align, reindex_like_indexers
from .common import AbstractArray, BaseDataObject
from .computation import apply_ufunc
from .coordinates import (DataArrayCoordinates, LevelCoordinatesSource,
Indexes)
from .dataset import Dataset, merge_indexes, split_indexes
Expand Down Expand Up @@ -1958,6 +1959,42 @@ def quantile(self, q, dim=None, interpolation='linear', keep_attrs=False):
interpolation=interpolation)
return self._from_temp_dataset(ds)

def rank(self, dim):
"""Ranks the data.
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
Returns
-------
ranked : DataArray
DataArray with the same coordinates and dtype 'float64'.
Examples
--------
>>> arr = xr.DataArray([5, 6, 7], dims='x')
>>> arr.rank('x')
<xarray.DataArray (x: 3)>
array([ 1., 2., 3.])
Dimensions without coordinates: x
"""
import bottleneck as bn
axis = self.get_axis_num(dim)
func = bn.nanrankdata if self.dtype.kind is 'f' else bn.rankdata
return apply_ufunc(func, self,
dask='parallelized',
keep_attrs=True,
output_dtypes=[np.float_],
kwargs=dict(axis=axis)).transpose(*self.dims)


# priority most be higher than Variable to properly work with binary ufuncs
ops.inject_all_ops_and_reduce_methods(DataArray, priority=60)
17 changes: 16 additions & 1 deletion xarray/tests/test_dataarray.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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):
Expand Down

0 comments on commit 8cc142b

Please sign in to comment.