Question

我有两个矩阵（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 ？

Answer 1

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

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

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

Answer 2

#!/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]]])

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

2 个答案: