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: 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.

