Search code examples
pythonarraysnumpy

Zero pad a numpy n-dimensional array

Not a duplicate of Zero pad numpy array (that I posted 9 years ago, ouch!) because here it's about n-dimensional arrays.

How to zero pad a numpy n-dimensional array, if possible in one line?

Example:

a = np.array([1, 2, 3])
zeropad(a, 8)  # [1, 2, 3, 0, 0, 0, 0, 0]

b = np.array([[1, 2], [3, 4], [5, 6]])
zeropad(b, (5, 2))  # [[1, 2], [3, 4], [5, 6], [0, 0], [0, 0]]

When using b.resize((5, 2)), here it works, but in some real cases, it gives:

ValueError: cannot resize this array: it does not own its data

How to zero pad numpy nd arrays no matter if it owns its data or not?

Solution

  • Instead of using pad, since you want to pad after, you could create an array of zeros and assign the existing values:

    out = np.zeros(pad, dtype=arr.dtype)
out[np.indices(arr.shape, sparse=True)] = arr

    Or, if you only want to pad the first dimension, with resize. Just ensure that the array owns its data with copy:

    out = arr.copy()
out.resize(pad)

    Output:

    array([[1, 2],
       [3, 4],
       [5, 6],
       [0, 0],
       [0, 0]])

    I really want a one-liner resize:

    IMO there is no good reason for that, but you could always use an assignment expression (python ≥ 3.8):

    (out:=arr.copy()).resize(pad)

    Output for a different pad

    arr = np.array([[1, 2], [3, 4], [5, 6]])
pad = (5, 3)

# output zeros + assignment
array([[1, 2, 0],
       [3, 4, 0],
       [5, 6, 0],
       [0, 0, 0],
       [0, 0, 0]])

# output resize
array([[1, 2, 3],
       [4, 5, 6],
       [0, 0, 0],
       [0, 0, 0],
       [0, 0, 0]])