Vectorizing numpy code

Question

I'm trying to vectorize the following code -

for i in range(imL.shape[0]):
    for j in range(imL.shape[1]):
        if j - disp[i][j] >= 0:
            imR[i, j - disp[i][j], :] = imL[i, j, :]

Essentially, I want to do this -

A[i, j - B[i, j], k] = C[i, j, k] for all (i, j, k)

I looked into boolean indexing and came up with this -

tmp = np.arange(imR.shape[1])
tmp = np.repeat(tmp[np.newaxis, :], imR.shape[0], axis=0)
tmp = np.repeat(tmp[:, :, np.newaxis], imR.shape[2], axis=2)

imR[(tmp - disp) >= 0] = imL

However, this throws up the error -

TypeError: NumPy boolean array indexing assignment requires a 0 or 1-dimensional input, input has 3 dimensions

I believe integer indexing is the solution, however, I am unable to come up with anything. Is there a way to vectorize this efficiently?

NOTE: imL and imR are N x M x 3 ndarrays. disp is an N x M ndarray.

w-m · Accepted Answer · 2018-06-15 16:05:18Z

1

You can simplify the creation of a temporary indices array with np.indices:

i_indices, j_indices = np.indices(disp.shape)

Then you can create an updated index array with your custom formula:

# j - B[i, k]
new_j_ind = j_indices - disp

Replace the updated j values with the original js on the condition:

# if j - disp[i][j] < 0, keep original j index
new_j_ind[new_j_ind < 0] = j_indices[new_j_ind < 0]

And write the array.

# A[i, j - B[i, j] if j - B[i, j] > 0 else j, k] = C[i, j, k] for all (i, j, k)
imR[i_indices, new_j_ind] = imL

Careful though: this is not well defined if [i, j - B[i, j]] ever map to the same coordinates for any (i, k). In the for loop it's well defined (the last written value wins), in the numpy vectorization code it is not.

edited Jun 15, 2018 at 16:05

answered Jun 14, 2018 at 15:58

w-m

11.3k1 gold badge46 silver badges51 bronze badges

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

4 Comments

Abhijeet Krishnan Over a year ago

Thanks for the tip regarding using np.indices for the temporary indices array. However, the step imR[select] = imL[select] is not semantically what I want to achieve. I believe it translates to A[i, j - B[i, j], k] = C[i, j - B[i, j], k] for all (i, j, k). It definitely does not give me the same result as compared to a for-loop.

w-m Over a year ago

Indeed, I misread the question. The solution originally presented was A[i, j, k] = C[i, j, k] for all (i, j, k) where j - B[i, j] > 0. I updated the answer.

Abhijeet Krishnan Over a year ago

Thank you, this worked perfectly. I did not understand how the indexing worked in imR[i_indices, new_j_ind] = imL. Could you point me to some resource where I might understand it further?

w-m Over a year ago

The documentation has a couple of good examples, you can start there: docs.scipy.org/doc/numpy-1.13.0/reference/…. I often find it helpful to take a simple example (like the one from the doc) and play around with it in an iPython session.

Collectives™ on Stack Overflow

Vectorizing numpy code

1 Answer 1

4 Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

4 Comments

Your Answer

Sign up or log in

Post as a guest

Related