Суммирование по массиву с модулями

Question

Рассмотрим следующую настройку:

import numpy as np
import itertools as it
A = np.random.rand(3,3,3,16,3,3,3,16)  # sum elements of A to arrive at...
B = np.zeros((4,4))  # a 4x4 array (output)

У меня есть большой массив 'A', который я хочу суммировать, но очень специфическим способом. «A» имеет форму (x, x, x, 16, x, x, x, 16), где «x» является некоторым целым числом. Желаемый результат - матрица 4x4 «B», которую я могу вычислить с помощью цикла for следующим образом:

%%timeit
for x1,y1,z1,s1 in it.product(range(3), range(3), range(3), range(16)):
    for x2,y2,z2,s2 in it.product(range(3), range(3), range(3), range(16)):
        B[s1%4, s2%4] += A[x1,y1,z1,s1,x2,y2,z2,s2]

>> 134 ms ± 1.27 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

, где элементы B - это «модуль-4» двух осей с 16 элементами в этом измерении в «A», здесь индексированные s1 и s2.

Как я могу добиться того же самого путем трансляции или иным образом? Очевидно, что при больших значениях «x» (размеры в «A») цикл for будет экспоненциально длиннее для вычисления, что не является идеальным.

EDIT:

C = np.zeros((4,4))
for i,j in it.product(range(4), range(4)):
    C[i,j] = A[:,:,:,i::4,:,:,:,j::4].sum()

Кажется, это тоже работает. Но все еще включает 1 цикл. Есть ли способ сделать это быстрее?

Answer 1

Вот более чистое и быстрое решение. К сожалению, они не одинаковы ...

def clean(A):
    return A.reshape(4*n*n*n, 4, 4*n*n*n, 4).sum(axis=(0, 2))

def fast(A):
    return np.bincount(np.tile(np.arange(16).reshape(4, 4), (4, 4)).ravel(), A.sum((0,1,2,4,5,6)).ravel(), minlength=16).reshape(4, 4)

При n==6 fast примерно в три раза быстрее.