Zero pad numpy array

Question

What\'s the more pythonic way to pad an array with zeros at the end?

def pad(A, length):
    ...

A = np.array([1,2,3,4,5])
pad(A, 8)    # expected : [1,2,3,         


        

Answer 1

This should work:

def pad(A, length):
    arr = np.zeros(length)
    arr[:len(A)] = A
    return arr

You might be able to get slightly better performance if you initialize an empty array (np.empty(length)) and then fill in A and the zeros separately, but I doubt that the speedups would be worth additional code complexity in most cases.

To get the value to pad up to, I think you'd probably just use something like divmod:

n, remainder = divmod(len(A), 1024)
n += bool(remainder)

Basically, this just figures out how many times 1024 divides the length of your array (and what the remainder of that division is). If there is no remainder, then you just want n * 1024 elements. If there is a remainder, then you want (n + 1) * 1024.

all-together:

def pad1024(A):
    n, remainder = divmod(len(A), 1024)
    n += bool(remainder)
    arr = np.zeros(n * 1024)
    arr[:len(A)] = A
    return arr