How to code a function that accepts float, list or numpy.array?

Question

I have the following simple Python function:

def get_lerp_factor( a, x, b ):
    if x <= a: return 0.
    if x >= b: return 1.
    return (x - a) / (b - a)

Many numpy functions, like numpy.sin(x) can handle a float or an array.

So how can I extend this in the same manner, so that it can also handle a numpy array for x?

def get_lerp_factor( a, x_maybe_array, b ):
    out = (x_maybe_array - a) / (b - a) # this should work...
    # but now I have to clamp each element of out between 0 and 1

Would I have to specifically check the type of x, and branch accordingly?

How about:

def get_lerp_factor( a, x_anything, b ):
    x = np.array( x_anything )
    out = ...(x)
    # now typecast out back into the same type as x... will this work?

?

Answer 1

You need numpy.asarray. This takes as its first argument:

Input data, in any form that can be converted to an array. This includes lists, lists of tuples, tuples, tuples of tuples, tuples of lists and ndarrays.

and it returns:

Array interpretation of a. No copy is performed if the input is already an ndarray.

So you can implement your function like this:

import numpy as np

def get_lerp_factor(a, x, b):
    a, x, b = np.asarray(a), np.asarray(x), np.asarray(b)
    return ((x - a) / (b - a)).clip(0, 1)

This works for scalars:

>>> get_lerp_factor(0, 9, 16)
0.5625

and also for iterables:

>>> get_lerp_factor(2, range(8), 6)
array([ 0.  ,  0.  ,  0.  ,  0.25,  0.5 ,  0.75,  1.  ,  1.  ])