The overlapping blocks on the 2nd dimension require something like as_strided
. B.M.
has an earlier solution with as_strided
, but I find it hard to understand. It also looks dangerous, since it displays out of bounds data. I'm approaching the task in smaller steps.
Selection along the last dimension for 'cols' [0,1] and [2,3] is relatively easy. Reshaping makes it easier.
In [27]: A=np.arange(36).reshape(3,3,4)
In [28]: A1=A.reshape(3,3,2,2)
In [29]: A2=A1[[0,1,2],:,[0,1,1],:]
In [30]: A2
Out[30]:
array([[[ 0, 1],
[ 4, 5],
[ 8, 9]],
[[14, 15],
[18, 19],
[22, 23]],
[[26, 27],
[30, 31],
[34, 35]]])
Focusing on one subarray, I found I could see it as 2 overlapping arrays:
In [59]: as_strided(A2[0],shape=(2,2,2),strides=(8,8,4))
Out[59]:
array([[[0, 1],
[4, 5]],
[[4, 5],
[8, 9]]])
np.lib.stride_tricks.as_strided
is a tricky function to use. I've seen it mostly used for moving-windows applications.
Applied to the whole array:
In [65]: A3=as_strided(A2,shape=(3,2,2,2),strides=(24,8,8,4))
In [66]: A3
Out[66]:
array([[[[ 0, 1],
[ 4, 5]],
...
[[30, 31],
[34, 35]]]])
And the target:
In [71]: A3[[0,1,2],[0,1,0]]
Out[71]:
array([[[ 0, 1],
[ 4, 5]],
[[18, 19],
[22, 23]],
[[26, 27],
[30, 31]]])
This can be put together in a way that allows assignment (shape and strides are adapted from the values for A1
)
In [105]: A1 = A.reshape(3,3,2,2)
In [106]: A1s = as_strided(A1, shape=(3,2,2,2,2), strides=(48,16,16,8,4))
In [107]: A1s[[0,1,2],[0,1,0],:,[0,1,1],:]
Out[107]:
array([[[ 0, 1],
[ 4, 5]],
[[18, 19],
[22, 23]],
[[26, 27],
[30, 31]]])
Test of assignment:
In [108]: A1s[[0,1,2],[0,1,0],:,[0,1,1],:] = 99
In [109]: A
Out[109]:
array([[[99, 99, 2, 3],
[99, 99, 6, 7],
[ 8, 9, 10, 11]],
[[12, 13, 14, 15],
[16, 17, 99, 99],
[20, 21, 99, 99]],
[[24, 25, 99, 99],
[28, 29, 99, 99],
[32, 33, 34, 35]]])
If you don't require assignment, it's easier to understand the logic without striding
:
np.array([A2[0,:-1], A2[1,1:], A2[2,:-1]])
slices=(slice(-1),slice(1,None))
np.array([A2[i,slices[j]] for i,j in zip([0,1,2],[0,1,0])])
===========================
Earlier stumbling around for an answer:
I think it can be extracted with advanced indexing.
Ind = np.ix_([0,0,1,1,2,2], [0,1,1,2,0,1], [0,1,2,3,2,3])
a[ind]
The values may need tweaking, since I can't test it here.
The idea is to enumerate which rows , cols are needed in each dimension, and use ix_
(I think that's the right function) to add newaxis
so they broadcast together.
Your idea of generalizing the 2d case is right. The trick is figure out when you need np.array([0,1,2])
, and when you need to rotate it with ix_
or [:,None]
,etc. I end up testing various ideas.
It may need to be
Ind = np.ix_([0,1,2], [0,1,1,2,0,1], [0,1,2,3,2,3])
No this isn't quite right. It will produce a 3x6x6; you want 3x2x2.
Reshaping to 3x3x2x2 may make indexing one the last dimension easier.
The striding trick mentioned in another answer may help turn the overlapping selection on 2nd into a similar 2x2 block. But I'd have to play around with this.
What I'm imaging is an index tuple of the form ([1,2,3],[?,?,?],:,[?,?,?],:)