pythonnumpymultidimensional-arraypaddingnumpy-ndarray

Numpy slicing on a zero-padded 2D array


Given a 2D Numpy array, I'd like to be able to pad it on the left, right, top, bottom side, like the following pseudo-code. Is there anything like this already built into Numpy?

import numpy as np

a = np.arange(16).reshape((4, 4))
# [[ 0  1  2  3]
#  [ 4  5  6  7]
#  [ 8  9 10 11]
#  [12 13 14 15]]

pad(a)[-4:2, -1:3]   # or any syntax giving the same result

#[[0 0 0 0]
# [0 0 0 0]
# [0 0 0 0]
# [0 0 0 0]
# [0 0 1 2]
# [0 4 5 6]]

pad(a)[-4:2, -1:6]

#[[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 1 2 3 0 0]
# [0 4 5 6 7 0 0]]

Solution

  • Some trickery:

    import numpy as np
    
    class PaddedArray:
        def __init__(self, arr):
            self.arr = arr
        def __getitem__(self, idx):
            idx = tuple(
                slice(s.start or 0, s.stop or size)
                for size, s in zip(self.arr.shape, idx)
            )
            slices = tuple(
                slice(max(0, s.start), min(size, s.stop))
                for size, s in zip(self.arr.shape, idx)
            )
            paddings = tuple(
                (max(0, -s.start), max(0, s.stop - size))
                for size, s in zip(self.arr.shape, idx)
            )
            return np.pad(self.arr[slices], paddings)
    

    Testing:

    >>> PaddedArray(np.arange(16).reshape((4, 4)))[-4:2, -1:3]
    array([[0, 0, 0, 0],
           [0, 0, 0, 0],
           [0, 0, 0, 0],
           [0, 0, 0, 0],
           [0, 0, 1, 2],
           [0, 4, 5, 6]])
    >>> PaddedArray(np.arange(16).reshape((4, 4)))[-4:2, -1:6]
    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, 1, 2, 3, 0, 0],
           [0, 4, 5, 6, 7, 0, 0]])
    >>> PaddedArray(np.arange(32).reshape((4, 4, 2)))[-4:2, -1:3, :]
    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,  1],
            [ 2,  3],
            [ 4,  5]],
    
           [[ 0,  0],
            [ 8,  9],
            [10, 11],
            [12, 13]]])