Как объединить массив записей NumPy (сумма, мин, макс и т. Д.)?

Question

Рассмотрим простую структуру записи массива:

import numpy as np 
ijv_dtype = [ 
    ('I', 'i'), 
    ('J', 'i'), 
    ('v', 'd'), 
] 
ijv = np.array([ 
    (0, 0, 3.3), 
    (0, 1, 1.1), 
    (0, 1, 4.4), 
    (1, 1, 2.2), 
    ], ijv_dtype) 
print(ijv) # [(0, 0, 3.3) (0, 1, 1.1) (0, 1, 4.4) (1, 1, 2.2)] 

я хотел бы aggregate некоторые статистические данные (суммы, мин, макс, и т.д.) из v путем группировки уникальных комбинаций I и J. Думая из SQL, ожидаемый результат:

select i, j, sum(v) as v from ijv group by i, j; 
i | j | v 
---+---+----- 
0 | 0 | 3.3 
0 | 1 | 5.5 
1 | 1 | 2.2 

(порядок не важен)

Лучшее, что я могу придумать для NumPy некрасиво, и я не уверен, что я заказал результат правильно (хотя, кажется, здесь работает):

# Get unique groups, index and inverse 
u_ij, idx_ij, inv_ij = np.unique(ijv[['I', 'J']], return_index=True, return_inverse=True) 
# Assemble aggregate 
a_ijv = np.zeros(len(u_ij), ijv_dtype) 
a_ijv['I'] = u_ij['I'] 
a_ijv['J'] = u_ij['J'] 
a_ijv['v'] = [ijv['v'][inv_ij == i].sum() for i in range(len(u_ij))] 
print(a_ijv) # [(0, 0, 3.3) (0, 1, 5.5) (1, 1, 2.2)] 

Хотелось бы подумать, что есть лучший способ сделать это! Я использую NumPy 1.4.1.

Answer 1

numpy является слишком низким уровнем для таких задач. Я думаю, что ваше решение хорошо, если вы должны использовать чистую numpy, но если вы не возражаете, используя что-то с более высоким уровнем абстракции, попробуйте pandas:

import pandas as pd 

df = pd.DataFrame({ 
    'I': (0, 0, 0, 1), 
    'J': (0, 1, 1, 1), 
    'v': (3.3, 1.1, 4.4, 2.2)}) 

print(df) 
print(df.groupby(['I', 'J']).sum()) 

Выход:

I J v 
0 0 0 3.3 
1 0 1 1.1 
2 0 1 4.4 
3 1 1 2.2 
     v 
I J  
0 0 3.3 
    1 5.5 
1 1 2.2