Search code examples
pythonnumpyarray-broadcasting

How to force a function to broadcast without invoking `np.vectorize`


I want to look for a way to force a function to broadcast.

There are scenarios in which the function/method may be overwritten in a later instance, to constant function. In such case if

arr = np.arange(0, 1, 0.0001)
f = lambda x: 5
f(arr) # this gives just integer 5, i want [5, 5,..., 5]

I am aware of methods like np.vectorize which force the function to broadcast, but the problem is this is inefficient, as it is essentially for loop under the hood. (see documentation)

We can also use factory methods like np.frompyfunc which allows us to transform python function to numpy universal function ufunc See here for instance. This outperformed np.vectorize, but still is way less efficient than builtin ufunc methods.

I was wondering if there is any efficient numpy way of handling this, namely to force the function to broadcast?


Solution

  • If there was a better way to make arbitrary Python functions broadcast, numpy.vectorize would use it. You really have to write the function with broadcasting in mind if you want it to broadcast efficiently.

    In the particular case of a constant function, you can write a broadcasting constant function using numpy.full:

    def f(x):
        return numpy.full(numpy.shape(x), 5)
    

    numba.vectorize can also vectorize functions more effectively than numpy.vectorize, but you need Numba, and you need to write your function in a way that Numba can compile efficiently.