numpy.where

[Neil Crighton]

Carl Banks <pavlovevidence <at> gmail.com> writes:

> > >>> condition = (min_time <= time) & (time <= max_time)
> > >>> new_time = time[condition]
> > >>> new_energy = energy[condition]
> 
> Won't work: condition is an array of ones and zeros, but you need to
> index the arrays with indices.  So, add a call to nonzero to get the
> indices of the elements.
> 
> elements = nonzero(logical_and(min_time<=time,max_time>=time))
> 
> new_time = time[elements]
> new_energy = energy[elements]
> 

It will work - try it :)

You can index numpy arrays with either an array of indices (dtype integer)
,or a boolean array (dtype bool). In the case above condition is a boolean
array. In general using a boolean array is faster, because you skip the call
to nonzero().

I find indexing with boolean arrays mindbogglingly useful, and one of the
great things about Numpy. You can do things like this (age, height and 
people are all arrays of the same length):

c0 = age < 15
c1 = height > 1.5
c2 = height < 1.8
people_subset1 =  people[c0 & c1]
people_subset2 =  people[(c1 & c2) | c0 ]
people_subset3 =  people[(c1 & c2) & ~c0 ]

It gets even more intuitive if you use record arrays (aka structured 
arrays).


Neil