Padding an N-dimension Matrix with different paddings for each element in the matrix - Python

Question

I am trying to pad two dimensions of an N-dimensional matrix with different paddings and override the values. Consider the Following example:

def determineShifts(layer):
    u = range(0, 2*layer + 1)
    b = range(0, 2*layer + 1)
    shifts = []
    mat = np.zeros((2 * layer + 1, 2 * layer + 1), dtype=object)
    for x, i in enumerate(u):
        for y, j in enumerate(b):
            up = (j, 2*layer - j)
            left = (i, 2*layer - i)
            mat[x, y] = (left, up)
    return mat

layer = 1

b = np.ones((3,3,3,3))
shifts = determineShifts(layer)

I want to pad the second last and final dimension of the array b such that the resulting shape is (3,3,5,5) and override the element of that matrix and repeat the process for all nodes, which in this case is (3,3). I would prefer to override the values (currently I receive a broadcasting error) rather than making a copy of the desired shape and iterating through the first and second dimension. A sample is included below:

c = np.ones((3,3,5,5))
for i in range(np.shape(c)[0]):
    for j in range(np.shape(c)[1]):
        c[i,j] = np.pad(b[i,j], shifts[i,j])

Is there some way to apply a function to the matrix to apply all the shifts to each of the elements (3,3, 3, 3) -> (3, 3, 5, 5) such that the code is computationally efficient?

Answer 1

np.pad() accepts different padding values for each axis, but not different ones within each axis, as per your example.

One general approach is to do a bit of arithmetic for the relocation of elements and then use fancy indexing. In your case, it looks like you are trying to stagger the 2D blocks of the last two dimensions in such a way that they move by 1: vertically for axis 0 and horizontally for axis 1.

You can do the same with the following arithmetic:

def stagger_ix(s):
    r = np.arange(np.prod(s))
    block = r // np.prod(s[-2:])
    shift_i, shift_j = block // s[1], block % s[1]
    i, j = r // s[-1] % s[-2], r % s[-1]
    
    newshape = np.array(s)
    newshape[-2:] += newshape[:2] - 1
    
    ix = (
        block * np.prod(newshape[-2:])
        + (i + shift_i) * newshape[-1]
        + (j + shift_j)
    )
    return newshape, ix

def stagger(b):
    newshape, ix = stagger_ix(b.shape)

    # now insert b in a zero(newshape), as per shift logic
    c = np.zeros(np.prod(newshape), dtype=b.dtype)
    c[ix] = b.ravel()
    c = c.reshape(newshape)
    return c

Your c array can be obtained as:

c = stagger(np.ones((3,3,3,3)))

Other examples -

# for example matrices
def rp1(s):
    return (np.arange(np.prod(s)) + 1).reshape(s)

>>> stagger(rp1((2,2,2,2)))
array([[[[ 1,  2,  0],
         [ 3,  4,  0],
         [ 0,  0,  0]],

        [[ 0,  5,  6],
         [ 0,  7,  8],
         [ 0,  0,  0]]],


       [[[ 0,  0,  0],
         [ 9, 10,  0],
         [11, 12,  0]],

        [[ 0,  0,  0],
         [ 0, 13, 14],
         [ 0, 15, 16]]]])

>>> stagger(rp1((2,3,2,5)))
array([[[[ 1,  2,  3,  4,  5,  0,  0],
         [ 6,  7,  8,  9, 10,  0,  0],
         [ 0,  0,  0,  0,  0,  0,  0]],

        [[ 0, 11, 12, 13, 14, 15,  0],
         [ 0, 16, 17, 18, 19, 20,  0],
         [ 0,  0,  0,  0,  0,  0,  0]],

        [[ 0,  0, 21, 22, 23, 24, 25],
         [ 0,  0, 26, 27, 28, 29, 30],
         [ 0,  0,  0,  0,  0,  0,  0]]],


       [[[ 0,  0,  0,  0,  0,  0,  0],
         [31, 32, 33, 34, 35,  0,  0],
         [36, 37, 38, 39, 40,  0,  0]],

        [[ 0,  0,  0,  0,  0,  0,  0],
         [ 0, 41, 42, 43, 44, 45,  0],
         [ 0, 46, 47, 48, 49, 50,  0]],

        [[ 0,  0,  0,  0,  0,  0,  0],
         [ 0,  0, 51, 52, 53, 54, 55],
         [ 0,  0, 56, 57, 58, 59, 60]]]])