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

what are you doing here? Setting m[i,j] += m[i,j] and then setting m[i,j] = 0?? — user3684792
– user3684792, Commented Nov 11, 2014 at 8:23

score 2 · Accepted Answer · 2014-11-11 09:51:42Z

2

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.

edited Nov 11, 2014 at 9:51

answered Nov 11, 2014 at 8:29

user2379410

Sign up to request clarification or add additional context in comments.

2 Comments

Jonas K Over a year ago

Thanks. I have fixed a typo in the code. How can I leverage this for the operations on m?

user2379410 Over a year ago

@JonasKlemming: It's included in the edit to my answer. The multiplication by a Boolean array and subsequent summation seems a bit inefficient, but can't think of a better way right now.

Collectives™ on Stack Overflow

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

1 Answer 1

2 Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

2 Comments

Your Answer

Sign up or log in

Post as a guest

Related