如何使用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.logical_and，但在这种情况下它们是等价的，因为所有操作数都是布尔数组。