Search code examples
numpyarray-broadcasting

Broadcasting a numpy array into an array of larger size using an index array


I have a very large 2D numpy array (A) and a smaller 2D array (B) that is smaller in both dimensions. B is square. I have an index array that is the same length as B. Like this:

A = np.array([[0, 0, 0, 0, 0, 0, 0, 0, 0],
             [0, 0, 0, 0, 0, 0, 0, 0, 0],
             [0, 0, 0, 0, 0, 0, 0, 0, 0],
             [0, 0, 0, 0, 0, 0, 0, 0, 0],
             [0, 0, 0, 0, 0, 0, 0, 0, 0],
             [0, 0, 0, 0, 0, 0, 0, 0, 0],
             [0, 0, 0, 0, 0, 0, 0, 0, 0],
             [0, 0, 0, 0, 0, 0, 0, 0, 0],
             [0, 0, 0, 0, 0, 0, 0, 0, 0]])
B = np.array([[2, 4, 5, 7],
              [2, 7, 1, 3],
              [6, 1, 3, 5],
              [1, 6, 3, 2]])
Index = np.array([0, 1, 5, 6])

I want to add the contents of B to A in the rows and columns defined by Index. A is not necessarily full of 0's so it needs to be addition. The result would look like this:

R = np.array([2, 4, 0, 0, 0, 5, 7, 0, 0],
             [2, 7, 0, 0, 0, 1, 3, 0, 0],
             [0, 0, 0, 0, 0, 0, 0, 0, 0],
             [0, 0, 0, 0, 0, 0, 0, 0, 0],
             [0, 0, 0, 0, 0, 0, 0, 0, 0],
             [6, 1, 0, 0, 0, 3, 5, 0, 0],
             [1, 6, 0, 0, 0, 3, 2, 0, 0],
             [0, 0, 0, 0, 0, 0, 0, 0, 0],
             [0, 0, 0, 0, 0, 0, 0, 0, 0]])

I've tried doing it with this:

A[Index, Index] += B

Which gives me the following error ValueError: shape mismatch: value array of shape (4,4) could not be broadcast to indexing result of shape (4,)

I don't understand why this doesn't work, because I can do this:

A[Index, Index[0]:Index[0]+4] = B

And then values will all appear in the correct rows but not the correct columns. Which makes sense, because I'm just giving it a block of 4 columns for it to appear in, but I would like it to appear in the columns defined by Index.

I also tried np.add.at(A[Index], Index, B) which gives me IndexError: index 6 is out of bounds for axis 0 with size 4. I don't think I understand why it's trying access index 6 in B, which has size 4.


Solution

  • Indexing with arrays uses broadcasting. Using 2 (4,) arrays as indices, produces s (4,) shape result:

    In [205]: A[Index,Index].shape
    Out[205]: (4,)
    
    In [206]: np.ix_(Index,Index)
    Out[206]: 
    (array([[0],
            [1],
            [5],
            [6]]),
     array([[0, 1, 5, 6]]))
    

    ix_ makes a (4,1) and (1,4) array. Those broadcast together to index a (4,4) shape:

    In [207]: A[np.ix_(Index,Index)].shape
    Out[207]: (4, 4)
    

    which matches your B:

    In [208]: A[np.ix_(Index,Index)] = B
    

    There's a major numpy indexing page, and these kinds of details are handled in the advanced indexing section.