此Numpy双循环的矢量化

Question

如何矢量化以下双循环？

我有一个N乘A矩阵和一个N乘B矩阵，其中A和B可能不同，N比A和B小得多。我想要如下生成A乘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]) 

任何想法？

Answer 1

首先vectorise 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])