CODEWITHSUNDEEP
pythonnumpy
itthrill
itthrill

Reputation: 1376

Slice a 3D array in to multiple 2D arrays

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

Upvotes: 1

Views: 1737

Answers (2)

hpaulj
hpaulj

Reputation: 231385

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.

Upvotes: 2

Paul Panzer
Paul Panzer

Reputation: 53029

You can use moveaxis to move the to split axis all the way to the left and then use sequence unpacking:

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

Upvotes: 3

Related Questions