Split and Swap Items in Nth Dimension to Nth+1 Dimension

Question

Let's say I had a numpy array like this:

[[[ 1,  2,  3,  4],
  [ 5,  6,  7,  8],
  [ 9, 10, 11, 12],
  [13, 14, 15, 16]]]

Which has the dimensions of (1, 4, 4)

and I wanted to convert it to:

[[[[ 1,  3], [ 2,  4]],
  [[ 5,  7], [ 6,  8]],
  [[ 9, 11], [10, 12]],
  [[13, 15], [14, 16]]]]

Which has the dimensions of (1, 4, 2, 2).

Is there any function within numpy which I could use to modify the array as such?

array.reshape() is close, but not entirely because it cannot swap items.

I want to avoid using straight python as this code is being used to format some data for a neural network input, and the speed of python (given the size of the actual arrays are) is likely to be a bottleneck.

Divakar · Answer 1 · 2019-10-05T20:43:00.300

Split the last axis, then permute axes -

In [41]: a
Out[41]: 
array([[[ 1,  2,  3,  4],
        [ 5,  6,  7,  8],
        [ 9, 10, 11, 12],
        [13, 14, 15, 16]]])

In [42]: a.reshape(-1,2,2).swapaxes(1,2)
Out[42]: 
array([[[ 1,  3],
        [ 2,  4]],

       [[ 5,  7],
        [ 6,  8]],

       [[ 9, 11],
        [10, 12]],

       [[13, 15],
        [14, 16]]])

If you need the final output shape to be (1,4,2,2), add a new axis at the front with None/np.newaxis -

In [43]: a.reshape(-1,2,2).swapaxes(1,2)[None]
Out[43]: 
array([[[[ 1,  3],
         [ 2,  4]],

        [[ 5,  7],
         [ 6,  8]],

        [[ 9, 11],
         [10, 12]],

        [[13, 15],
         [14, 16]]]])

In [44]: _.shape
Out[44]: (1, 4, 2, 2)

For a generic case with the number of elements per list being known, say n, it would be -

a.reshape(-1,n,a.shape[-1]//n).swapaxes(1,2)[None]

More info onnd-arraystransformation.

score 1 · Accepted Answer · answered Oct 05 '19 at 20:37

1

You are quite close. The only thing you need to do is then swap the third and the fourth axis:

>>> np.swapaxes(a.reshape(1,4,2,2), 2,3)
array([[[[ 1,  3],
         [ 2,  4]],

        [[ 5,  7],
         [ 6,  8]],

        [[ 9, 11],
         [10, 12]],

        [[13, 15],
         [14, 16]]]])

answered Oct 05 '19 at 20:37

Willem Van Onsem

443,496
30
428
555

Though this answer directly solves my issue, Divakar's answer is more generic and likely to be significantly more useful for people who are reading this question. – delyeet Oct 05 '19 at 20:45

score 1 · Answer 3 · answered Oct 05 '19 at 20:52

You can ask reshape to use FORTRAN order:

a = np.arange(1,17).reshape(1,4,4)

a.reshape(*a.shape[:-1],2,2,order="F")
# array([[[[ 1,  3],
#          [ 2,  4]],
# 
#         [[ 5,  7],
#          [ 6,  8]],
# 
#         [[ 9, 11],
#          [10, 12]],
# 
#         [[13, 15],
#          [14, 16]]]])

Split and Swap Items in Nth Dimension to Nth+1 Dimension

3 Answers3