CODEWITHSUNDEEP
pythonnumpyindexingslice
zenna
zenna

Reputation: 9176

Numpy advanced indexing

I have arrays with the following shapes

voxel_grids : (128, 32, 32, 32)
indices : (128, 3, 1024)

I want to construct an array scalar (128, 1024) such that

scalar[i,j] = voxel_grids[i, indices[i, 0, j], indices[i, 1, j], indices[i,2,j]]

Is there a straightforward way to do this using numpy (advanced) indexing?

Upvotes: 1

Views: 439

Answers (1)

Divakar
Divakar

Reputation: 221514

You could do something like this -

m = voxel_grids.shape[0]
out = voxel_grids[np.arange(m)[:,None],indices[:,0],indices[:,1],indices[:,2]]

Other way would be to extract the three slices from indices into three variables and use them for indexing. This might not be any more efficient than the previous one, but might be a bit easier to eyes. It's shown below -

m = voxel_grids.shape[0]
x,y,z = indices.swapaxes(0,1)
out = voxel_grids[np.arange(m)[:,None],x,y,z]

Upvotes: 1

Related Questions