pythonnumpyarray-broadcastingnumpy-indexing

Numpy indexing broadcasting introduces new dimension


I have an array I wan to use for mapping. Let's call it my_map ,type float shape (m,c).

I have a second array with indexes, lest call it my_indexes, type int size (n,c), every value is between 0 and m.

Trying to index my_map doing my_ans = my_map[my_indexes] I get an array of shape (n,c,c), when I was expecting (n,c). What would be the proper way to do it?

Just to be clear, what I am trying to do is something equivalent to:

my_ans = np.empty_like(touch_probability)
for i in range(c):
    my_ans[:,i] = my_map[:,i][my_indexes[:,i]]

Solution

  • To illustrate and test your problem, define simple, real arrays:

    In [44]: arr = np.arange(12).reshape(3,4)
    In [45]: idx = np.array([[0,2,1,0],[2,2,1,0]])
    In [46]: arr.shape
    Out[46]: (3, 4)
    In [47]: idx.shape
    Out[47]: (2, 4)
    

    Your desired calculation:

    In [48]: res = np.zeros((2,4), int)
    In [49]: for i in range(4):
        ...:     res[:,i] = arr[:,i][idx[:,i]]  # same as arr[idx[:,i], i]
    
        ...:     
    In [50]: res
    Out[50]: 
    array([[0, 9, 6, 3],
           [8, 9, 6, 3]])
    

    Doing the same with one indexing step:

    In [51]: arr[idx, np.arange(4)]
    Out[51]: 
    array([[0, 9, 6, 3],
           [8, 9, 6, 3]])
    

    This is broadcasting the two indexing arrays against each other, and then picking points:

    In [52]: np.broadcast_arrays(idx, np.arange(4))
    Out[52]: 
    [array([[0, 2, 1, 0],
            [2, 2, 1, 0]]), 
     array([[0, 1, 2, 3],
            [0, 1, 2, 3]])]
    

    So we are indexing the (m,c) array with 2 (n,c) arrays


    The following are the same:

    arr[idx] arr[idx, :]

    It is using idx to select whole rows from arr, so the result is shape of idx plus the last dimension of arr. Where as what you want is just the ith element of the idx[j,i] row.