Slice a 3D array in to multiple 2D arrays

Question

I have a numpy array of shape(64,64,3). How can we get three arrays of dimension (64,64) each.

Answer 1

您可以使用moveaxis将拆分轴一直移动到最左侧，然后使用序列拆包：

x,y,z = np.moveaxis(arr, -1, 0)

Answer 2

This is another case where iteration, for a few steps, is not bad.

In [145]: arr = np.ones((64,64,3))

Unpacking a list comprehension:

In [146]: a,b,c = [arr[:,:,i] for i in range(3)]
In [147]: a.shape
Out[147]: (64, 64)

Unpacking a transposed array:

In [148]: a,b,c = np.moveaxis(arr,-1,0)
In [149]: a.shape
Out[149]: (64, 64)

Timings for such a small example aren't ground breaking, but they hint at the relative advantages of the two approaches:

In [150]: timeit a,b,c = [arr[:,:,i] for i in range(3)]
3.02 µs ± 10.6 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

In [151]: timeit a,b,c = np.moveaxis(arr,-1,0)
18.4 µs ± 946 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

The array unpacking (on the 1st dimension) requires converting the array into a list of 3 subarrays, as can be seen by these similar timings:

In [154]: timeit a,b,c = list(np.moveaxis(arr,-1,0))
17.8 µs ± 15.1 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
In [155]: timeit a,b,c = [i for i in np.moveaxis(arr,-1,0)]
17.9 µs ± 14.6 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

It's not the unpacking or iteration that taking time. It's the moveaxis :

In [156]: timeit np.moveaxis(arr,-1,0)
14 µs ± 4.41 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

Looking at its code, we see that it uses transpose , after first constructing an order from the parameters:

Calling transpose directly is fast (since it just involves changing shape and strides):

In [157]: timeit arr.transpose(2,0,1)
688 ns ± 1.87 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

And unpacking that transpose is a bit faster than my original list comprehension.

In [158]: timeit a,b,c = arr.transpose(2,0,1)
2.78 µs ± 9.3 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

So moveaxis has a significant overhead relative to the rest of the task. That said, it probably does not increase with the side of the array. It's a fixed overhead.