CODEWITHSUNDEEP
pythonnumpylist-comprehension
FrenchKheldar
FrenchKheldar

Reputation: 485

Elegant list comprehension to extract values in one dimension of an array based on values in another dimension

I'm looking for an elegant solution to this:

data = np.loadtxt(file)
# data[:,0] is a time
# data[:,1] is what I want to extract
mean = 0.0
count = 0
for n in xrange(np.size(data[:,0])):
    if data[n,0] >= tstart and data[n,0] <= tend:
        mean = mean + data[n,1]
        count = count + 1

mean = mean / float(count)

I'm guessing I could alternatively first extract my 2D array and then apply np.mean on it but I feel like there could be some list comprehension goodness to make this more elegant (I come from a FORTRAN background...). I was thinking something like (obviously wrong since i would not be an index):

np.mean([x for x in data[i,1] for i in data[:,0] if i >= tstart and i <= tend])

Upvotes: 1

Views: 983

Answers (1)

DSM
DSM

Reputation: 353209

In numpy, rather than listcomps you can use lists and arrays for indexing purposes. To be specific, say we have a 2D array like the one you're working with:

>>> import numpy as np
>>> data = np.arange(20).reshape(10, 2)
>>> data
array([[ 0,  1],
       [ 2,  3],
       [ 4,  5],
       [ 6,  7],
       [ 8,  9],
       [10, 11],
       [12, 13],
       [14, 15],
       [16, 17],
       [18, 19]])

We can get the first column:

>>> ts = data[:,0] 
>>> ts
array([ 0,  2,  4,  6,  8, 10, 12, 14, 16, 18])

And create a boolean array corresponding to the terms we want:

>>> (ts >= 2) & (ts <= 6)
array([False,  True,  True,  True, False, False, False, False, False, False], dtype=bool)

Then we can use this to select elements of the column we're interested in:

>>> data[:,1][(ts >= 2) & (ts <= 6)]
array([3, 5, 7])

and finally take its mean:

>>> np.mean(data[:,1][(ts >= 2) & (ts <= 6)])
5.0

Or, in one line:

>>> np.mean(data[:,1][(data[:,0] >= 2) & (data[:,0] <= 6)])
5.0

[Edit: data[:,1][(data[:,0] >= 2) & (data[:,0] <= 6)].mean() will work too; I always forget you can use methods.]

Upvotes: 4

Related Questions