Question

I’m playing with NumPy and Scipy and I’m having trouble finding a feature in the documentation. I was thus wondering if anyone could help.

Suppose I have an array in NumPy with two columns and k rows. One column serves as an numerical indicator (e.g. 2 = male, 1 = female, 0 = unknown) while the second column is perhaps a list of values or scores.

Lets say that I want to find the standard deviation (could be mean or whatever, I just want to apply a function) of the values for all rows with indicator 0, and then for 1, and finally, 2.

Is there a predefined function to composite this for me?

In R, the equivalent can be found in the plyr package. Does NumPy and/or Scipy have an equivalent, or am I stuck creating a mask for this array and then somehow filtering through this mask and then applying my function?

As always, thanks for your help!

Answer 1

If I understand your description, you have a dataset something like this:

In [79]: x=np.random.randint(0,3,size=100)

In [80]: y=np.random.randint(0,100,size=100)

In [81]: d=np.vstack([x,y]).T

In [88]: print d[:5,:]
[[ 0 43]
 [ 1 60]
 [ 2 60]
 [ 1  4]
 [ 0 30]]

In this situation numpy.unique can be used to generate an array of unique “key” values:

In [82]: idx=np.unique(d[:,0])

In [83]: print idx
[0 1 2]

and those values used to drive a generator expression like this:

[113]: g=(d[np.where(d[:,0]==val),1].std() for val in idx)

The generator g will emit the standard deviation of all the entries in d which match each entry in the index. numpy.fromiterator can then be used to collect the results:

In [114]: print np.vstack([idx,np.fromiter(g,dtype=np.float)]).T
[[  0.          26.87376385]
 [  1.          29.41046084]
 [  2.          24.2477246 ]]

Note there is conversion of the keys to floating point in the last step during stacking, you might not want that depending on your data, but I did just it for illustrative purposes to have a “nice” looking final result to post.