Широковещательная передача с множественной матрицей

Question

У меня есть массив, состоящий из массивов N 3x3 (набор матриц, хотя тип данных - np.ndarray), и у меня есть массив, состоящий из массивов N 3x1 (набор векторов). То, что я хочу сделать, - умножить каждую матрицу на каждый вектор, поэтому я ожидаю получить N 3x1 массивы.

Простой пример:

A = np.ones((6,3,3)) 
B = np.ones((6,3,1)) 
np.dot(A,B) # This gives me a 6x3x6x1 array, which is not what I want 
np.array(map(np.dot,A,B)) # This gives me exactly what I want, but I don't want to have to rely on map 

Я устал все виды изменения формы, исследовали einsum и т.д., но не может заставить это работать так, как я хочу его. Как мне заставить это работать с многократным трансляцией? Эта операция, в конечном счете, должна быть выполнена много тысяч раз, и я не хочу map или перечислить операции распознавания, чтобы замедлить работу.

Answer 1

Вы можете использовать np.einsum для вычисления скалярных произведений и создать матрицу нужной формы:

np.einsum('ijk,ikl->ijl', A, B) 

Answer 2

A = np.random.rand(6, 3, 3) 
B = np.random.rand(6, 3, 1) 
C = np.array(map(np.dot, A, B)) 
D = np.sum(A*B.swapaxes(1, 2), axis=2)[..., None] 
assert np.allclose(C, D) 
assert C.shape == D.shape == (6, 3, 1) 

«allclose» потому, что есть какая-то разница с плавающей точкой закругления между двумя методами на порядок 1е-16.

.sapapaxis и [..., None] предназначены только для того, чтобы массивы соответствовали указанным вами формам. Вы также можете представить его проще:

A = np.random.rand(6, 3, 3) 
B = np.random.rand(6, 3) 
C = np.array(map(np.dot, A, B)) 
D = np.sum(A*B[:, None, :], axis=2) 
assert np.allclose(C, D) 
assert C.shape == D.shape == (6, 3)