How to vectorize this python code?

问题:

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?

回答1:

The problem is caused by array broadcasting in the line:

A = numpy.minimum(B, B[:,k] + B[k,:])

B is size 6 by 6, B[:,k] is an array with 6 elements, B[k,:] is an array with 6 elements.

(Because you are using the numpy array type, both B[:,k] and B[k,:] return a rank-1 array of shape N)

Numpy automatically changes the sizes to match:

1. First B[:,k] is added to B[k,:] to make an intermediate array result with 6 elements. (This is not what you intended)
2. Second this 6 element array is broadcast to a 6 by 6 matrix by repeating the rows
3. Third the minimum of the original matrix and this broadcast matrix is computed.

This means that your numpy code is equivalent to:

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

The easiest way to fix your code is to use the matrix type instead of the array type:

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

The matrix type uses stricter broadcasting rules so in this case:

1. A[:,k] is extended to a 6 by 6 matrix by repeating columns
2. A[k,:] is extended to a 6 by 6 matrix by repeating rows
3. The broadcasted matrices are added together to make a 6 by 6 matrix
4. The minimum is applied

回答2:

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.