student010101
student010101

Reputation: 197

Confusion regarding resulting shape of a multi-dimensional slicing of a numpy array

Suppose we have

t = np.random.rand(2,3,4)

i.e., a 2x3x4 tensor.

I'm having trouble understanding why the shape of t[0][:][:2] is 2x4 rather than 3x2. Aren't we taking the 0th, all, and the first indices of the 1st, 2nd, and 3rd dimensions, in which case that would give us a 3x2 tensor?

Upvotes: 0

Views: 63

Answers (1)

hpaulj
hpaulj

Reputation: 231355

In [1]: t = np.arange(24).reshape(2,3,4)
In [2]: t
Out[2]: 
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]]])

Select the 1st plane:

In [3]: t[0]
Out[3]: 
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11]])

[:] selects everything - ie. no change

In [4]: t[0][:]
Out[4]: 
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11]])

Select first 2 rows of that plane:

In [5]: t[0][:][:2]
Out[5]: 
array([[0, 1, 2, 3],
       [4, 5, 6, 7]])

While [i][j] works for integer indices, it shouldn't be used for slices or arrays. Each [] acts on the result of the previous. [:] is not a 'placeholder' for the 'middle dimension'. Python syntax and execution order applies, even when using numpy. numpy just adds an array class, and functions. It doesn't change the syntax.

Instead you want:

In [6]: t[0,:,:2]
Out[6]: 
array([[0, 1],
       [4, 5],
       [8, 9]])

The result is the 1st plane, and first 2 columns, (and all rows). With all 3 dimensions in one [] they are applied in a coordinated manner, not sequentially.

There is a gotcha, when using a slice in the middle of 'advanced indices'. Same selection as before, but transposed.

In [8]: t[0,:,[0,1]]
Out[8]: 
array([[0, 4, 8],
       [1, 5, 9]])

For this you need a partial decomposition - applying the 0 first

In [9]: t[0][:,[0,1]]
Out[9]: 
array([[0, 1],
       [4, 5],
       [8, 9]])

There is a big indexing page in the numpy docs that you need to study sooner or later.

Upvotes: 1

Related Questions