Ранжирование numpy массивов для возврата массивов рангов

Question

У меня есть трехмерный массив NumPy, представляющий n матриц измерений x, x. Например, при n = 3 и x = 2, примером будет:

matrices = np.array([[[2,8],[1,7]],[[3,1],[5,4]],[[9,6],[2,3]]]

matrices
array([[[2, 8],
        [1, 7]],

       [[3, 1],
        [5, 4]],

       [[9, 6],
        [2, 3]]])

Я хочу создать новый массив идентичной формы, но со значениями ранжирования этих матриц для каждого элемента (т.е. ранжирование по оси 0). Результаты будут следующими:

array([[[1, 3],
        [1, 3]],

       [[2, 1],
        [3, 2]],

       [[3, 2],
        [2, 1]]])

Я вижу, как сортировать (np.argsort (matrices, axis = 0)), но не могу найти простой способ вернуть рейтинг. Размеры матрицы х могут быть высокими, поэтому важно время выполнения. Кроме того, распределение python используется не включает в себя scipy.

Я обнаружил, что следующие работы:

def rank_stations(input):
    output = [0] * (np.size(input))
    for i, x in enumerate(sorted(range(len(input)), reverse=True, key=lambda y: input[y])):
        output[x] = i+1
    return output

results = np.apply_along_axis(rank_stations, 0, matrices)

С пятью матрицами измерения 980 x 980, это работает около 6 секунд по сравнению с 20 SE c с программным обеспечением моделирования моделирования, так что разумное улучшение. Алгоритм сортировки выполняется за доли секунды, но есть ли способ получить подобное время выполнения с ранжированием?

Answer 1

Используется numpy.apply_along_axis с scipy.stats.rankdata:

from scipy.stats import rankdata

np.apply_along_axis(rankdata, 0, matrices)

Или с использованием np.argsort дважды:

f = lambda x: x.argsort().argsort()
np.apply_along_axis(f, 0, matrices) + 1

Выход:

array([[[1., 3.],
        [1., 3.]],

       [[2., 1.],
        [3., 2.]],

       [[3., 2.],
        [2., 1.]]])