I have a numpy array r and I need to evaluate a scalar function, let's say np.sqrt(1-x**2) on each element x of this array. However, I want to return the value of the function as zero, whenever x>1, and the value of the function on x otherwise.
The final result should be a numpy array of scalars.
How could I write this the most pythonic way?
You can use like numpy.where(condition,if condition holds,otherwise) so np.where(x>1,0,np.sqrt(1-x**2)) will be answer
With the plain evaluation:
In [19]: f=lambda x:np.sqrt(1-x**2)
In [20]: f(np.linspace(0,2,10))
C:\Users\paul\AppData\Local\Temp\ipykernel_1604\1368662409.py:1: RuntimeWarning: invalid value encountered in sqrt
f=lambda x:np.sqrt(1-x**2)
Out[20]:
array([1. , 0.97499604, 0.89580642, 0.74535599, 0.45812285,
nan, nan, nan, nan, nan])
This produces nan and a warning for x>1. np.where as suggested by other answers can change the nan to 0, but you'll still get the warning. The warning can be silenced. But an alternative is to use the where parameter of np.sqrt (and other ufunc):
In [21]: f=lambda x:np.sqrt(1-x**2, where=x<=1, out=np.zeros_like(x))
In [22]: f(np.linspace(0,2,10))
Out[22]:
array([1. , 0.97499604, 0.89580642, 0.74535599, 0.45812285,
0. , 0. , 0. , 0. , 0. ])
where=(x<=1), since it does not work on my machine (python version 3.8.16, numpy version 1.21.6)y = np.where(
r>1, #if r[i]>1:
0, #y[i]=0
np.sqrt(1-r**2) #else: y[i] = (1-r[i]**2)
)
np.where. See examples