Search code examples
pythonarraysnumpymaskingarray-broadcasting

How does "Fancy Indexing with Broadcasting and Boolean Masking" work?


I came across this snippet of code in Jake Vanderplas's Data Science Handbook. The concept of using Broadcasting along with Fancy Indexing here wasn't clear to me. Please explain.

In[5]: X = np.arange(12).reshape((3, 4))
 X
Out[5]: array([[ 0, 1, 2, 3],
 [ 4, 5, 6, 7],
 [ 8, 9, 10, 11]])

In[6]: row = np.array([0, 1, 2])
 col = np.array([2, 1, 3])

In[7]: X[row[:, np.newaxis], col]
Out[7]: array([[ 2, 1, 3],
               [ 6, 5, 7],
              [10, 9, 11]])

It says: "Here, each row value is matched with each column vector, exactly as we saw in broadcasting of arithmetic operations. For example:"

In[8]: row[:, np.newaxis] * col
Out[8]: array([[0, 0, 0],
               [2, 1, 3],
               [4, 2, 6]])

Solution

  • If you use an integer array to index another array you basically loop over the given indices and pick the respective elements (may still be an array) along the axis you are indexing and stack them together.

    arr55 = np.arange(25).reshape((5, 5))
    # 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, 24]])
    
    arr53 = arr55[:, [3, 3, 4]]  
    # pick the elements at (arr[:, 3], arr[:, 3], arr[:, 4])
    # array([[ 3,  3,  4],
    #        [ 8,  8,  9],
    #        [13, 13, 14],
    #        [18, 18, 19],
    #        [23, 23, 24]])
    

    So if you index an (m, n) array with an row (or col) index of length k (or length l) the resulting shape is:

    A_nm[row, :] -> A_km
    A_nm[:, col] -> A_nl
    

    If however you use two arrays row and col to index an array you loop over both indices simultaneously and stack the elements (may still be arrays) at the respective position together. Here it row and col must have the same length.

    A_nm[row, col] -> A_k
    array([ 3, 13, 24])
    
    arr3 = arr55[[0, 2, 4], [3, 3, 4]]  
    # pick the element at (arr[0, 3], arr[2, 3], arr[4, 4])
    
    

    Now finally for your question: it is possible to use broadcasting while indexing arrays. Sometimes it is not wanted that only the elements

    (arr[0, 3], arr[2, 3], arr[4, 4])
    

    are picked, but rather the expanded version:

    (arr[0, [3, 3, 4]], arr[2, [3, 3, 4]], arr[4, [3, 3, 4]])
    # each row value is matched with each column vector
    

    This matching/broadcasting is exactly as in other arithmetic operations. But the example here might be bad in the sense, that not the result of the shown multiplication is of importance for the indexing. The focus here is on the combinations and the resulting shape:

    row * col  
    # performs a element wise multiplication resulting in 3 
    numbers
    row[:, np.newaxis] * col 
    # performs a multiplication where each row value is *matched* with each column vector
    

    The example wanted to emphasis this matching of row and col.

    We can have a look and play around with the different possibilities:

    n = 3
    m = 4
    X = np.arange(n*m).reshape((n, m))
    row = np.array([0, 1, 2])  # k = 3
    col = np.array([2, 1, 3])  # l = 3
    
    X[row, :]  # A_nm[row, :] -> A_km
    # array([[ 0,  1,  2,  3],
    #        [ 4,  5,  6,  7],
    #        [ 8,  9, 10, 11]])
    
    X[:, col]  # A_nm[:, col] -> A_nl
    # array([[ 2,  1,  3],
    #        [ 6,  5,  7],
    #        [10,  9, 11]])
    
    X[row, col]  # A_nm[row, col] -> A_l == A_k
    # array([ 2,  5, 11]
    
    X[row, :][:, col]  # A_nm[row, :][:, col] -> A_km[:, col] -> A_kl 
    # == X[:, col][row, :]
    # == X[row[:, np.newaxis], col]  # A_nm[row[:, np.newaxis], col] -> A_kl 
    # array([[ 2,  1,  3],
    #        [ 6,  5,  7],
    #        [10,  9, 11]])
    
    X[row, col[:, np.newaxis]]
    # == X[row[:, np.newaxis], col].T
    # array([[ 2,  6, 10],
    #        [ 1,  5,  9],
    #        [ 3,  7, 11]])