python - 向量化连续 numpy 计算

Question

我在 numpy 中有许多操作，我可以在循环中完美地执行这些操作，但我无法在一次 numpy 调用中对它们进行矢量化。

# data matrix
d = np.random.rand(1496, 1, 2)

# boolean matrix
r = np.random.rand(5, 1496, 1, 2) > 0.5

# result matrix
x = np.empty((5,))

# How can I avoid this loop?
for i in xrange(r.shape[0]):
  x[i] = d[r[i]].sum()

是否可以通过某种方式矢量化循环来加快速度？

Answer 1

@Psidom 方法的变体:

np.tensordot(d, r, axes=((0,1,2), (1,2,3)))

这依赖于 tensordot 函数。来自其documentation :

Given two tensors (arrays of dimension greater than or equal to one), a and b, and an array_like object containing two array_like objects, (a_axes, b_axes), sum the products of a's and b's elements (components) over the axes specified by a_axes and b_axes.

从根本上来说，tensordot 计算两个数组的乘积(就像@Psidom 所做的那样)并计算乘积数组的所有元素的总和。与 @Psidom 的方法相比，此方法的唯一“优点”是，它可以更灵活地指定在两个数组中执行乘积和求和的轴。。与 @Psidom 的方法相比，它并没有提供更高的性能。

另请参阅Understanding tensordot .