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)

但是对于 x0*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

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

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

python

numpy

nan

sqrt