Indexing on ndarrays result in wrong shape

Question

The following snippet

x = np.ones((10,10,10))
x = x[2,:,[2,3,7]]
print(x.shape)

results in x.shape = (3,10) instead of (10,3). How do I use a list to index the 3rd dimension to get a shape (10,3)?

Answer 1

NB: I changed your array to x = np.arange(125).reshape((5, 5, 5)) and adjusted the indexes/slices to y = x[2, :, [0, 2, 4]] to make it easier to see is being selected.

---

TL;DR You can either transpose your result, or represent the index on the first axis as a slice, and squeeze that result:

>>> x[2, :, [0, 2, 4]].T
array([[50, 52, 54],
       [55, 57, 59],
       [60, 62, 64],
       [65, 67, 69],
       [70, 72, 74]])

>>> x[2:3, :, [0, 2, 4]].squeeze()
array([[50, 52, 54],
       [55, 57, 59],
       [60, 62, 64],
       [65, 67, 69],
       [70, 72, 74]])

---

I think a more interesting question is why this happens. For the longest time, I internalized this as "numpy does weird stuff sometimes, just memorize it", but it does have an explanation that you can apply to any general case. From the link @Chrysophylaxs shared, combined advanced and basic indexing is handled by first slicing, then advanced indexing. If you do that here, you get:

>>> x[2, :]
array([[50, 51, 52, 53, 54],
       [55, 56, 57, 58, 59],
       [60, 61, 62, 63, 64],
       [65, 66, 67, 68, 69],
       [70, 71, 72, 73, 74]])

>>> x[2, :][:, [0, 2, 4]]
array([[50, 52, 54],
       [55, 57, 59],
       [60, 62, 64],
       [65, 67, 69],
       [70, 72, 74]])

Which is of the shape (5, 3), as we expected.

However, this is not what numpy is doing. As mentioned in that link, when a single index on one dimension is requested with advanced indexing on another, the single index is treated as advanced indexing instead of a slice.

Two cases of index combination need to be distinguished:

• The advanced indices are separated by a slice, Ellipsis or newaxis. For example x[arr1, :, arr2].
• The advanced indices are all next to each other. For example x[..., arr1, arr2, :] but not x[arr1, :, 1] since 1 is an advanced index in this regard.

When this is the case,

the dimensions resulting from the advanced indexing operation come first in the result array, and the subspace dimensions after that.

The dimensions resulting from the advanced indexing operation would be (1, 3) (3,) (The scalar index 2 contributes no dimensions, and the list index [0, 2, 4] contributes one dimension of size 3)^{Thanks @Mad Physicist and @Chrysophylaxs for the clarification}. This is why we get a (3, 5) array (you get a (3, 10) array)

>>> x[2, :, [0, 2, 4]]
array([[50, 55, 60, 65, 70],
       [52, 57, 62, 67, 72],
       [54, 59, 64, 69, 74]])

If you used a similar advanced indexing operation with two indices on the first dimension, you'd get a result with a shape of (2, 3, 5):

>>> x[[[2], [3]], :, [0, 2, 4]]
array([[[50, 55, 60, 65, 70],
        [52, 57, 62, 67, 72],
        [54, 59, 64, 69, 74]],

       [[75, 80, 85, 90, 95],
        [77, 82, 87, 92, 97],
        [79, 84, 89, 94, 99]]])

>>> x[[[2], [3]], :, [0, 2, 4]].shape
(2, 3, 5)

To get a (5, 3) array (or a (10, 3) array in your case), change the first index to a slice, and then squeeze the result. This allows numpy to switch back to the first case, where the slice is done before the advanced indexing.

>>> x[2:3, :, [0, 2, 4]]      # (1, 5, 3)
array([[[50, 52, 54],
        [55, 57, 59],
        [60, 62, 64],
        [65, 67, 69],
        [70, 72, 74]]])

>>> x[2:3, :, [0, 2, 4]].squeeze()
array([[50, 52, 54],
       [55, 57, 59],
       [60, 62, 64],
       [65, 67, 69],
       [70, 72, 74]])

---

Alternatively, just transpose (or move/swap axis of) the result you get from your regular indexing:

>>> x[2, :, [0, 2, 4]].T
array([[50, 52, 54],
       [55, 57, 59],
       [60, 62, 64],
       [65, 67, 69],
       [70, 72, 74]])

Answer 2

I think the strides and base give some added insight, if not an actual explanation.

Make a 3d array:

In [77]: x = np.arange(24).reshape(2,3,4); x.shape, x.strides
Out[77]: ((2, 3, 4), (48, 16, 4))

The slice-in-the-middle case that gives 'weird' shape:

In [78]: y = x[1,:,[1,2]]; y.shape, y.strides
Out[78]: ((2, 3), (12, 4))

In [79]: y
Out[79]: 
array([[13, 17, 21],
       [14, 18, 22]])

That y is a copy (own base), and the strides are normal for that shape.

Now try the two-step indexing that gives the expected shape:

In [80]: z = x[1][:,[1,2]]; z.shape, z.strides
Out[80]: ((3, 2), (4, 12))

The strides is "reversed", what we get from a transpose.

In [81]: z
Out[81]: 
array([[13, 14],
       [17, 18],
       [21, 22]])

And in fact its base is the same y array.

In [82]: z.base
Out[82]: 
array([[13, 17, 21],
       [14, 18, 22]])

Advanced index of columns (for 2d array) produces this transpose view.

The y case looks a lot like it the the [:,[1,2]] indexing, but couldn't (for one reason or other) perform the transpose that we see in z.

With the fully broadcasted version:

In [87]: w = x[1,np.arange(3)[:,None],[1,2]]; w.shape, w.strides
Out[87]: ((3, 2), (8, 4))

In [88]: w
Out[88]: 
array([[13, 14],
       [17, 18],
       [21, 22]])

This has the shape and values as z, but the strides are different.

And if we replace the first scalar index with a slice:

In [103]: u = x[1:2,:,[1,2]][0]; u.shape, u.strides
Out[103]: ((3, 2), (4, 12))

In [104]: u.base
Out[104]: 
array([[[13, 17, 21]],

       [[14, 18, 22]]])

In [105]: u.base.shape
Out[105]: (2, 1, 3)

Note that the third dimension [1,2] index is first, just as in the 'slice-in-the-middle' case; it transposes the (2,1,3) base to a (1,3,2), which I reduced to (3,2).