如何编写既兼容纯Python值又兼容NumPy数组的条件代码？

Question

对于在Python中编写“分段函数”，我通常使用if（以控制流或三元运算符形式）。在

def spam(x):
    return x+1 if x>=0 else 1/(1-x)

现在，为了避免价值观的单一价值观而工作。因此，我认为这样的事情是最好的：_{正如Leon所说，以下是错误的}

^{pr2}$

（如果我漏掉了什么，请纠正我，因为坦率地说，我觉得这很难看。）

当然，如果x实际上是一个NumPy数组，eggs只会起作用，因为否则{}只会生成一个布尔值，不能用于索引（至少做得不对）。在

有没有一种好方法可以编写看起来更像spam但在Numpy数组上使用惯用的代码，还是应该只使用vectorize(spam)？在

Answer 1

如果我想同时处理数字数组和numpy数组，我会使用numpy.asarray（如果参数已经是numpy数组，这是一个no-op）

def eggs(x):
    x = np.asfarray(x)
    m = x>=0
    x[m] = x[m] + 1
    x[~m] = 1 / (1 - x[~m])
    return x

（这里我使用asfarray来强制浮点类型，因为您的函数需要浮点计算）。在

对于单个输入，这比垃圾邮件函数的效率要低，而且可以说更丑。然而，这似乎是最容易的选择。在

EDIT：如果您想确保x不被修改（正如Leon所指出的），那么可以将np.asfarray(x)替换为np.array(x, dtype=np.float64)，默认情况下数组构造函数会复制。在

Answer 2

使用^{}。但是，即使是纯数字输入，也会得到一个数组作为输出。在

def eggs(x):
    y = np.asarray(x)
    return np.where(y>=0, y+1, 1/(1-y))

这对数组和普通数字都有效：

^{pr2}$

正如ayhan所说，这引发了一个警告，因为1/(1-x)将被计算为整个范围。但警告就是：警告。如果你知道你在做什么，你可以忽略这个警告。在本例中，您只需从索引中选择1/(1-x)，因为它永远不会是{}，所以您是安全的。在