2

The sample code is as below.

I want to get dataNew(h, w, length) according to data(h, w, c) and ind(h, w). Here length < c, it means dataNew is sliced from data.

Here, length and ind[i, j] is made sure to suit the c value.

I have realize it through for loops, but I wnat the python way. Please help, thanks!

import numpy as np

h, w, c = 3, 4, 5

data = np.arange(60).reshape((h, w, c))
print(data)

length = 3
ind = np.random.randint(0, 3, 12).reshape(h, w)
print(ind)

dataNew = np.empty((h, w, length), np.int16)
for i in range(h):
    for j in range(w):
        st = ind[i, j]
        dataNew[i, j] = data[i, j][st : st + length]

print(dataNew)
ToughMind
  • 987
  • 1
  • 10
  • 28

2 Answers2

1

We can leverage np.lib.stride_tricks.as_strided based scikit-image's view_as_windows to get sliding windows. More info on use of as_strided based view_as_windows.

from skimage.util.shape import view_as_windows

# Get all sliding windows along the last axis
w = view_as_windows(data,(1,1,length))[...,0,0,:]

# Index into windows with start indices and slice out singleton dims 
out = np.take_along_axis(w,ind[...,None,None],axis=-1)[...,0]

Last step is basically using advanced-indexing into the windows with those start indices. This could be made a bit simpler and might be easier to understand. So, alternatively, we could do -

m,n = ind.shape
I,J = np.ogrid[:m,:n]
out = w[I,J,ind]
Divakar
  • 218,885
  • 19
  • 262
  • 358
1

One way would be creating an indexing array using broadcasting and use np.take_along_axis to index the array:

ix = ind[...,None] + np.arange(length)
np.take_along_axis(data, ix, -1)
yatu
  • 86,083
  • 12
  • 84
  • 139