cut a pice from an array and pad it with zeros

Question

I need to cut a part of given size , and given location, from an N-dim array. If the part is to large, I need to pad it with zeros to achieve the given size.

The examples are in 2D for simplicity.

The given matrix:

[[1 8 3 3 8]
 [5 8 6 7 6]
 [8 3 5 6 5]
 [2 6 2 4 6]
 [6 5 3 7 4]]

I want to cut [2,4] part, starting from index (1,2), The part I cut is not big enough for the size, so padding with zeros are needed. The wanted result:

[[6 7 6 0]
 [5 6 5 0]]

I manage to write ugly and not N-dim code to do that.

# set example numbers
matrix =  numpy.random.randint(low=1, high=9, size=(5,5))
matrix_size = np.array(matrix.shape)

# size of the part we want to have in the end
size = np.array([2, 4])
# starting point of the cut
mini = [1, 2]

#calculating max index (in the given matrix) for the part we want to cut
maxi = np.add(size - 1 , mini)
cut_max_ind = np.minimum(maxi, matrix_size - 1) + 1

# copy from matrix to cut
# ??? a way to generalize it for N-dim ???
cut = matrix[mini[0]:cut_max_ind[0], mini[1]:cut_max_ind[1]]

#culculate the padding size
padding =  np.add(matrix_size - 1, maxi*-1)
padding_size = np.minimum(np.zeros((matrix.ndim), dtype=np.uint8), padding) * -1

for j in range(0, matrix.ndim):

    if (padding_size[j]):
        pad_width = size
        pad_width[j] = padding_size[j]
        pad_pice = np.zeros((pad_width), dtype = np.uint8)
        cut = np.append(cut, pad_pice, axis = j)

print "matrix"
print matrix
print "cut"
print cut

Any Ideas for improvement and generalization ?

Answer 1

You can solve it easier by preallocating an array of zeroes and then modifying the slices to fit your needs:

a = numpy.array([
    [1, 8, 3, 3, 8],
    [5, 8, 6, 7, 6],
    [8, 3, 5, 6, 5],
    [2, 6, 2, 4, 6],
    [6, 5, 3, 7, 4],
])

def extract_piece(array, in_idx):
    # make sure number of dimensions match
    assert array.ndim == len(in_idx)
    # preallocate output array
    out = numpy.zeros([i.stop - i.start for i in in_idx], dtype=array.dtype)

    # modify reading slices to not exceed bounds
    in_idx = [slice(i.start, min(i.stop, s), i.step) for s, i in zip(array.shape, in_idx)]
    # modify writing slices to fit size of read data
    out_idx = [slice(0, i.stop - i.start) for i in in_idx]

    # Copy data
    out[out_idx] = array[in_idx]
    return out

print a
print extract_piece(a, (slice(1, 3), slice(2, 6)))

Or in a 4D example

extract_piece(
    numpy.random.rand(4, 4, 4, 4),  # 4D data
    (
        slice(0,2),  # 1st dimension
        slice(0,4),  # 2nd dimension
        slice(0,6),  # 3rd dimension
        slice(0,1),  # 4th dimension
    )
)

if you are unfamiliar with slices:

a[1:2]

is the same as

a[slice(1,2)]