Numpy einsum внешняя сумма 2d-массивов - PullRequest
Новая
       19

Numpy einsum внешняя сумма 2d-массивов

0 голосов
/ 13 сентября 2018

Я пытался найти ответ, но не смог найти то, что мне было нужно.Извиняюсь, если это повторяющийся вопрос.

Предположим, у меня есть 2d-массив с формой (n, n*m).То, что я хочу сделать, это внешняя сумма этого массива для его транспонирования, что приводит к массиву с формой (n*m, n*m).Например, предположим, что у меня есть 

A = array([[1., 1., 2., 2.],
           [1., 1., 2., 2.]])

Я хочу сделать внешнюю сумму A и A.T так, чтобы получилось: 

>>> array([[2., 2., 3., 3.],
           [2., 2., 3., 3.],
           [3., 3., 4., 4.],
           [3., 3., 4., 4.]])

Обратите внимание, что np.add.outerне работает, потому что он распределяет входы в векторы.Я могу добиться чего-то подобного, выполнив 

np.tile(A, (2, 1)) + np.tile(A.T, (1, 2))

, но это не представляется разумным, когда n и m достаточно велики (n > 100 и m > 1000).Можно ли написать эту сумму, используя einsum?Я просто не могу понять, einsum.

Ответы [ 2 ]

0 голосов
/ 14 сентября 2018

Чтобы использовать broadcasting, нам нужно разбить его на 3D, а затем переставить оси и добавить - 

n = A.shape[0]
m = A.shape[1]//n
a = A.reshape(n,m,n) # reshape to 3D
out = (a[None,:,:,:] + a.transpose(1,2,0)[:,:,None,:]).reshape(n*m,-1)

Пример прогона для проверки - 

In [359]: # Setup input array
     ...: np.random.seed(0)
     ...: n,m = 3,4
     ...: A = np.random.randint(1,10,(n,n*m))

In [360]: # Original soln
     ...: out0 = np.tile(A, (m, 1)) + np.tile(A.T, (1, m))

In [361]: # Posted soln
     ...: n = A.shape[0]
     ...: m = A.shape[1]//n
     ...: a = A.reshape(n,m,n)
     ...: out = (a[None,:,:,:] + a.transpose(1,2,0)[:,:,None,:]).reshape(n*m,-1)

In [362]: np.allclose(out0, out)
Out[362]: True

Сроки с большим n, m - 

In [363]: # Setup input array
     ...: np.random.seed(0)
     ...: n,m = 100,100
     ...: A = np.random.randint(1,10,(n,n*m))

In [364]: %timeit np.tile(A, (m, 1)) + np.tile(A.T, (1, m))
1 loop, best of 3: 407 ms per loop

In [365]: %%timeit
     ...: # Posted soln
     ...: n = A.shape[0]
     ...: m = A.shape[1]//n
     ...: a = A.reshape(n,m,n)
     ...: out = (a[None,:,:,:] + a.transpose(1,2,0)[:,:,None,:]).reshape(n*m,-1)
1 loop, best of 3: 219 ms per loop

Дальнейшее повышение производительности с numexpr

Мы можем использовать multi-core с numexpr модулем для больших данных и для повышения эффективности памяти и, следовательно, производительности - 

import numexpr as ne

n = A.shape[0]
m = A.shape[1]//n
a = A.reshape(n,m,n)
p1 = a[None,:,:,:]
p2 = a.transpose(1,2,0)[:,:,None,:]
out = ne.evaluate('p1+p2').reshape(n*m,-1)

Синхронизация с такими же большими настройками n, m - 

In [367]: %%timeit
     ...: # Posted soln
     ...: n = A.shape[0]
     ...: m = A.shape[1]//n
     ...: a = A.reshape(n,m,n)
     ...: p1 = a[None,:,:,:]
     ...: p2 = a.transpose(1,2,0)[:,:,None,:]
     ...: out = ne.evaluate('p1+p2').reshape(n*m,-1)
10 loops, best of 3: 152 ms per loop
0 голосов
/ 13 сентября 2018

в одну сторону - 

(A.reshape(-1,*A.shape).T+A)[:,0,:]

Я думаю, что это займет много памяти с n>100 и m>1000.

но разве это не то же самое, что 

np.add.outer(A,A)[:,0,:].reshape(4,-1)
...