Using Numpy.where() with a function on each element - numpy

Question

I have a rather complicated function, say:
def func(elem):
// blah blah blah
return True
// blah blah blah
return False
I wish to use the numpy.where() function along the lines of
arr2 = np.where(func(arr1), arr1, 0)
But when I try this syntax and debug, I see that in func, that the entire array is passed rather than individual elements. The typical use cases I see in the documentation/examples only rely on simple comparators like arr < 5, but I need something quite a bit fancier that I don't want to try and write in a single line.
If this is possible, or if there is some vectorized substitute (emphasis on efficiency), any insights appreciated.

Answer 1

I figured out how to do it by using np.vectorize, followed by a list comprehension, not np.where. Maybe from this, one can find out a way to use numpy rather than of a list comprehension.
func_vec = np.vectorize(func)
[arr1 if cond else 0 for cond in func_vec(arr1)]
Anyways, by using func_vec(arr1) you get the True/False values per element.
Note: If you want a new array like arr1, replacing by 0 the elements that return False in your function, then this should work:
arr2 = np.where(func_vec(arr1), arr1, 0)
Edit:
Indeed, np.vectorize is not optimized to performance (by bad), is essentially a for loop under the hood. So I would recommend trying to write your function in a vectorized way rather than trying to vectorize it thereafter.
For example, try converting a function like this:
def func(elem):
if elem > 5:
return True
else:
return False
to something like this:
def func(elem):
return elem > 5
so that you can easily apply func(arr1) without error.
If you really have a function that returns just True or False, I'm pretty sure you can do it, regardless of its complexity. Anyways we're here to help you out!

Answer 2

It seems like you try to get the elements of arr1 you want using func function, but judging from the definition func works for a single element. You need a True/False array of the same shape as arr1 to do so.
If I get it right, a potential solution would be to modify func to operate on the whole array and not only on one element and return the True/False array of shape arr1.shape you need for np.where, since you want to do it in a single line that way.