What is the broadcasting order here?-CodePudding

Given the following:

polygons = np.array([
       [[1, 0],
        [1, 1],
        [0, 0],
        [0, 1]],

       [[3, 2],
        [2, 2],
        [4, 4],
        [2, 3]]])

sort_idx = np.array([
       [2, 0, 1, 3],
       [1, 0, 2, 3]])

polygons[sort_idx]  #< The problematic part

The expected output is:

array([[[0, 0],
        [1, 0],
        [1, 1],
        [0, 1]],

       [[2, 2],
        [3, 2],
        [4, 4],
        [2, 3]]])

I expect polygons[sort_idx] to return the "sorted" rows of polygons, with the sorting order given by the values in sort_idx. polygons[0][sorted_idx[0]] (and the equivalent for [1]) give the correct values, but running polygons[sort_idx] raises:

IndexError: index 2 is out of bounds for axis 0 with size 2

I estimate the issue is connected to me not understanding something of how the operation is broadcasted, but I don't really know what to search for or how to phrase the question. I saw similar questions recommending to use np.take() or polygons[sort_idx,...], but both raise the same error. What am I missing?

CodePudding user response：

Use np.take_along_axis:

np.take_along_axis(polygons, sort_idx[..., None], 1)
Out[]: 
array([[[0, 0],
        [1, 0],
        [1, 1],
        [0, 1]],

       [[2, 2],
        [3, 2],
        [4, 4],
        [2, 3]]])

CodePudding user response：

take_along_axis is supposed to 'streamline' indexing that we had to do (and still can do) with:

In [233]: polygons[np.arange(2)[:,None], sort_idx]
Out[233]: 
array([[[0, 0],
        [1, 0],
        [1, 1],
        [0, 1]],

       [[2, 2],
        [3, 2],
        [4, 4],
        [2, 3]]])

The first dimension index is (2,1), and then second (2,4), which broadcast together to select (2,4) in the first 2 dimensions (and the last dimension unchanged).

The values of sort_idx are meant to apply to the 2nd axis, the size 4 one. polygons[sort_idx] applies them to the first, size 2, dimension, hence the error.

Using slice for the first dimension returns too many values, a block. We want something more like a diagonal:

In [235]: polygons[:, sort_idx,:].shape
Out[235]: (2, 2, 4, 2)

So yes, it is a broadcasting issue. When using multiple indexing arrays, we want to think about how they broadcast against each other. Same rules as for operators apply.

In [236]: np.array([10,100])[:,None]* sort_idx
Out[236]: 
array([[ 20,   0,  10,  30],
       [100,   0, 200, 300]])