Vectorizing loops in NumPy

Question

I am trying to vectorize a loop iteration using NumPy but am struggling to achieve the desired results. I have an array of pixel values, so 3 dimensions, say (512,512,3) and need to iterate each x,y and calculate another value using a specific index in the third dimension. An example of this code in a standard loop is as follows:

for i in xrange(width):
    for j in xrange(height):
        temp = math.sqrt((scalar1-array[j,i,1])**2+(scalar2-array[j,i,2])**2)

What I am currently doing is this:

temp = np.sqrt((scalar1-array[:,:,1])**2+(scalar2-array[:,:,2])**2)

The temp array I get from this is the desired dimensions (x,y) but some of the values differ from the loop implementation. How can I eliminate the loop to compute this example efficiently in NumPy?

Thanks in advance!

Edit:

Here is code that is giving me differing results for temp and temp2, obviously temp2 is just the calculation for one cell

temp = np.sqrt((cb_key-fg_cbcr_array[:,:,1])**2+(cr_key-fg_cbcr_array[:,:,2])**2)
temp2 = np.sqrt((cb_key-fg_cbcr_array[500,500,1])**2+(cr_key-fg_cbcr_array[500,500,2])**2)
print temp[500, 500]
print temp2

The output for the above is

12.039

94.069123521

The scalars are definitely initialized and the array is generated from an image using

fg = PIL.Image.open('fg.jpg')
fg_cbcr = fg.convert("YCbCr")
fg_cbcr_array = np.array(fg_cbcr)

Edit2:

Ok so I have tracked it down to a problem with my array. Not sure why yet but it works when the array is generated with np.random.random but not when loading from a file using PIL as above.

Answer 1

Your vectorized solution is correct.

• in your for loop temp is a scalar value that will take only the last value
• use np.sqrt istead of math.sqrt for vectorized inputs
• you should not use array as a variable since it can shadow the np.array method

I checked using the following code, which may give you some tip about where the error may be:

import numpy as np

width = 512
height = 512
scalar1 = 1
scalar2 = 2

a = np.random.random((height, width, 3))
tmp = np.zeros((height, width))
for i in xrange(width):
    for j in xrange(height):
        tmp[j,i] = np.sqrt((scalar1-a[j,i,1])**2+(scalar2-a[j,i,2])**2)

tmp2 = np.sqrt((scalar1-a[:,:,1])**2+(scalar2-a[:,:,2])**2)

np.allclose(tmp, tmp2)