The most efficient way to assign several small matrices onto a large matrix in numpy

Question

I have a big matrix A with shape (10, 10)

array([[2, 1, 2, 1, 1, 4, 3, 2, 2, 2],
       [3, 2, 1, 2, 3, 3, 2, 3, 2, 4],
       [1, 3, 3, 4, 2, 4, 4, 3, 4, 1],
       [1, 3, 1, 3, 3, 1, 4, 2, 1, 2],
       [3, 3, 1, 3, 3, 2, 3, 4, 3, 2],
       [2, 4, 1, 4, 2, 1, 1, 2, 1, 1],
       [2, 3, 2, 3, 1, 4, 3, 1, 2, 3],
       [3, 1, 3, 2, 2, 4, 2, 3, 3, 3],
       [1, 2, 3, 2, 1, 3, 4, 4, 1, 3],
       [3, 1, 3, 2, 4, 3, 1, 1, 1, 1]])

and an array of positions B with shape (5, 2)

array([[4, 5], # row 4, column 5
       [2, 1],
       [2, 5],
       [4, 1],
       [6, 7]])

and several small matrices C with shape (5, 2, 2)

array([[[7, 9],
        [6, 7]],

       [[6, 6],
        [9, 6]],

       [[9, 6],
        [8, 9]],

       [[8, 7],
        [8, 7]],

       [[8, 6],
        [7, 7]]])

Now, I want to assign these 5 small matrices to the large matrix. The positions are the position for the up-left corner of the small matrix. If there exists overlapping area, we can use the last one, maximum or just sum it up. The effect I want looks like

A[B] += C

A for loop implementation looks like:

for i in range(B.shape[0]): 
    A[B[i][0]:B[i][0]+2,B[i][1]:B[i][1]+2] += C[i]

The expected result looks like

array([[ 2,  1,  2,  1,  1,  4,  3,  2,  2,  2],
       [ 3,  2,  1,  2,  3,  3,  2,  3,  2,  4],
       [ 1,  9,  9,  4,  2, 13, 10,  3,  4,  1],
       [ 1, 12,  7,  3,  3,  9, 13,  2,  1,  2],
       [ 3, 11,  8,  3,  3,  9, 12,  4,  3,  2],
       [ 2, 12,  8,  4,  2,  7,  8,  2,  1,  1],
       [ 2,  3,  2,  3,  1,  4,  3,  9,  8,  3],
       [ 3,  1,  3,  2,  2,  4,  2, 10, 10,  3],
       [ 1,  2,  3,  2,  1,  3,  4,  4,  1,  3],
       [ 3,  1,  3,  2,  4,  3,  1,  1,  1,  1]])

Is there a solution without for loop?

Can you post your expected output as well as your attempt at a solution using a loop? — user3483203

hpaulj hpaulj · Accepted Answer · 2021-07-26T15:59:54

Your arrays:

In [58]: A = np.array([[2, 1, 2, 1, 1, 4, 3, 2, 2, 2],
    ...:        [3, 2, 1, 2, 3, 3, 2, 3, 2, 4],
    ...:        [1, 3, 3, 4, 2, 4, 4, 3, 4, 1],
    ...:        [1, 3, 1, 3, 3, 1, 4, 2, 1, 2],
    ...:        [3, 3, 1, 3, 3, 2, 3, 4, 3, 2],
    ...:        [2, 4, 1, 4, 2, 1, 1, 2, 1, 1],
    ...:        [2, 3, 2, 3, 1, 4, 3, 1, 2, 3],
    ...:        [3, 1, 3, 2, 2, 4, 2, 3, 3, 3],
    ...:        [1, 2, 3, 2, 1, 3, 4, 4, 1, 3],
    ...:        [3, 1, 3, 2, 4, 3, 1, 1, 1, 1]])
In [59]: B=np.array([[4, 5], # row 4, column 5
    ...:        [2, 1],
    ...:        [2, 5],
    ...:        [4, 1],
    ...:        [6, 7]])
In [60]: C=np.array([[[7, 9],
    ...:         [6, 7]],
    ...: 
    ...:        [[6, 6],
    ...:         [9, 6]],
    ...: 
    ...:        [[9, 6],
    ...:         [8, 9]],
    ...: 
    ...:        [[8, 7],
    ...:         [8, 7]],
    ...: 
    ...:        [[8, 6],
    ...:         [7, 7]]])

