How can I optimize this double-loop over an array using numpy?

Question

This code is quite slow in Python. How can I optimize this using numpy. The array m and z are already numpy arrays but I assume I need to vectorize this somehow though I can't find an example where the vectorized function also has access to the value coordinates.

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

You can translate your code to vectorized Numpy using arrays of indices:

import numpy as np

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

But since you already know what i != j will evaluate to, it's faster to just use a diagonal array. This also makes it easier to do the slightly awkward operation on 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

Note that the & operator does np.bitwise_and and not np.logical_and but they're equivalent in this case, because all operands are Boolean arrays.