如何对以下双循环进行矢量化?
我有一个N乘A矩阵和一个N乘B矩阵,其中A和B可能不同,N比A和B小得多。我想按如下方式生成A by B矩阵,但理想情况下没有循环:
import numpy as np
def foo(arr):
# can be anything - just an example so that the code runs
return np.sum(arr)
num_a = 12
num_b = 8
num_dimensions = 3
a = np.random.rand(num_dimensions, num_a)
b = np.random.rand(num_dimensions, num_b)
# this is the loop I want to eliminate:
output = np.zeros( (num_a, num_b) )
for i in xrange(num_a):
for j in xrange(num_b):
output[i,j] = foo(a[:,i] - b[:,j])
有什么想法吗?
答案 0 :(得分:7)
第一个矢量化foo()
,即修改foo()
,使其能够在形状(N, A, B)
的数组上正确操作,返回形状(A, B)
的数组。这一步通常是困难的。如何完成这完全取决于foo()
的作用。对于给定的示例,它很容易做到:
def foo(arr):
return np.sum(arr, axis=0)
现在,使用broadcasting rules创建一个包含所有向量差异的(N, A, B)
数组,并将其传递给foo()
:
foo(a[:, :, np.newaxis] - b[:, np.newaxis])