我能否在这个计算中消除所有的Python循环:
result[i,j,k] = (x[i] * y[j] * z[k]).sum()
其中x[i]
、y[j]
和z[k]
是长度为N
的向量,x
、y
和z
的第一个维度分别具有长度A
、B
、C
,输出的形状为(A,B,C)
,每个元素是三重积的总和(逐个元素相乘)。
我可以将代码从3个循环减少到1个循环(下方代码),但我卡在尝试消除最后一个循环上。
如果必要,我可以通过轻微填充来使A=B=C
。
# Example with 3 loops, 2 loops, 1 loop (testing omitted)
N = 100 # more like 100k in real problem
A = 2 # more like 20 in real problem
B = 3 # more like 20 in real problem
C = 4 # more like 20 in real problem
import numpy
x = numpy.random.rand(A, N)
y = numpy.random.rand(B, N)
z = numpy.random.rand(C, N)
# outputs of each variant
result_slow = numpy.empty((A,B,C))
result_vec_C = numpy.empty((A,B,C))
result_vec_CB = numpy.empty((A,B,C))
# 3 nested loops
for i in range(A):
for j in range(B):
for k in range(C):
result_slow[i,j,k] = (x[i] * y[j] * z[k]).sum()
# vectorize loop over C (2 nested loops)
for i in range(A):
for j in range(B):
result_vec_C[i,j,:] = (x[i] * y[j] * z).sum(axis=1)
# vectorize one C and B (one loop)
for i in range(A):
result_vec_CB[i,:,:] = numpy.dot(x[i] * y, z.transpose())
numpy.testing.assert_almost_equal(result_slow, result_vec_C)
numpy.testing.assert_almost_equal(result_slow, result_vec_CB)