Adding different sized/shaped displaced NumPy matrices

Question

In short: I have two matrices (or arrays):

import numpy

block_1 = numpy.matrix([[ 0, 0, 0, 0, 0],
                        [ 0, 0, 0, 0, 0],
                        [ 0, 0, 0, 0, 0],
                        [ 0, 0, 0, 0, 0]])

block_2 = numpy.matrix([[ 1, 1, 1],
                        [ 1, 1, 1],
                        [ 1, 1, 1],
                        [ 1, 1, 1]])

I have the displacement of block_2 in the block_1 element coordinate system.

pos = (1,1)

I want to be able to add them (quickly), to get:

[[0 0 0 0 0]
 [0 1 1 1 0]
 [0 1 1 1 0]
 [0 1 1 1 0]]

In long: I would like a fast way to add two different shape matrices together, where one of the matrices can be displaced. The resulting matrix must have the shape of the first matrix, and the overlapping elements between the two matrices are summed. If there is no overlap, just the first matrix is returned unmutated.

I have a function that works fine, but it's kind of ugly, and elementwise:

def add_blocks(block_1, block_2, pos):
    for i in xrange(0, block_2.shape[0]):
        for j in xrange(0, block_2.shape[1]):
            if (i + pos[1] >= 0) and (i + pos[1] < block_1.shape[0])
               and (j + pos[0] >= 0) and (j + pos[0] < block_1.shape[1]):
                block_1[pos[1] + i, pos[0] + j] += block_2[i,j]
    return block_1

Can broadcasting or slicing perhaps do this?

I feel like maybe I'm missing something obvious.

Answer 1

An easy solution that looks like MATLAB solution is:

import numpy as np

block1 = np.zeros((5,4))
block2 = np.ones((3,2))

block1[1:4,2:4] += block2  # use array slicing

print(block1)

[[0. 0. 0. 0.]
 [0. 0. 1. 1.]
 [0. 0. 1. 1.]
 [0. 0. 1. 1.]
 [0. 0. 0. 0.]]

---

So package it as a reusable function:

import numpy as np

def addAtPos(mat1, mat2, xypos):
    """
    Add two matrices of different sizes in place, offset by xy coordinates
    Usage:
      - mat1: base matrix
      - mat2: add this matrix to mat1
      - xypos: tuple (x,y) containing coordinates
    """
    x, y = xypos
    ysize, xsize = mat2.shape
    xmax, ymax = (x + xsize), (y + ysize)
    mat1[y:ymax, x:xmax] += mat2
    return mat1

block1 = np.zeros((5,4))
block2 = np.ones((3,2))
pos = (2,1)
print(addAtPos(block1, block2, pos))

[[0. 0. 0. 0.]
 [0. 0. 1. 1.]
 [0. 0. 1. 1.]
 [0. 0. 1. 1.]
 [0. 0. 0. 0.]]

Answer 2

You just have to find the overlapping range, and then add the arrays using slicing.

b1 = np.zeros((4,5))
b2 = np.ones((4,3))
pos_v, pos_h = 2, 3  # offset
v_range1 = slice(max(0, pos_v), max(min(pos_v + b2.shape[0], b1.shape[0]), 0))
h_range1 = slice(max(0, pos_h), max(min(pos_h + b2.shape[1], b1.shape[1]), 0))

v_range2 = slice(max(0, -pos_v), min(-pos_v + b1.shape[0], b2.shape[0]))
h_range2 = slice(max(0, -pos_h), min(-pos_h + b1.shape[1], b2.shape[1]))

b1[v_range1, h_range1] += b2[v_range2, h_range2]

They're added in-place, but you could also create a new array. I might have missed some corner cases, though, but it seems to work fine.