NumPy Get elements based on starting indices and stride

Question

I am looking for the numpythonic way to accomplish the following:

A = np.arange(1000)
x = np.array([0, 10, 20, 30], dtype=int)
dx = np.array([3, 4, 5, 6], dtype=int)

for x_, dx_ in zip(x, dx):
    print(A[x_:x_+dx_])

Looking at similar answers and the as_strided documentation it doesn't seem like I can provide varying stride/window lengths. Is there a pure numpy way to accomplish this?

Answer 1

As @NickOdell pointed out, Numpy does not support jagged arrays. However if you can afford to waste some memory, there are workarounds. For example you could use np.nan as a fill value:

import numpy as np

A = np.arange(1000)
x = np.array([0, 10, 20, 30], dtype=int)
dx = np.array([3, 4, 5, 6], dtype=int)

A_padded = np.append(A, np.nan)

idx = x[:, np.newaxis] + np.arange(max(dx))
idx = np.where(idx >= (x + dx)[:, np.newaxis], -1, idx)

print(A_padded[idx])

Which prints:

[[ 0.  1.  2. nan nan nan]
 [10. 11. 12. 13. nan nan]
 [20. 21. 22. 23. 24. nan]
 [30. 31. 32. 33. 34. 35.]]

Whether this is a valid workaround depends on your context.

Otherwise you might want to take a look at https://awkward-array.org, which supports jagged arrays.

I hope this helps!