python - 如何使用 numpy 优化数组上的双循环？

Question

这段代码在 Python 中相当慢。我如何使用 numpy 来优化它？数组 m 和 z 已经是 numpy 数组，但我认为我需要以某种方式对其进行矢量化，尽管我找不到矢量化函数也可以访问值坐标的示例。

for i in xrange(states):
    for j in xrange(states):
        if i != j and z[i, j] != 0 and m[i, j] < 0.0:
            m[i, i] += m[i, j]
            m[i, j] = 0.0
            z[i, j] = 0

Answer 1

您可以使用索引数组将代码转换为矢量化 Numpy:

import numpy as np

i, j = np.indices([states, states])
bool_index = (i != j) & (z != 0) & (m < 0.0)
z[bool_index] = 0

但是由于您已经知道 i != j 的计算结果，因此仅使用对角数组会更快。这也使得在 m 上执行稍微尴尬的操作变得更加容易:

I = np.eye(states, dtype=bool)
bool_index = ~I & (z != 0) & (m < 0.0)

m[I] += (bool_index * m).sum(axis=1)
m[bool_index] = 0.0
z[bool_index] = 0

请注意，& 运算符执行的是 np.bitwise_and 而不是 np.logic_and，但在这种情况下它们是等效的，因为所有操作数是 bool 数组。