Python np.sqrt(x-a)*np.heaviside(x,a)

Question

我正在尝试根据一篇研究论文进行计算。在该计算中，函数的值应该是

• 0，对于 x
• sqrt(x-a)*SOMETHING_ELSE，对于 x>=a

在我的模块中，x 和 a 是一维 numpy 数组(长度相同)。在我的第一次尝试中，我将函数实现为

f = np.sqrt(x-a)*SOMETHING*np.heaviside(x,a)

但对于 x 0*NaN = NaN 中也是如此。

之后我也可以将结果数组中的所有 NaN 替换为 0，但这会导致 numpy.sqrt() 的警告输出用于我需要抑制的负值。另一种解决方案是通过添加 0j 并在之后取实部，将平方根的参数视为虚数:

f = np.real(np.sqrt(x-a+0j)*SOMETHING*np.heaviside(x,a))

但我觉得这两种解决方案都不是很优雅，而且第二种解决方案读起来不必要地复杂。有没有我在这里缺少的更优雅的 Python 方法？

Answer 1

在这种情况下，您可以使用 np.maximum 作弊，以不计算负数的平方根。此外，请注意，np.heaviside 不使用 a 作为阈值，而是使用 0(第二个参数在某些情况下是 heaviside 的输出)。您可以改用 np.where。

这是一个例子:

f = np.where(x<a, 0, np.sqrt(np.maximum(x-a, 0))*SOMETHING)

请注意，在这种特定情况下，可以简化表达式，甚至不需要 np.where(因为 np.sqrt(np.maximum(x-a, 0)) 给出 0)。因此，您可以简单地写:

f = np.sqrt(np.maximum(x-a, 0))*SOMETHING