make scalar python function accept list and arrays

Question

In python, is there an easy and efficient way to make function f(x: float) accept both lists and numpy arrays as arguments (in which case I would want to apply f element-wise and return the result in the same format as it was sent in)? For now, I need only 1-dimensional arrays.

As an illustration (my real f is more complex), let's say that I have:

def f(x):
    return math.log(x) if x > 0 else 0.0

Then this one works, but is not that elegant - and possibly not that efficient either, because of the recursion (which I use as it allows me to have just one function):

def f(x):
    if np.isscalar(x):
        return math.log(x) if x > 0 else 0.0
    elif isinstance(x, np.ndarray):
        return np.array([f(i) for i in x], dtype=float)
    else:
        return [f(i) for i in x]

Is there a better way?

Answer 1

Using a decorator function would be a good option if this is something you'd need to for more than just a single function definition in your code-base. It does mean that you have to wrap your code, but I'm assuming that what you mean by "having just one function" is that you do want the function expression in f(x) to be the same for all cases.

Using your existing code, the decorator function would look like this:

def elementwise_possible(func):
    def wrapper(x):
        if np.isscalar(x):
            return func(x)
        elif isinstance(x, np.ndarray):
            return np.array([func(i) for i in x], dtype=float)
        else:
            return [func(i) for i in x]
    return wrapper 

And you would write your function like this:

@elementwise_possible
def f(x):
    return math.log(x) if x > 0 else 0.0

The resulting output becomes

In[2]: A = f(2)
In[3]: A
Out[3]: 0.6931471805599453

In[4]: B = f(np.array([2,3,4]))
In[5]: B
Out[5]: array([0.69314718, 1.09861229, 1.38629436])

In[6]:C = f([5,6,7])
In[7]:C
Out[7]: [1.6094379124341003, 1.791759469228055, 1.9459101490553132]

I think efficiency should be the same.