I'm attempting to use numpy.select to conditionally assign values to an array. However, use of select requires calculation of all possible assignments rather than just the relevant assignment, which can cause div/0 errors, e.g:
import numpy as np
def testfunc(z):
conditionlist = [z < 0, z == 0, z > 0]
choicelist = [1 / z, 0, 1 + z]
return np.select(conditionlist, choicelist)
if __name__ == "__main__":
print testfunc(np.array([0]))
This code will fail with a div/0 error, although 1 / z where z = 0 never actually needs to be assigned to the returned array.
How can I assign values to a numpy array conditionally without running into div/0 errors? Is a loop the only option?
np.select([z < 0, z == 0, z > 0], [1 / (z + (z == 0)), 0, 1 + z])
z == 0 gives an array of booleans. Adding this to z gives an array without zero elements that is equal to z for the indices that np.select will use.