Функциональное назначение в NumPy

Question

Предположим, у меня есть два массива

A = [ 6, 4, 5, 7, 9 ]
ind = [ 0, 0, 2, 1, 2 ]

и функция f.

Я хочу построить новый массив B размером с число различных элементов в ind с B [i] результатом f с параметром подмассива A, проиндексированного i.

Для этого примера, если я возьму f = сумму, тогда

B = [10, 7, 14]

или f = max

B = [6, 7, 9]

Есть ли более эффективный способ, чем цикл for в numpy?

Спасибо

Answer 1

Для особого случая f = sum:

In [32]: np.bincount(ind,A)
Out[32]: array([ 10.,   7.,  14.])

---

Предполагается:

• f - это ufunc
• У вас достаточнопамять для создания двумерного массива формы len(A) x len(A)

Вы можете создать двумерный массив B:

B=np.zeros((len(A),max(ind)+1))

и заполнить различные места в B значениямиот A, так что первый столбец B получает значения только от A, когда ind == 0, а второй столбец B получает значения только от A, когда ind == 1 и т. д .:

B[zip(*enumerate(ind))]=A

в результате вы получите массив типа

[[ 6.  0.  0.]
 [ 4.  0.  0.]
 [ 0.  0.  5.]
 [ 0.  7.  0.]
 [ 0.  0.  9.]]

Затем вы можете применить f вдоль оси = 0, чтобы получить желаемый результат.Здесь используется третье предположение:

• Дополнительные нули в B не влияют на желаемый результат.

Если вы можете принять эти предположения, то:

import numpy as np

A = np.array([ 6, 4, 5, 7, 9 ])
ind = np.array([ 0, 0, 2, 1, 2 ])

N=100
M=10
A2 = np.array([np.random.randint(M) for i in range(N)])
ind2 = np.array([np.random.randint(M) for i in range(N)])

def use_extra_axis(A,ind,f):
    B=np.zeros((len(A),max(ind)+1))
    B[zip(*enumerate(ind))]=A
    return f(B)

def use_loop(A,ind,f):
    n=max(ind)+1
    B=np.empty(n)
    for i in range(n):
        B[i]=f(A[ind==i])
    return B

def fmax(arr):
    return np.max(arr,axis=0)

if __name__=='__main__':
    print(use_extra_axis(A,ind,fmax))
    print(use_loop(A,ind,fmax))

Для определенных значений M и N (например, M = 10, N = 100) использование дополнительной оси может быть быстрее, чем использование цикла:

% python -mtimeit -s'import test,numpy' 'test.use_extra_axis(test.A2,test.ind2,test.fmax)'
10000 loops, best of 3: 162 usec per loop

% python -mtimeit -s'import test,numpy' 'test.use_loop(test.A2,test.ind2,test.fmax)'
1000 loops, best of 3: 222 usec per loop

Однако, когда N увеличивается (скажем, M = 10, N = 10000), использование цикла может быть быстрее:

% python -mtimeit -s'import test,numpy' 'test.use_extra_axis(test.A2,test.ind2,test.fmax)'
100 loops, best of 3: 13.9 msec per loop
% python -mtimeit -s'import test,numpy' 'test.use_loop(test.A2,test.ind2,test.fmax)'
100 loops, best of 3: 4.4 msec per loop

---

Включение прекрасной идеи использования разреженногоматрица:

def use_sparse_extra_axis(A,ind,f):
    B=scipy.sparse.coo_matrix((A, (range(len(A)), ind))).toarray()
    return f(B)

def use_sparse(A,ind,f):
    return [f(v) for v in scipy.sparse.coo_matrix((A, (ind, range(len(A))))).tolil().data]

Какая реализация лучше всего зависит от параметров N и M:

N=1000, M=100
·───────────────────────·────────────────────·
│ use_sparse_extra_axis │ 1.15 msec per loop │
│        use_extra_axis │ 2.79 msec per loop │
│              use_loop │ 3.47 msec per loop │
│            use_sparse │ 5.25 msec per loop │
·───────────────────────·────────────────────·

N=100000, M=10
·───────────────────────·────────────────────·
│ use_sparse_extra_axis │ 35.6 msec per loop │
│              use_loop │ 43.3 msec per loop │
│            use_sparse │ 91.5 msec per loop │
│        use_extra_axis │  150 msec per loop │
·───────────────────────·────────────────────·

N=100000, M=50
·───────────────────────·────────────────────·
│            use_sparse │ 94.1 msec per loop │
│              use_loop │  107 msec per loop │
│ use_sparse_extra_axis │  170 msec per loop │
│        use_extra_axis │  272 msec per loop │
·───────────────────────·────────────────────·

N=10000, M=50
·───────────────────────·────────────────────·
│              use_loop │ 10.9 msec per loop │
│            use_sparse │ 11.7 msec per loop │
│ use_sparse_extra_axis │ 15.1 msec per loop │
│        use_extra_axis │ 25.4 msec per loop │
·───────────────────────·────────────────────·

Answer 2

Я не думаю, что вы можете выйти из цикла, но, возможно, использование разреженных матриц Сципи будет более эффективным.

[f(v) for v in scipy.sparse.coo_matrix((A, (ind, range(len(A))))).tolil().data]

Answer 3

По крайней мере, для добавления, это работает

import numpy as np

def op_at(f, ind, vals):
    base = np.zeros(np.max(ind)+1)
    f.at(base, ind, vals)
    return base

print op_at(np.add, [ 0, 0, 2, 1, 2], [ 6, 4, 5, 7, 9])

> [ 10.   7.  14.]

К сожалению, это не работает для макс.

Answer 4

Другая возможность,

from operator import itemgetter
from itertools import groupby

A = [ 6, 4, 5, 7, 9 ]
ind = [ 0, 0, 2, 1, 2 ]

z = zip(ind,A)
z.sort()

fst,snd = itemgetter(0), itemgetter(1)
g = groupby(z,fst)

f = sum
# or 
# f = max

for i in g:
    print i[0],f(snd(j) for j in i[1])