254 lines
11 KiB
Python
254 lines
11 KiB
Python
import pytest
|
|
import numpy as np
|
|
import pandas as pd
|
|
from pandas import DataFrame, concat
|
|
from pandas.util import testing as tm
|
|
|
|
|
|
def test_rank_apply():
|
|
lev1 = tm.rands_array(10, 100)
|
|
lev2 = tm.rands_array(10, 130)
|
|
lab1 = np.random.randint(0, 100, size=500)
|
|
lab2 = np.random.randint(0, 130, size=500)
|
|
|
|
df = DataFrame({'value': np.random.randn(500),
|
|
'key1': lev1.take(lab1),
|
|
'key2': lev2.take(lab2)})
|
|
|
|
result = df.groupby(['key1', 'key2']).value.rank()
|
|
|
|
expected = []
|
|
for key, piece in df.groupby(['key1', 'key2']):
|
|
expected.append(piece.value.rank())
|
|
expected = concat(expected, axis=0)
|
|
expected = expected.reindex(result.index)
|
|
tm.assert_series_equal(result, expected)
|
|
|
|
result = df.groupby(['key1', 'key2']).value.rank(pct=True)
|
|
|
|
expected = []
|
|
for key, piece in df.groupby(['key1', 'key2']):
|
|
expected.append(piece.value.rank(pct=True))
|
|
expected = concat(expected, axis=0)
|
|
expected = expected.reindex(result.index)
|
|
tm.assert_series_equal(result, expected)
|
|
|
|
|
|
@pytest.mark.parametrize("grps", [
|
|
['qux'], ['qux', 'quux']])
|
|
@pytest.mark.parametrize("vals", [
|
|
[2, 2, 8, 2, 6],
|
|
[pd.Timestamp('2018-01-02'), pd.Timestamp('2018-01-02'),
|
|
pd.Timestamp('2018-01-08'), pd.Timestamp('2018-01-02'),
|
|
pd.Timestamp('2018-01-06')]])
|
|
@pytest.mark.parametrize("ties_method,ascending,pct,exp", [
|
|
('average', True, False, [2., 2., 5., 2., 4.]),
|
|
('average', True, True, [0.4, 0.4, 1.0, 0.4, 0.8]),
|
|
('average', False, False, [4., 4., 1., 4., 2.]),
|
|
('average', False, True, [.8, .8, .2, .8, .4]),
|
|
('min', True, False, [1., 1., 5., 1., 4.]),
|
|
('min', True, True, [0.2, 0.2, 1.0, 0.2, 0.8]),
|
|
('min', False, False, [3., 3., 1., 3., 2.]),
|
|
('min', False, True, [.6, .6, .2, .6, .4]),
|
|
('max', True, False, [3., 3., 5., 3., 4.]),
|
|
('max', True, True, [0.6, 0.6, 1.0, 0.6, 0.8]),
|
|
('max', False, False, [5., 5., 1., 5., 2.]),
|
|
('max', False, True, [1., 1., .2, 1., .4]),
|
|
('first', True, False, [1., 2., 5., 3., 4.]),
|
|
('first', True, True, [0.2, 0.4, 1.0, 0.6, 0.8]),
|
|
('first', False, False, [3., 4., 1., 5., 2.]),
|
|
('first', False, True, [.6, .8, .2, 1., .4]),
|
|
('dense', True, False, [1., 1., 3., 1., 2.]),
|
|
('dense', True, True, [1. / 3., 1. / 3., 3. / 3., 1. / 3., 2. / 3.]),
|
|
('dense', False, False, [3., 3., 1., 3., 2.]),
|
|
('dense', False, True, [3. / 3., 3. / 3., 1. / 3., 3. / 3., 2. / 3.]),
|
|
])
|
|
def test_rank_args(grps, vals, ties_method, ascending, pct, exp):
|
|
key = np.repeat(grps, len(vals))
|
|
vals = vals * len(grps)
|
|
df = DataFrame({'key': key, 'val': vals})
|
|
result = df.groupby('key').rank(method=ties_method,
|
|
ascending=ascending, pct=pct)
|
|
|
|
exp_df = DataFrame(exp * len(grps), columns=['val'])
|
|
tm.assert_frame_equal(result, exp_df)
|
|
|
|
|
|
@pytest.mark.parametrize("grps", [
|
|
['qux'], ['qux', 'quux']])
|
|
@pytest.mark.parametrize("vals", [
|
|
[-np.inf, -np.inf, np.nan, 1., np.nan, np.inf, np.inf],
|
|
])
|
|
@pytest.mark.parametrize("ties_method,ascending,na_option,exp", [
|
|
('average', True, 'keep', [1.5, 1.5, np.nan, 3, np.nan, 4.5, 4.5]),
|
|
('average', True, 'top', [3.5, 3.5, 1.5, 5., 1.5, 6.5, 6.5]),
|
|
('average', True, 'bottom', [1.5, 1.5, 6.5, 3., 6.5, 4.5, 4.5]),
|
|
('average', False, 'keep', [4.5, 4.5, np.nan, 3, np.nan, 1.5, 1.5]),
|
|
('average', False, 'top', [6.5, 6.5, 1.5, 5., 1.5, 3.5, 3.5]),
|
|
('average', False, 'bottom', [4.5, 4.5, 6.5, 3., 6.5, 1.5, 1.5]),
|
|
('min', True, 'keep', [1., 1., np.nan, 3., np.nan, 4., 4.]),
|
|
('min', True, 'top', [3., 3., 1., 5., 1., 6., 6.]),
|
|
('min', True, 'bottom', [1., 1., 6., 3., 6., 4., 4.]),
|
|
('min', False, 'keep', [4., 4., np.nan, 3., np.nan, 1., 1.]),
|
|
('min', False, 'top', [6., 6., 1., 5., 1., 3., 3.]),
|
|
('min', False, 'bottom', [4., 4., 6., 3., 6., 1., 1.]),
|
|
('max', True, 'keep', [2., 2., np.nan, 3., np.nan, 5., 5.]),
|
|
('max', True, 'top', [4., 4., 2., 5., 2., 7., 7.]),
|
|
('max', True, 'bottom', [2., 2., 7., 3., 7., 5., 5.]),
|
|
('max', False, 'keep', [5., 5., np.nan, 3., np.nan, 2., 2.]),
|
|
('max', False, 'top', [7., 7., 2., 5., 2., 4., 4.]),
|
|
('max', False, 'bottom', [5., 5., 7., 3., 7., 2., 2.]),
|
|
('first', True, 'keep', [1., 2., np.nan, 3., np.nan, 4., 5.]),
|
|
('first', True, 'top', [3., 4., 1., 5., 2., 6., 7.]),
|
|
('first', True, 'bottom', [1., 2., 6., 3., 7., 4., 5.]),
|
|
('first', False, 'keep', [4., 5., np.nan, 3., np.nan, 1., 2.]),
|
|
('first', False, 'top', [6., 7., 1., 5., 2., 3., 4.]),
|
|
('first', False, 'bottom', [4., 5., 6., 3., 7., 1., 2.]),
|
|
('dense', True, 'keep', [1., 1., np.nan, 2., np.nan, 3., 3.]),
|
|
('dense', True, 'top', [2., 2., 1., 3., 1., 4., 4.]),
|
|
('dense', True, 'bottom', [1., 1., 4., 2., 4., 3., 3.]),
|
|
('dense', False, 'keep', [3., 3., np.nan, 2., np.nan, 1., 1.]),
|
|
('dense', False, 'top', [4., 4., 1., 3., 1., 2., 2.]),
|
|
('dense', False, 'bottom', [3., 3., 4., 2., 4., 1., 1.])
|
|
])
|
|
def test_infs_n_nans(grps, vals, ties_method, ascending, na_option, exp):
|
|
# GH 20561
|
|
key = np.repeat(grps, len(vals))
|
|
vals = vals * len(grps)
|
|
df = DataFrame({'key': key, 'val': vals})
|
|
result = df.groupby('key').rank(method=ties_method,
|
|
ascending=ascending,
|
|
na_option=na_option)
|
|
exp_df = DataFrame(exp * len(grps), columns=['val'])
|
|
tm.assert_frame_equal(result, exp_df)
|
|
|
|
|
|
@pytest.mark.parametrize("grps", [
|
|
['qux'], ['qux', 'quux']])
|
|
@pytest.mark.parametrize("vals", [
|
|
[2, 2, np.nan, 8, 2, 6, np.nan, np.nan],
|
|
[pd.Timestamp('2018-01-02'), pd.Timestamp('2018-01-02'), np.nan,
|
|
pd.Timestamp('2018-01-08'), pd.Timestamp('2018-01-02'),
|
|
pd.Timestamp('2018-01-06'), np.nan, np.nan]
|
|
])
|
|
@pytest.mark.parametrize("ties_method,ascending,na_option,pct,exp", [
|
|
('average', True, 'keep', False,
|
|
[2., 2., np.nan, 5., 2., 4., np.nan, np.nan]),
|
|
('average', True, 'keep', True,
|
|
[0.4, 0.4, np.nan, 1.0, 0.4, 0.8, np.nan, np.nan]),
|
|
('average', False, 'keep', False,
|
|
[4., 4., np.nan, 1., 4., 2., np.nan, np.nan]),
|
|
('average', False, 'keep', True,
|
|
[.8, 0.8, np.nan, 0.2, 0.8, 0.4, np.nan, np.nan]),
|
|
('min', True, 'keep', False,
|
|
[1., 1., np.nan, 5., 1., 4., np.nan, np.nan]),
|
|
('min', True, 'keep', True,
|
|
[0.2, 0.2, np.nan, 1.0, 0.2, 0.8, np.nan, np.nan]),
|
|
('min', False, 'keep', False,
|
|
[3., 3., np.nan, 1., 3., 2., np.nan, np.nan]),
|
|
('min', False, 'keep', True,
|
|
[.6, 0.6, np.nan, 0.2, 0.6, 0.4, np.nan, np.nan]),
|
|
('max', True, 'keep', False,
|
|
[3., 3., np.nan, 5., 3., 4., np.nan, np.nan]),
|
|
('max', True, 'keep', True,
|
|
[0.6, 0.6, np.nan, 1.0, 0.6, 0.8, np.nan, np.nan]),
|
|
('max', False, 'keep', False,
|
|
[5., 5., np.nan, 1., 5., 2., np.nan, np.nan]),
|
|
('max', False, 'keep', True,
|
|
[1., 1., np.nan, 0.2, 1., 0.4, np.nan, np.nan]),
|
|
('first', True, 'keep', False,
|
|
[1., 2., np.nan, 5., 3., 4., np.nan, np.nan]),
|
|
('first', True, 'keep', True,
|
|
[0.2, 0.4, np.nan, 1.0, 0.6, 0.8, np.nan, np.nan]),
|
|
('first', False, 'keep', False,
|
|
[3., 4., np.nan, 1., 5., 2., np.nan, np.nan]),
|
|
('first', False, 'keep', True,
|
|
[.6, 0.8, np.nan, 0.2, 1., 0.4, np.nan, np.nan]),
|
|
('dense', True, 'keep', False,
|
|
[1., 1., np.nan, 3., 1., 2., np.nan, np.nan]),
|
|
('dense', True, 'keep', True,
|
|
[1. / 3., 1. / 3., np.nan, 3. / 3., 1. / 3., 2. / 3., np.nan, np.nan]),
|
|
('dense', False, 'keep', False,
|
|
[3., 3., np.nan, 1., 3., 2., np.nan, np.nan]),
|
|
('dense', False, 'keep', True,
|
|
[3. / 3., 3. / 3., np.nan, 1. / 3., 3. / 3., 2. / 3., np.nan, np.nan]),
|
|
('average', True, 'no_na', False, [2., 2., 7., 5., 2., 4., 7., 7.]),
|
|
('average', True, 'no_na', True,
|
|
[0.25, 0.25, 0.875, 0.625, 0.25, 0.5, 0.875, 0.875]),
|
|
('average', False, 'no_na', False, [4., 4., 7., 1., 4., 2., 7., 7.]),
|
|
('average', False, 'no_na', True,
|
|
[0.5, 0.5, 0.875, 0.125, 0.5, 0.25, 0.875, 0.875]),
|
|
('min', True, 'no_na', False, [1., 1., 6., 5., 1., 4., 6., 6.]),
|
|
('min', True, 'no_na', True,
|
|
[0.125, 0.125, 0.75, 0.625, 0.125, 0.5, 0.75, 0.75]),
|
|
('min', False, 'no_na', False, [3., 3., 6., 1., 3., 2., 6., 6.]),
|
|
('min', False, 'no_na', True,
|
|
[0.375, 0.375, 0.75, 0.125, 0.375, 0.25, 0.75, 0.75]),
|
|
('max', True, 'no_na', False, [3., 3., 8., 5., 3., 4., 8., 8.]),
|
|
('max', True, 'no_na', True,
|
|
[0.375, 0.375, 1., 0.625, 0.375, 0.5, 1., 1.]),
|
|
('max', False, 'no_na', False, [5., 5., 8., 1., 5., 2., 8., 8.]),
|
|
('max', False, 'no_na', True,
|
|
[0.625, 0.625, 1., 0.125, 0.625, 0.25, 1., 1.]),
|
|
('first', True, 'no_na', False, [1., 2., 6., 5., 3., 4., 7., 8.]),
|
|
('first', True, 'no_na', True,
|
|
[0.125, 0.25, 0.75, 0.625, 0.375, 0.5, 0.875, 1.]),
|
|
('first', False, 'no_na', False, [3., 4., 6., 1., 5., 2., 7., 8.]),
|
|
('first', False, 'no_na', True,
|
|
[0.375, 0.5, 0.75, 0.125, 0.625, 0.25, 0.875, 1.]),
|
|
('dense', True, 'no_na', False, [1., 1., 4., 3., 1., 2., 4., 4.]),
|
|
('dense', True, 'no_na', True,
|
|
[0.25, 0.25, 1., 0.75, 0.25, 0.5, 1., 1.]),
|
|
('dense', False, 'no_na', False, [3., 3., 4., 1., 3., 2., 4., 4.]),
|
|
('dense', False, 'no_na', True,
|
|
[0.75, 0.75, 1., 0.25, 0.75, 0.5, 1., 1.])
|
|
])
|
|
def test_rank_args_missing(grps, vals, ties_method, ascending,
|
|
na_option, pct, exp):
|
|
key = np.repeat(grps, len(vals))
|
|
vals = vals * len(grps)
|
|
df = DataFrame({'key': key, 'val': vals})
|
|
result = df.groupby('key').rank(method=ties_method,
|
|
ascending=ascending,
|
|
na_option=na_option, pct=pct)
|
|
|
|
exp_df = DataFrame(exp * len(grps), columns=['val'])
|
|
tm.assert_frame_equal(result, exp_df)
|
|
|
|
|
|
@pytest.mark.parametrize("pct,exp", [
|
|
(False, [3., 3., 3., 3., 3.]),
|
|
(True, [.6, .6, .6, .6, .6])])
|
|
def test_rank_resets_each_group(pct, exp):
|
|
df = DataFrame(
|
|
{'key': ['a', 'a', 'a', 'a', 'a', 'b', 'b', 'b', 'b', 'b'],
|
|
'val': [1] * 10}
|
|
)
|
|
result = df.groupby('key').rank(pct=pct)
|
|
exp_df = DataFrame(exp * 2, columns=['val'])
|
|
tm.assert_frame_equal(result, exp_df)
|
|
|
|
|
|
def test_rank_avg_even_vals():
|
|
df = DataFrame({'key': ['a'] * 4, 'val': [1] * 4})
|
|
result = df.groupby('key').rank()
|
|
exp_df = DataFrame([2.5, 2.5, 2.5, 2.5], columns=['val'])
|
|
tm.assert_frame_equal(result, exp_df)
|
|
|
|
|
|
@pytest.mark.parametrize("ties_method", [
|
|
'average', 'min', 'max', 'first', 'dense'])
|
|
@pytest.mark.parametrize("ascending", [True, False])
|
|
@pytest.mark.parametrize("na_option", ["keep", "top", "bottom"])
|
|
@pytest.mark.parametrize("pct", [True, False])
|
|
@pytest.mark.parametrize("vals", [
|
|
['bar', 'bar', 'foo', 'bar', 'baz'],
|
|
['bar', np.nan, 'foo', np.nan, 'baz']
|
|
])
|
|
def test_rank_object_raises(ties_method, ascending, na_option,
|
|
pct, vals):
|
|
df = DataFrame({'key': ['foo'] * 5, 'val': vals})
|
|
with tm.assert_raises_regex(TypeError, "not callable"):
|
|
df.groupby('key').rank(method=ties_method,
|
|
ascending=ascending,
|
|
na_option=na_option, pct=pct)
|