Numpy reshape 2D array into 3D using mapping for axes 0 and 1-CodePudding

For my testcase, I have an array with 32 rows and 5 columns.

test = np.arange(0,160).reshape(32,-1, order='F')
print(test)
[[  0  32  64  96 128]
 [  1  33  65  97 129]
 [  2  34  66  98 130]
 [  3  35  67  99 131]
 [  4  36  68 100 132]
 [  5  37  69 101 133] ... and so on

In my case columns are timeseries and rows indicate a recording location. I now would like to rearrange the recording locations to their true position (idxs is this mapping), and instead expand the timeseries into a third dimension.

idxs = np.array([[ 0,  7,  8, 15, 16, 23, 24, 31],
       [ 1,  6,  9, 14, 17, 22, 25, 30],
       [ 2,  5, 10, 12, 18, 21, 26, 29],
       [ 3,  4, 11, 12, 19, 20, 27, 28]])

Desired output: a shape (4,8) array for each column in test, arranged following the mapping in idxs, stacked along axis 2 to obtain:

after_reshape.shape == (4,8,5) #for this testcase only 5 
timepoints/columns


print(after_reshape[:,:,0])
[[ 0  7  8 15 16 23 24 31]
 [ 1  6  9 14 17 22 25 30]
 [ 2  5 10 12 18 21 26 29]
 [ 3  4 11 12 19 20 27 28]]

print(after_reshape[:,:,1])
[[32 39 40 47 48 55 56 63]
 [33 38 41 46 49 54 57 62]
 [34 37 42 44 50 53 58 61]
 [35 36 43 44 51 52 59 60]]

If I use a for loop I could repeat test[:,i][idxs] to obtain the (4,8) arrays, but would still need to stack them. I assume there is an option without looping, too?

CodePudding user response：

This should work:

test = np.arange(0,160).reshape(32,-1, order='F')
idxs = np.array([[ 0,  7,  8, 15, 16, 23, 24, 31],
                 [ 1,  6,  9, 14, 17, 22, 25, 30],
                 [ 2,  5, 10, 12, 18, 21, 26, 29],
                 [ 3,  4, 11, 12, 19, 20, 27, 28]])
after_reshape = test[idxs.flatten()].reshape([4, 8, -1])

Then after_reshape[:, :, 0] gives you:

array([[ 0,  7,  8, 15, 16, 23, 24, 31],
       [ 1,  6,  9, 14, 17, 22, 25, 30],
       [ 2,  5, 10, 12, 18, 21, 26, 29],
       [ 3,  4, 11, 12, 19, 20, 27, 28]])

and after_reshape[:, :, 1]:

array([[32, 39, 40, 47, 48, 55, 56, 63],
       [33, 38, 41, 46, 49, 54, 57, 62],
       [34, 37, 42, 44, 50, 53, 58, 61],
       [35, 36, 43, 44, 51, 52, 59, 60]])

as requested