numpy数组与任意维度数组的乘法

Question

我有一个 numpy 数组 A，其形状为 (10,)。

我也有，截至目前，一个形状为 (10,3,5) 的 numpy 数组 B。我想在这两者之间进行乘法以获得 C，使得 C[0,:,:]=A[0]*B[0,:,:], C[1]=A[1]*B[1 ,:,:],等等

我不想用循环来解决这个问题，一个原因是它的美学，另一个原因是这段代码需要非常通用。我希望用户能够输入几乎任何形状的任何 B，只要前导尺寸为 10。例如，我希望用户也能够输入形状为 (10,4) 的 B。

所以：如何使用 numpy 实现这个乘法？谢谢

附录：例如，已被问到。会变小。假设 A 是 numpy 数组 [1,2,3]，B 是 numpy 数组 [[1,2],[4,5],[7,8]]。我希望两者相乘得到 [[1,2],[8,10],[21,24]]。 ...

>>> a
array([1, 2, 3])
>>> b
array([[1, 2],
       [4, 5],
       [7, 8]])
>>> #result
>>> c
array([[ 1,  2],
       [ 8, 10],
       [21, 24]])
>>>

Answer 1

您可以使用 None（或 np.newaxis）扩展 A 以匹配 B:

>>> A = np.arange(10)
>>> B = np.random.random((10,3,5))
>>> C0 = np.array([A[i]*B[i,:,:] for i in range(len(A))])
>>> C1 = A[:,None,None] * B
>>> np.allclose(C0, C1)
True

但这只适用于第二种情况。从@ajcr 借用，通过足够的转置，我们可以获得适用于一般情况的隐式广播：

>>> C3 = (A * B.T).T
>>> np.allclose(C0, C3)
True

或者，您可以使用 einsum 来提供通用性。回想起来，与转置路线相比，它在这里可能有点矫枉过正，但当乘法更复杂时它很方便。

>>> C2 = np.einsum('i,i...->i...', A, B)
>>> np.allclose(C0, C2)
True

和

>>> B = np.random.random((10,4))
>>> D0 = np.array([A[i]*B[i,:] for i in range(len(A))])
>>> D2 = np.einsum('i,i...->i...', A, B)
>>> np.allclose(D0, D2)
True

Answer 2

虽然我喜欢 einsum 符号，但我会添加一些变化....

您可以向 a 添加足够多的额外维度，这样它将 broadcast 横跨 b。

>>> a.shape
(3,)
>>> b.shape
(3,2)

b 的维度比 a

多

extra_dims = b.ndim - a.ndim

将额外维度添加到 a

new_shape = a.shape + (1,)*extra_dims    # (3,1)
new_a = a.reshape(new_shape)

乘

new_a * b

作为函数：

def f(a, b):
    '''Product across the first dimension of b.

    Assumes a is 1-dimensional.
    Raises AssertionError if a.ndim > b.ndim or
     - the first dimensions are different
    '''
    assert a.shape[0] == b.shape[0], 'First dimension is different'
    assert b.ndim >= a.ndim, 'a has more dimensions than b'

    # add extra dimensions so that a will broadcast
    extra_dims = b.ndim - a.ndim
    newshape = a.shape + (1,)*extra_dims
    new_a = a.reshape(newshape)

    return new_a * b

numpy数组与任意维度数组的乘法

numpy array multiplication with arrays of arbitrary dimensions

python

arrays

numpy

multiplication