用 numpy 函数替换 python 列表理解生成 3D 数组答案

【问题标题】：Replace python list comprehension generating 3D array with numpy functions用 numpy 函数替换 python 列表理解生成 3D 数组
【发布时间】：2020-10-08 15:37:44
【问题描述】：

我有两个矩阵（numpy 数组），mu 和 nu。从这些我想创建第三个数组如下：

new_array_{j, k, l} = mu_{l, k} nu_{j, k}

我可以天真地使用列表推导：

[[[mu[l, k] * nu[j, k] for k in np.arange(N)] for l in np.arange(N)] for j in np.arange(N)]

但很快就会变慢。

如何使用更快的 numpy 函数创建 new_array？

【问题讨论】：

标签： python arrays list numpy list-comprehension

【解决方案1】：

#!/usr/bin/env python                                                                               

import numpy as np

# example data
mu = np.arange(10).reshape(2,5)
nu = np.arange(15).reshape(3,5) + 20

# get array sizes
nl, nk = mu.shape
nj, nk_ = nu.shape
assert(nk == nk_)

# get arrays with dimensions (nj, nk, nl)

# in the case of mu3d, we need to add a slowest varying dimension
# so (after transposing) this can be done by cycling through the data
# nj times along the slowest existing axis and then reshaping
mu3d = np.concatenate((mu.transpose(),) * nj).reshape(nj, nk, nl)

# in the case of nu3d, we need to add a new fastest varying dimension 
# so this can be done by repeating each element nl times, and again it
# needs reshaping
nu3d = nu.repeat(nl).reshape(nj, nk, nl)

# now just multiple element by element
new_array = mu3d * nu3d

print(new_array)

给予：

>>> mu
array([[0, 1, 2, 3, 4],
       [5, 6, 7, 8, 9]])

>>> nu
array([[20, 21, 22, 23, 24],
       [25, 26, 27, 28, 29],
       [30, 31, 32, 33, 34]])

>>> nj, nk, nl
(3, 5, 2)

>>> mu3d
array([[[0, 5],
        [1, 6],
        [2, 7],
        [3, 8],
        [4, 9]],

       [[0, 5],
        [1, 6],
        [2, 7],
        [3, 8],
        [4, 9]],

       [[0, 5],
        [1, 6],
        [2, 7],
        [3, 8],
        [4, 9]]])

>>> nu3d
array([[[20, 20],
        [21, 21],
        [22, 22],
        [23, 23],
        [24, 24]],

       [[25, 25],
        [26, 26],
        [27, 27],
        [28, 28],
        [29, 29]],

       [[30, 30],
        [31, 31],
        [32, 32],
        [33, 33],
        [34, 34]]])

>>> new_array
array([[[  0, 100],
        [ 21, 126],
        [ 44, 154],
        [ 69, 184],
        [ 96, 216]],

       [[  0, 125],
        [ 26, 156],
        [ 54, 189],
        [ 84, 224],
        [116, 261]],

       [[  0, 150],
        [ 31, 186],
        [ 64, 224],
        [ 99, 264],
        [136, 306]]])

【讨论】：

您可以在创建 3d 数组时利用广播规则。不需要重复。
@hpaulj 我不愿意使用广播，因为担心它会使用一些启发式方法来决定沿哪些轴进行广播，如果存在相同大小的维度，这可能会给出错误的答案.我不知道这些担心是否没有根据。
你仍然需要转置。广播规则是明确的。
@Mashy 不确定我有很多要添加到 cmets 的内容。我们知道 mu 有轴 (L,K) 而 nu 有轴 (J,K)，我们想要一个带有轴 (J,K,L) 的结果，所以 nu 会想要转置，然后（如果不是使用隐式广播）我们将不得不添加这些额外的轴。幸运的是，因为在一种情况下添加了最慢的变化轴，这仅意味着将整个数组重复正确的次数，而在另一种情况下，添加了最快的变化轴，这意味着将每个元素重复正确的次数。在中间添加一个轴可能会更棘手......
@Mashy 试错是可以的，但在您进行实验时确保所有维度都不同确实很有帮助。

【解决方案2】：

两个快速解决方案（没有我通常的证明和解释）：

res = np.einsum('lk,jk->jkl', mu, nu) 

res = mu.T[None,:,:] * nu[:,:,None]     # axes in same order as result

【讨论】：

我喜欢它，看看你现在说的广播是什么意思。我没有意识到您在相乘之前明确地向两个数组添加了 size-1 维度，只是以为您在谈论将 3d 数组乘以 2d 数组并依靠它来确定要广播的维度。现在一切都说得通了。
@hpaulj 这些都是非常有用的方法。直到现在我才知道 numpy.einsum，但我想我会一直使用它。我最初的物理问题还有一些我用 numpy.matmul 明确写过的术语，但对于其他一些，einsum 似乎是一个更简洁的解决方案。干杯。
@Mashy einsum 看起来很有用。（我也没有遇到过。）虽然从上面不清楚如何表示除乘法之外的运算。