CODEWITHSUNDEEP
numpyvectorizationnumpy-ndarray
Abhijeet Krishnan
Abhijeet Krishnan

Reputation: 13

Vectorizing numpy code

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.

Upvotes: 1

Views: 121

Answers (1)

w-m
w-m

Reputation: 11232

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.

Upvotes: 1

Related Questions