这种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])