python - 接受标量或 numpy 数组作为参数的 python 函数

Question

正如标题所说，假设我想写一个符号函数(让我们暂时忘记 sign(0))，显然我们期望 sign(2) = 1 和 sign(array([-2,-2,2] )) = 数组([-1,-1,1])。但是，以下函数将不起作用，因为它无法处理 numpy 数组。

def sign(x):
    if x>0: return 1
    else: return -1

下一个函数也不会工作，因为如果 x 只是一个数字，它就没有形状成员。即使使用像 y = x*0 + 1 这样的技巧，y 也不会有 [] 方法。

def sign(x):
    y = ones(x.shape)
    y[x<0] = -1
    return y

即使使用另一个问题(how can I make a numpy function that accepts a numpy array, an iterable, or a scalar?)的想法，当 x 是单个数字时，下一个函数也不会工作，因为在这种情况下 x.shape 和 y.shape 只是 () 并且索引 y 是非法的。

def sign(x):
    x = asarray(x)
    y = ones(x.shape)
    y[x<0] = -1
    return y

唯一的解决方案似乎是首先确定 x 是数组还是数字，但我想知道是否有更好的方法。如果您有很多像这样的小函数，编写分支代码会很麻烦。

Answer 1

np.vectorize 可用于实现此目的，但会很慢，因为当使用数组调用装饰函数时，它所做的一切都是循环遍历数组元素并应用标量函数对每个，即不利用 numpy 的速度。

我发现对涉及 if-else 的函数进行矢量化很有用的方法是使用 np.choose :

def sign_non_zero(x):
    return np.choose(
        x > 0,  # bool values, used as indices to the array
        [
            -1, # index=0=False, i.e. x<=0
            1,  # index=1=True, i.e. x>0
        ])

这在 x 是标量或数组时有效，并且比在 python 空间中循环更快。

使用 np.choose 的唯一缺点是，以这种方式编写 if-else 逻辑并不直观，而且代码可读性较差。每当我使用它时，我都会包含上面的注释，以使读者更容易理解正在发生的事情。