Can I select arbitrary windows from the last dimension of a numpy array?-CodePudding

I'd like to write a numpy function that takes an MxN array A, a window length L, and an MxP array idxs of starting indices into the M rows of A that selects P arbitrary slices of length L from each of the M rows of A. Except, I would love for this to work on the last dimension of A, and not necessarily care how many dimensions A has, so all dims of A and idxs match except the last one. Examples:

If A is just 1D:

A = np.array([1, 2, 3, 4, 5, 6])

window_len = 3

idxs = np.array([1, 3])

result = magical_routine(A, idxs, window_len)

Where result is a 2x3 array since I selected 2 slices of len 3:

np.array([[ 2, 3, 4],
          [ 4, 5, 6]])

If A is 2D:

A = np.array([[ 1, 2, 3, 4, 5, 6],
              [ 7, 8, 9,10,11,12],
              [13,14,15,16,17,18]])

window_len = 3
idxs = np.array([[1, 3],
                 [0, 1],
                 [2, 2]])

result = magical_routine(A, idxs, window_len)

Where result is a 3x2x3 array since there are 3 rows of A, and I selected 2 slices of len 3 from each row:

np.array([[[ 2, 3, 4], [ 4, 5, 6]],
          [[ 7, 8, 9], [ 8, 9,10]],
          [[15,16,17], [15,16,17]]])

And so on.

I have discovered an number of inefficient ways to do this, along with ways that work for a specific number of dimensions of A. For 2D, the following is pretty tidy:

col_idxs = np.add.outer(idxs, np.arange(window_len))
np.take_along_axis(A[:, np.newaxis], col_idxs, axis=-1)

I can't see a nice way to generalize this for 1D and other D's though...

Is anyone aware of an efficient way that generalizes to any number of dims?

CodePudding user response：

For your 1d case

In [271]: A=np.arange(1,7)
In [272]: idxs = np.array([1,3])

Using the kind of iteration that this questions usually gets:

In [273]: np.vstack([A[i:i 3] for i in idxs])
Out[273]: 
array([[2, 3, 4],
       [4, 5, 6]])

Alternatively generate all indices, and one indexing. linspace is handy for this (though it's not the only option):

In [278]: j = np.linspace(idxs,idxs 3,3,endpoint=False)
In [279]: j
Out[279]: 
array([[1., 3.],
       [2., 4.],
       [3., 5.]])
In [282]: A[j.T.astype(int)]
Out[282]: 
array([[2, 3, 4],
       [4, 5, 6]])

for the 2d

In [284]: B
Out[284]: 
array([[ 1,  2,  3,  4,  5,  6],
       [ 7,  8,  9, 10, 11, 12],
       [13, 14, 15, 16, 17, 18]])

In [285]: idxs = np.array([[1, 3],
     ...:                  [0, 1],
     ...:                  [2, 2]])

In [286]: j = np.linspace(idxs,idxs 3,3,endpoint=False)

In [287]: j
Out[287]: 
array([[[1., 3.],
        [0., 1.],
        [2., 2.]],

       [[2., 4.],
        [1., 2.],
        [3., 3.]],

       [[3., 5.],
        [2., 3.],
        [4., 4.]]])

With a bit of trial and error, pair up the indices to get:

In [292]: B[np.arange(3)[:,None,None],j.astype(int).transpose(1,2,0)]
Out[292]: 
array([[[ 2,  3,  4],
        [ 4,  5,  6]],

       [[ 7,  8,  9],
        [ 8,  9, 10]],

       [[15, 16, 17],
        [15, 16, 17]]])

Or iterate as in the first case, but with an extra layer:

In [294]: np.array([[B[j,i:i 3] for i in idxs[j]] for j in range(3)])
Out[294]: 
array([[[ 2,  3,  4],
        [ 4,  5,  6]],

       [[ 7,  8,  9],
        [ 8,  9, 10]],

       [[15, 16, 17],
        [15, 16, 17]]])

With sliding windows:

In [295]: aa = np.lib.stride_tricks.sliding_window_view(A,3)

In [296]: aa.shape
Out[296]: (4, 3)

In [297]: aa
Out[297]: 
array([[1, 2, 3],
       [2, 3, 4],
       [3, 4, 5],
       [4, 5, 6]])

In [298]: aa[[1,3]]
Out[298]: 
array([[2, 3, 4],
       [4, 5, 6]])

and

In [300]: bb = np.lib.stride_tricks.sliding_window_view(B,(1,3))

In [301]: bb.shape
Out[301]: (3, 4, 1, 3)

In [302]: bb[np.arange(3)[:,None],idxs,0,:]
Out[302]: 
array([[[ 2,  3,  4],
        [ 4,  5,  6]],

       [[ 7,  8,  9],
        [ 8,  9, 10]],

       [[15, 16, 17],
        [15, 16, 17]]])

CodePudding user response：

I got it! I was almost there:

def magical_routine(A, idxs, window_len=2000):
    col_idxs = np.add.outer(idxs, np.arange(window_len))
    return np.take_along_axis(A[..., np.newaxis, :], col_idxs, axis=-1)

I just needed to always add the new axis to A's second to last dim, and then leave remaining axes alone.