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

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