python - 如何编写接受 float 、列表或 numpy.array 的函数？

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) 可以处理 float 或数组。

那么我如何以相同的方式扩展它，以便它也可以处理 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 .这作为它的第一个参数:

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.

它返回:

Array interpretation of a. No copy is performed if the input is already an 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.  ])