search

scikit-hep / awkward-0.x

Manipulate arrays of complex data structures as easily as Numpy.
BSD 3-Clause "New" or "Revised" License
215 stars 39 forks source link

Equivalent of numpy.where for awkward arrays #84

Open jpivarski opened 5 years ago

jpivarski commented 5 years ago

numpy.where is not a ufunc, but it's useful!

nsmith- commented 5 years ago

The non-ternary usage is equivalent to nonzero, e.g. np.where(array<1.) and (array < 1.).nonzero() So that would be nice to also implement, as it is not currently present in JaggedArray. A (admittedly concise) workaround is array.index[array<1.]

nsmith- commented 5 years ago

The ternary usage is necessary since fancy indexing doesn't work for __setitem__, e.g.

array[array<0] = 0

is supported in numpy but not awkward because of immutability.

array = awkward.where(array<0, 0, array)

could broadcast as appropriate.

jpivarski commented 4 years ago

You know, if the type is a (possibly jagged) number, like Clemen's problem on https://stackoverflow.com/a/58345579/1623645 , awkward.where could be implemented like this:

def where(predicate, iftrue, iffalse):
    predicate = predicate.astype(numpy.bool)   # just to make sure they're 0/1
    return predicate*iftrue + (1 - predicate)*iffalse