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

【问题标题】：numpy array multiplication with arrays of arbitrary dimensionsnumpy 数组乘法与任意维度的数组
【发布时间】：2015-08-24 18:26:48
【问题描述】：

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

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

我不想用循环来解决这个问题，一个原因是事物的美感，另一个原因是这段代码需要非常通用。只要前导维度为 10，我希望用户能够输入几乎任何形状的 B。例如，我希望用户也能够输入形状为 (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]])
>>>

【问题讨论】：

请包含 small 示例数组和预期输出。
如果B 是(3,5,10)，则A*B 有效。 numpy 会根据需要在开头自动添加维度（MATLAB 在末尾添加维度）。
如上所述，广播对离开数组最右边的维度。这意味着 (A*B.T).T 将匹配长度为 10 的乘法运算，这将是一个通用解决方案，但我认为 DSM 建议的 einsum 方法可以说更好。

标签： python arrays numpy multiplication

【解决方案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

但这仅适用于第 2 种情况。借用@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

【讨论】：

哇，一些复杂的东西，恕我直言！实际上，您的第一个与编辑配合得很好： >>> np.array([A[i]*B[i] for i in range(len(A))]) 看到有什么问题吗？谢谢！
@bob.sacamento：不，这应该可行——我只是在复制你的暴力示例。但由于这会从缓慢的 listcomp 创建数组，因此对于大型形状会很慢，并且应该仅用于测试目的。
这很无聊且不雅，但a[(slice(None),) + (None,) * (b.ndim - 1)] * b 也能搞定......
@DSM 终于开始尝试您的 einsum 方法。我直接复制并粘贴并得到错误“ValueError：操作数 0 没有足够的维度来匹配广播，并且无法扩展，因为在开始和结束时都指定了爱因斯坦和下标”有什么想法吗？另外，了解有关 einsum 的任何好的教程。我发现文档页面非常不透明。谢谢。
@DSM 如果您有兴趣，我还从 einsum 文档中复制并粘贴了一个示例 -- "np.einsum('...j,j', a, b)" --并得到同样的错误。知道这可能是什么吗？我可能有一个旧版本的 scipy/numpy 不能处理这种事情吗？

【解决方案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

【讨论】：