pythonnumpyarray-broadcasting

Expand numpy array to be able to broadcast with second array of variable depth

I have a function which can take an np.ndarray of shape (3,) or (3, N), or (3, N, M), etc.. I want to add to the input array an array of shape (3,). At the moment, I have to manually check the shape of the incoming array and if neccessary, expand the array that is added to it, so that I don't get a broadcasting error. Is there a function in numpy that can expand my array to allow broadcasting for an input array of arbitrary depth?

def myfunction(input_array):
    array_to_add = np.array([1, 2, 3])
    if len(input_array.shape) == 1:
        return input_array + array_to_add
    elif len(input_array.shape) == 2:
        return input_array + array_to_add[:, None]
    elif len(input_array.shape) == 3:
        return input_array + array_to_add[:, None, None]
    ...
Solution

  • One option would be to transpose before and after the addition:

    (input_array.T + array_to_add).T

    You could also use expand_dims to add the extra dimensions:

    (np.expand_dims(array_to_add, tuple(range(1, input_array.ndim)))
 + input_array
)

    Alternatively with broadcast_to on the reversed shape of input_array + transpose:

    (np.broadcast_to(array_to_add, input_array.shape[::-1]).T
 + input_array
)

    Or with a custom reshape:

    (array_to_add.reshape(array_to_add.shape+(1,)*(input_array.ndim-1))
 + input_array
)