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?
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!