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]]
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]]])