这个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

第一个矢量foo(),即foo()以一种可以在形状数组上正确操作的方式进行修改(N, A, B),返回一个形状数组(A, B).这一步通常是困难的.如何做到完全取决于做什么foo().对于给定的示例,它很容易做到:

def foo(arr):
    return np.sum(arr, axis=0)

现在,使用广播规则创建一个(N, A, B)包含所有向量差异的数组,并将其传递给foo():

foo(a[:, :, np.newaxis] - b[:, np.newaxis])