Question

here we go again:

I would like a function much like filter, but it ought to return arguments rather than values:

data = numpy.array([0.0, 35., 2., 44., numpy.pi, numpy.sqrt(2.)])
args_pass_A = some_sort_of_f(lambda x: x<4., data)
print(args_pass_A)
[0, 2, 4, 5]
args_pass_B = some_sort_of_f(lambda x: x>=44., data)
print(args_pass_B)
[3]

I tried to lookup in scipy.stats and scipy.stats.mstats, where I’ve been recently directed towards mquantiles (thank you aganders3). Does this ring any bell? I also tried to lookup in numpy documentation, but with no luck.

Thank you in advance.

Answer 1

For this case I’d just use where [docs]:

>>> import numpy
>>> data = numpy.array([0.0, 35., 2., 44., numpy.pi, numpy.sqrt(2.)])
>>> data
array([  0.        ,  35.        ,   2.        ,  44.        ,
         3.14159265,   1.41421356])
>>> numpy.where(data < 4)
(array([0, 2, 4, 5]),)
>>> numpy.where(data > 44.)
(array([], dtype=int64),)
>>> numpy.where(data >= 44.)
(array([3]),)

Note that what you’re looking for is really the first element of the returned tuple:

>>> numpy.where(data < 4)[0]
array([0, 2, 4, 5])

and you can use these indices to index data again:

>>> data[numpy.where(data < 4)[0]]
array([ 0.        ,  2.        ,  3.14159265,  1.41421356])