Замените python понимание списка, генерирующее 3D-массив, функциями numpy - PullRequest
Новая
       20

Замените python понимание списка, генерирующее 3D-массив, функциями numpy

1 голос
/ 18 июня 2020

У меня есть две матрицы (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)]

но это быстро становится медленным.

Как я могу создать new_array с использованием numpy функций, которые должны быть быстрее?

Ответы [ 2 ]

3 голосов
/ 18 июня 2020

Два быстрых решения (без моих обычных доказательств и пояснений): 

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

res = mu.T[None,:,:] * nu[:,:,None]     # axes in same order as result
1 голос
/ 18 июня 2020 
#!/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]]])
...