Получить среднее значение 2D среза 3D массива в numpy


У меня есть массив numpy с формой:

(11L, 5L, 5L)

Я хочу вычислить среднее значение по 25 элементам каждого "среза" массива [0, :, :], [1, :, :] etc, возвращающий 11 значений.

Это кажется глупым, но я не могу придумать, как это сделать. Я думал, что функция mean(axis=x) сделает это, но я перепробовал все возможные комбинации axis, и ни одна из них не дала мне желаемого результата.

Я, очевидно, могу сделать это с помощью цикла for и нарезки, но, конечно, есть лучший способ путь?

3 11

3 ответа:

Используйте Кортеж для оси:

>>> a = np.arange(11*5*5).reshape(11,5,5)
>>> a.mean(axis=(1,2))
array([  12.,   37.,   62.,   87.,  112.,  137.,  162.,  187.,  212.,
        237.,  262.])

Edit: это работает только с numpy версии 1.7+.

Всегда может использовать np.einsum :

>>> a = np.arange(11*5*5).reshape(11,5,5)
>>> np.einsum('...ijk->...i',a)/(a.shape[-1]*a.shape[-2])
array([ 12,  37,  62,  87, 112, 137, 162, 187, 212, 237, 262])

Работает с массивами более высокой размерности (все эти методы будут работать, если метки оси будут изменены):

>>> a = np.arange(10*11*5*5).reshape(10,11,5,5)
>>> (np.einsum('...ijk->...i',a)/(a.shape[-1]*a.shape[-2])).shape
(10, 11)

Быстрее загружаться:

a = np.arange(11*5*5).reshape(11,5,5)

%timeit a.reshape(11, 25).mean(axis=1)
10000 loops, best of 3: 21.4 us per loop

%timeit a.mean(axis=(1,2))
10000 loops, best of 3: 19.4 us per loop

%timeit np.einsum('...ijk->...i',a)/(a.shape[-1]*a.shape[-2])
100000 loops, best of 3: 8.26 us per loop

Масштабируется немного лучше, чем другие методы, поскольку размер массива увеличивается.

Использование dtype=np.float64 существенно не меняет вышеуказанные тайминги, поэтому просто перепроверьте:

a = np.arange(110*50*50,dtype=np.float64).reshape(110,50,50)

%timeit a.reshape(110,2500).mean(axis=1)
1000 loops, best of 3: 307 us per loop

%timeit a.mean(axis=(1,2))
1000 loops, best of 3: 308 us per loop

%timeit np.einsum('...ijk->...i',a)/(a.shape[-1]*a.shape[-2])
10000 loops, best of 3: 145 us per loop

Также кое-что интересное:

%timeit np.sum(a) #37812362500.0
100000 loops, best of 3: 293 us per loop

%timeit np.einsum('ijk->',a) #37812362500.0
100000 loops, best of 3: 144 us per loop

Вы можете reshape(11, 25), а затем вызвать mean только один раз (быстрее):

a.reshape(11, 25).mean(axis=1)

В качестве альтернативы, вы можете вызвать np.mean дважды (примерно в 2 раза медленнее на моем компьютере):

a.mean(axis=2).mean(axis=1)