如何编写接受float,list或numpy.array的Python函数？

Question

我有以下简单的Python函数:

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

许多numpy函数,如numpy.sin(x)可以处理浮点数或数组.

那么如何以相同的方式扩展它,以便它还可以处理x的numpy数组？

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

我是否必须专门检查x的类型,并相应地进行分支？

怎么样:

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

你需要numpy.asarray.这是第一个参数:

输入数据,可以转换为数组的任何形式.这包括列表,元组列表,元组,元组元组,列表元组和ndarray.

它返回:

数组解释a.如果输入已经是ndarray,则不执行复制.

所以你可以像这样实现你的功能:

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)

这适用于标量:

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

以及迭代:

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