How to vectorize this python code?

Question

I am trying to use NumPy and vectorization operations to make a section of code run faster. I appear to have a misunderstanding of how to vectorize this code, however (probably due to an incomplete understanding of vectorization).

Here's the working code with loops (A and B are 2D arrays of a set size, already initialized):

for k in range(num_v):
    B[:] = A[:]
    for i in range(num_v):
        for j in range(num_v):
            A[i][j] = min(B[i][j], B[i][k] + B[k][j])
return A

And here is my attempt at vectorizing the above code:

for k in range(num_v):
    B = numpy.copy(A)
    A = numpy.minimum(B, B[:,k] + B[k,:])
return A

For testing these, I used the following, with the code above wrapped in a function called 'algorithm':

def setup_array(edges, num_v):
    r = range(1, num_v + 1)
    A = [[None for x in r] for y in r]  # or (numpy.ones((num_v, num_v)) * 1e10) for numpy
    for i in r:
        for j in r:
            val = 1e10
            if i == j:
                val = 0 
            elif (i,j) in edges:
                val = edges[(i,j)]
            A[i-1][j-1] = val 
    return A

A = setup_array({(1, 2): 2, (6, 4): 1, (3, 2): -3, (1, 3): 5, (3, 6): 5, (4, 5): 2, (3, 1): 4, (4, 3): 8, (3, 4): 6, (2, 4): -4, (6, 5): -5}, 6) 
B = []
algorithm(A, B, 6)

The expected outcome, and what I get with the first code is:

[[0, 2, 5, -2, 0, 10] 
 [8, 0, 4, -4, -2, 9]
 [4, -3, 0, -7, -5, 5]
 [12, 5, 8, 0, 2, 13]
 [10000000000.0, 9999999997.0, 10000000000.0, 9999999993.0, 0, 10000000000.0]
 [13, 6, 9, 1, -5, 0]]

The second (vectorized) function instead returns:

[[ 0. -4.  0.  0.  0.  0.]
 [ 0. -4.  0. -4.  0.  0.]
 [ 0. -4.  0.  0.  0.  0.]
 [ 0. -4.  0.  0.  0.  0.]
 [ 0. -4.  0.  0.  0.  0.]
 [ 0. -4.  0.  0. -5.  0.]]

What am I missing?

Answer 1

Usually you want to vectorize code because you think it is running too slow.
If your code is too slow, then I can tell you that proper indexing will make it faster.
Instead of A[i][j] you should write A[i, j] -- this avoids a transient copy of a (sub)array.
Since you do this in the inner-most loop of your code this might be very costly.

Look here:

In [37]: timeit test[2][2]
1000000 loops, best of 3: 1.5 us per loop

In [38]: timeit test[2,2]
1000000 loops, best of 3: 639 ns per loop

Do this consistently in your code -- I strongly believe that solves already your performance problem!

Having said that...

... here's my take on how to vectorize

for k in range(num_v):
    numpy.minimum(A, np.add.outer(A[:,k], A[k,:]), A)
return A

numpy.minimum will compare two arrays and return element-wise the smaller of two elements. If you pass a third argument it will take the output. If this is an input array the whole operation is in place.

As Peter de Rivay explains, there is a problem in your solution with broadcasting -- but mathematically what you want to do is some kind of outer product over addition of two vectors. Therefore you can use the outer operation on the add function.

NumPy’s binary ufuncs have special methods for performing certain kinds of special vectorized operations like reduce, accumulate, sum and outer.