Your iteration, cleaned up a bit:

In [72]: for cnt,(i,j) in enumerate(B):
    ...:     A[i:i+2, j:j+2] += C[cnt]
    ...: 
In [73]: A
Out[73]: 
array([[ 2,  1,  2,  1,  1,  4,  3,  2,  2,  2],
       [ 3,  2,  1,  2,  3,  3,  2,  3,  2,  4],
       [ 1,  9,  9,  4,  2, 13, 10,  3,  4,  1],
       [ 1, 12,  7,  3,  3,  9, 13,  2,  1,  2],
       [ 3, 11,  8,  3,  3,  9, 12,  4,  3,  2],
       [ 2, 12,  8,  4,  2,  7,  8,  2,  1,  1],
       [ 2,  3,  2,  3,  1,  4,  3,  9,  8,  3],
       [ 3,  1,  3,  2,  2,  4,  2, 10, 10,  3],
       [ 1,  2,  3,  2,  1,  3,  4,  4,  1,  3],
       [ 3,  1,  3,  2,  4,  3,  1,  1,  1,  1]])

And to make the action clearer, lets start with a 0 array:

In [76]: A = np.zeros_like(Acopy)
In [77]: for cnt,(i,j) in enumerate(B):
    ...:     A[i:i+2, j:j+2] += C[cnt]
    ...: 
In [78]: A
Out[78]: 
array([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       [0, 6, 6, 0, 0, 9, 6, 0, 0, 0],
       [0, 9, 6, 0, 0, 8, 9, 0, 0, 0],
       [0, 8, 7, 0, 0, 7, 9, 0, 0, 0],
       [0, 8, 7, 0, 0, 6, 7, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 8, 6, 0],
       [0, 0, 0, 0, 0, 0, 0, 7, 7, 0],
       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0]])

I don't see overlap, so I think we could construct an index array from B, that would allow us to:

A[B1] += C

and if there was a overlap, it would write the last C value.

If we don't like that, there is the np.add.at ufunc that can perform unbuffered addition (or even np.max.at).

But it will take some time to work out the required B1 indices.

edit

Here's a way of using +=. I'm using linspace to construct a multidimensional index, which will be used inplace of the slices. Getting shapes right took a lot of trial and error and testing (in an interactive session). As long as the blocks don't overlap this is fast and correct. But as documented with np.add.at, this won't match the iterative approach when there are duplicate indices.

In [125]: B1 = B+2
In [126]: I = np.linspace(B,B1,2,endpoint=False).astype(int)
In [127]: A1 =np.zeros_like(Acopy)
In [128]: A1[I[:,:,0][:,None], I[:,:,1]] += C.transpose(1,2,0)
In [129]: np.allclose(A1,A)
Out[129]: True

I is a (2,5,2) shape array, where the first 2 in the number of "steps" In [130]: I Out[130]: array([[[4, 5], [2, 1], [2, 5], [4, 1], [6, 7]],

       [[5, 6],
        [3, 2],
        [3, 6],
        [5, 2],
        [7, 8]]])

And since the C subarrays are (2,2), this is the same as: np.stack([B,B+1])

The C transpose is needed since this indexing of A1 produces a (2,2,5) array:

In [134]: A1[I[:,:,0][:,None], I[:,:,1]]
Out[134]: 
array([[[7, 6, 9, 8, 8],
        [9, 6, 6, 7, 6]],

       [[6, 9, 8, 8, 7],
        [7, 6, 9, 7, 7]]])
In [135]: _.shape
Out[135]: (2, 2, 5)

If some blocks overlap, np.add.at can be used to sum the overlaps:

In [137]: A1 =np.zeros_like(Acopy)
In [138]: np.add.at(A1, (I[:,:,0][:,None], I[:,:,1]), C.transpose(1,2,0))
In [140]: np.allclose(A1,A)
Out[140]: True

or for the largest

In [143]: np.maximum.at(A1, (I[:,:,0][:,None], I[:,:,1]), C.transpose(1,2,0))
In [144]: np.allclose(A1,A)

The most efficient way to assign several small matrices onto a large matrix in numpy

3 Answers

edit