Assigning numpy array based on condition

Question

What is a better way to write this numpy python code?

age[age < 20.0] = 0.0
age[age > 0.0] = 1.0
mature = age

Here, mature contains 1.0 for all values of age > 20.0, else 0.0

Answer 1

mature = age = (age > 20.0).astype(float)

age > 20.0 is a boolean array. The astype(float) converts the array to float dtype, which changes True to 1.0 and False to 0.0. Note that this also converts NaNs to 0.

---

To preserve NaNs, like your original code, you could use np.clip :

mature = age = np.clip(age-20, 0, 1)

---

For example,

In [90]: age = np.array([np.nan, 30, 20, 10])

In [91]: (age > 20.0).astype(float)
Out[91]: array([ 0.,  1.,  0.,  0.])

In [92]: np.clip(age-20, 0, 1)
Out[92]: array([ nan,   1.,   0.,   0.])