Search code examples
pythonnumpynansqrt

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

I am trying to implement a calculation from a research paper. In that calculation, the value of a function is supposed to be

  • 0, for x<a
  • sqrt(x-a)*SOMETHING_ELSE, for x>=a

In my module, x and a are 1D numpy-arrays (of the same length). In my first attempt I implemented the function as

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

But for x<a, np.sqrt() returns NaN and even though the heaviside function returns 0 in that case, in Python 0*NaN = NaN.

I could also replace all NaN in my resulting array with 0s afterwards but that would lead to warning outputs from numpy.sqrt() used on negative values that I would need to supress. Another solution is to treat the argument of the squareroot as an imaginary number by adding 0j and taking the real part afterwards:

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

But I feel like both solutions are not really elegant and the second solution is unnecessarily complicated to read. Is there a more elegant way to do this in Python that I am missing here?

Solution

  • You can cheat with np.maximum in this case to not compute the square root of negative numbers. Moreover, please note that np.heaviside does not use a as a threshold but 0 (the second parameter is the output of the heaviside in some case). You can use np.where instead.

    Here is an example:

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

    Note that in this specific case, the expression can be simplified and np.where is not even needed (because np.sqrt(np.maximum(x-a, 0)) gives 0). Thus, you can simply write:

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