Getting only those values that fulfill a condition in a numpy array
Question:
There must a be a (very) quick and efficient way to get only elements from a numpy array, or even more interestingly from a slice of it.
Suppose I have a numpy array:
import numpy as np
a = np.arange(-10,10)
Now if I have a list:
s = [9, 12, 13, 14]
I can select elements from a:
a[s] #array([-1, 2, 3, 4])
How can I have an (numpy) array made of the elements from a[s] that fulfill a condition, i.e. are positive (or negative)?
It should result
np.ifcondition(a[s]>0, a[s]) #array([2, 3, 4])
It looks trivial but I was not able to find a simple and condensed expression. I’m sure masks do but it’s doesn’t look really direct to me.
However, neither:
a[a[s]>0]
a[s[a[s]>0]]
are in fact good choices.
Answers:
How about:
In [19]: b = a[s]
In [20]: b[b > 0]
Out[20]: array([2, 3, 4])
You should definitely accept unutbu’s answer, and it is what I generally use for this kind of situation within numpy. But in the interests of having multiple ways of doing things, having a method that works outside of numpy, or in case the intermediate array is offensively huge, I’ll add this alternative:
In [3]: [a[S] for S in s if a[S]>0]
Out[3]: [2, 3, 4]
Again, unutbu’s method is significantly faster. But I like this method because it can generalize still further. If you have an expensive funky function (e.g., not indexing), and want to test the result of that function, you might want to do this:
In [5]: [f for S in s for f in [FunkyFunction(a[S])] if f>0]
Out[5]: [2, 3, 4]
The weird part about this is that you make a list inside the other list, but this internal list only contains one item. Basically what you’re doing is saving the value to the variable f, and then using that value twice: once to test the value (f>0), and once to use that value in the list if the test passes.
I’m aware that this is an old question, and I’m not sure if the following was available at the time it was asked, but since I got here with a similar question, this is what I found that worked for me:
np.extract(a[s] > 0, a[s])
There must a be a (very) quick and efficient way to get only elements from a numpy array, or even more interestingly from a slice of it.
Suppose I have a numpy array:
import numpy as np
a = np.arange(-10,10)
Now if I have a list:
s = [9, 12, 13, 14]
I can select elements from a:
a[s] #array([-1, 2, 3, 4])
How can I have an (numpy) array made of the elements from a[s] that fulfill a condition, i.e. are positive (or negative)?
It should result
np.ifcondition(a[s]>0, a[s]) #array([2, 3, 4])
It looks trivial but I was not able to find a simple and condensed expression. I’m sure masks do but it’s doesn’t look really direct to me.
However, neither:
a[a[s]>0]
a[s[a[s]>0]]
are in fact good choices.
How about:
In [19]: b = a[s]
In [20]: b[b > 0]
Out[20]: array([2, 3, 4])
You should definitely accept unutbu’s answer, and it is what I generally use for this kind of situation within numpy. But in the interests of having multiple ways of doing things, having a method that works outside of numpy, or in case the intermediate array is offensively huge, I’ll add this alternative:
In [3]: [a[S] for S in s if a[S]>0]
Out[3]: [2, 3, 4]
Again, unutbu’s method is significantly faster. But I like this method because it can generalize still further. If you have an expensive funky function (e.g., not indexing), and want to test the result of that function, you might want to do this:
In [5]: [f for S in s for f in [FunkyFunction(a[S])] if f>0]
Out[5]: [2, 3, 4]
The weird part about this is that you make a list inside the other list, but this internal list only contains one item. Basically what you’re doing is saving the value to the variable f, and then using that value twice: once to test the value (f>0), and once to use that value in the list if the test passes.
I’m aware that this is an old question, and I’m not sure if the following was available at the time it was asked, but since I got here with a similar question, this is what I found that worked for me:
np.extract(a[s] > 0, a[s])