python - 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

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.

Answer 2

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 3

This is great, and here's how to extend the addition to a 3D matrix by adding a few lines to jorgeca's code:

import numpy as np

#two 3d arrays, of different size.
b1 = np.zeros((5,5,5), dtype=np.int) # a 5x5x5 matrix of zeroes
b2 = np.ones((3,3,3), dtype=np.int)  # a 3x3x3 matrix of ones

pos_v, pos_h, pos_z = 2, 2, 2  # a 3d offset -> to plonk b2 in the corner of b1

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))
z_range1 = slice(max(0, pos_z), max(min(pos_z + b2.shape[2], b1.shape[2]), 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]))
z_range2 = slice(max(0, -pos_z), min(-pos_z + b1.shape[2], b2.shape[2]))

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

This might help someone who wants to do the same in 3d (like me).

Answer 4

Here's @jorgeca's great code as a function, with some tests - I expanded the slices to try to make it a little more readable:

import numpy as np


def addAtPos(matrix1, matrix2, xypos, inPlace=False):
    """
    Add matrix2 into matrix1 at position xypos (x,y), in-place or in new matrix.
    Handles matrix2 going off edges of matrix1.
    """

    x, y = xypos
    h1, w1 = matrix1.shape
    h2, w2 = matrix2.shape

    # get slice ranges for matrix1
    x1min = max(0, x)
    y1min = max(0, y)
    x1max = max(min(x + w2, w1), 0)
    y1max = max(min(y + h2, h1), 0)

    # get slice ranges for matrix2
    x2min = max(0, -x)
    y2min = max(0, -y)
    x2max = min(-x + w1, w2)
    y2max = min(-y + h1, h2)

    if inPlace:
        # add matrix2 into matrix1, in place
        matrix1[y1min:y1max, x1min:x1max] += matrix2[y2min:y2max, x2min:x2max]
    else:
        # create and return a new matrix
        matrix1copy = matrix1.copy()        
        matrix1copy[y1min:y1max, x1min:x1max] += matrix2[y2min:y2max, x2min:x2max]
        return matrix1copy


def test_addAtPos():

    matrix1 = np.zeros((2,2))
    matrix2 = np.ones((2,2))

    test(addAtPos(matrix1, matrix2, ( 0, 0)), [[1,1],[1,1]])
    test(addAtPos(matrix1, matrix2, ( 2, 2)), [[0,0],[0,0]])
    test(addAtPos(matrix1, matrix2, (-1,-1)), [[1,0],[0,0]])
    test(addAtPos(matrix1, matrix2, ( 1,-1)), [[0,1],[0,0]])
    test(addAtPos(matrix1, matrix2, ( 1, 1)), [[0,0],[0,1]])
    test(addAtPos(matrix1, matrix2, (-1, 1)), [[0,0],[1,0]])


def test(actual, expected, message=''):
    "Compare actual and expected values and print OK or FAIL"
    passed = (actual == expected)
    if type(passed) == np.ndarray:
        passed = passed.all()
    actual = str(actual).replace('\n', '')
    expected = str(expected).replace('\n', '')
    if passed:
        print('[OK]  ', message, actual)
    else:
        print('[FAIL]', message, actual, ' != expected value of', expected)


test_addAtPos()

Output:

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

Answer 5

I'm sure there is a fast NumPy way to do this, but there is a more efficient way to do it even in normal Python:

block_1 = [ [ 0, 0, 0, 0, 0],
            [ 0, 0, 0, 0, 0],
            [ 0, 0, 0, 0, 0],
            [ 0, 0, 0, 0, 0]]

block_2 = [ [ 1, 1, 1],
            [ 1, 1, 1],
            [ 1, 1, 1],
            [ 1, 1, 1]]

pos = (1, 1)

x, y = pos

# width of the rows in block_2
length = len(block_2[0])

# skip the first y rows
for row_1, row_2 in zip(block_1[y:], block_2):
    # set length elements offset by x to the sum.
    row_1[x:length + x] = map(sum, zip(row_2, row_1[x:length + x]))

print '\n'.join(' '.join(map(str, row)) for row in block_1)

"""
0 0 0 0 0
0 1 1 1 0
0 1 1 1 0
0 1 1 1 0
"""