CODEWITHSUNDEEP
numpy
IntegrateThis
IntegrateThis

Reputation: 962

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

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.

Upvotes: 0

Views: 68

Answers (2)

Cain&#227; Max Couto da Silva
Cain&#227; Max Couto da Silva

Reputation: 4929

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!

Upvotes: 1

lezaf
lezaf

Reputation: 512

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.

Upvotes: 1

Related Questions