sum uneven segments of an array in numpy

Question

Given an ndarray x and a one dimensional array containing the length of contiguous slices of a dimension of x, I want to compute a new array that contains the sum of all of the slices. For example, in two dimensions summing over dimension one:

>>> lens = np.array([1, 3, 2])
array([1, 3, 2])
>>> x = np.arange(4 * lens.sum()).reshape((4, lens.sum())).astype(float)
array([[  0.,   1.,   2.,   3.,   4.,   5.],
       [  6.,   7.,   8.,   9.,  10.,  11.],
       [ 12.,  13.,  14.,  15.,  16.,  17.],
       [ 18.,  19.,  20.,  21.,  22.,  23.]])
# I want to compute:
>>> result
array([[  0.,   6.,   9.],
       [  6.,  24.,  21.],
       [ 12.,  42.,  33.],
       [ 18.,  60.,  45.]])
# 0 = 0
# 6 = 1 + 2 + 3
# ...
# 45 = 22 + 23

The two ways that come to mind are:

a) Use cumsum and fancy indexing:

def cumsum_method(x, lens):
    xc = x.cumsum(1)
    lc = lens.cumsum() - 1
    res = xc[:, lc]
    res[:, 1:] -= xc[:, lc[:-1]]
    return res

b) Use bincount and intelligently generate the appropriate bins:

def bincount_method(x, lens):
    bins = np.arange(lens.size).repeat(lens) + \
        np.arange(x.shape[0])[:, None] * lens.size
    return np.bincount(bins.flat, weights=x.flat).reshape((-1, lens.size))

Timing these two on large input had the cumsum method performing slightly better:

>>> lens = np.random.randint(1, 100, 100)
>>> x = np.random.random((100000, lens.sum()))
>>> %timeit cumsum_method(x, lens)
1 loops, best of 3: 3 s per loop
>>> %timeit bincount_method(x, lens)
1 loops, best of 3: 3.9 s per loop

Is there an obviously more efficient way that I'm missing? It seems like a native c call would be faster because it wouldn't require allocating the cumsum or the bins array. A numpy builtin function that does something close to this could likely be better than (a) or (b). I couldn't find anything through searching and looking through the documentation.

Note, this is similar to this question, but the summation intervals aren't regular.

Answer 1

You can use np.add.reduceat:

>>> np.add.reduceat(x, [0, 1, 4], axis=1)
array([[  0.,   6.,   9.],
       [  6.,  24.,  21.],
       [ 12.,  42.,  33.],
       [ 18.,  60.,  45.]])

The list of indices [0, 1, 4] means: "sum the slices 0:1, 1:4 and 4:". You could generate these values from lens using np.hstack(([0], lens[:-1])).cumsum().

Even factoring in the calculation of the indices from lens, a reduceat method is likely to be significantly faster than alternative methods:

def reduceat_method(x, lens):
    i = np.hstack(([0], lens[:-1])).cumsum()
    return np.add.reduceat(x, i, axis=1)

lens = np.random.randint(1, 100, 100)
x = np.random.random((1000, lens.sum())

%timeit reduceat_method(x, lens)
# 100 loops, best of 3: 4.89 ms per loop

%timeit cumsum_method(x, lens)
# 10 loops, best of 3: 35.8 ms per loop

%timeit bincount_method(x, lens)
# 10 loops, best of 3: 43.6 ms per loop