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
)