I have a 3-D numpy
array (let's call it a
) with shape (74, 74, 4563)
, and I want to extract a length-n
sub-array from each location in the first two dimensions. However, each of those sub-arrays starts in a different place, depending on the indices in the first two dimensions, i
& j
.
For example, if n=1000
, I may want a[0, 0, 0:1000]
, but also a[0, 1, 2:1002]
, etc... I have a 2-d array (called ix0
) which is a 2-d array that tells me where each sub-array starts for each i
/j
position. Finally, I am guaranteed that there will not be any "overflow"--that is, all the values in ix0 + n
are smaller than the dimension-2 length of a
(so we don't need to worry about asking for an index beyond the range that is present).
For example...
a = np.arange(74*74*4563).reshape(74, 74, 4563) ix0 = np.arange(74*74).reshape(74,74)/2 + 50 a[:, :, ix0:ix0+n]
which produces
IndexError: failed to coerce slice entry of type numpy.ndarray to integer
Is there a way to do this without looping through all the i
/j
index combinations or creating a big mask array?
Something along this line has been asked before, but for 2d. I may try to look that up.
But here's quick example of what was going on in the 2d case
In [1463]: x=np.arange(12).reshape(3,4) In [1464]: ix0=np.array([0,2,1]) In [1465]: N=2
We could iterate over each row of x
, collecting the desired N
length slice, and then join them into a list or array. A more general problem varies the length of slices, in which case they can't be reassembled into an array.
In [1466]: [x[i,ix0[i]:ix0[i]+N] for i in range(3)] Out[1466]: [array([0, 1]), array([6, 7]), array([ 9, 10])]
and then wrap that list in np.array
.
An alternative is to concatenate the indexes first:
In [1467]: x[np.arange(3)[:,None], np.array([np.r_[ix0[i]:ix0[i]+N] for i in range(3)])] Out[1467]: array([[ 0, 1], [ 6, 7], [ 9, 10]])
The last index array is:
In [1468]: np.array([np.r_[ix0[i]:ix0[i]+N] for i in range(3)]) Out[1468]: array([[0, 1], [2, 3], [1, 2]])
To apply to the 3d case we have two options. One is reshape it to 2d, apply one of these strategies, and reshape back. The other is to generalize the action I took to create these - that shouldn't be too hard, but will take some experimenting.
That last array shouldn't be hard to create with broadcasting.
In [1469]: ix0[:,None]+np.arange(N) Out[1469]: array([[0, 1], [2, 3], [1, 2]]) In [1470]: x[np.arange(3)[:,None], ix0[:,None]+np.arange(N)] Out[1470]: array([[ 0, 1], [ 6, 7], [ 9, 10]])
Now it should be even easier to generalize to 3d
In [1487]: X=np.arange(2*3*10).reshape(2,3,10) In [1488]: ix0=np.arange(2*3).reshape(2,3) In [1489]: ix0[...,None]+np.arange(N) Out[1489]: array([[[0, 1], [1, 2], [2, 3]], [[3, 4], [4, 5], [5, 6]]]) In [1490]: I,J,_=np.ix_(range(2),range(3),range(N)) In [1491]: I.shape Out[1491]: (2, 1, 1) In [1492]: J.shape Out[1492]: (1, 3, 1) In [1493]: X[I, J, ix0[...,None]+np.arange(N)] Out[1493]: array([[[ 0, 1], [11, 12], [22, 23]], [[33, 34], [44, 45], [55, 56]]])
I should make sure the values are right, but the shapes match, which in this sort of thing is 80% of work.