поэлементно умножение массивов различной формы в питона

Question

Скажем, у меня есть два массива a и b,

a.shape = (5,2,3) 
    b.shape = (2,3) 

тогда c = a * b даст мне массив c фигуры (5,2,3) с c[i,j,k] = a[i,j,k]*b[j,k].

Сейчас ситуация,

a.shape = (5,2,3) 
    b.shape = (2,3,8) 

и я хочу c иметь форму (5,2,3,8) с c[i,j,k,l] = a[i,j,k]*b[j,k,l].

Как это сделать эффективно? Мои a и b на самом деле довольно большие.

Answer 1

Это должно работать:

a[..., numpy.newaxis] * b[numpy.newaxis, ...] 

Использование:

In : a = numpy.random.randn(5,2,3) 

In : b = numpy.random.randn(2,3,8) 

In : c = a[..., numpy.newaxis]*b[numpy.newaxis, ...] 

In : c.shape 
Out: (5, 2, 3, 8) 

Ref: Array Broadcasting in numpy

Edit: Обновленный ссылка URL

Answer 2

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

import numpy as np 
a = np.random.normal(size=(5,2,3)) 
b = np.random.normal(size=(2,3,8)) 
c = np.einsum('ijk,jkl->ijkl',a,b) 

и:

In [5]: c.shape 
Out[5]: (5, 2, 3, 8) 

In [6]: a[0,0,1]*b[0,1,2] 
Out[6]: -0.041308376453821738 

In [7]: c[0,0,1,2] 
Out[7]: -0.041308376453821738 

np.einsum может быть немного сложнее в использовании, но достаточно мощный для этого рода проблемы индексации:

http://docs.scipy.org/doc/numpy/reference/generated/numpy.einsum.html

Также обратите внимание, что t его требует numpy> = v1.6.0

Я не уверен в эффективности для вашей конкретной проблемы, но если он не работает так же хорошо, как нужно, обязательно изучите использование Cython с явным для циклов и, возможно, его распараллелите используя prange

UPDATE

In [18]: %timeit np.einsum('ijk,jkl->ijkl',a,b) 
100000 loops, best of 3: 4.78 us per loop 

In [19]: %timeit a[..., np.newaxis]*b[np.newaxis, ...] 
100000 loops, best of 3: 12.2 us per loop 


In [20]: a = np.random.normal(size=(50,20,30)) 
In [21]: b = np.random.normal(size=(20,30,80)) 

In [22]: %timeit np.einsum('ijk,jkl->ijkl',a,b) 
100 loops, best of 3: 16.6 ms per loop 

In [23]: %timeit a[..., np.newaxis]*b[np.newaxis, ...] 
100 loops, best of 3: 16.6 ms per loop 

In [2]: a = np.random.normal(size=(500,20,30)) 
In [3]: b = np.random.normal(size=(20,30,800)) 

In [4]: %timeit np.einsum('ijk,jkl->ijkl',a,b) 
1 loops, best of 3: 3.31 s per loop 

In [5]: %timeit a[..., np.newaxis]*b[np.newaxis, ...] 
1 loops, best of 3: 2.6 s per loop