numpy 中 Kronecker 积的双循环矢量化和整形

Double loop vectorization and reshape for Kronecker product in numpy

我有两个矩阵

import numpy as np
n = 10
a = 2*np.ones((n,n,3))
b = 3*np.ones((n,n,3))

我想用提醒Kronecker乘积的方式将它们相乘,然后求和

s = 0
for i in range(n):
    for j in range(n):
        s +=  a*b[i,j]

有没有在numpy中向量化它的方法?

您的代码可以重写为:

因此,这应该有效:

s = a * np.sum(np.sum(b,axis=1),axis=0)

也许可以用np.einsum()写得更优雅:

import numpy as np


n = 10
a = 2 * np.ones((n, n, 3))
b = 3 * np.ones((n, n, 3))

s = 0
for i in range(n):
    for j in range(n):
        s +=  a * b[i, j]

print(s.shape)
# (10, 10, 3)


ss = a * np.einsum('ijk->k', b)
print(ss.shape)
# (10, 10, 3)
print(np.all(s == ss))
# True

甚至 np.sum():

sss = a * np.sum(b, axis=(0, 1))
print(sss.shape)
# (10, 10, 3)
print(np.all(s == sss))
# True

np.einsum() 似乎更快:

n = 100
a = 2 * np.ones((n, n, 3))
b = 3 * np.ones((n, n, 3))

%timeit f_with_loops(a, b)
# 1 loop, best of 3: 787 ms per loop
%timeit a * np.einsum('ijk->k', b)
# 10000 loops, best of 3: 121 µs per loop
%timeit a * np.sum(b, axis=(0, 1))
# 1000 loops, best of 3: 254 µs per loop