Numpy To List Over 2nd Axis

I would like to split a n-d numpy array based on a internal axis. I have a array of shape (6,150,29,29,29,1) I would like a list of arrays as - [150 arrays of shape (6,29,29,29,1)

Solution 1:

arr.transpose(1,0,2,3,4,5) or np.swapaxes(arr,0,1) put the 150 dimension first. Then you can use list.

Or you could use a list comprehension

[a[:,i] for i in range(150)]

The transpose is somewhat better

In [28]: timeit list(arr.transpose(1,0,2,3,4,5))
47.7 µs ± 47.1 ns per loop (mean ± std. dev. of7 runs, 10000 loops each)
In [29]: timeit [arr[:,i] for i inrange(150)]
88.7 µs ± 22.2 ns per loop (mean ± std. dev. of7 runs, 10000 loops each)
In [32]: timeit list(np.swapaxes(arr,0,1))
49.2 µs ± 51.1 ns per loop (mean ± std. dev. of7 runs, 10000 loops each)