Welcome, guest | Sign In | My Account | Store | Cart

This is a groupby function for arrays. Given a list of arrays and a `key` function, it will group each array based on the value of `key(args[0])`. The returned arrays will be two dimensional. The size of the first dimension is equal to the number of groups, and the size of the second dimension is equal to the size of the largest group. All of the smaller groups are padded with the value of the keyword argument `fill_value`.

There's also a short recipe in here for functional composition.

Python, 52 lines
 ``` 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52``` ```from numpy import array, vectorize, unique1d, ones from operator import itemgetter from itertools import imap, groupby, izip # functional composition def compose(*args): def composed(arg): for f in reversed(args): arg = f(arg) return arg return composed def agroupby(*args, **kwds): """A groupby function which accepts and returns arrays. All passed arrays are expected to be one dimensional and of the same shape. All of the arrays are grouped by `key(arg[0])` and then returned. The returned arrays will be two dimensional with each row corresponding to a group. The size of the first dimension is equal to the number of groups, and the size of the second dimension is equal the the size of the largest groups. All smaller groups are padded with the value of the keyword argument `fill_value`.""" keyfunc = kwds.get('key', lambda a: a) fill_val = kwds.get('fill_value', 0.0) args = [a.copy() for a in args] argsort = sorted(enumerate(args[0]), key=compose(keyfunc,itemgetter(1))) indexsort = [index for index, item in argsort] args = [a.take(indexsort) for a in args] # calculate groups g_mask = keyfunc(args[0]) g_set = unique1d(g_mask) g_max = max([g_mask[g_mask==g].shape[0] for g in g_set]) g_args = [fill_val * ones((len(g_set), g_max), dtype=a.dtype) for a in args] for gix, gval in enumerate(g_set): for ga, a in izip(g_args, args): b = a[g_mask==gval] ga[gix,:len(b)] = b return tuple(g_args) if __name__ == "__main__": from numpy import arange, set_printoptions, random set_printoptions(precision=2, suppress=True, linewidth=60); b = arange(100, 200) c = agroupby(b, key=lambda x: x%10) print c a = random.geometric(0.01, 20) b = a + 20 c, d = agroupby(a, b, key=lambda x: x%10) print c print d ```

I wrote this to group an array of values by the dates on which the values were recorded. So, if `dates` were an array of `datetime` instances, and `vals` were an array of values recorded on each of those dates, you could group `dates` and `vals` by the month in which they were recorded by calling:

agroupby(dates, vals, key=lambda dt: dt.month)

 Created by Alexander Ross on Tue, 11 Jul 2006 (PSF)