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

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.

The entire array is passed, because you supply the entire array as an input argument. One solution is for your function to go over the input array and return another array of the same size. — Jussi Nurminen
– Jussi Nurminen, Commented Nov 19, 2020 at 21:29

Cainã Max Couto da Silva · Accepted Answer · 2020-11-19 22:49:28Z

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!

edited Nov 19, 2020 at 22:49

answered Nov 19, 2020 at 21:44

Cainã Max Couto da Silva

4,9691 gold badge15 silver badges39 bronze badges

Sign up to request clarification or add additional context in comments.

4 Comments

IntegrateThis Over a year ago

Although now that I try this it seems it is just as slow as simply iterating through the array and applying func. I wonder if there are other methods/resources to apply a function efficiently to each element of an array.

Cainã Max Couto da Silva Over a year ago

Hey, which one of the code chunk worked for you? The first one with list comprehension or the last one with np.where? Just to know your goal...

IntegrateThis Over a year ago

The np.vectorize followed by a list comprehension. Well it all "worked", it's just a question of speed.

Cainã Max Couto da Silva Over a year ago

I have edited the answer. I think the best way you to solve your problem is by working directly within your function. Let me know if it works. If no, maybe it's better to remove this answer as accepted so that someone else may come with a better solution.

lezaf · Accepted Answer · 2020-11-19 21:33:10Z

1

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.

answered Nov 19, 2020 at 21:33

lezaf

5122 silver badges11 bronze badges

Collectives™ on Stack Overflow

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

2 Answers 2

4 Comments

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

4 Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Related