Understanding the behaviour of advanced multi-dimensional indexing on a 4D ndarray

Scenario

I have a 4D ndarray consisting of multiple 3D images/voxels with dimensions (voxels, dim1, dim2, dim3), let's say (12 voxels, 96 pixels, 96 pixels, 96 pixels). My goal is to sample a range of n slices from the middle of the volume of m voxels.

I have reviewed the Numpy documentation on (advanced) indexing, as well as this answer that explains broadcasting, and this answer that explains the insertion of a newaxis by numpy, but I am still unable to understand the underlying behaviour in my scenario.

Problem

Initially, I attempted to achieve the above by indexing the array in one go using the following code:

import numpy as np

array = np.random.rand(12, 96, 96, 96)

n = 4
m_voxels = 6
samples_range = np.arange(0, m_voxels)

middle_slices = array.shape[1] // 2
middle_slices_range = np.arange(middle_slices - n // 2, middle_slices + n // 2)

samples_from_the_middle = array[samples_range, middle_slices_range, :, :]

Instead of obtaining an array of shape (6, 4, 96, 96), I encountered the following IndexError:

IndexError: shape mismatch: indexing arrays could not be broadcast together with shapes (6,) (4,)

when I attempted to index the array explicitly or in two steps, it worked as expected:

explicit_indexing = array[0:6, 46:50, :, :]
temp = array[samples_range]
samples_from_the_middle = temp[:, middle_slices_range, :, :]
explicit_indexing.shape # output: (6, 4, 96, 96)
samples_from_the_middle.shape  # output: (6, 4, 96, 96)

Alternatively, as mentioned in this answer, another approach would be:

samples_from_the_middle = array[samples_range[:, np.newaxis], middle_slices_range, :, :]  
samples_from_the_middle.shape # output: (6, 4, 96, 96)

I have the following questions:

1. Why does the np.arange approach fail to produce the expected result while the explicit indexing (with a colon) works correctly, even though we're practically indexing with the same range of integers?
2. Why does the addition of a newaxis to the first indexing 1D-array seem to resolve the issue?

Any insights would be greatly appreciated.

So, numpy handles indexing differently depending on whether you use slices, which are what is created when you do my_array[a:b], or numpy arrays instead. A helpful way to think about it is as cartesian products. Look at this demo:

In [1]: import numpy as np

In [2]: x = np.array([[1,2,3],[4,5,6],[7,8,9]])

In [3]: x
Out[3]:
array([[1, 2, 3],
       [4, 5, 6],
       [7, 8, 9]])

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

In [5]: x[np.arange(3), np.arange(3)]
Out[5]: array([1, 5, 9])

Notice that when we use slices, we get the output you wanted. When we use numpy arrays, we instead get a 1D array with only 3 elements instead of 9. Why? This is because slices are automatically used to create cartesian products. Python is automatically generating indices of the form [0, 0], [0, 1], [0, 2], [1, 0], ... for all possible pairs of values from the two slices.

When using numpy arrays for indexing, this is not the case. Instead, the arrays are matched elementwise. That means only the pairs [0, 0], [1, 1], [2, 2] are created, and we get only the 3 diagonal elements. This has to do with numpy not treating 1D arrays as proper row or column vectors unless we explicitly state how many rows and columns an array has. When we do this, we enable numpy to do broadcasting, where, in essence, arrays are "repeated" along axes where their length is 1. This lets us do things like

In [10]: x = np.array([1,2,3,4,5])

In [11]: y = np.array([6,7,8])

In [12]: from numpy import newaxis as nax

In [13]: x = x[:, nax]

In [14]: y = y[nax, :]

In [15]: x + y
Out[15]:
array([[ 7,  8,  9],
       [ 8,  9, 10],
       [ 9, 10, 11],
       [10, 11, 12],
       [11, 12, 13]])

where you can see we obtained exactly the behavior you were looking for when indexing! Each element from the x array was paired with each element from the y array.

Now we can use this knowledge as follows:

In [16]: x = np.array([[1,2,3],[4,5,6],[7,8,9]])

In [17]: x
Out[17]:
array([[1, 2, 3],
       [4, 5, 6],
       [7, 8, 9]])

In [18]: x[0:3, 0:3]
Out[18]:
array([[1, 2, 3],
       [4, 5, 6],
       [7, 8, 9]])

In [19]: x[np.arange(3), np.arange(3)]
Out[19]: array([1, 5, 9])

In [20]: x[np.arange(3)[:, nax], np.arange(3)[nax, :]]
Out[20]:
array([[1, 2, 3],
       [4, 5, 6],
       [7, 8, 9]])

And we're done!

Just for completeness, note that the numpy.ix_ function exists precisely to take care of this for you. Here's an example:

In [21]: x = np.array([1,2,3,4,5])

In [22]: y = np.array([6,7,8])

In [23]: x, y = np.ix_(x,y)

In [24]: x
Out[24]:
array([[1],
       [2],
       [3],
       [4],
       [5]])

In [25]: y
Out[25]: array([[6, 7, 8]])

And finally, all of these are equivalent to using the numpy.meshgrid function, which explicitly creates arrays with every possible pairing of elements from x and y. You don't want to use this for indexing, however, since it's very wasteful regarding memory to explicitly create these pairings all at the same time and keep them in your RAM. Better let numpy work its magic for you.