pure numpy expression for selecting same-length subarrays with different starting indices from 3-D array

匿名 (未验证) 提交于 2019-12-03 10:10:24

问题:

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?

回答1:

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.



标签
易学教程内所有资源均来自网络或用户发布的内容,如有违反法律规定的内容欢迎反馈
该文章没有解决你所遇到的问题?点击提问,说说你的问题,让更多的人一起探讨吧!