Numpy: Most computationally efficient way to get the mean of slices along an axis where the slices i-CodePudding

For a 2D array, I would like to get the average of a particular slice in each row, where the slice indices are defined in the last two columns of each row.

Example:

sample = np.array([
    [ 0,  1,  2,  3,  4,  2,  5],
    [ 5,  6,  7,  8,  9,  0,  3],
    [10, 11, 12, 13, 14,  1,  4],
    [15, 16, 17, 18, 19,  3,  5],
    [20, 21, 22, 23, 24,  2,  4]
])

So for row 1, I would like to get sample[0][2:5].mean(), row 2 I would like to get sample[0][0:3].mean(), row 3 sample[0][1:4].mean(), etc.

I came up with a way using apply_along_axis

def average_slice(x):
    return x[x[-2]:x[-1]].mean()

np.apply_along_axis(average_slice, 1, sample)```

array([ 3. ,  6. , 12. , 18.5, 22.5])

However, 'apply_along_axis' seems to be very slow.

numpy np.apply_along_axis function speed up?

From from source code, it seems that there are conversions to lists and direct looping, though I don't have a full comprehension on this code

https://github.com/numpy/numpy/blob/v1.22.0/numpy/lib/shape_base.py#L267-L414

I am wondering if there is a more computationally efficient solution than the one I came up with.

CodePudding user response：

Bit of hacky, but one way using numpy.cumsum about 200x faster:

def faster(arr):
    ind = arr[:, -2:]
    padded = np.pad(arr.cumsum(axis=1), ((0, 0), (1, 0)))
    res = np.diff(np.take_along_axis(padded, ind, axis=1))/np.diff(ind)
    return res.ravel()

faster(sample)

Output:

array([ 3. ,  6. , 12. , 18.5, 22.5])

Benchmark:

large = sample[np.random.randint(0, 5, 10000)]

%timeit np.apply_along_axis(average_slice, 1, large)
# 47 ms ± 166 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit faster(large)
# 305 µs ± 2.36 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

Validation:

np.array_equal(faster(large), np.apply_along_axis(average_slice, 1, large))
# True