Split last dimension of arrays in lower dimensional arrays

Question

Assume we have an array with NxMxD shape. I want to get a list with D NxM arrays.

The correct way of doing it would be:

np.dsplit(myarray, D)

However, this returns D NxMx1 arrays.

I can achieve the desired result by doing something like:

[myarray[..., i] for i in range(D)]

Or:

[np.squeeze(subarray) for subarray in np.dsplit(myarray, D)]

However, I feel like it is a bit redundant to need to perform an additional operation. Am I missing any numpy function that returns the desired result?

Answer 1

Try D.swapaxes(1,2).swapaxes(1,0)

>>>import numpy as np
>>>a = np.arange(24).reshape(2,3,4)
>>>a
array([[[ 0,  1,  2,  3],
        [ 4,  5,  6,  7],
        [ 8,  9, 10, 11]],

       [[12, 13, 14, 15],
        [16, 17, 18, 19],
        [20, 21, 22, 23]]])

>>>[a[:,:,i] for i in range(4)]
[array([[ 0,  4,  8],
       [12, 16, 20]]),
 array([[ 1,  5,  9],
       [13, 17, 21]]),
 array([[ 2,  6, 10],
       [14, 18, 22]]),
 array([[ 3,  7, 11],
       [15, 19, 23]])]

>>>a.swapaxes(1,2).swapaxes(1,0)
array([[[ 0,  4,  8],
        [12, 16, 20]],

       [[ 1,  5,  9],
        [13, 17, 21]],

       [[ 2,  6, 10],
        [14, 18, 22]],

       [[ 3,  7, 11],
        [15, 19, 23]]])

Edit: As pointed out by ajcr (thanks again), the transpose command is more convenient since the two swaps can be done in one step by using

D.transpose(2,0,1)

Answer 2

np.dsplit uses np.array_split, the core of which is:

sub_arys = []
sary = _nx.swapaxes(ary, axis, 0)
for i in range(Nsections):
    st = div_points[i]; end = div_points[i+1]
    sub_arys.append(_nx.swapaxes(sary[st:end], axis, 0))

with axis=-1, this is equivalent to:

[x[...,i:(i+1)] for i in np.arange(x.shape[-1])]  # or
[x[...,[i]] for i in np.arange(x.shape[-1])]

which accounts for the singleton dimension.

So there's nothing wrong or inefficient about your

[x[...,i] for i in np.arange(x.shape[-1])]

Actually in quick time tests, any use of dsplit is slow. It's generality costs. So adding squeeze is relatively cheap.

But by accepting the other answer, it looks like you are really looking for an array of the correct shape, rather than a list of arrays. For many operations that makes sense. split is more useful when the subarrays have more than one 'row' along the split axis, or even an uneven number of 'rows'.