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

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

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