view of a NumPy array with non uniform stride

Question

I have 1D array of N elements. I need to create a PxR view (where PxR<N) on this array according to strides that re not uniform but rather specified in a aux list. Observe that aux list has exactly P elements. R is a constant.

Example

INPUTS:
L=[10,14,4,12,3,85,43,22,11,7,10]
R=3
aux=[2,5,7]

output= [[4,12,3],[85,43,22],[22,11,7]]

Comprehensive lists is an option.

My code:

output = [[L[a+p] for p in range(R)] for a in aux]

However, since L,R and aux are big in my case, I was wandering if there is a numpy function that do the job in a more efficient way.

I was exploring numpy.lib.stride_tricks.as_strided as an option. This would be great since it also enables me reduce memory occupation. But it seems that as_stride only works for constant strides....

Any suggestion to solve my problem?

Answer 1

With lists, that kind of list comprehension is the only option.

But if we start with arrays:

In [298]: L=np.array([10,14,4,12,3,85,43,22,11,7,10])
     ...: R=3
     ...: aux=np.array([2,5,7])

We can make an indexing array with:

In [299]: idx = aux[:,None]+np.arange(R)
In [300]: idx
Out[300]: 
array([[2, 3, 4],
       [5, 6, 7],
       [7, 8, 9]])
In [301]: L[idx]
Out[301]: 
array([[ 4, 12,  3],
       [85, 43, 22],
       [22, 11,  7]])

edit

Using a safer, and easier version as `as_strided:

In [323]: np.lib.stride_tricks.sliding_window_view(L,(3,))
Out[323]: 
array([[10, 14,  4],
       [14,  4, 12],
       [ 4, 12,  3],
       [12,  3, 85],
       [ 3, 85, 43],
       [85, 43, 22],
       [43, 22, 11],
       [22, 11,  7],
       [11,  7, 10]])
In [324]: np.lib.stride_tricks.sliding_window_view(L,(3,))[[2,5,7]]
Out[324]: 
array([[ 4, 12,  3],
       [85, 43, 22],
       [22, 11,  7]])

While [323] is a view, the subset (and [301]) is a copy.